ESPFormer：基于预期切片传输计划的双随机注意力
机器学习
2025-07-15 v2
摘要
虽然自注意力机制对Transformer的成功至关重要，但它可能导致训练期间注意力集中于少数标记，从而导致信息流次优。强制注意力矩阵为双随机约束已被证明可改善注意力分布的结构和平衡性。然而，现有方法依赖于迭代Sinkhorn归一化，这在计算上代价高昂。本文引入一种基于切片最优传输、利用预期切片传输计划（Expected Sliced Transport Plans, ESP）的全新、可并行化的双随机注意力机制。与先前方法不同，我们的方法无需迭代Sinkhorn归一化即可实现双随机性，显著提升了效率。为确保可微性，我们采用基于温度的软排序技术，实现了无缝集成到深度学习模型中。实验在多个基准数据集上进行，包括图像分类、点云分类、情感分析和神经机器翻译，结果表明我们的注意力正则化在各种应用领域 consistently改进性能。我们的实现代码可在https://github.com/dariansal/ESPFormer 查找。
关键词
引用
@article{arxiv.2502.07962,
title = {ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans},
author = {Ashkan Shahbazi and Elaheh Akbari and Darian Salehi and Xinran Liu and Navid Naderializadeh and Soheil Kolouri},
journal= {arXiv preprint arXiv:2502.07962},
year = {2025}
}