PaperlayerView exact revision on arXiv

Source-linked AI summary

SympFormer: Inertial Attention Blocks from Damped Hamiltonian Dynamics on Spaces of Probability Measures

Viktor Stein, Levin Maier, Wuchen Li, Gabriele Steidl

arXiv:2610.31600v1math.SGmath.APmath.DGmath.DSstat.ML

TL;DR

Transformer dynamics remain only partially understood, and rigorous geometric foundations for their mean-field descriptions are still needed. This paper develops local-mobility and quotient Stein geometries for accelerated attention and connects continuous linear-attention moment growth to discrete token norms. Geometric corrections yield small validation-loss reductions in some settings, with increased computational cost.

  • Problem

    A rigorous geometric justification for softmax attention’s gradient-flow interpretation, and foundations for accelerated attention dynamics on density spaces, remain open.

    Source: §1 · p. 2

  • Method

    The paper develops local-mobility and quotient Stein geometries with differential calculus to formulate gradient and damped Hamiltonian attention dynamics on specified domains.

    Source: §1 · p. 3

  • Results

    Geometric corrections to token and momentum updates give small reductions in mean validation loss in some settings, at increased computational cost.

    Source: §1 · p. 3

  • Takeaways & Limitations

    For accelerated linear attention at fixed token count, continuous moment blow-up implies every prescribed token-norm bound is exceeded in discrete layers for sufficiently small step sizes.

    Source: §1 · p. 3

    The trained architectures test the usefulness of geometrically motivated corrections, not numerical accuracy of a solver for the Newton equation.

    Source: §6 · p. 50

Abstract

from arXiv · show

The attention mechanism underpins modern large language models, yet their dynamics remain only partially understood. Recent perspectives interpret attention blocks as interacting particle systems, whose mean-field limits correspond to gradient flows of energy functionals on spaces of probability densities on $\mathbb R^d$ equipped with Wasserstein-type metrics. We extend this viewpoint by deriving accelerated attention blocks from inertial, Nesterov-type damped Hamiltonian dynamics on these density spaces, developing the geometry and differential calculus needed to make the construction rigorous. In particular, we introduce a local-mobility geometry for the Otto-Wasserstein and softmax-Wasserstein settings, and a quotient Stein geometry that resolves the degeneracy of the linear-attention kernel to yield a well-defined metric and differential calculus. We provide rigorous foundations when softmax attention admits a gradient-flow formulation and when its inertial counterpart admits a damped Hamiltonian on explicitly characterized domains. For linear attention, we obtain a finite-dimensional solution theory with sharp conditions for global existence versus finite-time second-moment blow-up. We further prove a quantitative link between continuous moment blow-up and the discrete-layer depth at which token norms cross a prescribed threshold. Finally, motivated by the finite-token geometry, we propose geometric corrections to token and momentum updates. In decoder-only, nanoGPT-style experiments on TinyStories and OpenWebText, these corrections lower validation negative log-likelihood on next-token prediction relative to momentum-only controls.

1 Introduction

The paper asks whether geometry and PDE analysis of attention’s mean-field limit can clarify finite-token, finite-layer transformer dynamics. It develops rigorous geometric foundations and connects continuous linear-attention blow-up to discrete token growth, while testing geometric update corrections.

1 Introduction

  • The paper is motivated by the partial understanding of transformer behavior and the open rigorous justification of softmax attention’s geometric gradient-flow interpretation.The introduction asks whether analysis of the mean-field limit can yield insights into transformers with finitely many tokens and layers.Source: §1 · pp. 2–3
  • The paper develops local-mobility geometry for softmax attention and quotient Stein geometry for linear attention, with differential calculus on specified domains.These constructions provide rigorous foundations for gradient-flow and damped Hamiltonian formulations.Source: §1 · p. 3
  • Coordinate reductions yield finite-dimensional matrix ODEs and strong-solution existence and uniqueness for admissible non-Gaussian initial data, with potentially unbounded support.The accelerated case requires a quadratic initial momentum potential; the reductions also give conditions for global existence and finite-time second-moment blow-up.Source: §1 · p. 3
  • For accelerated linear attention at fixed token count, continuous moment growth predicts the first explicit-Euler layer crossing of a prescribed directional second-moment threshold, with bounded layer-count error as the step size tends to zero.Continuous moment blow-up implies every prescribed token-norm bound is exceeded in discrete layers for sufficiently small step sizes.Source: §1 · p. 3
  • Geometric corrections to token and momentum updates give small reductions in mean validation loss in some settings, at increased computational cost.The comparison uses matched momentum-only controls on TinyStories and OpenWebText.Source: §1 · p. 3

2 Transformers as Dynamical Systems

The section recasts attention layers as token dynamics and, under suitable limits, as a continuity equation for the token distribution. It motivates a geometric interpretation while identifying conditions and scope that constrain its rigorous formulation.

Transformers as Dynamical Systems

  • With shared weights, the formal continuous-depth limit turns residual updates into an interacting particle system, and a suitably normalized mean-field limit yields a continuity equation for the token distribution.This is the passage from attention layers to transport dynamics illustrated in Figure 2.Source: Figure 2 · p. 5
  • The section assumes symmetric A = K^TQ with m = d for its geometric constructions and sets aside masking, normalization, scaling, output projection, and multi-head structure.The continuous-depth description also uses shared Q, K, and V across layers; the particle and density limits focus on the attention update.Source: §2 · pp. 5–6
  • The paper uses “linear self-attention” for unnormalized dot-product attention, while noting that softmax attention’s evaluation cost scales quadratically with sequence length.Kernel attention replaces softmax with a more general kernel to address this cost.Source: §2 · p. 6
  • Under suitable assumptions, the mean-field equation admits a formal gradient-flow interpretation specified by both an energy and a Riemannian metric.The metric determines the direction of steepest descent, so identifying an energy alone does not complete the construction.Source: §2 · p. 7
  • For softmax attention, rigorous geometric justification on non-compactly supported densities in Rd remained open; such densities permit analysis of token-magnitude growth excluded by sphere-constrained models.Additional integrability and coercivity conditions are needed on the whole space.Source: §2 · p. 7

3 Nesterov’s Acceleration Method and Attention Models

The section motivates Nesterov acceleration for attention as a damped Hamiltonian construction on density spaces. It explains the geometric prerequisites, the Stein obstruction, and the kernel-dependent function-space questions that motivate later frameworks.

3.1 Nesterov Acceleration on Rd

  • Nesterov acceleration can be written as a damped Hamiltonian system, providing the finite-dimensional model for introducing momentum into attention dynamics.The section recalls its continuous-time formulation and the damped Hamiltonian form.Source: §3 · p. 7, §3.1 · p. 9

3 Nesterov’s Acceleration Method and Attention Models

  • Attention kernels induce different mobility operators and formal geometries; making their accelerated systems rigorous requires resolving function-space and differentiability questions.These questions motivate the local-mobility calculus and quotient Stein geometry developed later.Source: §3 · p. 8

3.1 Nesterov Acceleration on Rd

  • For convex objectives with a minimizer, the cited discrete scheme has convergence rate F(x(k)) − F* ≤ O(∥x0 − x*∥^2/(τk^2)), while its continuous dynamics have a corresponding O(∥x0 − x*∥^2/t) rate.The discrete rate is stated for step sizes τ in (0, 1/L], where L is the Lipschitz constant of ∇F.Source: §3.1 · p. 8

3.2 Nesterov Acceleration on Probability Measures on Compact Sets

  • On compact density spaces, accelerated Otto–Wasserstein dynamics use an invertible mobility operator, whereas the Stein mobility may fail to be surjective on the full formal tangent space.Restricting the tangent space is necessary for surjectivity; on general compact manifolds, intrinsic formulation also requires an operator-valued kernel to average tangent vectors.Source: §3.2 · pp. 9, 11

3.3 Nesterov Acceleration and Attention Models on Rd

  • For attention on Rd, the kernel determines the mobility and its inverse induces the metric on suitable tangent and cotangent spaces.For softmax mobility, smooth positivity and finite second moments alone do not ensure finite normalization or coercivity; for the bilinear kernel, the range and nontrivial kernel require restricting and quotienting the spaces.Source: §3.3 · pp. 11–13
  • The section identifies differentiability as a separate requirement because the chosen probability-density classes are not given a smooth manifold structure.The function spaces must support well-defined variations with respect to both density and potential.Source: §3.3 · p. 13

4 A Non-Compact Functional Framework for Otto Calculus

The section builds a rigorous local-mobility calculus for smooth probability densities on R^d, where formal Otto-style constructions require suitable domains and analytic assumptions. A weighted Poincaré inequality supports the operator, metric, and gradient constructions used for Otto–Wasserstein and softmax–Wasserstein geometries.

4 A Non-Compact Functional Framework for Otto Calculus

  • Smooth positivity and finite second moments alone do not provide the regular tangent and cotangent spaces, operator invertibility, or differentiation needed for density calculus on R^d.The framework specifies the function spaces and analytic assumptions needed for these constructions.Source: §4 · pp. 13–14

4.2 Weighted Local-Mobility Geometry on P∞(Rd)

  • A weighted Poincaré inequality makes the centered potential space complete and supplies coercivity for solving the potential equation.The framework starts from a weighted Sobolev space of potentials.Source: §4 · p. 14
  • The local-mobility operator is an isomorphism between the chosen regular cotangent and tangent spaces, and its inverse defines the metric.The resulting calculus identifies gradients for suitable first variations and specializes to Otto–Wasserstein and softmax–Wasserstein geometries.Source: §4 · p. 14, §4.2 · p. 18

4.1 Background on Infinite Dimensional Geometry

  • Weak Riemannian metrics can lack gradients for smooth functionals when their differentials fall outside the metric map’s range.This range condition is a central analytic issue for Otto-type calculus on non-compact domains.Source: §4.1 · pp. 14–15

5 Stein Calculus for Bilinear Kernels

The bilinear Stein kernel makes the usual Stein–Onsager operator non-invertible on the full regular spaces because of its infinite-dimensional kernel. The section resolves this degeneracy with a quotient construction that yields an effective isomorphism and a positive-definite metric.

5 Stein Calculus for Bilinear Kernels

  • The bilinear Stein kernel’s operator has finite-dimensional range and an infinite-dimensional kernel, so it cannot be inverted on the full regular tangent and cotangent spaces.The kernel consists precisely of zero-moment cotangent classes.Source: §5 · pp. 21, 23
  • The quotient Stein geometry restricts tangent directions to those generated by the mobility operator and identifies cotangent elements that generate the same variation.This construction addresses the obstruction to making the operator an isomorphism on the full spaces.Source: §5 · pp. 21, 23
  • The calculus is developed for smooth bilinear Stein probability densities with a finite second moment and a logarithmic condition.The corresponding notion of differentiation uses local regular first variations and their regular extensions.Source: §5 · pp. 21, 24
  • On the effective spaces, the Stein–Onsager operator is an isomorphism whose inverse defines a well-defined positive-definite metric.The section also develops the corresponding gradient calculus for admissible regular first variations.Source: §5 · pp. 23–24

6 Softmax Attention: Geometry and Particle Dynamics

This section develops geometric interpretations of first-order and accelerated softmax attention and uses them to derive finite-token corrections and implementable attention-layer updates. It also describes adapting the updates to causal multi-head attention for numerical experiments.

6.1 Geometry and First-Order Softmax Dynamics

  • The softmax transformer PDE is a regular gradient flow on densities satisfying the finite-energy and regular-variation conditions characterized in Proposition 6.2 and Corollary 6.4.The results identify the gradient at each admissible density, but do not establish general gradient-flow trajectories or preservation of the admissible domain.Source: §6.1 · pp. 26–27

6.2 Accelerated Softmax Dynamics and Hamiltonian Structure

  • The accelerated softmax system is pointwise well-defined on its stated density–potential domain, while a regular damped Hamiltonian interpretation requires stronger conditions on the Hamiltonian’s regular variations.Examples illustrate both the regular Hamiltonian domain and a pair where pointwise well-definedness and the first-order gradient criterion hold but Hamiltonian regularity fails.Source: §6.2 · pp. 27, 29, 30

6.3 From Softmax Geometry to Attention Layers

  • The finite-token softmax metric makes attention the gradient of the interaction energy and yields an inertial update with a position-dependent Levi–Civita connection correction.The correction is explicitly computable using matrix products and row rescalings, at O(N^2d+Nd^2) arithmetic cost for fixed dense A, B and precomputed B^-1; geodesic expansion and its differential produce position and velocity updates.Source: §6.3 · pp. 30–31
  • Both updates use the explicit connection correction in (73), so they can be implemented without solving geodesic or parallel-transport equations.The correction is computed explicitly as a layer update.Source: §6.3 · p. 31
  • Section 9 adapts the construction to causal multi-head attention with learned weights, replacing the connection correction with proxy bΓ using B = Id and native attention weights.Combining (74) and (75) with LayerNorm modifications yields update (111); coordinate variant (110) inserts the correction directly into the velocity update, following (71).Source: §6.3 · p. 32
  • These adaptations provide the architectures compared for predictive performance and computational cost in numerical experiments.Source: §6.3 · p. 32

7 Linear Attention: Geometry, Inertial Dynamics and Effective Depth

This section develops a geometric and dynamical analysis of linear attention, using quotient Stein geometry and first- and second-order dynamics to study covariance and moment growth. It connects continuous blow-up to discrete-layer depth and token-norm behavior, including the threshold that persists under bounded nonnegative damping for unnormalized updates.

7.1 First-Order Linear Attention: Geometry and Analysis

  • Covariance coordinates induce a quotient Stein metric and finite-dimensional gradient flow, with an exact first-order PDE–moment correspondence for C1 initial densities having finite second moment.The correspondence extends beyond elliptical families; moment evolution and maximal existence time depend on the initial second-moment matrix.Source: §7 · p. 32, §7.1 · p. 36

7.2 Accelerated Linear Attention: Geometry and Hamiltonian Dynamics

  • The damped matrix system characterizes continuation: finite-time breakdown requires second-moment divergence, while V ⪯ 0 gives global existence; under commuting data, positive eigenvalues of V force blow-up with bounded nonnegative damping.The commuting-data result assumes zero initial momentum.Source: §7 · p. 32
  • The accelerated dynamics have a matrix Hamiltonian and damped Newton formulation, with exact ODE–PDE correspondence, local strong well-posedness, and continuous dependence for admissible initial densities and quadratic initial potentials.The framework treats acceleration as adding inertia; convergence-rate improvement is not established, and extending the Hamiltonian formulation to non-quadratic potentials remains open.Source: §7 · p. 32, §7.2 · pp. 38–39

7 Linear Attention: Geometry, Inertial Dynamics and Effective Depth

  • For finite-token attention, the continuous matrix system governs token dynamics, and its blow-up time predicts a critical effective depth separating bounded token norms below it from unbounded norms at or above it.The layer count needed to reach a prescribed second-moment threshold is predicted from the continuous threshold-crossing time.Source: §7 · p. 32, Figure 5 · p. 33

7.1 First-Order Linear Attention: Geometry and Analysis

  • For first-order linear attention, finite-time breakdown requires unbounded second moments, V ⪯ 0 ensures global existence, and commuting A, S0, and V allow explicit blow-up criteria.The explicit criteria apply to C1 initial densities with finite second moment, not only Gaussian data.Source: §7.1 · pp. 37–38

7.2 Accelerated Linear Attention: Geometry and Hamiltonian Dynamics

  • The accelerated PDE reduces to a finite-dimensional covariance system that can be written as a damped Hamiltonian system and a damped Newton equation for the linear-attention metric.The quadratic reduction uses the covariance matrix and quadratic-potential coefficient; the Hamiltonian formulation uses the momentum coordinate Π = P/2.Source: §7.2 · pp. 40–41
  • The covariance ODE and PDE have matching maximal existence intervals, with unique strong solutions corresponding exactly to the matrix-system solutions.The result also establishes preservation of geometric admissibility for the constructed strong solutions.Source: §7.2 · pp. 41–42
  • With commuting A, S0, and V, zero initial potential coefficient, and bounded damping, any positive eigenvalue of V forces finite-time moment blow-up.The corresponding directional second moments are strictly increasing, and sufficiently large moments grow superlinearly despite bounded damping.Source: §7.2 · p. 43

7.3 Finite-Token Layers and Effective Depth

  • The continuous covariance system exactly represents the finite-token dynamics: S(t) = X(t)^T X(t)/N, and unbounded S forces unbounded token norms for fixed N.When S becomes unbounded at a finite maximal time, the finite-token dynamics have that same maximal existence time.Source: §7.3 · p. 44
  • hLR(h) = tR + O(h): the continuous crossing time predicts the first discrete layer reaching any fixed second-moment threshold.For increasingly large thresholds as the discretization is refined, these effective depths approach the continuous blow-up time, even with a fixed number of tokens.Source: §7.3 · p. 46
  • For sufficiently small step size, a prescribed token-norm bound is exceeded before accumulated layer time reaches Tmax, even with a fixed number of tokens.This connects token growth in the continuous dynamics to a threshold-crossing layer in the discrete network.Source: §7.3 · pp. 46–47
  • At fixed effective depth T ≥ Tmax, the largest token norms diverge as the number of layers increases.Here effective depth is T = Lh, so fixing T while increasing L means reducing the step size h = T/L.Source: §7.3 · pp. 46–47
  • At fixed effective depth T < Tmax, increasing layers improves the continuous-dynamics approximation while token norms remain uniformly bounded.The refinement uses h = T/L and approximates a bounded trajectory.Source: §7.3 · p. 47
  • Bounded nonnegative damping does not remove the critical effective-depth threshold for unnormalized updates.This conclusion holds under the assumptions of Corollary 7.30.Source: §7.3 · p. 47

8 Discrete SympFormer architectures

This section discretizes inertial attention dynamics into token and auxiliary-velocity updates, then introduces decoder-only SympFormer architectures with geometric corrections. In most settings, the corrected discretizations outperform the comparison architectures in mean NLL, with the largest gains in looped architectures, but require more computation time.

8 Discrete SympFormer architectures

  • In most settings, geometrically corrected discretizations outperform the nanoGPT2-style baseline and an accelerated architecture without geometric correction in mean NLL, especially in looped architectures.The geometric correction increases computational time compared with the other methods.Source: §8 · p. 48

8.1 Time Discretizations of Riemannian Accelerated Gradient Flows

  • The time discretizations update token positions and carry a displacement or velocity variable between layers, with alternative schemes evaluating the connection at different displacements.The fixed-step damping and force coefficients are β_k = 1−hα(t_k) and γ_k = h^2.Source: §8.1 · p. 48

8.2 SympFormer Architecture

  • The decoder-only, causal, multi-head nanoGPT-style implementation modifies attention updates while retaining the momentum baseline’s MLP update.Attention and MLP updates have independently learned coefficients; the architecture also carries an auxiliary velocity stream between layers.Source: §8.2 · p. 49
  • Two SympFormer architectures replace the exact connection term with an approximation when discretizing the accelerated updates.The connection term is determined by the position-dependent metric.Source: §8.2 · p. 49

6. Look-ahead SympFormers HLA 1 and RLA

The section describes look-ahead SympFormer variants and a causal multi-head connection proxy used in their updates. The proxy is an approximation, and its added computations incur training cost.

6. Look-ahead SympFormers HLA 1 and RLA

  • Setting the correction to zero in H1 or R2 recovers H0, and all variants retain the MLP update.Source: §6 · p. 50
  • H1 and R2 use a computable connection proxy rather than the exact Levi–Civita connection, with LayerNorm-altered updates and causal multi-head attention.Source: §6 · p. 50

8.3 Causal Multi-Head Realization and Connection Proxy

  • The causal multi-head correction adapts the symmetric particle metric’s algebraic form to attention weights and sets B=Id.Source: §8.3 · p. 51
  • The proxy uses normalized tokens, update directions, causal attention weights, directional score derivatives, and weighted value aggregation.Source: §8.3 · p. 51
  • Additional directional projections, score contractions, and weighted value aggregation incur training cost without another softmax evaluation or matrix inverse.Source: §8.3 · p. 51

9 Numerical Results

The experiments compare geometric corrections with matched correction-off controls on TinyStories and OpenWebText using validation NLL and computational cost. Outcomes vary by attention mechanism and dataset, while the softmax results show small mean differences relative to seed variability.

9.1 Softmax Attention

  • For softmax attention, R2 has the lowest mean NLL on both datasets, but differences are small relative to seed variability and do not establish significance.Source: §9.1 · p. 52, Table 2 · p. 52
  • Geometry-corrected softmax models match H0’s parameter count but require more computation.Source: §9.1 · p. 52, Table 2 · p. 52

9.1.1 Parameter Sharing Across Depth

9.2 Linear Attention & Other Attention Mechanisms

  • R2 has the lowest mean NLL for linear attention on both datasets and native-L2 attention on OpenWebText, but geometric corrections are not uniformly beneficial.Source: §9.2 · p. 53, Table 4 · p. 54

10 Related results

This section situates the paper’s attention geometries and inertial dynamics in related work and develops regular-variation, Hamiltonian, and functional-analytic foundations for local-mobility and quotient Stein geometries, without assuming density spaces are smooth manifolds. It establishes results on discrete blow-up depth, Stein-operator limitations, fixed-ambient variations, and the well-posedness and unique solvability of the local-mobility gradient identity.

10 Related results

  • Theorem 7.26 quantitatively predicts the first discrete layer reaching a second-moment threshold from continuum blow-up time, up to bounded layer-index error as step size decreases.At fixed effective depth, maximum token norms remain bounded below the critical depth and diverge at or above it; Corollary 7.28 translates the prediction into arbitrarily large token norms.Source: §10 · p. 55
  • The paper supplies analytic domains for local-mobility Otto–Wasserstein and softmax–Wasserstein geometry, a quotient Stein calculus for linear attention, and Hamiltonian well-posedness and blow-up results.Theorem 4.3 identifies weighted function spaces where mobility is invertible and Corollary 4.4 establishes the induced metric; for linear attention, Theorem 7.18 proves strong well-posedness and Theorem 7.23 proves finite-time moment blow-up under bounded damping. Theorem 6.7 gives additional domain and compatibility conditions for regular softmax Hamiltonian derivatives, whose trajectory-domain preservation requires separate analysis.Source: §10 · pp. 55–56
  • Unlike prior momentum-based attention architectures including YuriiFormer [91], the proposed updates add token-dependent attention-metric connection corrections to position and momentum updates.Lemma 6.11 identifies the correction, and Tables 2 and 4 compare its predictive benefit and computational cost against matched momentum controls.Source: §10 · p. 56

A Non-Surjectivity of Stein’s Operator on Compact Domains

  • On the flat torus, an integrally strictly positive-definite kernel yields an injective but non-surjective smooth Stein operator, preventing a general metric or calculus on the full smooth tangent space.This limitation does not rule out the restricted spaces constructed in.Source: §10 · p. 55, A Non-Surjectivity of Stein’s Operator… · p. 64

B Differentiability in the Local Mobility Calculus

  • The local-mobility calculus defines admissible variations and regular first variations without assuming a smooth-manifold structure, with admissibility established for Otto–Wasserstein and softmax–Wasserstein geometries.A C1 functional alone need not have a smooth representative of its differential; the calculus applies when such a representative exists.Source: B Differentiability in the Local… · pp. 65, 67, 68
  • Regular variation classes and their extensions are unique, but the fixed-ambient convention does not define an intrinsic differential on the cotangent bundle.Changing how cotangent vectors at different densities are compared can change the density variation.Source: B Differentiability in the Local… · p. 69
  • Otto–Wasserstein variations satisfy admissibility automatically, whereas softmax–Wasserstein variations require an additional compatibility condition; on closed manifolds, regular variations reduce to ordinary partial differentials and damped Otto–Wasserstein equations.On a closed, connected manifold, the regular cotangent space is independent of density and every tangent vector is compatible with every smooth potential.Source: B Differentiability in the Local… · pp. 70–71

C Differentiability in the Quotient Stein Geometry

  • The regular Stein gradient exists uniquely in the effective tangent space, while the local regular first variation is unique only modulo constants.Directional limits are required only for admissible variations and do not alone imply a chain rule along arbitrary curves.Source: C Differentiability in the Quotient… · pp. 72–73
  • In Stein Hamiltonian calculus, a regular kinetic variation requires an additional H1(ρ) regularity condition, and its fixed-ambient density variation can depend on the full potential class rather than the effective cotangent class.The framework specifies damped Stein Hamiltonian equations for trajectories where the regular variations exist.Source: C Differentiability in the Quotient… · pp. 74–75

D Proofs of Section 4

  • The proof of Section 4’s functional-analytic lemma establishes equivalence between the local-mobility norm and the weighted H1 norm.The argument uses the local-mobility Poincaré inequality.Source: D Proofs of Section 4 · p. 75
  • The proof establishes that L_{mρ,B} is well-defined, surjective, and injective, using elliptic regularity for surjectivity and positivity for injectivity.The weak solution is shown to be smooth and to satisfy L_{mρ,B}Φ̇ρ = ρ̇; a vanishing energy integral forces Φ to be constant, hence its equivalence class is zero.Source: D Proofs of Section 4 · pp. 78–79
  • The proof of Corollary 4.6 verifies that the proposed expression satisfies the defining gradient identity, with uniqueness following from positive definiteness of G_{mρ,B}.Source: D Proofs of Section 4 · p. 79

E Proofs for Section 5

The proofs identify the quotient cotangent space with the finite-dimensional range of the linear-attention mobility operator. They also establish the operator’s kernel properties and verify the gradient identity by approximation.

Proofs for Section 5

  • A cutoff argument and integration by parts justify the gradient identity by passing to the limit in the gradient integral.The cutoff and score bound support the limiting argument.Source: E Proofs for Section 5 · p. 80
  • Positive definiteness of A yields Mρ(Φ) = 0 when the linear-attention mobility operator sends [Φ] to zero.This establishes the converse inclusion in the kernel characterization.Source: E Proofs for Section 5 · p. 81
  • The quotient cotangent space is isomorphic to the finite-dimensional range of the linear-attention mobility operator.The operator factors through the covariance matrix Mρ, whose range lies in Rd×d.Source: E Proofs for Section 5 · p. 81
  • The mobility operator’s kernel is infinite-dimensional because it has finite codimension in an infinite-dimensional cotangent space.The cotangent space is infinite-dimensional, while the operator’s range is finite-dimensional.Source: E Proofs for Section 5 · p. 81
  • The quotient map is well-defined and is an isomorphism because equivalent classes differ by elements of the operator’s kernel.Its kernel is zero by the quotient definition, and it is onto by the definition of the effective cotangent space.Source: E Proofs for Section 5 · p. 81

F Proofs for Section 6

This appendix proves the Section 6 variation, regularity, and lifespan claims, including the covariance-matrix reduction for linear-attention geometry and dynamics and the link between continuous moment blow-up and finite-token threshold-crossing depth. It also documents the implementation, data-preprocessing, training, evaluation, cost, layout, and update-coefficient conventions.

F Proofs for Section 6

  • The interaction energy has regular first variation [−Zρ] exactly under condition (59), and the proof verifies the required weighted-gradient condition.Compactly supported zero-integral tests establish necessity; cutoff estimates and finite energy establish the converse.Source: F Proofs for Section 6 · pp. 81–82
  • The metric identity (60) yields the softmax gradient-flow formulation at densities satisfying its assumptions.The proof combines Proposition 6.2’s first variation with differentiation of the normalizer.Source: F Proofs for Section 6 · p. 82
  • Under the stated integrability and positivity conditions, the nonlocal term in (61) is smooth and finite, giving its pointwise formulation whenever α(t) is finite.The argument uses strict positivity of ρ and invertibility of A to control the nonlocal integral, while the remaining spatial terms are smooth.Source: F Proofs for Section 6 · p. 83
  • Theorem 6.7’s variation combines kinetic and interaction contributions to give (61), with the potential equation understood modulo spatial constants.Condition (64) ensures the density variation defines a regular cotangent class.Source: F Proofs for Section 6 · pp. 83–84

G.2 Accelerated Linear Attention: Geometry and Hamiltonian Dynamics

  • The accelerated linear-attention equations in covariance and momentum coordinates are equivalent to the damped Hamiltonian and Euler–Lagrange formulations.The proof identifies the associated Lagrangian and Legendre transform, connecting the matrix dynamics to the damped geometric equation.Source: §G.2 · pp. 87–88

2. Since det Gt > 0, the pushforward in (92) has density

  • The pushforward construction gives a strong solution, and reducing any strong solution to matrix and scalar linear ODEs establishes uniqueness.The proof uses an invertible fundamental matrix to recover the density and potential representation.Source: §2 · pp. 88–89
  • The pushforward representation preserves smoothness, strict positivity, and unit mass, while moment estimates place ρt in P4 throughout the existence interval.Uniqueness of the matrix ODE system yields the representation; the estimates establish the fourth-moment condition.Source: §2 · p. 89
  • Global existence holds when every v_i ≤ 0; if some v_i > 0, a positive-forcing component blows up in finite time, with the matrix and PDE lifespans equal.The positive-forcing components are decoupled, and the maximal lifespan is the minimum of their component lifespans.Source: §2 · p. 92

G.3 Finite-Token Layers and Effective Depth

  • On compact time intervals, the finite-token Euler approximation has O(h) error for token matrices and second moments.The proof applies an Euler estimate to the finite-dimensional system and transfers its error bound to second moments.Source: §G.3 · p. 93
  • The discrete threshold-crossing time satisfies hLR(h) = tR + O(h), and the continuous crossing times tR approach Tmax as R increases.The proof locates the first discrete crossing within a time window proportional to h around tR.Source: §G.3 · pp. 93–94

H Implementation Details

  • Experiments use GPT-2 BPE tokenization with 50,257 token IDs and a padded 50,304-token vocabulary; the appendix also specifies paired data sampling and shared optimizer settings.TinyStories uses its supplied splits; OpenWebText is split 99.5%/0.5% for training and validation. AdamW uses peak learning rate 6 × 10−4 and momentum coefficients (0.9, 0.95).Source: H Implementation Details · p. 95
  • Learned attention and MLP update coefficients are trained separately, use ten times the base learning rate with no weight decay, and are not constrained to a continuous-time damping law.Coefficients for distinct trainable blocks are also learned separately.Source: H Implementation Details · p. 96
  • Both Table 3 layouts execute eight block applications; tied models reuse two alternating blocks, while untied models start as identical copies that train independently.Streams continue without resets or repeated input injection, and training backpropagates through all applications; repeated-use gradients accumulate on shared parameters. Sharing covers attention and MLP weights, LayerNorms, and learned update coefficients.Source: H Implementation Details · p. 96
  • Paired methods share batches, tokenizer, data splits, and validation targets; reported endpoints use the final evaluation rather than the best checkpoint.Peak memory is measured over complete dataset/seed blocks.Source: H Implementation Details · p. 96
  • GPU-hours are mean allocated time per completed run, averaged over paired seeds separately by dataset; GPU-hours are reported separately by dataset and hardware.Runs use one GPU; sigmoid- and L2-attention experiments use A100s for TinyStories and RTX A6000s for OpenWebText, so compute comparisons remain within each attention family. Peak MiB is peak PyTorch-allocated tensor memory, and parameter counts include trainable embeddings, blocks, and final normalization, counting shared weights once.Source: H Implementation Details · pp. 96–97

Related papers

Loading 2610.31600v1…