Hamiltonian Momentum Attention Blocks

• Hamiltonian Momentum Attention Blocks are attention layers augmented with momentum variables, enabling accelerated and inertial dynamics for token representation.
• They reinterpret self-attention as a time-discretized, damped Hamiltonian flow on probability density manifolds, bridging energy-based and mean-field perspectives.
• Empirical results show that HMABs improve stability and accuracy through geometric phase-boundary analysis and adaptive discretization compared to traditional methods.

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" (Stein et al., 17 Mar 2026), 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 (Bhattacharjee et al., 1 Jul 2025). 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 (Stein et al., 17 Mar 2026).

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 (Stein et al., 17 Mar 2026).

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 (Bhattacharjee et al., 1 Jul 2025).

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 (Bhattacharjee et al., 1 Jul 2025).

For a head $h$ with learned projections $W_Q^{(h)}$ and $W_K^{(h)}$, the effective exchange tensor is

$X$0

and the corresponding two-body attention Hamiltonian is

$X$1

This is equivalent to the scaled dot-product attention energy

$X$2

For a prompt with $X$3 tokens and last position $X$4, the head computes

$X$5

with $X$6 and $X$7. The last-token scores and weights are

$X$8

The head then forms

$X$9

Within this interpretation, $Y=\nabla \Phi(X)$0 acts as an effective mean field on candidate token spins. A single-head energy for a candidate token spin $Y=\nabla \Phi(X)$1 is

$Y=\nabla \Phi(X)$2

The summary mapping is therefore: two-body coupling $Y=\nabla \Phi(X)$3 corresponds to $Y=\nabla \Phi(X)$4, the mean-field energy on a candidate token is $Y=\nabla \Phi(X)$5, the spins are continuous, and softmax implements Boltzmann statistics. The output temperature $Y=\nabla \Phi(X)$6 and the internal scaling $Y=\nabla \Phi(X)$7 play the roles of thermodynamic temperatures (Bhattacharjee et al., 1 Jul 2025).

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 $Y=\nabla \Phi(X)$8 and $Y=\nabla \Phi(X)$9, a head’s theoretical contribution to next-token logits is modeled as

$\Phi$0

with theoretical logit gap

$\Phi$1

In energy variables, $\Phi$2, hence $\Phi$3. Under softmax with generation temperature $\Phi$4,

$\Phi$5

The phase boundary is $\Phi$6; $\Phi$7 predicts that $\Phi$8 dominates, whereas $\Phi$9 predicts that $O(1/k)$0 dominates (Bhattacharjee et al., 1 Jul 2025).

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, $O(1/k)$1 was computed from the head’s full $O(1/k)$2–$O(1/k)$3–$O(1/k)$4–$O(1/k)$5 output and compared with the full model’s actual next-token logit gap $O(1/k)$6. Across all 144 heads and 20 prompts, the most predictive head, Layer 3 Head 5 (L3H5), exhibited a strong negative correlation between $O(1/k)$7 and $O(1/k)$8, namely $O(1/k)$9, $O(1/k^2)$0, and $O(1/k^2)$1. 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 $O(1/k^2)$2, which can invert the sign of alignment and turn an internally “correct” detector into a negative contributor (Bhattacharjee et al., 1 Jul 2025).

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 $O(1/k^2)$3. Ablating L3H5 produced a small change to $O(1/k^2)$4, whereas ablating L0H0, used as a control with low correlation overall, degraded the gap substantially to $O(1/k^2)$5. The stated conclusion is that head contributions form context-dependent coalitions with both cooperative and antagonistic roles (Bhattacharjee et al., 1 Jul 2025).

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 $O(1/k^2)$6, the dominant token stabilizes; if iterations reduce $O(1/k^2)$7 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 $O(1/k^2)$8, let $O(1/k^2)$9 denote smooth positive probability densities with unit mass, and let $$H_{\mathrm{Heis}}=-\sum_{i<j} J_{ij}\,\mathbf{S}_i\cdot \mathbf{S}_j,$$0 be an invertible operator mapping tangent to cotangent spaces. For an energy functional $$H_{\mathrm{Heis}}=-\sum_{i<j} J_{ij}\,\mathbf{S}_i\cdot \mathbf{S}_j,$$1, the gradient flow is

$$H_{\mathrm{Heis}}=-\sum_{i<j} J_{ij}\,\mathbf{S}_i\cdot \mathbf{S}_j,$$2

Acceleration is introduced by a damped Hamiltonian formulation on the density manifold with cotangent variable $$H_{\mathrm{Heis}}=-\sum_{i<j} J_{ij}\,\mathbf{S}_i\cdot \mathbf{S}_j,$$3 and Hamiltonian

$$H_{\mathrm{Heis}}=-\sum_{i<j} J_{ij}\,\mathbf{S}_i\cdot \mathbf{S}_j,$$4

The accelerated dynamics are

$$H_{\mathrm{Heis}}=-\sum_{i<j} J_{ij}\,\mathbf{S}_i\cdot \mathbf{S}_j,$$5

equivalently

$$H_{\mathrm{Heis}}=-\sum_{i<j} J_{ij}\,\mathbf{S}_i\cdot \mathbf{S}_j,$$6

Here $$H_{\mathrm{Heis}}=-\sum_{i<j} J_{ij}\,\mathbf{S}_i\cdot \mathbf{S}_j,$$7 is a damping schedule, with examples including $$H_{\mathrm{Heis}}=-\sum_{i<j} J_{ij}\,\mathbf{S}_i\cdot \mathbf{S}_j,$$8 and the log-linear form $$H_{\mathrm{Heis}}=-\sum_{i<j} J_{ij}\,\mathbf{S}_i\cdot \mathbf{S}_j,$$9 used in implementation (Stein et al., 17 Mar 2026).

Two metric choices are emphasized. In the Wasserstein case,

$h$0

while in the Stein case,

$h$1

with $h$2 a symmetric positive definite kernel. The first equation is a continuity equation, $h$3, with velocity determined by the metric, the momentum potential, and the current density (Stein et al., 17 Mar 2026).

The linear self-attention HMAB is derived by taking $h$4, assuming $h$5 symmetric positive definite, and taking $h$6 symmetric. The mean-field PDE is

$h$7

which is identified as a Stein Variational Gradient Flow with bilinear kernel $h$8 and potential

$h$9

Under the empirical-measure ansatz $W_Q^{(h)}$0 and defining $W_Q^{(h)}$1, the inertial interacting particle system becomes

$W_Q^{(h)}$2

or in matrix form,

$W_Q^{(h)}$3

These equations define HMAB-linear (Stein et al., 17 Mar 2026).

The softmax HMAB is formulated through the energy

$W_Q^{(h)}$4

with $W_Q^{(h)}$5. If $W_Q^{(h)}$6 is assumed symmetric positive definite, the mobility is encoded by

$W_Q^{(h)}$7

yielding an accelerated Wasserstein-$W_Q^{(h)}$8-type gradient flow with nonlinear mobility. With $W_Q^{(h)}$9 and $W_K^{(h)}$0, the particle dynamics are

$W_K^{(h)}$1

and

$W_K^{(h)}$2

with $W_K^{(h)}$3. In block-matrix form,

$W_K^{(h)}$4

where $W_K^{(h)}$5 and $W_K^{(h)}$6. The conservative part is Hamiltonian with

$W_K^{(h)}$7

The kinetic energy is therefore non-separable, which distinguishes HMABs from Euclidean Nesterov accelerations (Stein et al., 17 Mar 2026).

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

$W_K^{(h)}$8

$W_K^{(h)}$9

where $X$00 or $X$01, depending on the integrator, and the step sizes are learned per layer. In the softmax case,

$X$02

$X$03

and

$X$04

$X$05

These schemes preserve the number of attention oracle calls per layer step (Stein et al., 17 Mar 2026).

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,

$X$06

then damped by $X$07 with $X$08, and then used in the position update

$X$09

AB-2 reuses previous function values,

$X$10

with $X$11, thereby achieving second-order accuracy with one new oracle call (Stein et al., 17 Mar 2026).

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

Regime	Core structure	Per-layer state
Linear HMAB	Bilinear kernel $X$12; accelerated SVGF	Tokens $X$13 and momentum $X$14
Softmax HMAB	$X$15; accelerated Wasserstein-type flow with nonlinear mobility	Tokens $X$16 and momentum $X$17

In implementation, $X$18 is initialized to zero at the first layer. Step sizes $X$19 and $X$20 are learned per layer, the damping schedule $X$21 with learnable $X$22 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 (Stein et al., 17 Mar 2026).

A separate momentum-augmented Hamiltonian blueprint, developed alongside the spin-bath analysis, proposes a canonical construction with coordinates $X$23, momenta $X$24, Hamiltonian

$X$25

potential

$X$26

Hamilton’s equations

$X$27

and leapfrog updates

$X$28

After $X$29 leapfrog steps, the evolved state $X$30 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 (Bhattacharjee et al., 1 Jul 2025).

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 $X$31, then the accelerated linear attention PDE preserves this class:

$X$32

where $X$33, $X$34, and $X$35 solve

$X$36

with $X$37 and

$X$38

For centered initial data $X$39, the dynamics reduce to

$X$40

$X$41

The stated significance is qualitative stability: Gaussian-like ensembles remain Gaussian-like under the accelerated flow (Stein et al., 17 Mar 2026).

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 (Stein et al., 17 Mar 2026).

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$42 and $X$43 through $X$44 in the softmax case. This produces different stability and accuracy tradeoffs (Stein et al., 17 Mar 2026).

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 $X$45–$X$46 mapping abstracts away value mixing and residual interactions; $X$47 and $X$48 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 (Bhattacharjee et al., 1 Jul 2025).

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.