Attention = Optimal Transport? Yes.

The attention module (Vaswani et al., 2017) is the computational engine behind modern LLMs, while entropic optimal transport (EOT) studies the efficient way to map one probability distribution to another. At first glance, the two areas seem unrelated—but inspired by Elon’s work (Litman, 2025), it is quite mind-blowing to discover that these two mechanisms are, in fact, mathematically equivalent.

Forward Attention Pass is Optimal Transport

In this post, we show that the optimal coupling $\mathrm{P\in \mathbb{R}^{n\times n}}$ in EOT with cost matrix $\mathrm{C = -QK^{\intercal}\in \mathbb{R}^{n\times n}}$ and entropy regularizer $\mathrm{\varepsilon=\sqrt{d}}$ is identical to the attention score $\mathrm{A\in \mathbb{R}^{n\times n}}$ in Transformers:

where $\mathrm{Q, K\in \mathbb{R}^{n\times d}}$ and $\mathrm{H(P)=\sum_{i,j=1}^n p_{ij} \log p_{ij}}$ is the discrete entropy.

\(\textbf{Proof}\) To handle both the simplex (column-sum) constraints for the RHS, we introduce Lagrange multipliers \(\lambda \in \mathbb{R}^n\). The Lagrangian is written as

where the first term represents the transport cost, the second term adds an entropy regularizer to ensure smoothness and positivity. The third term enforces the column-sum simplex constraint.

At the optimum, the derivative of \(\mathcal{L}\) with respect to \(\mathrm{P_{ij}}\) must vanish by the KKT condition:

which leads to the solution:

Since each column of \(\mathrm{P}\) must satisfy \(\mathrm{\sum_j P_{ij} = 1}\), we get

Hence, the optimal coupling with the cost matrix $\mathrm{C = -QK^{\intercal}}$ and the entropy regularizer \(\mathrm{\varepsilon=\sqrt{d}}\) is

Additionally, the Softmax operator in \eqref{EOT_attention} corresponds to the Shannon Entropy in EOT.

Locality-Biased EOT – Attention with Linear Biases (AliBi)

(Press et al., 2022) introduced the ALiBi embedding, which imposes a distance-based penalty on attention scores. This is equivalent to augmenting \(\mathrm{H(P)}\) in \eqref{EOT_attention} with an additional linear penalty.

PriorSoftmax: A Bayesian Extension

We can further generalize the entropic regularizer to Kullback-Leibler divergence to stay close to the prior $\mathrm{\pi}$. Such extension enables to encode inductive biases or location preference more smoothly.

We can also verify that \(\mathrm{L_{C, KL}=L_{C_{KL}}}\), where \(\mathrm{C_{KL}=C-\varepsilon \log \pi }\) in the EOT formulation \eqref{EOT_attention}.

Backward Attention Gradient is Policy Gradient

Assume $\mathcal{L}\mathrm{(c)}$ is the training loss of the attention module $\mathrm{Attention(c)=\sum_{j=i}^n Softmax(\frac{c}{\varepsilon})_j v_j}$, where $\mathrm{c\in \mathbb{R}^{n}, v\in \mathbb{R}^{d}}$. For convenience, we define the scalar $\mathrm{u_j=Softmax(\frac{c}{\varepsilon})_j}$.

By checking two cases $\mathrm{k=j}$ and $\mathrm{k\neq j}$ using quotient rule (Litman, 2025), we can verify that

We also denote $\mathrm{g=\dfrac{\partial \mathcal{L}(c)}{\partial u}}$ and a delta function $\mathrm{\delta_{ij}}$ at $\mathrm{i=j}$. Now the back-propagate gradient follows

The term $\mathrm{g_j - E[g]}$ denotes the advantage function of choosing token $j$, which quantifies the relative gains or losses compared to the mean utility $\mathrm{E_u[g]}$ induced by the current attention policy. This aligns with the REINFORCE policy-gradient formulation that incorporates a baseline. The baseline serves as a reference value to reduce gradient variance and make the training more stable.

Citation

@misc{deng2025ot_attention,
  title   ={{Attention = Optimal Transport? Yes.}},
  author  ={Wei Deng},
  journal ={waynedw.github.io},
  year    ={2025},
  howpublished = {\url{https://www.weideng.org/posts/attention_ot/}}
}