Sinkformer：具有双随机注意力的 Transformer
机器学习
2022-01-25 v2 机器学习
摘要
基于注意力的模型（如 Transformer）涉及数据点之间的成对交互，由可学习的注意力矩阵建模。重要的是，该注意力矩阵用 SoftMax 算子归一化，使其成为行随机矩阵。在本文中，我们改为提出使用 Sinkhorn 算法使注意力矩阵成为双随机矩阵。我们将所得模型称为 Sinkformer。我们表明，经典 Transformer 中的行随机注意力矩阵随着训练轮数增加趋近于双随机矩阵，这证明了使用 Sinkhorn 归一化作为信息性先验的合理性。在理论方面，我们表明与 SoftMax 操作不同，该归一化使得将自注意力模块的迭代理解为 Wasserstein 度量的离散化梯度流成为可能。我们还表明在样本数趋于无穷的极限下，当同时重新缩放注意力矩阵与深度时，Sinkformer 执行热扩散。在实验方面，我们表明 Sinkformer 提升了视觉与自然语言处理任务中的模型精度。特别是在 3D 形状分类上，Sinkformer 带来显著改进。
引用
@article{arxiv.2110.11773,
title = {Sinkformers: Transformers with Doubly Stochastic Attention},
author = {Michael E. Sander and Pierre Ablin and Mathieu Blondel and Gabriel Peyré},
journal= {arXiv preprint arXiv:2110.11773},
year = {2022}
}
备注
Accepted at AISTATS