---
title: Hamiltonian Momentum Attention Blocks
url: https://www.emergentmind.com/topics/hamiltonian-momentum-attention-blocks
type: topic
---

# Hamiltonian Momentum Attention Blocks

Searching arXiv for the cited papers and related work on Hamiltonian Momentum Attention Blocks.
Hamiltonian Momentum Attention Blocks (HMABs) are attention layers in which each token carries not only a position or feature state but also an explicit momentum or velocity variable, so that attention is realized as a time discretization of inertial, damped Hamiltonian dynamics on probability density manifolds equipped with Wasserstein-$2$-type metrics and related geometries. In the recent literature, the designation is associated most directly with "SympFormer: Accelerated attention blocks via Inertial Dynamics on Density Manifolds" [2603.16535], while its conceptual background is sharpened by Hamiltonian analyses of Transformer self-attention that map Query–Key couplings to effective spin-system energies and derive phase-boundary criteria for token dominance in GPT-2 [2507.00683]. Taken together, these works place self-attention within a continuum spanning static energy-based interpretation, mean-field particle dynamics, and momentum-augmented accelerated architectures.

## 1. Definition and conceptual lineage

HMABs are attention layers endowed with explicit momentum or velocity variables per token. In the formulation of SympFormer, tokens carry positions or features $X$ and momenta or velocities $Y=\nabla \Phi(X)$, where $\Phi$ is a learned, layer-evolving velocity field. The resulting dynamics are second-order rather than purely gradient-flow dynamics, so the token cloud undergoes accelerated transport in feature space rather than only dissipative aggregation [2603.16535].

This construction is motivated by a broader view in which classical self-attention can be related to gradient-flow dynamics on density space. HMABs lift that first-order picture to a damped Hamiltonian formulation on the manifold of probability densities and then discretize the resulting dynamics in time. The stated motivation parallels the Euclidean role of Nesterov acceleration: in Euclidean optimization, a momentum variable improves the rate from $O(1/k)$ to $O(1/k^2)$ for convex objectives, and the density-manifold analogue introduces inertial transport while preserving the number of oracle calls, meaning one attention evaluation per layer step [2603.16535].

A distinct but complementary lineage comes from the Hamiltonian analysis of GPT-2 attention heads. There, attention is interpreted as an interacting two-body spin system with continuous, real-valued Heisenberg-type spins, and the learned Query–Key maps define effective exchange couplings. That framework is static rather than momentum-augmented, but it provides the energy-based vocabulary from which momentum-augmented attention can be understood as a dynamical generalization [2507.00683].

## 2. Spin-bath Hamiltonian interpretation of self-attention

In the spin-system analogy, each attention head is modeled as a classical spin system in which tokens or token embeddings play the role of spins. The spins are continuous vectors, not binary variables, and live in the embedding or head subspace. The point of departure is the classical Heisenberg magnet,
$$
H_{\mathrm{Heis}}=-\sum_{i<j} J_{ij}\,\mathbf{S}_i\cdot \mathbf{S}_j,
$$
with the microscopic exchange tensor mapped to a learned Query–Key coupling [2507.00683].

For a head $h$ with learned projections $W_Q^{(h)}$ and $W_K^{(h)}$, the effective exchange tensor is
$$
W_{\mathrm{eff}}^{(h)}=\frac{W_Q^{(h)}\bigl(W_K^{(h)}\bigr)^\top}{\sqrt{d_{\mathrm{head}}}},
$$
and the corresponding two-body attention Hamiltonian is
$$
H^{(0)}(S_j,S_i)=-\,S_j\,W_{\mathrm{eff}}^{(h)}\,S_i^\top.
$$
This is equivalent to the scaled dot-product attention energy
$$
-\frac{(S_jW_Q^{(h)})\cdot (S_iW_K^{(h)})}{\sqrt{d_{\mathrm{head}}}}.
$$
For a prompt with $k$ tokens and last position $k$, the head computes
$$
Q=XW_Q^{(h)},\qquad K=XW_K^{(h)},\qquad V=XW_V^{(h)},
$$
with $X\in \mathbb{R}^{k\times d_{\mathrm{model}}}$ and $\tau=\sqrt{d_{\mathrm{head}}}$. The last-token scores and weights are
$$
\Omega_j=\frac{Q_k\cdot K_j}{\tau},\qquad \alpha_j=\mathrm{softmax}(\Omega)_j.
$$
The head then forms
$$
\mathbf{N}_0^{(h)}=\sum_{j=1}^{k}\alpha_j\,V_j,\qquad
\mathbf{N}_{\mathrm{proj}}^{(h)}=\mathbf{N}_0^{(h)}\bigl(W_O^{(h)}\bigr)^\top.
$$

Within this interpretation, $\mathbf{N}_{\mathrm{proj}}^{(h)}$ acts as an effective mean field on candidate token spins. A single-head energy for a candidate token spin $S$ is
$$
H^{(h)}(S)=-\,\mathbf{N}_{\mathrm{proj}}^{(h)}\cdot S^\top.
$$
The summary mapping is therefore: two-body coupling $J$ corresponds to $W_{\mathrm{eff}}^{(h)}$, the mean-field energy on a candidate token is $H^{(h)}(S)$, the spins are continuous, and softmax implements Boltzmann statistics. The output temperature $T$ and the internal scaling $\tau$ play the roles of thermodynamic temperatures [2507.00683].

This Hamiltonian picture is not yet an HMAB, because no canonical momenta are defined. Its significance lies in identifying an energy landscape implicit in standard self-attention. A plausible implication is that HMABs can be viewed as augmenting this static landscape with explicit inertial state variables and time evolution rather than replacing it.

## 3. Phase boundaries, logit gaps, and causal diagnostics

A central result of the GPT-2 Hamiltonian analysis is the derivation of phase-boundary criteria for token competition. For candidate tokens with embeddings $S_A$ and $S_B$, a head’s theoretical contribution to next-token logits is modeled as
$$
L_{\mathrm{theory}}^{(A)}=\mathbf{N}_{\mathrm{proj}}^{(h)}\cdot S_A^\top,\qquad
L_{\mathrm{theory}}^{(B)}=\mathbf{N}_{\mathrm{proj}}^{(h)}\cdot S_B^\top,
$$
with theoretical logit gap
$$
\Delta L_{\mathrm{theory}}=
L_{\mathrm{theory}}^{(A)}-L_{\mathrm{theory}}^{(B)}
=\mathbf{N}_{\mathrm{proj}}^{(h)}\cdot (S_A-S_B)^\top.
$$
In energy variables, $E_i=-L_i$, hence $\Delta E=-\Delta L$. Under softmax with generation temperature $T$,
$$
\frac{P_{\mathrm{good}}}{P_{\mathrm{bad}}}
=\exp\!\left(\frac{\Delta L}{T}\right).
$$
The phase boundary is $\Delta L=0$; $\Delta L>0$ predicts that $A$ dominates, whereas $\Delta L<0$ predicts that $B$ dominates [2507.00683].

The reported empirical study used GPT-2 small with 12 layers and 12 heads per layer, for 144 heads total, evaluated on 20 factual-recall prompts with one “good” and one “bad” continuation per prompt. For each head and prompt, $\Delta L_{\mathrm{theory}}$ was computed from the head’s full $Q$–$K$–$V$–$O$ output and compared with the full model’s actual next-token logit gap $\Delta L_{\mathrm{actual}}$. Across all 144 heads and 20 prompts, the most predictive head, Layer 3 Head 5 (L3H5), exhibited a strong negative correlation between $\Delta L_{\mathrm{theory}}$ and $\Delta L_{\mathrm{actual}}$, namely $r\approx -0.70$, $p<10^{-3}$, and $r^2\sim 0.48$. The interpretation given is that L3H5 acts antagonistically: its theoretical preference often opposes the model’s final output. The antagonism is traced to the output projection $W_O^{(h)}$, which can invert the sign of alignment and turn an internally “correct” detector into a negative contributor [2507.00683].

Targeted ablations were performed with PyTorch hooks that zeroed out selected heads during the forward pass. On the prompt “Lions are carnivores. Cows are …”, the model incorrectly favored “omnivores” over “herbivores,” with $\Delta L_{\mathrm{actual}}=-0.85$. Ablating L3H5 produced a small change to $-0.88$, whereas ablating L0H0, used as a control with low correlation overall, degraded the gap substantially to $-1.72$. The stated conclusion is that head contributions form context-dependent coalitions with both cooperative and antagonistic roles [2507.00683].

For HMABs, these phase-boundary constructions remain relevant because the evolved head state still induces logits through an effective field. In the momentum-augmented extension associated with the supplied Hamiltonian blueprint, if iterative dynamics increase $|\Delta L|$, the dominant token stabilizes; if iterations reduce $|\Delta L|$ or flip its sign, the head’s dynamics are antagonistic. This suggests that logit-gap analysis functions as both an interpretability diagnostic and a control criterion for momentum-augmented attention.

## 4. HMABs on density manifolds

The formal HMAB construction in SympFormer begins from a generalized gradient flow on the probability density manifold. Let $\Omega=\mathbb{R}^d$, let $\mathcal{P}(\Omega)$ denote smooth positive probability densities with unit mass, and let $G_\rho$ be an invertible operator mapping tangent to cotangent spaces. For an energy functional $F:\mathcal{P}(\Omega)\to \mathbb{R}$, the gradient flow is
$$
\partial_t \rho_t = -\,G^{-1}_{\rho_t}\!\left[\frac{\delta F(\rho_t)}{\delta \rho_t}\right].
$$
Acceleration is introduced by a damped Hamiltonian formulation on the density manifold with cotangent variable $\Phi_t$ and Hamiltonian
$$
H(\rho,\Phi)=T(\rho,\Phi)+F(\rho),\qquad
T(\rho,\Phi)=\frac{1}{2}\int \Phi\,G_\rho^{-1}[\Phi]\,dx.
$$
The accelerated dynamics are
$$
\partial_t
\begin{bmatrix}
\rho_t\\ \Phi_t
\end{bmatrix}
+
\begin{bmatrix}
0\\ \alpha_t \Phi_t
\end{bmatrix}
-
\begin{bmatrix}
0&I\\ -I&0
\end{bmatrix}
\begin{bmatrix}
\delta H/\delta \rho_t\\ \delta H/\delta \Phi_t
\end{bmatrix}
=0,
$$
equivalently
$$
\begin{cases}
\partial_t \rho_t - G^{-1}_{\rho_t}[\Phi_t]=0,\\
\partial_t \Phi_t + \alpha_t \Phi_t
+ \dfrac{1}{2}\dfrac{\delta}{\delta \rho_t}\int \Phi_t G_{\rho_t}^{-1}[\Phi_t]\,dx
+ \dfrac{\delta F}{\delta \rho_t}=0.
\end{cases}
$$
Here $\alpha_t\ge 0$ is a damping schedule, with examples including $\alpha_t=3/t$ and the log-linear form $\alpha(t)=r/t+m$ used in implementation [2603.16535].

Two metric choices are emphasized. In the Wasserstein case,
$$
(G_\rho^W)^{-1}\Phi=-\,\mathrm{div}(\rho \nabla \Phi),
$$
while in the Stein case,
$$
(G_\rho^S)^{-1}\Phi
=
-\,\mathrm{div}\!\left(\rho \int k(\cdot,y)\nabla \Phi(y)\rho(y)\,dy\right),
$$
with $k$ a symmetric positive definite kernel. The first equation is a continuity equation, $\partial_t \rho+\mathrm{div}(\rho u)=0$, with velocity determined by the metric, the momentum potential, and the current density [2603.16535].

The linear self-attention HMAB is derived by taking $A=K^\top Q\in \mathbb{R}^{d\times d}$, assuming $A$ symmetric positive definite, and taking $V\in \mathbb{R}^{d\times d}$ symmetric. The mean-field PDE is
$$
\partial_t \rho_t
=
-\,\nabla\cdot\!\left[\rho_t(x)\int (y^\top A x)V y\,\rho_t(y)\,dy\right],
$$
which is identified as a Stein Variational Gradient Flow with bilinear kernel $k(x,y)=y^\top A x$ and potential
$$
F_{\mathrm{lin}}(\rho)
=
-\frac{1}{2}\int y^\top V y\,\rho(y)\,dy.
$$
Under the empirical-measure ansatz $\rho_t\approx \frac{1}{N}\sum_i \delta_{X_i(t)}$ and defining $Y_i(t)=\nabla \Phi_t(X_i(t))$, the inertial interacting particle system becomes
$$
\dot X_i=\frac{1}{N}\sum_j Y_j (X_j^\top A X_i),\qquad
\dot Y_i=
-\,\alpha_t Y_i
-\frac{1}{N}\sum_j \langle Y_i,Y_j\rangle A X_j
+V X_i,
$$
or in matrix form,
$$
\dot X=\frac{1}{N}XAX^\top Y,\qquad
\dot Y=-\,\alpha(t)Y-\frac{1}{N}YY^\top XA+XV.
$$
These equations define HMAB-linear [2603.16535].

The softmax HMAB is formulated through the energy
$$
F_{\mathrm{SM}}(\rho)
=
-\frac{1}{2}\iint \exp(y^\top A x)\rho(x)\rho(y)\,dx\,dy,
$$
with $\delta F_{\mathrm{SM}}/\delta \rho(x)=-\int \exp(y^\top A x)\rho(y)\,dy$. If $B=VA^{-1}$ is assumed symmetric positive definite, the mobility is encoded by
$$
(G_\rho^{SM})^{-1}\Phi
=
-\,\nabla\cdot \left[
\rho(x)\,
\frac{V A^{-1}\nabla \Phi(x)}{\int \exp(z^\top A x)\rho(z)\,dz}
\right],
$$
yielding an accelerated Wasserstein-$2$-type gradient flow with nonlinear mobility. With $M_{i,j}(t)=\exp(X_i(t)^\top A X_j(t))$ and $S_i(t)=\sum_j M_{i,j}(t)$, the particle dynamics are
$$
\dot X_i=\frac{N B Y_i}{S_i},
$$
and
$$
\dot Y_i
=
-\,\alpha_t Y_i
+\frac{N}{2}A\sum_j X_j M_{i,j}
\left[
\frac{\|Y_i\|_B^2}{S_i^2}
+\frac{\|Y_j\|_B^2}{S_j^2}
+2
\right],
$$
with $\|u\|_B^2=u^\top B u$. In block-matrix form,
$$
\dot X = N\,\mathrm{diag}(M)^{-1}YB^\top,\qquad
\dot Y = -\,\alpha(t)Y + \frac{N}{2}\bigl(\{R,M\}+2M\bigr)XA,
$$
where $R=\mathrm{diag}\!\bigl((Y\odot (YB^\top))_d \oslash (M\odot M)\mathbf{1}\bigr)$ and $\{R,M\}=RM+MR$. The conservative part is Hamiltonian with
$$
H_{SM}(X,Y)
=
\frac{N}{2}\mathrm{tr}\!\bigl(\mathrm{diag}(M)^{-1}YB^\top Y^\top\bigr)
-\frac{N}{2}\mathbf{1}^\top M \mathbf{1}.
$$
The kinetic energy is therefore non-separable, which distinguishes HMABs from Euclidean Nesterov accelerations [2603.16535].

## 5. Discretization, architecture, and implementation

HMABs are implemented by time discretization of the particle systems above. In the linear case, the simplest one-oracle forward Euler layer is
$$
X^{(k+1)} = X^{(k)} + \frac{h_k}{N} X^{(k)} A X^{(k)\top} Y^{(k)},
$$
$$
Y^{(k+1)} = \zeta_{1,k} Y^{(k)} - \frac{h_k}{N} Y^{(k)} Y^{(k)\top} X^{(k)} A + h_k X^{(k)} V,
$$
where $\zeta_{1,k}\approx 1-\alpha(t_k)h_Y$ or $\exp(-\int_{t_k}^{t_{k+1}}\alpha)$, depending on the integrator, and the step sizes are learned per layer. In the softmax case,
$$
M^{(k)}=\exp\!\bigl(X^{(k)}A X^{(k)\top}\bigr),\qquad
S^{(k)}=\mathrm{diag}(M^{(k)})\mathbf{1},
$$
$$
X^{(k+1)} = X^{(k)} + h_X N\,\mathrm{diag}(M^{(k)})^{-1}Y^{(k)}B^\top,
$$
and
$$
R^{(k)}
=
\mathrm{diag}\!\left(
\bigl(Y^{(k)}\odot (Y^{(k)}B^\top)\bigr)_d
\oslash
\bigl(M^{(k)}\odot M^{(k)}\bigr)\mathbf{1}
\right),
$$
$$
Y^{(k+1)}
=
\zeta_{1,k}Y^{(k)}
+\frac{h_Y N}{2}\bigl(\{R^{(k)},M^{(k)}\}+2M^{(k)}\bigr)X^{(k+1)}A.
$$
These schemes preserve the number of attention oracle calls per layer step [2603.16535].

The paper also derives geometry-aware variants: conformally symplectic Euler, described as “kick-then-damp,” exponential Euler, and Adams–Bashforth AB-2 multistep updates. In the softmax conformally symplectic Euler step, momentum is first kicked by the conservative force,
$$
\hat Y = Y + h_Y \frac{N}{2}(\{R,M\}+2M)XA,
$$
then damped by $Y^+=\sigma \hat Y$ with $\sigma=\exp(-\int_t^{t+h_Y}\alpha(s)\,ds)$, and then used in the position update
$$
X^+=X+h_X N\,\mathrm{diag}(S)^{-1}Y^+ B^\top.
$$
AB-2 reuses previous function values,
$$
Y^+=Y+h_Y\bigl[(2+\theta)G(X,Y)-\theta G_{\mathrm{prev}}\bigr],\qquad
X^+=X+h_X\bigl[(2+\theta)F(X,Y)-\theta F_{\mathrm{prev}}\bigr],
$$
with $\theta=h_Y/(2h_{Y,\mathrm{prev}})$, thereby achieving second-order accuracy with one new oracle call [2603.16535].

A compact comparison of the two principal HMAB regimes is given below.

| Regime | Core structure | Per-layer state |
|---|---|---|
| Linear HMAB | Bilinear kernel $k(x,y)=y^\top A x$; accelerated SVGF | Tokens $X$ and momentum $Y$ |
| Softmax HMAB | $M=\exp(XAX^\top)$; accelerated Wasserstein-type flow with nonlinear mobility | Tokens $X$ and momentum $Y$ |

In implementation, $Y$ is initialized to zero at the first layer. Step sizes $h_X$ and $h_Y$ are learned per layer, the damping schedule $\alpha(t)=r/t+m$ with learnable $r,m>0$ is reported to work robustly, and LayerNorm is applied to momentum and tokens to stabilize training. The optimizer is AdamW with cosine learning rate schedule, warm-up, and gradient-norm clipping. Momentum is carried forward to the next HMAB layer and, in SympFormer, can also be fed to the MLP block as a look-ahead direction [2603.16535].

A separate momentum-augmented Hamiltonian blueprint, developed alongside the spin-bath analysis, proposes a canonical construction with coordinates $x$, momenta $p$, Hamiltonian
$$
H(x,p)=U(x)+\frac{1}{2}p^\top M^{-1}p,
$$
potential
$$
U(x)=-\frac{1}{2}x^\top Jx-h^\top x,
$$
Hamilton’s equations
$$
\dot x=M^{-1}p,\qquad \dot p=-\,\nabla U(x),
$$
and leapfrog updates
$$
p_{t+\frac12}=p_t-\frac{\Delta t}{2}\nabla U(x_t),\qquad
x_{t+1}=x_t+\Delta t\,M^{-1}p_{t+\frac12},\qquad
p_{t+1}=p_{t+\frac12}-\frac{\Delta t}{2}\nabla U(x_{t+1}).
$$
After $S$ leapfrog steps, the evolved state $x_S$ is used as a refined query. This canonical formulation is presented as a principled momentum augmentation of Hamiltonian attention and is conceptually adjacent to, but structurally distinct from, the density-manifold HMABs of SympFormer [2507.00683].

## 6. Theoretical guarantees, empirical performance, and limitations

For linear HMABs, SympFormer proves an invariant-family result for elliptically contoured distributions. If the initial density is elliptically contoured $E(m_0,\Sigma_0,g)$, then the accelerated linear attention PDE preserves this class:
$$
\rho_t=E(m_t,\Sigma_t,g),\qquad
\Phi_t(x)=\frac{1}{2}x^\top P_t x,
$$
where $m_t$, $\Sigma_t$, and $P_t$ solve
$$
\dot m_t = C_t m_t,\qquad
\dot \Sigma_t = \Sigma_t C_t^\top + C_t \Sigma_t,\qquad
\dot P_t = -\,\alpha_t P_t - (C_t^\top P_t + P_t C_t) + V,
$$
with $C_t=P_t(\kappa_g \Sigma_t + m_t m_t^\top)A$ and
$$
\kappa_g = \frac{1}{d}\int \|y\|^2 g(\|y\|^2)\,dy.
$$
For centered initial data $m_0=0$, the dynamics reduce to
$$
\dot\Sigma_t
=
\kappa_g(\Sigma_t A\Sigma_t P_t + P_t\Sigma_t A\Sigma_t),
$$
$$
\dot P_t
=
-\,\alpha_t P_t
-\kappa_g(A\Sigma_t P_t^2 + P_t^2\Sigma_t A)
+V.
$$
The stated significance is qualitative stability: Gaussian-like ensembles remain Gaussian-like under the accelerated flow [2603.16535].

Empirically, SympFormer evaluates decoder-only nano-GPT-style models on next-token prediction. On TinyStories after 10k steps, the reported best and last validation losses are as follows.

| Model | Best validation loss | Last validation loss |
|---|---:|---:|
| Baseline Transformer | 2.4473 | 2.4687 |
| YuriiFormer | 2.3872 | 2.4041 |
| HMAB Plain Euler | 2.3234 | 2.3247 |
| HMAB Presymplectic Euler | 2.4592 | 2.4728 |
| HMAB Presymplectic ExpEuler | 2.2523 | 2.3579 |
| HMAB Presymplectic AB-2 | 2.6546 | 2.6653 |
| HMAB Presymplectic ETD-AB2 | 1.8386 | 1.8386 |

The reported conclusion is that the oracle-preserving HMABs, especially ETD-AB2, substantially outperform the baseline Transformer and the Euclidean-geometry alternative. Similar advantages are stated for linear HMABs in separate tables. The paper also reports that momentum layer normalization and the log-linear damping schedule stabilize training, and that learned step sizes per layer allow adaptive control of transport and momentum magnitudes [2603.16535].

The comparison with Euclidean momentum methods is explicit. Recent Euclidean Nesterov transformers, exemplified here by YuriiFormer, introduce momentum in ambient space rather than on the density manifold, and their Hamiltonians are separable. HMABs instead use a non-separable kinetic energy depending on attention geometry, coupling $X$ and $Y$ through $M$ in the softmax case. This produces different stability and accuracy tradeoffs [2603.16535].

Limitations are also explicit. The strongest theoretical invariant is proved only for the linear or SVGD case. In the softmax case, rigorous global well-posedness and convergence remain future work. More broadly, the GPT-2 Hamiltonian analysis emphasizes that a pairwise $Q$–$K$ mapping abstracts away value mixing and residual interactions; $V$ and $W_O$ can invert or reshape contributions, as seen in the antagonistic role of L3H5. The analytic mapping also treats heads independently, whereas in practice heads cooperate and compete nonlinearly through residual streams. Finally, the empirical validation of the spin-bath picture is reported for GPT-2 small, so generalization to larger models or different training objectives remains to be tested [2507.00683].

From these results, a coherent picture emerges. HMABs are not merely attention layers with an added velocity buffer; they are a family of accelerated attention architectures derived from damped Hamiltonian dynamics on density space, with one-oracle discretizations, explicit geometric structure, and measurable empirical consequences. The static spin-bath Hamiltonian view supplies an interpretable energy language and phase-boundary diagnostics, whereas SympFormer supplies the fully dynamical realization in which token states evolve under momentum, damping, and attention-induced geometry.

Source: https://www.emergentmind.com/topics/hamiltonian-momentum-attention-blocks