Measuring Optimal Transport in Transformer Depth
Not on the map. Why?
Abstract (arXiv)
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.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 15 of its 22 references that are in the map. First posted on arXiv 1 Sep 2026.
Builds on
15 of 22 references.
References · 15
- 2022Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flowcited by 3,550
- 2023Pythia: A Suite for Analyzing Large Language Models Across Training and Scalingcited by 1,734
- 2023Eliciting Latent Predictions from Transformers with the Tuned Lenscited by 579
- 2020How to train your neural ODE: the world of Jacobian and kinetic regularizationcited by 228
- 2019Model Fusion via Optimal Transportcited by 217
- 2019Optimal transport mapping via input convex neural networkscited by 202
- 2023A mathematical perspective on Transformerscited by 160
- 2021Sinkformers: Transformers with Doubly Stochastic Attentioncited by 136
- 2022Neural Optimal Transportcited by 123
- 2018Gromov-Wasserstein Alignment of Word Embedding Spacescited by 119
- 2020Faster Wasserstein Distance Estimation with the Sinkhorn Divergencecited by 73
- 2020A Mathematical Theory of Attentioncited by 43
Built on by
Not cited yet.