SDPA统一原理：从最优传输到强化学习与信息几何的深度联系

大模型演进至今，注意力机制的架构也开始出现诸多变化。工业界的设计多基于直觉与实验，但当我们试图从第一性原理出发设计新机制时，面对诸如线性注意力是否真正有效等问题，需要一个坚实的理论基础。最近的一篇论文《Scaled-Dot-Product Attention as One-Sided Entropic Optimal Transport》为缩放点积注意力机制提供了首个基于数学第一性原理的论证，创造性地将其与最优传输理论、强化学习和信息几何联系起来。

文章的核心观点可以总结为三点：

1. 前向传播是最优传输：注意力机制通过Softmax生成权重的过程，等价于一个单边熵最优输运问题的精确解。
2. 反向传播是最优控制：通过标准反向传播计算出的梯度，在数学形式上等同于基于优势函数的策略梯度算法。
3. 内在联系是信息几何：前向的EOT问题定义了注意力分布空间上的信息几何，其曲率由费雪信息矩阵刻画，该矩阵同时决定了反向传播的梯度形式，从而将前向优化与反向学习统一。

一、引言：从经验解释到第一性原理

SDPA机制的核心是计算Query与一系列Key的点积相似度，再通过Softmax函数归一化为概率分布。这通常被解释为“点积是自然相似度度量”和“Softmax是可微分的argmax”。但这些是经验性、功能性的解释。论文旨在回答两个根本问题：

• 有效性：为何点积+Softmax如此有效？
• 本质：这一计算过程是否是某个有明确定义的数学问题的解？

论文的答案是肯定的，它证明SDPA并非启发式构造，而是一个具有完整理论依据的算法：其前向传播解决一个明确定义的变分问题，反向传播则执行一个复杂的、最优的学习更新。

二、前向传播：作为单边熵最优输运的解

最优传输理论概览

最优传输理论研究以最小成本将一种概率分布的质量转移到另一种分布。其发展历程中的重要节点构成了一个丰富的理论体系：

• Monge问题 (1781)：寻求从一个分布μ到另一个分布ν的确定性映射T，以最小化运输成本。其局限性在于确定性约束可能无解。
• Kantorovich问题 (1942)：通过引入“耦合”概念，允许质量拆分，将问题松弛为一个凸优化问题，从而与线性规划建立联系。
• Brenier定理 (1991)：在二次成本下，当μ绝对连续时，最优传输映射由某个凸函数的梯度给出，从而与梯度流和凸分析相联系。
• Sinkhorn算法 (2013)：引入熵正则化，使最优传输问题变得可高效求解，推动了其在机器学习中的广泛应用。
• Wasserstein距离：基于最优传输成本定义的概率分布间的距离度量，在GAN、文本相似度比较等领域应用广泛。

熵正则最优输运框架则是在经典目标函数中加入熵正则项，在最小化成本与最大化熵（分布平滑度）之间取得平衡。

SDPA对应的EOT问题

论文将SDPA的前向计算形式化为一个单边EOT问题。设有Query向量q，n个Key向量k_j，以及温度参数τ（通常为sqrt(d_k)），注意力分数定义为s_j = q^T k_j / τ。

定义（单边熵最优输运问题）：寻求一个概率分布p（即注意力权重），最小化如下目标泛函：

min_{p ∈ Δ_n} { -∑_j p_j * s_j + τ * ∑_j p_j log p_j }

其中Δ_n是概率单纯形。目标函数的第一项是最大化期望相似度（负号使得高相似度对应低成本），第二项是最大化分布的香农熵。

定理：当EOT的正则化参数ε设置为注意力温度参数τ时，上述优化问题存在唯一解，且该解正是缩放点积注意力权重：

p_j = exp(s_j) / (∑_i exp(s_i))

这为Softmax函数提供了坚实的数学基础：它不是任意的选择，而是特定优化问题的自然解。

标准EOT与SDPA的对比

SDPA对应的EOT问题与标准双边EOT（如Sinkhorn Attention所用）有本质区别：

特性	标准EOT	SDPA (One-sided EOT)
源测度	一般分布 μ	点质量 δ_q (单一Query)
目标测度	固定分布 ν	不受约束
传输计划	矩阵 γ ∈ ℝ^(n×m)	向量 p ∈ ℝ^n (注意力权重)
约束类型	双边约束 (行和、列和固定)	单边约束 (仅行和为1)

正是单边约束的宽松性，使得问题存在Softmax这一闭式解，而双边约束则需要Sinkhorn等迭代算法求解。

三、一个通用的注意力变分框架

基于上述EOT视角，可以提出一个通用的注意力设计框架。其核心是求解如下变分问题：

p = argmin_{p ∈ Δ_n} { -s^T p + Ω(p) }

其中Ω(p)是一个凸正则化项。不同的Ω(p)选择对应不同的注意力机制。

1. Softmax (Ω(p) = -τ H(p))

当正则化器为负的香农熵时，问题的解即为Softmax函数。

2. Sparsemax (Ω(p) = 1/2 ||p||₂²)

当正则化器为平方欧几里得范数时，问题的解是Sparsemax函数，其解析形式为 p_j = max(0, s_j - τ)，其中阈值τ使得概率和为1。这会产生稀疏的注意力分布。

3. α-entmax (Ω(p) = Tsallis正则化)

当使用Tsallis熵作为正则化器时，得到α-entmax族，它通过参数α控制稀疏程度，是Softmax (α→1) 和 Sparsemax (α=2) 的推广。其解的形式为 p_j = max(0, s_j - τ)^(1/(α-1))。

不同的正则化器选择决定了注意力分布的性质，总结如下：

α值	正则化器	变换名称	解的形式
α → 1	-τH(p)	Softmax	∝ exp(s_j/τ)
α = 2	½	p	₂²	Sparsemax	(sj − τ)+
α > 2	Tsallis正则化	α-entmax	(s_j − τ)^(1/(α−1))

4. ALiBi (Ω(p) = -τ H(p) + γ ∑_j p_j * |i - j|)

ALiBi机制在线性位置偏差的基础上进行Softmax，可以通过在标准熵正则化项上增加一个关于注意力距离的线性惩罚项来推导。这为ALiBi提供了一种基于最优传输理论的第一性原理解释，而非启发式设计。

5. PriorSoftmax (Ω(p) = τ KL(p || π))

KL散度正则化器将标准Softmax推广到贝叶斯框架。其解为 p_j ∝ π_j * exp(s_j/τ)，其中π是先验分布。这可以看作一个贝叶斯更新过程：先验π乘以似然exp(s_j)，再归一化得到后验p。

综合来看，不同的注意力机制对应于不同设计选择：

机制	正则化项 Ω(p)	结果/关键性质
Softmax	-τH(p)	通过香农熵最大化产生稠密、平滑分布
Sparsemax	1/2 ∑_j p_j^2	通过L2正则化产生具有稀疏零值的稀疏分布
α-entmax	1/(α(α-1)) ∑_j (p_j^α - p_j)	由α控制稀疏性，从Tsallis熵推导
ALiBi	-τH(p) + γ∑_j p_j |i - j|	通过线性惩罚引入结构化局部偏倚
PriorSoftmax	τ KL(p || π)	包含先验π的softmax，是贝叶斯推广

根据任务特征选择合适的机制：

• 需要考虑所有位置 -> Softmax (密集分布，无偏假设)
• 需要明确特征选择 -> Sparsemax (精确稀疏性)
• 可调节稀疏程度 -> α-entmax (参数化控制)
• 有序列结构 -> ALiBi (局部性偏差)
• 有先验知识 -> PriorSoftmax (贝叶斯整合)

四、反向传播：作为基于优势的策略梯度更新

论文进一步证明，SDPA的反向传播过程在数学上与强化学习中的策略梯度更新等价。

设通过注意力权重p加权Value向量得到上下文向量c = ∑_j p_j * v_j，下游任务的损失函数为L。定义第j个Key的边际效用为 a_j = -∂L/∂p_j，它衡量了增加对该Key的注意力权重所带来的“收益”。

定理：损失L关于注意力分数s_j的梯度为：

∂L/∂s_j = (1/τ) * p_j * (a_j - Ā)

其中 Ā = ∑_i p_i * a_i 是在当前注意力策略下的期望边际效用。

这个公式具有深刻的含义：(a_j - Ā) 正是强化学习中的优势函数，它度量了选择动作j（关注Key j）相对于平均表现的优劣。因此，标准的反向传播梯度更新，实质上是在执行一个带有基线（baseline，即Ā）的REINFORCE策略梯度算法：

• 正优势 (a_j > Ā)：鼓励增加该Key的分数。
• 负优势 (a_j < Ā)：鼓励减少该Key的分数。

这揭示了注意力机制的学习动态是一种理性的、方差缩减的最优控制过程。

五、统一桥梁：信息几何与Log-Sum-Exp势函数

前向是最优传输，反向是最优控制，这两者之间的联系并非巧合，而是由一个统一的几何结构——信息几何所决定。

关键发现：前向EOT问题的对偶势函数是Log-Sum-Exp函数：

V(s) = τ * log( ∑_j exp(s_j/τ) )

这个势函数的梯度恰好是注意力分布：∇V(s) = p。更重要的是，它的海塞矩阵（二阶导数）H = ∇²V(s) 直接正比于注意力分布p的费雪信息矩阵 F(p)：

H(s) = (1/τ) * F(p)

费雪信息矩阵定义了概率分布空间（统计流形）上的黎曼度量，它描述了该空间的局部曲率。在优化理论中，海塞矩阵描述了目标函数的曲率。

因此，论文中揭示的统一性可以概括为：同一个矩阵F(p)（或H(s)）扮演了三重角色：

1. 最优输运视角：它是EOT对偶势函数的曲率 (H)。
2. 信息几何视角：它是注意力分布流形的黎曼度量 (F)。
3. 强化学习视角：它是将边际效用向量a转换为标准梯度g的几何传感器 (g = (1/τ) F(p) * a)。

这个统一框架可以用下图关系表示：

前向传播 (EOT) ——(定义了对偶势函数 V(s))——> 势函数的曲率 H(s)
                                             ||
                                             ∝
                                             ||
反向传播 (RL) <——(几何传感器 F(p))—— 信息几何的度量 F(p)

标准梯度g与自然梯度g_nat在这个框架下形成了对偶关系：g_nat = F(p)^{-1} g ∝ a。这意味着标准反向传播隐式地利用了分布空间的几何结构，执行了一种流形感知的高效学习更新。

六、结论与展望

论文为缩放点积注意力机制建立了一个深刻的数学基础，将其从经验性构造提升为基于优化、控制和几何原理的理论体系。主要贡献包括：

1. 证明了SDPA前向传播是单边熵最优传输问题的唯一解。
2. 揭示了SDPA反向传播等价于带有基线减去的策略梯度更新。
3. 通过信息几何和对偶理论，将前向优化与反向学习统一在一个框架下，指出Log-Sum-Exp函数和费雪信息矩阵是核心的桥梁。

这种统一不仅加深了对现有注意力机制的理解，也为设计新机制提供了系统性的理论指导。未来的研究方向可能包括：将这一几何框架推广到多头注意力；探索整个Transformer架构的端到端几何性质；利用该理论分析大规模训练中的动力学现象；以及为新一代注意力机制（如MLA, MoR等）提供理论解释与设计原则。这一工作标志着对注意力机制的理解从经验走向原理，是深度学习理论化、几何化道路上的重要一步。

参考资料

[1] 大模型时代的数学基础(9)- SDPA和最优传输, 强化学习及信息几何的联系, 微信公众号：mp.weixin.qq.com/s/4PQI5pRVOpPHUvmzDLm2dg

版权声明：本文由 云栈社区 整理发布，版权归原作者所有。