Transformers With a Velocity Ledger

· 15 min read

#ml#transformers#residual-stream#momentum#nGPT#dynamical-systems#depth#architecture#physics

Part 2 of 5Networks as Integrators
  1. 1Your Skip Connection Is Half of Newton
  2. 2Transformers With a Velocity Ledgeryou are here
  3. 3A Network That Conserves Energy
  4. 4Backprop Without the Memory
  5. 5Depth on Demand
Runnable JAX companionA Velocity Ledger for Transformers, in JAX/Flax NNXPrefer to read the code? This post has a hands-on JAX / Flax NNX implementation.Open the JAX companion

The last post read a residual block x+F(x)x + F(x) as one forward-Euler step and gave a planet’s worth of stability to whatever second ledger you were willing to carry. It ended on a promise that felt too convenient to be true: if a ResNet is an integrator, then every residual network is, and the most consequential residual network of the decade is the Transformer. So the invitation on the table is to walk up to a real language model and ask the same question we asked of a toy ResNet. Does giving the residual stream a velocity change what the model does?

Look at a pre-norm Transformer layer, the standard modern kind, and the correspondence is already sitting there in plain sight:

x+=Attn(normx),x+=MLP(normx).x \mathrel{+}= \mathrm{Attn}(\mathrm{norm}\,x), \qquad x \mathrel{+}= \mathrm{MLP}(\mathrm{norm}\,x).

Two sub-updates per layer, each one a state plus an increment, each one a forward-Euler step on the residual stream xx. Depth is time, the stream is position, the attention and MLP branches are the vector field, exactly the dictionary from D1, now pointed at a two-hundred-dimensional stream of token representations instead of a width-2 plane. Which means the missing half of Newton is missing here too. Nothing accumulates. Each sub-update moves the stream directly, forgetting the direction of the last one.

So we did the obvious experiment. Take a small character-level GPT, hold everything fixed, and change only the residual-stream update: plain (the Euler layer above), ledger (D1’s exact velocity stream, one line), ngpt-lite (a first-order variant that lives on a sphere, no velocity), and ngpt-ledger (the ledger and the sphere). Three seeds each, best-val early-stopped, all in scripts/velocity_ledger.py.

Here is the headline, up front, because the story is more interesting once you stop hoping for the wrong thing. On quality, the four tie. Best validation loss lands at 1.433 (plain), 1.435 (ledger), 1.420 (ngpt-lite), 1.456 (ngpt-ledger), a spread of 0.036 against a per-seed noise of about 0.005. The velocity ledger did not make the Transformer better. But it did something the scoreboard cannot see, and it is exactly what the physics predicted: it changed the path. The residual stream reaches the same answer by a journey that is roughly a third as long and less than half as sharp. Same destination, calmer road. That, and not a win, is the true finding, and it is the one D1 told us to expect.

The dictionary still holds

Does the D1 correspondence really survive the move from a width-2 ResNet to a real Transformer, or does the analogy quietly break somewhere in the attention block? It survives, and it is worth pinning down as a table, because everything below reads off it:

pre-norm Transformernumerical integrator
layer index ll (and each sub-update within it)time tt
residual stream xlx_lposition
Attn(normx)\mathrm{Attn}(\mathrm{norm}\,x) and MLP(normx)\mathrm{MLP}(\mathrm{norm}\,x)the vector field being integrated
x+=F(normx)x \mathrel{+}= F(\mathrm{norm}\,x)one forward-Euler step, Δt=1\Delta t = 1
stacking LL layers (two sub-updates each)integrating for total time T=2LhT = 2L\,h
a velocity stream vlv_l (usually absent)the momentum ledger
its damping μ\muhow long inertia is remembered

The only genuinely new wrinkle is the norm: a pre-norm block reads a normalized copy of the stream, integrates the field there, and writes the raw increment back. That is a preconditioning of the field, not a change to the integrator; the update is still position-plus-increment, still Euler. This reading is not new either. Lu et al. (2019) read the whole Transformer as a multi-particle dynamical system and its layers as a Lie-Trotter splitting of an ODE; the ODE Transformer of Li et al. (2022) replaced the Euler step with a Runge-Kutta one and got a better integrator of the same field. Depth-as-time is an established lens on Transformers. What D1 adds is the specific missing piece, the velocity, and a specific prediction about what carrying it does to the trajectory.

And the prediction transfers verbatim. In D1, the velocity ledger’s signature was not accuracy, it was geometry: consecutive Euler steps forget each other and the path kinks, while a velocity that persists across steps makes the path coast, smooth and short. The plain ResNet on rings turned 53° per layer; the momentum net turned 10° on a path of nearly the same length. So the D1 question for a Transformer is not “does the ledger raise accuracy,” it is “does the ledger straighten the residual-stream trajectory.” Hold that question. It is the one the data answers loudly.

This idea has precedent, and that is fine

Before training anything, a question of provenance: is a velocity in the residual stream a new idea? It is not. This is a clean, instrumented demonstration of a known effect plus a deliberate contrast, not a novelty claim.

Where this sits in the literature. Closest is Momentum Streams for Optimizer-Inspired Transformers (2605.24425), which builds exactly this: an auxiliary velocity (and moment) stream running alongside the residual stream, with heavy-ball, Nesterov, and Adam-like variants. Their framing is ours: the velocity stream turns the first-order residual map into a second-order recurrence in depth, which is D1’s thesis stated inside a Transformer. YuriiFormer (2601.23236) reads GPT layers as gradient descent on a token energy and adds Nesterov/Polyak momentum in the representation at no extra attention or MLP evaluations, the same acceleration from a different door. The Momentum Transformer (2208.00579) puts momentum inside linear attention rather than the residual stream. The older ODE lineage (Lu et al., 2019; ODE Transformer, 2203.09176) is the depth-as-dynamics reading we are standing on. And nGPT (2410.01131) is the deliberate contrast: a first-order variable-metric optimizer on the hypersphere, with normalized displacements and a retraction step but no velocity at all. Our ngpt-lite is a simplified nGPT (the sphere and per-dimension step scales, without the weight and logit normalization); it is the first-order-on-a-sphere corner that makes the ledger’s effect legible by contrast.

So the experiment is not “we invented momentum for Transformers.” It is “the depth dynamics D1 predicted are real and measurable in a Transformer, and here is exactly what the velocity buys and what it does not, with the nGPT sphere as a control.”

The scope, before the claims

One caveat has to come before the results, because the results are about dynamics, and dynamics claims are easy to overstate.

What this experiment is and is not. It is one small character-level GPT (6 layers, 192-wide, 4 heads) on one corpus (the complete works of Shakespeare, character level), trained for 12000 steps with dropout, best-val early-stopped, three seeds. It is an illustration of a dynamical effect at fixed parameter count, not a scaling claim, not a recommendation, and not evidence that a velocity ledger helps at any size. All four variants are matched to within 0.1% of parameters (the ledger adds zero parameters; it is a depth-state, not a weight). The depth telemetry is measured on a fixed held-out probe batch. Read every number below as “on this model, this is what the ledger does to the trajectory,” and nothing wider.

The tie on quality

Every hope you might bring to a new architecture knob dies on the first figure, and it is better to look straight at it. Best validation loss, three seeds, four variants, on a zoomed axis so the whole spread is visible:

There is one more number worth stating so the earlier version of this story is not misremembered: the over-training blow-up is gone. In an earlier, smaller-corpus run the ledger and plain variants over-trained at very different rates, which looked like a robustness signal. With dropout and a five-times-larger corpus, every variant’s final val loss sits essentially on top of its best (blow-up ratio 1.00 across the board), so there is no hidden robustness win either. Just a clean four-way tie on quality. Which is exactly why the interesting part is not on this axis at all.

Where the signal actually lives

If the four models end up equally good, are they doing the same thing to get there? Not remotely. The whole payoff of the depth-as-time reading is that it hands us an axis the scoreboard does not have: the geometry of the path the residual stream takes from the embedding to the readout. Measure that path on the probe batch and the four variants separate cleanly.

Two things to read carefully here, because they are the crux. First, the ledger path is dramatically shorter: the residual stream does not need to travel 128 units to land where plain lands; carrying a velocity gets it to an equally-good answer in 48. Second, it is dramatically straighter: 29 degrees of bend per step against plain’s 75. This is precisely the D1 signature. In the ResNet, the momentum net made the same-length journey without the kinks; in the Transformer the momentum stream makes a shorter journey and still without the kinks. The velocity is doing the same job it did on the planet: accumulating direction across steps so the state stops re-deciding where to go at every layer.

The ngpt contrast sharpens it. ngpt-lite, first-order on the sphere, bends the most of all, 98 degrees a step, because retraction to the sphere after every sub-update keeps re-orienting the stream; its path is short only because the sphere caps how far it can wander. Add the ledger to it and the bend drops to 47 and the path to 32: the velocity smooths even the sphere’s re-orientations. Momentum straightens the path whether the geometry is flat or spherical.

You can see the trajectory itself. The residual stream is 192-dimensional, so this is a 2D shadow, but it is a faithful shadow: every segment you see has the real measured length, and every bend has the real measured angle. The verdict your eye reads, shorter and straighter, is exactly the model’s measured verdict.

Why the ledger’s path is short

Where does the shortening come from, mechanically? The velocity update is one line, D1’s exactly:

vl=μvl1+(1μ)F(normxl1),xl=xl1+hvl,v0=0,v_{l} = \mu\, v_{l-1} + (1-\mu)\, F(\mathrm{norm}\,x_{l-1}), \qquad x_l = x_{l-1} + h\, v_l, \qquad v_0 = 0,

with μ=0.9\mu = 0.9, h=1h = 1 fixed, one velocity stream shared across both sub-updates of every layer, zero-initialized at the embedding. Set μ=0\mu = 0 and it collapses exactly to the plain layer x+=F(normx)x \mathrel{+}= F(\mathrm{norm}\,x), so the ledger is one dial from Euler to heavy-ball, at identical parameter count.

The (1μ)(1-\mu) factor is the whole story. A plain layer writes the full block output FF into the stream; the ledger writes a running average of FF, damped by 0.90.9. Early in depth, before the average has filled, each step is tiny; the velocity accumulates gradient across layers and only opens up near the readout. Watch the actual step sizes:

The plain net’s step size balloons toward the readout, dumping a 25-unit increment in the final sub-update, which is where much of its 128-unit path and its stream-norm blow-up to 72 come from. The ledger’s steps stay small and even, because the (1μ)(1-\mu) damping refuses to let any single block output move the state far, and the velocity that funds each step is a smoothed accumulation, not a raw output. That is the whole shortening: the ledger integrates the field instead of chasing it.

The path straightens as training proceeds

Is the straight path there from initialization, or does the model learn to keep it short? The telemetry was dumped at four checkpoints (init, ten percent, fifty percent, and the end of training), so we can watch the gap open.

At initialization everything is short, because an untrained field is small. Training is what pushes the plain net’s path out: as the plain net learns to separate characters, its per-layer steps grow, its stream-norm climbs, and its path lengthens toward 128 and stays there. The ledger, over the same training, holds its path near 48. The straightness is a trained property of carrying a velocity, not a coincidence of the starting point, and it survives to the checkpoint that scores best. Same best-val, and one path a third the length of the other, at every point past the first tenth of training.

Same answer, gentler journey

So what did the velocity ledger buy a Transformer? Not accuracy. On this model the four variants tie on best validation loss to within seed noise, the over-training blow-up is gone, and the ledger is not a quality upgrade. That is the headline, and it is a tie.

What it bought is exactly what D1 said a second ledger buys: a calmer integrator. The residual stream reaches an equally-good answer along a path that is roughly a third as long and less than half as sharp, and it does so whether the geometry is flat (ledger vs plain) or spherical (ngpt-ledger vs ngpt-lite). The velocity accumulates direction across depth, so the stream stops re-deciding where to go at every sub-update and starts coasting, the way a thing with inertia moves. The plain net gets to the same place by lurching, a fresh 25-unit jump at the readout; the ledger glides.

Whether a shorter, straighter residual-stream path is worth anything is a genuinely open question this experiment does not answer. It might matter for very deep stacks, where a first-order flow’s accumulated wander could be a real cost; it might matter for interpretability, since a straight path is easier to read; it might not matter at all at this scale, which is exactly why the word here is tie. The one thing the data settles is that D1’s prediction transferred and held: a velocity in the residual stream is a smoother integrator of the same field, and the calm it produces is real, measurable, and precisely the calm the physics promised. A skip connection is half of Newton in a Transformer too, and the other half, once again, is a straighter path to the same destination.


Every number in this post comes from scripts/velocity_ledger.py (JAX + Flax NNX, four parameter-matched char-level GPTs, three seeds, best-val early-stopped, with depth telemetry on a fixed probe batch), reshaped for the panels above by scripts/export_velocity_viz.py. The runnable companion builds all four residual-stream updates in Flax NNX and renders the depth telemetry from the same bundle.

Cite as

Bouhsine, T. (). Transformers With a Velocity Ledger. Records of the !mmortal Data Scientist. https://tahabouhsine.com/blog/transformers-with-a-velocity-ledger/

BibTeX
@misc{bouhsine2026transformerswithavelocityledger,
  author       = {Bouhsine, Taha},
  title        = {Transformers With a Velocity Ledger},
  year         = {2026},
  month        = {jul},
  howpublished = {\url{https://tahabouhsine.com/blog/transformers-with-a-velocity-ledger/}},
  note         = {Blog post, Records of the !mmortal Data Scientist}
}

References

  1. Anonymous (2026). Momentum Streams for Optimizer-Inspired Transformers. preprint.arXiv:2605.24425
  2. Anonymous (2026). YuriiFormer: Transformer Layers as Momentum Gradient Descent on Token Energies. preprint.arXiv:2601.23236
  3. Loshchilov, I., Hsieh, C.-P., Sun, S., Ginsburg, B. (2024). nGPT: Normalized Transformer with Representation Learning on the Hypersphere. preprint.arXiv:2410.01131
  4. Nguyen, T. M., Baraniuk, R. G., Kirby, R. M., Osher, S. J., Wang, B. (2022). Momentum Transformer: Closing the Performance Gap Between Self-Attention and Its Linearization. preprint.arXiv:2208.00579
  5. Lu, Y., Li, Z., He, D., Sun, Z., Dong, B., Qin, T., Wang, L., Liu, T.-Y. (2019). Understanding and Improving Transformer From a Multi-Particle Dynamic System Point of View. ICLR 2020 Workshop.arXiv:1906.02762
  6. Li, B., Du, Q., Zhou, T., Jing, Y., Zhou, S., Zeng, X., Xiao, T., Zhu, J., Liu, X., Zhang, M. (2022). ODE Transformer: An Ordinary Differential Equation-Inspired Model for Sequence Generation. ACL 2022.arXiv:2203.09176
  7. Sander, M. E., Ablin, P., Blondel, M., Peyré, G. (2021). Momentum Residual Neural Networks. ICML 2021.arXiv:2102.07870