Recurrent Self-Attention Dynamics: An Energy-Agnostic Perspective from Jacobians
Authors: Akiyoshi Tomihari, Ryo Karakida
Organizations: Artificial Intelligence Research Center, AIST, Japan · Department of Computer Science, The University of Tokyo, Japan · RIKEN Center for Advanced Intelligence Project
The theoretical understanding of self-attention (SA) has been steadily progressing. A prominent line of work studies a class of SA layers that admit an energy function decreased by state updates. While it provides valuable insights into inherent biases in signal propagation, it often relies on idealized assumptions or additional constraints not necessarily present in standard SA. Thus, to broaden our understanding, this work aims to relax these energy constraints and provide an energy-agnostic characterization of inference dynamics by dynamical systems analysis. In more detail, we first consider relaxing the symmetry and single-head constraints traditionally required in energy-based formulations. Next, we show that analyzing the Jacobian matrix of the state is highly valuable when investigating more general SA architectures without necessarily admitting an energy function. It reveals that the normalization layer plays an essential role in suppressing the Lipschitzness of SA and the Jacobian's complex eigenvalues, which correspond to the oscillatory components of the dynamics. In addition, the Lyapunov exponents computed from the Jacobians demonstrate that the normalized dynamics lie close to a critical state, and this criticality serves as a strong indicator of high inference performance. Furthermore, the Jacobian perspective also enables us to develop regularization methods for training and a pseudo-energy for monitoring inference dynamics.
Figures & tables
Figure 1: Normalization improves the spectral norm of SA’s Jacobian.
Figure 2: Normalization layers play a crucial role in controlling the Jacobian’s eigenvalues. (a) and (c) show results on the Sudoku dataset.
Figure 3: Normalization drives the Lyapunov exponents toward zero and enables high accuracy.
Figure 4: ItrSA consistently improves accuracy as the number of loops T increases, regardless of the value of oscillator dimension N . Note that N=512 corresponds to the setting where tokens are not split into oscillators. Error bars indicate the standard deviation.
Figure 5: Energy-based regularization (“E-single” and “E-multi”) underperforms the original methods, while Jacobian spectral regularization (“Spec”) outperforms. We set H=8 except for E-single ( H=1 ). For AKOrN, we used N=4 .
Figure 6: Test-time inference of AKOrN on the Sudoku dataset.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Parameter
Sudoku
CIFAR-10
Hidden dimension D
512
384
Number of heads H
8
8
Initial value of η
1.0
1.0
Batch size
100
128
Number of epochs
100
200
Appendix
Table S.1 : Training and model configurations.
Method
Value
E-single
1e-6
E-multi
1e-4
Spec (ItrSA)
1e-4
Spec (AKOrN)
1e-5
Appendix
Table S.2 : Regularization coefficients used in Figure 5 .
Figure S.1 : Effect of the number of attention heads H on mean Lyapunov exponent and accuracy in the Sudoku dataset.
N
∥γ∥
maxj∣γj∣
4
0.229±0.000
0.229±0.000
8
0.031±0.002
0.052±0.000
16
0.098±0.006
0.098±0.006
32
0.348±0.000
0.348±0.000
64
0.489±0.001
0.489±0.001
128
0.841±0.000
0.841±0.000
Appendix
Table S.3 : γ in RMSNorm with different N .
Figure S.2 : Effect of multi-head energy on the mean Lyapunov exponent and accuracy for the Sudoku dataset. λ denotes the regularization coefficient of Eq. ( 20 ), and “Hard” indicates the hard constraint in Proposition 4.2 .
Figure S.3 : Energy-based regularization (“E-single” and “E-multi”) underperforms the original methods, while Jacobian spectral regularization (“Spec”) outperforms. We used N=8 for AKOrN.
Figure S.4 : Distribution of the Lyapunov exponent.
Model
# Parameters
Training loss
Validation loss
MLE
ItrSA
96.2M
4.31
5.67
0.196±0.016
GPT-2
124M
1.49
5.58
-
Appendix
Table S.4 : Language modeling results. MLE denotes the maximum Lyapunov exponent.
Figure S.5 : Sensitivity to initial conditions investigated using the ItrSA model.
Figure S.6 : Contribution and the quandratic form.
Figure S.7 : Lyapunov exponent on the CIFAR-10 dataset.
Standard self-attention compares tokens by similarity but does not account for how much of a past state survives to the present. We introduce Dynamical Kernel Attention (DKA), which derives attention from robust estimation of a latent state evolving under stochastic dynamics. As a past state decays toward equilibrium under the dynamics, the information it carries about the present vanishes, bounding the total evidence a head can draw. Unlike recency biases, DKA can abstain when relevant evidence is weak and guarantees a bounded error when old tokens are discarded. DKA preserves the parallel structure and computational cost of standard attention. We present the model and its theory; empirical evaluation is deferred to a subsequent version.
Transformer self-attention can be interpreted as a gradient flow on the unit sphere, in which tokens evolve under softmax interaction potentials and tend to form clusters. While prior work has established clustering behavior for single-head attention, the multi-head setting remains less understood due to geometric interference between heads, which invalidates standard monotonicity arguments. In this work, we develop a theoretical framework for multi-head self-attention dynamics and resolve several open questions. We show that, under suitable conditions on the score matrices, a natural multi-head energy functional is non-decreasing along both flat and spherical dynamics. We identify the key obstruction to per-head monotonicity as radial shadow terms, which are projections of each head's output onto token directions, persisting even under orthogonality assumptions. We introduce a sufficient condition ensuring monotonicity and establish robustness to approximate orthogonality. In a simplified scalar-head regime with equiangular token configurations, we derive a closed-form expression for the critical inverse temperature governing clustering behavior, and show that heterogeneous heads exhibit super-additive clustering rates. In this regime, we also prove a separation in clustering time between ReLU and softmax attention in the linearized dynamics. Finally, we establish an entropy production identity and show that attention entropy increases monotonically toward equilibrium as clustering progresses. Our results provide a unified perspective on the dynamics of multi-head attention and clarify the mechanisms underlying clustering and stability in transformer models.
Softmax attention maps every query--key interaction into a probability distribution, but the underlying structure remains largely unexplored. We define the energy field, the row-centered attention logit, and show that it exhibits invariant properties across models, architectures, and inputs. Two classes of invariants emerge. Mechanism-level invariants follow from the algebraic structure of softmax attention. They include a per-row zero-sum constraint, a rank bound determined by the head dimension, and spectral signatures that follow from them. Model-level regularities are not required by the mechanism, yet hold in every autoregressive language model we test, spanning several architecture families. The energy field distributes its variance over key positions without concentrating at a few. This delocalization traces to a property of the key matrix we call key incoherence. These invariants have practical consequences. The rank bound confines the energy field to a low-dimensional subspace. Key incoherence yields a per-head training monitor. All results are verified at multiple context lengths and input texts.