arXiv · 2410.19931
Provable optimal transport with transformers: The essence of depth and prompt engineering
Abstract
Despite their empirical success, the internal mechanism by which transformer models align tokens during language processing remains poorly understood. This paper provides a mechanistic and theoretical explanation of token alignment in LLMs. We first present empirical evidences showing that, in machine translation, attention weights progressively align translated word pairs across layers, closely approximating Optimal Transport (OT) between word embeddings. Building on this observation, we prove that softmax self-attention layers can simulate gradient descent on the dual of the entropy-regularized OT problem, providing a theoretical foundation for the alignment. Our analysis yields a constructive convergence bound showing that transformer depth controls OT approximation accuracy. A direct implication is that standard transformers can sort lists of varying lengths without any parameter adjustment, up to an error term vanishing with transformers depth.
Explore related subjects
Keep this discovery
Hadi Daneshmand. 2024-10-25. Provable optimal transport with transformers: The essence of depth and prompt engineering. https://arxiv.org/abs/2410.19931
Cite the original work for its findings. Save a collection to share your selection of sources.