Toward Linearly Regularizing the Geometric Bottleneck of Linear Generalized Attention
Authors: Jiaxu Liu
Xinping Yi
Xiangyu Yin
Yuhang Song
Gaojie Jin
Xiaowei Huang
Venue: Accepted by TMLR ()

Abstract: Transformers excel across domains, yet their full self-attention carries a prohibitive $\mathcal{O}(n^2)$ cost for long sequences with length $n$. Existing \textit{efficient} attention methods either restrict the attention pattern (local/sparse attention) or approximate the softmax kernel with certain drawbacks. The former suffers from attention bottlenecks (over-squashing of long-range dependencies) and invalidates the use of global tokens in autoregressive tasks, while the latter often requires sequential processing that can degrade in accuracy when approximations fall short. In this work, we introduce the \textit{Bottleneck Regularized Linear Attention (BRL-Attention)}, uniting the strengths of pattern-based and kernel-based techniques to enable efficient, global information flow with linear complexity. BRL-Attention extends a local attention pattern with a small set of compressed tokens that serve as a global information reservoir, ensuring long-range interactions without quadratic cost. This bottleneck regularization strategy effectively alleviates the geometric attention bottleneck and retains full expressiveness; that is, it matches the sequence modeling capacity of full softmax attention while mitigating over-squashing across layers. Moreover, it integrates global tokens without breaking causal masking, making it applicable to both encoder-only and autoregressive decoder architectures. Extensive experiments on sequence and graph benchmarks demonstrate that BRL-Attention matches or surpasses the predictive performance of standard Transformers with full attention, while substantially reducing memory usage and computation time to levels comparable with linear sparse attention.

Published in Transactions on Machine Learning Research (06/2025)
Toward Linearly Regularizing the Geometric Bottleneck of Linear Generalized Attention
Jiaxu Liu², Xinping Yi Xiangyu Yin', Yuhang Songl.3, Gaojie Jin', Xiaowei Huang! (fazzu.lu, z-yin22, sgyson10, ziomvei.huang huang) Sliverpool. ac.uk, edw..cn, fin.kim@ggmal.com com Universityy of Liverpool, 2 SSutheast University, Natiioal Tsing Hua University, UUniversity of Exeter
Reviewed on OpenReview: https: openreview. net/ farum? id= Vpyg3fgXbl
Abstract
Translormers excel across domains, yet their full self-altention carries prohibitive O cost for long sequences with length n. Existing efficient attention methods either restrict the attention pattern (local/sparse attention) or approximate the softmax kernel with cer- tain drawbacks. The former suffers from attention bottlenecks (over-squashing of long-range dependencies) and invalidates the use of global tokens in autoregressive lasks, while the lat- often requires sequential processing that can degrade accuracy when approximations fall short. In this work, we introduce the Bottleneck Regularized Lincar Attention (BRL- Attention), uniting the strengths of pattern-based and kernel-based techniques to enable efficient, global information flow with linear complexity. BRI-Attention extends local attention pattern with a small set of compressed tokens that serve as global information reservoir, ensuring long-range interactions without quadratic cost. This bottleneck regu- larization strategy effectively alleviates the geometric attention bottleneck and retains full expressivenes;; that is, it matches the sequence modeling capacity of full softmax attention while mitigaling over-squashing across layers. Moreover, integrates global tokens without breaking causal masking, making it applicable to both encoder-only and autoregressive de- coder architectures. Extensive experiments on sequence and graph benchmarks demonstrate that BRL- Attenttnn matches or surpasses the predictive performance of standard Transform- ers with full attention, while substantially reducing memory usage and computation time to levels comparable with linear sparse attention.
Transformers [Vaswani et al. 2017) have substan- tially advanced the stale-of-the-arl in areas such as natural language processing, computer vision, and graph learning (Dosovitskiy et al. 2020 Touvron
Arretice Pazara
(2-
Raxeval Allenies by Patter
Ta 91
al. 2021 Devlin et al. 2018 Touvron et al. 2023).
Their core strength lies in the attention mecha- nism, which models global token-to-token interac- tions. However, full self-attention involves comput- ing pairwise relationships among tokens, leading formidable O(n²) time and memory complex- ity, given the sequence length of n. This limitation becomes major obstacle for tasks involving long sequences or large-scale datasets, where computa- tional overhead can escalate dramatically.
Takan 64
Visualization of Allention Bottleneck in- duced by pattern-based generalized attention. Infor- mation of token #1 and token #5 struggle to propagate to token #6 within one self-attention block.
Figure 1:
In response, various efficient transformers (ET) have been proposed to replace full attention with mech- anisms that cost at most linear time the sequence length (Tay et al., 2022). These can be broadly catogorized into pattern-based and kernel- based methods (Sec. 1.1h. Pattern Pattern-basod transformers (Soc. 1.1
*Corresponding author

---

Published in Transactions on Machine Learning Research (06/2025)
ETC (Ainslie et al, 2020), and BigBird Zaheer et al. 2020), restrict each query token to attend to a local or blockwise subset of the entire sequence, reducing complexity to near ((n) by exploiting structured sparsity. Despite their efficiency, these localized patterns can cause an attention bottleneck (a.k.a. over-squashing (Alon & Yahav, 2020: Topping et al.) 202T)), as illustrated in Fig. [ where distant. tokens fail to interact ellectively within the limited receptive fields of cach layer. While introducing global tokens or memories can alleviate this for encoder-only tasks, it often breaks causal masking and thus remains impractical for autoregressive decoding. In contrast, kernel-based transformers (Sec. 1.1] approximate the softmax via low-rank projections or random feature maps (Katharopoulos et al., 2020 Choromanski et al. 2020). also aiming for linear or near-linear time. However, these methods can degrade performance if the chosen approximation rank is insufficient and often struggle in auloregressive settings: the need for causal masking lypically forces sequential al len- tion computation, forfeiting any parallelizable speedups. They also tend to be sensitive to random- variance, especially in sccnarios involving domain shifts or noisc.
such Sparse Transformers CChild et al. 2019
Longformer [Beltagy et al. 2020).
In this work, we propose a new paradigm, Bolile- neck Regularized Linear Attention (BRL-Attention), that combines the advantages of pattern- and kernel- based approaches while sidestepping their key lim- itations (discussed in Rmk. B.1. BRI-Attenntoon extends any pparse/patter:: based based attention mecha- nism with a small set of compressed tokens, which serve as a global information reservoir. These tokens can be integrated without invalidating causal struc- ture in auloregressive tasks. We formalize this de- sign as a bottleneck regularizer (Sec. 2.3) that chan- nels distant dependencies through these compressed tokens. Particularly, we introduce two key functions to facilitate an expressive attention, Compression and Propagation, which are two commuicattin functions defined between main tokens and compressed tokens. By design, these two functions help alleviate the attention bottleneck induced by over-squashed patterns, while offering benefits e.g., preserving the expressibility and robustness noisy attention. While linearizing attentions inevitably sacrifices expressibility (Hua et al. 2022), we theoretically prove the compa- rable expressibility of BRL-Attention towards full-attention, while mitigating over-squashing between con- tiguous layers. The resulting approach scales linearly with the sequence length on inference time, memory, and autoregressive training time.
Artention 7ya0
BRL-Atertio FaRt.
Receptive
Heceppoce
(Thi Layoo
Dhas Ligert
(lod Leyo
Mila Tiken
Componot Tikers
Compusdon
Percio Vay Sprap
Teup
Figure 2: (left)) Quadratic Full-Attention; (right) Proposed BRI-Attention (Linear) with a win- dow /chunk size (w) of 3 (w is typically set to 2 64 in our experiments).
Our main contributions are highlighted as follows: (1) Starting with sensitivity onalysis (Sec. 2.2 we demonstrate that conventional sparse patterns can fail to propagate long-range information within small number of layers. In contrast, adding our bottleneck regularizer effectively recovers global context with minimal overhead. (2) We propose the Bottleneck- Bottleneak-Regularized Linear Attention (BRL-Attention) mecha- nism, which augments any sparse/ patrse//pateee-bbasee efficient O(n) complexity and alleviates information bottlenecks in strictly local attention. Theoretically, we justify that with all introduced techniques in BRI-Attention, the resulting BRI-Former facilitates wider sensitivity bound between distant tokens in intersective layers (Sec. 2.4], which helps alleviate attention bottleneck. Meanwhile, we justify in Sec. 2.32.4 2.4 that BRI-Former is as expressive as the full-attention- transformer. (3) Through extensive experiments (Sec.33) 3), including long-sequence modeling and large-graph node classification, we show that BRL-Attention not only matches or surpasses full-attention transformers but also substantially reduces memory usage and computational cost.
allention with compressed lokens. This mechanism ensures
1 Preliminary
Definition 1.1 (Generalized Attention Mechanism Zaaeerr et al. 2020 ). Given the input token to the layers x E Raxd with soquence length n and embedding dimension the generalized attention mechanism is described by a directed graph D (a.k.a Attention Pattern (Tay et al., 2 whose vertex set [I....,n.. The set of ares (ie., directed edges) represent the set of inner products that the attention mechanism will
2

---

Published in Transactions on Machine Learning Research (06/2025)
Globe
Diibaal Slilin Wanden
Figure 3: (left) Sliding Window (SW) pattern Dawi (mid) Visualization of graph under different configura- tions of window sizes (right) Various attention pattern D employed in literature, where grey areas are either scores for individual token-pairs or sub-block of full attention.
consider. Let (i) denote the nei-neigs set of node in D, then the i-th output vector at layer E [1, of the generalized attention mechanism is defined as Eq.
<(9: (1) (D))
"
<(q: (1)
kK)))
F(l gen (x
Σ Σ
(1)
Full (x
11)
(2)
(EN ExEND()) r(q .k
> x(q, kk)
where q xW xWa and xW, are query, key, values and W Raxd are learnable projection weights. We use Rd x x - R to denote the softmax kernel x(a, b) == exp(ab (where a, b E R (x)).
For brevity, state the results with the batch size B = 1 and the number of heads H = without loss of generality. Tf D is the complete graph with adjacency matrix S, we recover the full attention mechanism of O(n²) complexity (Vaswani et al. 2017), expressed as in Eq. 2
1.1 Efficient Transformers
Existing works seek to improve memory efficiency in transformers through weight pruning (Michel et al. 2019,, weight. factorization (Lan et al., 2019), weight quantization (Zafrir et al. 2019). elficient pretraining Clark 2020), attention optimization [Lample et al. 2019), or knowledge distillation. Reducing the memory or computational requirements with these methods leads to training or inference time speedups, but fun- damentally, the time complexity is still quadratic w.r.t. the sequence length which hinders scaling long sequences. In this paper, we mainly focus the two lines of works that achieves near linear complexity transformers, namely, the pattern-basod approach and kernel-based approach. We defer more detailed discussion related works to Appendix D.4.
Pattern-Based Partial Attention. Pattern- based efficient transformers mitigate the quadratic complexity of ovaluating Frull by imposing structured sparsity pattorns. Early methods, such as Sparse Transformors Child et al. 2019), employ blockwise or strided atlention mechanisms to reduce computations, achieving complexity of O())R). Longformer (Beltagy al. 2020) and ETC dAinslie et al. 2020) further optimize this by combining local attention windows with global memory tokens at custom locations, resulting in complexity of O(u)), where is the window size. More of this category includes axial (Ho et al.) 2019 learnable patterns through hashing (Kitaev et al. 2020) clustering (Roy et al. 2021). These approaches strike balance between efficiency and coverage, making them ideal for long-sequence tasks. However, one koy problem with this class of methods is that they involve blocking operations, which are not parallel-frienlly and could induce potential allendion bolllenecks,, as discussed in Sec. 2.2
Kernel-Based Linear Attention. Kernel-based attentions approximate self-attention to achieve linear complexity. Specifically, according to Mercer's theorem (Mercer 1909 Aizerman, 1964), the cigenfunctions corresponding to the non-zero eigenvalues are continuous on Ed and K can be represented as <(x,y) Eee (xys))) With such property, the kernelized-full attention is typically given
F(0) (x $1q "1a"))" ビグ= 10k=100k00 ) TTV
where m(x,y) (T(x), +(y))v = T(x)o(y)T
(3)
0(q (1) Σリ ()
With certain error gap, Eq. (3)
essentially tells F (x())
(x!1) (1) With such an idea, approaches full
e.g. (Choromanski et al.) 2020 Katharopoulos et al. 2020 Wang et al. 2020 Peng et al.] 2021 reduce
the quadratic cost of F to roughly Omax (n, dd)). Such bernel-based approximations climinate the need
compute or store full attention matrices, making them well-suited for resource-constrained scenarios and real-time processing. However, kernel-based methods may exhibit limitations in parallelizing across the time
3

---

Published in Transactions on Machine Learning Research (06/2025)
dimension during training in an autoregressive teacher forced setting. As a result, there exists considerable gap between the theoretical complexity and actual running time. Similar to the findings in (Hua et al. 2022), we find that directly computing the full quadratic attention matrix is even faster than the kernal- based approaches on GPUs (sce Sec. 3.1
2
Method
In this paper, present a new attention paradigm replacing conventional full attention, namely, the Bottle- neck Regularized Linear Attention (BRL-Attention). We employ the notation of F as matrix-valued function and f as vector-valued function. The transformer block with BRI-Attention is formulated as
F (xx)) Generalized Attn
F()) (xx0 + prop (xx) x
(4)
gen
(n) Regularizer
( out BRL "1 xx)
LN(FFRL
(+1) LN( (f (it f (x ))) + Xout
!1)
(5)
The key difference of FBRL. to (oom D) is the introduction of regularization term F which is a propagation of compressed information [(ct]; regarded as patch to F (: D) that alleviate over-squashing The equalions in Eq. (5) are standard transformer layers with skip-connection and SwiGLU (Shazeer 2020) where @ is the Hadamard product, 0-elu(-) is the ReLU activation and LN denotes LayerNorm.
As below, we outline the core components and thcoretical foundations of our approach. (1) We first establish the necessity of imposing a structured attention pattern D to achieve linear complexily, and demonst how certain sparse patterns in (=: can induce attention information bottlenecks, wherein token-level interactions fail to propagate effectively across limited transformer depths. (2) We then introduce the bot- tleneck regularizer Fprop: constructed via set of compressed tokens X(ct]. Despite its linear complexity, we show that (tr A) maintains an expressiveness comparable to that of kernelized attention mech- anisms. (3) Next, we theoretically justify that, when the compression mapping Feomp is appropriately instantiated, switching the Full-Attentinn Fiul to the BRI-Attention FBRL results in a broader sonsitivity bound. With our customization, FBRL alleviates over-squashing and improves robustness noisy attention weights. (4) Finally, under any arbitrary sparse attention pattern D., we prove that with a trivial per- head regularization coefficient À = {1) H, the BRL-Attention operator FBRL, remains provably as expressive as Frull (see Thm. 2.9 Collectively, these results establish ORt-Attention tion mechanism that retains the full expressive capacity of standard transformers while offering improved scalability for long-sequence tasks. ( All proofs are deferred to Appendir
linear-complexity atten-
2.1 Necessity of Pattern D to Linear Generalized Attention
As detailed in Sec. 1.1 and TTay al. 2022), the paltern D that facilitates O(n) attention essentially encompasses blockwise and strided pattern attention. The blockwise approaches QQuu et al. 2019) chunk input sequences into blocks that reduces the complexity from n² to no where Ablock < n. then the x n score is computed intermediately by computing and combining divided Thlook X Ablock blocks, selectively picking blocks leads to linear complexity. For the strided attention patterns (Beltagy et al. 2020 Ainslie et al. 2020) approach, we illustrate in Fig. 3 a typical pattern D, í.e. the sliding window pattern, which is a fixed-size window mask surrounding each token. Given a fixed window size w, cach token attonds to w/2 tokens each side. As show in Fig. leftt, as only gray areas are need for computation, we can essentially group the gray tokens/blocke
X w size matrix, resulting in Oww)) complexity score computation. (Zaheer et al. 2020) further combined the blockwisc and strided approach, facilitating the block-diagonal, windkw local, and block-random patterns for efficient and expressive linear attention.
2.2 Sensitivity Analysis of Generalized Attention
As the layer-wise attention x; x(*+1 , F(i D); is continuous and differentiable, the bottleneck of atten- (I+1)
tion information can then be understood in terms of one token embedding x (l+1) failing to be affected by
4

---

Published in Transactions on Machine Learning Research (06/2025)
another (previous layer-) feature x of token p at distance M from node i. Hence, we employ the Jacobian /ox) ) as an explicit and formal way of asscssing the bottlenock-ed attention.
Desinition 2.1 (Attention Boltleneck). Under the definition of generalized attention, Fgon D) with particular D within I layers is said to be bottlenecked when there exists a token pair ~ p such that of gon (x(); D))e ()) O. This mcans no information is flow from i to p with / layers of attention blocks. I is nearly I,, then the whole transformer suffers information bottleneck similar to MPNNs.
[Thcory] Failing Cases Under Certain Patterns D. We reveal theoretically that the generalized allentions suffer from attention bottleneck with some patterns defined by certain D. Recall strided sliding window empowered linear transformers according to Sec. 2.1 Using multiple stacked layers of such windowed attention results in a receptive field, so that top layers have access a board range of input locations and have the capacity to build representations that incorporate information across the input, similar to MPNNs. The computation complexily of this pattern is ((wn), which scales linearly with input sequence length Now, letting the adjacency matrix derived from pattern Dew be S we state the following proposition.
Proposition 2.2. Let D Dew where the sliding window size is and S has eigenvalues bounded by T'w- Considering the source loken i and larget loken p that are M distance away, we have the sensilivily bound
OF (:; Dow)s
10
ff
2M w-1
1,
(6)
0xp
(U)
(rswtw)+
otherwise.
Prop. 2.2
indicates the particular == Daw leads to squashed information propagation, f.e. for the top w
layer token i's embedding to perceive token p, need at least
2M general attention layers.
Remark 2.3. We regard Daw as a simple example to illustrate how certain restricted D induces substantial bottlenecks in attention propagation. In fact, many other D with limited connectivity and sparse interaction patterns (e.g. diag chunk, dilated sliding window) exhibit similar issues, where the Jacobian with respect to distant tokens becomes effectively negligible at intermediate layers. This makes the attention mechanism behave analogously to a communication limiteed message passing network, which hinders the model's capacity to integrate information across long sequences.
2.3 The Bottleneck Regularized Linear Attention
This section tackles the RQ: How can a small, Irainable set of global memory capture long-range contert in linear complerity while preserving the cxpressibility. We investigate it by (1) introducing the compressed tokens that evolve by layer to accumulate global summarics, and then (2) coupling them to the main tokens through a bottleneck regularizer Fprop that propagate the reserved information context embedding, thereby marrying linear-attention efficiency with full-attention reachability.
[Practice] Compressed Token as Information Reservoir. We first introduce a key concept that facilitates the bottleneck regularizer, namely, the compressed token. Given m as the sequence length of the compressed token and det as the embedding dimension, we represent the compressed tokens as X[ct] E mxdee. The layer-wise update process of the trainable Xtet] is then defined
(1+1) [ct]
(F (FOmp (x (xx)x Trelu ((1 (11 (1-
(1) (1) x0)) (xt) MD) E Rmxd4
where the initialization (x(e))ss )ij 1). M E Ra is feature transformation matrix for residual connection and B is a control factor on compressed token evolution. Notably, setting / > 0 to evolve the compressod token is crucial for achieving good result, which will be shown in Cor. 2.5 and in experiment
3.2 The function Feomp (RR) (IR(==)))))) + (Rda), m is a matrix mapping that defines the compression of information from main tokens to compressed tokens, giving the dynamic of X(et]. This parameterized compression mapping must be at least as expressive as MLPs and is critical in our latter analysis. We provide instantiation and intuitions regarding how relates attention bottleneck in Sec. 2.4 Below, we detail how the X(ct] is formalized to regularizer via information propagation.
5

---

Published in Transactions on Machine Learning Research (06/2025)
Figure 4: The BRI-Attention block (enc/dec-only). The algorithm receives main token X, compressed token X[ct] for layer 4, and qucued history tokens (on training phase only when causal training is enabled, otherwise the history are the same as main tokens) as input. Within each layer, we first compute the multi- head generalized attention by Eq. ① under the- oretical guarantee Thm. 2.9 the instantiation of D can be quite flexible. In our case, we employ sparse local-attention e.g. (Beltagy et al., 2020) as a linear complexity instantiation. Then the context x another branch is manipulated with X[ct] by Fcomp to get the evolved X[ct]- We then construct. the regularizer Frrop by and the up- dated X[et]. Notably, the x and X(et] are both processed by their respective gated FNN before the next layer. Optionally, the x in Fprop could be processed by feature map / //pppiky to facilitate a lossless regularizer (see Cor. 2.5). We detail the algorithm in Alg. in Appendix.
Next Layer
x x
Add de Norm
Add & Norm
Feed Forward
Feed Forward Network
Netwerk
(1-3) Teal
Key
Propugation F ())
Add & Norm
Query
Add & Norm
Causal Mask
Key
Muti-Head Local-Attentinn (xp)
Compression Poing (1)
Query
Lengh::
Previous
[Practice] Regularizer via Compressed Information Propagation. With the evolution of compressed token defined layer-wise according to Eq. 7 we are ready to instantiate the Bottleneck Regularizer, formu- lated as a compression operation Feomp in Eq. (4), defined by
z -(xx))
E Ruxe zt zK,z"= = [cttt W!KV) E R Rmxd
(8)
F (x10; X A) ADatnn (2)()))T)
2Y E Buxd
For both encoder-only and causal models, x()) is the current input. 6 Rd - R a parameterized feature transformation for aligning feature spaces. Optionally, we could introduce & as spikiness enforcer (Zhang al., 2024). , @spiky (v) [exp(vw1...e + b)] (which has close representation as Eq. (10) such that theoretically benefits Prop. 2.6 The Fprop, defined by Fprop: (Rd),, Jn x (IR Jm (RR),, Jn, can be viewed as a cross sequence message passing that pulls the information from X(ct] back to main tokens. W E RR Xe
and W E Rdda are KV weights. The attention function (u) = TLapiace/Softann (u + brel) can be implemented as Softmax for regular attention, or Squared ReLU LLLU/LLplaace function (Ma et al. 2022 for better convergence speed and training stability. The bral E RR is the relative positional bias, which can
be drawn from approaches e.g. (Raffel et al.) 2020 Su al., 2024 Ke al., 2020 Press et al.; 2021).
We employ X(ct] = [ct]
(I+1) for each layer of Eq. 19. We defer the explanation of employing the evolved
X[ct] layer-wise in Cor. 2.5 The overall complexity of Fq. (9) is roughly O(mn), which scales linearly with sequence length n.
[Theory] Propagation Mapping as Regularizer Preserves Comparable Expressibility to Kernel- ized Attention. Tn the following, we theoretically justify that: Adding regularizer F to the generalized attention function Fren induces no degradation on attention expressibility despite positive X. Define the random feature map (Choromanski et al. 2020) Óeem : Rd + R for v E Rd as
Prtm (v) = e exp(elxp]] |v [exp(vw1), exp(vw2), ... e..p.v //Ve.
(10)
where WI...c are random transformations drawn from IO) According to (Choromanski et al., 2020 Prtm is one of the instantiation of Φ in Eq. 33) that facilitate a softmax kernel-approximation In the analysis below, we make the following assumptions on Eq. (9) (CI) Φ(-) ~ (+)m(+) and (C2) det c(1 + d). With these assumptions, we arrive at the following proposition:
Proposition 2.4. Al each layer I, let x x((+1) = =IEj Óvem ( |flatten (Ej= Ortm (rtm(()) ( TT (T where || is
the concatenation operator. The flatten(.) operation reshapes input Bexd - Rod to raster order. Then, if we ((+1) (1+1) Eq. ( sufficiently appronimate the kernalized self-altention (i.e., Eg. 3))
force x [ct]
6

---

Published in Transactions on Machine Learning Research (06/2025)
11+1) Recall that x [ct]] given in Prop. 2.4 is also a vector matching the shape Iddes of compressed tokens in Eq. (since c(1 + d) det under (C2)). T Therefore, the proposition essentially tells the following Corollary:
7
Corollary 2.5. With some particular instantiation of function Fprop(-) powered by particulor choice of Φ(-), there erists some fixed compressed token x [(tt] as input to Frrop, such that Forop (x; x (cts) is essentially an inlention-upprorinenn
Cor. 2.5 reflects the benefit of evolving X(ct] dynamically instead of sharing across layers: (1) [et)] with evolution Eq. 7 is capable on approximating x((+)) of Prop. 2.4 for all I; and (2) with X[ct] more expressive
((+1)
than x1++1) Prop is at least as expressive as attention approximation. Therefore, adding such particular Froop under Prop. 2.4 results in degradation in expressibility. Recall the instantiation of F Fprop in Eq. (9) We are then interested in how such a border definition of Frop could impact the expressibility. Denote I| for vectors as l2 norm and for matrices as spectral norm, we state the following proposition.
Proposition 2.6. Under assumptions that: (A1) Given feature space X' C Rd, for all i E [], the token feature x, satisfies X; E X and ||x;|| S Tz. Similarly, we assume the compressed token [xceils for j E [1, m] is bounded by Tet- (A2) All weight malrices e-g- that fealure transformations like allention paramelers salisfy ||Will S rw. We have: A parometerized network f(-) that is as expressive MLP with (L)) width and depth can approzimate @rfm (q:), E; E,0rtm(ke (ky) and Σ; Orim (k,)) arbitrarily well on the compact domain.
Prop. 2.6 essentially tells that (1) ((.), if simply defined MLP, à R - R for instance, can approximate @rim (q) arbitrarily well (a spiky c depiky facilitate better approximation); and (2) in Eq. (7), if disregarding
the skip connection ie., 3 = 0, then Fcomp with some simple MLP instantiation (e.g., E, (x, where
MLP=: Rda + Rd,),, can approximate x [et]]
((+1) arbitrarily well. Therefore, if we essentially approximate
both é and (1+1) then by Prop. 2.4 the propagation mapping Fprop, also defined as a function of
(q:) and x
((+1),
[ct) > can approximate the kernelized form of self-attention (Eq. (3) arbitrarily w Finally,
the broad definition of Frop as in Eq. (9) could also inherit the benefits instructed after Cor. 2.5 In the
following, we discuss how appropriate instantiation of the compression mapping Feomp impacts the attention bottleneck bound.
2.4
The Information Compression and Sensitivity Analysis of BRL-Attention
This section tackles the RQ: Given squashed pattern D, how concrete compression. mapping lets the FRen approzimate full attention yet keep the linear complerity and offers eztended sensitivity bownd. We investigate this by instantiating Fromp as differential cross-attention that uses rescaled dual score matrices. This construction yields the wider bounded-error Jacobian and preserves linear complexity.
[Practice] Information Compression. We now introduce an instantiation of Feomp (X(c)) x) in Eq. Let x be the history of x on time of sequence modeling. Define
he = X(ct) w E m.xes
h KI h K2 Rhhhhhh = xW / KL,22V) E <Raxd
(11)
Foomp (x(ot) = (SKI SV mxdds where " Jatti h (h(
(12)
/
Notably, for self-attention, x is identical to x, while for autoregressive training, X is the last history x of the same block size (more details will be illustrated in Appendix B.2). Compared to the formulation of Fprop a b) in Eq. 99) we regard X(ct] as the query and construct keys and values with main tokens x. The Jattn now have relative bias brel E Rmxn In addition, we make the following modifications: (M1) The attention scaler is changed from /d to my/et. (M2) Inspired by differential transformer (Ye al. 2024). we construct two attention scores SK1 and k with same X(et] as query and different transformations of x as keys. We regard their re-scaled subtraction as the final attention score.
[heory] Alleviated Attention Bottleneck (Information. Squashing) via Compression Map. The formulation as Eq. 12 offers the following properties:
Remarkably, the error gap between kernelized and full attention
-iqq// / kyF)Il bounded by
O /exp(6rrrw ht)/TTmmmm))) according to (Thm. 1, Wi et al. 2022h)

---

Published in Transactions on Machine Learning Research (06/2025)
Proposition 2.7. Assume for simplicity in Eq. 77 that B 0, and Feomp is defined according to Eq. 12
Then in terms of the sensitivily bound of Fprop; the modification M1 resolves the bound scoling with x [et] sequence length m, and M2 provides a controlled bound with factor y compared to vanilla cross-attention. The sensitivity bound will be
prop aF (x()) xx1)xx (1t2)); );
( (" Vddet (ddeec
(1 (1+h) h11)
S SO
(13)
Oxp
compared to the bound of O( ("=tarth) without M2. This bound is applicable to I E [0, which helps the Jdde:
alleviation of allention boltleneck compared to with Fren only (us shown in Prop. 2. the bound is consistently zero when M is large). For details w.r.t. the lower-bound, we refer the readers to Appendix E.7
Remark 2.8. Augmenting the compression Fcomp Eq- 12 to differential form (according to M2) could also potentially improve the robustness of learning useful information to those irrelevant, which benefits the retrieval of information in a long noisy sequence. While augmenting the propagation mapping Fpoop does not offer such a benefit for a row-wise Cattn, we keep it in regular cross-attention form (however, one may still augment prop to differential form if column- wise normalization is favored). We justify in detail regarding the above phenomenon via the concept of Signal-toNNoise Ratio (SNR) in Remark D.3.
Consider the Attennttnioon layer in the form Fll is equivalent to initializing Eq. (4) by D
BRL (x(()) Fllh (x(); Dew): + F F (x(1,, x [(t)] )i which Daw and X {1) Recalling Prop. 2.2 with the result of
0F pen (::Daw) ||=
8F (::P,n) || O. This implies that the information axp
Ineq. (6), we have 4I <
2M w-1
0xp
layers of attention propagation. When adding the regularizer Froop Foomp to gggn; we essentially extend the sensitivity bound to non-zero, which implies FBRL of token is more capable of receiving (potentially bottlenecked- Jinformation from p-th token from provious layers, which alleviates the squashness of attention propagation.
of token p cannot be captured by token i within [.2
[Theory] Carefully Tailored Attl-Attention. To show that Fprop o F Fcomp does not degrade the performance, finally, we state Thm. 2.9 to show that the regularized generalized attention FRon + Froop o Fromps with wwwpressiooopopppgaateo for- general long sequence tasks.
Ceappression-Propggaate
Offers Comparable Eupressibilltt
12), is as expressive full-attention Theorem 2.9 [Expressibility of FBRL). Given S < ec and e >> 0, for any continuous functions Foon [0. 1]mxd + RRaxd, let the atlention function Jaten of Fprop be column-wise softmar. Then let dp be the lp distance, there exists a BRI-Former with arbitrary sparse-attention D such that
mulated as Eqs. 9
dp(Foon: dp(Foon:BBBJ) FBRL) = =Foooo D) + Fprop(:: :roop(:FFoomple Foomp())) S..
(14)
Skelch Proof. (1) We first show Prop. E.2 where FBRL. (- D) + F F: Fcomp(-)) can simulate
generalized attention with a star graph S (given as Def. E.Ip, Fgen (x; D U 5), with any sparse pattern D where DnS = 0. This means FBRL = Fgen xx; DUS). (2) We employ Thm. E.3 Zaheer et al. 2020) that
stales: for any continuous function Fron and e > O, there exists a sparse Translormer Fgen( DUDD DUS) such that cFmm Fgen(:: S)) S e. (3) Since FBRL can approximate F (: U S), and Fgen ( DU S) can approximate any Foons it follows that (Foo, S with an arbitrary sparse pattern D.
2.5 Implementation
We illustrate the framework (self-attention version) of our proposed attention block in Fig, Meanwhile, we detail the batched algorithm of Eq. 4 in Alg. [ in Appendix. The local-attention part (the Fgen function in Alg. 1 is implemented based on Phil Wang's implementation
Initialiaation of Parameters. We allow for a flexible number of compressed tokens (m where m 2 1). Typically, an m < 256 is good enough for all tasks in our evaluations. We set c = da = in Eq. 9 for simplicity. For hyperparameter setups: We empirically find setting / = 0.5 performs well gen- erally. We let 7 exp(лιлк1) cxp(л0лкι) where default values for the learnable 1Q/K1/K2 are
Local Attention github repo: httpt//g///ttiiiieo////eetaaaaaa

---

Published in Transactions on Machine Learning Research (06/2025)
torch.normal (H, (Zhang et al. 2023) ) to facilitate cold start, which does not affect the embedding quality if the differential [b) that setting disabling the learnability of
Practically, it is good to initial mean of learnable Y to 0 (similar to
form offers negative impact. Additionally, we observed in Fig. [bb
offers negative impact. Finally, we select the regularizer weight in Eq. (4 from 0.1 ~ 1.0*torch. ones (H)
where we empirically find 0.5 generally works well. Still notably, as also learnable, the choice of which is not sensitive to the final result (which can be observed Fig. aa)).
Computational Complexity Analysis. Deferred to Appendix B
3
Experiments
3.1 Time and Memory Complexity Evaluation
M80
CIT
Eacolin Sapunir Lengle 00
Excoling Saypsce Laggese
Thodding Sagussie Lenglt (0)
(a) Forward GPU Mem (Training)
(b) Backward GPU Mem
(c) Forward Speed (Training) f
Figure 5: Comparison of the computational requirements for Forward/ Backward pass for Encoder- Decoder based models (batch size 16, 1-layer of encoder and 8-layers of decoder with 8 heads and 512 hidden dims) on simple copy task. Dotted lines denote out of memory corresponding models. Our BRL- Attention scales linearly with the sequence length, unlike Full-Attention (Softmax), which scales with the square of the sequence length both in memory and time.
1790
Inference Speed (lipoch
178
Liaferen 2012
GPU Mennory (Peak, MB)
Laval (50)) Asaatoa whd50
126
Laval (39) Ansaton v-d12
1258
Paicess
3900
Lefase
Listeem
12
150
990
LeC-Arratia wn15 and
Aderiie widlt
11
MiloAteaan and andi
3800
150008 Encoding Sequenee Length (n)
300000
20000
2903
5000
1000
1500
20000
25000
Eacoding Sequenee Length (o)
(a) Inference Speed (Batch Size 1) t
(b) Inference GPU Mem (Batch Size 1) 4
Figure 6: Comparison of inference efficiency of various Encoder- Only transformers (single batch, 4-layers of transformers, 4 heads and 256 hidden dims). Our model achieves 13x (further evaluation leads to the OOM of Full-Attention) of inference speed long sequence while scaling linearly with sequence length.
We evaluate transformers with various attentions Autoregressive Encoder-Decoder Copy Task quence modeling experiment in which encoder first processes given input sequence, and then toregressive decoder is trained to replicate that same sequence as its output. Notably, all basclincs are capable in convergence, and we only employ it as a debugging baseline for evaluating the memory and time consumption for encoder-decoder based models training. Specifically, we evaluate Ftte-ttten Performer, Linformer (k E (64, 128, 256)), Local-Attention (w E (50, 100}) and our BRI-Attention (w e (50, 100), m E (10,50 50, 100)). The experiment ran under batch size 16, with a 1-layer of encoder, 8-layer of decoder, heads, and 512 hidden dims. We benchmark with sequence length [200, 2000) with

---

Published in Transactions on Machine Learning Research (06/2025)
step size 100. As demonstrated in Fig. the Full-Attention Transformer, Performer, and Linformer with all settings all suffer from out-of-memory when sequence length exceeds 1100 ~ 1300. Our BRI-Attention, for the same window size w as Local-Attention, the BRL regularization only adds negligible computation, which also results in a theoretically (nn) memory complexity according to Sec. 2.5 Verified in Fig. 5 sim- ilar to Local-Altentton, the BRI--Attention scales linearly with the sequence length, unlike Full-Attention (Softmax), which scales with the square of the sequence length both memory and time.
Additionally, benchmark the inference-only speed and memory performance with a single batch. In this experiment, evaluate Full-Attention, Performer, Linformer (k {64, 128, 256)), Local-Attention (w E (256, 512)) and our BRI-Attention (w E {256, 512), coder/decoder only, with only self-attention. We test 4-layers of enc/dec, 4 heads, and 256 hidden dims. As no backward required, we can scale n up to 2 length with a single 24GB RTX-4090 GPU. In essence, we range n {210, 211 ..... 218). We demonstrate the inference speed and memory in Fig. The Full-Attention, similar to Linformer, gives out-of tintof--emory on > 2'4, which is incapable of very long sequence inference. The Performer, different from its autoregressive training performance (worse than BRI-Attention both time and memory cost when scaling to longer sequences), performs well on inference-speed. Our models achieve > 13x of inference speed on sequence length n 2 2 compared to the Full-Attention counterpart. This verifies the efficiency of our method.
{64, 128)). All the architectures are en-
Table 1: Results on Long Range Arena Benchmark. We compare our method to three major classes of efficient transformers, namely, the full attention, low-rank kernel, and pottern-based as mentioned in the preliminaries. The best results in boldface, and the second bests are underlined. X-marks in Path-X denote chance accuracy.
ListOps (2k)
Text (4k)
Retrieval (8k)
Image (1k)
Pathfinder (1k)
Path-X (166)
Average
Class
Linear
Full-Arieation
Transforner
64.27
57.46
42.44
54.35
Transformce (our-imp))
X
47.90
79.08
$2.31
75.04
76.64
84.72
72.19
Linformer (Wang et
2020
35.70
$0.94
52.27
38.56
76.34
51.36
Linser Traze
Katlagoponks
2010
16.13
65.90
53.09
42.34
75.30
50.55
Law- Kemeel
Petformee
18.01
65,40
53.82
42.77
77.05
51.41
Laurs
bah m1
35.33
65.11
59:61
38.67
77.80
55.30
cueFormurr
Mi
7722
90
63:41
61.96
43.17
70.33
55,
Flowformer
Wil
70
02.24
43.200
73.95
56.48
Local-Atta
15.82
$2.98
53.39
41.46
66.63
46.06
Sparse Trans
Child
17.07
71.21
Longformer
Bellagy m.
2014
35.03
62.85
56.89
42.22
69.71
53. $3,06
Losrnabbee / Fixed
BigBird
Zounez
$5.01
Patterus
Sliceformer
Yuaaa 82 X1
2024
37.05
61.60
02.23
48.02
82.04
58.
Reformer
Kilgey eq
Sinkchrrn Trans ( Iay ml m
20700a
333.67
61.20
50.83
41.23
67:45
51.39
Syuthesizer
[Tay
36.99
61.68
54.67
69:45
52.88
Fxxd Patterns I Low-Ranis Regularizer
BRI-Former (m 64)
47.37
80.20
$2.00
75.75
76.94
85.26
74.72
BRL Former (m 128)
49.14
$0.33
82:98
76.47
76.06
86
75.19
[RI-Formet (in 256)
49.98
80.90
83.22
76.20
77.48
86.89
75.78
3.2 Experiments on Sequence Modeling
Only Sequence Modeling. We evaluate BRL-Former on the Long Range Arena (LRA) TTy
ct al. 2020b,, and compare against established baselines in both full-attention and efficient-attention cat- egories. (1) Full-Atha Atll-Attention. We regard the standard Transformer as the baseline, which has quadratic complexity in sequence length. The 'our-imp' denotes our implementation with RoPE (Su et al. 2024)
based on (Amos et al., 2023) (2) Low-Rand & Kernel Approzimation, which reduces complexity by project- ing queries/ keys by approximating the attention matrix with random features. Learnable/Fized Attention Patterns, which leverages windowed or dilated sparse patterns to reduce attention complexity from (nn?) to approximately (nn) ~ On n.). In our experiments, we employ the LRA training procedure as instructed in (Amos et al. 2023 as baselines for comparison.
Main hyperparametors: (1) The window size w of the local pattern D, which scales the local receptive field. (2) The compressed token length m that aids in mitigating the attention over-squashing We explore three settings of (m (64, 128, 256)) alongside 512 for all local-attention-hasel models (e.g.
10

---

Published in Transactions on Machine Learning Research (06/2025)
Local-Attn, Longformer). The hyperparameter settings for each subtask for the orighnal-mmllemnnttioo and our-implementation are delegated to Tab. 11 and Tab. 12 respectively.
Tab. reports the final LRA scores. As shown, the BRL-attention with m 2 128 achieves either the best or runner-up performance across all LRA sub- tasks, which validates our capability to encode long sequences. In Tab. [ we additionally compare the performance of BRL-attention to the variant with 0 (disabling the Attn-Diff). We observe that the Atfn-Diff contributes positively to the final per-
Table 2: Ablation study on LRA w.r.t the effect of Attention Differentiation.
Model
ListOps
Text
Retrieval
Image 1
Pathindde f
Ponsfrmee
36,37
64.27
57.46
42.44
T1.40
Trundeun (ous-imp)
47.00
79.08
60.31
75.04
76.04
Local-Arto (m=1)
15.82
52.98
59.30
41.46
05.63
BRL For-oor (9// Atto Attn-Def)
00
78.70
81.90
74.33
76.70
BRI--Romeer
77.48
formance, which is also validated in the following autoregressive sequence modeling experiments.
Table 3: Results on WikiText-103. The best/sccond are bold/underlined. (our-imp) denotes our implementation with customized model configurations based on nanoGPT, fewer parameters than Qin et al. 2022a
2.50
Muthod
PPL (val)
(len))
Dvel Stp 5002
Seet 1x0000
Tracaferner
23:63
31.01
Trandoemre-18 2ho et al W01
32:37
32.30
(a) WikiText-103 validation loss (b) Tix validation loss over steps. The window size and over steps. The window size and compression length are both 64. compression length are both 128.
FLASH Hus at mi 532
31.15
34.53
Lincer Trans
h. harongrou of
2020
32.63
34.20
Prefurmse Cloe
75.20
77.60
TranNNoroo
29.57
Transformer tmp)
21.31
Loral (mm>imp)
Figure 7: Losses under various setups of BRL-Autention.
DRL
20.36
22.31
Decoder-Only Autoregressive Sequence Modeling. We study the autoregressive language modeling WikiTeext-033 (Merity el al. 2016). We detail the model parameter configuration in Tab. 13 We defer the study with OpenWebTText [Gokaslan et al. 2019 to Appendix A.3
In Fig. 7 we compare the performance of BRL-attention to local/full-attention under various settings. With various initial values, which is a coefficient of propagation mapping Fprop. in (a), we additionally study the impact of B in Eq. (7) where B = 0 indicates no residual connection on the evolution of X[et]- Wo observe that with only m = w = 64, the BRI-altention significantly surpasses the local attention, which indicates the attention bottleneck is relieved. Different setups of Amai do not heavily impact the final result as they are trainable. However, good Amnt is around 0.1 0.5. In (b), we additionally study the impact of in compression mapping Fcomp. With m comparable and better results against full-atiention. helethe setting 7 as trainable indicates enabling Diff, we observe that Attn-Diff benefits the optimization, which leads to lower losses compared to those with = 0. Using perplexity (PPL) as the evaluation metric, the final results are reported in Tab. 3 where baseline results are partially derived from (Qin et al., (022a) for reference. The BRI-Aitention obtains comparable or better perplexity to the vanilla attention and outperforms all existing linear models with clear margin. Compared to linear methods, BRL- achieves substantially lower perplexity, demonstrating the effectiveness of our method in causal models.
128 (< block size 512), the BRL-attention can achieve
3.3 Experiments on Large Graph Modeling
We evaluate BRI-Former on node classification tasks using DBLP, ACM, IMDB, and Freebase datasets from the HGB benchmark. DBLP, ACM, and IMDB follow HGB (Lv et al. 2021 guidelines, while Freebase uses the split from (Mao et al. 2023 Dataset details are in Tab. 10 Evaluation metrics include micro/macros
Bascline models span four categories: (1) Simple MPNNs (GCN, GAT); (2) Messago-passingg heterogencous GNNs (RGCN, HAN, HetGNN, SimpleHHGN); (3) Transformer-basse based models (GTN, HGT, NodeFormer, H[Normer); and (4) Pure transformers with Poly-Token, evaluated with Full-Attention (PHGT), Local-
11

---

Published in Transactions on Machine Learning Research (06/2025)
Table 4: Results on heterogeneous node classification datasets. Vacant positions (x) indicate the models run OOM on the corresponding datasets. We report the average results in 3 runs.
DBLP
IMDB
ACM
Freebase
Methods
Micro-F1
Macro-F1
Micro-F1
Macro-F1 f
Micro-F1
Macro-F1
Micro-F1 1
Macro-F1
GCN
Kipf Welling
201G
91.17+0.31
90.81=0.32
61:82=0.61
57.8881.18
92.1240.88 23
92:17=0.22
60. 00.23:00.92
27.8143.13
GAT
Velcckovic et al.
2
93.3910.30
93.83 27
61.8420.49
58.94111335
92.1910.15
92 02.26
65.26..0.80
40.741.2.58
RGON
al.
2018
92.07+0055
91 01552-0.50
62.05-00115
08.85++0.6
91.41+0.70
01.55-00 01.55-0.74
60-82-1.42
59.08+1.44
HAN
Wang ei al 2019
92.065+0.62 62
91. 01.67=0.10 49
64 0.58
67.74:0096 96
90.79: 43
90 0.43
61/42=3.50 3.56
57.05++2006
HetGNN
Zhang
92.3341041
91.7620.4 43
08.2510.67
86,0540.25 25
85.91=0125
Simple ICN
[LV
2121
94.46:10.22
94.01 2.
07.36
63.53 1.36
93.35:10.45
93.42:0.44
67.49:
62.49.1
GTN
Yun al.
2019
03.97:+0.54
93 52:00.005
66. 14:0.45
01447+0.56
91, 01.20:0.71
11.31=0.70
X
HGT
Ho cl
212U
93,19:0:25
93.01=0.22 23
67 007
63.00=:111
11.0020.77 76
91 12-0.76
66, 06.43:1188
60.08+2.21
NodeFFrmee
Wi et al
21/2b
93.68.10.42
93.05
65.86.00.22
62.15.00.77
91.80110.31
92. 02.72-00.4
67.01
00.8311.11
HINormer
91.91:0.21
91.37=0.25
07.8320.31
61.65+0.58
93.15400.08
08.28:0:13
67.78:0039
6.167=1.10
Fhill- Transformer
La al.
2024
95.33310.18
94.96.20.11 0.17
68.81:00.88
65.91110.00
93.72.10.0
00.70=0.39
68.74:11422
6.
Xotn
04.96+0024
94.87
07.93+0.14
05.45++0:32
03.333+0.30
03.58-00:0 0.224
67.78-00.3
60 0.04
BRI--ooruer
95.67=00.0
95.35:10.18
68. 08.99:0.112
06.29:0046
93,78=0.21
93. 08..1250.5
6.0
618802..00 10
Loce Inveev Vt
Loa ce DEL F0wee-00)
Food
(a) Performance Cmp-
(b) Memory Cmp.
Figure 8: Performance and mem-cost comparison of Figure 9: Convergence of training/valideen valliaaiio losses Transformers with different attention backbones on on heterogencous graph datasets. We compare Trans- various setups. We observe that BRL-Attn is superior in performance and memory efficiency (neglectfully more costly than Local-Attn) across all datasets.
former with Pull/Tocal/BRII- /BRI-AAttetioon. We observe that BRL-Attn typically converges faster than the others on their respective minimum.
Attention, and our BRI-AAttention. Baseline results from HGB (L et al. 2021 are quoted directly; others are re-evaluated via OpenIIGNN.
Tab. presents the node classification results. Our proposed BRL-Former with Poly-Token Token consistently demonstrates superior performance, outperforming other baselines, including the Full-Attention Transformer, across most scenarios. For the Freebase dataset, (Lu et al. 2024h notes that local structures are particularly significant, as Freebase is a knowledge graph composed of individual facts or triples. BRL-Attention naturally restricts the receptive field to local subgraphs, providing notable advantages on this dataset, surpassing the current statc-of-the-ars results. Morcover, BRL-Attention outperforms message-passing (H)GNN basclines in nearly all cases, indicating that compressed token propagation elfectively resolves the receptive field lim- itations inherent in
and enhances model performance on heterogeneous graphs. Additional analyses in Fig. 8 reveal BRI-Attention is suporior in performance, being memory efficient (neglectfully more costly than Local-Attn) while consistently converges faster across all datasets to lower loss values, confirming its efficiency and effectiveness on large graphs modeling.
4 Conclusion
We presented Bottleneck- Regiiaiiiidd Linear Attention, a mechanism that augments sparse/patterr--asee attention with a small set of compressed tokens to capture long-range dependencics at lincar cost. Our theoretical analysis and extensive experiments on sequence and graph benchmarks show that BRI-Attention consistently matches or outperforms full-attention baselines, while being more efficient. Unlike kernel-based methods, BRL-Attention avoids challenging kernel approximations and supports parallel training for both encoders and autoregressive decoders. In essence, the compressed tokens serve as a global reservoir that mitigates over-squashing without requiring costly quadratic attention or specialized global tokens that break causality. These findings position BRL-Attention as an efficient and scalable alternative for regular attention.
12

---

Published in Transactions on Machine Learning Research (06/2025)
References
Joshua Ainslic, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qilan Wang, and Li Yang. Etc: Encoding long and structured inputs in trans- formers. arXiv preprint arXix:2004.084 08483, 2020.
A Aizerman. Theoretical foundations of the potential function method in pattern recognition learning. Automation and remote control, 25:821-837, 1964.
Cri Alon and Eran Yahav. On the bottleneck of graph neural networks and its practical implications. arXiv preprint arXiv:2006.066 05200, 2020.
Ido Amos, Jonathan Berant, and Ankit Gupta. Never train from scratch: Fair comparison of long-sequence models requires data-driven priors. arXiv preprinl arXiv:2310. 000:2300.00900, 2023.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint ar 2rXi:2004.051500 2020.
Kangjie Cao, Ting Zhang, and Jueqiao Huang. Advanced hybrid Istm-transformer architocture for real-time multi-task prediction in engineering systems. Scienlific Reports, 14(1):4890, 2024.
Yu Chen, Lingfei Wu, and Mohammed Zaki. Ieraativ deep graph learning for graph neural networks: Better and robust node embeddings. Advances in neural information processing systems,
2020.
Rewon Child, Scott Gray, Alec Radford, and Tya Sutskever. Generating long sequences with sparse trans- formers. Xixiv preprint S.4::19944.00000 2019.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint Xxie::2009.149/ 2020.
K. Clark. Electra: Pre-training text encoders as discriminators rather than generators. arXio preprint arXiv:2003. 10555, 2020.
George Cybenko. Approximation by superpositions of sigmoidal function. Mathematics of control, signals and systems, 2(4):303 -14, 1989.
Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Sristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mixing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402.190 194277, 2024.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprinl (xXiv:1810.04800 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth Transformers for image recognition at scale. arXiv preprint arXiv:2010. 11929, 2020.
16x16 words:
Luca Franceschi, Mathias Niepert, Massimiliano Pontil, and Xiao Hc. Learning discrete structures for graph neural networks. In International conference on machine learning, pp. 1972-1982. PMLR, 2019.
Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus.
http://
Skylion007.c gi github. Lo/0pennWee/TextOoppuus 2019.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint wwXiv:2312.0075. 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. Xie preprint orXiv:21 1.00096, 2021.
13

---

Published in Transactions on Machine Learning Research (06/2025)
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces as effective structured state spaces. Advances in Newral Informalion. Processing Systems, 01222982-229942 2022.
Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidimensional transformers. arXiv preprint arXi:: 12180, 2019.
Ziniu Hu, Yuxiao Dong, Kuansan Wang, and Yizhou Sun. Hcterogenoous graph transformer. Proceedings of the web conference 2020, pp. 2704 2710, 2020.
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 1099-9111 PMLR, 17-23 Jul 2022. URL https://proccedinngs mmlr press peress/vile v11/2/ua222a..... html
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rmns: Fast autorogressive transformers with lincar attention. In International conference machine learning, pp. 5156 5165. PMLR, 2020.
Guolin Ke, Di He, and Tie-Yan Liu. Rethinking positional encoding in language pre-training, arXiv preprinl arXiv:2006. 15595, 2020.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. arXiv preprint Xrv:: 02907,, 2016.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint wxXiv:201.04451 2020.
Guillaume Lample, Alexandre Sablayrolles, Marc' Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys. Advances in Neural Information Processing Systems, 32, 2019.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Al- bert: A lite bert for self-supervised learning of language representations. arXiv preprint Stt::9999.19444 2019.
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Sbalev-Shwarlee et al. Jamba: A hybrid transformer-mambl language model. arXiv preprint XXx:22403.19887 2024.
Zhiyuan Lu, Yuan Fang, Cheng Yang, and Chuan Shi. Heterogeneous graph transformer with poly- tokenization. International Joint Conferences Artificial Intelligence, 2024.
Qingsong Lv, Ming Ding, Qiang Liu, Yuxiang Chen, Wenzheng Feng, Siming He, Chang Zhou, Jianguo Jiang, Yuxiao Dong, and Jie Tang. Are we really making much progress? revisiting, benchmarking and refining heterogeneous graph neural networks. In Proceedings of the 27th ACM SIGKDD conference on knowledge discovery data mining, pp. 1150-1 1160-116, 2021.
Mingqian Ma, Guoqing Täu, Chuan Cao, Pan Deng, Tri Dao, Albert Gu, Peiran Jin, Zhao Yang, Yingoe Xia, Renqian Luo, et al. Hybridna: A hybrid transformer-mamte mamba2 long-range dna language model. arXiv preprint Xrr:2200 10807, 2025.
Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer. Luna: Linear unified nested attention. Advances in Neural Information Processing Systems, 34:2441-245 24533, 2021.
Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: moving average equipped gated attention. arXiv preprint orXiv:2209.1065. 2022.
14

---

Published in Transactions on Machine Learning Research (06/2025)
Qiheng Mao, Zemin Liu, Chenghao Liu, and Jianling Sun. Hinormer: Representation learning heteroge- neous information networks with graph transformer. In Proceedings of the ACM Web Conference 2023, pp. 599 610, 2023.
Harsh Mchta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. arXiv preprint ar (rXtv:2206.19447 2022.
James Mercer. Xvi. functions of positive and negative type, and their connection the theory of integral equations. Philosophical transactions of the royal society of London. Series A, containing papers of mathematical or physical character, 209(441-488):415 1909.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint (aoo:160.07844. 2016.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads rcally better than onc? Advances in neural information processing systems, 32, 2019.
Fabian Pedregosa, Gaäl Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, et al. Scikit-learn: Machine learning in python. the Journal of machine Learning research, 12:2825- 2011.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprinl or N1i::2100302144, 2021.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. Tn International Conference on Machine Learning, pp. 25043-280777 PMLR, 2023.
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arxiv preprint arXiv:2108. 0000008..2400 12409, 2021.
Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. ai preprint ar arXiv:2210. 0.10440, 2022a.
Zhen Qin, Weixnan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. arXie preprint arXim:2002.0877 2022b.
Jiezhong Qiu, Hao Ma, Omer Levy, Scott Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self-attention for long document understanding, ar preprint 1r1:v:191..02972, 2019.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Colin RafTel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21 (140)):1-67, 2020.
Yu Rong, Wenbing Huang, Tingyang Xu, and Junzhou Huang. Dropedge: Towards deep graph convolutional networks on node classification. aix preprint ar Xiv: 10903, 2019.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53 68, 2021.
Michael Schlichtkrull, Thomas N Kipf, Peter Bloem, Rianne Van Den Berg, Tvan Titov, and Max Welling, Modeling relational data with graph convolutional networks. In The semantic web: 15th international conference, ESWC 2018, Heraklion, Crete, Greece, June 3-7, 2018, proceedings 15, pp. 593-607. Springer, 2018.
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXi::2022.05002 2020.
15

---

Published in Transactions on Machine Learning Research (06/2025)
Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. Xix preprint 000:2288.04933, 2022.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024.
Yi Tay, Dara Bahri, Liu Yang, Donald Meizler, and Da-Cheng Juan. Sparse sinkhorn atlention. In Interna- tional Conference on Machine Learning, pp. 9438 9447. PMLR, 2020a.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Scbastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. arXiv preprinl ar XXiv:2011.04000 2020b.
Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention for transformer models. In International conference on machine learning, pp. 10183 10192. PMLR, 2021.
Yi Tay, Mostala. Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Compul. Surv., 55),), December 2022. ISSN 0360-0300. doi: 10.1145/3530811. URL https: orreii.rrg/1... 1045// 3530811
Jake Topping, Francesco Di Giovanni, Benjamin Paul Chamberlain, Xiaowen Dong, and Michael M Bronstein. Understanding over-squashing and bottlenecks on graphs via curvature. arXiv preprint 1xXiv:2111.14330 2021.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient. image transformers & distillation through attention. In International conference on machine learning, pp. 10347 10357. PMLR, 2021.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaune Lample. Llama: Open and efficient foundation language models, 2023. URL
https://arxiv http://aaaiiv.oorraa//////////// 13971
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30,
2017.
Petar Veličkovié, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint 1:XX:::710.1900 2017.
Oriol Vinyals, Charles Blundell, Timothy Lillicrap, Daan Wierstra, et al. Matching networks for one shot learning. Advonces in neural information processing systems, 29, 2016.
Roger WalefTe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, et al. An empirical study of mamba- based language models. arXiv preprint arxiv:2406.00 07887, 2024.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. XrXiv preprint arxi::2006. 04 768, 2020.
Xiao Wang, Houye Ji, Chuan Shi, Bai Wang, Yanfang Ye, Peng Cui, and Philip S Yu. Heterogeneous graph attention network. In The world wide web conference, pp. 2022-2032, 2019.
Haixu Wu, Jialong Wu, Jiehui Xu, Jianmin Wang, and Mingsheng Long. Flowformer: Linearizing trans- formers with conservation flows. arXiv preprint arXi::2202.0622 2022a.
Qitian WU, Wentao Zhao, Zenan Li, David P Wipf, and Junchi Yan. Nodeformer: scalable graph structure learning transformer for node classification. Advances in Neural Information Processing Systems, 35:27387- 27401, 2022b.
16

---

Published in Transactions on Machine Learning Research (06/2025)
Tianzhu Ye, Li Dong, Yuqing Xia, Yutao Sun, Yi Zhu, Gao Huang, and Furu Wei. Differential transformer. arXiv preprint ar Xix:224 10. 05258, 2024.
Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, YX Wei, Lean Wang, Zhiping Xiao, al. Native sparse attention: Hardware-aligned and natively trainable sparse attention. arxiv preprint arXiv:2502. 100002..10000 11089, 2025.
Shen Yuan and Hongteng Xu. Sliceformer: Make multi head- attention as simple as sorting in discriminative tasks. arXiv preprint arXiv:2310.1160 2023.
Seongjun Yun, Minbyul Jeong, Raehyun Kim, Jaewoo Kang, and Hyunwoo J Kim. Graph transformer networks. Advances in neural information processing systems, 32, 2019.
Ofir Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat. Q8bert: Quantized 8bit bert. In 2019 Fifth Workshop on Energy Efficiintt Machine Learning and Cognilive Compuling-NeurIPS Edilion (EMC2- NIPS), pp. 36-39 19-39. IEEE, 2019.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural informalion processing systems, 33:17283 17297, 2020.
Chuxu Zhang, Dongjin Song, Chao Huang, Ananthram Swami, and Nitesh V Chawla. Heterogeneous graph neural network. In Proceedings of the 25th ACM SIGKDD international conference on knowledge discovery & data mining, pp. 793 -033, 2019.
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Ré. The hedgehog & the porcupine: Expressive linear attentions with softmax mimicry. arXie preprint orXiv:2402.0434 2024.
Renrui Zhang, Jiaming Han, Chris Liu, Peng Gao, Aojun Zhou, Xiangfei Hu, Shilin Yan, Pan Lu, Hongsheng Li, and Yu Qiao. Llama-adapter: Efficient fine-tuning of language models with zero-init attention. arviv preprint 2r03 16199, 2023.
Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long-short transformer: Efficiintt transformers for language and vision. Aduances in neural information processing systems, 34:17723 17736, 2021.
17

---

Published in Transactions on Machine Learning Research (06/2025)
A Additional Evaluations
A.1 Experiments Text and Vision Datasets
We evaluate our model on two datasets without graph structure: Grouus (Pedregosa et al., 2011) and Mini-ImageNett (Vinyals et al. 2016). The 20News dataset is a collection of approximately 20,000 newsgroup documents (nodes), partitioned (nearly) evenly across 20 different newsgroups. We take 10 classes from 20 newsgroups and use words (TF-IDD) with frequency of more than 5% fea- tures. The Mini- ImageNet dataset consists of x4 RGB images from 100 different classos with 600 sam- ples per class. For our experiment use, we choose 30 classes from the dataset, each with 600 images (nodes) that have 128 features extracted by CNN. Since there is no input graph, we use k-NN (over input node features) for artificially constructing graph for enabling GNN's message passing and the graph-based component. We report the results un- der the best k E [5.20] setup for each GNN base- line. summary of the statistics of each dataset is provided in Tab. g
Table 5: Quantitative results on semi-supervised classification with Mini-ImageNet. and 20News- Groups. We use k-NN (with different ks) for arti- ficially constructing an imput graph. The best and second-best. results are highlighted in bold and un- derlined, respectively, where models with and without graph are compared separately.
Class
Methods GCN
20News-Group
Mini-ImageNNe
65.96:10.88
85.961
Graph-bused (ANN & [5, 20]]
GAT DropEdge
61.06.00.6 44
85.41.0. 13
61.46:00.4
85.81+0.85
IDGL
65.0911.22
65.66.00. 42
LDS-GNN
66.15:10.36
00M
Gumbel-Sofaans
61.71:1.33
87.45220.55
+
Pull-Artn
64.04+0.16
87.46+0.14 54
NadeFormer Framewook (w/o graph)
Local-Att-mm 2
64.54+0.23
86.62+0091
Local-Atin-m =
64.36 4:4
87 07.088+0000 52
Local-Att--m 24
64.61+0..77
87.17+0049
DRI-Atttom to
64.81+00.35
87.37=+0. 80
BRt-Att
65.02-10.48 06
87.4610068
BRI-Aite-m
65.19:10:09
0..5
As depicted in Tab. 5 the NodeFormer with BRL-Attention
(without imput graph and edge loss Le) achieves
competitive performance against its opponents, including GNN-based baselines (Kipf & Welling) 2016
Veličković et al. 2017 Rong et al.) 2019 Chen et al. 2020 Franceschi et al. 2019) and NodeFormer (Wul et al. 2022bb with Gumbel-Softmax// /auxx////////a Lool//Local attention. For the smaller 20News, our method achieves the second best and outperforms the full-attention method, while the local-attention fails to achieve de- cent performance. Notably, the BRL-Attention offers 4x of GPU momory reduction compared to LDS on 20News. On the long-seque seggeeeue dataset Mini- ImageNe,,,,,, the BRL-Attention achieves the best performance among all groups; in this case, 24GB of memory insufficient to run LDS and full-attention due to their heavy computation on learning/ approximatinn global structures. Overall, the experiment suggests that the k-NN graphs are not necessarily informative, and besides, the BRL-Attention can learn useful latent graph structures from data while maintaining a memory-effcient nature.
A.2 Complexity Evaluation on Real-World Textual Dataset
Table 6: Memory usage comparison on autoregressive training and inference under various context length.
Model
Params
128
256
512
1024
2048
3072
4096
Autoregressive train
Full-Attention
50.930 00.93M
3.030GB
3.301CB
3.921CB
5.760GB
12,772CB
23. 869GB
00M(>-2GGB)
Local-Atteation
50.93M
3.030GB
3,215GB
3.563GB
4.07GBB
5.740GB
7,131GB
8.732GB
BRI-Attention
80.96M
3.809GB
3.979GB
4.256GB
4.842GB
6.428GBB
7.948GB
9.613CB
Inference (no grad)
Attenittnn
50.93M
2.260GB
2.260GB
2.264GB
2.266GB
2.590GBB
3.676GB
4.873GB
Locul-Attention
50.93M
2.260GB
2.260GB
2.264GB
2.266GB
2.340GB
2.950GB
3.063GB
BRL-Atteenionn
80.96M
2.397GB
2.401GB
2/44GBB
2,456CBB
2.735GB
3.280GB
3.588CB
We provide results of our BRI-AAttnntion compared to Local- and Full-Attention on WikiText-103 datasct under various context length (block size) E [128, 2096). Specifically, we evaluate the model with 128 since as shown Sec. 3.2 BRL-w 128-m 128 already outperforms the standard Transformer. The memory
18

---

Published in Transactions on Machine Learning Research (06/2025)
(history) length is the same imput length For the backbone settings, we employ model of 6 layers, 8 heads, and 512 embed dimensions. In Tab. G we show the VRAM usage of different models. Tn comparison, the BRL variant scale linearly with n similar to Local-Attention, while requires significantly less VRAM than Full-Attention when n > 1024 on both train and inference phase. In the time efficiency ovaluation Tab. 7 we observe the same trend where the efficiency of BRL outperforms Full-Autention on longer sequences.
Table 7: Time efficiency comparison on decoder inference under various context length.
Model
Params
128
n = 256
72 =
512
1024
n = 2048
11 = 3072
n = 4096
Full-Attention
50.93M
4.24ms
4.26ms
4.40ms
10.28ms
29.19ms
61.95ms
92.27ms
50.93M
4.24ms
4.26ms
4.37ms
5.82ms
9.54ms
13.53ms
18. 18.19ms
BRL-Attention
80.96M
10.26ms
10.26
10.98ms
11.15ms
14.10ms
19.57ms
24.85ms
A.3 On Larger Textual Dataset
Lecal Slintar 228tion
4.0
#6. 200 328
329
Roing
2.0
3.3
1.2
Train Validalion
2.7
332
Evall Stop On 1000)
Eval Step (x (x1000)
Figure 10: OpenWebText train and validation loss over steps. (Ieft) The window size and compression length are both 128; (right) The window size and compression length are both 256.
In addition to WikiText- 103, we study the autoregressive language pretraining on OpenWebText (Gokaslan et al. 2019 Radford et al.) 2019), which is scaled up datasets with ~9B/4M tokens in train train/validation set. We detail the model parameter configuration in Tab. 14
In Fig. 10 we compare the train and validation losses of BRL-attention to local/fullaattention. We employ the sctup of and m = w = 128 and m = w = 256 (< block size 1024), and employ the best practice as Sec. 3.2 where 3 = 0.5 and Atiit = 0.5. Similar to the results in Fig. 7 the BRL-attention with m = w = 128 can achieve comparable performance against full-attention, and BRL with m w 256 can achieve better results against full-attention while the losses of local-attention under both w 128 and 256 deviate. This demonstrates that our method can scale with larger dataset and is effective reducing the attention bottleneck.
B Theoretical Computational Complexity Analysis
B.1 Inference Memory Complexity Analysis
We write for sequence length, B for batch size, H for the number of heads, and L for the number of layers. We let d = dmodel = din for simplicity. Then the time complexity of vanilla Transformers ((3Bnd + BHn?) according to (Kitaev et al., 2020), which briefly O(n?) considering the constant nature of H, L, d. Write m for compressed token sequence length, assume ((x("), W,,, x(W (by assuming the linearity of Φ) for simplicity, our proposed regularizer Fprop then consists of three major parts: (1) Eq. 9 requires respectively Bnd + 2Bmd and Bnm computation of feature transformations and of the
19

---

Published in Transactions on Machine Learning Research (06/2025)
Table 8: Time Memory efficiency comparison of various efficient transformers. N/A entries are for encoder- only models. EEffciiett memory complexity might not equate a faster or more efficient model in practice.
Methods
Memory Compleitt
Complexity
Autoregrossive
Complexity
Matrix
Botclenecked
Tranofomue
((a))
Ol+²)
Explicil
Spare
Trnteormus
Dlax/()
Yeu
Explisii
Patentially
Loagforma Altesiin)
((8)
Dia)
Yes (w/o Global Atta)
O(nu)
Explisit
Potentially
ETC
Oin]
No
N/A
Explicit
No
DigDied
O(n)
Explacit (landoon)
Reformer
O(nlogn)
Oliilgg))
Yes
O(n Knolog n]
Explict
Patentially Potentially
Symthesteer
OlaF)
Oint
Yes
O(n))
Explict
Prefoeuerr
Olox")
Yes
Ol(222) (cumsux)
Implisi
N/A
Linformer
Oui)
No
Implict
N/A
BRI-Forace (Ous)
((a)
O(i)
(1()
Purtially Explicit
n dimensional attention logits. (2) Eq. requires Bmddet for skip connection M and (3) Foomp requires
3Bnda for token encoding, Bmdei for compressed token encoding and 2Bmn for attention logits. In total, the complexity is ((( + Bnm++ Brnde + 3Bdde + Bmndet + BmB))L) R O(n) which shows regularizer Prop scales lincarly with n. Under Thm. 2.9 FBRL = Fgen (:; D) + Frop approximates the expressibility of any continuous functions. Hence choose any sparse D, e.g. the attention with blockwise/ /strided pattern Belltgy et al. 2020), that gives a O(n) complexity [: then FBRL theoretically of approximately ((n) complexity.
B.2 Training Memory Complexity Analysis
Oliss) Cosplesit
Otsm) Complesiy with Carsi Made
Devele
Teadaa Forceg
Haneey Sogmemillick
Sugnrst
Fall
Raceuino
Receptice
(a) Autoregressive Training
(b) Token Generation
To facilitate linear complexity autoregressive training, the compressed token for the current segment/bloo /lock is derived from history tokens in former blocks such that it does not violate the causal structure of the current block. On generation, all in-context tokens are employed for computing compressod tokens.
Figure 11:
One may wonder what hindered us from explicitly formulating Fprop as kernalized self-attention since they are of the same inference complexity. To justify, we have already shown in Prop. 2.7 and Rmk. 2.8 that M2 applied to Fprop gives wider sensitivity bound while offering better robustness on noisy long input. Below, we give another viewpoint that focuses on the memory cost of autoregressive training compared the prominent kernalized approach- Performer (Choromanski al. 2020).
The regularizer essentially approximates the kernalized self-attention under Prop. 2.4 2.6
The complexity of Performer encoder-only training (bidirectional FAVOR) is of O(ed + nd + en) O(n). While as noted in (Tay et al. 2022 Hua et al. 2022). the unidirectional variations cannot be causally masked in an efficient. linear-time fashion. Training Performer for autoregressive tasks, which rely on parallelization and teacher forcing, requires a sequential left-to-right scan similar to RNNs. This makes it significantly slower under the hard requirement for manifesting the c x e KV matrix at every time step, recovering a On?)) complexity model.
For BRL-Attention (with pattern D explicitly set (Balated-)Slidingg Window), the largest matrix constructed via Fgen (+:) D) is the
QK matrix where w is the size of the window. Being compatible
For sliding window attention [Blltagy et al. computation O(Bnd + (.). Supposed y, sliding window attention scales linearly with sequence length, which can be viewed as O(n) complexity.
2020 each token attends to w tokens within the window size, hence por-layer
20

---

Published in Transactions on Machine Learning Research (06/2025)
Full-Atta
(b) Local-16
(c) BRL-16-16
(a) Mini-ImgNet (b) Mini-ImgNet (c) 20News-Grp (d) 20News-Grp (Local)
(BRL) (Tocal)
(BRL)
Figure 13: Feature embeddings (after t-SNE and normalization) and edge connections produced by Local-Attention and BRL- (e) Local-16 (f) 1RI111616 Attention on graph-enhanced application datasets. We mark the nodes with a particular class with one color. The compressed to- Figure 12: Visualizations of reconstructed kens are in black color, red and green lines are, respectivcly, the attention maps. We sample subset of attentive edges constructed by Feomp and F sequences where the attentions from all heads are averaged.
(d) Full-Attm
with causal masking, the memory complexity is reduced from O(n²) to O(wn) = O(n). For the BRL- Atlention components, for Frop, we allow compressed tokens to allend to every token, for Fromps on teacher forcing x, the input x for X(ct] f Femmp (x) the history block of x (where x the current block), which does not violate the causal structure (see Fig. I[a).. In this way, only a causal mask on the main token-to-main token interactions is necessary. Consequently, the largest matrix constructed in Eq. (9 12 is max n) max(d, des), m), which is also of ((n) memory complexity considering mi and d/det constant. Putting together, the BRL-Attention is of & O(n) complexity compared to the (nn²) complexity of kernelized methods in training.
Remark B.1. As displayed earlier, the pattern-base attentions suffer from attention information bottlenecks when distant tokens cannot effectively communicate (Prop. 2.2 While adding in-contert global tokens (Ainslie et al. 2020 Zaheer et al. 2020) potentially alleviates the bottleneck, it violates the causal structure, making them infeasible for autoregressive decoding (discussed in Sec. D.1j. Kernel-based methods typically suffer from the quadratic complexity on autoregressive training (discussed in Sec. D.2h and can degrade performance if the kernel approximation is insufficient. In contrast, our method not only capable of linear complexity autoregressive training (will be displayed in Sec. 3.1 but approximate model that is theoretically expressive as full transformer (Thm. 2.9
C Visualizations
Visualization with Graphs. Fig. 12 displays averaged attention (reconstructed from q, k) matrices on ACM and DBLP, showing how each model variant connects different tokens. In contrast to local attention. which primarily focuses neighborhood blocks and may overlook global structure, BRL-Attention consis- tently places varying weights on distant tokens. In Fig. 13 we visualize node embeddings (via t-SNE) and their attentively induced edges on PONews-Groupe and Mini-ImageNee. The local-only variant assigns fewer inter-cluster edges, often concentrating on nodes within the same neighborhood. BRL-Attention instead in- creases cross-cluster edges, forming additional links that act as pivots for global propagation. By gathering and distributing context, the BRL-Attention with compressed tokens reduces over-squashing, allowing the model to learn better-separated node embeddings.
21

---

Published in Transactions on Machine Learning Research (06/2025)
Butl Ait Scare
(Lac MSS Scar
128 Ais Soae Lock
(8AA) Raconezucied Pagela Store
X5J Fax www.rttteeeeeeeeeeee Aor Scar Aepulacar
134
140
360
140
300
442
400
2000
300
Puul Aut
Guac MES Soiee
jenu Alts Soam Lack
1904 Ancorncructe Paga farteer hoore
3501 dae www.etttteeeeeeeeeeee
200
360
300
102-
120
200
205
Tull Akt
Gucall Score
(2RJ Aim Scrm
(004) Recommroucted Regularierr Score
IBU hecerarrutee dir Scare Ragularon
340
20-
140
2M0
300
100
200
200
1.25
Scores
Tull Aktn Score
(Lacal Scure
(2AJ) Alrs Scom Loca
(00L) Recomeruuttdd Pagularizrr Score
IGU facentrrucdddd Acp Scare Regularan
300
200
300
100
300
Layer Scores
Bull Auti Scare
Scare
(28.) Alss Scom Laca
1856.1 cudd Pags lorter Soore
350 fac wwd.etttd
305
100
180
200
300
Layor Attintton Scores
Pudl JAE salane
Guacil MEE boiae
(enu Alt Soam Lack
1906 Anclnncructed Paya fortorr hoore
3501 dac www.tttteed dOY Scare
130
340
112-
300
50
Figure 14: Attention score visualization (pretrained on WikiText-103). We compare Full, Local, and BRL- respectively, we show the full attention score and the sliding window
Attention. For Full and attention score. For BRL-Attention, we show the score from its local part, the reconstructed regularizer score, and their weighted (by cocfficient X) sum reconstructed BRL-Attention score.
22

---

Published in Transactions on Machine Learning Research (06/2025)
Attention Score Visualization with Textual Dataset Pretrained Model. In Fig. 14 we show the attention of different. Transformer baselines pretrained on WikiText-103. In specific, we compare Full, Local, and BRL-Attention. For Full and Local-Attention, respectively, we show the full attn score and the sliding window attention score. For BRI-Attention, as the rogularizer is affected directly to embedding rather than domain, it difficult to visualize the actual additional information constructed by F To give a indirect measurement of supplemental information flow, we show the score from its local part, the Reconstructed Regularizer Score (RRS), and their weighted (by coefficient A) sum as reconstructed BRL-Attention score. Particularly, the RRS is derived via the score in propagation map, averaging all compressed tokens:
(SRRs) = (S prop Spomp)i.d
(15)
where S i, and Sprop is the n score aotnn (*((=) in Eq. 9 and Scomp is the m x score (gk)))
Va
in Eq. [2], denoting the propagation of information through compression mapping and propagation mapping on main tokens. In Fig. 14 we can observe that for some queries, the keys to attend (vertical lines) align
with that of the full-attention. This indicate that the regularizer is learning useful information from global context, and can properly inject information to the receptive field-limited embeddings.
Final Remarks
D.1 On the Infeasibility of Global Token for Autoregressive Decoding
In autoregressive decoding, at time step f, token x, should only attend to the tokens from the past (i.e., X,..,... xt-1), but not future tokens (i.e., X442, ..,,,.......))).).. This is enforced through causal masking during the attention compulation. Recall the attention operation in an autoregressive model defined as: S.y = fsoftmax /Vd). In autorogressive mode, the attention matrix must be causal, meaning that Vi j j 0. This ensures that the model only attends previous tokens and not future ones. The masked positions Siy are set to zero for future tokens to avoid pecking.
Longformer, ETC and BigBird introduced global tokens to help scale attention for long sequences. These global tokens are selected tokens (in context) that can attend to every token in the sequence, regardless of the token position. For example, suppose we introduce a global token Xg. In the attention matrix for a sequence {x1, X2... XN. xg); the global token x, will attend to all tokens in the sequence: Vj E [1, S., =1. =1
However, in the autoregressive mode, we want each token x, to attend only previous tokens {xi,..., x However, with the global token in play, the model has to compute attention matrix that includes cross- token interactions, the resulting attention matrix for the sequence xx.,, X,} would look like: S = Oxxe
1xx1
where the global token Xg interacts with all tokens, violating the causal structure.
1ixt
1
D.2 On the Quadratic Nature of (Linearized) Kernel Attention for Autoregressive Training
Recall the c x KV matrix in Linearized Attention (Performer, Linear Transformer, Linformer etc.) is constructed by S kT Re-arranging the computation reduces the complexity w.r.t from quadratic to linear. In autoregressive decoding (generation). at timestep I, define S, V,, notice that the computation of S, can be fully incremental, i.e., S, S. + k, v.. This means we only need to maintain a cache with constant O(c²) memory and whenever new input arrives at time stamp t, only constant O(c²) computation is required to accumulate k, v, into S., and get S..
However, on autoregressive training (with teacher foreing), re-arranging the computation in linearized alten- tion leads to a severe inefficiency. Due to the causal constraint for regressgeeesvve training, the query vector at each time step qi corresponds to different cache value S, = k...... This requires the model to compute and cache n different. values ( instead of only one kT in the non-auloregressive mode. In theory, the sequence (8,) can be obtained in O(ne)) by first computing (k, v,)a and then performing a large cumulative sum (cumsum) over n tokens. But practice, the cumsum introduces an RNN-style sequential dependency of steps, where O(22) state needs to be processed each step. The sequential dependency
23

---

full attn always holds, long as the scaling factor is Published in Transactions on Machine Learning Research (06/2025)
not only limits the degree of parallelism, but more importantly requires n memory access in the hard loop, which increase the complexity to quadratic.
D.3 Differential Form in Compression Mapping Improves Attention SNR
To justify the impact of M2 on Signal-to-Noise Ratio (SNR) in attention mechanisms, we first categorize the keys into two sets relative to given query: relevant keys and irvelevant keys. Let k = k2.... k,, be
the set of all keys, Let Kyel S k be the set of relevant keys, and kir = k Krel be the set of irrelevant keys. Let = S2,...,.... S,. be the attention weights corresponding to keys in k, where s; is the attention weight for key k;.
Definition D.1 (Attention SNR). We define the Attention SNR as the ratio of the average attention weight assigned to relevant keys to the average attention weight assigned to irrelevant keys:
Tattn
S1,
(16)
Kkrel
Kir
ky Ekaal
Ii Ekerr
where | and [Kree the number of relevant and irrelevant keys respectively. If [ - 0, we consider SNR to be infinitely high, indicating perfect attention. II keei - 0 while |kin > 0, SNR is zero,
indicating no signal. A higher SNR indicates a better ability of the attention mechanism to focus on relevant information while suppressing irrelevant information.
Attention SNR full and differential Feomp (X(et): Fxomp(X(et)] 8 is constructed by 9[ct) E Rexd and k E Raxd where the former is from compressed tokens and the latter is from main tokens. Denote the similarity scores in attention S as sim k), where a. = sim(q(et)) k;).
x). Ignoring the feature mappings in Eq. 12 the
Attention weight S is then formulated as s م expla")) Now, assume there is only one relevant
exp(a.")
= (where is the base similarity and 8 is small key and one irrelevant key, let a = + and air perturbation). Then the SNR is expressed as
K
exp(x+d)
full attn
rel K
+ exp(8).
(17)
exp(a)
Sir
exp(a+0)+exxpa)
For the differential case, have attention scores SK1 and S respectively constructed by similarity scores a sim(qrec,,k and a K2 sim sim(qret),k'' Similar to full-attention, let a a + 6, a which exactly mimic the situation in full-attention, then let R2 = + 81, a + 82 where 8(1,2) are learned
shift parameters by M2. Then with s/x
exp(a," we have Σ,= exp(af))
exp(8)
exp(61
diff attn
Syl
K2 K2
7/5
exp(6)+1 - exp(d (x)(()2)
(18)
K
K2
exp(
7/
exp(4)+1 7 )xx((1)+oxpp (xxp((b))
Proposition D.2. Under 81 -8 S 5, posilive.
diff astn
2 exp(5) = r
Proof. Proved in Appendix E.6
Therefore, improving the atlention SNR suffices by learning learning a differential form where d - 82 < 8. Compared to the fixed SNR in full-attention, the differential form as Eq. (12) offers more flexibility in countcracting the attention noise. For the propagation mapping Fprops since the keys are constructed by compressed tokens X(ct) which assumed to be permutation invariant, it would be meaningless index the noise, hence requirement for noise suppression.
24

---

Published in Transactions on Machine Learning Research (06/2025)
D.4 Additional Discussions Recent Works
D.4.1 Relation to Native Sparse Attention
Native Sparse Attention (NSA) is a recent efficient attention mechanism proposed in Yuan et al. 2025 Specifically, NSA can be classified as learnable pattern-based attention that modifies the Fgen(;; D). As F gen is independent to FBRL, NSA is fully compatible to our method. We omit the comparison here since NSA requires compiling fused-allention kernel and is only optimized for hopper GPUs (e.g., H100).
D.4.2 Relation to State Space Models
Recent work on deep State Space Models SSMs, most prominently S4 (Gu et al. 2021 and its successors S5 (Smith el al. 2022), DSS (Gupta et al. 2022). and Mamba (Gu & Dao 2023), demonstrates that learned state state-space layers can model dependencies across tens of thousands of steps while retaining O(z) complexity. These models encode a sequence by evolving a latent state according to linear ime-innvriintt dynamics (plus lightweight nonlinearities). and can be executed either as fast FFT-based convolution (encoder style) or recurrent update (decoder style). Empirically, S4 established new state of the art LRA and other long-sequence benchmarks and outperforming many Transformer variants.
Architectural Contrast with ARL-Attention. While both SSMs and BRL-Attention target the quadratic space/time bottleneck, their mechanisms are orthogonal. BRL-Attention keeps the content-based routing of self-attention but reduces cost by mediating global exchange through a fixed, trainable pool of compressed tokens, thus preserving the Transformer's dynamic contextual interactions and autoregressive training. In contrast, SSMs dispense with attention entirely: information mixed through fixed kernel convolutions or state updates whose influence depends on position rather than token content. This design gives SSMs excellent memory compression and extrapolation to sequence lengths far beyond training, yet it lacks the explicit query-key reweighting that underpins in-context reasoning and fine-grained token-to-token alignment in attention models.
On the Fairness of Direct Comparison. SSMs summarise the entire past into fired- width latent state, so the model cannot re-query earlier vokens on demand; once detail is not encoded it is effectively forgotten. This makes SSMs ideal for tasks whose signal is distributed and uniform along the sequence (c.g. long-range copy, audio, or algorithmic path-finding) but less suited to problems hinging on selective, content-dependenl retrieval, such as long-contexi QA or few-shot in-context learning, where a loken must attend to a rare, distant cue. Att-Attention retains query-ke key similarity scores (vía compressed-toeen bottleneck) and full causal masking, enabling sparse but crucial interactions at generation time. Moreover, since BRL's goal is to close the gap between full self-attention and pattern-based linear attention, direct comparison to fundamentally different, paradigms like SSMs can be misleading; we therefore reference SSM results as context but focus quantitative evaluation on attention based baselines.
D.4.3 Relation to Hybrid Architectures
Our proposed method (as illustrated Fig. can be viewed as a mixture of two paradigms within a single layer, f.e. the gencralized attention section and compressed token update soction, which relates this paper the field of hybrid models. Hybrid models generally aim achieve a better trade-off between performance, computational efficiency, and cong-context modeling than either approach can achieve alone.
Transformer-Recurre Hybrids. This category uses Transformer attention the primary building block and SSM/RNN parts auxiliary. For instance, Griffin (De et al. 2024] mix repeating recurrent blocks (similar to GSS (Mehta et al. (022 with local/global Multi-Query Attention (MQA) blocks across different residual blocks, matching the performance of Transforoorr despite being trained on over 6x fewer tokens. Similar models, e.g., (Cao et al. 2024) combine LSTM layers with Transformer parts to get good soquence handling and context understanding FLASH (Hua et al. (022) split the context into chunks, letting full transformer to process each chunk and RNN to handle chunk wise relationships. However, bringing in sufficiently many Transformer layers will also introduce quadratic complexity. Moreover, while
25

---

Published in Transactions on Machine Learning Research (06/2025)
RNNs efficient for inference due to their sequential state updates, this very sequentiality makes training parallelization across time steps challenging.
SrA Hybrids. These models are fundamentally built upon SSM or long convolutio prin- ciples but may incorporate attention-like features or be used alongside attention in broader architectures. For instance, Jamba (Lieber et al. 2024) alternates between blocks of Transformer layers and Mamba layers. A Jamba block has 1 7 ratio Attention-to-Mamble layers, which is proven to achieve optimal performance in the NLP domain (Waleffe et al. 2024]. This setup lets Mamba manage long sequenees of information, while the attention layers focus on specific token-level interactions. Simialrly, HybriDNA (Ma et al. 2025) also insert 1x of Transformer block into 7x of Mamba blocks for effectively balancing the advantages of both block types. Hyena (Poli et al. 2023 instead of using an SSM approach like Mamba, hybridizes implicitly parameterized long convolutions with data-controlled gating (mimiccinn the input-dependent na- ture of sell-attention). In essence, these models are lypically good for modeling extremely long sequences. However, fixed state memory can be limitation in practical applications: once token is processed, it's either included in the state memory or ignored, making it challenging to re-query an ignored token. This results in a lack of precision on lasks requiring strong associalive recall, e.g., copy tasks.
Attention-Centric Hybrids. The BRI-Attention modifies the allention operation itself to become glob- ally aware and efficient, rather than replacing entire Transformer blocks with distinet SSM RNN blocks. It avoids vertically stacking different layer types, which can be computationally heavy as each distinct full- sequence layer that processes the entire sequence embedding. Instead, BRL-Attention augments base linear generalized attention mechanism with a global context mechanism (the compression-propagaaiioo dition incurs only marginal computational cost (as shown Sec. A.2). To enable such guidance of global information, BRI-Atlention must employs distinct stralegies for training and generation in auloregressive scenarios. During training, the compressed tokens are updated using an attention- based method that learns to retrieve information from the available context, including memory buffer of past tokens. During gen- eration, this mechanism allows for step-by-step processing where the compressed tokens act evolving global summary, informed by the incrementally growing context, thus maintaining causal consistency.
This ad-
D.5 Discussion on Limitation
While BRL-Attention offers distinct advantages for processing long sequences, its relative benefits dimin- ish for inputs of shorter length. For instance, Vision Transformers (ViT's frequently operate on sequences composed of a limited number of patch-based tokens (e.g., approximately 197 tokens for ViT-B /16). Tn such scenarios, the computational overhead of standard quadratic self-attention is often manageable. Con- sequently, the introduction of BRL-Attention's compressed tokens might yield only marginal performance gains while potentially adding system complexity. For sequences comprising only few hundred tokens, con- ventional full self-attention or alternative, simpler efficient attention mechanisms may prove more practical.
Additionally, an empirical evaluation of speech processing tasks was not undertaken in the current work. State-of-the-art munsformer-based models in speech recognition and synthesis typically rely on extensive pre- training on very large-scale speech corpora. Such pro-training endeavors demand considcrable computational resources, often spanning weeks or months on multi-GPU setups, which were beyond the scope and available resources for this study.
E Proofs and Derivations
E.1 Proof of Prop. 2.2
Proof. Recall the update of token according to the sliding window attention
sw 11-
sew (1-
(111-)
(19)
=
kj i1
=1
26

---

Published in Transactions on Machine Learning Research (06/2025)
which can be expanded as
SaW Mi
S Sinz j2
W W (
(20)
J==1
Saw ij
S jn (x j2
11-2 2) 2
(21)
J2m1
Expanding until Xp (0)
we have x;
equal to
Σ
Saw Saw
S S (x)
(0) W! w(/-)).
(22)
M1 Jii
f1=1
==1
Therefore, the derivative can be expressed as
0x (1)
Saw S
Saw
-1- (W ... w((-D)
(23)
0xp
(0)
Jo--1-111
Σ
Saw II
Rhh1 ()))
(24)
all paths from to i 1=0
= -(S5w) 1ip
w
(25)
1=0
where the bound is
0x (2)
(1((S hxpp
w!")
(26)
0xp= (0)
L-1
|I((S5w))
|II
(27)
1=
S Srourw
(28)
which suffice to derive how many layers are needed for token p to reach For Saw the maximum direct w+1 neighbor distance dmddad 2 1 w-1 where recall w 2 3 is odd number window size. Essentially,
the dmudnd tells how far away we can "jump" up to from the current position in single application of the
adjacency. For token i to commute with j in the L-th layer, it must satisfy
2M
Ldmdud 2 M = L
(29)
W 1
otherwise, (s will be 0. Since L is an integer, we take Lmin = bound. Hence
2M
for obtaining the non-zero
Ox (-)
0
if
2M w-1
<
(30)
0xp
sw
T
if
L
w-11
2M
Replacing L by I ++] concludes the proof.
E.2 Proof of Prop. 2.4
Proof. For softmax kernel, a choice of approximator & is the drim defined as Eq. C1 and C2, the first input to Fprop ②(q) xs @rem (q) and the second input is
10
With Eq.
3
under
x
(11) [ct]
= Σ
Deta (k (K flatten Detu (k. (K(2) vv
(31)
j=1
27

---

Published in Transactions on Machine Learning Research (06/2025)
in Ruxda R mx Then, instantiating the Frop
o(q "d (exxand(x((e]]
prop (x x
(32)
((q [xtet] [ cl
orfm (q (1) Jexpand (eattan( Prem (k) Tv!
(33)
vrtn(q Ej= Or (k) (kg)) DT
Örfm (q ()) プラー Drem (k)
) Tv
(34)
Drim (q Ej= Drim (k) )T
gives the kernalized attention equivalent to Fkernel.
E.3 Proof of Prop. 2.6
Proof. By assumptions on input feature X' and on transformation |W]| we know that for all token of index i, qi ki, Vi lies in compact domain. As each component of Φ is continuous, can be approximated arbitrarily well by MLP with (1)) width and depth (Cybenko 1989)- The contimuity of ø also implics that Φ(q:), = p(k)) Tv, v lies in a compact domain, therefore the numerator lies in compact domain. Lastly, since all operations do not involve n, the depth and width are constant in n.
E.4 Proof of Prop. 2.7
Proof. We perform sensitivity analysis on the output of BR pxoo((x!) X ((t))); of token i after (1)-th layer propagation with respect to token & by
F(U) prop (x(1)
,
x[o]]] >x|l = a(F
p
Ox!
a(F prop);
8x
0x (1+1)
tetl
(35)
(1) 0xp
Ox!"
fx
Ox
fot)
11+1)
0xp (!)
aFF (1)
aF (x(!).
xxxt (
prop):
0 +
(36)
0x
[ct)
(4+1)
0x
Term T1
Term T2
For term T1, let d(q) xW di simplify the formulation of Fprop in Eq.
[ peopl: smm( 1
W
(1+1) [et]
Wz)
(1+1) )x [ct] Wz,
(37)
where sm(+) denote the row-wise Softmax function. Therefore, term T1 can be simplified
F[ 0 propl: (1+1)
Osm
[ct]
W (++1) W.) [xt)
T (4+1)
(38)
((11)
Ox [ct]
(x(]
Osm
W (x
(x(tt)) (+)))
W.ll
(39)
0x (1+1)
[ct] [x(c)]
Let Hg va x (1) W. RIxdet we have
a (( (rop))
Osm (x [ct)
(1+1)
W. T) HH2 (1+1) W
W.l
(40)
0x (1+1)
(1+1)
(1+1)
[ct]
(ct)
[ct]
-10(1)
!1) W W.W.W.
(41)
00(" To'
(42)
vdd
28

---

Published in Transactions on Machine Learning Research (06/2025)
For term T2, let us first consider Feomp to be in simple cross-attention form (with scaling factor changed from vd to medet)
FD = comp
sin
(1+1) [ct]
W (x T
(43)
my
then similar to the derivations above, let ((6)
[++1
W Rmxx, have
mvda
[ct]
aJ comp (x(1),
x [ct]] || @sm( ((()) (s!! WX
(1)
(44)
0xp
Os
0xp
(1)
[1(m)) 1
x let]
(1+1
w ( T
(45)
mydet
S 0 m rarti
= 0 ("etrir)
(46)
m vdea
Vdet
Hence with Eq. 46 the bound Eq. 36 is eventually
prop (x
(x
x >ill S 0
[et]
FeTart Vddet
(47)
0xp
(1)
Now, if we consider Foomp to be differential cross-attention form, briefly defined
F(I) comp
(sm(s >ysm (S (5 x2))x (Wk.
(48)
where (0) K1
WX" and similarly for S Then the bound can be derived as
myda
aF
comp (x(1,
(1)
x (
= 0sm (5(t)
as!!)
dsm(s
as!
[ct)
K1
K2
W
(49)
0xp
(1)
as
Oxp
(1)
as
OXP
(!)
Osm (1))
K1
as
Kl
W || + 1/lll
dsm(
K2)
os
K2
wXI
(50)
a
0xp
(
K2
0xp
(1)
< 0 Tet'
(1 + hn)),
(51)
dd
plugging in Eq. 36 gives
|| (1)
Oxp (1) prop (x(1)
xx ((t)))al S 0 TxTetr W
(1 + In))-
(52)
vddct
This concludes the proof.
E.5 Proof of Thm. 2.9
Proof. We firstly define the star-graph:
Definition E.1. The star gar-graph S centered at 0 is the graph defined on no. The neighborhood of all vertices is A(ii) (0, i) for E (I,...,n] and N((0) [1,...,n] n).
Define imput x = {xo, X x,} E R(++1)xd where {x] are the main tokens and x is the center token introduced by star-graph S. Now, we define a simplified version of BRI-Attention using softmax attention and without differential form, where the compression mapping
X = F Feomp (x) = En exp(dtet]k) )v;
(53)
ch exp(d
29

---

Published in Transactions on Machine Learning Research (06/2025)
fv(x). Under column-wiit softmax, trivial propagation mapping can be where q = ffx==xx(xx,))))) fx(x), expressed as
exp(qik) EkEND(i) exp(qukT) V.
Frop (x, x);
(54)
where i E n... Notably, both Feomp and F prop does not require the usage of Xo. Then for
FBRL Fege (XI...n D) + proop(x x) E Ruxd
(55)
we have the following proposition.
Proposition E.2. For any pollern such that Dn = 0, the Fgen(Xo..) (Xo..n: n; D U (US)).... can be simeulated by FBRL arbitrarily well, where DUS can be regarded as any graph containing star-graph. S.
Proof- Recall the generalized attention which computes
exp(q.k))
F0 (x; DUS):
jENous(i) ERENDUe(1) exp(q:kk) Vi,
E
(56)
for pattern D.S. Recall the definition of star-graph, for any token i E (1....,m) n) the neighborhood under DUS is (11) = U We now write the generalized attention output for index E {1,..., n. under the pattern PUS S as
Fgen (x; DUS); exp(q:kk)
exp(q:k))
Vo
Vj,
(57)
Zi
Z,
JEND(i
Term
Term T2
where
Z,= = =xp((qkkk))))
Σ
exp(q:k)
(58)
(EAD(()
We use q, v instead of q, v to highlight that they are generated by different neural networks. Notice that T2 is exactly the contribution from pattern D, which is equivalantly Fgen (x; D) = Fgen (x1. mi D) as D Dns= = 6. Therefore, it suffice to show that T1 can be simulated by Fprop (x, X); for all i E (1.....n.)
Observe the difference between T1 and Eq. (54) that the numerator exp(q:kō) T1 can be easily simulated by exp(q,k) as k aggregates the information from X(ct] (which can be regarded as xo) by compression Eq. (33, and one may simply let q, = q. On the denominator, as exp(q,k,) ) can be simulated, it suffice to show that EkEND() exp(q:k))) can be simulated by ΣLENO exp(quk] According to the property of inner product, qk' produces a scalar, thus qk' = kqq. As Vk E (ND(1): : Kş can be exactly simulated by let Vk E (N) q: = kk, the problem reduced to proving k can simulate any q. Notice that by the design of compression Eq. (53) X aggregates information from all tokens {x,) thus may design k fx (x, i) simple decoder, which receives the i-th token as input and decode the corresponding token out of x. Concluding above, under Xtct] in Eq. (3 (which eliminate the use of xo for star-graph), we have shown that F(!) F (x; DUS) can be simulated by FBRL Fgen D) (::
Note that, when only main tokens {x:} are employod (just as the BRI-Attention or any regular allentions), we naturally have D n S = 0. Therefore, Prop. E.2 works for BRI-Attention with any pattern D. Next, leverage the result in (Zaheer et al.) 2020 to complete
proof.
Thcorem E.3 (Thm. 1 (Zaheer et al. 2020g). Given S < 2o and 0, for any continuous func- lions Fron [0, 1] 1xd + Ruxd, there exists a transformer with sparse-allention, do(Foon: Fgen (: : U US) < e where DUS is any graph containing star graph S.
(UDU) such that
Combining the result from Prop. E.2 and Thm. E.3 we arrives at our conclusion that d (Foon: Feen (:en::D) D) Fprop) & with arbitrary sparse-attention D.
30

---

Published in Transactions on Machine Learning Research (06/2025)
E.6
Proof of Prop. D.2
Proof. Given
exp(a+-))
K
attn exp(a+d)+expla) exp(8),
full
$ rel
(59)
exp(a)
exp(a+))+exxla)
and
K2
exp(o)
exp (6))
difff
Pel
78 mll
exp(6)+1
exp(á )+ (2))
(60)
actn
K2
exp(62)
Sire
exp(6)+1
exp(
let A = exp(6) + and = exp(61) + exp(62); then the SNR for differential form can be simplifiod as
exp(8)
(xp(8.)
Bexp(0) - 7 exp(())
diff attn
B
(61)
(xp
B 7 exp(02)
When pii 'ttn full atta:
2
assume 7 > O, we essent.ially have
B exp(() - 7 exp((1) 2 exp(1)
(62)
B exx(((02
< -7Aexp(0) 2 -7A - exp(o) exp())xxp(((
(63)
= exp(61) S eexp(ō + 62)
(64)
81 Sd+
(65)
which concludes the proof.
E.7 On the Lower-bound of Jacobian Bottleneck
We show this section that under certain assumptions on singular value of weight matrix, the lower-bound of the Jacobian bottleneck Prop. 2.7 is positive. Below, we give a proof on simplified model architecture.
Proof. We aim to demonstrate that the Jacobian corresponding to the regularization path in BRL-Attention can have a non-zero lower bound in general cases, thereby alleviating the information bottleneck observed in purely pattern-based based attention mechanisms. Recall the Jacobian bottleneck established in Prop. 2.2 for pattern-basod attention Fgen: the proposition implics that when <<
[ 2M 1, the Jacobian through
85 BRL ox,
a
Fgen zero. Under these conditions, the Jacobian of the full ARt-ntttntion is JBRL
ax,""
pe Next, we make the following assumptions to facilitate the justification that J/bLL 0: (JBRL) ox!")
(1) We assume 3 = 0 and À for simplicity. (2) Non-singular weight matrices: That is, learnable weight
matrices W (e.g., W W ho WV. in Feompi w2. wk Wy in Fprop) along the computational path have
their smallest singular value Tmin (W) 2 sw 0. (3) We assume for all 0, x() and (x(O)). are unit vectors.
Consider the Jacobian for l-th layer: (JBRL)sp = 0( 0x ())
By Proof. E.4 we have
(Ji)x = 8(Fobop) (1+1)
DsnnHH(((x[)) ((+1) )ı W.) OHg
(1+1)
[ct]
W, = 0(1)
W à W W...
HH, ((2((x(()) )x W (x
(1+1)
[et]
(111)
vd"
(66)
31

---

Published in Transactions on Machine Learning Research (06/2025)
By aforementioned assumptions, let We W. W...k for simplicity. Recall that for any matrix M, its smallest singular value mmin (M) is defined as: minn (M) min |= [uM]]- As x; is assumed to be unit vector, we have
[I = |xx 0(1) 1
W. W W2 w21>
Xi Xiinomin (W))
SW
|xil|-
(67)
vd
vd
[(x 0x,
0(x(1=11)
F
comp (x x
Therefore, can say |1
0. Next, define (J2)kp
then with
vd
similar procedure, we can derive that || 2 vd since (Xtet)) is also assumed to be unit vectors. Therefore, according to the chain rule, we have
ms > 0. 2
[[JBRL)) - Σ (J2)a
(68)
This concludes the proof.
Although given the above proof, we still note that establishing universally strictly greater than zero lower bound for sensitivity (information propagation) is generally problematic without any assumption, even for standard full-attention mechanisms. The actual flow of information, and thus the Jacobian, is inherently dala-dependent and context-specific:
Data Specificity: If, for a particular input sequence, there is no meaningful semantic relationship to be captured between two specific tokens i and p, the attention scores connecting them (even in a full-attention model) might naturally be zero or near-zero. In such an instance, the Jacobian representing the influence of token p token would also be zero, leading to zero lower bound.
On Full-Attention w/o Bottleneck: Even in a Full-Attention Transformer, if the attention mechanism learns that token i should not attend to token for a given context, the corresponding attention weight will be zero, and thus the propagated information (and its gradient) through that specific path will be zero.
Therefore, a lower bound sensitivity would likely still be nearly zero in many valid scenarios. The crucial aspect of our BRI-Attention is not to guarantec minimum amount of information flow between all token pairs at all times, but rather to remove the structural certainty of zero flow imposed by sparse patterns in shallow layers. The non-zero upper bound provided by BRI.-Attention demonstrates that informatio ion can flow under bottleneck conditions described in Prop. 2.2
F Datasets and Parameters
Table 9: Statistics of the image/text classification datasets.
Property
Datapoints
Classes
Mini- Ming-emageNNet
Image clsestfication
no graph/le /N-NN graph
18,000
128
30
20Nows-Gronge
Text classification
graph /k-NN graph
9,607
10
Table 10: Statistics of the hetcrogencous graph datascts.
Dataset
Nodes
Node types
Edges
Edge types
Target
Classes
DBLPP
26,128
239,566
author
21,420
86.642
movie
ACM
10,912
547, $47.872
paper
Freebase
43,851
1
151,034
6
movie
32

---

Published in Transactions on Machine Learning Research (06/2025)
Table 11: Training hyperparameters Standard Full-Traassi Transformer in the original LRA paper. Settings are derived from (Tay al. 2020b).
Depth
Pooling
Listpp
512
8
1024
32
C1.8
Se-2
Test
1024
32
CI.S
5e-2
Retriewd
128
4
512
32
CLS
Se-2
bonge
128
8
128
256
C1.8
3e-
Pathiiuder
Path-X
32
4
32
32
C1.8
1e-3
Table 12: Training hyperparameters of BRI--Former and Transformers tagged with out-imp the LRA experiment. Most settings are similar that of Vanilla Transformers in (Amos et al. 2023
Depth
Pooling
Littpp
512
8
1024
G4
Monn
10-4
Text
1024
64
Mean
lc:4
5e-d
Imngee
3
64
8
128
16
Max
Se-4
Puthfinder
Menn
Se-d
Path-X
128
32
Max
Se-4
Table 13: Training hyperparameters of BRL-Pormer in the WikiText-103 modeling. Our settings similar to that of Vanilla Transformers in nanoGPT.
Model
Tmplementation
Depth
Num Heads
Hidden
FF size
w
m
block size
memory (queue length)
BSZ
T.R.
Baseline
Full
Qia et al.
2022a
8
512
2018
512
128
Se-d
Full
8
512
512
512
16
5e-1
Our-imp
Local
/anooPP
8
512
512
64/128
512
16
Se-4
BRL
8
512
512
61/128
61/128
512
512
16
4
5e-1
Table 14: Training hyperparameters of BRI-Former in the OpenWebText modeling. Our settings are similar to that of Vanilla Transformers in nanoGPT.
Model
Implementation
Depth
Hidden
block
(queue length)
512
512
4
1021
5
5e-1
nindPT/namoGPT
Se-4
BRL
6
512
512 4
256
256
1024
1021
16
Ge-1
33

---

Published in Transactions on Machine Learning Research (06/2025)
Algorithm BRL--Pormer and its Constituent Functions
Given x (data, Rnxd), Xtet] (initialized as mn. Embeddings, I mxxd) and Xhist (initialized as 0, 1
function BRL__ORMER(X, X[et] Khists causal)
if not causal then
Xhist <- wwwcotoccccccccccco (Xhist, [-hinst :] 4: end if
5: for i= 1 to L do
6:
x, [[et] t [rl_block(x, x[ct], Khist)
end for
if causal then
Xhist f- concat(Xhist; [)[-nhistt :]
10: end if
11: return x, Xhist
12:
end function
1: function BRL__LLCK(X, Xfet) Xhist)
2:
x', "'x('t) <- FBRL (x, X(ct); Xhist)
3: h 1 x x
h[ee] f- X[at] + [ct
Xout h + ffn (tnneaayr_aaae (oryer_)
[[ct] h[ct) + ffn(eayer_aaaaa ffn(Layer_nom(((aa)))
return Xout, x
8:
end function
1: function FBRL (x, [[t]] Xhist)
2: Nout f Fgen (x; D) b Linear generalized attention. We employ Local-Attention with pattern Daw-
3: X[ct] f- X[ct]
4: Xhist is not 0 then
[[tt] +- (xom)(X(ct)) ?)
Output in Rexda
end if
[[t]] (1 (el]layrr_nom({X6
BB (x(ct)) (ct) A)))
8: Xreg f- rrelu (layer_noree (layer_norm(Fooot
Output in Roxd
Nout +- Xout + X2eg
10: return Xout [[c]]
11: end function
1: function (xomp(((et) Xhist, 7)
2:
D Split into R xm Xm (des hKl k2 - split_headss heads f KL,2 (Xhis): ) H) & Lin projects dat Each split into H x (des Lin projects to det. Split into R x x (de//H)
h 4- split_heads(( (x(et) H)
h + split_heads ( (Xhist), H)
SKI, k Jaun (h hh Kl
Jattn (h 2 (h*))T)
In R xm Each head-wise embodding has shape in T (dea// b Group norm and reshaped into Rmx
x out,h
[ct) < (SKIA 78 K2, K2,
Vi
7: [(tt] <- poshape(group norm( ( corm(fx (atg ) )=))
8: return xie]
9: end function
1: function Fprop X(at), A)
2:
z +- split split_heads((__e
H)
In shape xxx B Lin projects to d. Output shape xmx In R Mxuxm Each head-wise embedding has shape in (d/HH) Reshaped into Boxd
3: z²,²' + split heads ( JKV (xtet)), H)
St T-um (z (z²(z))TT)
xh X
Xout. reshape (x 1
7: return X.nt
8: end function
34


Reviews:
[{"review_id": "3mzJHyqTwv", "reviewer": "Reviewer_Pfto", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n1. The paper addresses an important challenge in the field: finding alternatives to traditional Transformers that maintain performance while reducing their quadratic complexity problem.\n2. The theoretical foundation is robust, with analysis of local attention limitations and kernel approximation techniques that directly inform the model's architecture.\n3. The empirical validation is comprehensive, comparing BRL-attention across multiple modalities against several efficient Transformer alternatives. The results demonstrate consistent improvements over standard attention mechanisms, particularly in handling long-range dependencies.\n\n\n**Weaknesses**\n1. The architecture is rather complex, which has several consequences:\n  - The intricate design may deter other researchers from implementing or building upon this work.\n  - Key illustrations (Algorithm 1 and Figure 4) lack clarity and accessibility.\n  - Having this many components may increase the computational overhead despite linear scaling (more detail on that in weakness 3).\n2. The theoretical section would benefit from clearer connections to the practical implementation. Currently, the technical results appear overly technical from the architecture choices, and fail to make the motivation of the design clear to the readers.\n3. While the authors study the memory and the speed of BRL-attention on a toy task, it is not clear whether the same trends still hold in the experiments of Section 3.2 to 3.4. Given the modest performance gains over standard attention, it is important that the costs of BRL are much lower than the ones of attention for it to be of practical interest.\n4. The relatively small scale of experiments raises questions about how BRL-attention would perform in large-scale deployments, where the advantages of Transformers on more efficient alternatives tend to become more pronounced.\n5. Comparison to the deep state-space model literature is missing. It tackles the same problem as this paper, has been thoroughly investigated in the last few years, and achieves competitive results. For example, S4 significantly beat all the methods reported on the LRA benchmark.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces BRL-Attention, a novel attention mechanism for Transformers that achieves linear complexity while maintaining strong performance. The authors combine local and global attention to tokens that compress context, justifying this design with a theoretical analysis of attention layer expressivity.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I would like to see all the weaknesses mentioned above addressed before recommending acceptance. In particular, this means:\n- improving the presentation of the paper, by making the motivation of the architecture and its description more accessible.\n- providing experimental evidence that BRL attention is much cheaper than standard attention on the different tasks and that there are reason to believe that it will scale decently well.\n- integrating a thorough discussion of the deep SSMs literature where adequate.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "okOEIvnNPA", "reviewer": "Reviewer_jtaY", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength: This paper proposes a new architecture that combines sparse attention and linear attention. The regularization term novelly combines the feature mapping used in linear attention and a set of carefully-crafted compressed tokens, to compute a small cross-attention term. This helps to alleviate the issue of long range dependence issue called by sparse attention, particularly the variant of sliding window. To theoretically capture this notion, the paper develops a notion of bottleneck which is defined as the Jacobian between two points across layers is approximately 0. They prove that sliding window indeed causes bottleneck unless the transformer has large enough attention layers. The regularization is introduced to address this issue, and they prove that indeed the regularization term does not introduce too many errors, and the BRL attention approximates the vanilla full attention. Experiments are performed to verify the new architecture works well in practice.\n\nWeakness: There are two main weaknesses. \n\n1. The connection between regularization and alleviation is not strong enough. The main selling point of the paper is that, since pattern-based attention is provably bad at capturing long range dependence in terms of Jacobian bottleneck. Hence, the regularization term is very carefully designed to address this issue. However, the only part mentioning the *theoretical* effect of regularization is a very short paragraph saying that the regularization term makes the sensitivity bound nonzero, and this implies that the information is more likely to flow from token $i$ to $p$ across layers. This connection is also not made clear in the appendix. It is worth noting that most theories in this paper are explanatory and not very deep, they mainly serve the purpose of providing explanations on the new architecture, so it's crucial for authors to articulate clearly how the regularization term helps to alleviate the bottleneck.\n\n2. The structure of the paper is a bit messy. Authors attempt to compress many contents including theoretical and empirical into 12 pages, so I do understand the difficulty, but this results in the paper requires multiple re-readings in order for me to fully understand. For example, I feel Section 2.4, **Instantiation of Information Compression** could be further condensed, as the main goal is to show that the regularization term is not too large? More for this part in requested changes.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper considers a new attention mechanism called BRL-attention, with the goal of combining the strength of both pattern-based sparse attention, and linear attention. The major issue of pattern-based attention, especially the variant of sliding window, restrict long range dependence between two points after a certain number of attention layers. Linear attention on the other hand has a significant gap between theoretical and practical runtime due to limited parallelization. The proposed method is a variation that allows pattern-based sparsity for attention, while ensure the long range dependency to be captured via a regularizer term. The regularization crucially uses feature mapping, as it selects a subset of compressed tokens, compute the feature mapping as the queries, and project the compressed tokens as keys and values, then compute the attention with these components.\n\nThe main theoretical result is that, if we let attention bottleneck be that the Jacobian at the $l$-th layer of the $i$-th token with respect to the $p$-th token at the initial layer is roughly 0, then one could prove that the sliding window is bottlenecked unless the number of attention layers is large enough. Under the assumption that the feature mapping used in regularization approximates the random feature mapping and some requirement on the embedding dimension of the compressed tokens, then the attention between feature mapped queries and compressed tokens key/values actually approximates the self-attention module. Assuming the tokens and weight matrices are bounded in norms, they subsequently show that an MLP with $O(1)$ layers could approximate random feature mapping, providing a practical solution to design the feature mapping. Finally, they show that under proper scaling, the cross attention term for the regularization has small norm, and the BRL attention can approximate the full attention. \n\nEmpirically, authors show that their algorithm is memory and time efficient, and obtains good metrics in various tasks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Major: \n\n1. This paper has a lot of notations that are used before defined, or are never defined. For example, in Eq. (5), LN is used without ever defined, I think that's LayerNorm? Also, what is $\\sigma$ and $\\odot$? I assume $\\sigma$ is ReLU, and $\\odot$ is elementwise product? While these notations are introduced at the beginning of Section 2, but they are only explained in Section 2.3, and $\\odot$ is never explained. Authors should at least include a preliminary section in the appendix, or explain these notations directly after they are used. \n\n2. I think authors should not mix math notations with pytorch. For example, in Eq. (7), authors should either adopt math notations for both (which I think is a better option), or use pytorch syntax. Also, it is quite unusual to write algorithm block in pytorch instead of pseudocode. If space is a concern, authors could defer the algorithm block in the appendix, but I do feel it might be better to write in pseudocode. \n\n3. It's important to explicitly spell out the connection between regularization term and alleviation of bottleneck. In particular, in addition to a sensitivity upper bound on the regularization term, it's important to also exhibit a nontrivial lower bound. In fact, I'm a bit confused in the statement of Proposition 2.8 where the upper bound means BRL is \"repellent\" to the bottleneck, I think you need a lower bound to reach the conclusion. This is particularly crucial since most of the theories in this paper are not deep and serve as explanatory purpose, and a lower bound on the regularization would ensure that, at least in the case of sliding window, the bottleneck is actually alleviated. Ideally, the lower bound should not scale with $1/n$, as one would expect the context length to be long (in fact, this is the main motivation to design efficient attention approximation or new architecture).\n\n4. I'm not an expert on experiments, I'll defer comments on the experiments to other reviewers with more expertise on this front.\n\n\nMinor:\n\n- paragraph above the captions for Table 1: \"performs well differently from the training time performance\", not sure what the authors try to express here, please clarify.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "pnfTN0F4CR", "reviewer": "Reviewer_Av48", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n\n* The idea of combining sparse attention with compressed global tokens in theoretically sound ways is novel.\n\n* Bottleneck regularization is well-justified and extensively explored.\n\n* Formal sensitivity analysis and bottleneck evaluation based on Jacobians.\n\n* Proofs of expressivity equivalence to full attention under the BRL setup.\n\n* Tests on both sequence modeling and graph node classification tasks.\n\n* Exhaustive ablations (effect of compressed tokens, attention distinction).\n\n* Recorded substantial improvements in memory and speed, particularly for extremely long sequences (up to $2^18$).\n\n* In contrast to some kernel-based methods, BRL maintains causal masking, and therefore can be used both for encoder and decoder scenarios.\n\n* The organization of sections makes the method simple to trace despite technical complexity.\n\n# Weaknesses\n\n* Heavily Dependent on Engineering Decisions: There are several hyperparameters (e.g., window size, number of compressed tokens, $\\lambda,\\gamma$) that need to be tuned, and the method's robustness to these choices is not examined thoroughly. How general is the design to other domains (e.g., vision tasks)?\n\n* Possible Complexity in Implementation: Newly introduced elements (compression and propagation mappings) complicate training pipelines compared to baseline attention surrogates like Performer.\n\n* Comparisons could be extended: The paper mentions \"our-implementation\" of some baselines (e.g., full attention) — more details on deviations from vanilla baselines would avoid fairness concerns. No experiments on machine translation or image classification tasks, where causal models are important and the gains would be vital.\n\n* Insufficient qualitative analysis: No qualitative visualizations of how BRL-Attention alleviates bottlenecks compared to local attention over multiple layers (e.g., information flow diagrams beyond Figure 1).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents BRL-Attention (Bottleneck Regularized Linear Attention), a novel attention scheme to counteract inefficiencies in existing efficient transformer models. Traditional full-attention Transformers enjoy quadratic complexity, while recent efficient approaches — pattern-based (e.g., sparse or local attention) and kernel-based (e.g., low-rank approximations) — fall short either in having bottlenecks or sacrificing accuracy. BRL-Attention enhances pattern-based attention with a limited set of compressed tokens as a global information pool, enabling long-range interaction without quadratic cost and without sacrificing causal masking (key to autoregressive usage). Theoretical inspection and extensive experimentation on long-sequence benchmarks (e.g., Long Range Arena, WikiText-103, large graph corpora) attest that BRL-Attention is capable of competing or beating full-attention transformers with significant reductions in memory and compute cost.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* Describe in greater detail how sensitive BRL's performance is to compressed tokens' numbers, $\\\\\\\\beta,\\\\\\\\gamma$, etc.\n\n* Include experiments beyond text and graphs, e.g., image or speech processing (even some early ones).\n\n* Enumerate explicitly any modifications of re-implemented baselines (e.g., \"Transformer (our-imp)\") so that readers can have confidence that improvements are genuine.\n\n* Include visualization showing how token representations evolve with and without BRL in deeper layers.\n\n* Some figures (e.g., Fig. 4, Fig. 5) are too busy and may be more readable with a cleaner layout.\n\n* Include sketch proof diagrams in the main paper (Appendix is okay for full proofs, but the main paper needs to give a summary of key steps naturally).", "broader_impact_concerns": "There is no concern.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)