Source-linked AI summary
The Depth Flow of Token Representations Is Nonlinear and Does Not Descend Its Own Density
Alexandre Quemy
TL;DR
The paper asks whether the depthwise flow of token representations is linear and whether it descends its own density. It fits discrete Langevin dynamics to corpus-mean trajectories from Pythia models and evaluates predicted steps on held-out tokens. The flow is nonlinear, does not follow the log-density gradient, and preserves angular rank while shuffling norm and reversing concentration rank.
Problem
The paper investigates whether layer-to-layer transformations of token representations can be described by linear drift and whether the flow descends its vocabulary log-density.
Method
The authors fit a discrete Langevin equation to corpus-mean vocabulary-token trajectories, comparing Kramers–Moyal, linear, and quadratic estimators on held-out tokens.
Results
The drift is nonlinear: quadratic drift beats the linear map across the studied Pythia models, while the drift does not descend log-density and 4 to 45% of explainable drift circulates.
Takeaways & Limitations
Token representations are pushed by a nonlinear field that preserves angular order while moving the vocabulary, rather than sliding down a single density landscape.
Takeaways & Limitations
The study covers only the Pythia family at small scale, uses a single corpus, and may be under-resourced for Kramers–Moyal estimation.
Abstract
from arXiv · showhide
A token's representation is carried through the network layer by layer. The whole vocabulary carried together forms a flow. We fit this flow's equation of motion as a discrete Langevin model over corpus-mean trajectories of Pythia-160M and Pythia-410M, and score the predicted steps on held-out tokens. Linear maps are often used as cheap surrogates for a layer. The flow they summarize is not linear: a quadratic drift beats the linear linear map at every transition of both models, and the Kramers--Moyal estimator agrees wherever its neighborhoods stay local. We then characterize the flow further. First, we show that it does not descend its own log-density. The drift instead descends a potential that is not the density. Second, the rotational component is not negligible, $4$ to $45\%$ of the explainable drift, and the circulation shows in what the flow preserves: a token keeps its angular rank across all thirteen layers while its norm rank is shuffled and its concentration rank is reversed by the last block.
1. Introduction
The paper treats corpus-averaged token representations as a depthwise flow and fits its layer-to-layer dynamics with Langevin models. It finds nonlinear drift, failure of the log-density-gradient view, and preservation of angular rank rather than norm or concentration rank.
- Token-representation flow: Each vocabulary token follows a corpus-mean trajectory across residual-stream layers, and the vocabulary collectively forms a flow.The trajectory averages states at positions preceding occurrences of the token, rather than beginning at its own embedding row.
- Equation of motion: The study fits a Langevin equation with drift and diffusion, comparing Kramers–Moyal, linear, and quadratic estimators for layer-to-layer transformations.Predicted steps are evaluated on held-out tokens.
- Nonlinear dynamics: Quadratic drift and Kramers–Moyal beat the linear map at every layer of Pythia-160M and Pythia-410M.This tests whether linear maps adequately summarize the depthwise flow.
- Gradient and rotation: The drift is not the gradient of the vocabulary’s log-density; the motion instead contains a gradient component alongside diffusion and a load-bearing rotational component.The paper therefore rejects interpreting the flow as simply diffusing toward dense regions.
- Preserved structure: Only cosine to the ensemble mean direction is strongly preserved among cosine, vector norm, and participation ratio.This supports the importance of the rotational part of the dynamics.
2. Related work
The paper distinguishes its depthwise, vocabulary-level mean-trajectory flow from prior analytical, empirical, and generation-time approaches. It also tests and extends the closest prior stochastic model by allowing nonlinear drift, state-dependent correlated noise, and a gradient interpretation that can fail.
- Positioning the object: This work studies one corpus-mean trajectory per vocabulary token across depth, fits its equation from those trajectories, and scores it on held-out tokens.This differs from approaches using continuous-limit derivations, individual prompts, or generation time.
- Closest prior work: The closest prior work fits layerwise stochastic trajectories with a linear rotation-and-stretch drift and isotropic, state-independent noise.Its gradient interpretation is not tested, whereas this paper tests all three assumptions.
- Extensions: The paper’s model differs by testing nonlinear drift, correlated state-dependent noise, and the gradient reading of the flow.These tests correspond to the paper’s central departures from the closest prior model.
- Baseline: Yom Din et al.’s linear hidden-state map is the baseline because it transfers accurately across models and scales.The paper compares its candidate laws of motion against this baseline.
- Methods in context: The paper uses established instruments including Kramers–Moyal drift-and-diffusion estimation and the classical Langevin–Fokker–Planck correspondence.These tools also appear in stochastic-gradient and score-based generative-model settings.
3. The flow and how we measure it
The paper models token representations as a layerwise flow using corpus-mean trajectories, reduced states, and layer-specific discrete Langevin dynamics. It evaluates local, affine, and quadratic estimators on held-out token transitions while accounting for occurrence-dependent noise.
- The trajectories: 25,268 token trajectories contain thirteen points in R768 after filtering tokens with fewer than five occurrences.They are collected from 25,000 WikiText-103 articles processed by Pythia-160M.
- The reduced state: 16 PCA components retain 86.7% of pooled variance, although retention varies from 11% at layer 0 to 79% mid-network and 52% at the output.A shared basis is fitted across layers after standardizing each layer’s cloud to account for residual-stream norm growth.
- The model: The model treats each token trajectory as a discrete-time Langevin process with layer-specific push and noise functions of the current state.The state is the 768-dimensional mean projected onto a shared 16-component basis; coefficients do not depend directly on token identity.
- The estimators: The Kramers–Moyal estimator uses the k = 64 nearest training states, taking their mean step as the push and covariance as the noise.The affine and quadratic alternatives use least-squares fits; the affine map is the linear baseline and the quadratic map adds pairwise coordinate products.
- The protocol: Held-out evaluation predicts each next-layer step as a Gaussian and scores drift with R2 and the full model with mean log-likelihood.All fits use one 80/20 token split.
- The noise floor: Rare-token trajectories are noisier by construction: the noise is 40% of typical inter-token distance at five occurrences, 21% at ten to a hundred, and 2% above a thousand.This occurrence-dependent noise remains a limitation of the protocol.
4. The drift is nonlinear
The drift is nonlinear: quadratic and local estimators outperform the affine baseline across transitions, although local estimation degrades when neighborhoods become nonlocal. Rollouts also reveal a limitation of the fitted noise model, which causes generated token clouds to collapse.
- Nonlinear drift: KM wins eleven of twelve variance-explanation transitions and all twelve matched-noise likelihood contests, with R2 values of 0.54 to 0.96 versus 0.49 to 0.91.The quadratic map wins all twelve transitions in the same working-width comparison.
- Noise and likelihood: 8.4 to 29.1 nats: KM beats per-transition and pooled state-blind Gaussian nulls on all twelve transitions.With full covariance for both models, KM still wins all twelve by 1.1 to 7.6 nats, with bootstrap intervals excluding zero.
- Estimator limits: KM wins twelve transitions at 8 components, five at 32, and one at 64 as its fixed neighborhood expands from 11% to 28% of the cloud radius.Increasing k to 256 or 1,024 gives zero wins, whereas the quadratic map wins twelve of twelve at all four widths.
- Visual evidence: Figure 2 compares affine, quadratic, and local drift fields against black arrows representing mean held-out steps per spatial bin.The plotted transition uses the leading principal plane of layer 2, where in-plane curvature is largest.
- Trajectory generation: Iterated fitted-law rollouts shrink token spread to 0.53 of real mid-network, 0.87 by the last raw layer, and one quarter after the final transition.The likely cause is repeatedly redrawing persistent token identity as if it were step-to-step noise.
5. The drift is not the gradient of the log-density
The paper tests whether token-representation drift follows the gradient of its own log-density and finds no such alignment. Instead, the drift combines descent on another potential with a substantial rotational component.
- Interpretive frame: The tested log-density interpretation treats layers as diffusing particles toward dense vocabulary regions, distinct from architecture-level energy descent.The two readings concern different objects: whole-vocabulary log-density versus context-conditioned architectural computation.
- Log-density test: The log-density-gradient reading fails at every transition, with a mean alignment gap of −0.008 against a shuffle null.Per-transition gaps range from −0.116 to +0.053 across state dimensions and both Pythia models.
- Potential and circulation: A best cubic potential explains 55 to 96% of the unconstrained quadratic field, leaving 4 to 45% as rotational drift.The rotational remainder cannot be produced by any potential.
- Interpretation: The flow is not summarized by one landscape: potential descent, circulation, and correlated noise jointly contribute to its dynamics.The paper connects circulation with preservation of angular structure across depth.
6. What the flow conserves
Across thirteen layers, angular rank is the only measured token observable strongly preserved; norm rank is shuffled and participation-ratio rank is reversed late.
- Among cosine to the mean direction, vector norm, and participation ratio, only cosine is strongly preserved.
- The cosine to the ensemble mean direction is preserved with Spearman rank correlation 0.76.
- The vector norm rank is shuffled across layers, with Spearman rank correlation −0.17.
- Participation-ratio rank is partly reversed at the last layer, with Spearman rank correlation −0.59.
7. Conclusion
The study fits a discrete Langevin law to corpus-mean token trajectories and finds nonlinear, non-density-gradient dynamics with circulation, while identifying scope and surrogate limitations.
- 25,268 corpus-mean trajectories across thirteen Pythia-160M and Pythia-410M states were fit with a discrete Langevin model and scored on held-out tokens.
- Both nonlinear estimators beat the linear map on Pythia-160M and Pythia-410M.
- The drift does not descend its own log-density; it mostly descends another landscape, with 4 to 45% circulating.
- Angular rank survives all thirteen layers, while norm rank is shuffled and concentration rank is reversed by the last block.
- The study covers only the Pythia family at small scale and uses a single corpus, while the fitted equation is not yet a trajectory surrogate.
- The paper concludes that token representations are pushed by a nonlinear field that keeps angular order while moving the vocabulary.
Appendix A. The gradient test and its calibration
The appendix calibrates a gradient-flow test against stationary Gaussian controls and compares its readings with the real token-representation flow.
- A stationary Gaussian control with known gradient drift provides a calibration floor for interpreting near-zero alignment.
- The estimated gradient direction agrees with the analytic direction at cosine 0.914.
- Across drift-to-noise ratios from 0.9 to 106, the control reads alignments of 0.61 to 0.91, while real-flow shuffle nulls stay within 0.004 of zero.
- The control places 43 to 81% of drift energy along the estimated gradient, compared with 6 to 29% on the real flow.
- A shared translation can hide a residual gradient component, reducing the control reading from 0.90 to 0.42 when injected.