Official Review · ICLR 2025

This investigates whether transformers can provide theoretical guarantees in solving the Sinkhorn optimal transport problem. The authors demonstrate that transformers, equipped with appropriately engineered prompts, can approximate optimal transport solutions by leveraging attention layers and simulating gradient descent with adaptive step sizes. The authors designed prompts to aid transformers in solving optimal transport by providing the required memory and statistics within the model's layers. This setup helps in implementing gradient descent and enables transformers to handle various problem instances efficiently. The study finds that the performance of transformers in solving optimal transport improves with increased depth, establishing an explicit approximation bound that converges as the model’s depth increases. Sorting is shown as a special one-dimensional case of optimal transport, where transformers can approximate the sorted output by solving the transport problem. The experiments validate the theoretical predictions, showing that transformers can effectively approximate sorted sequences. The paper rigorously proves that attention matrices in transformers can converge to a doubly stochastic matrix that approximates the optimal transport solution. This convergence avoids rank collapse, a common issue in attention layers, thus preserving the expressiveness of transformers. Originality The paper is original in its approach to connecting transformers with provable guarantees in solving the Sinkhorn OT problem. It leverages the attention mechanism of transformers in an innovative way, showing that transformers can approximate optimal transport solutions under entropic regularization. Quality The paper is rigorous in both its mathematical formalism and experimental validation. The theoretical development is well-supported by proofs, such as the gradient descent simulation using attention layers and the convergence analysis of attention matrices. Clarity The paper generally communicates complex ideas effectively, with clear explanations of both the background and the transformer mechanics used in OT tasks. All proofs are included in the main paper. Significance Theoretically, it provides new insights into how transformers can function beyond sequence modeling, linking them to iterative optimization methods such as gradient descent. This connection strengthens the iterative inference hypothesis, suggesting transformers can perform complex tasks in optimization by layer-wise refinement. 1. The paper's theoretical proofs are presented with a high level of mathematical density that may hinder accessibility for a broader audience, particularly readers less familiar with optimization theory or transformers' inner workings. To improve accessibility, the authors could add intermediate explanations or visual aids (such as flowcharts or diagrams) to clarify the key steps in the proofs. An appendix with expanded explanations of some techn