Recurrent Self-Attention Dynamics: An Energy-Agnostic Perspective from Jacobians
Authors: Akiyoshi Tomihari
Ryo Karakida
Venue: NeurIPS 2025 poster ()
Keywords: Self-attention
dynamical system
energy function
stability analysis
Jacobian
symmetric and antisymmetric weights
Primary area: deep_learning
TL;DR: Stability analysis reveals that appropriately controlled Jacobian's spectra lead to high inference performance

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.

Recurrent Self-Attention Dynamics: An Energy-Agnostic Perspective from Jacobians
Akiyoshi Tomihari
Ryo Karakida 1.3
1 Artificial Intelligence Research Center, AIST, Japan Department of Computer Science, The University of Tokyo, Japan RIKEN Center for Advanced Intelligence Project
tomihari0g ecc. u-tokyo.as ac.yyo..c.... ac.jp, tarakida.rro@aait. go.
Abstract
The theoretical understanding of self-attention (SA) has been steadily progress- ing. A prominent line of work studies class of SA layers that admit energy function decreased by state updates. While it provides valuable insights into in- herent biases in signal propagation, it often relies on idealized assumptions additional constraints not necessarily present in standard SA. Thus, to broaden understanding, this work aims 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 's complex cigenvalues, which correspond to the oscillatory components of the dynamics. In addition, the Lyapunov exponents computed from the Jacobians demonstrate that the normalized dynamics lic close 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.
1 Introduction
The theoretical understanding of self-attention (SA), central component of Transformer archi- tectures [Vaswani et al., 2017], has deepened recent years, including phenomena such as rank collapse [Noci et al., 2022] and expressive capacity [Yun et al., 2020, Kajitsuka and Sato, 2024, Kingma and Ba. 2025]. One major line of research formulates attention mechanisms as processes that minimize explicit implicit energy functions [Ramsauer al., 2021, Yang al., 2022, Hoover et al., 2023]. Since these energy functions serve as potential functions for gradient flows or as Lyapunov functions, they offer convergence guarantees and provide intuitive explanations for behaviors such as clustering and rank collapse in recurrent SA architectures [Geshkovski et al., 2025, 2024,2024, Bruno al., 2025]. They restrict attention dynamics to hypersphere, typically result of normalization. This facilitates theoretical analysis and can yield well-behaved functional properties [Castin et al., 2024).
2024.
However, energy-based formulations often rely on idcalized assumptions and require architectural modifications. These include imposing constraints on the weight matrices, limiting the number
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

of attention heads to one, and modeling state updates in the continuum limit. In addition, some architectures inspired by Hopfield networks replace SA with cross- coos--atenntion [Ramsauer et al., 2021] or require double softmax passes [Hoover et al., 2023, Hu et al., 2025]. Although they can be effective for exploring a frontier of new architectures, their utility still remains limited for quantitatively understanding or improving existing, realistic SA models.
In this work, we deepen the understanding of SA by extending energy-based analysis and employing a more general stability analysis from dynamical systems perspective.
First, we revisit the energy-based formulation and partially relax traditional architectural constraints, such symmetric weights and single-head assumptions. to better approximate realistic SA settings (Section 4). These relaxed constraints provide insights into designing regularization methods, which experimentally explore later in Section 6.2.
To study a broader class of SAs more flexibly, we next analyze the Jacobian matrix of SA dynamics (Section 5). The Jacobian approach is more general than the energy-based analysis (a.k.a. Lyapunov's s direct method) in the sense that it characterizes linear stability (a.k.a. Lyapunov's ' indirect method) and enables us to more easily detect non-stationary dynamics, including oscillations. We demonstrate that normalization layers, unique to discrete updates, play critical role in stabilizing dynamics. Specifically, they effectively suppress the Jacobian's spectral norm (Proposition 5.1) and control oscillatory behaviors by normalizing the complex eigenvalues of the Jacobian (Section 5.2). In addition to the understanding of the normalization role, we empirically reveal that 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.
Finally, we investigate test-time scaling (performance improvement as the number of itera- tions increases) through the lens of Jacobians. We show that regularizing the spectral norm of weight matrices in SA improves performance (Section 6.2), and that the Jacobian offers an interpretation of the pseudo-energy proposed in prior work, linking it to large eigenvalues (Section 6.3).
Thus, our work broadens the dynamical understanding of SA and highlights the usefulness of the Jacobians and the Lyapunov exponent as promising and fundamental tools for further exploration of realistic SA architectures.
2 Related work
Energy-based understanding. Geshkovski et al. [2025, 2023, 2024], Karagodin et al. [2024], Bruno et al. [2025] formulated recurrent SA dynamics as interactions among tokens ("particles")) enabling theoretical analysis of phenomena such as meta-stable clustering and rank-one collapse. Their continuous-time dynamics monotonically decrease an energy (Lyapunov) function, typically requiring constraints like single-head attention, hyperspherical token states, and symmetric weights. Yang et al. [2022] similarly interpreted Transformers as alternating minimization of energy functions, though with strieter conditions on step sizes and fixed-point proximity. Ramsauer et al. [2021] formalized the attention mechanism as modern Hopfield networks, and Hoover et al. [2023], Hu al. [2025] further developed energy functions for Transformers. We do not address approaches based on the Hopfield networks, as they require architectural modifications distinct from standard self-attention.
Jacobian-based analysis. The Jacobian of state updates is fundamental for characterizing neural network dynamics. For example, it has been used to analyze edge-of-chaos behavior for stable signal propagation and gradient control [Boedecker et al., 2012, Poole et al., 2016, Pennington et al., 2017], and to investigate discrete-time stability in dynamics with anti-symmetric matrices [Haber and Ruthotto, 2017, Chang al., 2019]. Several studies have explored Jacobian- based regularization, including for generalization (Yoshida and Miyato, 2017) [] and for continual learning [Lewandowski et al., 2025]. Regarding SA specifically, Noci et al. [2022] analyzed Jacobians to explain rank collapse, while Castin et al. [2024] evaluated their spectral properties mathematically. In this work, we use Jacobian analysis to understand inference dynamics in realistic SAs and also employ them as regularizers and performance indicators.
2

---

Looped architectures. Looped architectures in Transformers have been studied since their intro- duction by Dehghani al. [2018]. Yang al. [2024], Giannou [2023] showed that looped Transformers can learn algorithmic tasks, and Saunshi et al. [2025] further demonstrated their ef- fectiveness in enhancing reasoning via strong inductive bias. Geiping et al. [2025] increased the number of loop iterations to improve performance on reasoning benchmarks, and Bansal et al. [2022] showed that looped architectures generalize to harder problems. Miyato et al. [2025] proposed artificial Kuramoto oscillatory neurons (AKOrN), looped architecture that successfully solves tasks in neuroscience-inspired manner, demonstrating strong empirical results in unsupervised object discovery, adversarial robustness, calibrated uncertainty quantification, and reasoning. Weight tying in ALBERT [Lan et al., 2020] and fixed-pi exed-point computation in equilibrium models [Bai et al., 2019] are also interpreted looped architectures. For more detailed overview of previous work, the extended related work in Section C.
Preliminaries
Notations. For matrix A, we use the subseripts ABJ) Apa.) and A denote the (i, entry, the i-th row, and the j-th column of A, respectively. We denote the time index by x and the head index in multi-head attention by Wn. All derivatives are computed using the numerator layout.
Self-attention. Multi-head self-attention (MSA) is defined as
MSA(X) Concat(SA, ( SAn
(1)
................................
and each SA head (X) is defined as
SAn XX = (oftmax(3XXW)) W X x )
(2)
for h = 1...., H. Here, X E 5 denotes a sequence S tokens, each represented by a D- dimensional embedding. The weight matrices W
K W E Dx correspond to the query, key, and value projections for head h, and V E R DxD is the output projection matrix. Typically, the head dimension and scaling factor are Dn D/H and 1/VDn
h W.K
Self-attention with energy functions. Geshkovski et al. [2025], Karagodin et al. [2024] used continuous equations and particle interpretation of tokens model state-update dynamics of SA as:
X Projx (oftmaax(8XW 8 KT X (xW
(3)
To have an energy function, the previous work has assumed constraints such as
W W VVT or 8º W -W Ip.
depending on the analyses. [Bruno et al., 2025] further assumes an unnormalized version of softmax. Under these conditions, the SA update can decrease an energy function. That is, the dynamics evolve in way that monotonically decreases the energy, thereby ensuring the Lyapunov stability. Because these models suppose symmetric weights, we refer class of SA layers with symmetric weights and Lyapunov functions symmetric SA. Yang et al. [2022] also formalized updates of SA using symmetric matrix (Appendix C).
Spherical constraint. To facilitate theoretical analysis of SA, several studies [Geshkovski et al., 2025, Miyato et al., 2025] have introduced a spherical constraint on token vectors by enforcing that each token vector has unit norm. This constraint enables an interpretation of token interactions as dynamics of particles on a hypersphere, and plays key role controlling the Lipschitz continuity of SA [Castin et al., 2024]. There are two commonly used operators with spherical constraint: a normalization operator 1I that enforces the spherical constraint, and a projection operator Proj that projects onto the tangent space of the sphere. Given token matrix X, Y E SXD such that (xgll = 1, these operators are defined token-wisel) as:
mT = Y /..l/ (a.ll. Projx Projx(()).F) (Y)R.== (ID- (D-- Ra. (t.) Yιa--
(5)
Here, IT(Y) projects each token vector YJ..) onto the unit hypersphere. Projx (Y) projects Y(.)) orthogonally to X(t.), restricting updates to the tangent space of the sphere.
3

---

In practical Transformer architectures, the spherical normalization can be interpreted as a special case of Root Mean Square Normalization (RMSNorm) [Zhang and Sennrich, 2019]. which applied to the input matrix Y E SRD as:
RMSNorm (Y) (b.) = diaagg)) (i.)]
where 7 E Rº a trainable parameter vector. RMSNorm rescales each token to have unit nor and applies clement-wiss wise scaling using the learned parameter Y. while II can be interpreted as the special case with 7 1. As we will show in Section 5, the trainable parameter plays an important role in stabilizing Jacobians.
AKOrN. AKOrN [Miyato et al., 2025] integrates a generalized Kuramoto model into an artificial neural network by updating oscillatory neurons through looped structure. The connectivity among oscillators implemented in several ways. this work, we focus on one of their AKOrNs that uses SA. Given sequential input C E RSXD AKOrN initializes X RSDD using C. Each token vector [ (i = 1.... 5) is partitioned into N-dimensional vectors, referred to as the oscillators.
AKOrN iteratively updates states using Kuramoto layer as follows:
A (0) Omg (osc)
(x)))) Proj X (C+ (C+M)((X()))
(DSE)
X (++1) 11 (usc)
x (8) n
(8)
where 17 denotes a positive discrete step size. The Omega layer (Omg) is given by a linear trans- formation by anti-symmetric matrices and determines the rotational dynamics of oscillators. The projection operator Proj x (Y) and the normalization operator IIY) are applied independently to each oscillator. We use the notation (asc) to denote oscillator-wiz wise operations. We provide further details of AKOrN in Appendix C.
Although the existence of an energy function can be guaranteed under certain special conditions for Kuramoto models, practical implementations of the Kuramoto layer do not assume such conditions in order to achieve better performance.
Iterative self-attention. Previous studies on looped architectures [Saunshi et al., 2025, Bansal al., 2022] have shown that injecting the input C. E SxD into the loop important for achieving test-time scaling. As we show later, our theory (Proposition 5.1) indicates that the normalization layer plays a critical role in controlling the norm of the Jacobian matrix, particularly in the case of RMS normalization, which is widely used in practice, Based on these insights, we propose and investigate the following update rule, referred to as iterative self-attention (ItrSA):
( = ())( (++1) RMSNorm (x +
Since ItrSA does not involve oscillator-wise operations and the Omg layer, both of which are distinctive components of AKOrN, it is suitable as baseline to compare against energy-based SAs.
4 Energy-based analysis
As described in the previous section, energy-based symmetric SA involves three constraints: (i) the symmetry of weight matrices, (ii) single head, and (iii) a continuous-tin time limit. To accommodate more realistic SA architectures, we partially relax (i) and (ii) in the following.
Extension of weight symmetry. Symmetric SA imposes symmetric constraints both WQ KT and WY We relax these constraints the following proposition.
Proposition 4.1. Consider the continuous-tinn dynamics for single-head SA equipped with projec- tion (3). The energy function
Esingte (X) = - exp 3x R.d
WOWK' KT X
(10)
i.j
is monotonically decreasing as d.Errinte (X)dtt < 0 under the condition:
W ( K TT +w9 ///2.
(1I)

---

Multi-head energy. Although energy functions have been proposed for single-head SA, no conre sponding formulation exists for MSA, which commonly used practice. We extend the above result to the multi-head setting as follows.
Proposition 4.2. Consider the continuous-lime dynamics for mulli-head SA without projection: dx/dt = Ex= (X.. An energy function
Emula (X)
-EE cxp (3x W W KT ((a4) 2
(12)
monotonically decreasing as d.Emmttel (x)nt(()). (X) S under the condition
W -W W W w 1/2, 1/2,
Wo KT
U1,
(13)
where (12)). E DDD(((2)) (h E [1, H) satisfies the orthogonality condition U UN d ID((H)
=
Propositions 4.1 and 4.2 imply that certain structures of weight matrices desirable to ensure the existence of energy function. Specifically, W WPKT be asymmetric, whereas W should remain symmetric. the multi-head scenario, a low-rank structure in the QK product is required. This aligns with practical Transformers, as they typically exhibit a low-rank structure due to the small inner dimension (the width of W W. We refer architectures that incorporate these properties as generalized symmetric SA, and we will explore their effectiveness in our experiments (Section 6.2). The proofs are provided in Appendices A.2 and A.3.
Jacobian-based analysis
In general, energy functions are used to guarantee the convergence of dynamics to fixed points (a.k.a. Lyapunov's 's direct method) [Khalil, 2002]. While this is a concrete approach to achieving stable dynamics, the construction of energy functions is usually unsystematic, and thus it is not obvious whether we can handle more realistic SA dynamics (c.g., discrete updates with normalization). Furthermore, recent experimental results have reported non-stationary dynamics (e.g., oscillations) (Karagodin et al., 2024, Miyato et al., 2025], suggesting the need for more flexible approaches applicable richer dynamics. Thus, we to analyzing the Jacobian matrix of state updates. The Jacobian controls the Lipschitzness of the function and also naturally appears in linear stability analysis (a.k.a. Lyapunov's 's indirect method), where state updates are locally described by f(x Ax) Re f (x) + (2x, A with the Jacobian = 05 /0x.
5.1 Normalization and spectral norm
Normalization operators, which do not appear in continuous- time dynamics, essential in the dis- crete setting because discretizing state updates causes the state vector to deviate from the hypersphere. We find that the normalization operators suppress the Jacobian's 's eigenvalues.
Proposition 5.1. Suppose that, in the update of ItrSA (9), the input to the normalization layer satisfies || [. + 7/ (R. Z for all i E [1, Then, the spectral norm of the Jacobian satisfies
@ RMSSorm X + +*XX)
maxx((1ss)
(1 II (X)
(14)
0x
R
where JMSA (X) OMM(()))X denotes the Jacobian of MSA.
We show the proof in Appendix A.4. This proposition highlights the key stabilizing effect of normalization: the spectral norm is inversely proportional to R. This effect appears to be particularly significant for preventing signal explosion in looped architectures, where the same operation is repeatedly applied.
Figure 2a shows that, practical model, normalization reduces the maximum Lyapunov exponent. This exponent corresponds to a time-averaged maximum singular value of J. which is further explained and analyzed Section 5.3. This result supports the stabilizing effect of normalization implied by Proposition 5.1.
5

---

In addition, under the assumption that | 2 for some constant ε >> 0 for all i, we can further show that as 77 - oc,
i @RMSNorm(> X + mAX)
O(1).
(15)
ax
That is, the Jacobian norm remains bounded even for a large 77 (see Appendix A.4).
Jacobian eigenvalues SA. Castin et al. [2024] provide an upper bound on the maximum eigenvalue
(in the form of a Lipschitz constant) of the Jacobian of SA defined as JMSA = d MSA(X). /
Specifically, their results in Theorem 3.3 and Lemma 3.8 state that the Jacobian JMSA (X) satisfies
for input tokens X E R SxD such that | [.all I for all [].
(MMSA(X)1 S
Σ
V3 WX lay 33 W
(S
(16)
This inequality indicates that the right-hand side of Eq. (14) is bounded independently of the input X. It further implies that when the norms of the weight matri- ces [| 2 || [W ll2, or [ WTTWlll 112,, or the num- ber of tokens S becomes large, the spectral norm can also become large. This is precisely the issue that nor- malization techniques can address. Figure I demon- strates that the spectral norm of the untrained SA's Jacobian can be effectively reduced through normal- ization (see Appendix B.4 for experimental settings). Interestingly, empirically observed that the spectral norm is not only reduced by normalization but remains O(1) with respect the number of tokens. This sug- gests that the current theoretical bound (Eqs. (14) and (16) is conservative, and tighter bound remains future theoretical work.
10a
10a
101
Norm w/o Norm
Spectral
10-1
10-²
3 10
40
60
100
Number tokens
Figure 1: Normalization improves the spectral norm of SA's Jacobian.
5.2 Normalization of oscillatory components
To clarify differences between continuous and discrete updates and highlight the role of normalization, we analyze the following simplified dynamics and their associated Jacobians:
á f2x,
(xx 52
(continuous)
(17)
x (1+1) (ID + (0)
((x = ID + 7
(discrete w/o Norm) (18)
a ((+1 II[ + (0)) ((x()) Ip
yy
Ip + 7/2
(discrete w/ Norm)
(19)
|
|y]
Maximum Lyspunov exponent
InfINA
Im
Norm
w/o Norm
w/ Norm
w/o Norm
ItrSA (w/ Norm)
Norm)
0.4
10
10
10
10a
-2
Re
(a) Maximum Lyapunov exponent with and without normalization
(b) Effect of normalization on eigen-
(c) Eigenvalue distribution in the complex plane
values in oscillatory case
Figure 2: Normalization layers play a crucial role controlling the Jacobian's eigenvalues. (a) and (c) show results on the Sudoku dataset.

---

Maximum
Mean
0.5
0.5
It AKOrN
It Norm)
AKDrN
0.4
ItrSA (w/o Norm)
0.4
ItrSA tw/o Norm)
Test accuracy
Symmetric
Test accuracy
Symmetric
0.3
0.3
0.2
0.2
0.1
0.1
0.0
0.0
-Inf
-2
0
2 Inf/NA
-Inf
-2
0
2 Inf/NA
Lyapunov exponent
Lyapunov exponent
Figure 3: Normalization drives the Lyapunov exponents toward zero and enables high accuracy.
where Ω is an anti-symmetric matrix and (In + (1) Since an anti-symmetric matrix has purely imaginary eigenvalues, these dynamics represent simple oscillatory systems. The discrete dynamics with normalization can also be interpreted as isolating the Omega layer used in AKOrN.
is known that the pure imaginary cigenvalues in the continuous- limit are essential for capturing long-term signal dependencies, but can be significantly damaged by discretization [Chang et al., 2019]. Generally, in continuous-time systems t f(x). the equilibrium point is Lyapunov stable if all eigenvalues X; of the Jacobian J (x)) satisfy Re(2,) I O. In Eq. (17), the Jacobian J ((x) = Ω has purely imaginary eigenvalues and they are the boundary of stability, allowing persistent oscillations. In contrast, the equilibrium points of discrete-time systems (b+1) = f(:((1) are Lyapunov stable all cigenvalues A of the Jacobian satisfy [Ayl < 1. For Eq. (18), all eigenvalues of the Jacobian (x) ID + 1y52 take the form + imw, for wy 2 O. implying that 1 > 1. Therefore, the system becomes unstable. To avoid the fundamental instability arising from discretization, previous work architecture design inspired by dynamical systems proposed to add diffusion term to the anti-symmetric weight matrix, i.e., Ω - 7 (r > 0) [Haber and Ruthotto, 2017, Chang et al., 2019].
We find that a normalization layer (19) serves as an alternative way to mitigate this instability by effectively rescaling the system through division by |/yl]. For simplicity, suppose that all eigenvalues of Ω degenerate to the same Wj = w. After straightforward calculation, we obtain !사yl S 1 (see Appendix A.5 for the derivation). The effect of this normalization is illustrated in Figure 2b. Although this scenario represents an idealized setting, the normalization of imaginary components is empirically observed even in the case of SA, as shown in Figure 2.
5.3 Lyapunov exponent indicates criticality
The Lyapunov exponent measures the exponential rate at which trajectories locally converge or diverge in dynamical systems [Khalil, 2002] (see details Appendix B.3). Intuitively, corresponds to the time-averaged value of In aJJ),, where (J) denotes the Jacobian's 's singular values. A positive exponent indicates instability and sensitivity to initial conditions, whereas negative exponent implies convergence. An exponent close to zero characterizes a critical regime, often referred to the edge of chaos, where signals neither explode nor vanish and can propagate for a long period. This critical regime has been reported to correlate with the high performance of neural networks across various contexts [Boedecker& et al., 2012, Poole et al., 2016, Pennington et al., 2017].
Figure shows that SA models achieving high test accuracy empirically have Lyapunov exponents close zero, thus operating near criticality. The maximum and mean exponents display nearly identical behaviors. We vary
including 1) and the norm of weight matrices, and plot multiple points (see details in Appendix B.3). In models with normalization layers, the exponent tends to concentrate around zero, indicating the criticality. High accuracy is achieved only in these models. This supports the stabilizing effect of normalization implied by Proposition 5. and Section 5.2. In contrast, energy-based symmetric SA models show negative exponents, consistent with their Lyapunov stability, leading to lower accuracy. Interestingly, the maximum Lyapunov exponent of successful models is slightly positive (~ 0.1), indicating that the dynamical state resides near
7

---

AKOrN
ItrSA
0.6
N=4
N=8
0.4
Accuracy
Accuracy
N= 16
N=
0.2
N-64
N=128
0.0
N=256
N=512
128
256
512 1024
2048
126 256
1024
2048
Figure 4: ItrSA consistently improves accuracy the number of loops 7 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.
criticality from the chaotic side. Notably, we observed that the dynamics with this slightly positive maximum Lyapunov exponent indicate the sensitivity to initial conditions, implying chaotic behavior (see Figure S.5 the appendix). This observation aligns with previous reports of positive maximum Lyapunov exponents in some Transformers [Inoue et al., 2022, Liu et al., 2024, Tong et al., 2025]. observed similar Lyapunov exponents across the CIFAR-I0 dataset (Figure S.7) and the language modeling task (Table S.4). We further observed that the number of attention heads increases, the Lyapunov exponents tend to increase (Figures S.1 and S.2). This implies that multi-head attention would favor more dynamic state.
Quantitative insight into inference dynamics
Here, we experimentally investigate the test-time scaling of inference looped architectures. We mainly focus on evaluation on the Sudoku task using the SATNet [Wang et al., 2019] dataset for in-distribution (ID) data and RRN dataset [Palm et al., 2018) for out-of-distribttion (OOD) data. At test time, we increased the number of loops beyond the training setting of T = 16. Details are provided in Appendix B.
6.1 Test-time scaling and normalization
Miyato et al. [2025] showed that AKOrN exhibits test-time scaling, whereas ItrSA does not, sug- gesting the superiority of AKOrN over ItrSA. However, with our formulation of ItrSA, Figure demonstrates that ItrAA also exhibits test-time scaling. Moreover, it shows that AKOrN fails to main- tain test-time scaling when N becomes large, which is consistent with the observations by Miyato al. [2025]. This issue can be mitigated by applying RMS normalization, where we use oscillator wise RMS normalization. The learned scaling parameter 7 prevents the Jacobian from becoming exces- sively large. Empirically, we confirmed that the trained remains small (see Table S.3), eliminating the need for explicit clipping such as hna <
6.2 Application to regularization
To apply our energy-based and lacobian-based analysis, we add the following regularization term R, scaled by tunable parameter A.
Method (i): Energy-based regularization. Proposition suggests that SA architectures satisfying specific conditions inherently minimize an energy function. To investigate the practical utility of this property, we consider applying the constraints identified the proposition as a regularization term to the existing multi-head MSA. We add the following energy-based regularization term to the loss function during the training of ItrSA, which can be interpreted as an approximation of energy-based SA models. By defining the concatenation as WV = [WY, w DDD we introduce
Rg-muti W V wo (W (wV W
(20)
where both V and O implemented orthogonal matrices. Note that each W in the proposition is interpreted the product W w in ItrSA. If Rg-muti 0, W.Y w becomes

---

ID Data
00D Data
1.0
0.6
0.8
0.5
ItrSA
Accuracy
Accuracy
0.4
AKOrN
0.6
E-single
0.3
E-multi
0.2
Spec (ItrSA)
0.1
Spec (AKOrN)
0.0
0.0
16
256
512
2048
16
512
2048
Figure 5: Energy-based regularization ("E-single" and "E-multi") underperforms the original methods, while Jacobian spectral regularization ("Spec") outperforms. We set = 8 except for E-single (H 1). For AKOrN, used N 4.
symmetric as implied by Proposition 4.2. For a single- case, we can also propose Proposition 4.1 regularization term Rg-single (see Appendix B.2).
Method (ii): Jacobian spectral regularization. On the other hand, controlling the Jacobian spectra is an effective way to stabilize neural networks. Following the regularization proposed by Lewandowski al. [2025], we introduce the following regularization term:
Rspee = Ewesa - Σ BESA
|bll
(21)
where the summations are taken over weight matrices W and bias terms in the SA modules, and o (W) denotes the largest singular value of W. This regularization encourages the singular values to be close to 1, which has been shown to be beneficial for recursive architectures [Chang et al., 2019].
We apply RSeec to both ItrSA and AKOrN.
Limitation of energy-based regularization. Figure shows the effects of the regularization methods. The accuracy of E-multi is lower than that without regularization. The Lyapunov exponents shown in Figure S.2 indicates that multi-head energy regularization encourages more convergent dynamics and this does not necessarily yield better performance. E-single fails to reduce the training loss and fails even on tasks possibly due the single-head constraint. These results suggest that energy-basee regularization may be unnecessary, casting doubt on the validity of the energy-based perspective for practical applications.
Spectral regularization is particularly effective in AKOrN. Spectral regularization substantially enhances the performance both AKOrN and ItrSA, and the effect is especially pronounced in AKOrN, Equation (16) shows that the maximum eigenvalue of the Jacobian can grow significantly when the weight matrices have large norms. Spectral regularization addresses this issue by directly constraining the Jacobian spectrum. When N 8, AKOrN achieves the best performance and spectral regularization is also effective (Figure S.3).
6.3 An interpretation of pseudo-energy via Jacobian
While is non-trivial for general SA dynamics to have an energy function, Miyato et [2025] empirically found that AKOrN approximately decreases the quan- tity Epscus(() (t) = - Tr(x Y (0), which we refer to as eneugy- during the updates of Eqs. (7) and (8), where ( C MSA(X(9)) They also reported that the prediction with a lower pseudo-energy performs het-
1.0
1e6
Cosine similarity Contribution
Energyy
0.4
Under certain symmetry assumptions without SA, it reduces to the energy function of generalized Kuramoto model. However, interpretation in AKOrN with SA remains unclear.
0.2
10
15
We found that the Jacobian provides interpretation. Sup- pose that vec(MSA(X (+)) DS is well-approximated
Figure 6: Test-time inference of AKOrN the Sudoku dataset.

---

by J where E RDS denotes vec(X (x) and J. = M / E R DDSDDS Fig- ure 6 shows that while Epeudo significantly decreases, the cosine similarity between (e((MSS(X()) and J remains high throughout iterations. Under this approximation, neglecting a constant term C, we obtain the relation Eprcudo = x. / quadratic form involving symmetric matrix S, J + J," We further expand 2 in the orthonormal basis of S, as 2t Ebs a of and the eigenvectors are sorted descending order of their eigenvalues (\ 2 .... 2 XPS). ). The contribution index in the figure quantifies the extent to which the state is captured by the top 2% eigenvalues, k= DS (ap)². During inference, the proportion of the state in the top
specifically Ek<00.02s(( (at) /
2% eigenspace increases monotonically, eventually dominating nearly 40%. This suggests that the loop effectively performs a computation analogous to power iteration, but constrained to real space with positive eigenvalues. Thus, the Jacobian provides a meaningful and promising interpretation for the observed decrease of pseudo-energy during AKOrN inference.
As side note, in Section B.5, more detailed analysis of the Jacobian matrix implies that this alignment to the cigenspace is predominantly determined by a certain time-indepenee -en-independeet matrix.
7 Conclusion
In this paper, we advanced the understanding and control of recurrent self-attention (SA) from a dynamical systems perspective. First, we generalized energy-based formulations to weaker symmetry and multi-heaa configurations closer to realistic settings. However, experiments showed that energy- regularized SA underperformed standard SA, suggesting other factors play in practical models. Second, we analyzed the Jacobian matrix of standard SA architectures, revealing that normalization layers effectively regularize their spectral properties. We further clarify that Lyapunov exponents at criticality characterize gigh-performance inference, which indicates that the state update of SA is more dynamic than energy- cengyycoonttaineed cases. We also argued how Jacobians provide quantitative insights into performance-enhancinn regularization and pseudo-energy behaviors.
Limitations. In this work, focused recurrent SAs without positional encoding, masking, or MI.P blocks, unlike looped Transformers in practice. Investigating how these components alter state dynamics remains an interesting direction. Regarding theoretical limitations, the upper bound in Proposition 5.1 provides a conservative upper bound that is looser than empirical observations, an issue also noted in existing analyses of SA's Jacobians without normalization. Deriving tighter bounds remains an intriguing open theoretical problem. Additionally, analytically justifying empirical phenomena, such as the Lyapunov exponent concentration (Section 5.3) and the Jacobian- based approximation (Section 6.3), by solving state dynamics would be challenging yet exciting themes for theory. We expect our findings serve as foundation for further theoretical and practical advancements in looped SA architectures and their rich inference dynamics.
Acknowledgments and Disclosure of Funding
We thank the reviewers for their insightful and helpful feedback the manuscript, and Takeru Miyato for his valuable comments. The authors acknowledge the funding support from JST FOREST (Grant No. JPMJFR226Q). RK is also supported by JSPS KAKENHI (Grant Nos. 22H051 16005511171111 23K 16965)..
References
Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances in Neural Information Processing Systems, 32, 2019.
Arpit Bansal, Avi Schwarzschild, Fitan Borgnia, Zeyad Emam, Furong Huang, Micah Goldblum, and Tom Goldstein. End-to-end algorithm synthesis with recurrent networks: Extrapolation without overthinking. Advances in Neural Information Processing Systems, 35, 2022.
Joschka Boedecker, Oliver Obst, Joseph T Lizier, N Michael Mayer, and Minoru Asada. Information processing in echo state networks at the edge of chaos. Theory in Biosciences, 131:205-21 213, 2012.
Giuseppe Bruno, Federico Pasqualotto, and Andrea Agazzi. Emergence of meta-stable clustering in mean-field transformer models. In International Conference on Learning Representations, 2025.
10

---

Valérie Castin, Pierre Ablin, and Gabriel Peyré. How smooth is attention? In International Conference on Machine Learning. pages 5817-5840. PMLR, 2024.
Bo Chang, Minmin Chen, Eldad Haber, and Ed H Chi. Antisymmetriernn: A dynamical system view on recurrent neural networks. In International Conference on Learning Representations, 2019.
Arman Cohan, Franck Dernoncourt, Doo Soon Kim, Trung Bui, Seokhwan Kim, Walter Chang, and Nazli Goharian. A discourse-awan attention model for abstractive summarization of long docu- ments. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume 2 (Short Papers), pages 615-621, 2018.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations, 2018.
Richard Diehl Martinez, Zébulon Goriely, Hope McGovern, Christopher Davis, Andrew Caines, Paula Buttery, and Lisa Beinborn. CLIMB curriculum learning for infant-inspired model building. In Alex Warstadt, Aaron Mueller, Leshem Choshen, Ethan Wilcox, Chengxu Zhuang, Juan Ciro, Rafael Mosquera, Bhargavi Paranjabe, Adina Williams, Tal Linzen, and Ryan Cotterell, editors, Proceedings of the BabyLM Challenge at the 27th Conference Computational Natural Language Learning, pages 112-127. Association for Computational Linguisties, 2023.
Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian R Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, and Tom Goldstein. Scaling up test-time compute with latent reasoning: A recurrent depth approach. arXiv preprint arXiv:2502.05177 2025.
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of clusters in self-attention dynamics. Advances in Neural Information Processing Systems, 36, 2023.
Borjan Geshkovski, Hugo Koubbi, Yury Polyanskiy, and Philippe Rigollet. Dynamic metastability in the self-attention model. arXiv preprint arXiv:2410.0m 00.0633, 2024.
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical perspec- tive on transformers. Bulletin of the American Mathematical Society. 62(3):427-479 2025.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. In International Conference on Machine Learning pages 1 11398-1144 1442. PMLR, 2023.
Eldad Haber and Lars Ruthotto. Stable architectures for deep neural networks. Inverse problems, 34 (1):014004, 2017.
Benjamin Hoover, Yuchen Liang, Bao Pham, Rameswar Panda, Hendrik Strobelt, Duen Horng Chau, Mohammed Zaki, and Dmitry Krotov. Energy transformer. Advances in Neural Information Processing Systems, 36, 2023.
Yunzhe Hu, Difan Zou, and Dong Xu. Hyperspherical energy transformer with recurrent depth. arXiv preprint 1.XXv:25022.164: 2025.
Katsuma Inoue, Soh Ohara, Yasuo Kuniyoshi, and Kohei Nakajima. Transient chaos in bidirectional encoder representations from transformers. Physical Review Research, 4(1):013204, 2022.
Tokio Kajitsuka and Issei Sato. Are transformers with one layer self-attention using low-rank weight matrices universal approximators? In International Conference on Learning Representations, 2024.
Nikita Karagodin, Yury Polyanskiy, and Philippe Rigollet. Clustering in causal attention masking. arXiv preprint arXiv:2411.04990 2024.
Hassan K Khalil. Nonlinear Systems. Prentice Hall, 2002.
Diederick P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.

---

Diederick P Kingma and Jimmy Ba. Understanding factual recall in transformers via associative memories. In International Conference on Learning Representations, 2025.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 (canadian institute for advanced research). URL http:/ //ww:. cs. wowontoootoo..oooo ca.torontoo.edu//kriii/ cifar html.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self sulf-suue learning of language representations. In International Conference on Learning Representations, 2020.
Alex Lewandowski, Michat Bortkiewicz, Saurabh Kumar, András György, Dale Schuurmans, Mateusz Ostaszewski, and Marlos C. Machado. Learning continually by spectral regularization. In International Conference Learning Reprexentutions, 2025.
Shuhong Liu, Nozomi Akashi, Qingyao Huang, Yasuo Kuniyoshi, and Kohei Nakajima. Exploiting chaotic dynamics as deep neural networks. arXiv preprint arXiv:2406.02355 2024.
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. In International Conference on Learning Representations, 2018. 20188.
Takeru Miyato, Sindy Löwe. Andreas Geiger, and Max Welling. Artificial kuramoto oscillatory neurons. In International Conference on Learning Representations, 2025.
Lorenzo Noci, Sotrris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 35, 2022.
Rasmus Palm, Ulrich Paquet, and Ole Winther. Recurrent relational networks. Advances in Neural Information Processing Systems, 31, 2018.
Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice. Advances in Neural Information Processing Systems, 30, 2017.
Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. Advances in Neural Information Processing Systems, 29, 2016.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 188):9, 2019.
Hubert Ramsauer, Bernhard Schäff, Johannes Lehner, Philipp Seidl, Michael Widrich, Lukas Gruber, Markus Holzleitner, Thomas Adler, David Kreil, Michael K Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. Hopfield networks is all you need. In International Conference on Learning Representations, 2021.
Nikunj Saunshi, Nishanth Dikkala, Zhiyuan Li, Sanjiv Kumar, and Sashank Reddi. Reasoning with latent thoughts: On the power of looped transformers. arXiv preprint arXiv:2502.17416 2025.
Sidak Pal Singh, Gregor Bachmann, and Thomas Hofmann. Analytic insights into structure and rank of neural network hessian maps. Advances in Neural Information Processing Systems, 34.
Anh Tong, Thanh Nguyen-Tang. Dongeun Lee, Due Nguyen, Toan Tran, David Leo Wright Hall, Cheongwoong Kang, and Jaesik Choi. Neural ODE transformers: Analyzing internal dynamics and adaptive fine-tuning. In International Conference on Learning Representations, 2025.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Po-We Wei Wang, Priya Donti, Bryan Wilder, and Zico Kolter. Satnet: Bridging deep learning and logical reasoning using differentiable satisfiability solver. In International Conference on Machine Learning, pages 6545-6554. PMLR, 2019.
12

---

Liu Yang, Kangwook Lee, Robert D Nowak, and Dimitris Papailiopoulos. Looped transformers are better at learning learning algorithms. In International Conference on Learning Representations. 2024.
Yongyi Yang, David P Wipf, et al. Transformers from an optimization perspective. Advances in Neural Information Processing Systems, 35, 2022.
Yuichi Yoshida and Takeru Miyato. Spectral norm regularization for improving the generalizability 2017.
of deep learning. arXiv preprint
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are trans- formers universal approximators of sequence-to-sequene functions? In International Conference on Learning Representations, 2020.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances Neural Information Processing Systems, 32, 2019.
13

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We claim that our abstract and instruction succinctly summarize the main contributions, and the introduction provides thorough overview of the paper's scope with our motivation.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We explicitly discuss the limitations of our theoretical analysis the limitations section of our paper, highlighting the need for further investigations.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
14

---

Answer: [Yes]
Justification: We clearly state our assumptions alongside the propositions and provide complete proofs in the appendix. This ensures that theoretical results are well-supported and verifiable.
Guidelines:
The answer NA means that the paper does not include theoretical results, All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We include all essential details needed to replicate our main experimental results within the paper to ensure that our findings are reproducible.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
15

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [No]
Justification: The code and data used in the experiments not publicly available, and we do not plan to release them. As result, the supplemental material does not contain instructions for reproducing the main experimental results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/ Sawiiidee//odSuumiise
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurlPS code and data submission guidelines (https: SCooeeeueeiiiiaee cttppssccc/pulcc////i S wwwiiissssnsiiiiieoeoo.
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: We detail all necessary training parameters ensure that our experimental results can be faithfully reproduced.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We include error bars and standard deviations in our results where applicable, ensuring that the statistical significance of our findings is clear and well-documented.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
16

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: Detailed descriptions of the computational resources used, including hardware
specifies and implementation details, provided in the Appendix aid in reproducing experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: http///eerrpp.ppeee wwweeris..cppplliipeeeeeeeeeecee cctppccccp/eeeeeeecceecc htw.hi///eehcchiieeeleeeniee
Answer: [Yes]
Justification: After thoroughly reviewing the NeurIPS Code of Ethics, we confirm that our research adheres to all the specified guidelines.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer:[NA]
Justification: Given the theoretical nature of our work, we assess that it does not directly engage with socictal impacts.
17

---

Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considcrations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning conrectly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our research does not involve the release of data models that pose high risks for misuse, hence specific safeguards are not required.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-uss should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: Our study does not use any external assets, thus no licensing attribution issues are applicable.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
18

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No new assets are introduced in our paper, so there are no associated documen- tation requirements.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
19

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core methodology of this research does not involve the use of LLMs in any important, original, non-standard way.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https: htttps://neurip. http:///errpppcccooofeerrrrse/// cct//ooneeeencce/225//////////// for what should or should not be described.
20

---

A Proofs
A.1 Lemmas
We use the following lemmas in our derivation. Lemma A.1 (Singh et al.).
TAXB
A
(22)
0X
Lemma A.2. Let Y E RR be an input matrix. Then, the Jacobian of II with respect to is given by
a (Y) 0Y
blockdiag (1 (TY (ID
Yis (d 1 Y ||
(23)
Proof. Since II operates independently each of Y, the Jacobian is block-diagonal, with each block corresponding to the derivative single normalized row:
oIl(Y) 0Y
a11(Y (a. 1 jim))
blockdiag({
(24)
For each row, we compute the gradient of the normalized vector:
A
(aa R.R./ Yaal
(25)
0Y R4
YR. 6.
(In
(26)
Lemma A.3. Let Y 5xD be input matrix. Then, the Jacobian of RMSNorm with respect is given by
a RRMSNor((Y) blockdiag
(1ixiail diag(>)(I
Rd
M
(27)
OY
M
Proof. Since RMSNorm is expressed as
RMSNorm (Y) = diag(>)II(() (.))
(28)
the result follows from lemma A.2.
A.2 Proof of Proposition 4.1
Proposition 4.1 is restated.
Consider the continuous-ti dynamics for single-head SA equipped with projection (3). The energy function
Esingle (X) =- -
exp
BX [i
W
W KT X(33
(29)
1.3
is monotonically decreasing d'Eringte (X) < 0 under the condition:
WV ( QT wwW myRwky KT 1/2. 1/22.
(30)
Proof. Let A derivative of Exingle with respect to X
X JVWV and let A WOWKT The first-order
[s]]
is:
d.Einnt (X)
dexp
3X [.
AX
dexp
(XXT AX
(31)
dX [s.]]
dX
[i.]]
dX [..]]
J
jzN
21

---

exp
3x [i,:]
| AX [5.11
(32)
d'((t.)
exp
(.. (xx] (i:))
[1.:]
A A X
Σ exp
(1.) AA)
(i.)]
AX G.
(33)
j7i
exp
(8X [i..]
A
(6.:]
(AT
+ (A))
(34)
=
Σ exp
(x [..3] AX
(..)) (AT + AA B-l-
(35)
Under the given condition on the weights, W (AT + A))2, have:
A [ = (softmax (3 XAXX )
+ +A)) TA))Ta/2
(36)
(X [..] AX (A A) 11-1/2
exp
(37)
Zi
where Z₁ = Ex exp (3XT. [..1 (
is the normalization term.
Then, we have,
d.E. d.Esingle (X)
d. dEssigle( (X)
dX
(38)
dt
dX
dt
d. d.Essngle (X)
Σ
(ID (R. (p.))
(39)
dX
Σ
Σ exp 3x (..) (x (AT + A) (LD [[..)
(1.)) (1.))
(40)
Σ (Z, (ID [ X (1.))
(41)
= 2 Σ T Tg (ID (i. (
(42)
00,
(43)
where, the last inequality, we used the fact that the matrix In [(,.j [1.)] definite.
is positive semi-
A.3 Proof of Proposition 4.2
Proposition 4.2 is restated. Consider the continuous-time dynamics for multi-head SA without projection: dx/ dt -m=1= (X) An energy function
Emudia (X)
ΣΣ exp (9X W WOW KT X
(44)
is monotonically decreasing as d.Emmutil (X) (xmitt(X)/// < O under the condition
W.Y' -W W + w KT 1/2, w...k U
(45)
where U12)A E DD2/(22) (h E [1, satisfies the orthogonality condition U, Sxw D//2H)
U
k.
Proof. Let AA softmax ( W x ) and let A, W WKT The first-order
derivative of Emutii (X) with respect to X (R..) is, similarly to the single-head case, given by:
e (X)
ΣΣ
exp (BXT., T. AnXisa) (AT + An) GA
(46)
0x bdl
h
22

---

Under the given condition on the weights, similar to the single-head case, we have:
exp
[xx T X (A), + 1,42/
[.
(47)
Zn,
where Zh, = Ey exp (1X A
A
is the normalization term. Then, we have,
d'iinnti (X)
d.EmatiiX (X)
dx
(48)
dt
dX
dt
Emmotal (X)
Σ
Σ M
(49)
[6.1l (1.)
h
Σ
Σ exp (3XTT An ((a) ) (AN + Ab))
bil Σ Amtil)
(50)
= -2 Σ
Σ Zn, A Σ AmRa
(51)
= 2 Σ
Σ
Zn. Axp,al1?
(52)
h
00
(53)
where we the fact that for h + lt,
A, AN = U2 LIN 2, = =0, A, = 0.
(54)
A,AN UI,N =0;
U1,
(55)
and thus
Aaj AR(P.)] AMpis
(56)
exp (3X [6.1]
An.
Σ
X(ii]
(57)
Z
(A + (A AT Σ
exp
[3x]] T
[e.3]
An
An.Xiaa bia))
X
1/.1/44
(58)
Zia
= 0.
(59)
A.4 Proof of Proposition 5.1
Proposition 5.1 is restated. Suppose that, in the update of ItrSA (9), the imput the normalization layer satisfies [x..] + + [..]ll 2 for all i E S). Then, the spectral norm of the Jacobian
satisfies the upper bound
3 ORMSNorI(X X + AX
maxx(175))
(1 + 17 (M)) (2))
(60)
OX
R.
where Juss OSSA(X (AMSA(X)/2. // denotes the Jacobian of MSA.
Proof. First, for any vector a E RD, the eigenvalues of the matrix Ip || are 1 (with multiplicity
T
D --1 1) and 0 (with multiplicity 1). Hence,
0.
ID
1.
(61)
|| 2
2
23

---

Using Lemma A.3, we have
(11Vi-il diag(>)
RRMSNrm Y
blockdiag
In
(62)
0Y
Meall
1 Yrall
1o.1 Y.. [R,3]
diag(y)
1D
(63)
||
max, has
(64)
Therefore, setting = X + + = X + + (C + (S)) we have
RMSNorm X + 7AX)
@RMSNorm(Y' OY
(65)
0X
OY
0x
2
@RMSNorm Y
0Y
(66)
0Y
2
0x
max, his R
|IsD n.Jmsallz
(67)
(h7)()))))
(1+ + 17 (X)
(68)
R
Here, we can also show that, in the limit + oo, the Jacobian norm remains O(1). Define, for each i, the projection
(B. [c.]]
: ID
(69)
| [..]]]
so that |Pille = 1. Let D = diag(r), and define the block factors
A,(1) (id] + 17 AX 1.1l22
D DP.
A = blockdiag ((A((1))
(70)
By the triangle inequality, for each i,
[(.a] AX [d]] la 2 ||AX [AXp,all l22 [xalll2-
(71)
Recall that X is the output of the previous layer, so that
X = RMSNorm(2 Z)
(72)
for some Z SxD Therefore,
|X R.
RMSNorm/ Z (F.a))
(73)
= =I|diag(7)
Z
(74)
| 12
(75)
EE= 2
< max hy)
(76)
max hy|
(77)
for each = S. Assume that min; ||X (G.gl2 2 E for some constant 2 max, hil/e, we have
0. Then, for sufficiently large satisfying
X (.. l1
17 2 2
(78)
A Ba1l22
24

---

Hence,
[X k.l + A Balll2
2 7 | X t.ll2
||
|xxill
(79)
2 2 ||A R.
(80)
Using submultiplicatiiity and P.||| 1, we obtain
|| D Pa Pille
|A,(1))22= = [x + AXX kajl2
(81)
|| 12
(82)
(a.) + + 7 X
[1
2 DI12
(83)
r) X|AXpll p.1l22
for all I.
For block-diagonal matrix, the operator norm equals the maximum block norm; thus
||)(1)) max [n)l2))
(84)
D|ID1|2
max 1) | B.1l22
(85)
2 D112
(86)
1 min;
|| | XAXxxll x.1ll2
2 D|l
(87)
n
Therefore, we have
AA(r) (IND + JMSA)ll < (I) la (1) [| l2 + m|I[Msllle)
(88)
2
D 2 | l|2 + -1))
<
(89)
Since Eq. (16) makes constant upper-bound of ||JMSA 21,, for 00, obtain
a RMSNorm( X AX)
=01).
(90)
0x
A.5 Derivation of the eigenvalue bound in oscillatory eases (Section 5.2)
We show that all eigenvalues λy of the Jacobian
yy
[(x) = In
(ID + 112).
(91)
|
|l
satisfy |Ayl I 1, where y (In 1)52)ar.
We begin by computing the norm of y:
|| = x: (I + 1/52 [In + 1/522x
(92)
vaa Qr vaa in xc + 1 TTsTTs T Srx
(93)
+ 12
(94)
where we used the fact that for antisymmetric matrix Q x TT f2x 0. Note also that anti-
symmetric matrix has eigenvalues of the form =iwj, where Wj 2 (j = 1, 2....). ) For simplicity, assume all eigenvalues have identical magnitude W/ Then, we have
| = + 2
(95)
We also use the facts that
yy
Io
(96)
|/
25

---

and
D + 1/52/12 |1 dmFw
V1+ mpw?.
(97)
Combining these, we obtain the following bound on the spectral norm of J(a)):
| =
|
yy ID
(ID + 7/52)
(98)
y 12
llll
1
| + 1/12112
(99)
|
/ + w2ر
1.
(100)
/1 + rPww²
This implies that all eigenvalues of J (x) satisfy PA,l S 1.
B Experimental details
B.1 Experimental setup
We solved Sudoku task, which a puzzle played on a 9 grid, where some of the cells are pre-filled with digits from 1 to 9, and the remaining cells are left blank. The objective is to fill in the blank cells such that each 1) row, 2) column, and 3) subgrid contains each digit exactly once.
In our experiments, we used two Sudoku datasets: the SATNet [Wang et al., 2019] and RRN dataset [Palm et al., 2018]. The key differences between the two are that the RRN dataset is more difficult (with only 17-34 given digits compared to 31-42 in SATNet) and larger in size (198k samples vs. 10k samples). Following Miyato al. [2025], we used the SATNet dataset for training as in in-distribution (ID) data and the RRN dataset as out-of-distribution (OOD) data. This setup allows us evaluate the ability of models generalize more challenging settings.
We primarily followed Miyato et al. [2025] and used their official implementation! setting the dimen- sion of oscillators of AKOrN N 8. The readout module of AKOrN (described in Appendix C.2) was also incorporated into our ItrSA model. We used the Adam optimizer [Kingma and Ba, 2015] and trained for 100 epochs with batch size 100. For all settings, we tuned the learning rate over 1 x 10-3) 3] and, for regularization methods in Figure 5, the parameter over (1 10-8,1 1 10 ..... 10-1), selecting values based OOD accuracy iteration T 16. All experiments were conducted on NVIDIA H200 GPUs, and we run experiments with different random seeds.
{I x 10-6, 6 5 x 10-6
We also conducted experiments on the CIFAR-10 dataset [Krizhevsky et al.]. See table S. for training and model configurations.
Table S.1: Training and model configurations.
Parameter
Sudoku
CIFAR-10
Hidden dimension D Number of heads H
512
384
8
8
Initial value of 7)
1.0
1.0
Batch size Number of epochs
100
128
100
200
B.2 Single-head generalized symmetric SA
For single-head generalized symmetric SA, we define
single WW' W')
(101)
under the condition that = If RE.single = 0, setting W WVWO satisfies the condition on
W described in Proposition 4.1.
'https://githann htt.:h/w...cooottttoooooouviss.a
26

---

Table S.2: Regularization coefficients used in Figure
Method
Value
E-single
le-6
E-multi
le-4
Spec (ItrSA)
le-4
Spec (AKOrN)
le-5
B.3 Lyapunov exponent
The Lyapunov exponent quantifies the exponential rate which nearby trajectories in dynamical system diverge. For discrete time system x (+++)) =f =(x(()), the Lyapunov spectrum {A;} is defined as:
A lim
log |a 1
(102)
T 27
where ( is the i-th eigenvalue of the positive semi-definite matrix
(x
A (7)
(103)
dx
(x
and f denotes the T-fold composition of the function f. The maximum Lyapunov exponent is then defined as
Amax max Aze
(104)
To mitigate numerical sensitivity, we also use mean Lyapanov exponent
Amean
A
(105)
M
where M denotes the number of Lyapunov exponents.
In our experiments, we approximated the Lyapunov spectrum using a finite time horizon of T 16 on randomly selected sample. For models without normalization and symmetric SA, we trained evaluate how the
them for only one epoch, as full training was not feasible due instability. Lyapunov exponent varies, we adjusted the input scaling of X, the step size 1), and the norms of the value projection weights, [W || and [Wo) ||. the SA update of X.
B.4
Details of other figures
For Figure 2c, we computed the eigenvalues of the Jacobian matrix at T = 16 on a randomly selected sample from the Sudoku dataset. For the model with normalization, used the fully trained model. For models without normalization, we followed the same setup as in the Lyapunov experiments and used models trained for only one epoch.
For the computation of the SA's Jacobian in Figure 1, we used the CCDV arXiv summarization dataset [Cohan et al., 2018], it provides text data suitable for varying the number of tokens. We used an initialized SA and computed the Jacobian and SA followed by normalization over 500 randomly selected samples. The norm of tokens
to R 100 and their dimensions D 256.
B.5
Additional results
B.5.1
Effects of model structure and regularization
Number of attention heads. To further examine the effect of normalization, we computed Lyapunov exponents while varying the number of attention heads H. Because the maximum Lyapunov exponent numerically sensitive, we report the mean Lyapunov exponent instead. The results in Figure S. show that models with few heads (H E {1, 2)) perform poorly, whereas H = 8 achieve the highest accuracy. Moreover, the mean Lyapunov exponent is positively correlated with accuracy, suggesting that more dynamic states are associated with better performance.
27

---

1.0
Hml
H-2
0.5
H=2
0.8
Hizd
H=4
Test accuracy
HmB
Test accuracy
0.4
0.6
0.3
0.4
0.2
0.2
0.1
0.0
0.0
-1.5
-1.0
-0.5
-1.5
-1.0
-0.5
Lyapunov exponent
Lyapunov exponent
(a) ID
(b) OOD
Figure S.1: Effect of the number of attention heads H on mean Lyapunov exponent and accuracy in the Sudoku dataset.
y in RMSNorm. Table S.3 presents the values of the y parameter learned by ItrSA, The results indicate that the trained models exhibit small y values, with max, l7y1 < 1.
Table S.3: y RMSNorm with different N.
N
|||
max, hn|
0.229 + 0.000
0.229 0.000
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
0.489
0.001
0.489 0.001
128
0.841 + 0.000
0.841 + 0.000
256
0.738 + 0.000
0.738 0.000
512
0.811 + 0.000
0.811 + 0.000
Further results on regularization. In Figure S.2, we evaluate the effect of the proposed multi-head energy (Proposition 4.2) both the mean Lyapunov exponent and accuracy. As the regularization strength À increases, both metrics consistently decrease. Even at À = 0, accuracy is lower than the original model due to the orthogonality constraint. The "Hard" constraint yields the lowest accuracy and the smallest Lyapunov exponent. Overall, these results indicate that stronger regularization suppresses the Lyapunov exponent but also degrades accuracy, suggesting that while multi-head energy regularization encourages more convergent dynamics, it does not necessarily yield better performance.
In Figure S.3 we plot the case where we used N as the oscillator dimension of AKOrN. AKOrN achieves the best performance and spectral regularization is also effective.
B.5.2 Distribution of Lyapunov exponents
Figure S.4 shows the distibution of the Lyapunov exponent.
B.5.3 Lyapunov exponents in language modeling tasks
To evaluate in a more realistic scenario, we conducted language modeling experiments on the BabyLM Challenge dataset (2023, 10M version) [Diehl Martinez et al., 2023], comparing our ItrSA model with GPT-2 [Radford et al., 2019]. We trained for 30 epochs using the AdamW optimizer. As shown Table S.4, the maximum Lyapunov exponent (MLE) for ItrSA is slightly positive, which consistent with our other tasks and suggests mildly chaotic dynamics. We also show the loss values in the case GPT-2 as reference which confirms that the performance our ItrA can become comparable to that of GPT-2.
28

---

1.0
A=le-4
0.20
A=1e-4
A-1el
A=lel
0.8
Herd
Hard
0.15
Test accuracy
Test accuracy
0.6
0.10
0.4
0.05
0.2
0.0
0.00
1225-0.500-0.45---11 0.5-00..555-
Lyapunov exponent
Lyapunov exponent
(a) ID
(b) OOD
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 Proposition 4.2.
Data
0.5
0.8
ItrSA
Accuracy
0.6
Accuracy
0.3
E-single
0.4
0.2
E-multi
Spec IItrSA)
0.2
0.1
Spec (AKOrN)
16
64
128
256
512
1024
2048
16
54
128
256
512
1024
2040
Figure S.3: Energy-basse based regularization ("E-single" and "E-multi") underperforms the original methods, while Jacobian spectral regularization ("Spec"" outperforms. We used N = 8 for AKOrN.
B.5.4 Sensitivity to initial conditions
In study, define "criticality" the point at which the largest Lyapunov exponent takes In our experiments, when the maximum Lyapunov exponent gets slightly positive, we observed behaviors consistent with widely accepted notions of chaos, particularly sensitivity to initial conditions as follows.
6000
6000
5000
5000
Frequency
4000
Frequency
4000
3000
3000
2000
2000
1000
1000
0
%
-0.8
-0.6
-0.4
-0.2
0.0
-0.6
-0.4
-0.2
0.0
0.2
Lyapunov Exponent
Lyapunov Exponent
(a) AKOIN
(b) ItrSA
Figure S.4: Distribution of the Lyapunov exponent.
29

---

Table S.4: Language modeling results. MLE denotes the maximum Lyapunov exponent.
Model
# Parameters
Training loss
Validation loss
MLE
ItrSA
96.2M 124M
4.31
5.67
0.196 + 0.016
GPT-2
1.49
5,58
Starting from an input, we ran the loop of the ItrSA model for 128 steps to obtain Tt=0- We then repeated the from perturbed input æ+=0 where NO0 10-3) is added at ( 0. From each trajectory, we sampled 300 equally spaced points between 1 and / 10000. The results are shown in Figure S.5.
Figure S.5a plots the LI distance between the two trajectories. In Figures S,5b and S,5c, we visualize single coordinate of five trajectories. Figure S.5b shows the complete trajectories. Since observed that the trajectories rapidly oscillate between two clusters and the visualization was subtle, Figure S.5c provides zoomed-in view of one cluster (values > 0.1).
Overall, the distance between the original and perturbed trajectories increases exponentially over time, demonstrating chaotic behavior in the system.
0.3
10a
MA yort plalAll
0.2
distance
102
value
102.
0.0
10
2000
4000
6000
8000
10000
1000
2000
3000
4000
5000
Iteration
Iteration
(a) LI distance between trajectories
(b) All sampled trajectories
Value
0.25
Iteration
(c) Trajectories forming single cluster
Figure S.5: Sensitivity to initial conditions investigated using the ItrSA model.
B.5.5 Jacobian-based interpretation of pseudo-energy
We provide additional insights into the interpretation of pseudo-energy via the Jacobian discussed in Section 6.3. Using Lemma A.3 from Noci et al. [2022], the Jacobian of SA expressed as
(SA Ox
Op
(Is W
POWVT Ox
(106)
where softmax K X /VDH).
As our experimental result indicated, we observed that Jæ aligns well with vec(MSA(X)). Since each head can be expressed as vec(SA(X)) = (P & WVT) our observation implies that in the Jacobian (106), the last term is dominant, that is,
H
J ~ L
P(U
& WVT
(107)
In other words, the derivative of the attention matrix P is small while that of the value matrix remains significant. In addition, the derivative of the attention matrix P sufficiently small over the whole time, the Lipschitz continuity implies that P remains close to its initialization, suggesting
30

---

iime-nnepenneent Jacobian approximation:
H J.
P
W
(108)
h=1
Then, the pseudo-energy is approximated by the following quadratic form:
Pepseudo
Σ
P
W
(109)
Figure S.6 empirically confirmed that both the contribution index and the pseudo-energy behave simi- larly even under these approximations. PV indicates the contribution index using the approximation (107), and PV (T = 0) uses (108). The figure shows that the contribution index can be effectively explained using only the attention matrix at the initialization of inference. The blue curve shows the approximated pseudo-energy (109) and works similarly to the original one. Thus, our bacobian-based analysis interprets the pseudo-energy as quantifying exploration in eigenspaces corresponding to large eigenvalues of the attention matrix determined by the initial inference states.
0.45
40000
0.35
alde Darm
jacobien
-6000
0.30
---- P/
(T-0)
80000
100000
0.15
120000
0.10
10
Figure S.6: Contribution and the quandratic form.
B.5.6 CIFAR-10
For the experiments on the CIFAR-IO dataset [Krizhevsky et al.), we used the same architecture and setup in the Sudoku experiments. We used the Adam optimizer and tuned the learning rate across {1 10-6, 5, 10-6, 1 10-3) based on the test accuracy at the iteration T 16. We trained models for 200 epochs and set the batch size 128. We used N AKOON.
the dimension of oscillators of
Figure S.7 shows the Lyapunov exponent on the CIFAR-10 dataset. This result is in the same trend with that on the Sudoku dataset.
31

---

Maximum Lyapunov exponent
Inf/NA
0.8
AKOIN
Norm) Symmetric SA
Test accuracy
0.6
w/ Norm
w/o Norm
1
0.4
0.5
0.2
0
10-3
10-1
10
10a
-Inf -2
Inf/NA
Lyapunov exponent
(b) Maximum Lyapunov exponent with and without normalization
(a) Maximum Lyapunov exponent accuracy
Figure S.7: Lyapunov exponent on the CIFAR- 10 dataset.
32

---

C Other details
C.1 Extended related work
Energy-based understanding. The Transformer architecture has been a focus of efforts to provide theoretical grounding. Geshkovski et al. [2025, 2023, 2024] formulated recurrent SA dynamics as interactions among tokens ("particles"), enabling theoretical analysis of phenomena such meta- stable clustering and rank-one collapse. Their continuous- tottuuuuuuueeeuuuuueeeee dynamics monotonically decrease an energy (Lyapunov) function given by summation over exponential functions, commonly requir- ing constraints such as single-head attention, hyperspherical token states, and symmetric weights. Karagodin et al. [2024] extended this framework the case of causal attention masking. Bruno et al. [2025] succeeded in mathematically characterizing the meta-stable clustering as a Wasserstein gradient flow of mean-field token dynamics, with the energy serving as its potential function, although they replaced the softmax function with an unnormalized exponential function and restricted their analysis to identity weight matrices. Yang al. [2022] considered exponential energy function similar to that of Geshkovski et al. [2025], describing the Transformer as performing alternating majorzatioonmiiniizztooo minimization updates on distinct energy functions. Their approach also accommodates
discrete state updates and MLP layers, although it entails complex conditions, including constraints on step sizes and proximity to fixed points. Ramsauer al. [2021] formalized the cross-an -atss-attnntion
mechanism as modern Hopfield networks, Hoover et al. [2023]. Hu et al. [2025] further developed energy functions for Transformers including self- attentttenssns. We do not address approaches based on Hopfield networks in this work, as they require architectural modifications, such as adding auxiliary signal paths that are absent in standard Transformers, which are beyond our scope.
Jacobian-based analysis. The Jacobian of state updates is fundamental for characterizing neural network dynamics. For example, it has been used to analyze edge-of-chaos behavior for stable signal propagation and gradient control [Boedecker al., 2012, Poole et al., 2016, Pennington et al., 2017]. Haber and Ruthotto [2017] interpreted forward propagation in neural networks as continuous-time dynamical systems and analyzed their Jacobians to prevent exploding and vanishing gradients. Chang et al. [2019] extended the ODE-based perspective to recurrent neural networks and proposed using anti-symmetric weight matrices satisfy discrete-time stability conditions. Several studies have explored Jacobian-base regularization techniques. Yoshida and Miyato [2017] proposed spectral norm regularization to reduce sensitivity input perturbations and improve generalization. Miyato et al. [2018] applied spectral normalization to stabilize the training of generative adversarial networks. Lewandowski et al. [2025] introduced spectral regularization for continual learning, aiming to prevent the loss of plasticity and maintain trainability across tasks by keeping the maximum singular value of each layer close to one. Regarding SA specifically, Noci et al. [2022] analyzed Jacobians to explain rank collapse, while Castin et al. [2024] evaluated their spectral properties mathematically. In this work, we use Jacobian analysis to understand inference dynamics in realistic SAs and also employ them as regularizers and performance indicators.
Looped architectures. Looped architectures in Transformers have been explored since their intro- duction by Dehghani et al. [2018]. One example is weight tying, as seen in the ALBERT model [Lan et al., 2020]. Equilibrium models [Bai et al., 2019] 9] use fixed-point solutions, which can be interpreted as infinitely looped computations. Yang al. [2024], Giannou et al. [2023] showed that Transformers with looped structures are capable of learning algorithmic tasks. Saunshi et al. [2025] further showed that looped architectures enhance reasoning ability through strong inductive bias. As the number of recurrent updates (i.e., loops) increases, performance scales efficiently, phenomenon we refer to as test-time scaling. Geiping et al. [2025] successfully applied test-time scaling to reasoning benchmarks, and Bansal et al. [2022] showed that it enables models to solve problems at test time that are more difficult than those seen during training. Miyato et al. [2025] proposed artificial Kuramoto oscillatory neurons (AKOrN), looped architecture that successfully solves tasks in neuroscience- inspired manner, demonstrating strong empirical results in unsupervised object discovery, adversarial robustness, calibrated uncertainty quantification, and reasoning.
33

---

C.2 Details of preliminaries
Energy-based analysis by Yang al. [2022] alternating inexact minimization algorithm as:
Yang et [2022] formalized updates of SA using
X softmax, X W X ) W"
(110)
where W DDDD is a symmetric matrix and softmaxx is a function reweighted with coefficient vector B.
Operation oscillators We use Xi to refer the j-th oscillator of the i-th token of X, which defined as X., (i,0--)N++MM E RN. They are defined as:
Y.
Omg (osc) (x()));s = 22 Xi,, Proj (Y) = (Ix-Xx,) X. Vig. IT( IT(o))( =
(111)
AKOrN then uses a readout module to read out patterns independent of the phase.
C' (m) E DxN my
|| 2z Σ UuayXin E RN
(112)
where Ukis E R N is learned weight matrix, g is learned function and l = 1... DN.
34


Reviews:
[{"review_id": "JfaqcYw5kj", "reviewer": "Reviewer_u5hQ", "summary": "This paper digs into what happens when you loop the same self‐attention (SA) layer over and over (no energy function forced). In other words, they make energy-agnostic characterization of inference dynamics of SA by dynamical systems analysis. They drop the old symmetry + single‐head rules and use a Jacobian‐based trick to show how norm layers (RMSNorm etc) keep the loop stable but still edgy (edge of chaos). They measure Lyapunov exponents and find best performance when exponents are around  0. They also try two regression methods to bridge the theoretical findings and practical usefulness. Energy-symmetric regularization does not help, but Jacobian-spectral one improves the results.", "questions": "- How to extend this Jacobian stability proof to full Transformer (residual+MLP+LN)? What is your intuition?\n- Why is there no other dataset?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n- Broadens theory beyond simple energy‐only constraints, so multihead SA works now.\n- Cool Jacobian angle explains why norm layers keep it right at the “edge of chaos”.\n- Jacobian‐spectral regularizer seems to boost performance in loops. \n\nWeaknesses\n\n- As they also mentioned in their limitation section, they ignore position encodings, masks & MLPs in the Transformer layer  - it is a very simplified transformer layer so we are not sure if it really scales to a real transformer.\n- Energy‐based regression looks nice, but actually hurts accuracy drastically, so I am not sure if the theory developed in the paper can result in any practical improvement.\n- Experimental limitation: They tested only on a single toy task: the Sudoku dataset, no real-world dataset presented e.g, NLP task or CV task, etc.", "was_revised": true, "final_justification": "The paper is solid and relevant to Neurips community so I would recommend the accept.", "rebuttal": "Thank you for your kind reading and comments. It appears that our inadequate explanation led to some misunderstanding regarding our claims and research approach. We would like to address these points as follows:\n\n---\n> **W1**: ... it is a very simplified transformer layer so we are not sure if it really scales to a real transformer.\n\n> **Q1**: How to extend this Jacobian stability proof to full Transformer (residual+MLP+LN)? What is your intuition?\n\n Similar to prior works [1,2,3], our focus is on understanding SA layers, not to replicate a full Transformer.  \n\n[1] Castin et al. \"How Smooth Is Attention?\", ICML 2024.\n\n[2] Geshkovski et al. \"The emergence of clusters in self-attention dynamics\", NeurIPS 2023. \n\n[3] Ormaniec et al. \"What Does It Mean to Be a Transformer? Insights from a Theoretical Hessian Analysis\", ICLR 2025.\n\nWe believe that in individual research papers, it is fairly common to focus on and analyze a single component of the blocks that constitute the architecture. This perspective is particularly valid in the analysis of Jacobians because the Jacobian of the entire model is **factored as the product of the Jacobians** of its components (e.g., $J_{MLP \\circ SA}=J_{MLP} J_{SA}$). The combination of layers is left to subsequent research efforts, and we expect that our findings will serve as a useful foundation in that context.\n\nNote that for your Q1, we tested ItrSAs and AKOrNs, which include a skip connection (in eq. 8) and normalization layers. \nIt does not include an MLP block, but a single linear layer referred to as the Omega term. \n\n---\n> **W2**: Energy‐based regression looks nice, but actually hurts accuracy drastically, so I am not sure if the theory developed in the paper can result in any practical improvement.\n\nOur purpose was to provide quantitative **insight into the limitations of the energy‑based approach** and to reveal the superiorities of non‑energetic (i.e., energy agnostic) models. This was one of our key take‑home messages. Therefore, we do need to aim for any practical improvement in the energy‑based approaches. To avoid a potential misunderstanding, we will re-emphasize this purpose of research in Section 6 (experiments) of the revised manuscript. \n\n---\n> **W3**: Experimental limitation: They tested only on a single toy task: the Sudoku dataset, no real-world dataset presented e.g, NLP task or CV task, etc.\n\n> **Q2**: Why is there no other dataset?\n\n**We evaluated our method on CIFAR-10** (in Appendix, L258–L259). \n\nWe also argue that **the Sudoku dataset is not merely a toy task**, as it inherently captures structured reasoning. Several studies (e.g., [4]) have focused on it because achieving high performance is far from trivial, especially for the OOD case. Miyato et al. [5] also use the Sudoku dataset to evaluate test-time computation tasks. More recently, this interest in the Sudoku task has led to a new benchmark for evaluating LLM reasoning abilities [6]. We expect our work to provide a solid baseline for exploring such more complex reasoning tasks.\n\nMoreover, **we measured the maximum Lyapunov exponent in an NLP setting**.\nWe used ALBERT [7], which employs a looped architecture, and computed the Lyapunov exponent over 12 loops. We compared the exponents of the pretrained model and a randomly initialized model using five text samples.\n\nAs shown in the table below, the Lyapunov exponent of the pretrained model is close to 0, indicating that the pretrained model is closer to criticality.\n\n| Model        | Maximum Lyapunov Exponent (mean ± std) |\n|--------------|-------------------------------|\n| Initialized  | 0.160 ± 0.003                    |\n| Pretrained   | 0.142 ± 0.015                    |\n\n[4] Wang et al., \"SATNet: Bridging deep learning and logical reasoning using a differentiable satisfiability solver\", ICML 2019.\n\n[5] Miyato et al., \"Artificial Kuramoto Oscillatory Neurons\", ICLR 2025. \n\n[6] Seely et al., \"Sudoku-Bench: Evaluating creative reasoning with Sudoku variants\", Preprint.\n\n[7] Lan et al., \"ALBERT: A Lite BERT for Self-supervised Learning of Language Representations\", ICLR 2020.\n\n---\n\nWe welcome any further questions or requests for clarification to increase the significance of our paper."}, {"review_id": "bXUlF5LAxS", "reviewer": "Reviewer_nAwM", "summary": "This work explores the dynamical effects of adding normalization in the recurrent self-attention dynamics. The authors first formulate the dynamics governed by an energy function that is minimized over time and propose it as a regularization. The authors argue that this direct Lyapunov method is restrictive; therefore, they suggest using Lyapunov’s indirect method (studying the Jacobian matrix at different times) to study non-stationary and oscillatory dynamics. They find that adding normalization stabilizes the dynamics (spectral radius < 1) and makes the model operate close to “criticality,” which is a property found in highly performant networks. The authors also propose a regularization based on ensuring that the singular values of the Jacobian stay close to 1.", "questions": "How do the authors define “criticality” in the context of the dynamical system considered? The authors mention it’s the edge of stability and chaos, but do the systems that they consider exhibit chaos if sent beyond the tipping point? \n\nI think the caption of Fig. 5 is wrong. It should be the other way round, right?\n\nThere’s a typo on line 196: singular value.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**: Most of the conclusions are based on theoretical results, along with some empirical results. The authors were able to relax previous constraints about matrices being symmetric/diagonal and single-head attention. The authors draw an empirical connection between adding normalization and “critical” states observed in better-performing models.\n\n**Weaknesses**: \n\nThe authors say they use a Jacobian-based analysis instead of an energy-based analysis because of non-stationary dynamics (oscillations). I believe Lyapunov’s direct method can be used for studying non-stationary dynamics. In fact, the type of oscillatory dynamics they consider is stationary from the point of view of dynamical systems (there is no limit cycle behavior here). \n\nThe finding that adding normalization stabilizes the dynamics of a recurrent circuit has recently been made: [Rawat, S., et. al., (2024). NeurIPS 37, 14712-14750]. In fact, normalization has been linked to the “critical” phenomena [Morone, F., et. al., (2025). bioRxiv, 2025-05].\n\nI think it’s interesting that adding normalization leads to critical behavior, but the authors do not give an intuition of why this happens. Also, what is so special about criticality that it enables good performance in the model considered? It might be out of the scope of this paper, but I would appreciate more intuition or some empirical results on why it might be a desirable property.", "was_revised": true, "final_justification": "The authors have clarified the questions I had in my initial review. I recommend the paper for acceptance.", "rebuttal": "Thank you for your helpful comments and suggestions.\nFirst, for an upcoming discussion phase, let us emphasize that **our main claim is not**\n> explores the dynamical effects of adding normalization in the recurrent self-attention dynamics\n\nnor\n> They find that adding normalization stabilizes the dynamics (spectral radius < 1) and makes the model operate close to “criticality,”\n\nRather, our take-home message is a broader perspective: to move beyond the constraints of energy-based formulations and instead utilize the Jacobian and Lyapunov exponents as fundamental tools to understand self-attention (SA) dynamics more comprehensively. The above two results are just concrete examples of our perspective.\n\n---\n> **W1**: ... I believe Lyapunov’s direct method can be used for studying non-stationary dynamics. In fact, the type of oscillatory dynamics they consider is stationary from the point of view of dynamical systems (there is no limit cycle behavior here).\n\nWe fully agree with your perspective, and our work does not in any way contradict it. We argue that an indirect Jacobian-based method is **more useful** for analyzing the stability and oscillatory behavior of self-attention (SA). While Lyapunov functions can exist even for non-stationary states, constructing them is generally difficult and no universal method is known. For oscillatory analysis of SA, Lyapunov functions are only known under highly restricted settings (e.g., the weight matrix is the identity or symmetric matrix [1], and all tokens share the same natural frequency [2]).\n\n[1] Geshkovski et al. \"The emergence of clusters in self-attention dynamics\", NeurIPS 2023.\n\n[2] Miyato et al., \"Artificial Kuramoto Oscillatory Neurons\", ICLR 2025.\n\n---\n> **W2**: The finding that adding normalization stabilizes the dynamics of a recurrent circuit has recently been made: [Rawat, S., et al., (2024). NeurIPS 37, 14712-14750]. In fact, normalization has been linked to the “critical” phenomena [Morone, F., et al., (2025). bioRxiv, 2025-05].\n\nThank you for sharing these excellent works. \n\nWhile their models and the form of normalization differ from ours, we believe that the importance of normalization is a fundamental aspect common across various systems. Our contribution is to theoretically and empirically validate this importance **in SA** and to be the first to highlight its significance in modern SA models.\n\nWe have added the citations in the extended related work section as follows. \n\n\"Rawat et al. [2024] showed that dynamic divisive normalization ensures the local stability of ORGaNICs via Lyapunov’s indirect method, improving interpretability and trainability.\"\n \n\"Morone et al. [2025] showed that normalization keeps recurrent networks stable beyond the unit spectral radius, and that slower responses indicate an approaching loss of stability.\"\n\n---\n> **W3**: I think it’s interesting that adding normalization leads to critical behavior, but the authors do not give an intuition of why this happens. Also, what is so special about criticality that it enables good performance in the model considered? It might be out of the scope of this paper, but I would appreciate more intuition or some empirical results on why it might be a desirable property.\n\nThanks for your insightful comment. \nIn the revised version, we will add the following explanation:\n\nNote that the Jacobian’s spectral norm corresponds to the Lipschitz constant of the SA function. Normalization suppresses excessively large Lipschitz constants, but if the constant is too small, the SA function becomes overly smooth and loses expressive power. This suggests that, as long as training and inference can progress, it is desirable to allow the spectral norm to be as large as possible. Thus, the model is driven toward criticality when a model without a Lyapunov function can attain a non‑negative maximum Lyapunov exponent. This aligns with the picture suggested by mean‑field theory, where the model’s representational power is limited in the ordered phase [3].\n\n[3] Poole et al., \"Exponential expressivity in deep neural networks through transient chaos\", NeurIPS 2016.\n\n---\n> **Q1**: How do the authors define “criticality” in the context of the dynamical system considered? The authors mention it’s the edge of stability and chaos, but do the systems that they consider exhibit chaos if sent beyond the tipping point?\n\nWe understand that “criticality” and “chaos” are polysemous terms in general. In our study, we define “criticality” as the point at which the largest Lyapunov exponent takes zero. In our experiments, when the maximum Lyapunov exponent gets slightly positive, we observed behaviors consistent with widely accepted notions of chaos, particularly **sensitivity to initial conditions** as follows.\n\n**Experiments**\n\nTo evaluate sensitivity to initial conditions, we first ran the input through the loop for T=128 steps. After this, we introduced a small Gaussian perturbation to the input at the beginning of the loop.\n\nThe following table shows the element-wise L1 distance between the hidden representations of the original and perturbed samples over time. We take the element that maximizes the element-wise distance.  \n\nWe observed that the distance between the original and perturbed samples exponentially increased over time, indicating chaotic behavior in the system.\n\n| Step | 0       | 10       | 100      | 500      |1000   |\n|------|----------|----------|----------|----------|----------|\n| Distance | 3.83e-04 | 2.82e-04 | 2.22e-03 | 1.91e-01 |2.64e+01|\n\n---\n> **Q2**: I think the caption of Fig. 5 is wrong. It should be the other way round, right?\n\n> **Q3**: There’s a typo on line 196: singular value.\n\nThank you for catching our mistakes. We will correct them accordingly.\n\n---\nWe welcome any further questions or requests for clarification to increase the significance of our paper."}, {"review_id": "gfcIa6E0bq", "reviewer": "Reviewer_jJ1c", "summary": "This paper generalizes the theoretical understanding of self-attention in Transformers by overcoming the previous restrictive energy-based formulations to Jacobian-based analysis. This work first relaxes classical assumptions in energy-based models, e.g., symmetric weight matrices and single-head constraints, and experimentally shows that energy-based regularizers are harmful in practice. The key contribution is in analyzing the Jacobian matrix of self-attention updates, which shows the importance of normalization layers for stabilizing predictions by controlling the spectral properties and thus suppressing oscillatory instabilities in discrete-time systems. Through experimentation, mainly in the context of Sudoku puzzles, this work reveals that good-performing models are close to criticality with Lyapunov exponents close to zero, indicating that maximal performance is achieved at the `edge of chaos' where dynamics do not blow up nor tend to zero. Further, this work also gives a theoretical foundation for the necessity of normalization operations in Transformers, provides novel spectral regularization techniques that yield better performance and presents a Jacobian-driven explanation of observed \"pseudo-energy\" effects. However, the practical impact is limited by the focus on simplified recurrent self-attention architectures that exclude key components of real Transformers like positional encoding, attention masking, and MLP blocks, with evaluation restricted primarily to toy problems rather than realistic NLP tasks.", "questions": "1. The paper demonstrates that energy-based regularization impairs performance, raising fundamental questions about the entire energy-based research program for understanding Transformers. This negative result warrants much deeper investigation.\n\n2. The theoretical insights are primarily validated on Sudoku puzzles. Can the authors demonstrate that the criticality findings (Lyapunov exponents ≈ 0) and normalization stabilization effects hold for at least one realistic NLP task where Transformers excel? What predictions does the proposed framework make regarding performance on language modeling or other standard benchmarks?\n\n3. The authors acknowledge that Proposition 5.1 provides \"conservative\" bounds that are \"looser than empirical observations.\" Can they derive tighter bounds that match the empirical O(1) scaling behavior shown in Figure 1? What architectural design principles can practitioners extract from these loose bounds if not?\n\n4. The proposed Lyapunov analysis uses $T=16$ iterations to approximate asymptotic behavior. Can the authors demonstrate that this finite-time approximation reliably predicts long-term stability properties? How sensitive are the critical findings to the choice of $T$, and what theoretical justification supports using finite-time approximations?\n\n5. The authors’ analysis focuses on RMSNorm, but doesn't systematically compare other normalization schemes. Can the authors show how LayerNorm, BatchNorm, or other variants affect the Jacobian properties they study? What principled guidelines does their framework provide for choosing normalization parameters in practice?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 1, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\n1. Quality: This paper provides rigorous mathematical development supported by solid theoretical reasoning, formal claims with comprehensive proofs, and consistent empirical support across architectures.\n\n2. Clarity: This paper is well-structured progression from energy-based to Jacobian analysis with clear explanations of key concepts like criticality and spectral stabilization\n\n3. Significance: This paper addresses fundamental questions about Transformer dynamics, provides theoretical grounding for normalization importance, and offers valuable negative results that could redirect energy-based research. \n\n4. Originality: This paper introduces a novel Jacobian-based framework that moves beyond restrictive energy formulations, original extension to multi-head attention, and fresh perspective connecting Lyapunov exponents to self-attention performance.\n\n\n**Weaknesses:**\n\n1. Quality: The experimental validation in this paper is limited primarily to Sudoku puzzles and undermines generalizability claims, as this narrow combinatorial task may not represent the diverse computational demands of real-world Transformer applications.\n\n2. Clarity: Mathematical notation could be simplified; some figures need a more straightforward presentation; and experimental details scattered across appendices reduce readability. \n\n3. Significance: Small-scale experiments with unclear scalability constraints practical impact to modern large Transformers, lack of validation on realistic NLP tasks, and missing analysis of computational overhead.\n\n4. Originality: Insufficient differentiation from existing edge-of-chaos literature, limited connection to gradient flow dynamics, and somewhat contrived experimental design choices.", "was_revised": true, "final_justification": "I appreciate the authors' responses, which have resolved most of my concerns. I have decided to raise my score.", "rebuttal": "Thank you for your kind reading and helpful feedback.\n\n---\n> **W1**: Quality: The experimental validation in this paper is limited primarily to Sudoku puzzles and undermines generalizability claims, as this narrow combinatorial task may not represent the diverse computational demands of real-world Transformer applications.\n\n> **Q2**: The theoretical insights are primarily validated on Sudoku puzzles. Can the authors demonstrate that the criticality findings (Lyapunov exponents ≈ 0) and normalization stabilization effects hold for at least one realistic NLP task where Transformers excel? What predictions does the proposed framework make regarding performance on language modeling or other standard benchmarks?\n\nFollowing your advice, **we have measured the maximum Lyapunov exponent in an NLP setting**. Specifically, we used ALBERT [1], which employs a looped architecture, and computed the Lyapunov exponent over 12 loops. We compared the exponents of the pretrained model and a randomly initialized model using five text samples.\n\nAs shown in the table below, the Lyapunov exponent of the pretrained model is close to 0, indicating that the pretrained model operates closer to criticality.\n\nRegarding the Sudoku task, several studies (e.g., [2]) have focused on it because achieving high performance is far from trivial, especially for the OOD case. In particular, Miyato et al. [3] also use the Sudoku dataset to evaluate test-time computation tasks. More recently, this interest has even spawned a new benchmark for evaluating LLM reasoning abilities [4]. We expect our work to provide a solid baseline for exploring such more complex reasoning tasks.\n\n| Model        | Maximum Lyapunov Exponent (mean ± std) |\n|--------------|-------------------------------|\n| Initialized  | 0.160 ± 0.003                    |\n| Pretrained   | 0.142 ± 0.015                    |\n\n[1] Lan et al., \"ALBERT: A Lite BERT for Self-supervised Learning of Language Representations\", ICLR 2020.\n\n[2] Wang et al., \"SATNet: Bridging deep learning and logical reasoning using a differentiable satisfiability solver\", ICML 2019.\n\n[3] Miyato et al., \"Artificial Kuramoto Oscillatory Neurons\", ICLR 2025.\n\n[4] Seely et al., \"Sudoku-Bench: Evaluating creative reasoning with Sudoku variants\", Preprint.\n\n---\n> **Q1**: The paper demonstrates that energy-based regularization impairs performance, raising fundamental questions about the entire energy-based research program for understanding Transformers. This negative result warrants much deeper investigation.\n\nWe have thoroughly **searched the regularization coefficient**, and within the range of our experiments, the observed behavior in the figures remained consistent regardless of model size. Although we demonstrated this typical behavior across hyperparameters in the main manuscript, we will add such (minor but) exhaustive figure variants in the appendix of the revised manuscript.\n\n\n\n---\n> **W3**: Significance: Small-scale experiments with unclear scalability constraints practical impact to modern large Transformers, lack of validation on realistic NLP tasks, and missing analysis of computational overhead.\n\nOur primary goal is to reveal the limitations of existing energy-based or energy-inspired models, **focusing on understanding** rather than building large-scale models that achieve high performance without interpretability.\n\n---\n> **W4**: Originality: Insufficient differentiation from existing edge-of-chaos literature, limited connection to gradient flow dynamics, and somewhat contrived experimental design choices.\n\nWe position our work **as the first to explicitly connect SA to the edge-of-chaos literature**.\nModels such as ItrSA and AKOrN exhibit weak chaotic behavior, and time-series plots reveal oscillatory dynamics.\n Our systematic analysis introduces novel perspectives that, to the best of our knowledge, have not been explored before. \n\nRegarding gradient flow dynamics, it is theoretically known that in energy‑based SA models (under ideal constraints such as identity weights), the flow exhibits metastable cluster states but eventually converges to a rank‑collapse state (e.g., [5]). In contrast, our results suggest that SA models closer to reality possess fundamentally more dynamic behavior, including oscillations. Theoretical analysis of such dynamics remains limited, and our work aims to open new directions for theoretical research, which we hope will inspire future studies.\n\n[5] Bruno et al., \"Emergence of meta-stable clustering in mean-field transformer models\", ICLR 2025.\n\n---\n> **Q3**: The authors acknowledge that Proposition 5.1 provides \"conservative\" bounds that are \"looser than empirical observations.\" Can they derive tighter bounds that match the empirical O(1) scaling behavior shown in Figure 1? What architectural design principles can practitioners extract from these loose bounds if not?\n\nIt seems that our inadequate description causes confusion about the limitation on the upper bound in theory.\nHere, **\"the current theoretical bound\" refers not to ours** but to the spectral norm of SA evaluated in Castin et al. [6] (Eq. (15)). \n\nCastin et al. reported that for a single SA without normalizations, the spectral bound is of O((#tokens)^{1/2}) in the current theory, but the trained models in real experiments show the spectral norm of  O((#tokens)^{1/4}). That is, the quantitative gap between theory and experiment itself worked as their contribution. \n\nIn a similar way, we observed that for the normalized case, the theoretical upper bound is of O((#tokens)^{1/2}) and looser than the empirical evaluation. But, what's noteworthy here is that normalization keeps the spectral norm at **O(1), which is smaller than O((#tokens)^{1/4})**!  Thus, Figure 1 empirically highlights the architectural importance of normalization for the token size. \n Although the current theory of the SA spectral bound cannot yet explain these empirical results, we believe that the accumulation of such empirical facts stimulates both the understanding of SA and the development of its theory.\n\n[6] Castin et al., \"How Smooth Is Attention?\", ICML 2024.\n\n---\n> **Q4**: The proposed Lyapunov analysis uses iterations to approximate asymptotic behavior. Can the authors demonstrate that this finite-time approximation reliably predicts long-term stability properties? How sensitive are the critical findings to the choice of, and what theoretical justification supports using finite-time approximations?\n\nSince computing the infinite sum is infeasible, we approximate the Lyapunov exponent with a finite sum. Prior studies measuring deep neural networks as dynamical systems [7] have also used finite time intervals (typically 5–15 steps) for similar analyses. It is also noteworthy that, as you can see in Figure 3, the exponents of AKOrNs are concentrated around 0 with very small fluctuations. \n This implies that using T=16 steps is a rational estimation with low variance.\n\n[7] Liu et al., \"Exploiting chaotic dynamics as deep neural networks\", Physical Review Research, 7(3), 2025.\n\n---\n> **Q5**: The authors' analysis focuses on RMSNorm, but doesn't systematically compare other normalization schemes. Can the authors show how LayerNorm, BatchNorm, or other variants affect the Jacobian properties they study? What principled guidelines does their framework provide for choosing normalization parameters in practice?\n\nWhile we focus on RMSNorm, our analysis is motivated by normalization schemes commonly used in related models such as Hopfield networks and AKOrN, where RMSNorm serves as a general form. In practice, SA often employs RMSNorm or LayerNorm, which differs from the RMSNorm mainly by centering, so our conclusions largely carry over. BatchNorm is out of scope because it is not used in these related models.  \n\nRegarding the choice of normalization parameters, our framework suggests setting them so that the largest Lyapunov exponent remains close to zero, ensuring stability while preserving expressivity. \nHere, one interesting observation is that, as we reported in line 273, \"Empirically, we confirmed that the trained $\\gamma$ remains small (see Table S.2), eliminating the need for explicit clipping such as $|\\gamma_i| \\leq 1$\". That is, we empirically confirmed that the training of normalization parameters can automatically obey the guiding principle of the criticality of the largest Lyapunov exponent. \n\n---\nWe welcome any further questions or requests for clarification to increase the significance of our paper."}, {"review_id": "1MiDZWwImu", "reviewer": "Reviewer_VJDA", "summary": "The authors describe a study of self-attention dynamics in  self-attention (SA) architectures. The study goes beyond existing work which focused on symmetric projection weights, continuous time dynamics and multihead attention. The authors achieve this by studying the spectral norm of the Jacobian of the multi head attention state update. The results were tested on looped  transformer architectures, which are studied for learning iterative processes.", "questions": "That method (i) fails to improve performance: can you relate this to the same reason why symmetric SA also does not perform well? How would the resulting networks look like in figure 3?\n\nline 277: \"specific conditions\" for readability, specify what these conditions are here (refer to eq (13)). Proposition 4.2 makes no statement on W^o, but it is present in eq (19). Explain. \n\nThe extension to multi-heads is trivial, except for the orthogonality condition. Can you explain further why this orthogonalization is necessary? What does this imply on the expressive power of the model?\n\nRelatedly: \"practical Transformers, as they typically assume a (similar but different) low-rank structure as well.\" Why is this the case? What is this assumption? I assume this can be done for fine-tuning and efficiency (compression) but the statement requires support.\n\nEq. (9) how is adding the input different from the residual connection already present in token mixers?\n\nEq (10) and Eq(12) the terminology with/without projection is confusing. Which projection is meant there? The equations are exactly the same, except for a head subscript. \n\n\"Figure 1 demonstrates that the spectral norm of the untrained SA’s Jacobian can be effectively reduced through normalization\", what about trained SA?\n\n\"although these constraints reduce their trainability,\" Vague statement. I'm not convinced by this statement, possibly because trainability implies convergence. Lyapunov stability doesn't necessarily make training unstable. I think what the authors meant is that the expressive power of the model is reduced/constrained.\n\n\"test-time scaling\", please define at least once the axis of scaling (assumed loops, but could also be S number of tokens)\n\nSection 5.3. Results of experiments are described without stating what the experiment is.\n\n**Minor points (no need to address in rebuttal):**\nPlease define Proj_X in equation 3\n\n\"concatination\"\n\"interpretted\"\n\"singular vale\"\n\nThe caption of figure 5 is not consistent with the text lines 293-302.\n\nSubtitle of section 5.1 \"normalization of spectral norm\" do not make sense to me. Proposition 5.1 is the spectral norm of [rms] normalization, not the other way around.\n\n\"the setting where tokens are not split into oscillators\", Describe what splitting into oscillators mean.\n\nMissing x axis label in figure 3", "limitations": "\"we focused on recurrent SAs without positional encoding, masking, or MLP blocks, ... an interesting direction\" Understood, but would the theory be tractable, applicable in this case?", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- Overal an interesting study and an informative read. Simple and clear study of the effect of normalization on the eigenvalue of update Jacobians. ItrSA formulation driven by Jacobian. I'm convinced that the Jacobian approach is much more practical than energy-based modeling approach in practice.\n- Regularization following the spectral properties of the Jacobian is a useful contribution.\n\nWeaknesses:\n- Several statement lack sharpness or are confusing, impacting negatively the readability of the paper. Some typos would have been caught with a simple spellchecker...\n- The limitation in experiments and theory to iterative problems, so no empirical benchmarks for language modeling is limited.", "was_revised": true, "final_justification": "The authors clarified some of my concerns. Like other reviewers, I'm concerned with the limited scope in experimentation, but IMP nevertheless it is an interesting study worth publication. I have increased my score accordingly.", "rebuttal": "Thank you for your insightful and constructive comments, especially for acknowledging that the Jacobian approach is more practical than energy-based modeling.\n\n---\n> **Q1**: That method (i) fails to improve performance: can you relate this to the same reason why symmetric SA also does not perform well? How would the resulting networks look like in figure 3?\n\n\nWe empirically observed that both method (i) and symmetric SA exhibit optimization difficulties, and we hypothesize that constraining the weights reduces the expressive power of self‑attention. Accordingly, we will include the figure for method (i), corresponding to Figure 3, in the revised manuscript.\n\n---\n> **Q2**: line 277: \"specific conditions\" for readability, specify what these conditions are here (refer to eq (13)). Proposition 4.2 makes no statement on W^o, but it is present in eq (19). Explain.\n\nSpecifically, when $R_{\\text{E-single}} = 0$, setting $W^{V} = W_{1}^{V} W_{1}^{O}$ satisfies the condition on  $W^{V}$ in Proposition 4.2; that is, the $W^{V}$ in Eq. (13) corresponds to $W_{1}^{V} W_{1}^{O}$ in Eq. (19).\nAs stated in line 285, we provide a detailed explanation in Appendix B.2. \n\n---\n> **Q3**: The extension to multi-heads is trivial, except for the orthogonality condition. Can you explain further why this orthogonalization is necessary? What does this imply on the expressive power of the model?\n\nOur goal is to highlight the limitations of existing single‑head energy formulations, rather than to propose new energy constructions. As L822–L823 suggests, orthogonality is a natural **sufficient** condition to guarantee the energy, since it makes the interaction between different heads vanish and reduces to the calculation of the single‑head case. In other words, finding a condition without orthogonality seems to require a fundamentally different approach from the original energy construction of the single-head case, which is out of our scope.\n\nAlthough a theoretical analysis of expressive power is beyond our scope, the above derivation of the multi‑head energy implies that interactions among different heads, which generally arise when orthogonality is not enforced and can increase the original energy function, are important for enhancing the model’s expressivity. We expect that our analysis will stimulate subsequent work and provide opportunities for further analyses of SA.\n\n---\n> **Q4**: Relatedly: \"practical Transformers, as they typically assume a (similar but different) low-rank structure as well.\" Why is this the case? What is this assumption? I assume this can be done for fine-tuning and efficiency (compression) but the statement requires support.\n\nWhat we meant was that both our energy-based formulation with orthogonality (Eq. 13) and practical Transformers have a similar low-rank structure, in the sense that **the embedding dimension D is H times larger (where H is the number of heads)**. \nIn almost all practical Transformers, we set $W_{h}^{Q}, W_{h}^{K} \\in \\mathbb{R}^{d_h* H \\times d_h}$ where $D=d_h*H$ and $d_h$ denotes the embedding dimension of each head. Consequently, the multi-head SA naturally makes the product $W_{h}^{Q}, W_{h}^{K \\top}$  low-rank. \nIn our theoretical energy-based formulation, this low-rankness is similar. But, the orthogonality is not required in practice, and our formulation is different from practical models in this sense. This is what we meant in the sentence of \"(similar but different) low-rank structure\".  We will add the above detailed description of low‑rankness in the revised manuscript.\n\n\n---\n> **Q5**: Eq. (9) how is adding the input different from the residual connection already present in token mixers?\n\nThe point is that adding the input $C$ in Eq. (9) is conceptually different from the residual connection typically used in token mixers. While a standard residual connection only passes the transformed features forward, explicitly adding the input $C$ into the loop introduces a direct pathway for the original signal at each iteration, which has been shown to improve test-time performance [1].\n\n[1] Saunshi et al., “Reasoning with Latent Thoughts: On the Power of Looped Transformers,” ICLR 2025.\n\n---\n> **Q6**: Eq (10) and Eq(12) the terminology with/without projection is confusing. Which projection is meant there? The equations are exactly the same, except for a head subscript.\n\nThe terminology “with/without projection” refers to whether the projection matrices (as defined in Eq. (3)) are applied.  \nSpecifically, The single-head energy function in Eq. (10) corresponds to the case where the projection matrices are used, while the multi-head energy function in Eq. (12) corresponds to the unprojected form  ($\\frac{dX}{dt} = \\sum_{h=1}^{H} SA_{h}(X)$), where no projection is applied.\n\n---\n> **Q7**: \"Figure 1 demonstrates that the spectral norm of the untrained SA’s Jacobian can be effectively reduced through normalization\", what about trained SA?\n\nWe empirically observed that it is hard to decrease the training loss of the models without any normalization, so Figure 1 compares the untrained SA to highlight the effect of normalization on the spectral norm of the Jacobian. For trained SA, normalization remains crucial not only for stabilizing the Jacobian but also for enabling successful optimization, which is why we focus on the untrained case to isolate the impact of normalization.\n\n---\n> **Q8**: \"although these constraints reduce their trainability,\" Vague statement. I'm not convinced by this statement, possibly because trainability implies convergence. Lyapunov stability doesn't necessarily make training unstable. I think what the authors meant is that the expressive power of the model is reduced/constrained.\n\nThank you for catching this. Following your suggestion, we have removed the phrase “although these constraints reduce their trainability”. As you say, what we meant was that the expressive power of the model is constrained.\n\n---\n> **Q9**: \"test-time scaling\", please define at least once the axis of scaling (assumed loops, but could also be S number of tokens)\n\nWe define “test-time scaling” as the improvement in performance when the number of iterations (loops) during inference is increased. We will add this definition at its first occurrence to avoid ambiguity, clarifying that the scaling axis refers to the number of loops rather than the number of tokens $S$.\n\n---\n> **Q10**: Section 5.3. Results of experiments are described without stating what the experiment is.\n\nWe acknowledge the concern. In Section 5.3, we describe the experimental results shown in Figure 3, with detailed settings and methodologies provided in Appendix B.3. To improve clarity, we will revise the section to briefly summarize the experiment setup before presenting the results.\n\n---\nWe welcome any further questions or requests for clarification to increase the significance of our paper."}]

Author rebuttal:

Meta-review: This paper presents a dynamical systems analysis of recurrent self-attention (SA) architectures, moving beyond the more restrictive assumptions of prior energy-based analyses. As noted by the reviewers, the paper's main claims are that normalization layers play a crucial stabilizing role by controlling the eigenvalues of the state update Jacobian, and that high-performing SA models operate near a critical state, or the "edge of chaos," as evidenced by their Lyapunov exponents being close to zero. The primary strengths of this work are its novel theoretical perspective on a widely used mechanism, providing valuable insights into the behavior of recurrent SA. The theoretical results are complemented by a set of experiments that empirically validate the main claims and explore practical applications like regularization. The main weaknesses, highlighted by reviewers wxyz and abcd, are the focus on a simplified setting that omits components like MLP blocks or positional encodings, and the fact that some of the theoretical bounds are acknowledged to be conservative.

The discussion during the rebuttal period was productive. Reviewers raised important questions regarding the simplified architectural assumptions and the practical tightness of the theoretical bounds. Reviewer efgh also correctly pointed out a significant typo in the caption of Figure 5, which contradicted the plotted results and the main text. The authors effectively addressed these points by acknowledging the limitations as important directions for future work, clarifying the novelty of their analysis in the context of SA, and committing to correcting the erroneous figure caption. My decision to accept is based on the paper's conceptual contribution; it provides a fresh and insightful "energy-agnostic" lens for understanding the dynamics of self-attention. The connection drawn between normalization, Jacobian stability, and criticality is an important finding that advances our theoretical understanding of these complex systems.

Although the final scores are somewhat borderline, I believe the paper's contributions are of sufficient interest to the NeurIPS community. The work is well-motivated, the analysis is sound, and the findings could inspire new approaches to designing and training more stable and performant recurrent architectures. I strongly encourage the authors to incorporate the detailed feedback from all reviewers into the final version of their paper, paying special attention to clarifying the limitations of their analysis and correcting the figure caption as promised.
Decision: Accept (poster)