arXiv:2505.12044cs.LG2025-05NeurIPS被引 10

解决带偏置注意力计算慢的问题,实现高速精确计算。

FlashBias: Fast Computation of Attention with Bias

论文配图:FlashBias: Fast Computation of Attention with Bias
图 1 · 摘自论文原文
  • 基于低秩压缩感知理论,设计高效计算方法。
  • 在AlphaFold3中提速1.5倍,视觉语言模型超2倍无精度损失。
  • 适合需要高效偏置注意力的生物、视觉、语言模型应用。

带偏置的注意力通过在查询-键值得分中引入先验知识作为加性偏置矩阵,广泛应用于视觉、语言、蛋白质折叠等先进模型,成为该基础模块的关键演进。然而,偏置项引入导致注意力计算严重效率瓶颈,破坏加速器(如FlashAttention)中紧密融合的内存-计算流水线,使大部分性能优势丧失,使偏置注意力计算成本高昂。令人意外的是,尽管应用广泛,针对偏置注意力的专门效率优化仍缺失,严重阻碍其在复杂任务中的应用。深入分析FlashAttention计算机制后,我们证明其最优效率由注意力权重矩阵的秩决定。受此理论启发,本文提出FlashBias,基于低秩压缩感知理论,可对多种常用偏置实现快速精确计算,并对一般形式偏置提供快速准确近似。FlashBias能充分利用现代GPU中的高度优化矩阵乘法操作,在AlphaFold3的Pairformer中实现1.5×加速,在视觉与语言模型的偏置注意力中实现超过2×加速,且精度无损。代码已开源:https://github.com/thuml/FlashBias。

原文摘要 · Abstract (English)

Attention with bias, which extends standard attention by introducing prior knowledge as an additive bias matrix to the query-key scores, has been widely deployed in vision, language, protein-folding and other advanced scientific models, underscoring its status as a key evolution of this foundational module. However, introducing bias terms creates a severe efficiency bottleneck in attention computation. It disrupts the tightly fused memory-compute pipeline that underlies the speed of accelerators like FlashAttention, thereby stripping away most of their performance gains and leaving biased attention computationally expensive. Surprisingly, despite its common usage, targeted efficiency optimization for attention with bias remains absent, which seriously hinders its application in complex tasks. Diving into the computation of FlashAttention, we prove that its optimal efficiency is determined by the rank of the attention weight matrix. Inspired by this theoretical result, this paper presents FlashBias based on the low-rank compressed sensing theory, which can provide fast-exact computation for many widely used attention biases and a fast-accurate approximation for biases in general formalizations. FlashBias can fully take advantage of the extremely optimized matrix multiplication operation in modern GPUs, achieving 1.5$\times$ speedup for Pairformer in AlphaFold 3, and over 2$\times$ speedup for attention with bias in vision and language models without loss of accuracy. Code is available at this repository: https://github.com/thuml/FlashBias.

注意力机制高效计算AlphaFoldGPU优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。