FlashBias：高效计算带偏置的注意力机制
机器学习
2025-10-27 v3
摘要
带偏置注意力（attention with bias）通过在查询-键得分中引入先验知识作为可加偏置矩阵，扩展了标准注意力机制，广泛应用于视觉、语言、蛋白折叠等前沿科学模型中，成为这一基础模块的关键演进。然而，引入偏置项会导致注意力计算的严重效率瓶颈。它破坏了基于FlashAttention等加速器的高度紧密内存-计算流水线，从而削弱了大多数性能收益，使得带偏置注意力计算昂贵。令人惊讶的是，尽管该方法应用广泛，但针对带偏置注意力的针对性效率优化仍然缺失，这严重限制了其在复杂任务中的应用。深入分析FlashAttention的计算过程，我们证明其最优效率取决于注意力权重矩阵的秩。基于此理论结果，本文提出FlashBias，基于低秩压缩感知理论，为许多常用偏置提供快速精确计算，为一般形式的偏置提供快速准确近似。FlashBias充分利用现代GPU中高度优化的矩阵乘法运算，使AlphaFold 3中的 Pairformer 加速1.5倍，使视觉和语言模型中的带偏置注意力实现超过2倍的加速，且不损失精度。代码已开源于本仓库：https://github.com/thuml/FlashBias。
引用
@article{arxiv.2505.12044,
title = {FlashBias: Fast Computation of Attention with Bias},
author = {Haixu Wu and Minghao Guo and Yuezhou Ma and Yuanxu Sun and Jianmin Wang and Wojciech Matusik and Mingsheng Long},
journal= {arXiv preprint arXiv:2505.12044},
year = {2025}
}