Introduction
✨ GitHub repository:
mcbal/neqnn(work in progress)
Transformers are powerful driven dynamical systems, yet their internal computation is rarely discussed in terms of nonequilibrium thermodynamics. Building on the dynamical mean-field theory framework developed in Spin-Model Transformers (2023), we design a minimal parallel transformer-like module whose forward pass performs one or more mean-field updates following a drive quench, separating out fast state relaxation, changing external drive, and slow parameter learning.
We characterize the spin-model transformer module’s dynamical regimes to map out a design space of stateless and stateful variations of transformer-like and deep-equilibrium-like architectures and leverage its spin-model backing to compute differentiable proxies for housekeeping entropy production and post-quench relaxation mismatch. The spin-model transformer module acts as a controllable laboratory in which architectural dynamics, irreversibility diagnostics, and candidate local learning protocols can be studied together.
Driving a spin-model transformer module
In this section we design a minimal spin-model transformer module whose forward pass implements a controllable nonequilibrium quench-and-relax process. We identify three timescales and two dynamical regimes, leading to a natural categorization of the design space into stateless and stateful variations of transformer-like and deep-equilibrium-like architectures.
A minimal controllable drive-conditioned system
In Spin-Model Transformers (2023) we showed how to apply dynamical mean-field theory to approximate the time-dependent behavior of asymmetric vector-spin models. We started from a spin system of $N$ vector spins $\mathbf{s}_{i,t} \in \mathbb{R}^{D}$ talking to each other via an $N \times N$ pairwise coupling matrix $J_{ij}$ with the underlying synchronous parallel-updates stochastic dynamics characterized by a discrete-time Markov chain transition probability. External magnetic fields $\mathbf{x}_{i,t} \in \mathbb{R}^{D}$ bias the vector spins and act as local drives. We use the notation $\mathbf{X}_{t} \in \mathbb{R}^{N \times D}$ and $\mathbf{M}_{t,k} \in \mathbb{R}^{N \times D}$ to refer to matrices of stacked vectors.

Our vector-spin system of interest is defined by a first-order Plefka[t-1,t]1 mean-field recurrence relation for its spin expectation values, or magnetizations
where $\beta$ denotes the inverse temperature and $R^2=D/2 -1$ is a natural large-$D$ hyperspherical length scale2. The effective field
\begin{equation} \mathbf{h}_{i,t,k} = \mathbf{x}_{i,t} + f_{\boldsymbol{\theta}_{\mathrm{FFN}}}\left( \mathbf{x}_{i,t} \right) + \sum_{j} J_{ij} (\mathbf{X}_{t}) \mathbf{m}_{j,t,k} \end{equation}contains a parameterized drive-dependent coupling rule3
\begin{equation} \mathbf{J} (\mathbf{X}_{t}) = \mathrm{softmax}\left( \mathbf{X}_{t} \boldsymbol{W}_{Q} \boldsymbol{W}_{K}^{T} \mathbf{X}_{t}^{T} \right), \label{eq:softmax} \end{equation}and a parameterized position-wise non-linear drive-dependent field $f_{\boldsymbol{\theta}_{\mathrm{FFN}}}\left( \mathbf{x}_{i,t} \right)$ added to the local drive $\mathbf{x}_{i,t}$. Each vector spin effectively experiences a local mean-field that is the sum of a residual stream drive, a feed-forward-like drive, and attention-like couplings. The time index $t$ tracks changes in the external drive while $k$ indexes an internal relaxation step.
By making the effective drive as well as the couplings depend on the drive $\mathbf{X}_{t}$, a sudden shift $\mathbf{X}_{t} \to \mathbf{X}_{t+1}$ changes both the local fields as well as the interactions and quenches the system into a new instantaneous dynamics4. During internal relaxation iterations, the drive $\mathbf{X}_{t}$ and the parameters $\boldsymbol{\theta} = \{ \mathbf{W}_{Q}, \mathbf{W}_{K}, \boldsymbol{\theta}_{\mathrm{FFN}} \}$ are held fixed.
Approximation stack: At the stochastic level, the process looks something like $P_{\boldsymbol{\theta}, \mathbf{X}_{t}}(\mathbf{S}_{t, k+1} | \mathbf{S}_{t, k}) = \prod^{N}_{i=1} \frac{\operatorname{exp}\left[\beta \mathbf{s}_{i,t,k+1} \cdot \mathbf{h}_{i,t,k} \right]}{Z\left(\mathbf{h}_{i,t,k}\right)}$ with effective field $\mathbf{h}_{i,t,k} = \mathbf{x}_{i,t} + f_{\boldsymbol{\theta}_{\mathrm{FFN}}}\left( \mathbf{x}_{i,t} \right) + \sum_{j} J_{ij} (\mathbf{X}_{t}) \mathbf{s}_{j,t,k}$ and single-site normalization constant $Z\left(\mathbf{h}\right)$. The product-form kernel assumes parallel, conditionally independent site updates. An asynchronous spin process would need a different kernel and should not be conflated with this one. The stochastic process is the thermodynamic object, and the magnetization recurrence used as the neural-network module is a Plefka mean-field approximation to its one-point dynamics. In the next sections, we will, among other quantities, discuss deterministic mean-field fixed-point magnetizations $\mathbf{M}^{\star}_{t}(\mathbf{X}_{t})$ which approximate one-point marginals of the stochastic steady state $\pi_{\mathbf{X_{t}}}$. It is useful to keep these levels of approximation in the back of your mind.
We end up with a highly reconfigurable system that is dynamically shaped by the drive. These drive dependencies enable the mean-field system to encode subtle correlational structures. The system’s parameters can be shaped through training in an outer optimization loop to make an already responsive system adaptive by controlling how the system responds, fluctuates, and relaxes after a quench.
Building modules: three clocks, slow plasticity, and two relaxation limits
Looking at Eq. \eqref{eq:paralleltransformer} we notice its close resemblance to the forward pass of a parallel transformer block. To make this more precise, we need to specify how to implement spin-model transformer modules in practice. What is up with this weird internal state and internal relaxation dimension? How can this system even serve as a neural network module?
Let us begin by writing the forward pass Eq. \eqref{eq:paralleltransformer} more generally as
\begin{equation} \mathbf{M}^{(\ell)}_{t, k+1} = F_{\boldsymbol{\theta}^{(\ell)}_{n}} \left( \mathbf{X}^{(\ell)}_{t}, \mathbf{M}^{(\ell)}_{t, k} \right) \end{equation}and clearly state the clocks involved:
- $k$ indexes fast internal relaxation within the module
- $t$ indexes changes in the environmental drive or input context
- $n$ indexes slow parameter updates $\boldsymbol{\theta}^{(\ell)}_{n+1} = \boldsymbol{\theta}^{(\ell)}_{n} - \eta \nabla_{\boldsymbol{\theta}^{(\ell)}} \mathcal{L}^{\ell}_{t}$ for some learning rate $\eta \ll 1$ and (potentially layer-dependent) loss function $\mathcal{L}^{\ell}_{t}$ where slow here actually means small parameter changes since the optimizer clock $n$ often tracks the drive clock $t$ in practice
The layer index $\ell$ indexes network depth (number of stacked layers), which, in our framework, identifies members of a collective of different driven spin systems driving each other. The layer index does not (have to) correspond to external time nor to internal relaxation time; it is an additional axis labeling the simple feed-forward topology (depth) of the computational graph. For clarity, we drop it in the remainder of this section.
The three clocks define three characteristic timescales: $\tau_{\mathrm{relax}}$, $\tau_{\mathrm{drive}}$, and $\tau_{\mathrm{learn}}$. Throughout, we assume slow plasticity $\tau_{\mathrm{drive}} \ll \tau_{\mathrm{learn}}$. The relative size of $\tau_{\mathrm{relax}}$ and $\tau_{\mathrm{drive}}$, or, equivalently, the number of internal updates allocated before the next quench, determines the computational regime.
We frame our module-design intuition around a quench-and-relax scenario: when the input drive switches, i.e., $\mathbf{X}_{t-1} \to \mathbf{X}_{t}$, the spin system has to adapt to the sudden change. A general post-quench module then looks like
\begin{equation} \mathbf{M}_{t, K} = F^{K}_{\boldsymbol{\theta}_{n}} \left( \mathbf{X}_{t}, \mathbf{M}_{t, 0} \right), \end{equation}with two independent design choices: the number of internal relaxation steps $K$ (relaxation horizon) and the choice of $\mathbf{M}_{t, 0}$ (initialization policy), leading to the following design space:
| Reset or amortized initialization | Carried initialization | |
|---|---|---|
| Finite-step regime with $K < \infty$ | finite-depth, stateless, transformer-like module | recurrent stateful module |
| Fixed-point regime where $K \to \infty$ | implicit or deep-equilibrium-like (DEQ) module | identical if the fixed point is unique; path-dependent if mean-field branches coexist |
Finite-step regime
In this regime, only $K < \infty$ internal updates are allocated before the next quench, which may reflect genuine competition between relaxation and drive timescales, or simply deliberate computational truncation as in recent looped and recursive-reasoning approaches. The intuition here is that the system tracks a moving family of instantaneous stationary marginals with potentially nonzero lag. In this regime, there can be no strong separation between drive and relaxation if approaching the steady state takes more than $K$ steps. But, since the module is controllable, the outer loop can nudge the module’s parameters $\boldsymbol{\theta}_{n}$ towards more efficient and useful relaxation.
The initialization $\mathbf{M}_{t, 0} = \mathbf{M}_{t-1, K}$ makes the module architecture genuinely recurrent and stateful, but with a full context window of hidden states, situating it somewhere in between recurrent neural networks and transformers. Another option is to warm-start with a learned amortized initializer $\mathbf{M}_{t, 0} = \mathbf{X}_{t}\mathbf{W}_{V}$ for the post-quench relaxation, which estimates the drive-conditioned response to which the module should relax. If we interpret these initializations as a value stream, then, for $K=1$, the forward pass looks like a parallel transformer block.
Fixed-point regime
In this regime, $\tau_{\mathrm{relax}} \ll \tau_{\mathrm{drive}}$ so we consider $\mathbf{X}_{t}$ clamped and let $K \to \infty$ until the deterministic mean-field equations converge to fixed-point magnetizations $\mathbf{M}^{\star}_{t}(\mathbf{X}_{t})$ compatible with the frozen drive $\mathbf{X}_{t}$. These values approximate the mean-field marginals of the frozen-drive nonequilibrium steady state (NESS). The intuition here is that the clamped input fixes an instantaneous stochastic transition rule. Although its one-point marginals become stationary, asymmetric couplings can sustain probability currents and positive entropy production beneath those stationary marginals.
In case of a unique fixed point, the initial values $\mathbf{M}_{t, 0}$ are erased, and the module is stateless. But the deterministic mean-field equations may admit multiple stable fixed-point branches or basins. Warm-starting with $\mathbf{M}_{t, 0} = \mathbf{M}^{\star}_{t-1}$ can then produce path-dependent branch selection and hysteresis behavior.
The fixed-point regime has a simple sufficient stability condition that is directly related to the module’s fluctuations. One can show that the Jacobian eigenvalues of the radial mean-field update map $\phi_{\beta}$ introduced in Eq. \eqref{eq:paralleltransformer} reveal that the map is globally contractive whenever $\rho_{t} = \beta J_{t} /2<1$ with $J_{t} = \operatorname{max}_{i}\sum_{j}\left|J_{ij,t}\right|$. (For non-scalar couplings, a operator coupling norm factor appears instead.) In this regime the frozen drive selects a unique fixed point. For $\rho_{t} \geq 1$, convergence may still occur, but uniqueness, absence of cycles, and initialization independence are no longer guaranteed.
On the connection to transformers
The resemblance to a transformer forward pass should be understood as a plausibility bridge, not as evidence that trained transformers literally implement the nonequilibrium thermodynamics developed below. In the spin-model transformer module, a change in context quenches a drive-dependent transition rule; one or a few relaxation steps give a finite-depth transformer-like computation, while convergence gives an implicit or deep-equilibrium-like module. The practical value of this correspondence is that it turns transformer-shaped architectures into accelerator-friendly laboratories for driven many-body dynamics with structured, learnable couplings. Their mean-field observables must first be calibrated against exact stochastic systems at small scale, but the architectural mapping provides a route to controlled experiments beyond analytically tractable toy models.
We end this section with a cheat sheet mapping concepts between spin-model transformer modules and transformer modules.
| Spin-model analogue | Transformer component |
|---|---|
| Local drives $\mathbf{x}_{i,t}$ | Module input |
| “The drive” $\mathbf{X}_{t}$ | Current context window |
| Saturation of magnetization map $\varphi_{\beta}$ | Bounded response nonlinearity; analogous to normalization or gating |
| Parameterized couplings $J_{ij}$ | Attention matrix |
| Learned amortized initializer $\mathbf{M}_{t, 0} = \mathbf{X}_{t}\mathbf{W}_{V}$ | Amortized initial state; analogous to a value stream |
| Head-specific coupling matrices $J^{(h)}(\mathbf{X}_{t})$ acting on separate subspaces or parallel vector-spin systems | Multihead attention |
| Magnetizations $\mathbf{m}_{i,t}$ | Internal state and module output |
Steady-state circulation and post-quench mismatch
In this section, we show how the quench-and-relax process behind the forward pass of a spin-model transformer module relates to notions of irreversibility. We add just5 enough physical context to introduce and motivate differentiable proxies for housekeeping entropy production and post-quench relaxation mismatch, which we can compute at the same mean-field level as the spin-system dynamics.
Quench-and-relax protocol
During relaxation after a quench we hold the input drive $\mathbf{X}_{t}$ fixed and let the spin system settle. The average magnetizations $\mathbf{M}_{t}$ may stop changing, but its microscopic dynamics need not become reversible. Asymmetric couplings can sustain circulating probability currents when forward sequences of spin configurations remain more likely than their backward step reversals. We call this source of irreversibility steady-state irreversibility. Its entropy-production rate is the running cost of maintaining a nonequilibrium steady state under the current input drive. For our system, we can estimate this “housekeeping” entropy production from asymmetric couplings and delayed correlations (a vector-spin translation of Aguilera et al., 2020),
\begin{equation} \sigma^{\star}_{\mathrm{hk},t} = \beta \sum_{ij} \left(J_{ij}(\mathbf{X}_{t}) - J_{ji}(\mathbf{X}_{t})\right) D^{\star}_{ij,t} , \label{eq:sigma_hk} \end{equation}where $D^{\star}_{ij,t}$ denote the one-step delayed correlations evaluated under the stationary one-step joint law. Intuitively, this is like
\begin{equation} \sigma^{\star}_{\mathrm{hk},t} = \sum_{ij} \left[\operatorname{directionality}\right]_{ij} \times \left[\operatorname{delayed\ flow}\right]_{ij,t}. \end{equation}The couplings provide a directional bias; the delayed correlations report whether fluctuations actually propagate along that direction.
Now we quench again. After changing the input drive abruptly, $\mathbf{X}_{t} \to \mathbf{X}_{t+1}$, the system is still distributed approximately according to its old steady state $\pi_{t}$, while the new transition rule $P_{\boldsymbol{\theta}, \mathbf{X}_{t+1}}(\mathbf{s}_{t+1, k+1} | \mathbf{s}_{t+1, k})$ induced by the new input drive actually favors another steady state $\pi_{t+1}$. The relative-entropy distance to the new frozen-drive steady state
\begin{equation} \Delta_{t+1, k} = D_{\mathrm{KL}}\left(p_{t+1,k} \lVert \pi_{t+1} \right),\label{eq:vmfkl} \end{equation}measures how far the current distribution remains from the stationary distribution selected by the new frozen drive. The new housekeeping part remains after relaxation while $\Delta_{t+1, k} \to 0$ as the actual distribution relaxes to the new stationary distribution. At the exact Markov-process level, its decrease under a fixed transition rule is associated with nonadiabatic entropy production. Below, we replace both distributions by factorized mean-field approximations and use the resulting quantity as a diagnostic of post-quench lag.
A driven spin-model transformer module with asymmetric couplings can thus be approximately measured in two ways during the quench-and-relax process: the cost of “running” a nonequilibrium steady state after relaxation and the “catching up” during relaxation after its input drive changes. Housekeeping entropy production measures sustained asymmetric circulation under a fixed drive while mismatch measures the relaxation still required after a drive change.
Allocating budgets: Intuitively (at the level of the exact dynamics, not necessarily at the mean-field level), every quench-and-relax cycle has an “entropy budget”
\begin{equation} \Sigma_{\mathrm{cycle}} \approx \sum^{K}_{k=1} \sigma_{\mathrm{tot},t,k} = \underbrace{\sum^{K}_{k=1} \sigma_{\mathrm{hk},t,k}}_{\text{rent, transient rate}} + \underbrace{\left( \Delta_{t,0} - \Delta_{t,K} \right)}_{\text{moving costs}} \end{equation}where the transient housekeeping $\sigma_{\mathrm{hk},t,k} = \sigma_{\mathrm{tot},t,k} - (\Delta_{t,k} - \Delta_{t,k+1})$ can be obtained from the transient delayed correlations $D_{ij,t,k}$ (see Appendix A). The neural-network modules defined in the corners of the design space of the previous section allocate this “entropy budget” differently depending on the number of internal relaxation steps. The DEQ corner ($K \to \infty$) fully pays off the “moving costs” every cycle and then just pays rent. The transformer corner ($K=1$ with amortized initialization) barely pays rent and then carries residual mismatch into the next drive change. Looping layers pays off moving costs. With a carried state or a well-tuned amortized initializer, we can potentially reduce the opening balance of the next cycle.
Mean-field proxy for housekeeping entropy production
Evaluating Eq. \eqref{eq:sigma_hk} at the mean-field level (see Appendix A) leads to
\begin{equation} \sigma^{\star}_{\mathrm{hk},t} \approx \frac{\beta^2}{2} \sum_{ij} \left(J_{ij}(\mathbf{X}_{t}) - J_{ji}(\mathbf{X}_{t})\right)^2 C^{\star}_{ij,t} , \end{equation}where $C^{\star}_{ij,t} = \operatorname{Tr} \left( \Sigma^{\star}_{i,t} \Sigma^{\star}_{j,t} \right) \geq 0$ with $\Sigma_{i,t,k} = \operatorname{Cov} \left[ s_{i,t,k} \right]$ denoting the single-site covariances / susceptibilities. The latter is dominated by a $O(D)$ contribution that is angle-independent, so the proxy may mostly reflect coupling magnitudes and susceptibility rather than representational geometry, which is captured in subleading $O(1)$ terms. For unconstrained couplings, the proxy measures squared coupling nonreciprocity, weighted by how strongly the fluctuation spaces of the two sites overlap. Under strict causal masking, however, pairwise nonreciprocity is largely enforced by the architecture. The proxy then behaves more like a susceptibility-weighted attention-concentration statistic than a measurement of learned antisymmetry.
Mean-field proxy for post-quench relaxation mismatch
To evaluate Eq. \eqref{eq:vmfkl}, we use the fact that the mean-field vector-spin model machinery is built on the von Mises-Fisher distribution (see Appendix B), leading to the full factorized mean-field mismatch
\begin{equation}\Delta_{t,k}^{\mathrm{MF}}\approx\sum_i\left[R^2\log\frac{R^2-\lVert\mathbf m_{i,t}^{\star}\rVert^2}{R^2-\lVert\mathbf m_{i,t,k}\rVert^2}+\frac{2R^2\left(\lVert\mathbf m_{i,t}^{\star}\rVert^2-\mathbf m_{i,t,k}\cdot\mathbf m_{i,t}^{\star}\right)}{R^2-\lVert\mathbf m_{i,t}^{\star}\rVert^2}\right],\end{equation}where $\mathbf m_{i,t,k}$ is the current magnetization and $\mathbf m_{i,t}^{\star}$ is the instantaneous fixed-point magnetization. The expression vanishes when $\mathbf m_{i,t,k}=\mathbf m_{i,t}^{\star}$ at every site and penalizes both differences in magnetization norm and angular misalignment with the frozen-drive fixed-point response.
Numerical experiments
Mean-field proxy fidelity
Probing and measuring a spin-model transformer module
Do these dynamical quantities reveal anything useful about how such a system learns or adapts under a changing drive? Does mismatch drop because the module anticipates the response to new boundary conditions, or because the response landscape becomes dull and easy to predict?
Can relaxation supply a learning signal?
The spin-model transformer module defines a driven nonequilibrium module and two measurable aspects of its dynamics: sustained circulation under a fixed drive and relaxation following a change in the drive. Neither quantity, by itself, specifies a unique or useful learning rule. In particular, many different prediction, clamping, and parameter-update protocols could be built around the same dynamics.
Prediction and correction are two runs of the same dynamics under different clamping patterns: relax under a drive constructed before the next observation and you have a prediction; relax under the observed drive and you have a target. There are many ways to set up the interface protocols, and we suspect most of them collapse into trivially predictable dynamics.
Closed-loop adaptive behavior
Discussion and related work
Stochastic thermodynamics of driven steady states: …
Dynamical mean-field theory for asymmetric spin systems: …
Recurrent-depth and implicit neural computation: …
Prediction, physical learning, and driven adaptation: …
Conclusion and outlook
…
Use the mismatch measure $\Delta_{t, k}$ as a halting signal: relaxing until mismatch drops below threshold turns the relaxation horizon $K$ into adaptive test-time compute.
…
The same protocol extends naturally to collections of modules. A detached output from one module can serve as a time-dependent boundary condition for another, so modules can alternately act as system and environment. This would turn the single-module prediction–correction cycle into a sequence of local quenches through a layered or recurrent collective. We leave that extension until the behavior of the single-module objective is understood.
…
Acknowledgements
We acknowledge interesting back-and-forth discussions with Claude Opus 4.8, Claude Fable 5, GPT 5.5, and GPT 5.6.
References
A non-exhaustive list of references and inspiration includes:
- A unifying framework for mean-field theories of asymmetric kinetic Ising systems by Miguel Aguilera, S. Amin Moosavi, and Hideaki Shimazaki
- The thermodynamics of prediction by Susanne Still, David A. Sivak, Anthony J. Bell, and Gavin E. Crooks
- Three detailed fluctuation theorems by Massimiliano Esposito and Christian Van den Broeck
- Self-organized fine-tuned response in a driven spin glass by Jacob Mitchell Gold
If you happen to find this work useful, please consider citing it as:
@article{bal2026,
title = {Nonequilibrium Dynamics in Spin-Model Transformers},
author = {Bal, Matthias},
year = {2026},
month = {?},
url = {https://mcbal.github.io/post/nonequilibrium-dynamics-in-spin-model-transformers/}
}
Appendices
Appendix A: Mean-field delayed correlations and housekeeping entropy production
To evaluate Eq. \eqref{eq:sigma_hk}, we need stationary one-step delayed correlations $D^{\star}_{ij,t}$. To this end, let us first compute a first-order Plefka[t-1,t] mean-field approximation of the transient one-step delayed correlations $D_{ij,t,k}$,
Evaluating this expression as we did for the magnetizations in Spin-Model Transformers (2023), we end up with the mean-field approximation
\begin{align} D_{ij,t,k} \approx \beta J_{ij}(\mathbf{X}_{t}) \operatorname{Tr} \left( \Sigma_{i,t,k} \Sigma_{j,t,k-1} \right), \end{align}where $\Sigma_{i,t,k} = \operatorname{Cov} \left[ s_{i,t,k} \right]$ denotes the single-site covariance / susceptibility. The trace captures which directions on the vector-spin sphere are still available to fluctuate. If a spin is weakly magnetized, it has many soft directions. If it is strongly magnetized, many directions are suppressed because the spin is pinned close to its mean direction.
At stationarity, $p_{t,k-1} \to \pi_{t}$ and $\mathbf{m}_{i,t,k} \to \mathbf{m}^{\star}_{i,t}$ so that
\begin{align} D^{\star}_{ij,t} = \lim_{k\to\infty} D_{ij,t,k} = \mathbb{E}_{\mathbf{s} \sim \pi_{t}, \mathbf{s'} \sim P_{\boldsymbol{\theta}, \mathbf{X}_{t}}(\cdot | \mathbf{s})} \left[ \left( \mathbf{s'}_{i} - \mathbf{m}^{\star}_{i,t} \right) \cdot \left( \mathbf{s}_{j} - \mathbf{m}^{\star}_{j,t}\right) \right] \end{align}and the mean-field approximation becomes
\begin{align} D^{\star}_{ij,t} \approx \beta J_{ij}(\mathbf{X}_{t}) \operatorname{Tr} \left( \Sigma^{\star}_{i,t} \Sigma^{\star}_{j,t} \right), \end{align}Substituting the large-$D$ approximation from Spin-Model Transformers (2023),
\begin{align} \Sigma_{i,t} \approx \frac{\mathbf{I}_{D}}{1+\gamma_{i,t}} - \frac{\mathbf{m}_{i,t} \mathbf{m}_{i,t}^{T}}{R^2 \gamma_{i,t}}, \end{align}we end up with the $(i \leftrightarrow j)$-symmetric explicit expression
\begin{align} C^{\star}_{ij,t} &= \operatorname{Tr} \left( \Sigma^{\star}_{i,t} \Sigma^{\star}_{j,t} \right) \geq 0 \nonumber \\ &= \frac{D}{(1+\gamma^{\star}_{i,t})(1+\gamma^{\star}_{j,t})} \nonumber\\ &- \frac{\lVert \mathbf{m}^{\star}_{j,t} \rVert^2}{R^2 \gamma^{\star}_{j,t} \left( 1 + \gamma^{\star}_{i,t} \right)} \nonumber\\ &- \frac{\lVert \mathbf{m}^{\star}_{i,t} \rVert^2}{R^2 \gamma^{\star}_{i,t} \left( 1 + \gamma^{\star}_{j,t} \right)} \nonumber\\ &+ \frac{\left( \mathbf{m}^{\star}_{i,t} \cdot \mathbf{m}^{\star}_{j,t} \right)^2}{R^4 \gamma^{\star}_{i,t} \gamma^{\star}_{j,t}}, \end{align}where
\begin{align} \gamma^{\star}_{i,t} &= \sqrt{1 + \beta^2 \lVert \mathbf{h}^{\star}_{i,t} \rVert^2 / R^2 },\\ \mathbf{h}^{\star}_{i,t} &= \mathbf{x}_{i,t} + f\left( \mathbf{x}_{i,t} \right) + \sum_{l} J_{il} (\mathbf{X}_{t}) \mathbf{m}^{\star}_{l,t}. \end{align}Since $\gamma^{\star} \sim O(1)$, the one-step delayed correlations $D^{\star}_{ij,t}$ are dominated by a large angle-independent isotropic contribution $O(D)$ with the other three terms being $O(1)$. The housekeeping entropy production proxy in the Plefka[t-1,t] approximation becomes
or, using the symmetry of $C^{\star}_{ij,t}$,
\begin{equation} \sigma^{\star}_{\mathrm{hk},t} \approx \frac{\beta^2}{2} \sum_{ij} \left(J_{ij}(\mathbf{X}_{t}) - J_{ji}(\mathbf{X}_{t})\right)^2 C^{\star}_{ij,t} . \end{equation}Appendix B: The von Mises-Fisher distribution: Kullback-Leibler divergence
The natural distribution for a directional variable $\mathbf u$ on the unit hypersphere $S^{D-1}$ is the von Mises–Fisher distribution,
\begin{equation}q(\mathbf u\mid\boldsymbol\mu,\kappa)=C_D(\kappa)\exp\left(\kappa\boldsymbol\mu^\top\mathbf u\right),\end{equation}where $\boldsymbol\mu$ is a unit vector giving the mean direction and $\kappa\geq 0$ is the concentration. For two such distributions $q_a=q(\boldsymbol\mu_a,\kappa_a)$ and $q_b=q(\boldsymbol\mu_b,\kappa_b)$, the KL divergence is
\begin{equation}D_{\mathrm{KL}}(q_a\Vert q_b)=\log\frac{C_D(\kappa_a)}{C_D(\kappa_b)}+A_D(\kappa_a)\left(\kappa_a-\kappa_b\boldsymbol\mu_a^\top\boldsymbol\mu_b\right),\end{equation}with
\begin{equation}A_D(\kappa)=\frac{I_{D/2}(\kappa)}{I_{D/2-1}(\kappa)}.\end{equation}The divergence therefore measures both a mismatch in concentration and a mismatch in direction.
In the vector-spin model, the spins have fixed radius $R$, so we write $\mathbf s=R\mathbf u$. The single-site conditional distribution under an effective field $\mathbf h$ is
\begin{equation}q(\mathbf s\mid\mathbf h)\propto\exp\left(\beta\mathbf h \cdot \mathbf s\right)=\exp\left(\beta R\lVert\mathbf h\rVert \times \left(\frac{\mathbf h}{\lVert\mathbf h\rVert} \cdot \mathbf u\right)\right).\end{equation}Its vMF parameters are therefore
\begin{equation}\boldsymbol\mu_{\mathbf h}=\frac{\mathbf h}{\lVert\mathbf h\rVert},\qquad \kappa_{\mathbf h}=\beta R\lVert\mathbf h\rVert.\end{equation}Thus $\boldsymbol\mu$ is not the effective field itself, but its normalized direction. Likewise, the concentration is not generally $\beta$: it combines inverse temperature, spin radius, and field magnitude. Only in the special case $R=1$ and $\lVert\mathbf h\rVert=1$ do we obtain $\boldsymbol\mu=\mathbf h$ and $\kappa=\beta$. The corresponding mean magnetization is
\begin{equation}\mathbf m(\mathbf h)=R A_D(\kappa_{\mathbf h})\boldsymbol\mu_{\mathbf h}.\end{equation}For the spin-model transformer module, the effective field at site $i$, external step $t$, and internal relaxation step $k$ is
\begin{equation}\mathbf{h}_{i,t,k} = \mathbf{x}_{i,t} + f_{\boldsymbol{\theta}_{\mathrm{FFN}}}\left( \mathbf{x}_{i,t} \right) + \sum_{j} J_{ij} (\mathbf{X}_{t}) \mathbf{m}_{j,t,k}.\end{equation}Each mean-field state therefore defines a factorized vMF approximation,
\begin{equation}q_{t,k}(\mathbf s)=\prod_i q\left(\mathbf s_i\mid\boldsymbol\mu_{i,t,k},\kappa_{i,t,k}\right),\end{equation}with
\begin{equation}\boldsymbol\mu_{i,t,k}=\frac{\mathbf h_{i,t,k}}{\lVert\mathbf h_{i,t,k}\rVert},\qquad \kappa_{i,t,k}=\beta R\lVert\mathbf h_{i,t,k}\rVert.\end{equation}If $q_t^\star$ denotes the corresponding frozen-drive fixed-point distribution, the post-quench mean-field mismatch can be monitored through
\begin{equation}\Delta_{t,k}^{\mathrm{MF}}=D_{\mathrm{KL}}\left(q_{t,k}\Vert q_t^\star\right)=\sum_i D_{\mathrm{KL}}\left(q_{i,t,k}\Vert q_{i,t}^\star\right).\end{equation}This quantity compares the system’s current directional and concentration state with the stationary mean-field response associated with the current drive. It approaches zero as the mean-field state relaxes to that fixed point. In the zero-field case, $\kappa=0$ and the vMF distribution is uniform, so its mean direction is irrelevant.
Explicitly, the single-site contribution comparing the current state $q_{i,t,k}$ with the frozen-drive fixed-point state $q_{i,t}^{\star}$ is given by
\begin{equation}D_{\mathrm{KL}}\left(q_{i,t,k}\Vert q_{i,t}^{\star}\right)=\log\frac{C_D(\kappa_{i,t,k})}{C_D(\kappa_{i,t}^{\star})}+A_D(\kappa_{i,t,k})\left(\kappa_{i,t,k}-\kappa_{i,t}^{\star}\boldsymbol\mu_{i,t,k}\cdot\boldsymbol\mu_{i,t}^{\star}\right).\end{equation}Substituting $\boldsymbol\mu_{\mathbf h}=\mathbf h/\lVert\mathbf h\rVert$ and $\kappa_{\mathbf h}=\beta R\lVert\mathbf h\rVert$ gives the equivalent effective-field expression
\begin{equation}D_{\mathrm{KL}}\left(q_{\mathbf h_a}\Vert q_{\mathbf h_b}\right)=\log\frac{C_D(\beta R\lVert\mathbf h_a\rVert)}{C_D(\beta R\lVert\mathbf h_b\rVert)}+\beta R A_D(\beta R\lVert\mathbf h_a\rVert)\left(\lVert\mathbf h_a\rVert-\frac{\mathbf h_a\cdot\mathbf h_b}{\lVert\mathbf h_a\rVert}\right).\end{equation}Here $\mathbf h_a$ parameterizes the distribution in the first argument of the KL divergence and $\mathbf h_b$ parameterizes the reference distribution in the second argument. The original vMF form above remains well defined when $\mathbf h_a=\mathbf 0$, in which case $\kappa_a=0$ and its mean direction is irrelevant.
In the large-vector-dimension limit, we again set $R^2=D/2-1$ and introduce
\begin{equation}\gamma_{\mathbf h}=\sqrt{1+\frac{\beta^2\lVert\mathbf h\rVert^2}{R^2}}.\end{equation}Using the leading large-$D$ magnetization response $\mathbf m_{\mathbf h}\approx\beta\mathbf h/(1+\gamma_{\mathbf h})$, the KL divergence becomes
\begin{equation}D_{\mathrm{KL}}^{D\to\infty}\left(q_{\mathbf h_a}\Vert q_{\mathbf h_b}\right)\approx\frac{\beta^2\left(\lVert\mathbf h_a\rVert^2-\mathbf h_a\cdot\mathbf h_b\right)}{1+\gamma_{\mathbf h_a}}+R^2\left(\gamma_{\mathbf h_b}-\gamma_{\mathbf h_a}-\log\frac{1+\gamma_{\mathbf h_b}}{1+\gamma_{\mathbf h_a}}\right).\end{equation}The same expression can be written entirely in terms of the corresponding magnetizations. The large-$D$ response can be inverted as
\begin{equation} \mathbf{h}\approx\frac{2R^2}{\beta\left(R^2-\lVert\mathbf{m}\rVert^2\right)}\mathbf{m},\qquad \lVert\mathbf{m}\rVert < R. \end{equation}Substitution gives the magnetization-only form
\begin{equation}D_{\mathrm{KL}}^{D\to\infty}\left(q_{\mathbf m_a}\Vert q_{\mathbf m_b}\right)\approx R^2\log\frac{R^2-\lVert\mathbf m_b\rVert^2}{R^2-\lVert\mathbf m_a\rVert^2}+\frac{2R^2\left(\lVert\mathbf m_b\rVert^2-\mathbf m_a\cdot\mathbf m_b\right)}{R^2-\lVert\mathbf m_b\rVert^2}.\end{equation}For the post-quench relaxation considered here, $\mathbf m_a=\mathbf m_{i,t,k}$ is the current magnetization and $\mathbf m_b=\mathbf m_{i,t}^{\star}$ is the instantaneous fixed-point magnetization. The full factorized mean-field mismatch is therefore
\begin{equation}\Delta_{t,k}^{\mathrm{MF}}\approx\sum_i\left[R^2\log\frac{R^2-\lVert\mathbf m_{i,t}^{\star}\rVert^2}{R^2-\lVert\mathbf m_{i,t,k}\rVert^2}+\frac{2R^2\left(\lVert\mathbf m_{i,t}^{\star}\rVert^2-\mathbf m_{i,t,k}\cdot\mathbf m_{i,t}^{\star}\right)}{R^2-\lVert\mathbf m_{i,t}^{\star}\rVert^2}\right].\end{equation}This expression is asymmetric, as required for a KL divergence. It vanishes when $\mathbf m_{i,t,k}=\mathbf m_{i,t}^{\star}$ at every site and penalizes both differences in magnetization norm and angular misalignment with the frozen-drive fixed-point response.
Footnotes
A first-order expansion around the non-interacting system between consecutive (internal relaxation) time slices. See Spin-Model Transformers (2023). ↩︎
The large-$D$ approximation gets rid of dealing with the modified Bessel functions originating from the von Mises-Fisher distribution used in the ansatz for the decoupled mean magnetizations. It is mainly motivated by the empirical fact that the embedding dimensions in modern neural networks are large. See Spin-Model Transformers (2023) for full details. ↩︎
A drive-dependent coupling rule turns a fixed-size $N \times N$ coupling matrix into a parameterized rule that supports variable system size, drive-dependent routing, and a way to scale system size without learning new explicit parameters. Softmax attention is a convenient choice for a bounded positive row-stochastic coupling rule. Other possible choices include additive or multiplicative combinations with slower base coupling parameters $\mathbf{J}^{0}$ that are drive-independent, leading to a system with persistent interactions in the absence of a drive. ↩︎
In this case, it is more accurate to call the drive $\mathbf{X}_{t}$ an external protocol parameter configuring the instantaneous dynamics. ↩︎
We deliberately focus on our quench-and-relax use case because the nonequilibrium thermodynamics literature is, frankly, quite confusing at times and a daunting terminological minefield. ↩︎

