Scaled-Dot-Product Attention as One-Sided Entropic Optimal Transport

TL;DR

This paper mathematically justifies scaled-dot-product attention as the solution to a one-sided Entropic Optimal Transport problem, linking it to optimal inference and a variance-reduced reinforcement learning update.

Contribution

It provides a first-principles derivation of SDPA, connecting it to EOT and information geometry, and reveals the nature of its gradient as a manifold-aware update.

Findings

SDPA solves a degenerate EOT problem maximizing similarity with maximal entropy.

The backpropagation gradient matches an advantage-based policy gradient from reinforcement learning.

The EOT formulation induces a Fisher information geometry on attention distributions.

Abstract

The scaled-dot-product attention (SDPA) mechanism is a core component of modern deep learning, but its mathematical form is often motivated by heuristics. This work provides a first-principles justification for SDPA. We first show that the attention forward pass is the exact solution to a degenerate, one-sided Entropic Optimal Transport (EOT) problem, which seeks a distribution that maximizes similarity while being maximally entropic. This optimization perspective has a direct consequence for the backward pass. We prove that the standard gradient computed via backpropagation is mathematically identical to an advantage-based policy gradient, a variance-reduced update rule from reinforcement learning. Crucially, we demonstrate that the EOT formulation of the forward pass induces a specific information geometry on the space of attention distributions. It is this geometry, characterized by the Fisher Information Matrix, that dictates the precise form of the learning gradient, revealing the advantage-based update as a natural consequence of the optimization problem being solved. This unified view reveals SDPA as a principled mechanism where the forward pass performs optimal inference and the backward pass implements a rational, manifold-aware learning update.