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

## Abstract

Abstract Can we establish provable performance guarantees for transformers? Establishing 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 sizes up to an approximation factor. Our results rely on an engineered prompt that enables the transformer to implement gradient descent with adaptive stepsizes 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 as depth increases. Our findings provide novel insights into the essence of prompt engineering and depth for solving optimal transport. In particular, prompt engineering boosts the algorithmic expressivity of transformers, allowing them implement an optimization method.
With increasing depth , transformers can simulate several iterations of gradient descent.

## 1 Introduction

Language models with *theoretical guarantees* are reliable and, therefore, more practical. Extensive experiments confirm the striking capabilities of transformers, such as "multi-task learning" (Radford et al., ), "in-context learning" (Brown, ), generalization (Garg et al., ) to name but a few. But is it possible to theoretically ensure these capabilities and quantify their limits? Consider the simple example of sorting. We prompt the GPT-4 model to assess whether it can sort:

$$ $\text{prompt: sort}(2,1,4,3)\to\text{output: }(1,2,3,4)$ $$

While GPT-4 seems to be capable of sorting, prompting to verify sorting capability is computationally infeasible for two reasons: (i) the elements in the list can be arbitrary, and (ii) the list can be arbitrarily large. Thus, an infinite number of queries would be needed to verify that GPT-4 can sort. For theoretical verification, we need to study language models at a mechanistic level beyond black-box prompting. Here, we study how transformers can solve *optimal transport*.

Optimal transport is a fundamental optimization problem at the intersection of combinatorial and continuous optimization. Sorting is a special case of optimal transport in one dimension (Brockett, ). While efficient sorting algorithms are discrete, Brockett () raised the fundamental question of how to solve optimal transport with continuous dynamical systems. This question was motivated by the success of neural networks, which rely on continuous state dynamics across their layers for feature extraction. Brockett () proposes a continuous-state dynamical system over orthogonal matrices that iteratively solves optimal transport and, hence, can sort and diagonalize matrices. Building upon this study, we investigate whether the feature dynamics in transformers are capable of performing optimal transport.

Beyond statistical parametric models, language models are powerful computational machines. Recent studies show that language models are capable of learning to implement algorithms, including gradient descent for least squares (Ahn et al., ; Von Oswald et al., ; Akyürek et al., ; Gatmiry et al., ), non-parametric regression (Kim et al., ), Expectation-Maximization (EM) to solve mixture of linear models (Jin et al., ), the temporal difference method and other algorithms for reinforcement learning (Lin et al., ; Wang et al., ), and support vector mechanism (Tarzanagh et al., ). Surprisingly, simple components such as normalization layers can iteratively whiten data across the layers even at initialization (Meterez et al., ; Joudaki et al., ). This computational perspective has provided valuable insights into language models at a mechanistic level, linking them to algorithms with provable guarantees. Using this approach, we establish theoretical guarantees for solving optimal transport with transformers. Specifically, we prove that transformers can implement gradient descent (with adaptive step sizes) on the dual optimal transport objective regularized by entropy. In particular, each layer with two attention heads can simulate an iteration of gradient descent. Therefore, the induction of multiple attention heads can implement several iterations of gradient descent. This connection allows us to establish error bound for optimal transport with transformers that vanishes with *depth*: Given two sets of $n$ points in $\mathbb{R}^{d}$, a transformer can approximate the optimal transport map between the sets up to

$$ $\displaystyle O\left(\tfrac{\text{n}^{\nicefrac{{3}}{{2}}}}{\text{depth}^{ \nicefrac{{1}}{{2}}}}\right)\text{-error for all integer $n$}.$ (1) $$

Remarkably, the above bound holds for different choice of $n$ indicating that the transformer can solve multi-instances of optimal transport at the same time, an assertion for the capability of multi-task learning.
The established bound is achieved by prompt engineering.

Recent studies demonstrate that interacting with language models is an art: proper prompting can significantly enhance their performance. The seminal work of Wei et al. () shows that adding phrases such as "let’s think step by step" to the prompt encourages language models to produce more accurate reasoning. Prompting is becoming an essential skill in modern society, as prompt engineering positions are now being posted and well-paid by various companies. But what is the essence of prompt engineering? For the case study of optimal transport, we show that prompt engineering can significantly enhance the computational capabilities of transformers by providing the necessary memory and statistics.

## 2 Background

### 2.1 Optimal transport

Consider two sets of points ${x_{1},\dots,x_{n}\in\mathbb{R}^{d}}$ and ${y_{1},\dots,y_{n}\in\mathbb{R}^{d}}$, and define $C\in\mathbb{R}^{n\times n}$ such that $C_{ij}=\|x_{i}-y_{j}\|^{2}_{2}$. Finding the optimal transport map (in Wasserstein-2 metric) between these two sets of points casts to the optimization of a linear function over the set of permutation matrices (Brockett, ):

$$ $\displaystyle P_{*}:=\arg\min_{P\in\mathbb{R}^{n\times n}}\text{Tr}(PC),\quad \text{subject to }P\text{ being a \emph{permutation} matrix},$ (2) $$

where $\text{Tr}(M)$ denotes the trace of matrix $M$. Sorting lists is an example of optimal transport. Specifically, if $C_{ij}=(x_{i}-y_{j})^{2}$, where $y_{i}=i$ for $i=1,\dots,n$ and $x_{1},\dots,x_{n}\in\mathbb{R}$, then the linear transformation of $[x_{1},\dots,x_{n}]$ with $P_{*}$ sorts $x_{1},\dots,x_{n}$(see Remark 2.28 in (Peyré et al., ) for more details). Yet, the optimal transport problem is more general than sorting.

Optimal transport lies at the intersection of discrete and continuous optimization. There are various combinatorial algorithms for sorting and optimal transport, but our primary focus here is on continuous optimization methods. In particular, we review two fundamental methods: (i) constrained continuous optimization and (ii) Sinkhorn regularization.

While optimal transport involves optimization over the combinatorial set of permutation matrices, it can be relaxed to optimization over a continuous set. The state-of-the-art method is based on linear programming, specifically solving the following continuous convex optimization problem:

$$ $\displaystyle\widehat{P}=\arg\min_{P\in\mathbb{R}^{n\times n}}\text{tr}(PC), \quad\text{subject to }P\text{ being a \emph{doubly stochastic} matrix}.$ $$

Comparing the above problem with the original problem in (2), we notice that the constraint requiring $P$ to be a permutation matrix has been relaxed to allowing $P$ to be a doubly stochastic matrix. Recall the solutions to linear programs lie among the extreme points of the constraint set. Since the extreme points of doubly stochastic matrices are permutation matrices (Conte et al., ), the above linear program has the same solution as (2), i.e., $P_{*}=\widehat{P}$.

The above linear program requires the optimization of $O(n^{2})$ variables. Due to the quadratic growth with $n$, solving the linear program becomes computationally challenging for large $n$. Cuturi () proposes a computationally efficient alternative based on regularization with entropy:

$$ $\displaystyle P^{*}_{\lambda}:=\arg\min_{P\in\mathbb{R}^{n\times n}}\text{Tr}( PC)+\lambda\sum_{ij}P_{ij}\log(P_{ij}),\text{ subject to $P$ is a \emph{doubly stochastic} matrix}$ (3) $$

The Lagrangian dual of the above program reduces to the optimization of $O(n)$ variables which is considerably fewer than $O(n^{2})$ variables for the original linear program. Introducing the dual parameters $v\in\mathbb{R}^{n}$ and $u\in\mathbb{R}^{n}$, the Lagrangian function is defined as follows:

$$ $\displaystyle L(u,v,C)=\text{Tr}(PC)+\lambda\sum_{ij}P_{ij}\log(P_{ij})-u^{ \top}\left(P1_{n}-\frac{1}{n}1_{n}\right)-v^{\top}\left(P^{\top}1_{n}-\frac{1} {n}1_{n}\right).$ $$

It is easy to check that the minimizer of $L$ with respect to $P$ is $P_{ij}=e^{\frac{-C_{ij}+v_{j}+u_{i}}{\lambda}-1}$. This structure inspired the use of Sinkhorn’s fixed point iteration to find the solution of the dual problem. In particular, (Sinkhorn, ) proves that there exists a unique doubly stochastic matrix of the form $[P_{\lambda}^{*}]_{ij}=e^{\frac{-C_{ij}+v_{i}^{*}+u_{j}^{*}}{\lambda}-1}$ that is the solution of a simple fixed point iteration where $u^{*},v^{*}$ are unique up to scaling factors. Leveraging this fundamental theorem, Cuturi () proposes a fixed-point iteration that efficiently solves the dual problem. We will later elaborate on the fixed-point iteration and its convergence.

Apart from Sinkhorn’s fixed
point iteration, there are many different methods to solve Lagrangian dual problem such as first-order optimization methods. Recall the minimizer $P_{ij}=e^{\frac{-C_{ij}+u_{i}+v_{j}}{\lambda}-1}$. Plugging the minimizer into $L$ reduces the problem to the following optimization

$$ $\displaystyle\arg\min_{v,u\in\mathbb{R}^{n}}\left(L(u,v):=\lambda\left(\sum_{ ij}e^{(-C_{ij}+u_{i}+v_{j})/\lambda-1}\right)-\frac{1}{n}\sum v_{i}-\frac{1}{n }\sum_{i}u_{i}\right).$ $$

It is easy to check that $L$ is convex in $u$ and $v$ as its Hessian is diagonally dominant, hence positive semi-definite. Thus, standard first-order optimization can optimize $L$. In particular, one can use gradient descent with coordinate-wise adaptive stepsizes, such as the following recurrence

$$ $\displaystyle\begin{cases}u^{(\ell+1)}=u^{(\ell)}-D_{\ell}\nabla_{u}L(u^{(\ell )},v^{(\ell)})\\ v^{(\ell+1)}=v^{(\ell)}-D_{\ell}^{\prime}\nabla_{v}L(u^{(\ell)},v^{(\ell)}) \end{cases},$ (4) $$

where $\nabla_{u}L$ denotes the gradient of $L$ with respect to $u$ and $D_{\ell},D_{\ell}^{\prime}\in\mathbb{R}^{n\times n}$ are are diagonal matrices with positive diagonal elements, which are stepsizes for each coordinate. We will prove that self-attention layers can implement the above recurrence.

### 2.2 Self-attention layers

Attention layers are fundamental building blocks of neural networks, developed over decades of research. Hochreiter () pioneered this development by proposing an attention mechanism for Recurrent Neural Networks (RNNs) inspired by human cognition. Graves () employs the attention mechanism to develop a memory system for a parametric version of the Turing machine. Bahdanau () adapts this attention mechanism in neural Turing machines to design a powerful model for machine translation. While attention was originally introduced for recurrent models, Vaswani () introduced non-recurrent attention layers, combined with residual connections (He et al., ), thereby significantly enhancing the training of attention weights. This attention mechanism is the focus of our analysis.

Attention layers rely on on a convex combination. Let $Z\in\mathbb{R}^{m\times d}$. An attention layer is a function denoted by $\mathrm{atten}_{w}:\mathbb{R}^{m\times d}\to\mathbb{R}^{m\times d}$ with parameters $w:=\left[w_{k},w_{v},w_{q}\in\mathbb{R}^{d\times d}\right]$ defined as

$$ $\displaystyle\mathrm{atten}_{[w]}(Z)=A_{m\times m}Zw_{v},\quad A_{ij}=\tfrac{e ^{\langle w_{k}z_{i},w_{q}z_{j}\rangle}}{\sum_{j=1}^{m}e^{\langle w_{k}z_{i},w _{q}z_{j}\rangle}},$ (5) $$

where $z_{i}$ and $z_{j}$ are rows of $Z$.
The convex combination of data points imposes a local dependency that can simulate the focusing mechanism in neural networks.

Tay et al. () investigates whether attention layers are capable of sorting. Since self-attention layers cannot directly implement Sinkhorn’s fixed-point iteration, Tay et al. (); Sander et al. () propose a novel attention mechanism called "Sinkhorn attention". While it can not implement Sinkhorn’s recurrence, we demonstrate that standard attention layers can implement gradient descent with adaptive stepsizes on
$L$ with a specific prompt engineering.

## 3 Prompting and Model

#### The Engineered Prompt.

We propose a particular prompt structure to encode optimal transport in transformers:

$$ $\displaystyle Z_{n}=\setcounter{MaxMatrixCols}{11}\begin{bmatrix}{\color[rgb]{ 0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}x_{1}}&{\color[rgb]{0,0,1 }\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}y_{1}}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|x_{1}\|^{2}}&{\color[rgb]{ 1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|y_{1}\|^{2}}&{\color[ rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{ 1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}\\ {\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}x_{2}}&{ \color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}y_{2}}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|x_{2}\|^{2} }&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|y_{2}\|^ {2}}&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[ rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{ 1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}\\ {\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\vdots}&{ \color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}\\ {\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}x_{n}}&{ \color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}y_{n}}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|x_{n}\|^{2} }&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|y_{n}\|^ {2}}&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[ rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{ 1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}\\ {\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}0}&{\color[ rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}0}&{\color[rgb]{ 1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}-1/n}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}\end{bmatrix}\in\mathbb{R}^{( n+1)\times(2d+9)}.$ (6) $$

The highlighted elements in blue are the original prompts, which are sufficient for optimal transport. The elements highlighted in red are carefully engineered. We will prove that this particular prompt engineering allows attention layers to iteratively solve optimal transport.

#### Transformer.

We consider a specific transformer architecture composed of multiple attention and feedforward layers, all connected via residual connections. Let $Z^{(\ell)}$ denote the intermediate representation of the input $Z$ at layer $\ell$ which obeys the following recurrence

$$ $\displaystyle Z^{(0)}_{n}$ $\displaystyle=Z_{n},$ (7) $\displaystyle Z^{(\ell+1/2)}_{n}$ $\displaystyle=Z^{(\ell)}_{n}+\sum_{j=1}^{2}\mathrm{atten}_{w^{(\ell,j)}}(Z^{( \ell)}_{n})B_{j}^{(\ell)},$ $\displaystyle Z^{(\ell+1)}_{n}$ $\displaystyle=Z^{(\ell+1/2)}_{n}+(Z^{(\ell+1/2)}_{n}w_{f}^{(\ell)})_{+},$ $$

where $w_{f}^{(\ell)}\in\mathbb{R}^{d^{\prime}\times d^{\prime}}$ are the weight matrices for the feedforward layers, $B_{j}^{(\ell)}\in\mathbb{R}^{d^{\prime}\times d^{\prime}}$ are the mixing weights for the attention heads, and $(a)_{+}=\max(0,a)$ represents the ReLU activation function used in the feedforward layers. The model includes two attention heads and employs the standard softmax attention mechanism commonly used in practice.

## 4 Transformers as iterative algorithms

### 4.1 Adaptive gradient descent with transformer

The next theorem proves that transformers can implement iterations of gradient descent on $L$. The proof relies on the expressive power of attention layers combined with the engineered prompt, which provides the required memory to store iterates of gradient descent.

###### Theorem 1

There exists a configuration of parameters such that

$$ $\displaystyle\begin{cases}[Z^{(\ell)}_{n}]_{(1:n),(2d+7)}=u^{(\ell)}-D_{\ell} \nabla_{u}L(u^{(\ell)},v^{(\ell)})\\ [Z^{(\ell)}_{n}]_{(1:n),(2d+8)}=v^{(\ell)}-D_{\ell}^{\prime}\nabla_{v}L(u^{( \ell)},v^{(\ell)})\end{cases},$ $$

holds for all integer values of $n$, where $u^{(\ell)}$ and $v^{(\ell)}$ are gradient descent in  (4)
iterations starting from $u_{0}=v_{0}=0$ with the following adaptive stepsizes

$$ $\displaystyle[D_{\ell}]_{ii}=\frac{\gamma_{\ell}}{\sum_{j}e^{(-C_{ij}+u_{i}^{( \ell)}+v_{j}^{(\ell)})/\lambda-1}+1},\quad[D_{\ell}^{\prime}]_{jj}=\frac{ \gamma_{\ell}}{\sum_{i}e^{(-C_{ij}+u_{i}^{(\ell)}+v_{j}^{(\ell)})/\lambda-1}+1}.$ $$

Remarkably, the above result holds for an arbitrary
$n$ since transformers admit inputs with different number of tokens. We will elaborate on this important property by establishing convergence rate to $P^{*}_{\lambda}$, thereby proving that
a single transformer is capable of solving optimal transport for all $n$.

Notably, the above result supports the *"iterative inference hypothesis"* (Jastrzębski et al., ), that links the mechanism of deep neural networks to widely used optimization methods. This hypothesis postulates that residual connections enable deep networks to implicitly implement gradient descent across layers to tackle complex tasks. The hypothesis is based on striking observations on the underlying mechanisms of Convolutional Neural Networks (CNNs) (Alain, ). Previous studies have theoretically proven this hypothesis for solving least-squares using transformers (Ahn et al., ; Von Oswald et al., ; Akyürek et al., ). Building on these studies, we demonstrate that transformers can implement gradient descent for a different objective function to solve optimal transport, advancing our understanding of the iterative inference mechanism in deep networks

Prompt engineering is essential for the proof. Expanding the input size by adding columns and rows creates an extended data representation matrix across the layers. Attention layers can utilize a part of this expanded matrix as memory to store the iterates of gradient descent. Furthermore, the input dependent part of the prompt supplies the necessary statistics for the attention layers to implement gradient descent. To elaborate, we will present the proof and explain the essence of attention layers and prompt engineering.

###### Theorem 1

### 4.2 Proof of Theorem 1

The proof leverages the computational power of attention layers. We demonstrate that two attention heads can jointly implement a single step of gradient descent (with adaptive step sizes) on $L(u,v)$. By induction, multiple attention heads can implement several iterations of gradient descent with adaptive step sizes. The proof is constructive, explicitly determining the choice of parameters.

#### Parameters Choice.

Define $Q^{(\ell,j)}=w_{k}^{(\ell,j)}(w_{q}^{(\ell,j)})^{\top}$. Let $d^{\prime}=2d+9$ and $e_{i}\in\mathbb{R}^{d^{\prime}}$ denote the $i$-th standard basis vector $[e_{i}]_{j}=\begin{cases}1&i=j\\
0&\text{otherwise}\end{cases}$.
We choose parameters such that

$$ $\displaystyle\lambda Q^{(\ell,1)}$ $\displaystyle=\setcounter{MaxMatrixCols}{12}\begin{bmatrix}\textbf{0}_{d^{ \prime}\times d}&2e_{1},\dots,2e_{d}&\textbf{0}_{d^{\prime}}&-e_{2d+3}&-e_{2d+ 1}&e_{2d+7}&\textbf{0}_{d^{\prime}}&-\lambda e_{2d+5}&\textbf{0}_{d^{\prime}}& \textbf{0}_{d^{\prime}}&e_{2d+5}&\textbf{0}_{d^{\prime}}\end{bmatrix}\in \mathbb{R}^{d^{\prime}\times d^{\prime}}$ (8) $\displaystyle\lambda Q^{(\ell,2)}$ $\displaystyle=\setcounter{MaxMatrixCols}{11}\begin{bmatrix}2e_{d+1},\dots,2e_{ 2d}&\textbf{0}_{d^{\prime}\times d}&-e_{2d+3}&\textbf{0}_{d^{\prime}}&-e_{2d+2 }&e_{2d+8}&-\lambda e_{2d+5}&\textbf{0}_{d^{\prime}}&e_{2d+5}&\textbf{0}_{d^{ \prime}}&\textbf{0}_{d^{\prime}}\end{bmatrix}\in\mathbb{R}^{d^{\prime}\times d ^{\prime}}$ $\displaystyle[w_{v}^{(\ell,1)}]_{ij}$ $\displaystyle=\begin{cases}1&i=2d+6\;\text{and}\;j=2d+7\\ 0&\text{otherwise}\end{cases},\quad w_{f}^{(\ell)}=\textbf{0}_{d^{\prime} \times d^{\prime}}$ $\displaystyle[w_{v}^{(\ell,2)}]_{ij}$ $\displaystyle=\begin{cases}1&i=2d+6\;\text{and}\;j=2d+8\\ 0&\text{otherwise}\end{cases},\quad B_{j}^{(\ell)}=\gamma_{\ell}I_{d^{\prime} \times d^{\prime}}.$ $$

Notably, there are many choices for $w_{k}^{(\ell,j)}$ and $w_{q}^{(\ell,j)}$ that ensure the above equations hold.

#### Notations.

Consider the matrix $M\in\mathbb{R}^{n\times n}$ defined as $M_{ij}=e^{\frac{-C_{ij}+u_{i}^{(\ell)}+v_{j}^{(\ell)}}{\lambda}-1}$. By definition,

$$ $\displaystyle L(u^{(\ell)},v^{(\ell)})=\lambda\sum_{ij}\underbrace{\exp\left( \frac{-C_{ij}+u_{i}^{(\ell)}+v_{j}^{(\ell)}}{\lambda}-1\right)}_{=M_{ij}}- \frac{1}{n}\sum_{i}u_{i}-\frac{1}{n}\sum_{j}v_{j},\text{holds.}$ $$

#### A Proof Based on Induction.

Assuming that the statement holds for $\ell$, we prove that it holds for $\ell+1$. It is easy to check that induction base holds. The choice of $w_{f}^{\ell}$ and $w_{v}^{(\ell,j)}$ ensure that only the $2d+7$-th and $2d+8$-th columns of $Z^{(\ell)}_{n}$ change with $\ell$. Induction hypothesis concludes

$$ $\displaystyle Z_{n}^{(\ell)}=\setcounter{MaxMatrixCols}{11}\begin{bmatrix}{ \color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}x_{1}}&{ \color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}y_{1}}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|x_{1}\|^{2} }&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|y_{1}\|^ {2}}&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[ rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{ 1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{0,.5,.5} \definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}u^{(\ell)}_{1}}&{\color[rgb]{ 0,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}v^{(\ell)}_{1}}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}\\ {\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}x_{2}}&{ \color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}y_{2}}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|x_{2}\|^{2} }&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|y_{2}\|^ {2}}&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[ rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{ 1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{0,.5,.5} \definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}u^{(\ell)}_{2}}&{\color[rgb]{ 0,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}v^{(\ell)}_{2}}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}\\ {\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\vdots}&{ \color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{0,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}\vdots}&{ \color[rgb]{0,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}\\ {\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}x_{n}}&{ \color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}y_{n}}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|x_{n}\|^{2} }&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|y_{n}\|^ {2}}&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[ rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{ 1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{0,.5,.5} \definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}u^{(\ell)}_{n}}&{\color[rgb]{ 0,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}v^{(\ell)}_{n}}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}\\ {\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}0}&{\color[ rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}0}&{\color[rgb]{ 1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}-1/n}&{\color[rgb]{0,.5,.5} \definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}?}&{\color[rgb]{0,.5,.5} \definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}?}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}\end{bmatrix}\in\mathbb{R}^{( n+1)\times(2d+9)},$ $$

where elements highlighted in teal indicates the equality that follows from the induction assumption. Assuming the inequality above holds, we proceed to prove that

$$ $\displaystyle Z_{n}^{(\ell+1)}=\setcounter{MaxMatrixCols}{11}\begin{bmatrix}{ \color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}x_{1}}&{ \color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}y_{1}}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|x_{1}\|^{2} }&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|y_{1}\|^ {2}}&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[ rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{ 1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{0,.5,.5} \definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}u^{(\ell+1)}_{1}}&{\color[rgb ]{0,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}v^{(\ell+1)}_{1}}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}\\ {\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}x_{2}}&{ \color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}y_{2}}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|x_{2}\|^{2} }&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|y_{2}\|^ {2}}&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[ rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{ 1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{0,.5,.5} \definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}u^{(\ell+1)}_{2}}&{\color[rgb ]{0,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}v^{(\ell+1)}_{2}}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}\\ {\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\vdots}&{ \color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}&{ \color[rgb]{0,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}\vdots}&{ \color[rgb]{0,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}\vdots}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\vdots}\\ {\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}x_{n}}&{ \color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}y_{n}}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|x_{n}\|^{2} }&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\|y_{n}\|^ {2}}&{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[ rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{ 1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}1}&{\color[rgb]{0,.5,.5} \definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}u^{(\ell+1)}_{n}}&{\color[rgb ]{0,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}v^{(\ell+1)}_{n}}&{ \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}\\ {\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}0}&{\color[ rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}0}&{\color[rgb]{ 1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}-1/n}&{\color[rgb]{0,.5,.5} \definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}?}&{\color[rgb]{0,.5,.5} \definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}?}&{\color[rgb]{1,0,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}0}\end{bmatrix}\in\mathbb{R}^{( n+1)\times(2d+9)}.$ $$

Indeed, the extended prompt offers sufficient memory to store the vectors $u^{(\ell)}$ and $v^{(\ell)}$ obtained through gradient descent on $L(u,v)$.

#### Constructing Gradients With Attention Heads.

We begin by computing the output of the first attention head in layer $\ell+1$, step by step. Through straightforward algebra, we obtain the following:

$$ $\displaystyle Z^{(\ell)}_{n}Q^{(\ell,1)}=\setcounter{MaxMatrixCols}{12}\begin{ bmatrix}0&2x/\lambda&\textbf{0}_{n}&-\textbf{1}_{n}/\lambda&-\|x\|^{2}/\lambda &u^{(\ell)}/\lambda&0&-\textbf{1}_{n}&\textbf{0}_{n}&\textbf{0}_{n}&\textbf{1} _{n}/\lambda&0\\ 0&0&0&0&0&0&0&0&0&0&0\end{bmatrix}\in\mathbb{R}^{(n+1)\times d^{\prime}}$ $$

where $x=[x_{1},\dots,x_{n}],y=[y_{1},\dots,y_{n}]\in\mathbb{R}^{n\times d}$ and $x^{2}=[\|x_{1}\|^{2},\dots,\|x_{n}\|^{2}],y^{2}=[\|y_{1}\|^{2},\dots,\|y_{n}\|
^{2}]\in\mathbb{R}^{n}$. With these notations and the preceding equation, we proceed as follows:

$$ $\displaystyle Z^{(\ell)}_{n}Q^{(\ell,1)}(Z^{(\ell)}_{n})^{\top}$ $\displaystyle=\begin{bmatrix}\tfrac{1}{\lambda}\left(-x^{2}\textbf{1}_{n}^{ \top}+2xy^{\top}-\textbf{1}_{n}(y^{2})^{\top}+u^{(\ell)}\textbf{1}_{n}^{\top}+ \textbf{1}_{n}(v^{(\ell)})^{\top}\right)-\textbf{1}_{n}\textbf{1}_{n}^{\top}& \mathbf{0}_{n}\\ \mathbf{0}_{n}^{\top}&0\end{bmatrix}$ $\displaystyle=\begin{bmatrix}\log(M_{n\times n})&\mathbf{0}_{n}\\ \mathbf{0}_{n}^{\top}&0\end{bmatrix}$ $$

where $\log(M)$ for an input matrix $M$ is defined as a matrix with $[\log(M)]_{ij}=\log(M_{ij})$. Similarly, we define the matrix $\exp(M)$ such that $[\exp(M)]_{ij}=e^{M_{ij}}$. Thus,

$$ $\displaystyle\exp(Z^{(\ell)}_{n}Q^{(\ell,1)}(Z^{(\ell)}_{n})^{\top})=\begin{ bmatrix}M_{n\times n}&\textbf{1}_{n}\\ \textbf{1}_{n}^{\top}&1\end{bmatrix}$ (9) $$

Furthermore, the choice of parameters $w_{v}^{(\ell,1)}$ obtains

$$ $\displaystyle Z^{(\ell)}_{n}w_{v}^{(\ell,1)}=-\gamma\begin{bmatrix}\textbf{0}_ {n}&\dots&\textbf{0}_{n}&\textbf{1}_{n}&\textbf{0}_{n}&\textbf{0}_{n}\\ 0&\dots&0&-1/n&0&0\end{bmatrix}$ $$

Stitching all equations together yields

$$ $\displaystyle\mathrm{atten}_{w^{(\ell,1)}}(Z^{(\ell)}_{n})B_{1}^{(\ell)}= \begin{bmatrix}\textbf{0}_{n}&\dots&\textbf{0}_{n}&-D_{\ell}(M\textbf{1}_{n}- \frac{1}{n}\textbf{1}_{n})&\textbf{0}_{n}&\textbf{0}_{n}\\ 0&\vdots&0&n-1/n&0&0\end{bmatrix}$ $$

Similarly, we can show that

$$ $\displaystyle\mathrm{atten}_{w^{(\ell,2)}}(Z^{(\ell)}_{n})B_{2}^{(\ell)}= \begin{bmatrix}\textbf{0}_{n}&\dots&\textbf{0}_{n}&\textbf{0}_{n}&-D^{\prime}_ {\ell}(M^{\top}\textbf{1}_{n}-\frac{1}{n}\textbf{1}_{n})&\textbf{0}_{n}\\ 0&\vdots&0&0&n-1/n&0\end{bmatrix}$ $$

Replacing the last two equations into (7) concludes the induction proof.

## 5 Provable optimal transport with Transformers

By linking the intermediate data representation to a well-established algorithm, we gain access to powerful theoretical tools to prove that transformers can solve optimal transport. In particular, we utilize convergence analysis for gradient descent as well as Sinkhorn’s recurrence to prove transformers can find the optimal transport map.

#### Linking Attention Patterns to the Optimal Transport Map.

Recall the optimal transport map $P^{*}_{\lambda}$ defined in (3). We will prove that this matrix can be estimated by the attention matrices across the layers. Consider a block of *attention patterns* denoted by $A^{(\ell)}\in\mathbb{R}^{n\times n}$, defined as

$$ $\displaystyle A_{ij}^{(\ell)}=\frac{e^{\langle w_{k}^{(\ell,1)}z_{i}^{(\ell)}, w_{q}^{(\ell,1)}z_{j}^{(\ell)}\rangle}}{\sum_{j=1}^{n}e^{\langle w_{k}^{(\ell, 1)}z_{i}^{(\ell)},w_{q}^{(\ell,1)}z_{j}^{(\ell)}\rangle}},$ (10) $$

where $z_{i}^{(\ell)}$ is the $i$-th row of $Z^{(\ell)}$. We establish the convergence of $A^{(\ell)}$ to $P^{*}_{\lambda}$ in an appropriate metric.

#### Convergence Proof.

As discussed, the optimal transport matrix $P^{*}_{\lambda}$ has the following form (Cuturi, ):

$$ $\displaystyle P^{*}_{\lambda}=\text{diag}(w^{*})Q\text{diag}(q^{*}),\quad w^{* },q^{*}\in\mathbb{R}^{n}_{+},\quad Q\in\mathbb{R}^{n\times n}_{+},\quad Q_{ij} =e^{-\tfrac{C_{ij}}{\lambda}-1}.$ $$

It is easy to verify that replacing $w^{*}$ and $q^{*}$ with $cw^{*}$ and $q^{*}/c$ leads to the same matrix $P^{*}_{\lambda}$ for all $c\in\mathbb{R}_{+}$. Franklin and Lorenz () introduce a metric that accounts for this particular scaling invariance. Consider the following metric

$$ $\displaystyle\mu(w,w^{\prime})=\log\left(\max_{ij}\frac{w_{i}w^{\prime}_{j}}{w _{j}w^{\prime}_{i}}\right).$ (11) $$

Remarkably, $\mu$ is a metric that satisfies the triangle inequality (Franklin and Lorenz, ). However, $\mu$ is not a norm, as $\mu(w,w^{\prime})=0$ only implies that there exists a constant $c$ such that $w=cw^{\prime}$. The next theorem establishes an explicit convergence rate for the attention matrices $A^{(\ell)}$ to the solution of optimal transport in $\mu$.

###### Theorem 2

There exists a choice of parameters and an integer $k\leq\ell$ such that
$A^{(k)}$ can be expressed as:

$$ $\displaystyle A^{(k)}=\text{diag}(w^{(k)})Q\text{diag}(q^{(k)}),$ $$

where $w^{(k)},q^{(k)}\in\mathbb{R}^{n}_{+}$ obey:

$$ $\displaystyle\max\left\{\mu(q^{(k)},q^{*}),\mu(w^{(k)},w^{*})\right\}\leq\frac {36n^{\nicefrac{{3}}{{2}}}e^{\nicefrac{{r}}{{\lambda}}}\sqrt{r}}{\sqrt{\ell}(1 -\eta)},$ $$

for
$(\nicefrac{{1}}{{4}})r^{2}=\|w^{(1)}-w^{*}\|^{2}_{2}+\|q^{(1)}-q^{*}\|^{2}_{2}
,\eta=\nicefrac{{(\phi(Q)^{1/2}-1)}}{{(\phi(Q)^{1/2}+1)}},\phi(Q)=\max_{ijkl}
\nicefrac{{Q_{ik}Q_{jl}}}{{(Q_{jk}Q_{il})}},$
as long as $\ell\geq 64n^{3}\exp(\nicefrac{{3r}}{{\lambda}})r$.

The above theorem theoretically verifies that transformers can solve optimal transport for an arbitrary number of points $n$, with provable worst-case approximation guarantees. According to the theorem, the attention patterns converge to the optimal transport matrix $P^{*}_{\lambda}$ at a rate of $O\left(\frac{1}{\text{depth}^{\nicefrac{{1}}{{2}}}}\right)$, implying that the transformer performance improves with increasing depth.

An application of the last theorem is the provable sorting capability of transformers, achieved up to certain approximation factors. As discussed, sorting is a specific case of optimal transport for $d=1$ with $y_{1}\leq\dots\leq y_{n}$. By solving the optimal transport problem, transformers can effectively sort lists. Since convergence is guaranteed for regularized transportation, we expect the transformer to sort with an error that diminishes as $\lambda\to 0$. We will experimentally assess the sorting accuracy with the transformer in experiments.

The statement of the last theorem proves that it is possible to avoid the rank collapse of attention layers, which indicate that attention patterns converge to a certain low-rank matrix as depth increases (Geshkovski et al., ; Wu et al., ; Dong et al., ). Dong et al. () shows that attention layers without residual connection suffer from this collapse. Geshkovski et al. () proves even with residual connections attention with symmetric weights may suffer from the rank collapse. This rank collapse significantly reduces the expressivity of attention layers and poses challenges for training (Daneshmand et al., ). Consequently, a line of research investigates effective techniques to avoid the rank collapse issue (Meterez et al., ; Daneshmand et al., ; Joudaki et al., ). We argue that the last theorem demonstrates the possibility of avoiding the rank collapse of attention patterns with a specific prompt engineering and particular parameter choices. Recall that $P^{*}_{\lambda}$ serves as an approximation of the optimal transport matrix $P^{*}$, which is a full-rank permutation matrix. Consequently, we anticipate that the attention patterns will retain a high rank throughout the layers, as will be demonstrated in the experiments.

###### Theorem 2

### 5.1 Proof of Theorem 2

According to Thm. 1, a transformer can implement gradient descent. Therefore, the proof casts to analyzing gradient descent (with specific coordinate-wise stepsizes) on the convex $L$. However, we cannot directly apply existing convergence results from convex optimization. The existing convergence results for smooth convex optimization are in terms of function value $L$ when $L$ is not strongly convex(^1^11It is easy to check that $L$ is not strongly convex since it Hessian has a zero eigenvalue.). But, the theorem statement aims at the convergence to the optimal transport map.

We take inspiration from convex optimization proof. The proof consists of two key steps: (i) the convergence of attention patterns to a matrix that is approximately doubly stochastic, and (ii) a hypothetical simulation of Sinkhorn’s recurrence.

- (i)
As shown in Thm. 1, the transformer can perform gradient descent with an adaptive step size on the convex function $L$. Since $L$ is convex, gradient descent is guaranteed to converge to a stationary point where the gradient norm becomes zero. Specifically, it is straightforward to verify that $\nabla_{u}L=M\mathbf{1}-\tfrac{1}{n}\mathbf{1}$ and $\nabla_{v}L=M^{\top}\mathbf{1}-\tfrac{1}{n}\mathbf{1}$, where $M_{ij}=\exp\left(\tfrac{-C_{ij}+u_{i}+v_{j}}{\lambda}-1\right)$. Therefore, small gradients for $u$ and $v$ imply that $M$ is close to being doubly stochastic.
- (ii)
We demonstrate that when the matrix $M$ is approximately doubly stochastic, it is near the desired solution $P^{*}_{\lambda}$. To establish this, we (hypothetically) run Sinkhorn’s recurrence starting from $M$ and use its contraction property proven by Franklin and Lorenz ().

Before elaborating on the details of (i) and (ii), we present two propositions.

#### Preliminaries.

Define the functions $row:\mathbb{R}^{n\times n}_{+}\to\mathbb{R}^{n}_{+}$ and $col:\mathbb{R}^{n\times n}_{+}\to\mathbb{R}^{n}_{+}$ as

$$ $\displaystyle\mathrm{row}(A)_{i}=\frac{1}{n\sum_{j}A_{ij}},\quad\mathrm{col}(A )_{j}=\frac{1}{n\sum_{i}A_{ij}}.$ $$

We also introduce functions $f,g:\mathbb{R}^{n\times n}\to\mathbb{R}^{n\times n}$ defined as

$$ $\displaystyle f(A)=A\text{diag}(\mathrm{col}(A)),\quad g(A)=\text{diag}( \mathrm{row}(A))A.$ $$

Indeed, $f(A)$ (resp. $g(A)$) normalizes the columns (resp. rows) of $A$ by a scaling factor of their average. We will later use $f$ and $g$ to formulate Sinkhorn’s recurrence, which iteratively normalizes the rows and columns of a matrix with positive elements. The next proposition proves that an almost doubly stochastic matrix remains almost doubly stochastic under $f$ and $g$. To formulate the statement, we introduce a set containing matrices that almost doubly stochastic matrices:

$$ $\mathcal{S}_{\epsilon}:=\left\{A\in\mathbb{R}^{n\times n}_{+}\mid\|A\textbf{1} _{n}-\tfrac{1}{n}\textbf{1}_{n}\|_{\infty}\leq\epsilon\text{ and }\|A^{\top} \textbf{1}_{n}-\tfrac{1}{n}\textbf{1}_{n}\|_{\infty}\leq\epsilon\right\}.$ $$

###### Proposition 1

Suppose that $A\in\mathcal{S}_{\epsilon}$; then $f(A)\in\mathcal{S}_{3\epsilon}$ and $g(A)\in\mathcal{S}_{3\epsilon}$, as long as $\epsilon<\nicefrac{{1}}{{(3n)}}$.

Recall the metric $d$ defined in (11). The next proposition establishes a particular property of $f$ and $g$ with respect to $d$.

###### Proposition 2

Let $A\in\mathcal{S}_{\epsilon}$ be decomposed as $A=\text{diag}(w)Q\text{diag}(q)$, where $w,q\in\mathbb{R}_{+}^{n}$.

- (i)
For $f(A)=\text{diag}(w)Q\text{diag}(q^{\prime})$, $d(q,q^{\prime})\leq 4n\epsilon$ holds for $\epsilon<\tfrac{1}{4n}$.
- (ii)
For $g(A)=\text{diag}(w^{\prime})Q\text{diag}(q)$, $d(w,w^{\prime})\leq 4n\epsilon$ holds for $\epsilon<\tfrac{1}{4n}$.

###### Proposition 1

###### Proposition 2

#### (i) Convergence Analysis.

According to Theorem 1, there is a choice of parameters such that

$$ $A_{ij}^{(\ell)}=e^{\frac{-C_{ij}+u_{i}^{(\ell)}+v_{j}^{(\ell)}}{\lambda}-1},$ $$

where $u^{(\ell)}$ and $v^{(\ell)}$ are the iterates defined in (4). The following lemma establishes that, as the number of iterations $\ell$ increases, $A^{(\ell)}$ meets a neighborhood of doubly stochastic matrices.

###### Lemma 1

For $\gamma_{k}^{-1}=(n+2)e^{\nicefrac{{2r}}{{\lambda}}}$, there exists a $k\leq\ell$ such that

$$ $\displaystyle A^{(k)}\in\mathcal{S}_{\epsilon},\quad\text{ where }\epsilon^{2} :=\left(\frac{1}{\ell}\right)3ne^{\nicefrac{{3r}}{{\lambda}}}r.$ $$

Notably, the matrix $A^{(k)}$ has a specific structure that ensures $A^{(k)}\in\mathcal{S}_{\epsilon}$ is sufficient to approximate $P^{*}_{\lambda}$. To prove this statement, we leverage the contraction property of Sinkhorn’s recurrence.

###### Lemma 1

#### Contractive Sinkhorn’s Process.

According to the last lemma, there exists an iteration $k\leq\ell$ such that $A^{(k)}\in\mathcal{S}_{\epsilon}$. We then apply Sinkhorn’s recurrence starting from $A_{1}=g(A^{(k)})$ as:

$$ $A_{m+1/2}=f(A_{m}),\quad A_{m+1}=g(A_{m+1/2}).$ $$

Notably, we utilize the above recurrence solely for the proof; hence, there is no need for a transformer to implement this recurrence. According to the definition, $A_{m}$ can be decomposed as $\text{diag}(w_{m})Q\text{diag}(q_{m})$, where $Q_{ij}=e^{-C_{ij}/\lambda-1}$ and $q_{m},w_{m}\in\mathbb{R}^{n}_{+}$. Sinkhorn () proves that there exist unique vectors $w^{*},q^{*}\in\mathbb{R}^{n}_{+}$ such that $P^{*}_{\lambda}=\text{diag}(w^{*})Q\text{diag}(q^{*})$, where $w^{*}_{i}=e^{u^{*}_{i}/\lambda}$ and $q^{*}_{j}=e^{v^{*}_{j}/\lambda}$. (Franklin and Lorenz, ) establish the linear convergence of $(w_{m},q_{m})$ to $(w^{*},q^{*})$:

$$ $\displaystyle\begin{cases}\mu(w_{m+1},w^{*})\leq\eta\mu(w_{m},w^{*})\\ \mu(q_{m+1},q^{*})\leq\eta\mu(q_{m},q^{*})\end{cases},\quad\eta=\frac{\phi(A_{ 1})^{1/2}-1}{\phi(A_{1})^{1/2}+1}<1,$ (12) $$

where $\phi(A)=\max_{ijkl}\frac{A_{ik}A_{jl}}{A_{jk}A_{il}}$. Since $A_{1}=\text{diag}(w_{1})Q\text{diag}(q_{1})$, we have $\phi(A_{1})=\phi(Q)$.

#### (ii) Approximating the Optimal Solution.

Propositions 1 and 2 enable us to demonstrate that there exists a constant $c$ such that $cA_{1}$ lies within a neighborhood of $P^{*}_{\lambda}$.
Proposition 1 combined with Lemma 1 ensure $A_{1}\in\mathcal{S}_{3\epsilon}$. Thus, we can apply Proposition 2 to obtain:
$\mu(q_{2},q_{1})\leq 12n\epsilon.$
Using Proposition 1 again, we find that $A_{1+1/2}\in\mathcal{S}_{9\epsilon}$. Consequently, we can invoke Proposition 2 once more to yield:
$\mu(w_{2},w_{1})\leq 36n\epsilon.$ Applying the triangle inequality together with 12 completes the proof:

$$ $\displaystyle 36n\epsilon$ $\displaystyle\geq\mu(w_{2},w_{1})\geq\mu(w_{1},w^{*})-\mu(w_{2},w^{*})\geq(1- \eta)d(w_{1},w^{*})$ $\displaystyle 12n\epsilon$ $\displaystyle\geq\mu(q_{2},q_{1})\geq\mu(q_{1},q^{*})-\mu(q_{2},q^{*})\geq(1- \eta)\mu(q_{1},q^{*}).$ $$

## 6 Experiments

While our contribution is primarily theoretical, we connect the theoretical findings to practical observations through experiments. First, we experimentally validate our findings in Theorems 1 and 2, proving that transformers are capable of optimal transport. Second, we show that transformers can learn from data to solve optimal transport, combining the theoretical expressivity with data-driven learning used in practice. Finally, we elaborate on the significance of prompt engineering for optimal transport. Implementations are available at [https://github.com/hadidaneshmand/ot-transformer](https://github.com/hadidaneshmand/ot-transformer).

#### Data specification.

We consider optimal transport with $d=1$, varying $n$, which is the focus of our study. In particular, $x_{1},\dots,x_{n}$ are a random permutation of $[\nicefrac{{1}}{{n}},\nicefrac{{2}}{{n}},\dots,\nicefrac{{n}}{{n}}]$, and $y_{i}=\nicefrac{{i}}{{n}}$ in our experiments. We use a regularization constant of $\lambda=0.005$ in the related experiments.

#### Training protocol.

In Sections 6.2 and 6.3, we use Adam (Kingma, ) with a step size of $0.001$ and $10^{4}$ iterations for training. We reparameterize $Q_{\ell}:=w_{k}^{(\ell)}(w_{q}^{(\ell)})^{\top}$ and optimize $Q_{\ell}$. Parameters are initialized randomly from a Gaussian distribution with variance $\nicefrac{{1}}{{(2d+9)}}$. Notably, we set $w_{f}^{(\ell)}=0,B_{j}^{(\ell)}=(\nicefrac{{1}}{{20}})I_{d^{\prime}}$, and only optimize $w^{(\ell,j)}:=[w^{(\ell,j)}_{k},w^{(\ell,j)}_{q},w^{(\ell,j)}_{v}]$, as this does not limit the optimal transport capability (see (8)
). Training is conducted on the Colab platform using a T4 GPU.

### 6.1 Experimental validations

Since the proof of Theorem 1 is constructive, it provides an explicit choice of parameters for transformers. Specifically, we set $\gamma_{k}=0.01$ and use the weight matrices defined in (8). We examine the sorting accuracy of the transformer. As previously discussed, sorting is equivalent to optimal transport for $d=1$. Given the attention pattern $A^{(2000)}$, we compute the linear transformation of the input $x=[x_{1},\dots,x_{n}]$ using $A^{(2000)}$ to generate an estimate for the sorted $x$. An example of sorting is:

$$ $\displaystyle\text{Input}:[0.5,0.75,0.25,0.0]\quad\to\quad\text{Output}:[0.018 ,0.24,0.50,0.73]$ $$

We observe that generated output approximates the sorted list $[0.0,0.25,0.5,0.75]$. Interestingly, the same transformers can also approximately sorts larger lists:

$$ $\displaystyle[0.375,0.5,0.125,0.875,0.75,0.25,0.0,0.625]\quad\to\quad[0.02,0.1 2,0.25,0.37,0.5,0.62,0.75,0.84]$ $$

Figure: Figure 1: Convergence of Attention Patterns. The plotted matrices are $A^{(1)},A^{(300)}$ and $A^{(600)}$ defined in (10). We observe these matrices converge to $P^{*}_{\lambda}$ (the rightmost plot). Thm 2 proves this convergence.
Refer to caption: extracted/5972024/images/att0.png

Figure 1 illustrates the convergence of the attention matrices $A^{(\ell)}$ to $P^{*}_{\gamma}$, as established in Theorem 2. Notably, we observe that the attention patterns maintain high rank across in contrast to the observations on the rank collapse phenomenon; please refer to remarks in Section 5.

Figure 2 further illustrates that a single network can solve optimal transport on different sample sizes. In particular, this figure demonstrates that the transformer, with the parameter choices specified in (8), can find $P^{*}_{\gamma}$ for both $n=4$ and $n=8$ simultaneously without any changes to the parameters.

Figure: Figure 2: Optimal Transport of Different Sizes. left: $n=8$, right: $n=4$. The transformer weights remain exactly the same.
Refer to caption: extracted/5972024/images/att200008.png

### 6.2 Learning from data

So far, we proved and experimentally validate that a transformers is capable of optimal transportation. Now, we experimentally investigate whether a transformer can learn from data to solve optimal transport. Recall the hidden representation of transformer at layer $\ell$ denoted by $Z_{\ell}$. We optimize parameters $w^{(\ell,j)}$ to solve the following minimization problem

$$ $\displaystyle\min_{w^{(1,j)},\dots,w^{(20,j)}}\mathbb{E}\left[\|[Z_{20}]_{1:n, 2d+9}-\text{sorted}(x)\|^{2}\right]$ (training loss) $$

where sorted$(x)\in\mathbb{R}^{n}$ contains sorted $x_{i}$s and the expectation is taken over random data (details in Data specification). In order to approximate the expectation, we draw $500$ samples uniformly at random. Figure 3 compares the attention patterns —denoted by $A^{(\ell)}$ defined in (10)— across the layers where we observe that these patterns are converging to the optimal solution. This observation validates that transformers iteratively solve optimal transport across their layers (similar to gradient descent on $L$). While the the transformer is trained for $n=7$, we observe a good approximation for $n=8$ and $n=9$.

Figure: Figure 3: Optimal Transport After Training. *Rows:* $n=7,8$ and $9$; *Columns:* attention patterns $A^{(\ell)}$ defined in (10) for $\ell=5,10,15,20$. The last columns shows the optimal transport map, i.e., $P^{*}$ in (2) ; *Training:* optimizing the training loss on random data with $n=7$ (see Section 6.2 for more details).
Refer to caption: x1.png

### 6.3 The role of prompt engineering

We experimentally evaluate the impact of the engineered prompt (6) on solving optimal transport for $d=1$. Specifically, we reduce the number of columns in the prompt by removing additional ones as

$$ $\displaystyle Z^{\prime}=\begin{bmatrix}x_{1}&y_{1}&0\\ \vdots&\vdots&\vdots\\ x_{n}&y_{n}&0\end{bmatrix}\in\mathbb{R}^{n\times 3}.$ (13) $$

The last column is designated for the output. Let $Z^{\prime}_{20}$ represent the output of a transformer with 20 layers. We optimize the weights so that the last column of $Z^{\prime}_{20}$ predicts the sorted values of $x_{1},\dots,x_{n}$:

$$ $\displaystyle\min_{w^{(1,j)},\dots,w^{(20,j)}}\mathbb{E}\left[\|[Z_{20}^{ \prime}]_{:,3}-\text{sorted}(x)\|^{2}\right]$ (14) $$

where details on training data and process are presented in Data specification and Training protocol, respectively.

Figure 4 shows the clear impact of prompt engineering on the performance, where the above prompt (without engineering) leads to a significantly worse approximate of optimal transport matrix.

Figure: Figure 4: The Significance of Prompt Engineering. *Left:* attention pattern in the last layer denoted by $A^{(20)}$ after optimizing the training loss (14) for inputs without prompt engineering; *Middle:* attention pattern in the last layer ($A^{(20)}$) after training with the engineered prompt in (6); *Right:* optimal transport map associated with the input. $x_{1},\dots,x_{n}$ and $y_{1},\dots,y_{n}$ in prompts are fixed.
Refer to caption: x13.png

## 7 Discussions

We proved that transformers with fixed parameters can solve multiple instances of optimal transport on different number of points, with an explicit accuracy bound. Our analysis shows that transformers can implement gradient descent on a specific objective function using a specially engineered prompt. The engineered prompt provides additional memory to implement gradient descent. These findings open several avenues for future research, including: including: (i) depth-efficient guarantees, (ii) the analysis of training dynamics, and (iii) studying prompt engineering beyond optimal transport.

#### (i) Depth Efficient Guarantees.

According to Theorem 2, a transformer with $O(\epsilon^{-2})$-depth can obtain an $O(\epsilon)$-accurate solution. This is due to the established convergence rate for gradient descent with adaptive stepsizes. However, $O(\log(1/\epsilon))$ Sinkhorn iterations suffice for achieving $\epsilon$-accuracy. While there is a considerable gap between the established convergence analysis for gradient descent and the convergence rate of Sinkhorn’s iteration, our result is sufficient to prove deep transformers can provably solve optimal transport. We call for bridging this gap through a tighter convergence analysis.

#### (ii) Training for Optimal Transport.

We proved that transformers are able to provably solve optimal transportation and experimentally showed (in Section 6.2) that transformers can learn to solve optimal transport in $\mathbb{R}$ by training over random observations. Building upon this observation, we suggest a theoretical analysis of the training mechanism for optimal transport. To pursue this line of research, one can check whether parameters in (8) are local or global minimizers of the training loss. (Ahn et al., ; Wang et al., ) demonstrate that the properties of generative data distributions can be leveraged to analyze the stationary points of training dynamics. We believe that this technique can be used to analyze the landscape of training loss for optimal transport.

#### (iii) Prompt Engineering Beyond Optimal Transport.

Prompt engineering is essential for demonstrating that transformers are capable of solving optimal transport. In Section 6.3, we experimentally show that prompt engineering is also important in practice. Despite its widespread use, the underlying mechanisms of prompt engineering remain understudied. Studying prompt engineering for optimal transport is a step towards the broader goal of *understanding the role of prompt engineering in general*. We conjecture that prompt engineering enhances the computational power of transformers, enabling them to simulate a wider class of algorithms.

## Akhnowldgements

We thank Jacob Andreas for his helpful discussions and for inspiring the experiment in Section 6.3. We also thank Amir Reiszade for his insightful discussions on the related literature on in-context learning. We acknowledge
support from NSF TRIPODS program (award DMS-2022448).

## Appendix A Proof of Lemma 1

#### Notations.

Define the concatenated vector of iterates as

$$ $\theta_{k}=\begin{bmatrix}u^{(k)}\\ v^{(k)}\end{bmatrix},$ $$

and consider the following block diagonal matrix:

$$ $\Lambda_{k}=\begin{bmatrix}D_{k}&0\\ 0&D^{\prime}_{k}\end{bmatrix},$ $$

where $D_{k}$ and $D^{\prime}_{k}$ are diagonal matrices at iteration $k$ defined in Theorem 1:

$$ $\displaystyle[D_{\ell}]_{ii}=\frac{\gamma_{\ell}}{\sum_{j}e^{(-C_{ij}+u_{i}^{( \ell)}+v_{j}^{(\ell)})/\lambda-1}+1},\quad[D_{\ell}^{\prime}]_{jj}=\frac{ \gamma_{\ell}}{\sum_{i}e^{(-C_{ij}+u_{i}^{(\ell)}+v_{j}^{(\ell)})/\lambda-1}+1}.$ $$

Define the ball $B(r)=\{\theta\in\mathbb{R}^{n}\mid\|\theta\|\leq r\}$. If $\theta_{k}\in B(r)$, then

$$ $\displaystyle\frac{\gamma_{k}}{n\exp(r/\lambda)+1}I_{n}\preceq\Lambda_{k} \preceq\gamma_{k}I_{n}.$ (15) $$

Smoothness of $L$. The Hessian of $L$ has the following form

$$ $\displaystyle\nabla^{2}L:=\begin{bmatrix}\nabla^{2}_{uu}L&\nabla^{2}_{uv}L\\ \nabla^{2}_{vu}L&\nabla^{2}_{vv}L\end{bmatrix}=\begin{bmatrix}\text{diag}(\sum _{i}M_{ij})&M\\ &\\ M^{\top}&\text{diag}(\sum_{j}M_{ij})\end{bmatrix}$ (16) $$

We will prove that the Hessian bounded within the domain $\theta\in B(r)$.
For all $v:=\begin{bmatrix}s\in\mathbb{R}^{n}\\
s^{\prime}\in\mathbb{R}^{n}\end{bmatrix}$ such that $\|v\|^{2}=1$, we have

$$ $\displaystyle v^{\top}\nabla^{2}Lv=\|s\|^{2}_{\text{diag}(\sum_{i}M_{ij})}+2s^ {\top}Ms^{\prime}+\|s^{\prime}\|^{2}_{\text{diag}(\sum_{j}M_{ij})}$ (17) $$

where $\|v\|^{2}_{A}=v^{\top}Av$. Recall $M_{ij}=e^{\tfrac{-C_{ij}+u_{i}+v_{j}}{\lambda}-1}$, hence

$$ $\displaystyle\sum_{ij}s_{i}s^{\prime}_{j}M_{ij}$ $\displaystyle=\sum_{ij}s_{i}e^{\nicefrac{{u_{i}}}{{\lambda}}}s_{j}e^{\nicefrac {{v_{j}}}{{\lambda}}}e^{-\nicefrac{{C_{ij}}}{{\lambda}}-1}$ (18) $\displaystyle\leq\sum_{ij}s_{i}e^{\nicefrac{{u_{i}}}{{\lambda}}}s_{j}e^{ \nicefrac{{v_{j}}}{{\lambda}}}$ (19) $\displaystyle\leq\sqrt{\sum_{i}s_{i}^{2}e^{2\nicefrac{{u_{i}}}{{\lambda}}}} \sqrt{\sum_{i}(s^{\prime}_{i})^{2}e^{2\nicefrac{{v_{i}}}{{\lambda}}}}$ (20) $\displaystyle\leq e^{\nicefrac{{2r}}{{\lambda}}}$ (21) $$

It is easy to check that $\|\text{diag}(\sum_{i}M_{ij})\|$ and $\|\text{diag}(\sum_{j}M_{ij})\|$ are bounded by $ne^{\nicefrac{{r}}{{\lambda}}}g$. Replacing these inequalities into (17) yields

$$ $\displaystyle v^{\top}\nabla^{2}Lv\leq ne^{\nicefrac{{r}}{{\lambda}}}\left( \underbrace{\|s\|^{2}+\|s^{\prime}\|^{2}}_{=1}\right)+2e^{\nicefrac{{2r}}{{ \lambda}}}\leq(n+2)e^{\nicefrac{{2r}}{{\lambda}}}.$ (22) $$

Thus, $L(\theta)$ is $\zeta$-smooth for $\zeta:=(n+2)e^{\nicefrac{{2r}}{{\lambda}}}$ when $\theta\in B(r)$.

Boundedness of iterates. The recurrence relation of the iterates defined in (4) leads to the following inequality:

$$ $\displaystyle\|\theta_{k+1}-\theta^{*}\|^{2}_{\Lambda_{k}^{-1}}=\|\theta_{k}- \theta^{*}\|^{2}_{\Lambda_{k}^{-1}}-2\langle\theta_{k}-\theta^{*},\nabla L( \theta_{k})\rangle+\|\nabla L(\theta_{k})\|^{2}_{\Lambda_{k}},$ (23) $$

recall $\|v\|^{2}_{A}=v^{\top}Av$.
Since $L$ is $\zeta$-smooth within $B(r)$, by Theorem 2.1.5 of Nesterov [], we have

$$ $\displaystyle\langle\nabla L(\theta),\theta-\theta^{*}\rangle\geq\frac{1}{ \zeta}\|\nabla L(\theta)\|^{2}.$ (24) $$

Substituting the above inequality into (23) yields

$$ $\displaystyle\|\theta_{k+1}-\theta^{*}\|_{\Lambda_{k}^{-1}}^{2}\leq\|\theta_{k }-\theta^{*}\|_{\Lambda_{k}^{-1}}^{2}-\left(\frac{2}{\zeta}-\gamma_{k}\right) \|\nabla L(\theta_{k})\|^{2}.$ (25) $$

Let $\Delta_{k}:=\|\theta_{k}-\theta^{*}\|^{2}_{\Lambda_{k}^{-1}}$. For $\gamma_{k}=\frac{1}{\zeta}$, the above inequality ensures that $\Delta_{k}$ is monotonically decreasing:

$$ $\Delta_{k+1}\leq\Delta_{k}-\left(\frac{2}{\zeta}-\gamma_{k}\right)\|\nabla L( \theta_{k})\|^{2}\leq\Delta_{k}.$ $$

To maintain $\theta_{k}\in B(r)$ for all $k$, choose $r$ such that

$$ $\|\theta_{k}\|\leq\Delta_{k}+\|\theta^{*}\|_{\Lambda_{k}^{-1}}\leq\|\Delta_{1} \|_{\Lambda_{1}^{-1}}+\|\theta^{*}\|_{\Lambda_{1}^{-1}}\leq 2\left(\|\theta_{1 }-\theta^{*}\|+\|\theta^{*}\|\right)=r.$ $$

We now show that $\theta_{k}\in B(r)$ concludes the proof.

#### Convergence to a stationary point.

Since $\theta_{k}\in B(r)$, we can take the average of (25) over $k=1,\ldots,\ell$:

$$ $\displaystyle\sum_{k=1}^{\ell}\|\nabla L(\theta_{k})\|^{2}\leq\zeta\left(\sum_ {k=1}^{\ell}\Delta_{k}-\Delta_{\ell+1}\right)\leq\zeta\Delta_{1}\leq\zeta\left (ne^{\nicefrac{{r}}{{\lambda}}}+1\right)r.$ $$

The above inequality leads to the following bound on the minimum gradient norm:

$$ $\displaystyle\min_{k\leq\ell}\|\nabla L(\theta_{k})\|^{2}\leq\tfrac{1}{\ell} \sum_{k=1}^{\ell}\|\nabla L(\theta_{k})\|^{2}\leq\left(\frac{1}{\ell}\right) \zeta(ne^{\nicefrac{{r}}{{\lambda}}}+1)r.$ (26) $$

Closeness to Doubly Stochastic Matrices.
By definition,

$$ $\displaystyle\nabla L(\theta_{k})=\begin{bmatrix}M^{(k)}\mathbf{1}-\frac{1}{n} \mathbf{1}\\ (M^{(k)})^{\top}\mathbf{1}-\frac{1}{n}\mathbf{1}\end{bmatrix},$ (27) $$

where $\mathbf{1}$ denotes the vector of all ones.
Substituting the expression for $\nabla L(\theta_{k})$ into (26) gives

$$ $\displaystyle\min_{k\leq\ell}\left(\|M^{(k)}\mathbf{1}-\frac{1}{n}\mathbf{1}\| ^{2}+\|(M^{(k)})^{\top}\mathbf{1}-\frac{1}{n}\mathbf{1}\|^{2}\right)\leq\frac{ \zeta(n\exp(r/\lambda)+1)r}{\ell}.$ (28) $$

## Appendix B Proof of Proposition 1

We prove $f(A)\in\mathcal{S}_{3\epsilon}$ and the proof for $g(A)\in\mathcal{S}_{3\epsilon}$ follows exactly the same.
Since $A\in\mathcal{S}_{\epsilon}$, the following two inequalities hold

$$ $\displaystyle\left|\sum_{i}A_{ij}-\tfrac{1}{n}\right|\leq\epsilon\implies\sum_ {i}A_{ij}\geq\tfrac{1}{n}-\left|\frac{1}{n}-\sum_{i}A_{ij}\right|\geq\tfrac{1} {n}-\epsilon$ (29) $$

Using the above two inequalities, we proceed as

$$ $\displaystyle|\tfrac{A_{ij}}{n\sum_{i}A_{ij}}-A_{ij}|$ $\displaystyle=A_{ij}|1-\tfrac{1}{n\sum_{i}A_{ij}}|$ (30) $\displaystyle\leq\tfrac{A_{ij}\epsilon}{\sum_{i}A_{ij}}$ (31) $\displaystyle\leq\tfrac{A_{ij}\epsilon}{\frac{1}{n}-\epsilon}.$ (32) $$

We use the above inequality to complete the proof:

$$ $\displaystyle\left|\sum_{j}\tfrac{A_{ij}}{n\sum_{i}A_{ij}}-\frac{1}{n}\right|$ $\displaystyle\leq\left|\sum_{j}\tfrac{A_{ij}}{n\sum_{i}A_{ij}}-\sum_{j}A_{ij} \right|+|\sum_{j}A_{ij}-\tfrac{1}{n}|$ (33) $\displaystyle\leq\sum_{j}|\tfrac{A_{ij}}{n\sum_{i}A_{ij}}-A_{ij}|+\epsilon$ (34) $\displaystyle\leq\tfrac{\epsilon}{1/n-\epsilon}\sum_{j}A_{ij}+\epsilon$ (35) $\displaystyle\leq\epsilon\left(1+\tfrac{1/n+\epsilon}{1/n-\epsilon}\right)$ (36) $$

## Appendix C Proof of Proposition 2

We prove part (i), and the proof for part (ii) follows exactly the same. The following inequality holds for $A\in\mathcal{S}_{\epsilon}$:

$$ $\displaystyle\forall j:\left|\sum_{i}A_{ij}-\tfrac{1}{n}\right|\leq\epsilon.$ (37) $$

Using the above inequality, we get:

$$ $\displaystyle|q^{\prime}_{j}-q_{j}|$ $\displaystyle=\left|\tfrac{q_{j}}{n\sum_{i}A_{ij}}-q_{j}\right|$ (38) $\displaystyle=q_{j}|\frac{1}{n\sum_{i}A_{ij}}-1|$ (39) $\displaystyle\leq q_{j}\left(\tfrac{\epsilon}{\tfrac{1}{n}-\epsilon}\right)$ (40) $$

Plugging the above inequality into $\mu$ concludes the proof:

$$ $\displaystyle\frac{q_{i}q^{\prime}_{j}}{q_{j}q^{\prime}_{i}}\leq\frac{1}{1-2n \epsilon}\implies\mu(q,q^{\prime}_{i})\leq\log(\tfrac{1}{1-2n\epsilon})\leq \tfrac{2n\epsilon}{1-2n\epsilon}.$ (41) $$

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: - Could the authors construct transformers with low-rank and sparse attention matrices, which can also effectively implement optimal transport?
- I am interested in the insights derived from the prompt construction, and the paper would benefit from a more detailed discussion on this topic.
soundness: 3
strengths: This paper introduces a novel inquiry into whether transformers can effectively implement optimal transport. Additionally, the construction is independent of the number points, and can be used to solve optimal transport problems of different scales.
summary: This article constructs a prompt that enables transformers to perform gradient descent on dual optimal transport. Additionally, it provides an explicit approximation bound that improves with increasing depth. These results offer guarantees for the effectiveness of transformers in solving the optimal transport problem.
weaknesses: - The experiments in Section 6 only evaluate constructed transformers. However, it remains unclear whether training standard transformers can achieve the solutions constructed in the theory.

- Lines 375 to 384 in Section 5 and lines 479 to 483 in Section 6 imply that attention patterns maintain full rank. This assertion contradicts to the commonly observed phenomenon in practice, where attention matrices are highly low-rank and sparse. 

- The proof is based on the constructuion, so the required depth is only a sufficient condition. The required depth is presented as a sufficient theoretical condition but does not preclude the possibility that shallower transformers may also effectively perform optimal transport.

- Although this paper focuses on a novel task (optimal transport), the primary proof idea is similar to that of many previous works: using multi-layer transformers to simulate multi-step GD iterations.

---

confidence: 2
contribution: 3
initial_rating: 5
presentation: 2
questions: see weaknesses.
soundness: 3
strengths: 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.
summary: 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.
weaknesses: 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 technical choices and motivations, particularly around the parameter settings in Section 4, could further aid comprehension.

2. Prompt engineering is a significant aspect of the paper’s contribution, especially given its role in facilitating the transformer’s ability to approximate the optimal transport problem. However, the authors do not fully explore the sensitivity or robustness of their prompt design. For example, it is unclear if minor variations in the prompt structure or parameters would significantly impact the model's performance in solving OT, which is crucial for understanding the generalizability of this approach. Adding a sensitivity analysis that examines the impact of slight prompt modifications on performance would strengthen the prompt engineering component.

3. The experiments are relatively limited, focusing primarily on simple sorting tasks. For example, demonstrating performance on high-dimensional OT tasks, even with a modest dimensionality (e.g.,  𝑑=10), would provide valuable insights into the approach's scalability and robustness.  A head-to-head comparison with other OT algorithms (e.g., Sinkhorn or linear programming methods) on shared metrics (e.g., accuracy, runtime, or scalability) would provide a clearer picture of where the transformer-based approach stands.

4. This paper focuses on using transformers for OT.  How this framework could extend to broader optimization or machine learning problems?

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: To strengthen the paper, please consider:

- consider bridging the gap between whether transformers can perform gradient descent on optimal transport tasks and whether they
- conduct it in real-world scenarios;
- consider different OT  costs;
- adding more intermediate steps of matrix derivations in the Appendix;
- including dimensions directly in matrix definitions and accompanying explanations;
- extending notation section with more notations that are used in main sections;
- fixing mentioned mistakes;
- including a section that explicitly discusses the relevance of optimal transport theory to transformer models, explaining how this perspective can lead to new insights into the mechanics.
soundness: 4
strengths: The paper is rigorously constructed, demonstrating strong theoretical foundations, and supporting their claims of transformers being able to conduct gradient descent on optimal transport problems with explicit error bounds and convergence metrics. While there is a lack of available experiments specifically related to optimal transport, the authors find a way around this limitation by conducting experiments on sorting and optimal transport for different input sizes, which validate the theoretical assertions and strengthen the quality and reliability of the paper's contributions.
summary: The paper addresses the problem of theoretical guarantees of solving the entropic optimal transport (EOT) in the context of Wasserstein-2 distance using transformers. The authors provide theoretical assurance by demonstrating that transformers when equipped with specific prompts perform gradient descent on the dual formulation of the EOT problem.  The key contributions include proving that transformers with fixed parameters can implement gradient descent on the Wasserstein-2 EOT problem and establishing the error bound for this approach. The experimental results further demonstrate the ability of transformers with specific prompts to sort lists of numbers supporting the theoretical findings.

**Post rebuttal**: Score increase 3->6, see the discussion.
weaknesses: **Missing Generalization**. The authors establish a link between transformers and optimal transport, providing a specific and *very artificial* construction. While the choice of parameters, architecture, and prompts is justifiable within this framework, the construction remains limited, as there is no exploration beyond these specific Optimal Transport (OT) costs or an attempt to generalize the findings. Different applications and data types may require distinct cost functions to accurately reflect the underlying transport dynamics; for instance, a Manhattan cost better represents routing scenarios in transportation logistics, while a Gaussian cost can effectively model probabilistic asset movements in finance. By restricting their focus to specific OT costs, the authors miss potentially new application insights in the research field.

**Limited Scope**. Numerous studies [1], [2], [3] are bridging the gap between the capabilities of transformer models in artificial scenarios and their performance in real-world learning tasks, which is crucial for a comprehensive understanding of in-context learning. While the authors successfully establish a connection between transformers and optimal transport, they neglect to address whether these models are genuinely capable of solving optimal transport problems when trained in real-world scenarios. This omission limits the depth of their analysis and may hinder the reader's ability to fully grasp the implications of their findings within practical contexts. A discussion on this point would significantly enhance the paper’s contribution to the field.

[1] Von Oswald, Johannes, et al. "Transformers learn in-context by gradient descent." International Conference on Machine Learning. PMLR, 2023.

[2] Akyürek, Ekin, et al. "What learning algorithm is in-context learning? investigations with linear models." arXiv preprint arXiv:2211.15661 (2022).

[3] Wang, Jiuqi, et al. "Transformers Learn Temporal Difference Methods for In-Context Reinforcement Learning." arXiv preprint arXiv:2405.13861 (2024).

**Mistakes in Equations**. The equations contain several errors that make it difficult to follow the logical flow of the proofs. Some examples include: 
- at line 241 the definition of $\lambda Q^{(l, 1)}$ contains excessive $\mathbf{0}{d'}$ vector at the end; 
- at line 241 and 242 vectors $e{2d+5}$ in both definitions of $\lambda Q^{(l, 1)}$ and $\lambda Q^{(l, 2)}$  must be moved more left to get additional term $v^{(l)}$ or $u^{(l)}$ in the equation at line 291; 
- at line 275 the iteration index of $Z$ has to be $(l+1)$; 
- at line 286 some of the zeros must be zero vectors or matrix; 
- at 291 the $\lambda$ is missing and, additionally, the expression for the upper left block is wrong; 
- at line 304 there must not bee coefficient $\gamma$, instead it should be moved to line 309. 

These errors disrupt the coherence of the mathematical arguments and create unnecessary confusion. Carefully revising each equation to ensure consistency and accuracy would significantly enhance readability and strengthen the rigor of the presentation.

**Late or Missing Introduction of Notation or Definitions**. Notation is often introduced at various points of the paper, making it challenging for readers to locate definitions or fully understand terms when they first appear. In some cases, notation seems to be used without any prior introduction, leaving readers to infer meanings from context, which reduces the accessibility of the theoretical sections. It was mentioned:

- at line 202 the meaning of the subscript of $[Z^{(l)}_n]$ is not defined clearly;
- at line 299 it is not seen that this matrix is attention matrix $A$;
- at line 286 and 291 is used different notation of $[||x_1||^2, \dots, ||x_n||^2]$, $||x||^2$ and $x^2$;
there is no clear evidence of whether column or row vector notation is used;
- at line 237 authors introduce $d' = 2d + 9$ but they still use $2d + 9$ in lines 267, 275. 

*Suggestion:* A more effective approach would be to introduce all key notations systematically in a dedicated section at the beginning of the paper. Alternatively, providing a table summarizing all notations and definitions, with clear explanations and references to where they are first used, would help readers navigate the mathematical content more easily.

**Difficulty in Understanding Matrix Derivations.** In several parts of the paper, it is challenging to follow the derivation, for example, in lines 291 or 308. The transitions from definitions to their matrix forms are often abrupt, leaving readers uncertain about how certain matrices are constructed or why they take particular forms. In several instances like in lines 291, 304, 309, and 314, the dimensions of matrices are not specified, which is crucial for notation with blocks.

*Suggestion:* To improve clarity, consider adding more intermediate steps in the Appendix. Providing explicit dimensions for each matrix upon introduction would make it easier for readers to follow matrix operations. Including dimensions directly in matrix definitions or accompanying explanations would also help readers track the flow of derivations and verify the correctness of matrix manipulations.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. I think it would be helpful to comment on the choice of the adaptive step size in Theorem 1. My guess is that this is an artifact of using softmax attention, as it is difficult to express the gradients of the loss in terms of the softmax normalization. But it still seems strange that the step sizes depend on the parameters $u$ and $v$ which are being updated, as then it does not seem as though the transformer layers are really performing gradient descent. Does this affect the convergence of the transformer output to the true OT matrix? 
2. In your experiments, how does the prediction of the trained transformer compare to the solution obtained by running gradient descent on the dual OT loss?
3. Is your specific prompt engineering important for the numerics? Clearly it plays a central role in the proof. However, I would expect that you could first apply a learnable feature map to the original prompt and then apply attention to the featurized prompt and obtain similar results (i.e, make the prompt engineering learnable).
soundness: 3
strengths: 1. The paper provides guarantees for a broad class of discrete OT problems, which arise in a variety of applied fields.
2. The main results are supported by rigorous proofs, which are correct to the best of my knowledge.
summary: The paper proves the existence of a transformer which takes as an input a collection of points $(x_1,y_1, \dots, x_n,y_n),$ in $\mathbb{R}^d$ and outputs an approximate solution to the optimal transport matching problem between the sets of points. Their proof relies uses an explicit construction of a transformer which unrolls the gradient descent algorithm on the entropic-regularized OT problem for the dual potentials, and which takes a specifically-engineered prompt as an input. The convexity of the dual OT problem then implies a bound between the transformer output and the optimal coupling, which decreases as the depth of the transformer increases. Their result adds further evidence to the burgeoning interpretation of transformers as algorithm-learners.
weaknesses: 1. The idea of using transformers to unroll gradient descent algorithms is not new in the theory literature, it has been studied in [1], [2], and [3] extensively. In light of those results, it doesn't seem very surprising that you can also construct a transformer to unroll GD for the optimal transport loss.
2. The paper only bounds the approximation error of using transformers to solve OT problems and ignores other sources of error such as the pre-training generalization error and the training error of the transformer parameters.

Overall, I think the authors begin to tackle an interesting and important question in learning theory, but their analysis is not taken far enough.

[1] Bai, Yu, et al. "Transformers as statisticians: Provable in-context learning with in-context algorithm selection." Advances in neural information processing systems 36 (2024).
[2] Guo, Tianyu, et al. "How do transformers learn in-context beyond simple functions? a case study on learning with representations." arXiv preprint arXiv:2310.10616 (202
[3] Von Oswald, Johannes, et al. "Transformers learn in-context by gradient descent." International Conference on Machine Learning. PMLR, 2023.

Meta-review: In this paper, the authors show that there exists a particular prompt with a specific configuration of parameters for transformers that implement Sinkhorn's algorithm and thus approximatively solve regularized optimal transport problems. 

Having looked at the paper and the reviews, I see two main weaknesses in this work:

1. The experimental section could be improved: the authors do not address the question of whether the considered transformers are practically capable of solving optimal transport problems when trained in real-world data. At this point, it is not clear if the "out-of-distribution generalization" mentioned by the authors would occur in practice. Showing that transformers can practically learn the stinkhorn's algorithm would be a significant contribution (even on toyish data).

2. Lack of clarity and mistakes in equations: The paper still contains many typos and lacks some definitions after the revision. Actually, I am not sure the authors implemented what they claimed to in the rebuttal (e.g., the many comments of Reviewer mAff)

Although this paper's idea and result are interesting, I will recommend rejection for the above reasons. However, I believe the idea and the result have high potential, and I encourage the authors to address the abovementioned points in future submissions.

Decision: Reject