从最优传输角度解释注意力机制的数学原理。
Scaled-Dot-Product Attention as One-Sided Entropic Optimal Transport
- 将注意力前向传播视为一类退化的单边熵最优传输问题。
- 反向传播梯度等价于强化学习中的优势基策略梯度。
- 揭示了注意力分布空间的几何结构对学习更新的影响。
缩放点积注意力(SDPA)是现代深度学习的核心组件,但其数学形式常由启发式动机支持。本文首次从第一性原理出发,证明注意力前向传播是退化的一侧熵最优传输(EOT)问题的精确解,该问题旨在最大化相似性的同时保持最大熵。这一优化视角直接导致了反向传播的结果:通过反向传播计算的标准梯度在数学上等同于强化学习中的优势基策略梯度,一种方差降低的更新规则。关键地,我们证明了前向传播的EOT表述在注意力分布空间中诱导出特定的信息几何结构,该结构由费舍尔信息矩阵刻画,决定了学习梯度的精确形式,使优势基更新成为所求优化问题的自然结果。这一统一视角揭示了SDPA是一种有原则的机制:前向传播执行最优推理,反向传播实现理性且流形感知的学习更新。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。