Recurrent Self-Attention Dynamics: An Energy-Agnostic Perspective from Jacobians

arXiv:2505.19458 · cs.LG, cond-mat.dis-nn, cs.NE, stat.ML · Submitted 2025-05-26 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Hi, I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: We're looking at the paper: "Recurrent Self-Attention Dynamics".

Jane: Here's the gist of the paper. The gist The theoretical understanding of self-attention (SA) has been steadily progressing, and this work aims to relax energy constraints and provide an energy-agnostic characterization of inference dynamics by dynamical systems analysis Relaxing Energy Constraints The paper revisits the energy-based formulation and partially relaxes traditional architectural constraints such as symmetric weights and single-head assumptions to better approximate realistic SA settings These relaxed constraints provide insights into designing regularization methods, which are explored experimentally later in Section 6.2 The work extends the analysis to a broader class of SAs by analyzing the Jacobian matrix of SA dynamics, which is more general than energy-based analysis and enables the detection of non-stationary dynamics including oscillations Jacobian Analysis and Normalization Analyzing the Jacobian matrix is valuable when investigating more general SA architectures without necessarily admitting an energy function The analysis reveals that normalization layers play an essential role in suppressing the Lipschitzness of SA and the Jacobian’s complex eigenvalues, which correspond to the oscillatory components of the dynamics Specifically, normalization effectively suppresses the Jacobian’s spectral norm (Theorem 5.1) and controls oscillatory behaviors by normalizing the complex eigenvalues of the Jacobian (Section 5.2) Lyapunov Exponents and Criticality 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 High-performance SA models exhibit a maximum Lyapunov exponent close to zero, suggesting that rich non-stationary inference dynamics emerge at the boundary between convergence and instability Quantitative Insights into Inference Dynamics The Jacobian provides an interpretation for the pseudo-energy proposed in prior work, linking it to large eigenvalues The Jacobian analysis shows that the state update of SA is more dynamic than energy-constrained cases Specifically, the Jacobian provides a meaningful and promising interpretation for the observed decrease of pseudo-energy during AKOrN inference Regularization Methods The work develops regularization methods for training and a pseudo-energy for monitoring inference dynamics Two regularization methods are proposed: 1.

Tom: Energy-based regularization, which involves adding an energy-based regularization term to the loss function during the training of ItrSA, interpreted as an approximation of energy-based SA models 20 2.

Jane: Jacobian spectral regularization, which introduces a regularization term RSpec that encourages the singular values to be close to 1, which has been shown to be beneficial for recursive architectures Test-Time Scaling and Performance The analysis investigates test-time scaling through the lens of Jacobians, showing that regularizing the spectral norm of weight matrices in SA improves performance Empirical results indicate that ItrSA also exhibits test-time scaling, and this issue can be mitigated by applying RMS normalization Furthermore, spectral regularization is particularly effective in AKOrN, substantially enhancing the performance of both AKOrN and ItrSA Jacobian-based Interpretation of Pseudo-Energy The Jacobian provides an interpretation for the observed decrease of pseudo-energy during AKOrN inference by showing that the contribution index can be effectively explained using only the attention matrix at the initialization of inference This suggests that the loop effectively performs a computation analogous to power iteration, but constrained to real space with positive eigenvalues Lyapunov Exponents in Practice The maximum Lyapunov exponent measures the exponential rate at which trajectories locally converge or diverge, and an exponent close to zero characterizes a critical regime often referred to as the edge of chaos, which correlates with high performance in neural networks The results show that SA models achieving high test accuracy empirically have Lyapunov exponents close to zero, thus operating near criticality Sensitivity to Initial Conditions The study observed behaviors consistent with widely accepted notions of chaos, particularly sensitivity to initial conditions when the maximum Lyapunov exponent gets slightly positive The L1 distance between the two trajectories increases exponentially over time, demonstrating chaotic behavior in the system In this work, the authors advanced the understanding and control of recurrent self-attention (SA) from a dynamical systems perspective They argue that normalization layers effectively regularize their spectral properties, and Lyapunov exponents at criticality characterize high-performance inference The Jacobian provides quantitative insights into performance-enhancing regularization and pseudo-energy behaviors Limitations The authors focused on recurrent SAs without positional encoding, masking, or MLP blocks, unlike looped Transformers in practice A limitation is that the upper bound in Theorem 5.1 provides a conservative upper bound that is looser than empirical observations Deriving tighter bounds remains an intriguing open theoretical problem Additionally, analytically justifying empirical phenomena, such as the Lyapunov exponent concentration and the Jacobian-based approximation, by solving state dynamics would be challenging yet exciting themes for theory Experimental Details Experiments were conducted on NVIDIA H200 GPUs, using the Adam optimizer and training for 100 epochs with batch size 100 The experiments were performed on two Sudoku datasets: SATNet for in-distribution (ID) data and RRN dataset for out-of-distribution (OOD) data The dimension of oscillators of AKOrN was set to N = 8 The learned scaling parameter γ in RMSNorm was found to remain small, with maxj γj < 1 Code and Data Availability The code and data used in the experiments are not publicly available, and the authors do not plan to release them As a result, the supplemental material does not contain instructions for reproducing the main experimental results The paper does not provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results Institutional Review Board (IRB) Approvals The paper does not involve crowdsourcing nor research with human subjects, so no IRB approvals or equivalent reviews were obtained The paper does not involve crowdsourcing nor research with human subjects Declaration of LLM Usage The core methodology of this research does not involve the use of LLMs in any important, original, or non-standard way The core methodology of this research does not involve the use of LLMs in any important, original, or non-standard way Broader Impacts Given the theoretical nature of our work, the authors assess that it does not directly engage with societal impacts The paper does not discuss both potential positive societal impacts and negative societal impacts of the work performed Safeguards Our research does not involve the release of data or models that pose high risks for misuse, hence specific safeguards are not required The paper does not pose such risks Licenses for Existing Assets Our study does not use any external assets, thus no licensing or attribution issues are applicable Our study does not use any external assets, thus no licensing or attribution issues are applicable New Assets No new assets are introduced in our paper, so there are no associated documentation requirements<ref:2505.

Paper summary: Tom: We started by looking at SA layers that have an energy function decreasing by state updates, but this paper argues that relying on those functions is too restrictive because they often assume idealized assumptions or extra constraints not present in standard self-attention.

Jane: So, the core thesis here is to relax those energy constraints and provide an energy-agnostic characterization of inference dynamics by using dynamical systems analysis instead.

Lu: They first revisited the energy-based formulation and partially relaxed traditional architectural requirements like symmetric weights and single-head assumptions just to get a better approximation of realistic self-attention settings.

Meng: What’s the main takeaway from those relaxations? Are they just making it easier to fit data, or is there some real physical insight gained?

Jane: They say these relaxed constraints give insights into designing regularization methods, which they then explore experimentally later in Section six point two of the paper <ref:2505.19458#pg2,insights into designing regularization methods, which>.

Tom: So, that means the constraints aren't just academic; they lead directly to practical ideas for how we should regularize these things during training.

Lu: Then, to study a wider class of self-attention layers more flexibly, they analyze the Jacobian matrix of the SA dynamics instead of sticking strictly to energy-based analysis.

Jane: Analyzing the Jacobian approach is more general than the energy-based analysis because it characterizes linear stability, which is Lyapunov’s indirect method, and that's what they are leaning into.

Meng: So why go for a more complex matrix analysis when you could just stick with the simpler energy function approach? What’s the advantage there?

Jane: The Jacobian approach allows them to study a broader class of SAs without necessarily admitting an energy function, which is much more flexible.

Tom: And what does analyzing that Jacobian matrix actually reveal about the dynamics of these recurrent layers? Does it show us something new about their behavior?

Jane: It reveals that the normalization layer plays an essential role in suppressing the Lipschitzness of self-attention and controlling those complex eigenvalues of the Jacobian, which correspond to the oscillatory components of the dynamics.

Lu: That means we see how normalization directly impacts how those oscillations behave in these models.

Conclusion: Tom: So wrapping up Recurrent Self-Attention Dynamics: An Energy-Agnostic Perspective from Jacobians, the authors show that this Jacobian perspective offers a way to characterize inference dynamics without relying on restrictive energy functions.

Jane: They argue that analyzing the Jacobian matrix is highly valuable for investigating more general self-attention architectures because it characterizes linear stability and reveals how normalization controls the oscillatory components of those dynamics.

Lu: The paper demonstrates that Lyapunov exponents computed from these Jacobians show that normalized dynamics lie near a critical state, which serves as a strong indicator of high inference performance.

Meng: So, in simple terms for someone who isn't reading the math, what does this mean for deploying models?

Jane: It means that high-performance self-attention models operate near a critical state where the maximum Lyapunov exponent is close to zero, which is a strong indicator of good inference performance.

Tom: The paper also gives us tools—regularization methods and a pseudo-energy—that can be used for both training and monitoring during inference.

Lu: It’s about using the Jacobian approach to give quantitative insights into performance-enhancing regularization and pseudo-energy behaviors, showing how the state update is more dynamic than energy-constrained cases.

Akiyoshi Tomihari, Ryo Karakida

Artificial Intelligence Research Center, AIST · Department of Computer Science, The University of Tokyo · RIKEN Center for Advanced Intelligence Project

cs.LG, cond-mat.dis-nn, cs.NE, stat.ML

Submitted: 2025-05-26

Updated: 2026-10-18

Comments: NeurIPS 2025. Clarified some assumptions in the propositions

Code: https://github.com/autonomousvision/akorn

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 79/100

The gist: The gist The theoretical understanding of self-attention (SA) has been steadily progressing, and this work aims to relax energy constraints and provide an energy-agnostic characterization of

Key concepts

Energy-Agnostic Characterization: The paper moves beyond traditional energy functions used in some self-attention studies. Instead, it uses dynamical systems analysis—specifically the Jacobian matrix—to understand how the model's state evolves without needing a predefined energy function. This allows for a broader characterization of inference behavior.
Jacobian Matrix Analysis: Analyzing the Jacobian matrix of the state updates is key to understanding general self-attention architectures. It reveals non-stationary dynamics, including oscillations, and shows that normalization layers are crucial for suppressing these complex eigenvalues. This analysis provides a more general tool than energy-based methods.
Lyapunov Exponents and Criticality: Lyapunov exponents measure the exponential rate at which trajectories diverge or converge in the system's dynamics. The paper finds that high-performance models operate near a critical state, where the maximum Lyapunov exponent is close to zero. This criticality is identified as a strong indicator of excellent inference performance.
Normalization Layer Role: The normalization layer plays an essential role in stabilizing self-attention dynamics. It suppresses the Lipschitzness of the system and controls oscillatory behaviors by normalizing the complex eigenvalues of the Jacobian matrix, which are responsible for these oscillations.

Terminology

Regularization — 1× in this paper · explained in 3 episodes

Summary

The gist The theoretical understanding of self-attention (SA) has been steadily progressing, and this work aims to relax energy constraints and provide an energy-agnostic characterization of inference dynamics by dynamical systems analysis

Relaxing Energy Constraints

The paper revisits the energy-based formulation and partially relaxes traditional architectural constraints such as symmetric weights and single-head assumptions to better approximate realistic SA settings These relaxed constraints provide insights into designing regularization methods, which are explored experimentally later in Section 6.2 The work extends the analysis to a broader class of SAs by analyzing the Jacobian matrix of SA dynamics, which is more general than energy-based analysis and enables the detection of non-stationary dynamics including oscillations

Jacobian Analysis and Normalization

Analyzing the Jacobian matrix is valuable when investigating more general SA architectures without necessarily admitting an energy function The analysis reveals that normalization layers play an essential role in suppressing the Lipschitzness of SA and the Jacobian’s complex eigenvalues, which correspond to the oscillatory components of the dynamics Specifically, normalization effectively suppresses the Jacobian’s spectral norm (Theorem 5.1) and controls oscillatory behaviors by normalizing the complex eigenvalues of the Jacobian (Section 5.2)

Lyapunov Exponents and Criticality

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 High-performance SA models exhibit a maximum Lyapunov exponent close to zero, suggesting that rich non-stationary inference dynamics emerge at the boundary between convergence and instability

Quantitative Insights into Inference Dynamics

The Jacobian provides an interpretation for the pseudo-energy proposed in prior work, linking it to large eigenvalues The Jacobian analysis shows that the state update of SA is more dynamic than energy-constrained cases Specifically, the Jacobian provides a meaningful and promising interpretation for the observed decrease of pseudo-energy during AKOrN inference

Regularization Methods

The work develops regularization methods for training and a pseudo-energy for monitoring inference dynamics Two regularization methods are proposed:

Energy-based regularization, which involves adding an energy-based regularization term to the loss function during the training of ItrSA, interpreted as an approximation of energy-based SA models 20
Jacobian spectral regularization, which introduces a regularization term RSpec that encourages the singular values to be close to 1, which has been shown to be beneficial for recursive architectures

Test-Time Scaling and Performance

The analysis investigates test-time scaling through the lens of Jacobians, showing that regularizing the spectral norm of weight matrices in SA improves performance Empirical results indicate that ItrSA also exhibits test-time scaling, and this issue can be mitigated by applying RMS normalization Furthermore, spectral regularization is particularly effective in AKOrN, substantially enhancing the performance of both AKOrN and ItrSA

Jacobian-based Interpretation of Pseudo-Energy

The Jacobian provides an interpretation for the observed decrease of pseudo-energy during AKOrN inference by showing that the contribution index can be effectively explained using only the attention matrix at the initialization of inference This suggests that the loop effectively performs a computation analogous to power iteration, but constrained to real space with positive eigenvalues

Lyapunov Exponents in Practice

The maximum Lyapunov exponent measures the exponential rate at which trajectories locally converge or diverge, and an exponent close to zero characterizes a critical regime often referred to as the edge of chaos, which correlates with high performance in neural networks The results show that SA models achieving high test accuracy empirically have Lyapunov exponents close to zero, thus operating near criticality

Sensitivity to Initial Conditions

The study observed behaviors consistent with widely accepted notions of chaos, particularly sensitivity to initial conditions when the maximum Lyapunov exponent gets slightly positive The L1 distance between the two trajectories increases exponentially over time, demonstrating chaotic behavior in the system

Conclusion

In this work, the authors advanced the understanding and control of recurrent self-attention (SA) from a dynamical systems perspective They argue that normalization layers effectively regularize their spectral properties, and Lyapunov exponents at criticality characterize high-performance inference The Jacobian provides quantitative insights into performance-enhancing regularization and pseudo-energy behaviors

Limitations

The authors focused on recurrent SAs without positional encoding, masking, or MLP blocks, unlike looped Transformers in practice A limitation is that the upper bound in Theorem 5.1 provides a conservative upper bound that is looser than empirical observations Deriving tighter bounds remains an intriguing open theoretical problem Additionally, analytically justifying empirical phenomena, such as the Lyapunov exponent concentration and the Jacobian-based approximation, by solving state dynamics would be challenging yet exciting themes for theory

Experimental Details

Experiments were conducted on NVIDIA H200 GPUs, using the Adam optimizer and training for 100 epochs with batch size 100 The experiments were performed on two Sudoku datasets: SATNet for in-distribution (ID) data and RRN dataset for out-of-distribution (OOD) data The dimension of oscillators of AKOrN was set to N = 8 The learned scaling parameter γ in RMSNorm was found to remain small, with maxj γj < 1

Code and Data Availability

The code and data used in the experiments are not publicly available, and the authors do not plan to release them As a result, the supplemental material does not contain instructions for reproducing the main experimental results The paper does not provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results

Institutional Review Board (IRB) Approvals

The paper does not involve crowdsourcing nor research with human subjects, so no IRB approvals or equivalent reviews were obtained The paper does not involve crowdsourcing nor research with human subjects

Declaration of LLM Usage

The core methodology of this research does not involve the use of LLMs in any important, original, or non-standard way The core methodology of this research does not involve the use of LLMs in any important, original, or non-standard way

Broader Impacts

Given the theoretical nature of our work, the authors assess that it does not directly engage with societal impacts The paper does not discuss both potential positive societal impacts and negative societal impacts of the work performed

Safeguards

Our research does not involve the release of data or models that pose high risks for misuse, hence specific safeguards are not required The paper does not pose such risks

Licenses for Existing Assets

Our study does not use any external assets, thus no licensing or attribution issues are applicable Our study does not use any external assets, thus no licensing or attribution issues are applicable

New Assets

No new assets are introduced in our paper, so there are no associated documentation requirements<ref:2505.

Improvements for AI systems

Replace standard self-attention layers in recurrent architectures with Iterative Self-Attention (ItrSA) utilizing RMSNorm, which is shown to effectively suppress the Jacobian's spectral norm and control oscillatory behaviors by normalizing complex eigenvalues of the Jacobian. This makes recurrent models more stable and less prone to signal explosion during inference.
Implement a regularization term based on Jacobian spectral regularization, RSpec, in training loops rather than relying solely on energy-based regularization. This is because Spectral regularization is particularly effective in AKOrN and Jacobian spectral regularization substantially enhances the performance of both AKOrN and ItrSA.
Design recurrent models with multi-head attention to favor more dynamic states by increasing the number of attention heads, as experiments show that as the number of attention heads increases, the Lyapunov exponents tend to increase, which implies a more dynamic state.
Use the Jacobian-based interpretation of pseudo-energy during inference to quantify exploration in eigenspaces corresponding to large eigenvalues of the initial inference states. This allows for a better understanding of why the loop effectively performs a computation analogous to power iteration, but constrained to real space with positive eigenvalues.

Abstract

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.