Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This study investigates how transformer models transport token states across layers, assessing whether this movement aligns with optimal transport principles. By analyzing the Pythia-160m and Pythia-410m models, the authors find that while both models achieve optimal transport at the final layer, only the larger model does so at the optimal cost, with significant discrepancies observed at earlier layers. The findings reveal that the efficiency of token movement improves with training, highlighting the dynamics of layer interactions in transformer architectures.
Transformers may not be as efficient in token transport as optimal transport theory suggests, especially in early layers.
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.