Research map / Architectures and efficiency / Attention and transformers
FlashBias: Fast Computation of Attention with Bias
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes FlashBias, which factorizes attention bias matrices into low-rank factor tensors (exact, SVD, or neural decomposition) so FlashAttention-style kernels avoid reading the dense bias from HBM, yielding speedups for ALiBi, SwinV2, PDE solvers, and AlphaFold 3 Pairformer.
Problem
Attention with additive bias loses most of FlashAttention's IO efficiency because the dense bias matrix must be read from GPU HBM, creating a quadratic memory and runtime bottleneck.
Approach
Theoretically links FlashAttention IO cost to the rank of the attention/bias matrix, then replaces the dense bias with two low-rank factor tensors obtained by exact formulas (e.g., ALiBi, squared distance), SVD of learned parameters, or lightweight neural factor networks, and concatenates them onto queries and keys so standard FlashAttention kernels can be reused.
Key results
- On SwinV2-B ImageNet-1K, FlashBias reached 87.186% Acc@1 vs 87.144% official code, with inference time 0.190s vs 0.473s and memory 9429MB vs 12829MB.
- On AlphaFold 3 Pairformer with neural decomposition, FlashBias reduced runtime from 26.85s (open-source code) to 18.19s (about 32%) on PDB 7wux, with pTM 0.9498 vs 0.9500.
- For a 1.5B ALiBi GPT-2 style model, FlashBias cut bias-processing overhead in training from 5.0s to 2.3s relative to pure causal FlashAttention (per 100 iterations at N=2048) versus FlashAttention with Bias.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace dense bias reads with exact low-rank ALiBi decomposition | GPT-2 with ALiBi · GPT-2 ALiBi, training and inference | time (s): 124.3 -> 121.6; 40.32 -> 39.26 | Helped |
| Replace dense bias reads with exact low-rank ALiBi decomposition | GPT-2 with ALiBi · GPT-2 ALiBi, training and inference | time (s): 122.4 -> 121.5; 40.03 -> 39.78 | Helped |
| Use SVD-decomposed bias in the last 8 layers | SwinV2-B · ImageNet-1K | Time (s): 0.473 -> 0.190 | Helped |
| Use SVD-decomposed bias in the last 8 layers | SwinV2-B · ImageNet-1K | Acc@5: 98.232% -> 98.220% | No clear effect |
| Remove the attention bias (ablation) | SwinV2-B · ImageNet-1K | Acc@1: 87.144% -> 9.376% | Hurt |
| Use exact spatial-distance bias decomposition | Transformer PDE Solver · PDE solving, sequence length 32186 | C_D Error: 0.0173 -> 0.0113 | Helped |
| Approximate pair-representation bias with neural low-rank factors | AlphaFold 3 · PDB ID 7wux | Time (s): 26.85 -> 18.19 | Helped |
| Approximate pair-representation bias with neural low-rank factors | AlphaFold 3 · PDB ID 7wux | pTM: 0.9500 -> 0.9498 | No clear effect |
| Remove pair-representation bias (ablation) | AlphaFold 3 · PDB ID 7wux | pTM: 0.9500 -> 0.1713 | Hurt |
| Replace the bias matrix with decomposed factor functions | AlphaFold 3 · Training | Running Time (s / 10 iters): 165 -> 140 | Helped |
| Use SVD-decomposed relative-position bias in fine-scale layers | Pangu-Weather · ERA5, 100 samples | Time (s/100iters): 98.022 -> 76.779 | Helped |
| Apply FlashBias only to the last 8 layers rather than more layers (ablation) | SwinV2-B · ImageNet-1K | running time | Mixed |
Surprising findings
- FlexAttention degraded badly on SwinV2-B (2.885s vs 0.473s official code) because per-layer dynamic bias tensors force recompilation.
- Pure FlashAttention without bias drops SwinV2-B Acc@1 from 87% to 9%, showing the bias is essential even though it is costly to compute.
Limitations
- Neural and SVD decompositions are approximate and can introduce error, requiring per-layer rank selection and skipping layers where bias is not low-rank (only last 8 SwinV2 layers, first 16 Pairformer blocks).
- Speedups are modest for short sequences (e.g., Pangu-Weather N=144, about 20% runtime reduction) and the AlphaFold 3 training acceleration was not verified to full-model accuracy.
- Experiments are on a single A100 GPU, and the AlphaFold 3 model is a public reproduction (Proteinix) rather than the official model.
Experimental setup
Datasets: ImageNet-1K PDB weightedPDB_before2109_wopb_nometalc_0925 ERA5 Driving car simulation (Neural-Solver-Library)
Benchmarks: ImageNet-1K PDB test set recentPDB_1536_sample384_0925 ERA5 (Pangu-Weather output) Driving car simulation PDE benchmark
Metrics: Acc@1 Acc@5 running time GPU memory pLDDT loss pTM relative L2 error drag coefficient error output difference
Baselines: FlashAttention with Bias FlexAttention Pure FlashAttention Official Code INT8 PTQ FlashAttention w/o bias
Builds on: FlashAttention FlashAttention-2 FlexAttention ALiBi LoRA
Closest papers by meaning
- Delta Attention: Fast and Accurate Sparse Attention Inference by Delta Correction (NeurIPS 2025 · similarity 0.88)
- SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inference (ICML 2025 · similarity 0.87)
- FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (ICLR 2024 · similarity 0.86)
- Flash Invariant Point Attention (NeurIPS 2025 · similarity 0.86)
- SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Training (NeurIPS 2025 · similarity 0.86)
- FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (NeurIPS 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.