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

Questions:

Is your specific prompt engineering important for the numerics? Clearly it plays a central role in the proof.

Yes, please check added figure 4 in the revised document and [G.2] in the general response. The prompt engineering is indeed very important for solving optimal transport even for pretrained transformers. Thanks helping us to eleborate on this important aspect.

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).

Autmated engineering of prompt during training is the subject our follow-up in-progress research.

I think it would be helpful to comment on the choice of the adaptive step size in Theorem 1.

You are right: softmax limits the choice of stepsize. However, we prove there is a choice of stepsize that allows provable convergence in Theorem 3.

it does not seem as though the transformer layers are really performing gradient descent.

This is a variant of gradient descent with adapative coordinate-wise stepsizes. Adam and Adagrad are examples of gradient descent with coordinate-wise apative stepsizes.

Does this affect the convergence of the transformer output to the true OT matrix?

Yes. Theorem 3 proves this particular adaptive stepsize can still acheive a convergence rate. Notably, some part of engineered prompt are carefully chosen to obtain adaptive stepsizes that ensure the convergence.

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?

We do not claim that this specific stepsize choice comes with empirical benefits. However, the implementation of gradient descent was more challenging for softmax attentions due to the normalization in softmax. We added a paragraph in discussion to motivate future studies on this line of research (see Depth Efficient Guarantees in Discussions).

Weaknesses:

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.

Please check general response [G.1]

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.

We added figure 3 to eleborate on the role of training. Please check the general respone [G.3].

"Missing Generalization": using a specific cost function for optimal transport

This is not a limitation of our result, but it may be the limitation of transformers used in practice. Transformers have their own limitations and can not solve all problems. Focusing on $\ell_2$ cost for optimal transport is due the fact that we were only able to prove that transformers can solve optimal transprot for this specific cost function. Given $x_1$ and $y_1$, how can we consturct an attention layer that simply computes $\|x_1 - y_1\|_1$? It is not clear whether transformers can provably solve optimal transport with a cost function different from from norm 2. We clearly specify the cost in the abstract to avoid confusions about the choice of cost function for optimal transport.

To solve optimal transport with a differnt cost, we beleive the prompt must contain $C$ where $C_{ij} = cost(x_i,y_j)$ in instead of points $x_i$ and $y_i$. The main issue with this formulation is this formulation is beyond standard prompting. Consider example of sorting. We can give an unsorted list containing $x_1, \dots, x_n$ to transformers to sort. But, consturcting matrix $C_{ij} = cost(x_i,y_j)$ from $x_i$ and $y_i$ is not a standard prompting.

Limited Scope. Numerous studies [1], [2], [3] are bridging the gap between the capabilities of transformer models in ...

Please check the general response [G.1].

Mistakes in Equations

We appreciate your careful check of the proof. There is no mistake in our proof and there are only two typos.

at line 241 the definition of contains excessive vector $0$

This extra vector is needed for experiments in Figures 3 and 4 in the revised pdf. This is not a mistake but was a leveraged degree of freedom.

at line 241 and 242 vectors in both definitions of and must be moved more left to get additional term or in the equation at line 291;

There is no need for shift to left.

at line 275 the iteration index of has to be;

This typo is fixed.

at line 286 some of the zeros must be zero vectors or matrix;

We omit matrix/vector notions when it can be implied from the context.

at 291 the is missing and, additionally, the expression for the upper left block is wrong;

This is only a typo fixed in blue in the revised version. Notably, this typo does not change any following derviations.

at line 304 there must not bee coefficient, instead it should be moved to line 309.

Notably, matrix $D_\ell$ contains $\gamma_\ell$.

Difficulty in Understanding Matrix Derivations.

We will improve the presentation and apply your comments in the camera-ready version. Remarkably, matrix derivations are not the main components of our proof. Details on matrices dimensions can be concluded from the context. We beleive that all the reviewers were able to follow the proofs. For example, reivewers LUYE assess "The paper generally communicates complex ideas effectively, with clear explanations of both the background and the transformer mechanics used in OT tasks.". Furthermore, reviewer and YKZP confirm "The main results are supported by rigorous proofs...".

Questions

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.

One key insight that we bring is the possibility of out-of-distribution generalization with transformers. Suppose that we train a transformer to sort lists of size $7$. Is it possible to use the same transformer to sort lists of size $9$. Please check added figure 3 in the revised pdf to see this striking capability. This is an example of out-of-distribution generalization. In this paper, we take a step towards explaining this generalization, showing a transformer can sort lists of different sizes.

Interestingly, we use a specific prompt engineering for the proof. While prompt engineering is widely used for langauge models, it is not well understood how it can improve the performance of language models. Here, we prove that prompt engineering can enrich the computation power of transformers allowing them to solve fundamental optimal transport problem. Please check revised figure 4 showing the practical benefits of the proposed engineered prompt. This insight can be greatly leverage in future theoretical and practical studies.

consider bridging the gap between whether transformers can perform gradient descent on optimal transport tasks and whether they conduct it in real-world scenarios

Please check the general response [G.3].

consider different OT costs;

We question the possibility of solving OT with a general cost function with standard transformers. As explained above, transformers have their own limitations.

Prompt engineering is a significant aspect of the paper’s contribution ...

As mentioned in the general response [G.2], we have added the experimental results in Figure 4, which substantiate the important role of prompt engineering. In this experiment, we compare the performance of trained transformers with and without prompt engineering for solving optimal transport. We observed that the engineered part of the prompt (highlighted in red in Equation 6) significantly enhances the performance of transformers.

The paper's theoretical proofs are presented with a high level of mathematical density that may hinder accessibility for a broader audience

The choice of parameters/prompt is based on the subtle observation that gradient descent with adaptive coordinate-wise step sizes on $L(u,v)$ relies on the computation of the matrix $M_{ij} = e^{\left( \frac{-C_{ij} + u_i + v_j }{\lambda} - 1 \right)}$ as $\nabla_u L = 1_n M - \frac{1}{n} 1_n$. This computation can be performed by a single attention layer.

Recall that softmax attention consists of three main computations:

• (i) Computing inner products of tokens.
• (ii) Computing element-wise exponentials of inner products.
• (iii) Normalizing each row of the matrix.

Using (i) and (ii), along with including $u$ and $v$ in the engineered prompt, we construct $M$ in Equation 8.

In Equation 8, you will notice additional $1$ and $0$ elements after steps (i) and (ii). These elements allow us to derive the coordinate-wise step sizes after applying the normalization in softmax (iii). Notably, the step sizes are tuned so that we can achieve a convergence rate as shown in Theorem 2.

We use colors in the proof to highlight different components of the engineered prompt, which enable us to compute the gradient descent on $L$. We will include the above summary in the appendix or main text.

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

We demonstrate that prompt engineering can significantly boost the computational performance of transformers. This aligns with experimental observations, which show that adding sentences to the prompt often leads to improved responses. We added experiments in Figure 4 to further eleborate on the important role of prompt engineering. Our study highlights the need to revisit the role of prompt engineering in enhancing the computational power of transformers.

This study serves as a foundational step in a larger framework that establishes the computational power of language models. This framework connects transformers to optimization methods, thereby proving that transformers can solve a variety of problems, including regression and policy evaluation in Reinforcement Learning. While current literature predominantly relies on oversimplified models, such as linear attention, our results hold for softmax attention, which is widely used in practice. We show that linking transformers to optimization methods is effective, even for transformers deployed in real-world applications. We highlight prompt engineering can be leverage to boost the computational prower of transformers.

The experiments are relatively limited

We add two experiments in the revision: (i) experiments on trained transformers in figure 3 (ii) experiments showing the siginficant of the engineered prompt in figure 4. We will add experiments for $d>1$ in future revisions.

Questions.

I am interested in the insights derived from the prompt construction, and the paper would benefit from a more detailed discussion on this topic.

We have added Figure 4 to elaborate on the role of prompt engineering. We observed that training without prompt engineering results in poor performance on optimal transport. Furthermore, attention layers remain low rank after training. Thank you for helping us develop this insight.

Could the authors construct transformers with low-rank and sparse attention matrices, which can also effectively implement optimal transport

We beleive low-rank attention layer can not solve optimal transport since the solution of optimal transport is a full-rank permuation matrix.

Weaknesses.

The experiments in Section 6 only evaluate constructed transformers.

Please refer to the general response [G.1]. We have added experiments to further elaborate on (i) the mechanism of trained networks in Figure 3, and (ii) the significant role of prompt engineering in training:

• (i) Trained transformers can sort larger lists after being trained on smaller ones, similar to the constructive transformer (see [G.3] in the general response).
• (ii) Without the engineered prompt, transformers are unable to learn to sort in practice.

... 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 attention pattern remains sparse since the solution of optimal transport approximates a permutation matrix (see Figures 2 and 3).

One of our main contributions is proving that attention layers can maintain a high rank. We have added a new experiment in Figure 3 to demonstrate that even trained transformers maintain a high rank. Figure 4 presents a notable observation: without prompt engineering, the attention pattern does not maintain a high rank. This experiment suggests that specific prompt engineering, combined with the choice of parameters, enables attention layers to maintain a high rank across layers.

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

This limitation applies to shallow transformers, not to our theory. We observe that the performance of transformers improves with the addition of layers, as shown in Figure 3.

Increasing depth is an effective strategy to enhance the performance of language models. For example, GPT-2 XL, with 48 layers, achieves better performance compared to GPT-2, which has 12 layers. Providing insights into the importance of depth is one of our key contributions.

Related literature also confirms that deeper networks are computationally more powerful, including references [1-6] in the general response.

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.

We compared our results to the related literature in the general response [G.1]. What sets our results apart are:

• (i) Using attention layers with softmax rather than linear attentions.
• (ii) Demonstrating that prompt engineering can significantly boost the computational power of transformers for optimal transport (see Figure 4).
• (iii) Showing that transformers can solve problems of varying sizes.

Thank you so much for taking time and reviewing our response.

While it is evident that increasing network depth can improve performance in practive, I am concerned about the necessity of the required depth in the work. The construction relies on a large depth, yet I question whether such depth is necessary for this specific task.

Please check paragraph (i) Depth Efficient Guarantees in Discussions. We clearly explained that our convergence bound may not be tight and transformers may need less depth in practice. To our knowledge, there is no baseline or existing results that proves transformers achieve an approximation error with a certain number of layers. Thus, we can not compare with other bounds and assess whether our bound is tight. Yet, we establish the first bound and future studies can improve it.

The experiments primarily emphasize the necessity of prompt word engineering but fail to provide high-level insights into the underlying mechanisms.

The high level insight is that it may not be possible to unfold gradient descent iterations without prompt engineering. To implement gradient descent on $L$, we need to compute $C_{ij} = \|| x_i - y_j\||^2$ and store iterates $u_\ell$ and $v_\ell$. We engineered prompt to realize computing $C_{ij}$ using attention recurrence. Without prompt engineering, the transformer may not be expressive enough to implement gradient descent. Indeed, using prompt engineering to implement gradient descent on particular objective sets our results apart from other works that rely on implementing gradient descent with transformers.

Although the authors clarified the technical distinctions from related works, I still believe that the core idea of the proof closely follows these works: using multi-layer transformers to simulate multi-step gradient descent iterations. Consequently, the presented results lack sufficient novelty (,even if task-specific constructions are required).

There are many papers published and frequently cited that all rely on the idea of unfolding gradient descent with transformers layers, including references [1]-[8] in the general response. 4 references among these papers focus on the same task of linear regression. Still, these papers are published due to novelty and their distinguished settings and results. Using similar proof techniques is standard method in research and we believe it does not reduce the novelty. For example, statistical learning theory is based on proof techniques from empirical process theory, which itself heavy rely on concentration bounds in probability theory.

Again, we stress that we used softmax attention for a task completely different from [1]-[8]. Non of these papers are actively leveraging from prompt engineering.

It remains unclear whether the solution constructed by this work corresponds to the solution that standard transformers would converge to when trained on the optimal transport task.

We do not claim that we analyze trained transformers. We prove that transformers with prompt engineering can provably solve optimal transport on various input sizes.

Our results can provide insights on the mechanism of trained transformers. Consider the result in Figure 3 showing a transformer trained to solve sorting on lists with length 7 can sort lists of length 9. The question that we address is: How such striking generalization is even possible given that the distribution of inputs are completely different due to their length difference? Here, we prove this generalization is possible due to the implementation of gradient descent.

Notably, there are many different weight configurations that can implement the same algorithm (gradient descent with adaptive stepsizes). Thus, we can not compare the structure of weight matrices to conclude whether the algorithm implemented by trained transformers is gradient descent. However, we see that the process of solving optimal transport is also iteratively across the layers of transformers in Figure 3.

Thanks again for considering our response. Please let us know if we can provide further clarification.

This meta-review contains major misunderstandings that reflect a fundamental misinterpretation of both the results and the reviews:

• The paper does not claim that transformers implement Sinkhorn's algorithm, as incorrectly stated in the meta-review. Instead, it proves that transformers can implement gradient descent with adaptive step sizes on the dual Lagrangian function.

• Figure 3 clearly demonstrates that out-of-distribution generalization occurs in practice after training which is stated as the main weakness in the meta review.

• Although reviewer mAff raised their score from 3 to 6 after the rebuttal, the meta-review inaccurately claims that the response did not address the reviewer’s concerns.