技术摘要：作为时间依赖 Wasserstein 梯度流的多头 Transformer 架构

问题陈述

尽管 Transformer 架构彻底改变了机器学习，但现有的理论分析往往依赖于与实际实现不符的理想化假设。具体而言，先前的平均场理论和 Wasserstein 梯度流公式通常假设：

1. 单头设置：忽略了多头注意力机制。
2. 时间无关性：假设权重矩阵在整个演化过程中保持固定（自治系统）。

在实践中，现代 Transformer 利用具有不同权重的多个头，且这些权重随层（深度）变化。这导致了一个真正的非自治向量场，其中有效的注意力交互随深度而变化。此外，最近的证据表明，学习到的时间依赖权重动态可能使自治近似失效，从而可能阻止在静态设置中经常推导出的令牌（token）聚类现象。严格的自治梯度流数学理论与实际应用中使用的非自治、多头架构之间存在关键差距。

方法论

作者通过将多头 Transformer 中的数据流建模为时间依赖梯度流来弥合这一差距，该梯度流针对显式捕捉注意力机制设计的交互能量。

1. 数学公式

• 从常微分方程 (ODE) 到偏微分方程 (PDE)：论文从一个离散的多头 Transformer 模型开始，其中令牌 $x_i$ 根据具有时间依赖、头依赖权重矩阵 $\theta^{(h)}_t$ 的 softmax 自注意力机制进行演化。
• 平均场极限：作者通过将令牌表示为球面 $S^{d-1}$ 上的时间依赖概率测度 $\mu_t$（由于层归一化），并将 Transformer 参数表示为参数空间上的时间依赖概率测度 $\Theta_t$，推导出了平均场公式。
• Transformer PDE：这导出了一个连续性方程（Transformer PDE），描述了 $\mu_t$ 的演化，其驱动力是参数分布 $\Theta_t$ 平均后的速度场。权重 $\Theta_t$ 本身通过随机微分方程 (SDE) 演化，允许随机初始化和噪声注入。

2. 加权 Wasserstein 几何

为了将系统作为梯度流进行分析，作者引入了加权 Wasserstein 距离 $W_{m,2}(\cdot, \cdot; \Theta)$：

• 非线性迁移率：该度量包含了由注意力机制和层归一化诱导的非线性、各向异性迁移率 $m_\mu(x, \theta)$。
• 有效迁移率：一个关键的技术贡献是推导出了有效标量迁移率 $\hat{b}_{\Theta_t}^\mu(x)$，定义为各个头迁移率的 $\Theta_t$ 加权调和平均。这使得复杂的多头几何结构可以简化为令牌空间上的单个标量因子，在保持 Wasserstein 型结构的同时保留了时间依赖参数的影响。
• 等价性：作者证明了该加权距离等价于经典的 2-Wasserstein 距离，确保了空间仍然是完备且紧致的。

3. 梯度流特征

• 最小化运动：利用最小化运动方案（Jordan-Kinderlehrer-Otto），作者构建了时间依赖交互能量 $E(\mu, t)$ 的解。
• 最大斜率曲线：他们确立了 Transformer PDE 的解是相对于加权 Wasserstein 度量的时间依赖能量的最大斜率曲线。
• 能量耗散恒等式 (EDI)：推导出了严格的 EDI，将能量变化与度量导数及强上梯度联系起来，并考虑了能量泛函的显式时间依赖性。

主要贡献与结果

1. 理论框架

• 广义平均场模型：论文提供了一个统一的框架，将单头和多头公式作为特例进行恢复，并允许参数随时间演化（非自治）。
• 梯度流等价性：证明了 Transformer PDE 的解恰好是交互能量的梯度流，反之亦然，梯度流满足 PDE。
• 稳定性与鲁棒性：
  • 输入扰动：作者证明了 Grönwall 型估计，表明对初始数据的连续依赖性，从而证明了模型对噪声输入的鲁棒性。
  • 权重扰动：他们建立了扰动交互能量向未扰动能量的 $\Gamma$-收敛性，证明了在权重初始化策略受到扰动时梯度流的收敛性。这为权重衰减等常见正则化技术提供了理论基础。

2. 渐近行为

论文分析了系统的长期行为，根据权重演化的可积性区分了两种机制：

• 可积机制（例如 Ornstein-Uhlenbeck）：如果权重动态满足适当的可积条件（$\int_0^\infty |\dot{\Theta}_t| dt < \infty$），则梯度流的 $\omega$-极限集非空。该集合中的每个元素都是与权重收敛子序列相关的极限能量的驻点。
• 非可积机制（例如振荡权重）：如果权重表现出持续的振荡（违反可积条件），系统不一定松弛到稳态。论文表明，令牌聚类并非普遍特征，而是取决于参数动态的稳定性。

3. 数值验证

通过在 $S^2$ 上的数值实验验证了理论预测：

• Ornstein-Uhlenbeck (OU) 机制：模拟证实令牌云松弛到与极限权重分布一致的稳态配置。强上梯度相对于头数 $H$ 以 $O(H^{-1})$ 的速率衰减。
• 振荡机制：在存在不衰减的振荡权重的情况下，令牌云无法聚类或稳定，且梯度不衰减。这在数值上证实了收敛性所需的可积假设的尖锐性。
• 能量平衡：沿模拟轨迹验证了离散能量耗散恒等式，确认了数值方案与连续理论的一致性。

意义与主张

论文声称通过以下方式显著缩小了理论平均场分析与实际 Transformer 架构之间的差距：

1. 超越自治性：为时间依赖和多头动态提供了严格框架，这些是深度学习固有的，但此前被排除在严格的梯度流分析之外。
2. 阐明聚类：证明了令牌聚类并非所有深度 Transformer 动态的固有属性，而是取决于权重演化的稳定性和可积性。
3. 几何洞察：引入有效迁移率（调和平均）作为一种机制，将多头复杂性压缩为令牌空间上可处理的几何结构。
4. 稳定性论证：通过稳定性估计，为 Transformer 对输入噪声的鲁棒性以及权重初始化/正则化策略的有效性提供了数学依据。

该工作确立了"Transformer PDE"是非自治设置中推理动态的有效、严格描述，使得能够利用强大的最优传输工具来分析和潜在地设计未来的神经架构。