# Measuring Optimal Transport in Transformer Depth

arXiv: https://arxiv.org/abs/2609.00748 · HTML: https://mlsanity.com/paper/2609.00748/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 15 of its 22 references that are in the map.
- First posted on arXiv 1 Sep 2026.

## Builds on (most cited first, 15 of 15)

- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [Eliciting Latent Predictions from Transformers with the Tuned Lens](https://mlsanity.com/paper/2303.08112/) (2023)
- [How to train your neural ODE: the world of Jacobian and kinetic regularization](https://mlsanity.com/paper/2002.02798/) (2020)
- [Model Fusion via Optimal Transport](https://mlsanity.com/paper/1910.05653/) (2019)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [A mathematical perspective on Transformers](https://mlsanity.com/paper/2312.10794/) (2023)
- [Sinkformers: Transformers with Doubly Stochastic Attention](https://mlsanity.com/paper/2110.11773/) (2021)
- [Neural Optimal Transport](https://mlsanity.com/paper/2201.12220/) (2022)
- [Gromov-Wasserstein Alignment of Word Embedding Spaces](https://mlsanity.com/paper/1809.00013/) (2018)
- [Faster Wasserstein Distance Estimation with the Sinkhorn Divergence](https://mlsanity.com/paper/2006.08172/) (2020)
- [A Mathematical Theory of Attention](https://mlsanity.com/paper/2007.02876/) (2020)
- [OT-Transformer: A Continuous-time Transformer Architecture with Optimal Transport Regularization](https://mlsanity.com/paper/2501.18793/) (2025)
- [Transformers through the lens of support-preserving maps between measures](https://mlsanity.com/paper/2509.25611/) (2025)
- [Lines of Thought in Large Language Models](https://mlsanity.com/paper/2410.01545/) (2024)

## 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.