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

ICLR 2025 Conference Submission12353 Authors

Keywords: Theory of Language Models, Deep Learning Theory, Optimal Transport, Optimization

Abstract: Can we establish provable guarantees for transformer performance? Providing such theoretical guarantees is a milestone in developing trustworthy generative AI. In this paper, we take a step toward addressing this question by focusing on optimal transport, a fundamental problem at the intersection of combinatorial and continuous optimization. Leveraging the computational power of attention layers, we prove that a transformer with fixed parameters can effectively solve the optimal transport problem (in Wasserstein-2 with entropic regularization) for an arbitrary number of points. Consequently, the transformer can sort lists of arbitrary size up to an approximation factor. Our results rely on an engineered prompt that enables the transformer to implement gradient descent with adaptive step sizes on the dual optimal transport. Combining the convergence analysis of gradient descent with Sinkhorn dynamics, we establish an explicit approximation bound for optimal transport with transformers, which improves with increasing depth. Our findings provide novel insights into the essence of prompt engineering and depth for transformers.

Primary Area: learning theory

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 12353