A Geometric Theory for Nonequilibrium Transformer Dynamics

원문 초록

We develop a variational and geometric theory based on statistical physics for the joint evolution of Transformer attention and token representations at fixed model parameters. To this end, we introduce a free energy for the attention and model the evolution of the attention weights by its Fisher--Rao gradient flow. At fixed token representations, we prove that the softmax matrix is the unique free-energy minimiser and that every positive attention trajectory converges to it. For prescribed time-dependent representations, we show that the attention dynamics becomes linear after a change to logit variables. This result yields tracking bounds relative to the instantaneous softmax and sufficient conditions for attention to converge to a fixed matrix. Hamiltonian dynamics for the representations and their momenta are then coupled to the attention relaxation driven by the same energy. This coupling results in an energy-dissipation law. We also develop two alternative Hamiltonians incorporating the attention value map. We prove global well-posedness of the attention flow and all three coupled systems.