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

缩放点积注意力（SDPA）机制是现代深度学习的核心组成部分，但其数学形式通常基于启发式方法进行解释。本研究从第一性原理出发，首次为SDPA提供了理论依据。我们首先表明，注意力机制的前向传播是退化的一侧熵正则化最优传输（EOT）问题的精确解，该问题旨在寻找一个分布，在最大化熵的同时也最大化相似性。这种优化视角对反向传播过程具有直接影响。我们证明，通过反向传播计算的标准梯度在数学上等同于一种基于优势函数的策略梯度——这是强化学习中一种方差缩减的更新规则。关键在于，我们展示了前向传播的EOT表述在注意力分布空间上诱导出特定的信息几何结构。这种由费舍尔信息矩阵刻画的几何结构，决定了学习梯度的具体形式，揭示了基于优势的更新规则其实是所求优化问题的自然结果。这一统一的视角表明，SDPA是一个原则性的机制：前向传播执行最优推断，而反向传播则实现了一种理性的、考虑流形结构的学习更新。

论文试图从第一性原理出发，解释Transformer中广泛使用的Scaled-Dot-Product Attention（SDPA）机制的数学合理性。传统上，SDPA的形式是基于启发式设计的，缺乏从优化或学习理论角度的系统性推导。

论文的核心思想是将SDPA的前向传播过程解释为一个退化的、单侧的熵正则化最优传输（Entropic Optimal Transport, EOT）问题的解析解。通过这一视角，SDPA不仅被解释为在最大化相似性的同时最大化熵的最优分布，还揭示了其反向传播梯度与强化学习中基于优势的策略梯度之间的等价性。

1. 首次将注意力机制与最优传输理论建立直接联系，为SDPA提供了理论基础。 2. 揭示了标准注意力梯度等价于一种方差缩减的策略梯度更新，提供了对反向传播的新理解。 3. 通过EOT推导出注意力分布上的信息几何结构（由Fisher信息矩阵刻画），进一步解释了学习动态的本质。 4. 为注意力机制提供了统一视角：前向是推理，反向是符合几何结构的理性更新。 5. 该理论框架为设计新型注意力机制提供了潜在方向，例如引入不同的熵正则项或几何结构。

1. Vaswani et al., "Attention Is All You Need", NeurIPS 2017. 2. Cuturi, "Sinkhorn Distances: Lightspeed Computation of Optimal Transport", NeurIPS 2013. 3. Peyré & Cuturi, "Computational Optimal Transport", Foundations and Trends in Machine Learning, 2019. 4. Zhang et al., "Natural Policy Gradient and Optimality Conditions for Bounded-Rational Decision-Making", AISTATS 2020. 5. Amari, "Information Geometry and Its Applications", Springer, 2016. 6. 最近也有研究尝试将最优传输与神经架构搜索、图神经网络、生成模型结合，如"OT-GAN: A Generative Network based on Optimal Transport" (ICML 2019)。