Source-linked AI summary
Measuring Optimal Transport in Transformer Depth
Alexandre Quemy
TL;DR
The paper asks whether trained transformers move token-state clouds between layers at optimal cost and along the optimal token-pairing map. It measures these properties with exact assignments, sampling-floor correction, calibration, and cost decomposition on two Pythia models. Agreement with the optimal map rises with depth and training, while cost is optimal wherever transitions can be resolved except Pythia-160m’s final transition.
Problem
The paper asks whether a trained transformer moves its token states between consecutive layer clouds at the cheapest cost and along the optimal-transport map.
Method
The study compares Pythia-160m and Pythia-410m with exact optimal transport, correcting sampling floors, calibrating against known-optimal couplings, and separating common from token-specific motion.
Results
Agreement with the optimal map reaches 0.89 for Pythia-160m and 0.88 for Pythia-410m at the last transition, while network cost is optimal wherever resolved except Pythia-160m’s final transition.
Takeaways & Limitations
The last layer of a trained transformer acts as an optimal-transport map, the first does not, and middle-layer motion is often too small for single-layer judgment.
Takeaways & Limitations
With 4,000-token samples, only four of twelve Pythia-160m transitions and five of twenty-four Pythia-410m transitions are resolvable individually; middle layers require blocks.
Abstract
from arXiv · showhide
A transformer carries each token's state from layer to layer, and the whole vocabulary carried together forms a cloud that moves with depth. We ask whether a trained network moves this cloud the way optimal transport would: at the cheapest cost, and along the map that pairs each token with its optimal destination. We measure both on Pythia-160m and Pythia-410m, with an exact assignment between consecutive layer clouds, a measured sampling floor, calibration on couplings known to be optimal, and a split of the cost into the common shift of the cloud and the token-specific moves. At the last layer, both models move their tokens where the optimal-transport map sends them, at the optimal cost for Pythia-410m and slightly above it for Pythia-160m. At the first layer they do not. In between, single layers can be judged on cost at only two of ten transitions, and blocks of several layers move the cloud at close to the optimal cost. The agreement at the last layer is much weaker at initialisation (0.64 against 0.86) and grows with training.
1 Introduction
The paper asks whether trained transformers move token-state clouds between layers using optimal transport, both in cost and in token pairing. It introduces a calibrated measurement protocol and reports near-optimal behavior concentrated at the network’s exit.
- Optimal transport provides the cheapest pairing between layer clouds and, for squared distance, a unique map given by the gradient of a convex function.The map is y = ∇φ(x).
- The study tests whether Pythia-160m and Pythia-410m use optimal transport between consecutive layer clouds, both in cost and induced pairing.It evaluates token positions sampled from WikiText-103 and Piles.
- The distance from the cheapest move offers a label-free measure of wasted motion for comparing layers and models, while map agreement can track changes across training.
- The protocol corrects sampling error, calibrates against known-optimal couplings, separates common cloud shift from token-specific motion, and supplies an uncalibrated lower bound.
- The last layer follows the optimal-transport map and is cost-optimal except at the exit of Pythia-160m, while the first layer does not.
2 Related work
Related work studies token clouds, attention, regularized flows, and fitted depth dynamics, but this paper measures unregularized optimal transport between consecutive trained-transformer layer clouds. Its tests compare both transport efficiency and the network’s induced pairing with the exact optimum.
- Transformers as maps between measures: Prior theory models transformers as maps between token-state measures, whereas this paper empirically tests whether trained networks choose the cheapest map and pairing.
- Attention and optimal transport: Earlier attention studies place optimal transport inside attention with fixed regularization; this work tests unregularized transport between consecutive layer clouds.
- OT-regularised flows: OT-Transformer and rectified-flow methods impose optimal paths through kinetic-energy penalties, while this paper asks whether ordinary transformer training produces near-optimal coupling anyway.
- Depth as dynamics: Fitted linear dynamics describe depth with linear maps; the paper uses the optimal-transport linear baseline TG to test what network behavior extends beyond linearity.At the 160m exit, one third of the move lies beyond any linear map and 61% beyond TG while still following optimal transport.
3 Transport efficiency
The paper measures whether transformer layer transitions move token-state clouds at optimal-transport cost, correcting sampling floors and separating common shifts from token-specific movement. Where measurable, both models are usually near optimal, but the Pythia-160m exit is a clear exception and many small transitions remain unresolved.
- Setup: Optimal transport cost is the minimum over all pairings, while the network’s pairing yields efficiency eℓ; eℓ = 1 denotes a perfect mover.The analysis removes the common cloud shift because it is exactly optimal and often dominates total cost.
- Measurement protocol: The protocol samples 4,000 points, solves the assignment exactly, corrects the independent-sample floor, and calibrates readings on optimal couplings.Transitions with a floor exceeding 30% of the relevant cost are reported as unresolvable.
- Shift decomposition: The common shift dominates total cost: full-space A/P is 1.5–2.5 at most transitions and 84 at the Pythia-160m exit, while the camera changes A/P by up to threefold.Full-space results nevertheless show the same general behavior as the reduced-space analysis.
- Results: At three of four measurable Pythia-160m transitions and every measurable Pythia-410m transition, the centered network coupling is optimal within 10–15% instrument precision.At the Pythia-160m entry, a sliced bound gives true efficiency at least 0.94, placing the network within 6% of optimal.
- Results: At the Pythia-160m exit, token-specific efficiency is 0.86 versus a 1.00 calibration, implying 14–27% higher cost than the cheapest move; Pythia-410m shows no gap.Whether this difference reflects model properties or scale remains open.
- Resolution limits: Between layers 4 and 10, single-layer token-specific moves are unresolved, but accumulated blocks become measurable: 3→9 reads 0.97 versus 0.95 calibration and 8→11 reads 0.96.The Pythia-410m block 4→20 is the only resolvable block reported for that model.
4 Agreement with the Brenier map
The network’s token moves increasingly agree with the exact optimal-transport map toward the final layer, while agreement is weak at entry and varies across transitions.
- Agreement rises with depth, from barely matching the optimal-transport move at the first layer to matching it at the last.The shuffled-pairing null remains at ρ ≤0.08.
- At the Pythia-160m exit, ρ reaches 0.89 with median cosine 0.95; Pythia-410m ends at ρ 0.88 with median cosine 0.95.Seed ranges of ρ are ±0.01.
- Pythia-160m agreement is ρ 0.41 at entry, 0.38–0.62 in the middle, and 0.80 at transition 10→11.The 410m curve dips near 0.4 of depth where its tokens move least.
- At the Pythia-160m exit, the linear Gaussian Brenier map explains 39% of the token-specific move, while the remaining 61% still agrees with exact optimal transport.The remainder has ρ 0.83 and median cosine 0.94, with null 0.01.
- Exit agreement in Pythia-160m grows from 0.64 at initialization to 0.86 when trained, whereas entry agreement falls from 0.52 to 0.37.At transition L10→11, agreement rises from 0.18 to 0.77 across the same checkpoints.
5 Conclusion
The study finds that optimal-transport agreement strengthens with depth: the last layer closely follows the optimal map, while the first does not and middle layers are often unresolved individually.
- The study uses exact optimal transport between consecutive layers for two models, two corpora, and a 16-dimensional projection.
- The network’s token-specific move has the optimal cost at every resolvable transition except Pythia-160m’s last transition.
- The last-layer agreement is much weaker at initialisation and grows with training, while agreement at the entry decreases.
- 4,000 tokens per sample resolve only four of twelve Pythia-160m transitions and five of twenty-four Pythia-410m transitions.Middle layers therefore require assessment in blocks of several layers rather than one layer at a time.
- The study is limited to two small models from one family, two small corpora, and mean flows that raise agreement relative to raw states.
A Linear map and remainder
The linear-map control tests whether network transport can be explained by a Gaussian Brenier map fitted from cloud means and covariances, then evaluates the residual transport.
- The Gaussian Brenier map TG is fitted from the two clouds’ means and covariances and maps one Gaussian cloud to matching target statistics.
- The control measures how much TG explains the network’s token-specific move and whether the residual agrees with exact optimal transport.It uses R2 and rank agreement ρ for the linear component, then compares the remainder with exact transport using the same floor and shuffled null.
- 39% of the Pythia-160m exit move is explained by the linear map, while the remaining 61% of cost still agrees with optimal transport at ρ 0.83.
- At the Pythia-410m exit and Pythia-160m transition 10 →11, TG explains 90% and 73% respectively, leaving remainders below their measurement floors.Those remainder agreements cannot be measured reliably.
- The table reports linear-map R2 and ρ, remainder cost share and floor, remainder agreement, and full-map ρ for comparison.