Provable optimal transport with transformers: The essence of depth and prompt engineering

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.

💡 Research Summary

The paper “Provable Optimal Transport with Transformers: The Essence of Depth and Prompt Engineering” investigates why and how transformer models align tokens across layers, especially in machine translation, by linking the process to optimal transport (OT). The authors first present empirical evidence that attention weight matrices in a standard transformer gradually approximate the optimal permutation matrix P* that solves a discrete OT problem between source and target word embeddings. By visualizing attention heatmaps at layers 1, 6, and 12 on a bilingual corpus, they show a clear convergence toward the OT solution.

To explain this phenomenon, the authors formalize the soft‑max self‑attention operation and demonstrate that each attention layer can be interpreted as one iteration of gradient descent on the dual of the entropy‑regularized OT problem. The dual objective L(u, v) depends on two vectors u and v, and its gradient updates are of the form u_{ℓ+1}=u_ℓ−D_ℓ∇u L(u_ℓ, v_ℓ) and similarly for v. Theorem 3.1 proves that, with a specific (input‑independent) configuration of query, key, and value matrices, a stack of ℓ attention layers exactly reproduces ℓ steps of this gradient descent. Theorem 3.2 then provides a convergence bound: after ℓ layers the attention matrix A^{(ℓ)} is within O(ℓ^{-1/2}) of the regularized OT plan P*{λ}. Crucially, the bound does not depend on the problem size n, implying that sufficiently deep transformers can solve OT (and thus sorting) for arbitrarily large n without any parameter re‑training.

The paper also explores the role of prompt engineering. By augmenting the input matrix Z₀ with additional features—norms of the embeddings, constant vectors, and zero‑padding—the authors give the model extra “memory” that allows the attention mechanism to store intermediate gradient iterates. Experiments on synthetic OT instances (training on n=7 points, testing on n=9) show that the engineered prompt dramatically improves both speed of convergence and final alignment accuracy. A larger experiment on 10,000 real translation pairs (English–French) confirms that standard transformers, when supplied with the engineered prompt, can align tokens almost perfectly across layers, even though the model was never explicitly trained for OT.

The authors compare their findings to prior work on Sinkhorn‑Attention, which modifies the attention computation to directly implement the Sinkhorn algorithm. They demonstrate that ordinary soft‑max attention already approximates the Sinkhorn fixed‑point iteration when depth is sufficient, removing the need for specialized attention variants.

Limitations are acknowledged: the theoretical results rely on entropy regularization and assume exact gradient‑descent dynamics, which may not hold for very deep or noisy real‑world models. Computational cost grows linearly with depth, so extremely deep networks may be impractical. The analysis also does not cover non‑quadratic cost matrices or unregularized OT.

In conclusion, the paper provides a mechanistic and provable link between transformer depth, attention dynamics, and optimal transport. It shows that depth controls OT approximation accuracy, and that carefully engineered prompts act as external memory, enabling transformers to solve sorting and matching tasks without parameter changes. This work reframes transformers as implicit optimizers, offering a new theoretical lens for interpreting their success in language tasks and suggesting new avenues for prompt‑based algorithmic control.