FlashBias: Fast Computation of Attention with Bias

TL;DR

FlashBias addresses the IO bottleneck in attention with bias by leveraging a low-rank decomposition of bias terms. Building on the insight that attention weight matrices have low effective rank, the method replaces dense bias matrices with two factor tensors via exact, SVD, or neural decompositions, dramatically reducing on-chip storage and memory traffic. Across NLP, vision, and scientific transformers (e.g., AlphaFold 3 and PDE solvers), FlashBias achieves up to 2×–>× speedups without sacrificing accuracy, and enables fast training in bias-rich models. The approach enriches the attention mechanism with bias without compromising the accelerator-friendly, matrix-multiplication–centric computation, unlocking scalable deployment of bias-enhanced transformers. The work includes a theoretical IO/rank analysis and extensive experiments with practical deployment details (e.g., Triton backends) and shows substantial real-world impact for AI-for-science and large-scale language and vision models.

Abstract

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.

Figures (12)

• Figure 1: (a-b) Comparison between attention mask and bias, where the spatial bias is from Swin Transformer liu2021Swin for computer vision, ALiBi bias is used in language modeling press2022train and pair bias is from AlphaFold abramson2024accurate. (c) FlashAttention needs to read bias tensors from HBM to the on-chip SRAM.
• Figure 2: Comparison between FlashAttention and FlashBias. FlashBias utilizes low-rank decomposition to bypass the read of the whole bias matrix, successfully avoiding the quadratic IO overload.
• Figure 3: Efficiency comparison among FlashBias, FlashAttention w/ Bias and FlexAttention dong2024flex. Here "Pure FlashAttention" refers to canonical FlashAttention without a bias term, which can be viewed as an upper bound of computation efficiency. Dotted lines indicate out-of-memory situations.
• Figure 4: Efficiency ratio over "Pure FlashAttention", which is calculated by $\frac{\text{Method Efficiency}}{\text{Pure FlashAttention Efficiency}}$.
• Figure 5: Speed of different implementations on attention with 128 hidden channels and 8 heads. $R$ is set as 8 in FlashBias. Dotted lines indicate out-of-memory. More attention speed comparisons are in Appendix \ref{['appdix:speed']}.

Theorems & Definitions (15)

• Theorem 3.1: FlashAttention computation benefits from low rank
• Theorem 3.2: Compressed sensing complexity of low-rank dense matrix candes2010power
• Corollary 3.3: A "lower bound" for HBM access of attention with bias
• Example 3.4: ALiBi press2022train in language models
• Example 3.5: Spatial distance in scientific problems
• Remark 3.6: Understanding neural decomposition
• Corollary 3.7: HBM access of FlashBias
• Remark 3.8: Trade-off between approximation accuracy and efficiency
• Example 3.9: Comparison with FlashAttention
• proof