[论文解读] SympFormer: Accelerated attention blocks via Inertial Dynamics on Density Manifolds

Viktor Stein, Wuchen Li+1 more|arXiv (Cornell University)|Mar 17, 2026

Generative Adversarial Networks and Image Synthesis被引用 0

一句话总结

论文提出基于密度流上的惯性（Nesterov 型）动力学的加速注意力块，用于 Transformers，推导出哈密顿动量注意力块并在保持 oracle 调用数量不变的前提下实现更快收敛。通过 Stein 与 Wasserstein 梯度流分析线性与 softmax 注意力，并提供可实现的粒子化离散化方法。

ABSTRACT

Transformers owe much of their empirical success in natural language processing to the self-attention blocks. Recent perspectives interpret attention blocks as interacting particle systems, whose mean-field limits correspond to gradient flows of interaction energy functionals on probability density spaces equipped with Wasserstein-$2$-type metrics. We extend this viewpoint by introducing accelerated attention blocks derived from inertial Nesterov-type dynamics on density spaces. In our proposed architecture, tokens carry both spatial (feature) and velocity variables. The time discretization and the approximation of accelerated density dynamics yield Hamiltonian momentum attention blocks, which constitute the proposed accelerated attention architectures. In particular, for linear self-attention, we show that the attention blocks approximate a Stein variational gradient flow, using a bilinear kernel, of a potential energy. In this setting, we prove that elliptically contoured probability distributions are preserved by the accelerated attention blocks. We present implementable particle-based algorithms and demonstrate that the proposed accelerated attention blocks converge faster than the classical attention blocks while preserving the number of oracle calls.

研究动机与目标

• 用变分、密度空间的视角来解释 Transformer 的注意力块。
• 引入惯性（Nesterov 型）动力学以开发加速注意力块。
• 推导哈密顿动量注意力块及其相应的离散化。
• 给出线性注意力在理论上的性质，包括 Stein 梯度流行为及椭圆轮廓分布的保持。
• 提供可实现的粒子化算法，并将收敛性与经典注意力块进行比较。

提出的方法

• 将 Transformer 的注意力块建模为在概率密度空间上用 Wasserstein-2 或 Stein 度量的梯度流离散化。
• 推导在密度流形上的二阶惯性（哈密顿）动力学，得到带动量的注意力块。
• 对线性自注意力进行特化，得到对二次势能的 Stein 梯度流，势能为 k(x,y)=y^T A x。
• 证明在加速线性注意力动力学下椭圆轮廓分布的保持。
• 将扩展到带非线性的移动性的 Stein- Wasserstein 型梯度流的 softmax 自注意力，并推导相应的加速动力学。
• 给出可实现的粒子化离散化，形成有限维的线性阻尼哈密顿粒子系统。

Figure 2: Validation loss (circles) and training loss (lines) on the tinystories data set after 10000 optimization steps, illustrating the results from Table ˜ 1 .

实验结果

研究问题

• RQ1 注意力块能否在密度流上具备惯性动力学以加速收敛？
• RQ2 加速注意力块与概率空间上的梯度流之间的数学关系是什么？
• RQ3 在线性注意力下，加速块是否保持结构族（如椭圆轮廓分布）？
• RQ4 如何将加速流离散化为适合实践的高效粒子化算法？
• RQ5 相较于经典对应物，加速块在线性与 softmax 注意力中的表现如何？

主要发现

• 加速线性注意力给出对二次能量的 Stein 变分梯度流，且具有时域二阶惯性系统。
• 在线性注意力下，加速动力学保持椭圆轮廓分布。
• 给出一个具体的粒子化离散化，产生线性注意力的有限维线性阻尼哈密顿系统。
• 对于 softmax 注意力，加速流被形式化为带非线性迁移率的 Wasserstein-2 型梯度流，并给出相应的粒子系统。
• 该框架提供可实现的算法，在保持相同的 oracle 调用次数的同时比经典注意力块收敛更快。