FlashBias: Fast Computation of Attention with Bias
Authors: Haixu Wu
Minghao Guo
Yuezhou Ma
Yuanxu Sun
Jianmin Wang
Wojciech Matusik
Mingsheng Long
Venue: NeurIPS 2025 poster ()
Keywords: FlashAttention
Efficient Attention
Primary area: deep_learning
TL;DR: This paper presents FlashBias to speed up computation of attention with bias, which brings 1.5x speedup for AlphaFold and 2x speedup for SwinV2.

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.

FlashBias: Fast Computation of Attention with Bias
Haixu Wu', Minghao Guo?, Yuezhou Ma', Yuanxu Sun', , Jianmin Wang',
Wojeiech Matusik Mingsheng Long 1
School of Software, Tsinghua University, MIT CSAIL
(suhaixu98, guwmhh2014)@maaaiiia com, (mayz20, sunyuanx22))mmiils waingguuinnnaaaaaea edu. cn, {j i (es.nnghennngttinnniiii edu. en, <vojciech)@caail. edt.edu
Abstract
Attention with bias, which extends standard attention by introducing prior knowl- edge as an additive bias matrix to the query-key scores, has been widely deployed in vision, language, protein-folding and other advanced scientific models, underscor- ing its status as a key evolution of this foundational module. However, introducing bias terms creates severe efficiency bottleneck in attention computation. It disrupts the tightly fused memory-compe pipeline that underlies the speed of accelcrators like FlashAttenton,, thereby stripping away most of their performance gains and leaving biased attention computationally expensive. Surprisingly, despite its com- mon usage, targeted efficiency optimization for attention with bias remains absent, which seriously hinders its application in complex tasks. Diving into the computa- tion 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 fast-accurate approximation for biases in general formalizations. FlashBias can fully take advan tage of the extremely optimized matrix multiplication operation in modern GPUs, achieving 1.5x 1.5x speedup for Pairformer in AlphaFold 3, and over 2 speedup for attention with bias in vision and language models without loss of accuracy. Code is available at this repository: https: //tithubt com/coo/////o/ooo coom/thumm/Faaa /hmm///mmmFaaaia
Introduction
40 has enabled impressive achievements in extensive areas, including computer vision [25] 45) natural language processing 2 39. scientific applications [[ 42). etc. Especially, the core design of Transformer, the attention mechanism with powerful relation model- ing capacity has emerged as a foundation module deep learning models, making its optimization of vital importance. As seminal progress, FtashAttention [10 9 speeds up the computation of standard attention by successfully reducing the read-write cost of GPU high bandwidth memory (HBM) with 10-aware techniques. Afterwards, researchers further extend FlashAttention to support diverse kinds of masks 35 41, such as causal mask in decoder-only Transformers or sliding window mask.
In recent years, Transformer
While standard attention attention with masks has enjoyed elaborative efficiency optimization, notice that attention with bias is a similarly important extension in dealing with complex tasks, where useful prior knowledge is introduced model learning. For example, in vision and language Transformers, the relative distance among tokens well-established attention bias and has been proven essential to the final performance of various tasks 225 24 (3)) Also, for scientific problems with rich domain-specific prior 42 attention bias is
bias term of dot- dot-product attention weights to guide the
an indispensable component, such as the pair representation bias in AlphaFold [1] 34]. However, the targeted efficiency optimization for attention with bias is still lacking. All previous research in
extending standard attention is centered on the speedup of attention with masks 35 where the
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Spatial Bins
Poir Bios
GRAM
Parallel Compute on SRAM
[20 HB)
1918/s
Block
Block
Write
1.5 400
Query, Key, Value, Blas
Results
GPU High Bandwidth Hamory (HBB))
(a) Attention Mask
(b) Attention Bias in Famous Transformers
(c) Attention Computation Process
Figure 1: (a-b) Comparison between attention mask and bias, where the spatial bias is from Swin Transformer [25] for computer vision, ALiBi bias is used in language modeling [30] and pair bias is
from AlphaFold [ (c) FlashAttention needs to read bias tensors from HBM to the on-chip SRAM.
sparsity nature of masks enables possible computation reduction. However, unlike the attention mask that defines computation logic, the bias matrix describes the pairwise relation among tokens, which is inherently continuous and dense as showcased in Figure [ making previous sparsity-based speedup techniques inapplicable. Although FlexAttention [TT]. benefiting from compiler techniques, new feature in PyTorch 2.5 (27), can support general formalizations of bias terms, it still depends on element-wise operations that are less optimized than matrix multiplications and fails in speeding up dynamic bias. To date, the fast computation of attention with bias remains a nascent area to explore.
In FlashAttention [10]. researchers find that instead of computation FLOPS, the read-write (IO) overload of GPU high bandwidth memory (HBM) is the actual bottleneck of speed. As the bias matrix usually dense, it really hard to bypass the quadratic IO complexity, where the computation needs to read the whole bias term from HBM at least once, making its speedup intractable. In this paper, notice that the IO challenge that face here can be recast into the classical compressed sensing problem [12). whose basic assumption is that the "measurement" (corresponding to I0 overload here) expensive but the computation (corresponding to the fast on-chip computation) is cheap. This new perspective offers us valuable theoretical understandings. Specifically, for a dense matrix, the optimal "measurement" (IO overload) highly related to the matrix's 's rank [6]. This inspires us to dive into the rank of attention bias and surprisingly find that most of the widely used biases are inherently of low rank, eventually discovering one flexible way to enable fast computation of attention with bias.
Based on the above understandings, this paper presents FlashBias based on the low-rank compressed sensing theory. By formalizing commonly used attention biases into the multiplication of two factor tensors, FlashBias achieves fast and exact computation for many Transformers, covering vision, language and scientific domains. Going further, we present an approximation method for general bias formalization with a low-rank neural decomposition technique, which successfully speeds up more complicated attentions AlphaFold 3 . Our contributions are summarized as follows:
Based on an in-depth study of the computation bottleneck of attention with bias, this paper theoretically proves that for a dense matrix, e.g. dot-product attention weights or various attention biases, the optimal efficiency in GPU is determined by the rank of the matrix.
Inspired by theoretical analyses, this paper presents FlashBias based on the low-rank com- pressed sensing theory, which utilizes exact and SVD decompositions for fast computation of many widely-used attention biases, and an approximation version for general biases.
FlashBias can accelerate family of widely-used backbones without loss of accuracy, which brings 1.5x speedup for AlphaFold and over 2 speedup for vision and language models.
2 Preliminaries
2.1 Attention with Bias
Attention [0] contains queries E IR NxC keys k E IR MxC and values v E MxC where N, M denote the sequence length of queries and keys respectively and C represents the channel of hidden representations. Conventionally, attention weights are calculated based on the dot-product between queries and keys. Beyond solely relying on the representation dot-product, useful prior knowledge is also commonly introduced into the attention mechanism as a bias term to guide learning, which is:
o = softmax( qk"
+ b)v.
(1)
VC

---

Here o E RNNC denotes results and b E R N represents the prior weight for query-key pairs, which is also referred to as attention bias [8] 25 Actually, attention bias has been widely used in diverse domains and has been proven to be indispensable for model learning, especially for complex tasks. For example, in computer vision, Swin Transformer (25) adopts the relative distance among different pixels as attention biases to introduce the essential locality inductive bias into vision backbones. Similarly, relative position is also used to encode sequential information in language models 14 39). In addition, attention bias widely exists in Transformer-based scientific deep models to incorporate -speiii priors, such as the pair representation bias in AlphaFold [IS. 34]
It is worth noting that attention mask is also commonly used in attention mechanisms to enable the calculation under a specific logic, such as the upper triangle mask (aka. causal mask) in decoder-only Transformers [2] and sliding window mask in sparse Transformers [3]. Usually, the mask is also applied as an additive term to the query-key dot-product. However, attention biases that we focus on in this paper are usually dense continuous values that incorporate detailed prior knowledge into each query-key pair to reflect the degree of token relations, while masks only contain zeros and negative infinite values and are usually sparse. This difference makes their speed-up strategies distinct.
Additionally, large language models widely adopt rotary position embedding (RoPE) ky introduce relative position into attention, which can be considered as a special multiplicative bias and is already computation efficient. Considering that RoPE is outside the additive bias scope. we defer its discussion to Appendix D and will showcase that FlashBias can also be extended to multiplicative bias.
2.2 Fast Computation of Attention
As the core component in modern deep models 18 25| 29], the attention mechanism has been widely explored. Although it has demonstrated impressive performance, one serious drawback is its computational complexity, which is quadratic w.r.t. sequence length. To tackle the efficiency bottleneck, the fast and exact computation of the attention mechanism has also been widely inves- tigated. Based on tiling techniques, researchers formalized the softmax function in attention with a sequential calculation process (26] 31), successfully reducing memory cost to linear complexity. Afterwards, FlashAttention [00] and FlashAttention-22 [9] further demystify that the I0 overload of GPU HBM the actual efficiency bottleneck and enable significant speedup by utilizing the super-fast fast on-ehip SRAM. The FlashAttention family has served as a default acceleration option of modern Transformers and has been built-in and supported by PyTorch (27]
Further, to extend the model's capability in handling diverse types of attention, researchers have made a great effort to enable fast computation for attention with masks, such as the causal mask 2 and sliding window mask (3. etc. Technically, the design principle in speeding up attention with masks to utilize the sparsity masks to reduce the IO complexity. For instance, Binary Block Masking (35) splits the mask matrix as blocks and adopts binary values to indicate the blocks with non-zero masks, which enables the masking process to skip the I0 read of all-zero mask blocks. Subsequently, FlashMask [4I] utilizes the spatial continuity of non-zero mask values and proposes to only record the start and end indices of mask segments, achieving fast computation under rich mask types. Although the above works can speed up a wide range of mask types, as attention bias is usually continuous and dense, these sparsity-
methods cannot be applied to attention bias.
Recently, FlexAttention 11 taking advantage of deep learning compiler techniques, has been released in PyTorch 2.5, which supports wide range of masks (e.g. causal mask (2) and biases (e.g. ALiBi (30)) by pre-compiling clement-wise operations. However, since many on-chip calcula- tions are less optimized than matrix multiplications, FlexAttention cannot achieve a perfect speedup and fails to support data-dependent dynamic biases. In contrast, FlashBias can take full ad advantage of extremely optimized matrix multiplications and is also applicable to complex dynamic biases.
3 Method
As aforementioned, attention with bias is essential for model learning, while its natively dense and quadratic tensor shape results in a serious IO bottleneck. Inspired by the compressed sensing theory [12 6), we verify that many well-established attention biases are inherently of low rank, paving the way for fast computation of attention with bias. Specifically, we dive into the computation of FlashAttention and theoretically prove that the optimal efficiency is inherently determined by the rank of attention weights and bias terms. Further, we present FlashBias based on the low-rank decomposition, which achieves optimal efficiency with natural adaptation for on-chip computing.
3

---

(6) FlashAttention with Blss
(6) FlashBles Computation
K.C
(RHG]HM
RR RxM
Food
Inser Lonp
V: M.K
Soy
Asner Loup
V: V.MKC
Quary
Attention with
Value
Low-rank Decomp
Quety
Attention with
QUA Lann
Online SoftMax
Exact Neural
Online SoftMax
Compuite otichip
Compute onochig
N:N
N:N
u: NxC
0,
(R+C)
0: N.C
Figure 2: Comparison between FtashAttention and FlashBias, FlashBias utilizes low-rank decompo- sition bypass the read of the whole bias matrix, successfully avoiding the quadratic IO overload.
3.1 Rethinking FlashAttention Computation
We begin by analyzing the theoretical basis of FlashAttention's 's speedup (without bias mask) dense and continuous attention weights. Our analysis yields that the rank of dense matrices, such as the dot-product attention weight qk' _ N and the bias matrix b IR XM inherently decides the IO cost, which is formally stated as follows. All the proofs can be found in Appendix AA]
Theorem 3.1 (FlashAttentionn computation benefits from low rank). Suppose N = M and let be the rank of dot-product attention weight s. C aN be the channel dimension with constant and sequence length N, be the size of SRAM with S = 3NC and N S 3 S 1. Then, 1) the HBM access of FlashAttention is θ ((1 + (8) times smaller than the standard attention, and 2) 2 R
N
As demonstrated Theorem the speedup of FlashAttention proportional to 3 (SRAM size) and inversely proportional to (channel dimension). If consider the attention weight matrix, the optimal speedup of FlashAttention is obtained by reducing the channel dimension to R, ic. R. The same technique is also used in DeepScek-v3 [3) Multi-Head Latent Attention, which reduces the channel dimension by projecting q. k, v into a small latent space for acceleration.
low-rank
Theorem 3.2 (Compressed sensing complexity of low-rank dense matrix (). Given a N N dense matrix with rank R, the theoretically optimal compressed tensor is of storage complexity ( R).
Theorem 3.2 demonstrates that the optimal storage of bias linearly related to its rank, highlighting the essentiality of the low-rank property. Further, integrating with prior analyses of HBM access in exact attention [9) we can derive the 10 complexity property of HBM access in attention with bias.
Corollary 3.3 (A "lower bound" for HBM access of attention with bias). Given q E NxC. k. v MMC bias b of rank R and SRAM of size S where (C + R) < S S N + R), there does not exist an algorithm to compute exact attention with bias through of NM +R"))+ HBM access for all e [[C + NC, + R)). Here o(*) represents the strict asymptotic upper bound.
8
3.2 FlashBias
Inspired by the above theoretical results, present FlashBias based on low-rank decomposition techniques, with novel design to utilize the low-rank property of attention bias to reduce HBM access.
Overall design As shown Figure 2 instead of block-wise reading bias, FlashBias replaces the quadratic bias matrix b E N two factor tensors. Specifically, considering bias matrix calculated by b f(x: X1)), where Xg NCC" Xx E R MxC" represent the source information for generating the bias, which is set as spatial position each pixel in Swin Transformer 255 and the representation of protein residues in AlphaFold [1, if there exist factor functions Qq Ók satisfying:
f(xq xk) = &g (xq (xx) óq. O& RC + R
The computation of attention with bias can be equivalently formalized as follows:
qk' VC
[ql [alvCrba (xq) [k (klox (xx))x)]' )v,
0 = softmax
+ bbv softmax
(3)
VC

---

Table 1: The computation of FlashBias for widely-used attention biases, which includes three different types: (a) Exact decomposition by finding exact Φq. @k, (b) SVD decomposition for cases using model parameter as bias, (c) Neural decomposition for using model representation as dynamic bias.
Model
Original
(a) Exact Decomp
(b) SVD Decomp
[c) Neural Decomp
Domain
Bias / Model
Type
Exact Blas f
Approximated Blas
Approdimated Blan
Language
ALiBi
3D
(a)
Attention with Bias
Motmul
Matmul
Matmul
Vision
Swin Trans. [25
(b)
dx
dgg
OkLe f
Spatial Distance (a)
Science
Pangu- Weather [
Pangu-Weather
(b)
Attention with
Xg
Bias
X4
AlphaFold [I]
(c)
where [+]*] denotes the concatenation operation along the channel dimension. Notably, this design significantly reduces the storage cost for the attention bias from O MM) to O ((N + R)) Al-
though its design will require recalculating the bias weight, this computation is just simple matrix multiplication of (x (xx))x((x))))))))) T, an operation that has been extremely optimized on modern GPUs.
Such simple design is broadly applicable to wide range of variants for attention with bias. In practice, we implement through three concrete instantiations for 0q. 6k. as shown in Table [
Exact decomposition We find that some well-established attention biases can be directly decom- posed into factor functions, enabling fast and exact computation. Here are representative cases.
Example 3.4 (ALiBi [30] in language models). Given Xg = [1,... .N1,xx [1.... .MI the ALiBi bias is calculated as f(Xq.. Xk.)) = j. which can be directly decomposed into low-rank formalization by defining (Xq.) = [1, and Ór (xx.)) [-j, 11, corresponding to the case R 2. The original ALiBi also involves causal mask, while we only focus on the bias term here.
=
Example 3.5 (Spatial distance in scientific problems). Transformers has been used as surrogate models for PDE solving [433. especially for aerodynamic simulation. It is critical to introduce spatial distance to guide atlention learning among massive computational points. Let Xq Xk E RRN record the 3D spatial positions of N computation points, where Xq E RR' is the position of i-th point. For the spatial distance f(xq,;: xk.j) = ||x Xk., 112, can be exactly decomposed as:
(g(xx.)) [xq 1, -2x x 1, 2xg x2 x 1, 2,2
(4)
Φr (x(x)))) = [1, xk. j.: Xk.j.,, 1, xk.,,, Xk.j.1, 1, x
2. Xk.),2)-
SVD decomposition Some models such as Swin Transformer 255 and Pangu- Weather 44 adopt the Icarnable model parameters for relative position encoding. Specifically, each bias term in their model is an N x .M matrix of model parameters. As this type of bias is fixed once the model has been well trained, it is convenient to conduct Singular Value Decomposition (SVD) (0) for low-rank decomposition of these parameters. In practice, we precompute SVD once offlin,, incurring negligible runtime overhead. The resulting decomposed factor tensors can then be utilized to accelerate the subsequent inference process, thanks to their low-rank nature.
Neural decomposition The two mechanisms discussed above apply to static bias terms. Beyond these, some models introduce more complex or dynamically generated biases. For example, in AlphaFold [1, the bias term is projected from the intermediate pair representations. In this case, its low-rank decomposition cannot be explicitly or exactly derived. Also, due to the data-dependent property, SVD decomposition needs to be conducted for bias at every inference. To accelerate these complex biases, we present neural approximation version of FlashBias, which employs two lightweight neural networks Φq.0, 0 Øk.8, : E.C + IR approximate factor functions dq and ok. which is supervised by the following objective function:
min L(xq: xx) = I1āq.1 (xq)ōx., - l(xqq xx) 8,
(5)
Here 01,0, are learnable parameters, which can be optimized by fine-tuning 01, 02 on the training set. Note that FlashBias attempts to completely replace the original bias term, a design fundamentally different from LoRA [T6] which learns an additive term to the pretrained model parameters. Similar to the SVD decomposition version, once these two lightweight neural networks Φq.0, 6k.00 have been well optimized, they can be directly applied to all future inference.
5

---

(A) Train GPU Mamory (3R)
(0) Test GPU Momory (GB)
Train Running Time (A/100iter)
(4) Test Running Time (2/1 (Biters)
Atstion of Ban
Flavedtioe fine w/ Bon
eln Six
PloisA/wrdee Sixe
ResAttrtio
Fextttaticn
Hardizs t-d
Foskiio
Tio B-l6
Plardias 1-16
Fin Bx37
Fasidkes b+22
Fuaidiin
Fishiion
Plevikis 1464
Pasldtas
Haridtiss
Hosh Atttittttinn
Pure Hlaistherktn
Para cabAtweiion
Pure /lasbonecte
vid 20t get get with with gith year
5ot yiu you get yas une wnth wer
1.6
V014
2016
1020
1249
yotn
2000
Sequencee
Langth
Sequenee Langth
Sequence Length
Sequerice Length
Figure 3: Efficiency comparison among FlashBias, FlashAttention w/ Bias and FlexAttentio [. Here "Pure FlashAttention" refers canonical FlashAttention without bias term, which can be viewed as an upper bound of computation efficiency. Dotted lines indicate out-of-memory situations.
Remark 3.6 (Understanding neural decomposition). In Eq. (). define factor networks Φq,0 9k.82 as token-wise functions, instead of depending on the whole sequence. This is because elements in bias that reflect pair-wise relations are solely determined by corresponding token information, such as token position in relative position bias. Besides. rank R essentially controls the dimension of dot-product space, thereby should be set based on "complexity" prior knowledge that derives bias.
Corollary 3.7 (HBM access of FlashBias). Given q E RRN ,k, v E TRMXC, .runk- R bias b and size-S SRAM (C S S S NC), the HIM access complexity of FlashBias is O NR(()+R²))
S
Remark 3.8 (Trade-off between approximation accuracy and efficiency). Although SVD and neural decomposition will introduce approximation error, owing to the native low-rank property of attention bias in modern Transformers, it will not affect the final performance by setting R as a reasonable value. For example, in some layers in Swin Transformer, R = 32 can maintain over 99% energy of the original bias with shape 576 x 576. More case studies are included in Appendix Baaudixxxaa (ane
Example 3.9 (Comparison with FlashAtteniion). For FlashAttention with bias, its HBM access is 9
( NMC + N M) 19), Given a regular setting of Transformers and GPU (C = 64, is 100KB) and supposing R = 64, M.M >> C, R, then HBM 10 of FlashBias is R 6x smaller than FlashAttention.
Speed up training In FlashBias, the exact decomposition version can be directly applied both training and inference phases. As for SVD or neural decomposition, their computations are based on pretrained model; thereby, the inference speedup is obvious. Note that for large-scale pretrained models, such as AlphaFold [1. the inference speedup is already sufficiently valuable. Going further, it is also possible to speed up the training phase with these two types of biases. Specifically, for the SVD type, we can replace the N x M model parameter bias with N x R and M x R model parameter tensors at the initialization phase. As for the neural type, instead of using the whole bias matrix, can also be replaced with two lightweight neural networks when defining the model architecture.
4 Experiments
As summarized in T we extensively test FlashBias on Transformers for language mod- eling, computer vision and scientific problems, which can significantly speed up model effi- ciency at both training and test phases. This section will first present an overall efficiency comparison among different techniques and then integrate FlashBias to different domain- specific Transformers with various bias types. All experiments were performed in PyTooch222 2.5.0 27 and Triton 3.0.0 38 single A100 GPU.
Table 2: Summary of experiments and base models.
Base Model
Bias Type
Phase
Plain Transformer
40
Static
Train & Inference
GPT-2 62)
Static
Train & Inference Inference
Swin Trans.
24
Learnable
PDE Solver AlphaFold
Learnable Train & Inference Inference
Learnable
4.1 Overall Comparison
Setups To give a clear comparison among FlashBias, vanilla FlashAttention with Bias 99 and the latest FlexAttention [T]. make comprehensive efficiency evaluation based a plain

---

(9) Train GPU Mamory Ratio
(b) Test GPU Memory Ratio
(c) Train Running Time Ratio
(6) Tast Running Time Ratio
Fosh-ktentonn km
Hushittrtion Sizs
Plalvetention w/
Faclctuaton
Rec
Flxcttaaion
Frddiis 2-d
Rusitio I-
Resiblin
Floeddis R-52
Rasiitio
Plasitiia Brid
Flocidias Brdd
501
uil wiih years
get wit with uith yorti
50
90
1000
Sequenoe Length
Sequence Longth
Sequerccel Langth
Sequercce Lenghh
Figure 4: Efficiency ratio over "Pure FlashAttention", which calculated by
Method Efficiency
Pure FlashAAtention Efficiency
Transformer 44),, which consists of 8 layers. Each Transformer layer involves feedforward network with 1024 intermediate hidden channels and attention with 512 hidden channels, heads, as well as a static bias matrix of shape # heads x N x N. All the metrics are recorded with batch size of 1.
Results As shown in Figure and FlashAttention with Bias 9 and FlexAttenion (T]. demonstrating the effectiveness of our design.
find that FlashBias (red lines) consistently outperforms
According to Figure Ba-b), when sequence length N 16384, FlashBias can significantly reduce GPU memory usage, which is 5x smaller than FlexAttention and FlashAttention with Bias during training and 10x smaller during inference. This memory benefit can also be theoretically justified by Theorem 3.2 As for training and inference running time, FlashBias presents 18.0 and 44% speedup compared to vanilla usage of FlashAttention and is also better than FlexAttention in long sequences.
worth noticing that, although FlexAttention 11 which optimizes FlashAttention with deep learning compiler techniques, can significantly boost the running time compared to vanilla FlashAt- tention with Bias (Figure (c-d) orange lines), it still cannot reduce GPU memory due to the quadratic storage cost of the bias matrix ((a-b) orange lines). In addition, due to the HBM access overload of the bias matrix, it falls short in running time when inputting long sequences. These observations further demonstrate that the low-rank property utilized in FlashBias is the key to efficiency.
Implementation choices As formalized in Eq. (3) for implementation, researchers can concatenate the decomposed factor functions with query and key tensors in the original attention along the channel dimension, which means FlashBias can seamlessly adopt previous efficient imple- mentations for canonical attention, such as PyTorch SDPAN In pursuit of well- optimized operation, we also implement FlashBias based on the Triton backend (38). Figure 5S demonstrates that the Triton im- plementation is much more efficient in the forward pass and the PyTorch SDPA- based version is better in the training phase. Therefore, we adopt the Triton version for inference and PyTorch SDPA version for training throughout this paper. Besides, although vanilla PyTorch SDPA can accelerate the computa- tion, it still suffers from the out-of -om-ooo issue, which can be perfectly resolved by FlashBias.
(a) Train Speed (TFLOPs/s)
[b) Test Speed (TFLOPsA))
on (Trk
[fr)
Fuvidiio (PyTarch 50%
Flasktiis (D
PyTord--DDAA with Bos Matric
Pylorch-SPPA with Jin Motio
Sequence Longth
Sequence Length
Figure 5: Speed of different implementations on atten- tion with 128 hidden channels and 8 heads. R is set as 8 in FlashBias. Dotted lines indicate -femory More attention speed comparisons are in Appendix B
4.2 Large Language Model
Setups Here employ FlashBias to speed up ALiBi 30], which is well-known bias in language modeling and has been proven better than rotary position embedding [37] in handling input length extrapolation. We follow the configuration in GPT-2 (2) and evaluate based on large language model with ALiBi bias, which contains 48 decoder-only Transformer layers, 1.5B parameters total.
Torch implementation of torch, wwt.:/ffunntiioonac..lee..eooeee
7

---

R-12
R-1
R-194
R-17
R-114
R-3
R-1 R-8
R-309
R-42
R-1 R-4 R-1 R-1 R-12 R-6
Figure 6: Visualization of bias matrix in SwinV2-B. can maintain 99.5% energy of the bias matrix.
Each layer contains attention with 1600 channels and 50 heads, well feedforward network with 6400 hidden dimensions. Beyond large model size, this task also involves the causal mask.
stated Example 3.4 we can adopt the exact decomposition for ALiBi bias, where R is equal to and the result of FlashBias is exactly equivalent to the original computation. Since attention masking has been well-optimized in FlashAttention [9] and FlexAttention [T] we directly integrate them with our method support the decoder-only Transformer. This combination also demonstrates that FlashBias is orthogonal to attention attention-masking speed-up techniques. Notably, FlashAttention has ALiBi slopes feature, which does not load bias matrix from HBM but computes the bias weights through just-in-time compilation. Since this is not a generic technique, we do not employ this feature in the main text. More discussion can be found in Appendix [
Implementation
Results As mentioned above, since this task requires causal attention, we test FlashBias by integrating it with the attention mask speedup techniques in FlashAttention and FlexAttention.
Table 3: Experiment of GPT-2. #Time records the running time for 100 iterations when = 2048. refers to the time difference w.r.t. Pure Causal *Attention without bias, reflecting the additional cost in processing the bias matrix.
Results Table B demonstrate that FlashBias still outperforms FlashAt- tention and FlexAttention in process- ing the ALiBi bias term. Specifi- cally, in processing bias, FlashBias reduces over 50% (5.0+2.3) time cost of FlashAttention and over 25% (3.4+2.55) time cost of FlexAttention.
Training
Inference
Method
Time(s)
Time(s)
Pure Causal FlashAttention FlashAttention with Bias
119.3
38.77
124.3
5.0
40.32
1.55
FlashBias Bias (Ours)
121.6
2.3
39.26
0.49
Pure Causal FlexAttention FlexAttention with Bias FlashBias (Ours)
119.0
38.76
122.4
3.4
40.03
1.27
121.5
2.5
39.78
1.02
Note that ALiBi bias is relatively sim- ple, enabling the compiler-bri based FlexAttention with favorable speedup. However, it may degenerate in handling more complex bias matrices, such as the bias in Swin Transformer (24) in the next section.
4.3 Vision Transformer
Setups We test FlashBias on the image classification task based on Swin Transformer v2 24 Specifically, we adopt the open-source S2in mode This model contains layers with input resolution as 384 384 and window size 24 x 24, thus, the sequence length of its WindowAttention 576. Attentint in every layer contains a relative position bias with size # heads x576 x 576, which is set learnalle model parameter. We attempt to speed up the Attention computation with FlashBias. All the efficiency metrics are evaluated under a batch size of 64.
Implementation Since the biases are set as learnable parameters, we directly obtain them from pretrained model and adopt the SVD decomposition to generate decomposed factor tensors. However, we observed that in one layer, not heads present low-rank bias matrices, as shown in Figure Thus, we only adopt FlashBias to layers where most of heads' biases are of low rank, that is, the last layers of Swin while employing vanilla FlashAttention with Bias for the other layers.
Results In Table we can find serious performance drop "Pure FlashAttention" (without bias, Acc @1 1 87% % 176--999 39%), which verifies the essentiality of bias and also highlights the importance of
Model and training contigurations of SwinV2-B.

---

speedup attention with bias. Further, compared to the official implementation, FlashBias can reduce 60% running time and 27% memory cost, which is valuable for real-time development.
More importantly, FlashBias will not affect the performance, where the top- 5 accuracy drops less than 0.02% and the fluctuation of the top-1 accuracy (+0.042%) is within the standard de- viation. In particular, even using well- established quantization techniques will still cause a 0.64% top-1 1 accu- racy drop for 22% speedup according to the PasterTransformer document" This comparison demonstrates that the low-rank property a principal basis for fast computation of attention.
Table 4: Experiment of Swin on ImageNet- #Time and #Mem correspond to inference efficiency on A per batch. Offline calculation of SVD for all biases takes 4.79s.
Method
Acc@1 1
Acc@5
Time(s) Mem(MB)
Official Code
87.144%
98.232%
0.473
12829
Pure FlashAttention
9.376%
19.234%
0.180
3957
Flash with Bias
87.142%
98.232%
0.230
11448
FlexAttention [ INT8 PTQ
87.142%
98 98.232%
2.885
25986
86.46%
Around 22% speed up
FlashBias (Ours)
87.186%
98.220% 0.190
9429
Also, FlexAttention seriously degenerates in Swin V2-B. This is because, unlike experiments in Section 4. Swin Transformer's bias matrices are different in value and shape among different layers, which requires recompilation each time. This issue has also been mentioned by FlexAttention's 's author as "I you are adding (B, H, N. NJ bias tensor, then you honestly shouldn' 't be using FlexArtention"
4.4 Scientific Deep Models
In addition to language and vision tasks, scientific problems usually involve rich domain-specific prior knowledge; thereby, attention bias also widely exists in scientific Transformers. Here we evaluate FlashBias in two representative models: Transformer-baser PDE solvers 4 and AlphaFold M.
Transformer-based PDE solvers Attention mechanism has been proven equivalent to Monte- Carlo approximation of the integral operator justifyngg its theoretical foundation for solving partial differential equations (PDEs). However, in processing complex geometries, the attention mechanism may fall in perceiving the 3D space, encouraging the utilization of spatial distance prior. Here we follow the driving car simulation task in 43] whose input is the position of computation mesh points and output is the physics quantities on these points. We test FlashBias on an 8-layer Transformer solver with a 3D distance bias described in Example 3.5 Each layer contains attention with 128 hidden dimensions and 8 heads, as well as a feedforward network with 256 hidden channels.
To approximate the adaptive mesh in numerical methods 36), include token-wise learnable weight a; for the 3D distance bias in each head of every layer, i.e. fxx... Xk.j) = 0s [x||xxq; Xk.j 1
Unlike bias discussed in ALiBi 30 or Swin V2 (25), the learnable weights require the training process to record the gradient of the bias matrix, posing a challenging efficiency issue in backpropagation.
Table 5: Experiments of Transformer PDE solvers. The efficiency metrics are recorded under a batch size of in the format of #Mem (GB) / #Time (s/100iters). Accuracy comparisons are in Appendix G
Training Phase
Inference Phase
Method (1 (Learnable Bias)
8192
16384
32186
8192
16384
32186
FlashAttention
12.8 / 15.4
OOM
OOM
4.54 5.46
15.3 / 21.2
OOM OOM
FlexAttention
Not supported in current version
21.9 184.0
OOM
FlashBias (Ours)
1.46/ / 4.54
2.02 14.7 2.97
51.1
0.98 1.22
1.03/ 3.48
1.13 / 12.7
Results FlashBias is the only method that can support training of the Transformer solver on 32186 points (Table 5),, which also presents a significant memory and running time advantage compared with other methods. Notably, FlashAttention and well, as they need to record dense gradient matrices, further highlighting the practicality of FlashBias.
cannot support learnable bias training
AlphaFold 3 for protein folding As a significant progress of AI for science, AlphaFold 3 [I] employs an elaborate architecture. Specifically, its core design, Pairformer, contains 144 attention
TisterTrransoomee released by NVIDIA.
Discussion about FlexAttention with dynamic bias matrix.

---

Table 6: Experiment of AlphaFold 3. The left part illustrates folding examples. Note that AlphaFold based on a diffusion model; thereby, slight differences are normal. The right table's s efficiency is tested the one-time inference of Pairformer with 1218 protein residue tokens (PDB ID: 7wux).
Test Set
PDB ID 7wux
Method
PLLDDT Loss
|PTM Time(s) Mem(GB)
Open- Code
3.3724
0.9500 26.85
13.62
FlashAttention w/ Bias
4.3669
0.1713 8.27
12.89
FlashAttention Bias
3.3724
0.9500
20.39
13.62
POB ID. Pouss
PDB 7vex
FDB ID: 7pab
FlashBias (Ours)
3.3758
0.9499 18.19
13.62
Prediction
Speadup FlashBlas
blocks, all of which contain the bias term pair representations, Thus, its speedup is highly related to the fast computation of attention with bias. Since Alphafold is not officially open-sourced, we follow a public high-quality reproduction, Proteinix 8 which is implemented PyTorch [27].
After detailed analysis of AlphaFold 3, we find that its efficiency bottleneck is triangle self-attention which involves the bias matrix projected from intermediate pair representations, making vary among different samples, layers and heads. To approximate this complex bias, we employ the neural decomposition version of FlashBias, whose inputs Xq: X are set as the combination of pair and single protein representation. We fine-tune neural basis functions for 10,000 iterations on the PDB dataset. Since we only need to optimize the newly added parameters in 0q.0. .k. this process only takes about 10 hours on single A100 40GB GPU, then you can infer new protein with FlashBias. For comparison, the full training process of AlphaFold will take about days on 128 A A100 GPUs.
Results Table that, compared to the public code, FlashBias reduce the running time by 32%. In addition, removing bias (w/o bias) will significantly improve the efficiency (26.85s vs. 8.27s) but will seriously damage performance. This observation highlights the essentiality of accelerating attention with bias. Despite FlashBias being based on neural decomposition, it will not affect the final performance, whose metric fluctuation is within the standard deviation. Beyond inference, FlashBias can also be a promising tool for training speedup (see Appendix
E for details).
Neural decomposition visualization To give clear illustration of neural decomposition, we also plot the original pair representation bias and FlashBias approximated bias in Figure Neural decomposition can give a relatively accurate estimation of the bias matrix, which performs well in capturing the "texture" of the original bias. In addition, it is also observed that neural decomposition is not completely perfect in reconstructing the diagonal weights. Despite this deficiency, FlashBias still maintains the original accuracy of AlphaFold 3. This may benefit from the dot- dot-product eeffaatention and residual connection, which can give robust and dominating weight for relation modeling.
PDB (D: 7r8r
PDBID:: 7pzb
Neural Decomp
R+130
Ry-45
Rm83
R=1225
R-206
R-49
R-80
Re-227
Original Bios Matrix
Figure 7: Comparison between neural decomposed factor tensors' multiplication and original bias. Biases of 7r6r (245 residues) and 7pzb (600 residues) in the first layer of Pairformer are plotted, which contains 4 heads. We also mark the rank value that can maintain 99% energy of original biases.
Conclusions
This paper focuses on the fast computation of attention with bias, which is an essential extension of the attention mechanism and is widely used in language, vision and scientific domains. After in in-depth analysis of FlashAttention, we notice that the optimal efficiency depends on the low-rank property of the attention weight. This further inspires us to present FlashBias based on the low-rank compressed sensing theory. where we also present three practical methods for low-rank decomposition of attention bias, achieving theoretically favorable efficiency. Experimentally. FlashBias can seamlessly support the fast computation of wide range of famous Transformers without loss of accuracy.
10

---

Acknowledgments and Disclosure of Funding
This work was supported by the National Natural Science Foundation of China (U2342217 and 62021002), the BNRist Innovation Fund (BNR2024RC01010, and the National Engineering Re- search Center for Big Data Software.
References
[1] Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Rouneberger, Lindsay Willmore, Andrew J Ballard, Joshua Bambrick, ct al. Accurate structure prediction of biomolccular interactions with alphafold 3. Nature, 2024.
[2] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint arXiv:2303.08772 2023.
[3] Iz Beltagy, Matthew Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05//55 2020.
[4] Kaifeng Bi, Lingxi Xie, Hengheng Zhang, Xin Chen, Xiaotao Gu, and Qi Tian. Accurate medium-range global weather forecasting with 3d neural networks. Nature, 2023.
[5] Boris Bonev, Thorsten Kurth, Christian Hundt, Jaideep Pathak, Maximilian Baust, Karthik Kashinath, and Anima Anandkumar. Spherical fourier neural operators: Learning stable dynamics on the sphere. In ICML, 2023.
[6] Emmanuel J Candès and Terence Tao. The power of convex relaxation: Ncar-optimal matrix completion. IEEE transactions information theory, 2010.
[7] Venkat Chandrasekaran, Sujay Sanghavi, Pablo A Parrilo, and Alan S Willsky, Sparse and low-rank matrix decompositions. IFAC Proceedings Volumes, 2009.
[8] Xinshi Chen. Yuxuan Zhang, Chan Lu, Wenzhi Ma. Jiaqi Guan. Chengyue Gong, Jincai Yang, Hanyu Zhang, Ke Zhang, Shenghao Wu, Kuangqi Zhou, Yanping Yang, Zhenyu Liu, Lan Wang, Bo Shi. Shaochen Shi, and Wenzhi Xiao. Protenix advancing structure prediction through comprchensive alphafold3 reproduction. bioRxiv, 2025.
[9] Dao. Flashattention-2:: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv.230700669) 2023.
[10] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. NeurIPS, 2022.
[11] Juechu Dong, Boyuan Feng. Driss Guessous, Yanbo Liang, and Horace He. Flex attention: A programming model for generating optimized attention kernels. arXiv preprint arXiv:2412.05496, 2024.
[12] David L Donoho. Compressed sensing. IEEE Transactions on information theory, 2006.
[13] Robert Stratman Elliot et al. Electromagnetics: history, theory, and applications. 1993.
[14] Team GLM, Aohan Zeng, Bin Xu, Bowen Wang, Chenhui Zhang, Yin, Dan Zhang, Diego Rojas, Guanyu Feng, Hanlin Zhao, et al. Chatglm: A family of large language models from glm-130b to glm-4 all tools. arXiv preprint ar 2r0:6:2066. 12793, 2024.
[15] Hans Hersbach, Bill Bell, Paul Berrisford, Shoji Hirahara, András Horányi, Joaquín Muñoz-Sabater Julien Nicolas, Carole Peubey, Raluca Radu, Dinand Schepers, al. The cra5 global reanalysis. Quarterly Journal the Royal Meteorological Society, 2020.
[16] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen Zhu, Yuanzhi Li, Shean Wang. Lu Wang. Weizhu Chen, al. Lora: Low-rank adaptation large language models. ICLR, 2022.
[17] Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira. Perceiver: General perception with iterative attention. In ICML, 2021.
[18] John Jumper, Richard Evans, Alexander Pritzll, Tim Green, Michael Figumov, Olaf Ronneherger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 2021.

---

[19] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In JCLR, 2015.
[20] Virginia Klema and Alan Laub. The singular value decomposition: Its computation and some applications.
IEEE Transactions on automatic control, 1980.
[21] Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Noural operator: Learning maps between function spaces with applications to pdes. JMLR, 2023.
[22] Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang. Patrick Labatut, Daniel Haziza, Luca Wehrstedt, Jeremy Reizenstein, and Grigory Sizov. xformers: A modular and hackable transformer modelling library. https://githur: wtt////ccbboorreerrratttttoooora
2022.
[23] Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, al. Deepseek-v3 technical report. arXiv preprint 2xiv:2412.114451 2024.
[24] Liu, Han Hu, Yutong Lin, Zhuliang Yao. Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, et al. Swin transformer v2: Scaling up capacity and resolution. In CVPR, 2022.
[25] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Ching- Feng Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In ICCV, 2021.
[26] Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv preprint arXiv: arXiv:1805. 1805.02857.. 2018.
[27] Adam Paszke, S. Gross, Francisco Massa, A. Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Z. Lin, N. Gimelshein, L. Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high peegh-eeffrmacce deep learning library. In NeurIPS, 2019.
[28] SGOPAL Patro and Kishore Kumar Sahu. Normalization: A preprocessing stage. arXiv preprint 10Xiv.1503.06462, 2015.
[29] William Peebles and Saining Xie. Scalable diffusion models with transformers. In ICCV, 2023.
[30] Ofir Press. Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In ICLR, 2022.
[31] Markus N Rabe and Charles Staats. Self-attention does not need (n²²) memory. arXiv preprint arXiv:2112.05668 2021.
[32] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog. 2019.
[33] John L Russell. Kepler's laws of planetary motion: science, 1964.
The British journal for the history of
[34] Andrew W Senior, Richard Evans, John Jumper, James Kirkpatrick, Laurent Sifre, Tim Green, Chongli Qin, Augustin Zidek, Alexander WR Nelson, Alex Bridgland, ct al. Improved protein structure prediction using potentials from deep learning. Nature, 2020.
[35] Agniv Sharma and Jonas Geiping. Efficiently dispatching flash attention for partially filled attention masks. arXiv preprint arXiv:2409. 15097, 2024.
[36] Pavel Solin. Partial differential equations and the finite element method. John Wiley & Sons, 2005.
[37] Jianlin Su. Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 2024.
[38] Philippe Tillet, H. T. Kung, and David Cox. Triton: intermediate language and compiler for tiled neural network computations. In ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, 2019, 2019.
[39] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Ama Lachaux. Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint 137 2023.
12

---

[40] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez. L ukasz Kaiser, and Illia Polosukhin. Attention all you need. In NeurIPS, 2017.
[41] Guoxia Wang, Jinle Zeng. Xiyuan Xiao, Siming Wu, Jiabin Yang, Lujing Zheng, Zeyu Chen, Jiang Bian, Dianhai Yu, and Haifeng Wang. Flashmask: Efficient and rich mask extension of flashattention. ICLR, 2025.
[42] Hanchen Wang, Tianfan Fu, Yuanqi Du, Wenhao Gao, Kexin Huang. Ziming Liu, Payal Chandak, Shengehao Liu, Peter Van Katwyk, Andreea Deac, al. Scientific discovery in the age artificial intelligence. Nature, 2023.
[43] Haixu Wu, Huakun Luo, Haowen Wang, Jianmin Wang, and Mingsheng Long. Transolver: A transformer solver for pdes on general geometries. In ICML, 2024.
[44] Haixu Wu, Hang Zhou, Mingsheng Long, and Jianmin Wang. Interpretable weather forecasting for worldwide stations with a unified deep model. Nature Machine Intelligence, 2023.
[45] Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip HS Torr, and Vladlen Koltun. Point transformer. ICCV, 2021.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: This paper presents practical algorithm for fast computation of attention with bias. The theoretical results are derived based on the proofs in Appendix and the effectiveness of proposed algorithm is supported by results in Table
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We have discussed in Appendix K and also gave a possible solution.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
14

---

Answer: [Yes]
Justification:
All the proof' can be found in Appendix A
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We have clearly stated the implementation details in the first paragraph of every experiment section and further detail them Appendix[] []]
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
15

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: All the experiments based the open-sourced GitHub repository. The data is public. Code is available at this repository: https:///ithh: com/ tthmm/FlashBi
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://mips.cc.
coolicc/gulees/ooeeebbbasss www.eeeeuuumisssoooooooooooo for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ //nips. wwtps/cccppplgguuuiieeeeaec
Stteeee//ooSSSSeissSioica www.ssssiicggcooo for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
All the implementation details have been included in Appendix !
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Since we mainly focus on the computation efficiency in this paper, all the efficiency metrics are quite stable. Among different rounds of experiments, the GPU memory usage is fixed and the running times' difference is less than 0.01s for 100 iterations.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
16

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CL, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We have included this information in Appendix !
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer:
[Yes]
Justification:
We have strictly followed the NeurIPS Code of Ethics throughout this paper.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We have discussed broader impacts in Appendix
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
17

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Yes, have faithfully eited the corresponding paper and open-sourced code according to their licenses.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
18

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
19

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: We only try to speed up GPT-2 in one of our experiments, while LLMS are not included in our core method.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
20

---

A Proofs for Theorems in the Main Text
Due to the page limitation of the main text, present the proofs for theorems in Section 3 here.
Proof of Theorem 3.1 This theorem can be proven based on the theoretical analyses in FlashAtten- tion 41 and the basic low-rank compressed sensing knowledge.
Proof. According FlashAttention 41,
the theoretical complexity of HBM accesses in the standard
attention and FlashAttention are:
Nº S
FlashAttention: 1ish
StandardAAtenntion: 10 standardd
θ (NC (NC++N9) (6)
where + represents the asymptotic tight bound. Given C = aN and S =BNC, the ratio of HBM
access in standard attention over FlashAttention is:
1aach
(NC + N2)5 = = ( (((0+ ( 1)
= (01+ -
=
(7)
staanddard
NºC²
a²
Also, since R the rank of the attention weight s, which is the dot-product of queries and keys, thus & R < C. Further, we can derive that Z N
As cited in the main text, the least storage cost of an N x N dense matrix with rank R is equal to (2NR - R²), whose proof can be found in [6]. Since R < N, then
Proof Theorem 3.2
NR 2NR - R² < 2NR.
8
Thus, the optimal storage overload for an N x N matrix is equal to ((NR).
Proof of Corollary 3.3 This corollary is derived from Proposition
FlashAttention [0] Specifically, consider the extreme that S θ (N(C + R)), in this case, have
NM + R2) 1) (M(C + R)). N(C + R)
Since we have to read the M x C keys and M R decomposed bias factor tensor from HBM, there
does not exist an algorithm to finish the computation of attention with bias in (M'C' + R)) access.
Proof of Corollary 3.7 This corollary can be derived from Theorem in FlashAttention 10
Proof. As formalized in Eq. 3 FttshAttention needs to read queries, keys, values, as well as decom- posed factor tensors of biases into the on-chip SRAM for computation.
FlashBias's 's computation is based on the same tiling method as FlashAttention v2 (9. The N * case corresponds to the noss-attention, which is used for fusing additional information [17), thus we suppose M 2 N here. Suppose that the computation splits queries into T blocks. The overall HBM access complexity is θ (N(C R) + M(C R)T) θ (M(C R)T). Let the block size of keys and values as Bg. x (C + R) and the block size of queries as Bg x (C + R). Considering the size limitation of SRAM is S, following FlashAttention v2 (9) these hyperparameters are as:
S
S
5 (RRB))
S Thus, the overall HBM access complexity is
min C' (C+
Bg CC+R*
BR,v θ
min C+R' Bg
=
(10)
Then, the number of blocks T
N By
0
N(C+R)
NMMCC R)²
= (M + +R²)
θ
(11)
Recall Corollary 3.3 we can find that the above theoretical complexity is quite well-optimized as there does not exist an algorithm with o [NM(C*+ complexity.
(NMIC
Calculation in Example
3.9 Considering N, M > C, it easy calculate the HBM access
ratio between FlashAttention and FlashBias as follows:
(1 +1+50 C2 ) (1+ 100x 1024 )00 x 1024
2x64²
6,
(12)
C R2
2(64 642)
where we consider the half-precision float, whose storage cost is equal to 2B.
21

---

B Attention Speed Measurement
As practical technique, in the main text, we primarily focus on the efficiency benefits to the entire Transformer brought by FlashBias. As supplement, we conduct detailed comparison solely on the efficiency the attention mechanism's ' computation. Specifically, we compared several advanced implementations of attention: FlashAttention [9], PyTorch SDPA and xFFormers [22], as well as the vanilla PyTorch implementation. Since the official implementation of FlashAttention does not support backpropagation for bias matrix, we adopt open- soured Triton extension for experiments
Mon-Causal Attention In this part, consider an full attention with 4 heads and each head contains 32 channels. As shown in Figure & if we only consider the attention speed, FlashBias is over 2 faster than previous best implementation in inference and 1.3x faster in the training phase.
(a) Train Speed (TFLOPs/s)
(b) Test Speed (TFLOPs/s)
40
80
FlashBias (Ours)
FlashBias (Ours)
FlashAltention
FlashArention
Forrers
Speed (IROP//
PyTorch SDPA
1.3x
Speed (TRLPPAA))
PyTorch PyTorch
PyTorch Vapilla
2x
20
2000
4000
4000
abco
10000
12500
14000
16000
2000 4000 6000 8000 10000 12000 14000 16000
Sequanoe Length
Length
Figure 8: Non-causal attention speed, where attention is with 4 heads and 32 hidden channels per head. is set as The speed is measured with the batch size as
Causal Attention In this part, we consider the causal attention. Following the experiments in Section 4.2, we configure this attention as a GPT-2 style, which contains 50 heads and 32 channels per head. presented in Figure g FlashBias achieves 1.75x and 1.3x speedup over previous implementations for the inference and training phases, respectively.
It is also notable that all the previous methods suffer from the out-of-memory issue. PyTorch SDPA and xFormer are more memory-efficient but still fail when the sequence length reaches 16,384. This because they have to save the entire bias matrix with shape of num_headss x seq_len seq__eee__eeee len in the GPU HBM, which is memory-prohibited in handling extremely long sequences. In contrast, FlashBias utilizes the inherent low-rank property of the bias term, thereby only needing to save the decomposed factors. This experiment also highlights the memory advantage of FlashBias.
(a) Train Speed (TFLOPs/s)
(b) Test Speed (TFLOPs/s)
FlashBias (Ours)
FlashBias (Ours)
60
FlashAttention
100
FlashAttention
1.3x
xFormers
Speed (TRIPPA))
45
Palorch Vanilla
Prrc SDPA
Speed (TFLOPU/))
75
PyTorch Vanille
ToToch SDPA
1.75x
25
2000
4000
6000
8000
10000
12000
14000
16000
2000
4000
8000 10000 12000 14000 16000
Sequenoe
Sequence Length
Figure 9: Causal attention speed, where attention with 50 heads and 32 hidden channels per head. R is set 2. Dotted lines indicate out-of-memory. Speed is measured with the batch size 1.
Torch implementation torch. wttech.nnn.uuna.c functinniiinn onal www.scaaaeecec coreddddd_pppreeee htttt_poouutttttttttttt Triton implementation of FlashAttention with bias.
22

---

C Speedup of Pangu-Weather Weather
Pangu-Weather [4] is significant step in adopting Transformers for global weather forecasting. Specifically, its backbone a 3D Swin Transformer with a hierarchical structure, which contains two different scales. Its speedup is similar to our experiments in Section 4.3) As listed in Table I we consider speedup this model based on the SVD decomposition version of FlashBias.
Setups Since Pangu-Weather Weather is based on the 3D window (with shape 2 x 6 x 12), its bias for relative position encoding is slightly different from Swin Transformer. Especially, its bias is in the shape of # num x heads x 144 144 for each block, where # num represents the number of 3D windows. According to metcorological knowledge, different longitudes share the same bias.
Implementations Since Pangu-Weather Weather does not provide accessible model weights code, our experiments are based on an open-sourced PyTorch reproduction All the other implementations are the same as the descriptions in Section 4.3] Notably, fine scale are low-rank. Thus, we only apply FlashBias to the four 3D Swin layers in the fine scale, where set R 56 maintain 99% energy of the original bias matrix. As discussed in Section [.33] FlexAttention fails in processing such dynamic bias; we didn't 't compare with it in this task. Since the whole ERA5 data is over 150TB, we only test the model based on 100 samples in 2024.
find that only relative position biases in the
Results As presented in Ta- 7 FlashBias can also speedup Compared to the open-sourced code, FlashBias re- duces over 20% running time and over 50% GPU memory usage. However, due to the limited se- quence length (N = 144 in this case), the running time speedup not as considerable as V2. Just plotted in Figure FlashBias will bring more sigmif- icant speedup in rumning time and GPU memory for long sequences. As improving the spatial resolution of forecasting is a golden problem in weather prediction [44), we believe that FlashBias has strong advantage in supporting future research on higher-resolution reanalysis data.
Table 7: Experiments of Pangu- Weather on ERAS 15 Output difference measures of the L.2 distance between the outputs of FlashBias and the official code, averaged from 100 different inputs. This difference is calculated on the z-score normalized (8 model outputs to balance diverse meteorological factors.
Method
Outpur Difference Time(s/100iters) Mem(MB)
Open- Open-sourced Code
98.022
26552
FlashAttention w/o bias
0.0128
74.089
12141
FlashAttention w/ bias
79.649
13186
and and
FlashBias (Ours)
0.0003
76.779
12222
D More Results of ALiBi Bias 30
As we stated in the main text, the released Attentittn supports an ALiBi_slopes feature, where only ALiBi slope values in the shape of # heads are loaded from HBM, and the specific bias in the shape of Block_o Bolockgg_Blockk k is created in just-in-time (JIT) compilation Here Block_q q and Block_q denote the tilling size in FlashAttention [10 9].
Note that ALiBi is quite simple bias, and the imple- mentation of FlashAttention's 's ALiBi_slopes feature takes advantage of that the bias val- ues can be created JIT. In the following experiments, we also utilized this property of ALiBi in FlashBias, where we generate the decomposed factor tensors of shape Flock_qx2 and Block_kx* in JIT, instead of directly creating the Block_ q <Booc__ k matrix. As presented in Table g since the difference between FlashAttention and FlashBias is in JIT, these two implementations almost at the same speed.
Table 8: Experiments on GPT-2 size model with ALiBi bias. Here we experiment with FlashAttention's 's ALiBi_ sl sl olopss feature and FlashBias with JIT speedup. The running time for 100 iterations at both training and inference phases when N = 2048 are recorded.
Method
Training Phase (s) Test Phase (s)
FlashAttention w/o bias Flash with ALiBi ALiBi_slopes
119.3
38.77
119.8
38.98
Flash with decomposition in JIT 119.8
38.98
/http:://githubb htttp://gihub..om//aaassssa hww//hhsshha////gggguuuooo Implementation of ALiBi_st ssoopes feature in FlashAttention
23

---

E More Results in AlphaFold 3
Analysis model efficiency The theoretical complexities of triangle self-attention and triangle multiplication in AlphaFold 3 are both cubic w.r.t. the sequence length, while other components: data embedding, single attention with pair bias and feedforward layer, are linear or quadratic complexity. Thus, according to theoretical complexity, these "triangle operations" are the bottleneck.
We further record the time cost of each component in AlphaFold [ during the inference of PDB ID 7wux. As listed in Table 9 the triangle attention accounts for 53.3% of the total inference time.
Table Running time of each component in AlphaFold 1 for inferring PDB ID 7wux.
Component
Complexity w.n.t. Sequence Length
Running Time (s)
Ratio
Data Embedding
Linear
1.91
7.1%
Triangle Sclf-attention
Cubic
14.32
53.3%
Triangle Multiplication
Cubic
9.97
37.1% %
Single Attention with Pair Bias
Quadratic
0.48
1.8%
FeedForward
Linear
0.17
0.7%
Potential in speeding up training of AlphaFold 3 [ In the main text. we only evaluate FlashBias during the inference phase. Going further, as mentioned the end of Section 3.2,, if directly replace the bias matrix in AlphaFold 3 [1] with two decomposed factor functions, we can also
accelerate the training process. As shown in Table TO, FlashBias can save around 15.2% running time and 17.7% GPU memory usage compared to the open-sourced code.
Table 10: Experiments of FlashBias in accelerating the training process of AlphaFold T where the default setting is to crop all the residue sequence within 384 tokens.
Method
Running Time (s 10 iters)
Maximum GPU Memory (GB)
Open-sourced Code
165
23.572
FlashAttention with bias
153
23.572
FlashBias (Ours)
140
19.390
However, since the complete training of AlphaFold 3 will require around 128 GPUs for 7 days, considering the resource limitation, we would like leave the verification of the final performance of FlashBias-acceleaatte AlphaFold 3 as our future work.
F More Results in Swin Transformer V2 24
Statistics of Swin Transformer V2 bias We visualize the rank of bias matrices that can remain 95% energy across 24 layers. As presented in Figure T0] the bias of later layers of lower rank. To reduce the accumulation error, we only apply Flash- Bias to the last layers described in Sec- tion 4.3 Specifically, for the last 8 layers, set R in FlashBias 16. In fact, ap- plying FlashBias to more layers will bring more significant speedup, while to maintain the model accuracy, we only apply Flash- Bias to the last 8 layers in this work.
576
Maximum Rank Value
Rank Value
Ronk Value
400
200
100
12
16
20
24
Layer
Figure 10: Maximum and mean rank value to remain
95% energy of bias matrices across different heads of 24 layers Swin
SVD decomposition visualization
Figure the main text, we visualize the original bias matrix at the 20th layer with 16 heads. To deliver an intuitive understanding, we also take this layer's bias as an example and compare with the SVD decomposed factor tensors' multiplication. As shown in Figure []] SVD decomposed factor tensors accurately reconstruct the original bias.

---

EEEERVE EYEERRU ETEEENT
Driginal ton
Figure 11: Comparison between SVD decomposed factor tensors' multiplication and original bias.
G More Results in Transformer PDE Solver
Table 5 we only present the efficiency comparison. Further, to demonstrate the performance benefits brought by spatial distance bias, we also include the performance metric in Table [T] With spatial distance bias, the error of the estimated drag coefficient can be reduced by 65.3%, which is a significant progress in industrial design. This further confirms the importance of attention with bias.
Table 11: Performance comparison among attention w/o bias and bias in PDE solving. The relative L2 of surface pressure and surrounding velocity is recorded. We also calculate the drag coefficient Cn based on predecere physics fields, whose relative L2 w.r.t. ground truth also included.
Method (Sequence Length N 32186)
Surface Pressure Eror Surrounding Velocity Error Cp Error
Pure attention without spatial distance bias
0.0838
0.0278
0,0173
FlashAttention with spatial distance bias FlashBias with spatial distance bias
OOM
OOM
OOM
0.0706
0.0201
0.01
Relative Promotion
15.7%
27.7%
65.3%
H Generalization for Diverse Biases
In FlashBias, we present a neural decomposition version to complex and dynamic biases, which has been tested in speeding up AlphaFold 3 in Section 4.4 To further demonstrate the expressive capability of neural decomposition, in this section, will train neural factor functions Φq.0, 01 to approximate more diverse biases, which can be meaningful for scientific tasks.
Neural Decomp
Gravity Bias
Neural Decomp
Sphericall Bias
(8) Gravity Bias
(b) Spherical Bias
Figure 12: Adopting neural decomposition techniques in FlashBias for more diverse biases.
Gravity bias Many phenomena are inherently governed by underlying physical forces, where gravity is one of the basic factors. Accurately approximating the gravity force is essential for the
25

---

modeling of planetary motion [3 or electronic simulation [3.. Thus, we consider introducing the gravity bias into the attention mechanism. Specifically, this bias can be formalized follows:
(xx,, xk.j)
(13)
||x Xk.j
where X denotes the spatial positions of the i-th and j-th objects. We train @q.0, Pk.02 based on randomly sampled points from [0, 1] x [0, 1] in the 2D space. Since the bias is inversely proportional to spatial distance, we further add 0.01 to the diagonal bias for numerical stability.
Xk.j
When analyzing atmospheric circulation, it is intuitive to consider it as the dynamics on spherical surface. Therefore, previous research has attempted introduce spherical Fourier analysis into the model design [5]. Another alternative approach is to add spherical distance bias to attention in Transformer-based models. Thus, we also consider the spherical distance bias:
Spherical distance
[Xq.s;; Xk.)) aresin (y (Vste?+ Xq.i00 Xk.j.0 + cos x.0 COS Xk.jo sin ( xq Xk.j. 1
2
(14)
where Xq.i, Xk.j E 2 records the latitude and longitude of the i-th and j-th position respectively. Similarly, we also train 0q.0,. Ōk.0, based on randomly sampled points in [-x, T] x [0, 2m
For the biases mentioned above, we set R 32 and 6q.01 ŌK.0, as three linear layers with in-between tanh activation functions. Then we optimize these model parameters with the Adam [9 optimizer for 10,000 iterations, which will take less than 30 seconds neural decomposition performs very well in these two biases, especially for the spherical bias. As for gravity bias, since the numerical instabilitty of the inverse proportion function, it is more difficult for optimization, while our method still captures the locality of the bias matrix.
A100 GPU. As presented in Figure [2]
The above experiments further testify the capability of FlashBias in accelerating broader scenarios.
I Implementation Details
As a supplement to the main text, we include more details of implementation here. All the experiments conducted based on Py 2.5.0 [27] and Triton 3.0.0 (38) a single A100 GPU with 144 CPU cores. All efficiency metrics are averaged from 1,000 iterations. For example, some metrics are recorded s/100iters or s/1Oiters, where we divide the running time of 1,000 iterations by 10 or 100, respectively. In our experiments, all the algorithms' efficiency performance is quite stable.
ALiBi in Large Language Model Since the experiments for ALiBi speedup are based on directly replacing the ALiBi bias with the exact decomposition, FlashBias's output results are completely equal to the original version. The training efficiency is tested under the Adam [9 optimizer.
Relative position encoding in Swin Transformer This part of the experiments strictly follows the official code in Swin Transformer V2 (24) As illustrated in Figure [0] we only apply FlashBias to the last 8 layers. It is also possible to apply FlashBias to part of low-rank heads each layer. However, this will require some preprocessing steps, such rearranging the heads and model parameters to ensure that all low-rank heads are memory continuous. Considering the implementation simplification, we eventually adopt the current design, which only speeds up the final layers.
Spatial distance bias in PDE solver In this experiment, we follow the code base provided here" which includes the application of Transformers for the car design task. Specifically, the input contains the position of the pre-defined computation mesh and the output is the pressure and air velocity on these computation mesh points. In this task, we also adopt the exact decomposition for computation. The training phase is also based on the Adam 19 optimizer with a batch size of 1.
Pair representation bias in AlphaFold 3 In AlphaFold 0. adopt the neural decomposition version. This version involves the training of newly added neural layers, whose finetuning configura- tion is listed in Table 12, Here we find that optimizing the factor functions for 10,000 iterations can obtain a nearly converged version. This may be because the factor functions are token-wise, which
26

---

Table 12: Configuration for finetuning factor functions 6g.0, 2x.02.
Part
Configuration
Sum of and column in pair representation, with shape of N 128
Input Xq: Xk
Single representation model beginning with shape of N x 449
Output 60.0, (Xq). Pq.0, (xxx Decomposed factor tensors with shape of N x # heads x x96 (# heads
Input Dim: 577: Hidden Dim: 256; Output Dim: 384 (= 4 x 96)
Model 0q.01 6x.02
Three linear layers with in-hetween tanh activation function
Initial learning rate: 0.001: Optimizer: Adam;
Training
Learning rate decay: every 50 iterations, reduce to the origin's 0.95
Overall steps: 10,000 iterations
Train set: SttphhhccddDD__hoooo22e2222 __op___wwpppp__p_ htt____opppppppaa_eeee nometale_ 0925" Test set: "recemtPDB_1555 15t__1_350mmm584_______ _0925"
Dataset
every iteration, the model will receive N (sequence length) samples for training. To up, these layers have been optimized with around 3,840,000 "samples" after 10,000 iterations. Also, to reduce the cumulative error, we only apply FlashBias to the first 16 Pairformer blocks.
J Extension to Multiplicative Bias
In this section, we will discuss the extension to multiplicative bias in the following formalization:
qk
softmax
b)v,
(15)
where b E N represents the multiplicative bias and C' denotes the number of hidden channels. denotes the Hadamard product, namely the element- wise multiplication.
Rotary position embedding (RoPE) 37
As mentioned in Section 2.1 RoPE is special multi-
plicative bias. Specifically, its pre-softmax attention weight is defined as follows:
C/2-1
Omn
Re(
Σ 1=0
(16)
where Omn represents the pre-softmax weight between m-th query and n-th key. q: represents the 2i and 2i + 1-th element of the m-th query representation. By comparing the above Eq. (15) and Eq. (6). we can obtain the following observations:
RIx2
(i) RoPE does not follow the common formalization of multiplicatiie bias. As formalized in Eq. (6). RoPE does not directly multiply a bias weight to each attention value, which also involves a detailed reweighing along the channel dimension. Thus, we prefer to consider RoPE as a unique technique, not a "general formalization" for multiplicative bias.
(ii) RoPE is already FlashAntention friendly and does not need further modifications to fit the computation of FlashArtention. According to Eq. (34) in its official paper (371. RoPE can be accomplished by multiplying the rotary tensors with q and k before computing qk' With this design, it does not need load additional N x N bias matrix; thereby, it can be seamlessly integrated with FlashAttention. Thus, we do not think it needs further modifications to fit the computation flow of FlashAttention. That is also why we do not experiment with RoPE in this paper.
General multiplicative bias Here, we will demonstrate that the low-rank decomposition proposed in FlashBias can be extended to multiplicative bias defined in Eq. (5)).
Specifically, suppose that b E NRN can be decomposed as the multiplication of two rank It tensors: og and dk e R N f We can rewrite the calculation of attention with multiplicative bias as follows:
qk T
② b)y = softmax q'k
o = softmax
jv.
VC
VC
where q' = [q 2 "q.1""
② [q.1]
E NXCR k' = k ② @1,1," ... , @ @K.R) E NXCR
(17)
27

---

Here Qq Ok E RNXI represents the i-th channel of the decomposed factor tensors Qg and ok- whose channel dimension will be broadcasted C' times before multiply to q or k. Operation (... .... denotes the concatenation along the channel dimension. Besides, the above computation also requires repeating the original q and along the channel dimension for R times and multiplying the factor tensors for each repeat. It is easy to verify the correctness of the above reformalization.
Example J.1. Given the multiplicative bias bij- cos(i-j], which can also be viewed as simplified version of RoPE (Eq. (16), b can be decomposed into two factor tensors with R = 2:
Φ (x) = [cos(1), sin( E RIIx, sr (xx,) = [[cos()), sin(()) E RRIX2
(18)
Based on this decomposition, we can use the above Eq. {17) speedup, where the whole N x N bias matrix is transformed into the reweighted repeat of q qqand. k.
Corollary J.2 (Efficiency analysis). FlashBias's 's extension to attention with multiplicative bias
(Eg. 15 can reduce the storage complexity of vanilla FlashAttention when R S
V
C
+ 1.
Proof. According to (9) the HBM access of FlashAttention with bias is 6 NMCC²
NMM. it easy
5
to drive that the HBM access of Eq. 17 is O
NMCCRR) Thus, when R S
S
+ 1, we have:
R
( C + 0 NM).
NMC
(19)
S
O
Example J.3 (Design practice). Given regular setting C = 64 and S = 100KB, FlashBias can achieve acceleration for attention with multiplicative bias if there exists a decomposition with R S 27.
Since the multiplicative bias is not as common as the additive bias and RoPE has already been well implemented, we would like to leave more explorations of multiplicative bias as future work.
K Limitations
As mentioned in the main text, FlashBias's speedup is based on the low-rank nature of attention bias. Although the low-rank biases are observed in many well-established backbones, if the bias matrix is not of low rank, FlashBias will have to trade off between efficiency and final performance.
About this limitation, one possible solution is to introduce an additional sparsity matrix as a sup- plement to low-rank decomposition, which is well-known matrix approximation technique [7]. Specifically, given the attention bias matrix b, this method formalizes the decomposition process as a convex relaxation problem, which can be formalized as follows:
f + t arg min ||||| + /||t|/1 r.t
(20)
s.t. +artt=h b,
where ||]|. represents the nuclear norm and is expected to be low-rank and is optimized to be sparse. Afterwards, the low-rank matrix r can be sped up through FlashBias, and we can adopt the same technique as attention mask speedup (35] for accelerating the sparse component t. Further, this technique can also make up for the current shortage of FlashBias in approximating the diagonal value (Figure 7 where we can decompose the bias into the sum of low-rank and diagonal matrix.
Since we mainly focus on the low-rank property of attention bias in this paper and it will not affect the final performance in our experiments, would like to leave this as our future work.
L Broader Impacts
This paper presents FlashBias as practical method enable fast computation of attention with bias, which shows significant efficiency advantage over FlashAttention and FlexAttention [[]] FlashBias also enables 1.5x speedup for Pairformer in AlphaFold 3 [[] and over 2x speedup in vision and language models without performance drop. These efficiency benefits can facilitate future fine-tuning or development of Transformers in yeal-world applications. Also, this paper considers a new low-rank compressed sensing perspective in fast attention computation, which can be inspiring for future research. Since we purely focus on the algorithm design for fast computation of attention with bias, there are potential negative social impacts or ethical risks.
28


Reviews:
[{"review_id": "DB15XJVpOv", "reviewer": "Reviewer_xU4x", "summary": "This paper addresses the inefficiency of computing attention mechanisms with additive bias terms, which are essential for performance but suffer from quadratic I/O complexity in existing methods like FlashAttention. The proposed method based on low-rank compressed sensing that decomposes large bias matrices into smaller factor tensors. It uses exact decomposition for known low-rank biases, SVD for fixed learned biases, and neural decomposition for complex dynamic biases. Experimental results showed their method achieves substantial speedups—1.5× for AlphaFold and over 2× for vision/language models—while maintaining accuracy. It also drastically reduces GPU memory usage, enabling training on larger-scale cases.", "questions": "1. For dynamic biases (Section 3.3), what is the latency/memory cost of the lightweight networks during inference?  \n2. In mainstream models, which bias matrices are not of low rank? A study on corner cases would be beneficial.\n3. Could FlashBias dynamically adjust the decomposition rank \\(R\\) per layer/head (e.g., based on singular value thresholds)?", "limitations": "The authors have discussed most of the limitations of their work and provide the potential solutions. However, it would be beneficial to include more examples or studies related to bias matrices that are not of low rank.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. Proposed Flashbias provides 1.5–2× speedups in AlphaFold, Swin Transformer, and GPT-2 with ALiBi as well as memory saving.\n2. The paper is well-organized and clearly written. The figures and tables are illustrative and effectively support the overall narrative.\n3. Focus on the speedup of LLM with bias and presented a novel method.\n\nWeaknesses: \n1. Speedups are reported, but energy consumption (e.g., TDP reduction on A100) is unmeasured.\n2. In Section 4.4, \"After a detailed analysis of AlphaFold 3, we find that its efficiency bottleneck is triangle self-attention,\" there is no data to support this claim.\n2. Some of the images are too small to be easily readable.", "was_revised": true, "final_justification": "The Flashbias method shows significant speedups and memory savings in AlphaFold, Swin Transformer, and GPT-2. The paper is well-organized, clearly written, and supported by strong data. Recommended for acceptance.", "rebuttal": "We would like to sincerely thank Reviewer xU4x for providing a detailed review and insightful questions.\n\n> **Q1:** \"Speedups are reported, but energy consumption (e.g., TDP reduction on A100) is unmeasured.\"\n\nFollowing the reviewer's suggestion, we further tested the inference power cost of Swin Transformer V2. Since the TDP (Thermal Design Power) is usually predefined in the A100, which is 400W in our device, we record the power draw during experiments.\n\nAs shown below, compared to the official code of Swin V2, FlashBias can slightly reduce the average and maximum power draw. Benefitting from a significantly faster running speed, **FlashBias can reduce power costs by 61% (123.6$\\to$48.2) per batch.**\n\n| SwinV2 Inference | Avg Power Draw (W) | Max Power Draw (W) | Running Time (s/batch) | Power Cost per Batch (J) |\n| - | - | - | - | - |\n| Original Code    | 261.3              | 290.2              | 0.473                  | 123.6                    |\n| FlashBias        | **253.5**          | **285.4**          | **0.190**              | **48.2**                 |\n\n> **Q2:** There is no data to support \"After a detailed analysis of AlphaFold 3, we find that its efficiency bottleneck is triangle self-attention.\"\n\nThe theoretical complexities of triangle self-attention and triangle multiplication in AlphaFold 3 are both cubic w.r.t. the sequence length, while other components are of linear or quadratic complexity. Thus, in the theoretical complexity aspect, these \"triangle operations\" are the efficiency bottleneck.\n\nAs per your request, we have further recorded the time cost of each component in AlphaFold 3 during the inference of PDB ID 7wux. As shown below, the **triangle self-attention accounts for 53.5% of the total inference time.** For scientific rigor, we will revise this statement as *\"its efficiency bottleneck is 'triangle operations', where triangle self-attention is one of them.\"*\n\n| Vanilla Pairformer's Components | Theoretical Complexity w.r.t. Sequence length | Running Time (s) | Ratio of Whole Model |\n| - | - | - | - |\n| Data Embedding                  | Linear                                        | 2.26             | 7.2% |\n| **Triangle Self-attention**     | **Cubic**                                     | **16.83**        | **53.5%** |\n| Triangle Multiplication         | Cubic                                         | 11.68            | 37.1% |\n| Single Attention with Pair Bias | Quadratic                                     | 0.52             | 1.6% |\n| FeedForward                     | Linear                                        | 0.19             | 0.6% |\n\n> **Q3:** \"Some of the images are too small to be easily readable.\"\n\nThanks for the reviewer's valuable suggestion. Since we cannot update the submission during the rebuttal stage, we will enlarge $\\underline{\\text{Figs 3 and 4}}$ and the color bar for figures in the Appendix in the future revision.\n\n> **Q4:** \"For dynamic biases, what is the latency/memory cost of the lightweight networks during inference?\"\n\nSince \"lightweight networks\" used in FlashBias are just simple linear layers, their computation are quite efficient. \n\nFor example, experiments in AlphaFold 3 ($\\underline{\\text{Table 6  of main text}}$) are based on learnable lightweight networks. Even though FlashBias adds two newly linear layers to each Pairformer block, our method still enables a 1.5x speedup over the open-sourced code. Specifically, for the inference of PDB ID 7wux, the newly added \"lightweight networks\" will only account for around 0.2% (0.0417s vs. 20.89s) of the total running time and 3.5% of the total memory cost.\n\n| Inferring PDB ID 7wux             | Time (s) | Mem (GB) |\n| - | - | - |\n| Open-Sourced Code                 | 31.48    | 21.58    |\n| FlashBias                         | 20.89    | 10.37    |\n| Lightweight Networks in FlashBias | 0.0417   | 0.37     |\n\n> **Q5:** \"In mainstream models, which bias matrices are not of low rank? A study on corner cases would be beneficial.\" \"However, it would be beneficial to include more examples or studies related to bias matrices that are not of low rank.\"\n\nAs presented in $\\underline{\\text{Fig. 7 of Appendix D}}$, not all the relative position bias matrices in SwinV2 are of low rank, but we can select the low-rank layers or heads and apply FlashBias on these low-rank layers/heads for speedup.\n\nBesides, we also want to highlight that, as proved by $\\underline{\\text{Theorem 3.2}}$, the rank of the bias matrix determines the optimal storage complexity. Thus, FlashBias does not attempt to speed up all the mainstream models by low-rank decomposition, but presents a possible way to reach the optimal efficiency when the bias matrix presents a low-rank property. We will include these discussions in the $\\underline{\\text{Limitation section in Appendix H}}$.\n\n> **Q6:** \"Could FlashBias dynamically adjust the decomposition rank (R) per layer/head (e.g., based on singular value thresholds)?\"\n\nFor the SVD decomposition version, the rank can be adjusted for different layers/heads. For instance, in SwinV2 experiments, the selected rank values of different layers and heads are different, as stated in $\\underline{\\text{Appendix D}}$. \n\nIf you refer to the dynamic biases, such as AlphaFold 3, I do not think online SVD is an efficient design, since we need to compute SVD for every inference, as stated in $\\underline{\\text{Lines 184-186 of main text}}$. That is why we present the neural decomposition version in FlashBias for dynamic biases."}, {"review_id": "JrTx3xFGFy", "reviewer": "Reviewer_BRmd", "summary": "An additive bias term of attention is widely used to guide the learning process. This paper improves the efficiency of bias computation by realizing its low rank nature. Specifically, one doesn't need to store an $N \\times N$ bias matrix. Instead, the bias matrix is represented by the product of two $N \\times r$ matrices. Each element of the bias matrix is computed in-place of GPU high bandwidth memory which is similar to the characteristic of FlashAttention. Thus, the proposed method is called FlashBias. FlashBias achieves 1.5x speedup for AlphaFold, and over 2x speedup for attention with bias in vision and language models without loss of accuracy.", "questions": "In figure 3 and figure 4, which kind of bias is used?", "limitations": "The proposed method is not applicable to non-additive bias, such as rotary position embedding (RoPE).", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The proposed method is easy to implement and achieves significant speedup for attention with bias.\n2. The proposed method is verified by various tasks with various bias types, including spatial bias, alibi bias, and pair bias. This paper designs different strategies to obtain the low-rank representations of bias.\n3. This paper provides theoretical analysis of the IO complexity of bias. Its efficiency is supported by the theoretical results.\n\nWeaknesses:\n1. If my understanding is right, this method is applicable to additive bias but it doesn't work for non-additive bias, such as rotary position embedding (RoPE). However, the authors claim that the proposed method provide a fast-accurate approximation for biases in general formalization. It is better to discuss whether RoPE can be approximated by low-rank representations.\n2. The results of alibi bias, i.e. table 3 are not convincing. To my knowledge, the alibi bias has been optimized by FlashAttention in a similar way (although they doesn't formalize alibi bias as a low-rank representation). I believe that the improvements in table 3 have other reasons.", "was_revised": true, "final_justification": "All of my concerns have been addressed. So I improve my score from 4 to 5. First, the authors provide detailed analysis about RoPE and they show that FlashBias can be extended to multiplicative bias. Second, the authors update the experiments about alibi bias and they add `ALiBi_slopes` feature for FlashAttention.", "rebuttal": "We would like to sincerely thank Reviewer BRmd for providing valuable feedback and questions.\n\n> **Q1:** \"This method is applicable to additive bias, but it doesn't work for non-additive bias, such as rotary position embedding (RoPE).\" \"However, the authors claim that the proposed method provides a fast-accurate approximation for biases in general formalization. It is better to discuss whether RoPE can be approximated by low-rank representations.\"\n\nThanks for the reviewer's thoughtful question. We will revise \"a fast-accurate approximation for biases in general formalization\" as **\"for additive biases in general formalization\"** in the future revision for scientific rigor.\n\nNext, following the reviewer's comments, we would like to discuss more about RoPE and multiplicative bias in the following formalizations:\n\n$$\\mathbf{o}=\\operatorname{softmax}(\\frac{\\mathbf{q}\\mathbf{k}^\\top}{\\sqrt{C}}\\mathbf{b})\\mathbf{v},\\ \\ \\ (1)$$\n\nwhere $\\mathbf{b}\\in\\mathbb{R}^{N\\times N}$ represents the multiplicative bias and $C$ denotes the number of hidden channels.\n\n**(1) RoPE does not follow the formalization in the above Eq. (1) and is already FlashAttention-friendly.**\n\nOur paper identifies the efficiency bottleneck in computing attention with additive bias as: **they have to additionally load a $N\\times N$ bias matrix from HBM to SRAM**. However, in RoPE, according to their paper, its pre-softmax attention weight is defined as follows:\n\n$$a_{mn}=\\text{Re}[\\sum_{i=0}^{C/2-1}\\bar{\\mathbf{q}}_i\\bar{\\mathbf{k}}_ie^{i(m-n)\\theta_i}],\\ \\ (2)$$\n\nwhere $a_{mn}$ represents the pre-softmax weight between $m$-th query and $n$-th key. $\\bar{\\mathbf{q}}_i\\in\\mathbb{R}^{1\\times 2}$ represents the $2i$ and $2i+1$-th element of the $m$-th query representation.\n\nBy comparing the above Eq.(1) and Eq.(2), we can obtain the following observations:\n\n- **RoPE does not follow the common formalization of multiplicative bias ($\\underline{\\text{Eq. (1) above}}$).** As formalized in $\\underline{\\text{above Eq. (2)}}$, RoPE does not directly multiply a bias weight to each attention value, which also involves a detailed reweighing along the channel dimension. Thus, we prefer to consider RoPE as a unique technique, not a \"general formalization\" for multiplicative bias.\n- **RoPE is already FlashAttention-friendly and does not need further modifications to fit the computation of FlashAttention.** According to $\\underline{\\text{Eq.(34) in their paper}}$, RoPE can be accomplished by multiplying the rotary tensors with $\\mathbf{q}$ and $\\mathbf{k}$ before computing $\\mathbf{q}\\mathbf{k}^\\top$. With this design, it does not need to load an additional $N\\times N$ bias matrix; thereby, it can be seamlessly integrated with FlashAttention. Thus, we do not think it needs further modifications to fit FlashAttention.\n\nConsidering the special definition and native efficiency of RoPE, we did not discuss it in the original submission, but will add the above discussion as a separate section in a future revision.\n\n**(2) FlashBias can be extended to multiplicative bias formalized in the above Eq. (1).**\n\nNext, we will show that the low-rank decomposition proposed in FlashBias can be extended to multiplicative bias defined in $\\underline{\\text{Eq. (1) formalized above}}$. \n\n**Overall design:** Specifically, suppose that $\\mathbf{b}\\in\\mathbb{R}^{N\\times N}$ can be decomposed as the multiplication of two rank-$R$ tensors: $\\phi_{\\mathbf{q}}$ and $\\phi_{\\mathbf{k}}\\in\\mathbb{R}^{N\\times R}.$ We can rewrite the calculation of attention with multiplicative bias as follows:\n\n$$\\operatorname{softmax}(\\frac{\\mathbf{q}\\mathbf{k}^\\top}{\\sqrt{C}}\\mathbf{b})\\mathbf{v}=\\operatorname{softmax}(\\frac{\\mathbf{q}^\\prime{\\mathbf{k}^\\prime}^\\top}{\\sqrt{C}})\\mathbf{v},$$\n\n$\\text{where}\\ \\mathbf{q}^\\prime=\\sqrt{C}[\\mathbf{q}\\phi_{\\mathbf{q},1},\\cdots,\\mathbf{q}\\phi_{\\mathbf{q},R}]\\in\\mathbb{R}^{N\\times CR},\\  \\mathbf{k}^\\prime=[\\mathbf{k}\\phi_{\\mathbf{k},1},\\cdots,\\mathbf{k}\\phi_{\\mathbf{k},R}]\\in\\mathbb{R}^{N\\times CR}. \\ \\ \\ (3)$ \n\n$\\phi_{\\mathbf{q},i},\\phi_{\\mathbf{k},i}\\in\\mathbb{R}^{N\\times 1}$ represents the $i$-th channel of $\\phi_{\\mathbf{q}}$ and $\\phi_{\\mathbf{k}}$. \nThe above computation requires repeating the original $\\mathbf{q}$ and $\\mathbf{k}$ along the channel dimension for $R$ times and multiplying the factor tensors for each repeat. It is easy to verify the correctness of the above reformalization. And, the above formalization can reduce storage complexity of vanilla multiplicative bias when $R\\leq 1+\\frac{N}{2C}$. \n\n**Example:** Given the multiplicative bias $\\mathbf{b}_{ij}=\\cos(i-j)$, which can also be viewed as a simplified version of RoPE, $\\mathbf{b}$ can be decomposed as $R=2$ with \n\n$$\\phi_{\\mathbf{q}}(\\mathbf{x}_{\\mathbf{q},i})=[\\cos(i), \\sin(i)]\\in\\mathbb{R}^{1\\times 2}$$\n\n$$\\phi_{\\mathbf{k}}(\\mathbf{x}_{\\mathbf{k},j})=[\\cos(j), \\sin(j)]\\in\\mathbb{R}^{1\\times 2}.$$\n\n Thus, we can use the above $\\underline{\\text{Eq. (3)}}$ for speedup, where the whole $N\\times N$ bias matrix is transformed into the reweighted repeat of $\\mathbf{q}$ and $\\mathbf{k}$.\n\nSince the multiplicative bias is not as common as the additive bias and RoPE has already been well implemented, we would like to leave more explorations as our future work.\n\n> **Q2:** \"The results of alibi bias, i.e. table 3 are not convincing. To my knowledge, the alibi bias has been optimized by FlashAttention in a similar way (although they doesn't formalize alibi bias as a low-rank representation). I believe that the improvements in table 3 have other reasons.\"\n\nThanks for the detailed review and insightful question. \n\n**(1) Clarification about Table 3.**\n\nAfter carefully checking the implementation of FlashAttention and contacting their authors, we find that our experiments in ALiBi are misled by their official Triton implementation examples. During our initial submission, we found that the official Triton implementation example of FlashAttention with ALiBi is wrong (see Line 493 in `flash-attention/blob/main/flash_attn/flash_attn_triton.py` of the official GitHub repository; they should not directly broadcast the bias vector to the matrix). Thus, we mistakenly thought the current FlashAttention could not support ALiBi well.\n\nThis drove us to avoid using the ALiBi speedup feature in FlashAttention. Therefore, the FlashAttention results in Table 3 are still based on the full matrix-type bias. That is why FlashBias can reduce the time for processing bias by around 50%.\n\n**(2) New comparison with FlashAttention and FlashBias.**\n\nThanks for the reviewer's professional reminder. We find that the officially released FlashAttention is not like their Triton examples. Specifically, the released FlashAttention supports an `ALiBi_slopes` feature, where only ALiBi slope values in the shape of #$\\operatorname{heads}$ are loaded from HBM, and the specific bias in the shape of Block_q $\\times$ Block_k is created in Jit. Here Block_q and Block_k denote the tilling size in FlashAttention.\n\nNote that ALiBi is quite a simple bias, and the implementation of FlashAttention's `ALiBi_slopes` feature takes advantage of that the bias values can be created in Jit. In the following experiments, we also utilized this property of ALiBi in FlashBias, where **we generate the decomposed factor tensors of shape Block_q$\\times$2 and Block_k$\\times$2 in Jit**, instead of directly creating the Block_q $\\times$ Block_k matrix. However, since the difference between FlashAttention and FlashBias is in Jit, these two implementations are almost at the same speed.\n\n|                                            | Training (s) | Inference (s) |\n| - | - | - |\n| FlashAttention without Bias                | 119.3        | 38.77         |\n| FlashAttention with `ALiBi_slopes` feature | 119.8        | 38.98         |\n| FlashBias with decomposition in Jit        | 119.8        | 38.98         |\n\nFor clarity, we will include all the above discussions and the specific usage of FlashAttention in the revised paper.\n\nAt last, we would like to highlight that beyond ALiBi, FlashBias is much more useful in processing complex biases, such as SwinV2 or AlphaFold 3. We believe this generality of FlashBias is meaningful to the research community.\n\n> **Q3:** \"In figure 3 and figure 4, which kind of bias is used?\"\n\nAs stated in $\\underline{\\text{Line 219 of main text}}$, Figures 3 and 4 are just for an \"overall efficiency comparison\". Thus, these experiments do not relate to any specific type of bias. These experiments are used to compare the model efficiency between loading a #$\\operatorname{heads}\\times N\\times N$ bias matrix and only loading two factor vectors of shape #$\\operatorname{heads}\\times N\\times R$."}, {"review_id": "2HmgTXcQkc", "reviewer": "Reviewer_Bp2x", "summary": "This paper addresses a critical and previously overlooked efficiency bottleneck in modern deep learning: the computation of attention mechanisms that include an additive bias term. While attention with bias is essential for the performance of many state-of-the-art models in domains like vision (e.g., Swin Transformer) and biology (e.g., AlphaFold), its efficient computation has been largely unaddressed. The author finds that the primary bottleneck being the quadratic I/O overhead of reading the large, dense bias matrix from GPU memory.\n\nThe core contribution of this paper is that the author finds that the attention bias is low-rank. Based on this, they propose a method FlashBias, which performs an offline low-rank decomposition for the bias, and integrates them into the QK computation. FlashBias achieves comparable performance with obvsious decreasing of the computational time and memory.", "questions": "Have you performed any experiments to validate that training a model like Swin Transformer from scratch with factorized bias parameters can achieve the same final accuracy as the original model trained with a full-rank bias?", "limitations": "See weakness part.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "[Strength]\n1.This paper presents a clear and coherent logical flow—from identifying the limitations, to proposing FlashBias as a solution, and finally validating it with experimental results that strongly support the initial motivation. The overall reasoning is well-structured and complete.\n2.The core theoretical insight of the paper is novel and elegant. Instead of brute-force engineering, the authors reframe the I/O problem through the lens of compressed sensing, leading to the key hypothesis that the efficiency is determined by the rank of the bias matrix. This perspective is original and provides a strong theoretical foundation for their method. \n3.The proposed method, FlashBias, is simple, effective, and broadly applicable. The evaluation is exceptionally comprehensive and convincing. \n\n[Weakness]\n1.The paper's scope might seem limited as attention bias is not universally used. \n2.While the paper shows minimal performance loss at the chosen rank R, the inherent trade-off between rank, approximation error, and model accuracy warrants a more systematic analysis. Such a study would enhance the robustness of the results and offer clearer guidance for applying FlashBias in practice.\n3.The paper shows strong inference-time speedups, but training-time benefits are less validated. For example, while factorizing bias terms in models like Swin Transformer is promising, there are no results showing that training from scratch achieves similar accuracy, leaving training-time gains uncertain.", "was_revised": true, "final_justification": "The author's rebuttal is thorough, and the experimental results are solid, prompting me to assign a higher score. However, since Attention Bias is not a universally applicable setting, this remains the only limitation of the paper, preventing me from giving it a strong accept.", "rebuttal": "Many thanks to Reviewer Bp2x for providing the insightful review and questions.\n\n> **Q1:** \"The paper's scope might seem limited as attention bias is not universally used.\"\n\nWe agree with the reviewer's opinion that \"attention bias is not universally used in previous research\". But we would also want to point out that **the poor efficiency in computing attention with bias may be one of the reasons for this not-universally-used situation.**\n\nAs presented in our experiments of \"**Transformer-based PDE solvers** ($\\underline{\\text{Table 5 and Table 9}}$)\", spatial bias is significantly helpful in PDE solving, while to our best knowledge, none of the previous Transformer-based neural PDE solvers adopt this important bias. This may be because directly incorporating spatial bias into attention will cause out-of-memory error as presented in $\\underline{\\text{Table 5 of main text}}$. We believe that FlashBias can promote the future usage of attention bias, which has already been proven useful in many tasks.\n\n> **Q2:** \"While the paper shows minimal performance loss at the chosen rank R, the inherent trade-off between rank, approximation error, and model accuracy warrants a more systematic analysis. Such a study would enhance the robustness of the results and offer clearer guidance for applying FlashBias in practice.\"\n\nThanks for this valuable suggestion.\n\n**(1) Rank $R$ represents the lower bound that will not lose accuracy.**\n\nSince the **basic principle** in building fast computation methods of attention mechanism is to maintain its original performance, all the experiments in our original submission are based on the spirit of \"do not affect the model accuracy\".\n\nThus, we prefer to view the \"rank\" $R$ used in our paper as the lower bound that does not affect accuracy, not a hyperparameter that needs to be tuned for trading off performance and efficiency.\n\n**(2) More experiments about the performance-efficiency trade-off.**\n\nFollowing the reviewer's suggestion, we also tested FlashBias in SwinV2 with different rank values. Note that in our experiments, we only adopt FlashBias to **half of the attention heads** in SwinV2 to maintain 95% energy of bias matrices. More implementation details can be found in $\\underline{\\text{Appendix D}}$.\n\nAs shown below, decreasing $R$ will cause a consistent performance drop but bring efficiency advantages. Note that, for the $R=16$ case, as shown in $\\underline{\\text{Figure 7 in Appendix D}}$, this setting will cause some of the heads to not maintain 95% energy, thereby causing a more serious performance drop. However, even in this setting, FlashBias still achieves better accuracy than quantization techniques INT8 PTQ and brings a more significant speedup (over 60%) over their official code.\n\n| SwinV2                                                       | Running Time (s/batch)                | Top1 Acc |\n| - | - | - |\n| Official Code                                                | 0.473                                 | 87.144%  |\n| FlashBias R=96                                               | 0.210                                 | 87.142%  |\n| FlashBias R=64                                               | 0.198                                 | 87.137%  |\n| FlashBias R=32 (Default setting, see $\\underline{\\text{Appendix D}}$) | 0.190                                 | 87.126%  |\n| FlashBias R=16 (Cannot consistently maintain 95% energy)     | 0.188                                 | 86.982%  |\n| INT8 PTQ                                                     | Around 22% speed up over official code | 86.46%   |\n\n> **Q3:** \"The paper shows strong inference-time speedups, but training-time benefits are less validated. For example, while factorizing bias terms in models like Swin Transformer is promising, there are no results showing that training from scratch achieves similar accuracy, leaving training-time gains uncertain.\" \"Have you performed any experiments to validate that training a model like Swin Transformer from scratch with factorized bias parameters?\"\n\nThanks for this insightful question.\n\n**(1) Clarify the verified effective scope of FlashBias.**\n\nFirst, we want to highlight that **FlashBias attempts to ensure exact or nearly-exact computation of attention with bias.** Since in SVD and neural decomposition, \"replacing bias with factorized parameters\" will change the model architecture, which will also affect the backward properties of the model, we only test the inference speedups of SwinV2 and AlphaFold 3, as stated in $\\underline{\\text{Table 2 of main text}}$. Specifically, as discussed in $\\underline{\\text{Lines 205-212 of main text}}$, the verified effective scope of FlashBias can be summarized below. **Note that speeding up the inference of widely deployed models (e.g. Swin Transformer and AlphaFold 3) is already very significant.**\n\n| Decomposition | Verified Effective Scope |\n| - | - |\n| Exact         | Training & Inference |\n| SVD           | Inference            |\n| Neural        | Inference            |\n\nFor clarity, we will highlight that **the effectiveness of SVD and neural decomposition for the training phase is under exploration** in the $\\underline{\\text{Limitation and Future Work}}$ section of the revised paper.\n\n**(2) New experiments on training Swin V2 with factorized bias.**\n\nSince both SwinV2 Large and AlphaFold 3 experimented in our paper will take massive computation costs (around 32 A100 GPUs for several days), **we made a quick test on training SwinV1 Tiny, whose training costs 8 GPUs for 2 days.** Here are the results. Since the sequence length is quite small (seq_len = 49), the efficiency advantage is not significant, but this experiment can be used to test the training performance of FlashBias. Following our findings in $\\underline{\\text{Table 7 of Appendix D}}$, we only replace biases in half of the attention heads at each layer with low-rank factors.\n\nAs shown below, training with low-rank bias will not affect the final performance, where the Top1 Acc difference (0.09%) is under the standard deviation (*note: the Top 1 Acc difference on ImageNet within 0.2% is usually treated as normal performance fluctuation*).\n\nDue to the time and resource limitations, we cannot make more large-scale experiments at the rebuttal stage. Considering the scope of this paper, we would like to leave the exploration of FlashBias training performance as our future work.\n\n| SwinV1-Tiny (Attn `seq_len = 49`) | Top1 ACC | Training Time | Training Memory |\n| - | - | - | - |\n| Official Code | 81.29 | 0.3372        | 9191MB          |\n| FlashBias (R=8)                 | 81.20 | 0.3303        | 8892MB          |"}, {"review_id": "fhgmOLMaDg", "reviewer": "Reviewer_UJUf", "summary": "FlashBias addresses the long-standing efficiency gap for attention with bias by leveraging low-rank properties, providing a generalizable, high-performance solution across diverse domains. It seamlessly integrates with existing Transformers, enabling faster training/inference without compromising accuracy, and paves the way for scaling complex bias-aware models.", "questions": "- Can the authors provide more comprehensive experiments across a wider variety of architectures and tasks to demonstrate the generality of FlashBias?\n- How easy is it to integrate FlashBias into existing model architectures? Are there plans or existing implementations?", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n- Rigorous theory links bias rank to GPU efficiency; comprehensive cross-domain experiments validate speedups without accuracy loss.\n- Fills a critical gap in bias-aware attention optimization, enabling 1.5×–2× speedups in AlphaFold, Swin Transformer, etc., with real-world utility.\n- Novel application of low-rank compressed sensing to attention bias; neural decomposition for dynamic biases offers unique adaptability.\n\nWeaknesses\n- While the paper reports significant speedups, it would benefit from a broader evaluation across diverse tasks and models to establish the generality of FlashBias.", "was_revised": true, "final_justification": "The authors effectively addressed key concerns. Expanded experiments (Sandwich bias in GPT-2, PatchTST for time series) plus initial tests on 5 backbones confirm FlashBias’s generality. Clarifications on integration and explanations for limited applicability to Qwen/Llama etc. are reasonable. No critical unresolved issues. Its strengths—efficiency, usability, cross-domain value—outweigh minor gaps, justifying a positive recommendation.", "rebuttal": "Many thanks to Reviewer UJUf for providing a detailed review and insightful suggestions.\n\n> **Q1:** \"While the paper reports significant speedups, it would benefit from a broader evaluation across diverse tasks and models to establish the generality of FlashBias.\" \"Can the authors provide more comprehensive experiments across a wider variety of architectures and tasks to demonstrate the generality of FlashBias?\"\n\nIn the original submission, we have tested FlashBias in **5 well-established backbones, such as GPT-2, Swin Transformer, Pangu-Weather and AlphaFold 3, covering language, vision and science**. We think these experiments can provide strong support for the generality of FlashBias.\n\nFollowing the reviewer's suggestion, we further evaluated FlashBias in speeding up the **Sandwich bias** [1], which was originally proposed for natural language processing. Here are the experiment details.\n\n[1] Chi et al., Dissecting Transformer Length Extrapolation via the Lens of Receptive Field Analysis, ACL 2023\n\n**(1) Sandwich bias:** Specifically, Sandwich bias matrix $\\mathbf{b}\\in\\mathbb{R}^{N\\times N}$ is defined as follows:\n\n$$b_{ij}=\\sum_{k=1}^{C/2}\\cos(\\frac{i-j}{10000^{2k/C}}),$$\n\nwhere $C$ represents the number of channels of deep representation, such as queries, keys and values.\n\n**(2) FlashBias design:** Sandwich bias described above can also be sped up through the \"Exact decomposition\" version of FlashBias, where the decomposed factor functions $\\phi_{\\mathbf{q}},\\phi_{\\mathbf{k}}$ are defined as follows:\n\n$$\\phi_{\\mathbf{q}}(\\mathbf{x}_{\\mathbf{q},i})=[\\cos(\\frac{i}{10000^{2/c}}),\\sin(\\frac{i}{10000^{2/c}}),\\cdots,\\cos(\\frac{i}{10000^{2(\\frac{C}{2})/c}}),\\sin(\\frac{i}{10000^{2(\\frac{C}{2})/c}})]\\in\\mathbb{R}^{1\\times C}$$\n\n$$\\phi_{\\mathbf{k}}(\\mathbf{x}_{\\mathbf{k},j})=[\\cos(\\frac{j}{10000^{2/c}}),\\sin(\\frac{j}{10000^{2/c}}),\\cdots,\\cos(\\frac{j}{10000^{2(\\frac{C}{2})/c}}),\\sin(\\frac{j}{10000^{2(\\frac{C}{2})/c}})]\\in\\mathbb{R}^{1\\times C}.$$\n\nWith simple trigonometric function knowledge, it is easy to verify that based on the above definition of factor functions, $\\mathbf{b}=\\phi_{\\mathbf{q}}\\phi_{\\mathbf{k}}^\\top$.\n\n**(3) Efficiency comparison:** The efficiency comparisons are listed below. For this type of attention bias, we also examine the speed advantage of FlashBias in GPT-2. As shown below, FlashBias can also achieve a significant efficiency promotion compared to the vanilla FlashAttention with Bias, which is 23% for training and 17% for inference.\n\n| GPT-2 with Sandwich bias (Sequence Length=8192, using gradient checkpointing) | Train Time (s/100 iters) | Inference Time (s/100 iters) |\n| - | - | - |\n| FlashAttention with Bias                                     | 853.21                   | 196.01                       |\n| FlashBias                                                    | 658.23                   | 163.33                       |\n| Relative Promotion                                           | 23%                      | 17%                          |\n\n> **Q2:** \"How easy is it to integrate FlashBias into existing model architectures? Are there plans or existing implementations?\"\n\nIt is really convenient to implement FlashBias following the formalization in $\\underline{\\text{Eq. (3) of main text}}$. All the implementation details can be found in $\\underline{\\text{Appendix G of our paper}}$.\n\nSpecifically, the researchers just need to concatenate the decomposed factor functions with queries and keys following $\\underline{\\text{Eq. (3) of main text}}$, and then the researcher can seamlessly adopt the implementation of FlashAttention. For further speedup, you can also tune the tilling hyperparameter (such as the size of tilling blocks) to achieve better efficiency, while we do not specially tune these hyperparameters during our experiments. Since we cannot add links during the rebuttal phase, we promise to release the code upon publication."}]

Author rebuttal:

Meta-review: This paper introduces a method for faster attention computation when additive bias terms are involved. The reviews are uniformly positive, praising its simplicity, strong speedups across a variety of tasks, and GPU memory improvements. There are some remaining doubts on the universality of the application, although they rely on bias not being used in all models, which as the authors note may itself be caused by the inefficiencies being ameliorated in this work. Given the positive reviews and diversity of evaluation I recommend acceptance without reservations.
Decision: Accept (poster)