[论文解读] Sinkformers: Transformers with Doubly Stochastic Attention

Michael E. Sander, Pierre Ablin+2 more|arXiv (Cornell University)|Oct 22, 2021

Advanced Neural Network Applications参考文献 59被引用 14

一句话总结

本文提出Sinkformers，一种Transformer的变体，通过使用Sinkhorn算法替代标准自注意力中的SoftMax归一化，生成双随机注意力矩阵。该方法在视觉和自然语言处理任务中均提升了模型性能，尤其在3D形状分类任务中实现了显著的准确率提升，同时在无限样本极限下为注意力机制提供了关于Wasserstein梯度流和热扩散过程的理论见解。

ABSTRACT

Attention based models such as Transformers involve pairwise interactions between data points, modeled with a learnable attention matrix. Importantly, this attention matrix is normalized with the SoftMax operator, which makes it row-wise stochastic. In this paper, we propose instead to use Sinkhorn's algorithm to make attention matrices doubly stochastic. We call the resulting model a Sinkformer. We show that the row-wise stochastic attention matrices in classical Transformers get close to doubly stochastic matrices as the number of epochs increases, justifying the use of Sinkhorn normalization as an informative prior. On the theoretical side, we show that, unlike the SoftMax operation, this normalization makes it possible to understand the iterations of self-attention modules as a discretized gradient-flow for the Wasserstein metric. We also show in the infinite number of samples limit that, when rescaling both attention matrices and depth, Sinkformers operate a heat diffusion. On the experimental side, we show that Sinkformers enhance model accuracy in vision and natural language processing tasks. In particular, on 3D shapes classification, Sinkformers lead to a significant improvement.

研究动机与目标

• 通过使用Sinkhorn算法将自注意力中的行随机SoftMax替换为双随机注意力，以提升Transformer性能。
• 通过将Sinkformers解释为Wasserstein梯度流与热扩散过程，为自注意力机制提供理论基础。
• 通过实证验证，表明训练后的Transformer中注意力矩阵会自然收敛至双随机矩阵，从而证明使用Sinkhorn归一化作为有意义归纳偏置的合理性。
• 证明Sinkformers在多种任务中，包括3D形状分类、图像分类与序列建模，相比标准Transformer具有更高的准确率。
• 实现对现有Transformer架构中SoftMax的即插即用替换，仅带来极小的计算开销。

提出的方法

• 将自注意力中的SoftMax操作替换为Sinkhorn算法，将未归一化的注意力矩阵 $ K^0 = \exp(C) $ 归一化为双随机矩阵。
• 通过Sinkhorn算法的迭代行与列归一化，使结果矩阵的行和列均和为1。
• 引入一个超参数控制Sinkhorn迭代次数，实现从标准Transformer（1次迭代）到完全收敛Sinkformers之间的插值。
• 在深度学习框架中以可微分方式实现Sinkhorn算法，保持可微性与GPU兼容性。
• 利用双随机属性作为先验，促进所有token之间更均衡、更民主的注意力分配。
• 理论分析表明，Sinkformers在无穷小极限下对应于离散Wasserstein梯度流，并在平均场极限下收敛至热方程。

实验结果

研究问题

• RQ1 在优化过程中，训练后的Transformer中注意力矩阵是否会自然收敛至双随机矩阵？
• RQ2 将SoftMax替换为Sinkhorn归一化是否能提升视觉与自然语言处理基准测试中的模型准确率？
• RQ3 双随机注意力机制在连续极限下与最优传输及梯度流有何关系？
• RQ4 Sinkformers在无限深度与无限样本极限下的理论行为如何？
• RQ5 与标准Transformer相比，Sinkformers在实践中是否能以极小的计算开销高效实现？

主要发现

• 在ModelNet40 3D形状分类基准上，Sinkformers相比标准Transformer实现了显著的准确率提升，展现出在几何学习任务中的强大性能增益。
• 在IWSLT’14德语到英语翻译任务中，Sinkformer在35个训练周期后达到中位BLEU分数34.73，略高于标准Transformer的34.68，且训练协议相同。
• 在IMDb情感分类任务中，Sinkformers表现优于标准Transformer，表明其在序列建模中具有更好的泛化能力。
• 在视觉Transformer的猫狗图像分类任务中，Sinkformer的中位测试准确率为79.5%，高于标准Transformer的79.0%。
• 在MNIST上的一层、单头自注意力模型中，随着patch大小减小，Sinkformer的表现优于Transformer，表明其在低容量设置下具备更强的表征能力。
• Sinkformer每轮训练时间与标准Transformer几乎相同（3分25秒 vs. 3分20秒），证实其计算开销极低。