提出无需迭代的高效双随机注意力机制,提升模型信息流动
ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
- 用切片最优传输构造可并行的双随机注意力
- 在多类任务上显著提升性能,图像分类准确率最高增1.8%
- 适合需要高效注意力结构的深度学习应用
自注意力机制虽推动了Transformer的成功,但训练中易过度集中于少数标记,导致信息流不佳。强制注意力矩阵满足双随机约束已被证明能改善分布结构与平衡性。然而现有方法依赖迭代式Sinkhorn归一化,计算成本高。本文提出一种基于切片最优传输的新颖、完全可并行的双随机注意力机制,利用期望切片传输计划(ESP)。与以往方法不同,该方法无需迭代Sinkhorn归一化,显著提升效率。为保证可微性,引入基于温度的软排序技术,实现与深度学习模型的无缝集成。在图像分类、点云分类、情感分析及神经机器翻译等多个基准数据集上的实验表明,所提增强注意力正则化在多样化应用中均持续提升性能。代码已开源:https://github.com/dariansal/ESPFormer。
原文摘要 · Abstract (English)
While self-attention has been instrumental in the success of Transformers, it can lead to over-concentration on a few tokens during training, resulting in suboptimal information flow. Enforcing doubly-stochastic constraints in attention matrices has been shown to improve structure and balance in attention distributions. However, existing methods rely on iterative Sinkhorn normalization, which is computationally costly. In this paper, we introduce a novel, fully parallelizable doubly-stochastic attention mechanism based on sliced optimal transport, leveraging Expected Sliced Transport Plans (ESP). Unlike prior approaches, our method enforces doubly stochasticity without iterative Sinkhorn normalization, significantly enhancing efficiency. To ensure differentiability, we incorporate a temperature-based soft sorting technique, enabling seamless integration into deep learning models. Experiments across multiple benchmark datasets, including image classification, point cloud classification, sentiment analysis, and neural machine translation, demonstrate that our enhanced attention regularization consistently improves performance across diverse applications. Our implementation code can be found at https://github.com/dariansal/ESPFormer.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。