MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head
Authors: Kewei Zhang
Ye Huang
Yufan Deng
Jincheng YU
Junsong Chen
Huan Ling
Enze Xie
Daquan Zhou
Venue: ICLR 2026 Poster (2026)
Keywords: Linear Attention
Model Architecture
Efficiency
Primary area: generative models
TL;DR: Multi-Head Linear Attention addresses the performance degradation in linear attention by preserving representational diversity through head-wise token dimension computation.

Abstract: While the Transformer architecture dominates many fields, its quadratic self-attention complexity hinders its use in large-scale applications. **Linear attention** offers an efficient alternative, but its direct application often degrades performance, with existing fixes typically re-introducing computational overhead through extra modules (e.g., depthwise separable convolution and few self-attention blocks) that defeat the original purpose. In this work, we identify a key failure mode in these methods: **global context collapse**, where the model loses representational diversity. To address this, we propose **Multi-Head Linear Attention (MHLA)**, which preserves this diversity by computing attention within divided heads along the token dimension. We prove that MHLA maintains linear complexity while recovering much of the expressive power of softmax attention, and verify its effectiveness across multiple domains, achieving a **3.6%** improvement on ImageNet classification, a **6.3%** gain on NLP, a **12.6%** improvement in image generation tasks and a **41%** enhancement in video generation tasks with the same computational complexity.

Published as a conference paper at ICLR 2026
MHLA: RESTORING EXPRESSIVITY OF LINEAR ATTENTION VIA TOKEN-LEVEL MULTI-HEAD
Kewei Zhang"', Ye Huang", Yufan Deng', Huan Ling,, Enze Xie², Daquan Zhou't it !Peking University, Shenzhen Graduate School Project Page: paarrupppku.ggthuuboo//MEME
Jincheng Yu?, Junsong Chen?,
NNVIDIA
MHLA
(a) Generation results from SANA-MHLA
10131
Self Attention
Lincar Attension
CSR
460.2
MHLA
LongBench
Througbput (Images/sec)
FlnshAtcation
100
235.0
YBench
20.6
CZI
Image
10
Med4
(c) Multi-Domain Benchmark Performance of MHLA
19.7
A160
0.55
7.1x Speed
Flack Americs
MITLA
Lawan Amn
00M
11100
7.40
8.22x Speced up
256
512
1024
2048
4096
Resolution
Throughput (imagexise))
(b) Throughput Resolution
(d) Throughput of DIT-S/2 with 4096 resolution
Figure 1: (a) Generation results from our fine-tuned SANA model using MHLA. (b) Perfor- mance and efficiency comparison between the proposed MHLA and baselines. The throughput was tested on the NVIDIA H100 Tensor Core GPU. Following the previous method, we report the FID in the table resolution of 256 x 256. (c) Average rank and entropy of attention scores for DeîT-T with different attention types, showing MHLA yields richer and more focused attention. (d) Throughput of DiT-S/2 at 4096 resolution across different devices. All improvements are solely due to MHIA, and can be further combined with orthogonal techniques for even greater speedups.
ABSTRACT
While the Transformer architecture dominates many fields, its quadratic self- attention complexity hinders its use in large-scale applications. Linear atten- tion offers efficient alternative, but its direct application often degrades per- formance, with existing fixes typically re-introducing computational overhead through extra modules (e.g., depthwise separable convolution) that defeat the orig- inal purpose, In this work, we identify a key failure mode in these methods: global context collapse, where the model loses representational diversity. To address this, we propose Multi-Head Linear Attention (MHLA), which preserves this diversity by computing attention within divided heads along the token dimension. We prove that MHLA maintains linear complexity while recovering much of the
Equal contribution.
Correspond daquan, whoo@phu..cee... edu..c

---

Published as a conference paper at ICLR 2026
expressive power of softmax attention, and verify its effectiveness across multiple domains, achieving a 3.6% improvement on ImageNet classification, a 6.3% gain on NLP, a 12.6% % improvement image generation, and a 41 % enhancement on video generation under the same time complexity.
INTRODUCTION
Self-attention is the core module for the recent dominant model architecture, Transformer, for both computer vision (Dosovitskiy et al., 2021), natural language processing (Vaswani et al., 2017), and generative tasks (Rombach et al., 2022). However, its quadratic time and memory complex- ity severely limit its scalability to long sequence tasks such as high-resolution image generative and video generation tasks (Zhou et al., 2022; Kong et al., 2024; Zhou al., 2024).
To address the efficiency issue, growing line of research (Katharopoulos et al., 2020; Choroman- ski al., 2021) has developed linear attention mechanisms that replace the softmax kernel with associative feature maps. These approaches reduce the computational and memory complexity of attention from quadratic to linear by compressing all keys and values into a global summary. Al- though this improves efficiency, it eliminates one of the key advantages of softmax attention-its ability to adapt to each query individually. Consequently, linear attention often experiences notable accuracy degradation, particularly in long-sequence modeling tasks.
Recent works (Fan et al., 2025b; Han et al., 2023; 2024) have sought to mitigate the performance degradation of linear attention by integrating components such as depthwise convolutions and gating modules. However, this reliance on external modules introduces additional computational overhead and continues suffer from performance degradation as sequence length increases. In this paper, we present a solution to the performance bottleneck in linear attention that requires no additional self-attention modules. Our key insight is that, in conventional linear
depthwise convolution attention design, all tokens are compressed into a single global key-value summary (KV summary) that is shared by every query. This design could have reduced the model's representation capacity, as illustrated in Figs. 1b and 2. To evaluate diversity, we compare the rank of the attention weight matrices across different models. We find that using a shared global KV summary limits the model's capacity represent rich interactions, effectively capping it at a fixed rank. As sequences grow longer, this constraint tends to push the attention weights toward more uniform distribution. In practice, this reduces diversity and degrades performance tasks where queries must concentrate on a small subset of relevant tokens.
Our design goal is therefore simple: restore query-dependent diversity, the ability for different queries to retrieve different contexts, without sacrificing lincar-time behavior or introducing heavy auxiliary modules.
Thus, we introduce Multi-head Linear Attention (MHLA) to achieve the aforementioned charac- teristics. Specifically, MHLA partitions tokens into noveoeallg blocks ("heads" in the spa- tial dimension), computes local key-value summaries, and lets each query block compute a query- conditioned mixture over these summaries retrieve tailored context; within the selected blocks, token contributions are further refined by a query-dependent reweighting module. Thanks to the simplicity MHLA, the implementation only relies on standard GEMMs, keeping the overall com- putational overhead negligible with O(N) complexity, retaining compatibility with streaming/state- ful execution. It was clearly observed that adding MHLA raise the rank of the attention weights matrix significantly. as shown in Fig. 3b. The difference between previous linear attentions and MHLA is briefly illustrated in Fig. 2.
We validate MHLA on image classification, image generation, natural language processing, and video generation tasks. Experiments show that MHILA consistently outperforms existing linear at- tention baselines with negligible computational overhead. Our main contributions are summarized as follows:
We conduct in-depth analysis of linear attention and identify one of the root causes of its performance degradation: the absence of grouping along the token dimension during similarity calculation. This limitation can be quantified by examining the rank of the attention matrix.
2

---

Published as a conference paper at ICLR 2026
Other Linear Attention
MHLA
Local KV
KV Summary
d'x d
dx
dx
Output
Multi-Head Mixing
Output
Figure 2: Comparison between the proposed MHLA and other linear attentions. MHLA di- vides multiple heads on the token dimension. Through Multi-Head Mixing, MHLA restores query- conditioned selectivity by mixing KV summaries with query-specific weight, improving token-level diversity while keeping linear complexity.
We propose new formulation of linear attention that achieves state-of-the-an performance on both discriminative and generative tasks, while maintaining O(N) computational complexity and avoiding reliance on additional modules.
We conduct extensive experiments across various tasks, achieving sate-offthe-ar performance. On ImageNet, MHLA delivers 3.6% % accuracy gain self-attention while image gener- ation tasks it improves the performance of the DiT architecture by 12.6%. MHLA also achieves 6.3 % improvement natural language processing tasks and provides substantial 41 % im- provement compared to vanilla linear attention video generation tasks.
RELATED WORKS
Transformers (Vaswani et al., 2017) have advanced various fields (Devlin et al., 2019; Dosovitskiy al., 2021; Saharia et al., 2022), but their quadratic time and memory complexity due to self- attention limit scalability, especially for long sequences. To overcome this, linear attention mech- anisms (Katharopoulos et al., 2020; Choromanski et al., 2021) have been proposed, which replace softmax with kernel-based methods to achieve linear time complexity. While these mechanisms im- prove the efficiency, they often lose expressiveness, making them suffer from a performance drop in capturing complex token interactions. Several solutions (Fan et al., 2025b; Han et al., 2023), includ- ing adding convolutional layers or gating mechanisms, have attempted to recover performance but tend to introduce additional computational costs. See the detailed related works in the Appendix A.
3 ANALYSIS OF LINEAR ATTENTION
3.1 PRELIMINARY
We first formulate the calculation the attention weights for both self-attention and linear attention mechanism. Given an input token sequence X IR V'd we first compute queries, keys, and values via Q = XW2, K X WW,, = XWv, where Wo, WK, Wy |Rdxd are learnable projections. The attention output of the token i can be expressed as:
Sim K,),
Y
(1)
Sim( Km)
m==1
where Sim(-.) calculates the similarity between the input matrix. In softmax attention (Vaswani al., 2017), Sim(Q, K2) (sp(Q,K] /Vd), pairwise similarities need be calculated and normalized per query, resulting in O(22) complexity.
Linear attention replaces the exponential kernel with positive feature map ob-) such that Y, (QQ)) (K3)TVV)
Sim(Q K2) (())
(2)
(Q.) T(Km)T))
global key-value summary G Σ, ((K,) V, and normalizer = Σn & ((,,)) ) respectively. This reduces the complexity from O(22) to O(Nds), enabling linear-time scaling with sequence length.
where the numerator and denominator can be precomputed
3

---

Published as a conference paper at ICLR 2026
3.2 GLOBAL CONTEXT COLLAPSE
Linear attention achieves linear-time complexity by reusing a global key-value summary across all queries, which can be formulated as G Σ== o (K3) VV, E RRdxd But this fixed-size design introduces an intrinsic information bottleneck:
Observation
As the sequence length increases, the information requiring representation exceeds the capacity of the fixed-size d x d matrix, leading to performance saturation. We term this phenomenon global context collapse.
This observation can be quantified using two complementary metrics, which are the rank and the sparsity of the attention matrix:
Rank limitation. The rank of the attention matrix has been widely studied as a key indicator of feature diversity and representational capacity attention mechanisms (Fan et al., 2025b; Han et al., 2023; Bhojanapalli et al., 2020). Specifically, with Q = x(Q) and K = <KK), global linear attention produces
Alin = Q K E E
rank(Am) S min (Q), rank [)} < d.
Conclusion
Regardless of N, the representational capacity of Alin strictly bounded by Although several prior studies have attempted to increase the rank of 2025b; Cao Wang, 2025), this bound results severely rank-deficient approximation of the full n X attention matrix when >> d. constraining the model's ability to capture diverse, query-conditioned attention patterns.
summaries (Fan et al.,
We empirically verify this effect Fig. 3b, which shows that the rank attention scores in linear- cicuntion-based models is consistently capped by the head dimension (typically dn 72), and the relative expressivity of the attention map degrades as the sequence length increases.
Loss of sparsity. The sparsity of the attention matrix a critical factor influencing the perfor- mance of attention mechanisms. Sparse distributions generally exhibit lower entropy, concentrating probability mass on smaller set of informative tokens (Zhang et al., 2025; Deng et al., 2023), which benefits model optimization. Linear attention, however, computes scores by first compress- ing all key-value pairs into a single global summary, and each query interacts with this shared representation only once. In contrast, softmax attention leverages the exponential function to enable each query to produce a distinct distribution over tokens (see Appendix B), Because linear atten- tion relies on the same aggregated representation for all queries, cannot reweight individual keys according to query- query-specific relevance.
Conclusion 2
As the sequence length N increases, the contribution of each token becomes negligible. Consequently, the attention weight distribution approaches uniformity, reducing the sparsity and impairing the model's ability to selectively emphasize informative tokens.
To quantify this effect, we compute the average entropy of the attention scores over 500 random sam- ples for each attention variant. For each row of the attention score matrix, lower entropy indicates that the distribution closer to a one-hot vector, reflecting stronger concentration on single token. As shown in Fig. 3a and Fig. 3b, linear attention exhibits significantly higher entropy, confirming its lack of focus compared to softmax-bast attention.
Taken together, these findings reveal that the reliance single global key-value summary in linear attention leads to a severe collapse in representational capacity, manifested as both rank de- ficiency and elevated entropy in the attention map. We refer to this phenomenon global context
4

---

Published as a conference paper at ICLR 2026
Setteen Afietion
MIILA
Anetion
Softru Aneetio
MILLA
532
423
MiLL
1.4-% Mylk, 5.0
08-22 1425
Figure 3: (a) Visualization of attention score and attention maps of MHLA and baselines. (b) Average rank and entropy of attention scores for DeiT-T, showing MHLA yields richer and more focused attention.
collapse. Fig. 3a visualizes attention scores and maps, clearly illustrating the inability of linear at- tention to capture fine-grained information. This observation motivates the development of methods that restore query-conditioned token-level diversity while preserving the linear-time complexity of the attention mechanism, which was detailed in the next section.
MULTI-HEAD LINEAR ATTENTION
4.1 OVERVIEW
Here we formalize the proposed Multi-Head Linear Attention (MHLA). As shown in Fig. 4a. MHLA operates by splitting the sequence along the token dimension into multiple "heads" and running linear attention in parallel across these "heads". Let the input sequence be X E Nxd projected to queries, keys, and values: Q = XWe, K = XWK, V = XWv, with Q, K,V E gNxdd For efficiency, adopt kernelized formulation, denoting Q @(Q), K 6(K) for chosen feature map o(-).
Standard linear attention aggregates all tokens into single global d summary shared by every query, which reduces expressivity by collapsing token-level diversity. To mitigate this, we split the sequence into M non-overlapping blocks (the MHLA "heads"), with block b containing N₂ tokens and _x=1 Ng N. In practice on vision models, blocks are defined on spatial (2D) or spatiotemporal (3D) grids rather than by flattening to ID. For each block b we compute a local key-value summary and its normalizer:
S Σ R Rdddd jEb
26 = -K, Rdd
K2 JEb
(3)
To restore query adaptivity, MHLA constructs a distinct mixture of all key-value summaries for each query block i through Multi-Head Mixing. Queries in block can then attend to this mix- ture, where different key-value summaries are weighted according to the attention preferences of the current query block. Let m; E Rd denote the nonnegative, learnable mixing coeffi- cients for block i, which are optimized during training. The mixed summaries are then defined as 5, = -E= m,,6 Sp, and the corresponding normalizer is 2 = EM m4,6
The process be done with highly hardware-efficient GEMM operation between key-value summaries and coefficient matrix Me E R 1.MxM consisting of m,, Given a query vector ạ E Rd from block the output is
0 q G
M mib T S Σ m4 2b
(4)
Each output element can thus be interpreted query-specific, bloek-dependent recombination of the entire value sequence. tasks like language modeling and video generation, the normalizer term can be omitted for better training stability (Qin et al., 2022) when the sequence is getting longer.
4.2 MULTI-HEAD MIXING
The core of MHLA's adaptivity is a learned coefficient matrix Me E R M The element at
position (i. denotes the affinity between query-block i and the local key-value summary of block j. Equivalently, the i-th row of M,, denoted m;, specifies how query-block linearly combines the M local summaries into query-specific global summary.
5

---

Published as a conference paper at ICLR 2026
Kay-Value Summaries
Learnable Coefficiant Matrix
Output
M+25
Key Blocks
Value Blocks
Inner Block Matmul
Multi-Head Mixing
dxd
M
M No
Nomallzze
M x dxd
Input
Quary Blocks
Summaries
Inner Block
Matmall
M N.
X
Nv
M Mxdxd
X.d (V= H W)
(b)
Figure 4: (a) Overview of the proposed Multi-Heal Linear Attention. (b) We visualize two rows of the initialized Learnable Coefficient Matrix corresponding to Block and Block 14 separately when is 25. We reshape the two rows and the dimension in 2D for better understanding.
Each m, is produced and learned end-to-end; in practice enforce nonnegativity and malization. Because blocks are defined along spatial or spatiotemporal axes, we initialize M, to favor locality: for row i set initial coefficients m. (0) 1 dist(i, j)/ maxx (dist(i,k)) where dist(i, measures the Euclidean distance and max, dist Re) is the maximum distance from to any position k. The coefficients are then normalized such that Σ, m 1. A visualization of this initialization can found in Fig. 4b. This locality-biased initialization produces more stable and faster convergence while leaving Mc free adapt during training. To further ensure stability, clip the coefficients to the interval (0, 1) on every update.
The token-level effect of the token Writing each local summary as sum over its tokens, G, Elestock j K.V.T the mixture for query-block i expands
Mixing is transparent. Let b(t) denote the block index of
N Σ mi,t0) K.V.T E Rddd
5
Σ
Sj
Ø=1
t=1
For query vector Q(g) (from block i), the numerator of the kernelized update becomes
q S
mame (q V, E
(5)
Eq. makes the mechanism transparent: each query-block rescales the contribution of entire blocks via m,, and within each block the usual kernel inner product ğ" K, differentiates tokens. Thus, MHLA restores query-condittoned, token-level weighting two-stage manner (block selection x mntra-block reweighting). Importantly, all operations reduce to blockwise summary computation and linear combinations of M matrices of size d, asymptotic complexity remains linear in N while expressive capacity is substantially increased.
Chunkwise parallel form of MHLA. Linear attention commonly employs chunkwise parallel training (Hua et al., 2022; Sun et al., 2023) to maintain linear-time complexity under causal mask- ing, by partitioning the sequence into blocks and updating running summary per block. MHLA naturally fits this setting: each head can be directly mapped to a chunk, and we maintain one local summary S, per chunk. For a detailed derivation, see Appendix C.
4.3 ANALYSIS OF MULTI-HEAM LINEAR ATTENTION
Rank analysis. Partition the sequence into M non-overlapping blocks of size N,. Let the query matrix be = ....... TT. with Qo R From Eq. in the calculation of attention score,

---

Published as a conference paper at ICLR 2026
Table Comparison between Self Attention, Linear Attention, and MHLA. We report computa-
tion complexity, maximum achievable rank, memory complexity and query-conditioned selectivity.
Method
Time Complexity
Rank Bound
Memory Complexity
Cerry-Conditioned
Self Attention
O(N2) or O(N)
Lincar Attention
O(Nd)
Old*)
MHLA (ours)
O(NI + M2 Ee min(n..))
O(
the mixed key sequence seen by query-block can be expressed as
Y; = [m m4,02; .... m4.b(n) E Rdxx
where mi,b(() is the mixing coefficient selecting the block of token t. The attention submatrix contributed by query-block i is A, Q,Y; RNN, and the full attention matrix AMHLA = [A1 A2 ... Am] E R"x". Then applying standard rank inequalities gives
ran)(As) min (an) rank)))) S min(ns: d),
which yields the global bound rank(AMm..)) S min(n, n, Σm min(ns. d)).
This upper bound is attainable under mild, generic conditions: if each block product Q0%, has full row rank т min(ns; d) and the row spaces of {Q are linearly independent, then we get rank (AMM.AA = min(n, Eb=1 rt). Even when the independence assumption is not fully satisfied, the blockwise mixture still expands the diversity of the row spaces, causing rank(AAmI.a) to grow roughly additively with M. We empirically validate this behavior in Fig. 3b, where MHLA consistently achieves substantially higher attention-scon rank than other linear attention variants- and does so without relying on auxiliary components such as depth-wise convolutions. This confirms that MHLA natively restores much of the representational capacity lost in global linear attention, whose rank remains strictly limited by regardless of the sequence length N.
Sparsity analysis. The learned coefficient matrix Me allows each query-block to assign higher weights to a subset of blocks that are more relevant, effectively pruning irrelevant tokens at the block level. Within each selected block, the kernel inner products 7 K, further differentiate token contributions, leading to sharper and more concentrated attention distributions. We validate this effect empirically in Fig. 3b, where MHLA consistently yields lower attention entropy compared to other linear- attention basclines and even the softmax attention. This confirms that MHLA preserves query-conditioned selectivity and achieves substantially higher sparsity, enabling the model to attend to small, semantically relevant subset of tokens rather than spreading attention uniformly.
Efficiency analysis. The computation of MHLA consists of local Key-value summary computation, Multi-Head Mixing, and output computation, with a time complexity of O(MN,de + M2 + MN,d))
O(N M2 To better capture local information while ensuring efficiency, the number of blocks M is usually set to satisfy M s N. Therefore, N& becomes the leading term and the time complexity of MHLA is O(Nd²). The comparison of self at- tention, linear attention, and MHLA is summarized in Tab. 1. We also provide an empirical analysis of the scaling relationship between N and M in Appendix F.4 that verifies the induced complexity.
5 EXPERIMENTS
5.1 IMAGE CLASSIFICATION
Settings. We adopt the training configurations from prior work (Fan al., 2025b;a; Touvron al., 2021). The proposed MHLA is integrated into two representative architectures, DeiT (Touvron et al., 2021) and VLT (Fan et al., 2025b), across multiple model scales. The models are trained on ImageNet- 1K (Deng et al., 2009). For VLT, strictly follow the setup in (Fan et al., 2025b). All models are trained for 300 epochs with a batch size of 1024 and a peak learning rate of Ie-3, The head number M is set to 16 if there is no extra description. See Appendix E for more details.
Results. We evaluate the pretrained DeiT models described above and report the result in Tab. 2a, which clearly shows the superior performance of the proposed MHLA. We reach the best accuracy
7

---

Published as a conference paper at ICLR 2026
Table Comparison on Image Classification task. MHLA achieves the best accuracy with mini- mal parameter overhead on DeiT models, and outperforms Transformer- LA-, and Mamba-bar based SOTAs, Results marked with an are reproduced under the same training setup as MHLA-VIL
(a) Comparison of different attentions DeiT.
Comparison with SOTA models on ImageNet- K.
Attention Type
Params
FLOPss
Topl-ACC
Cost
Model
Params
FLOPs
TopI-ACC
Comparison Deit-T Setting
FL-PVT-T (Ham al., 2023)
2.0G
Self Attn
5.7M
1.1G
72.2
FL-PVTV2-B1 (Han 2023)
13M
2.2G
79.5
Lincar Attn
5.7M
L.1G
69.8
MSVMambe-M
2024)
12M
1.5G
79.8
Focused LA (Han al., 2023)
6.1M
1.1G
74.1
NAT.M (Hassami et al,, 2023)
20M
2.7G
81.8
Inline Attn (Han et al., 2024)
6.5M
1.1G
74.5
RAVLI-T (Fan al. 2025b)
15M
2.4G
82.3"
MALA (Fan et ol., 2025a)
6.3M
1.1G
75.1
MAVIT- (Fan et al., 2025a)
16M
2.5G
82.4"
MHLA (Ours)
5.7M
75.8
MHLA-VLI-T
16M
2.4G
82.6
Comparison on Deit-S Setting
FAT-B3 (Pan al., 2023)
29M
4.4G
83.6
Self Attn
RALA (Fan et al., 2025b)
Artn
24M
22M
4.2G
4.6G
77.6
79.8
80.4
Vmamha-T (I et al., 2024)
MaSMaamba-T T (Shi et al., 2024)
(llatamizaee
Kautz, 2025)
32M
30M
4.4G
4.9G
5.1G
83.0
82.6
MALA (Fan 2025a)
24M
4.6G
80.3
MAVIT-S (Fan al., 2025a)
4.6G
84.3"
MIILA (Ours)
81.0
MHLA-VLIS
27M
4.6G
84.6
in linear attention across all model sizes, while introducing the fewest extra parameters compared with baselines. We then port the proposed MHLA to VLT (Fan et al., 2025b) and evaluate the perfor- mance under the same settings. The results shown in Tab. 2b, illustrating the proposed MHLA's state-of-the-as performance with consistent improvements compared with baseline models.
5.2 IMAGE GENERATION
Table Mlass-to-mmage Generation. Across all model sizes, MHLA achieves the best perfor- mance. Notably, at and XL scales, it matches self-attention performance without relying on any extra modules.
Settings. 1) For Class-to-Image(CC genera- tion, we train DiT (Peebles Xie, 2023) and DiG (Zhu et al., 2025) from scratch for 400k steps on ImageNet-IK (Deng et al., 2009) with batch size 256 and learning rate 1c-4, following their original settings. We evaluate five variants in DiT and DiG, where the original self-attention (DIT)
(a) Comparison of attention types across models.
Model
Attention Type
Resolution
FID 1
Self
256
Lincar Atteation (Ours)
256
89.72
or GLA (Yang et al., 2024) (DiG) is replaced by our MHLA while keeping other components changed. The head number is set to 16 for both 256 and 512 resolutions. We try extra CPE (Chu et al., 2021) and the output gating module (Yang et al., 2024). Their effects are analyzed in Ap- pendix F.2. 2) For Text-to- (ext-tt-Imagg((22)) genera- tion, we finetune Sana-0.6B (Xie et al., 2024) model from official checkpoint. Both the original linear attention and our MHLA variant are trained for 40k steps with batch size of 256.
DIT-S/2
256
Self Attention
512
84.54
Linear
512
125.33
MHLA (Ours)
512
78.63
(Yang et 2024)
256
62.06
DIG-S/2
GLA
512
99.0 04.04
MHLA (Ours)
256
59.49
Self Attention
256
43.47
DIT-B/2
Linear
256
60.47
MHLA (Ours)
256
37.47
Self Linear Attention
256
256
32.35
DIT-1/2
MHI (Ours, w/None) MILLA (Ours, CPE)
256
25.37
256
24.21
MHLA (Ours, CPE+Gating)
256
21.37
C21 results. The main quantitative results are summarized Tab. 3a, where our method con- sistently achieves across all DiT model sizes. In addition, Fig. 1b compares the throughput of our MHLA with baseline attention mechanisms on DiT-S as the in- put resolution increases. Notably, MHLA main- tains throughput nearly identical to linear atten- tion while delivering performance on par with, or even surpassing, self-attention. At 512 reso- lution, MHLA achieves better FID scores while doubling the throughput of self-attention. To fur- ther demonstrate the fast adaptation ability of our approach to existing models, we fine-tune the pretrained DiT-XL/2 model for 400k steps under the same settings. As shown in Tab. 3b, our model achieves a lower FID score than CFG is applied. Full results of the experiments can be found in Appendix F.
Self Attention
256
19.47
Linear Atlention
256
28.63
titte-of-thearrt performance
DIT-XL2 MHLA (Ours, None)
256
20.32
MHLA (Ours, CPE)
256
22.79
MHLA (Ours, CPE4Garing)
256
(b) Fast adaptation results DIT-XL/2.
Model
Attention Type
FID
18
sFID |
Self Attention
9.62
121.50
6.85
DiT-XL/2
MHLA (Ours)
8,34
121.27
5.52
Self Attention
2.27
278.24
4.60
DiT-XL/(())
MHLA (Ours)
2.54
252.07
4.67
without CFG, and delivers comparable performance when
8

---

Published as a conference paper at ICLR 2026
Analysis. Although we add modules such DWConv (CPE) (Fan et al., 2025b) smaller DiT models, it is worth noting that their benefits diminish as model size increases (CPE even degrades performance DiT-XL). As shown in Tab. 3a, plain MHLA already matches the performance of self-attention on XI. models, while adding CPE leads to regression. These results highlight the intrinsic advantage of MHLA and suggest that, although modules like DWConv may offer gains small scales, their benefits do not scale with model size or sequence length.
Table 4: Comparison T2I models.
Tealning tose
SANA
SANA-MHIA
Model
FID4
CLIP f
GenEval f
0.35
PixArt-o (Chen et al., 2023)
6.14
27.55
0.48
PixAAt-E (Chen et al., 2024)
6.34
27.62
0.52
SANA* (Xie et al., 2024)
6.10
28.15
0.64
SANA-MHLA
5,90
28.26
0,68
5000
10000
12489
Step
20000
35000
3008
20600
Figure 5: Loss comparison.
shown in Tab. 4, replacing linear attention with MHLA consistently improves multiple evaluation metrics, surpassing not only the baseline Sana model but also the PixArt (Chen et al., 2023) series. Fig. 5 further visualizes the training loss curves. The MHLA-basee model rapidly adapts, matching the pretrained checkpoint within the first steps and subsequently converging to a lower loss. This demonstrates MHLA's fast adaptation capability and promising performance at larger model scale.
Fast adaptation to SANA.
5.3 VIDEO GENERATION
Video generation involves extremely long sequence lengths, where quadratic attention becomes prohibitively slow. To evaluate MHLA under such ultra-long contexts, we fine-tune pretrained Wan2. 1-1.3B 3B model by replacing its Flash Attention modules with MHLA. For comparison, also fine-tune a version with full vanilla linear attention (LA). The training uses 81 81-frame videos at 480x800 resolution, corresponding to sequence length of 31,500 tokens, with the mixing-head number M 105. In addition, we train a hybrid model with 2/3 of the layers replaced with MHLA.
We evaluate all models VBench (Tab, 5). MHLA achieves substantially stronger performance than vanilla LA at the same latency. Under such extreme sequence lengths, vanilla LA suffers from global context collapse, leading to severe degradation, whereas MHLA maintains linear-time com- plexity and recovers performance comparable to the original FlashAttention based Wan2.1-1.33 achieving 2.1x inference speedup. The hybrid model further offers an excellent trade-off, deliv- ering 1.6x speedup with even better overall performance. As shown in Fig. 6, MHLA also adapts rapidly during fine-tuni tine-tuning, quickly approaching the pretrained model's loss trajectory. In contrast, vanilla LA fails to train under ultra-long sequences, with its loss plateauing at a high level.
Table 5: MHLA in Video Generation. Wan-FA indicates pretrained Wan2.1-1. 1-1.3B. Wan-MHLA and Wan-LA re- place all layers with MHLA and Linear Attention, re- spectively. MHn-MHLA-H H only replaces 2/3 layers.
Wan-MMlAA
6.3
Win-- Maon-Linnrr-Attettioo
Win Sla--Flast-Atttntioo
Teaining
04-
Model
Quality
Semantic 1
Total
Latency (8)
2002
2300
Wan-FA
85.23
75.65
83.31
166
Step
Wan-LA
69.96
11.38
58.24
Loss comparison on Wan- 1.13B. MHLA shows much stronger convergence capability.
Figure
Wan-MHLA
84.26
16.16
82.62
81
Wan-MHLA-H
84.87
79.59
83.82
103
5.4
NATURAL LANGUAGE PROCESSING
To evaluate MHIAA under autoregressive modeling, we test its performance in language model- ing. Following GLA (Yang al., 2024), train a 0.3B model from scratch on 10B tokens from FineWeb-Edu (Penedo et al., 2024) with a batch size of 0.25M tokens, using a cosine learning rate schedule (max LR 3e-4), weight decay of 0.01, and gradient clipping of 1.0. The head number M is set to 32 for MHLA with a training context length of 2048.
In Tab.6, we report language modeling perplexity and zero-shot commonsense reasoning accuracy, where MHLA achieves performance comparable to Transformer++ et al., 2023) and state- of-the linear models such as GDN (Yang et al., 2025) and Mamba2 (Dao & Gu, 2024). On

---

Published as a conference paper at ICLR 2026
LongBench (Bai al., 2024) (Tab. 7), MHLA further demonstrates clear advantages over other recurrent models, particularly on Multi-Doc QA, summarization, and code tasks, achieving the best average score and highlighting its strong long-context understanding capability.
Table 6: MHLA NLP. We report results evaluated on models trained with 10B tokens. We high- light the best and second best entries.
Model
CSR avg.
Wino.
PIQA acc
ARC-c acc.n
OBQA scc_n 1
ARC-e
BoolQ acc
Wiki.
LMB.
ucc_n
ppl L
ppl
GLA (340M)
46.0
50.0
62.9
25.5
31,0
45.8
60.8
41.47
86.98
Transforme+++ (340M)
46.8
49.6
64.4
25.7
32.8
48.1
60.5
34.57
60.46
(390M)
46.4
50.5
24.9
32,4
48.3
38.32
Mamba2 (340M)
47.0
29.8
64.6
25.5
32.0
49.2
61.2
35,40
58.51
GDN (360M)
46.9
51.3
64.5
25.4
31.4
47.3
62.0
35.01
60.16
MHLAA (340M)
47.1
51.3
64.4
25.9
33.4
46.5
61.3
38.31
71.64
Table 7: MHLA on LongBench. We report results evaluated on 340M models trained with 10B tokens. We highlight the best and second best entries
Model
Multi-Duc HQA
Single-Duc QA
Few-shol SSM
Synthefic PEN
Summarization GvR
Code RBP
2WM
Mus
QQA
NQA
TQA
PZH
QMS
MNs
LCC
AVG
Manba(360M)
3.37
2.36
4.57
2.28
5.16
5.49
12.23
6.97
GLA(325M)
3.23
2.31
4.53
2.13
3.94
0.70
1,98
0.27
11.42
17.72
15.34
12:55
6.33
GDN(346MM
2.86
2.24
1.54
4.73
2.48
6.85
7.61
0.53
0.41
12.46
17.91
15.98
10422
9.8
6.86
4.97
2.13
2.22
4.45
2.15
6.24
7.47
0.76
1.18
11.75
16.81
15.11
11.56
9.92
6.92
Maniba2(330M)
3.56
2.38
1.69
4.70
2.20
4.97
7.03
0.72
1.51
12.57
17.65
14,00
10.15
9.49
6.62
MHLA(325M)
3.58
197
4.68
2.38
0.41
6.44
1.69
12.58
18.59
15.01
13.37
12.72
7.41
5.5 ABLATION STUDY
Multi-Head Mixing. To evaluate the impact of our initialization strategy and learnable design in Multi-Head Mixing, we consider two variants: (1) uniform initialization without locality bias and (2) locality-biased initialization with frozen coefficients. We train and evaluate these variants on DeÎT-T, with results shown in Tab. 8a. The results show that our locality- biased initialization provides strong prior, achieving competitive performance even without learning. Al- lowing the coefficients be learnable further adapts them to the dataset distribution, yielding additional performance gains.
Table 8: Ablation study of the proposed MHLA.
(a) Ablation of init strategy on DeiT-T. LB-init denotes Locality- biased Initialization.
LB-init Learnalle Topl-aco(*6)
75.4
75.1
75.8
(b) Token-level head number
Head number. We also analyze the choice of head number M.
ablation on DiT-S/2, 512px.
For DIT-S/2 at 512 resolution, the input sequence length 1024.
Head number FID.
Throughput]
As discussed in Sec. 4.3, MHLA retains linear complexity when
79.56
435
M < VI024= = 32. We evaluate M E {4, 16, 64), with results sum- marized in Tab. 8b. MHLL achieves excellent FID already at M=16
16
78.63
435
64
79.50
408
while maintaining the highest throughput, implying that MHIAA can reach best performance with relatively small M and thus leading to almost no overhead.
CONCLUSION
In this paper, we introduce a novel linear attention mechanism, termed Multi-Head Linear At- tention (MHLA). By partitioning tokens into multiple groups, MHLA effectively preserves token- wise diversity. Without relying on additional modules such as depthwise convolutions or hybrid self-attention layers, MHLA achieves performance comparable or even surpassing that of self- attention-based models. We envision this work as establishing a fundamental attention mechanism that can benefit wide range of downstream applications, such as high-quality image generation, long-horizon video synthesis, and large-scale language modeling.
Future improvement. Although MHLA demonstrates strong performance on image classification and AIGC, current NLP experiments are limited to 0.3B model and mainly serve prelim- inary validation. A key future direction is to evaluate MHLA at larger language model scales to better understand its scaling behavior in realistic LLM settings.
10

---

Published as a conference paper at ICLR 2026
ETHICS STATEMENT
This paper does not involve studies with human subjects, and it does not raise any concerns regarding harmful insights, discrimination, privacy issues. The methods employed focus improving the efficiency of transformer models machine learning tasks such as image classification, generation, and natural language processing. No conflicts of interest present, and the research adheres the highest standards of scientific integrity.
REPRODUCIBILITY STATEMENT
To ensure reproducibility, the authors have provided sufficient details on the methods and exper- imental setup. The MHLA implementation and experiments are described in detail in the paper. The code will be publicly available once the paper is accepted. All experiments. including image classification and generation tasks, are reproducible as they adhere to standard benchmarks (c.g., ImageNet- K) and configurations from previous work. The authors have also included the results for different model architectures and configurations, demonstrating consistency across various tasks. Further, the supplementary appendix provides additional implementation details to facilitate repro- duction.
REFERENCES
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 31 11-3333 2024.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. In arXiv preprint a5Xiv:2004.05150, 2020.
Srinadh Bhojanapalli, Chulhee Yun, Ankit Singh Rawat, Sashank J. Reddi, and Sanjiv Kumar. Low- rank bottleneck multi-head attention models, 2020. URL
0002.07022
Tom Brown et al. Language models are few-shot learners. Advances in Neural Information Pro- cessing Systems (NeurIPS), 2020.
Yuan Cao and Dong Wang. Saga: Selective adaptive gating for efficient and expressive linear atten- tion, 2025, URL httph:tt//// /ttp:sv///axvv.o// 25509 12817.
Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Yue Wu, Zhongdao Wang, James Kwok, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixart-a: Fast training of diffusion transformer for photorealistic text-to-image synthesis, 2023.
Junsong Chen, Chongjian Ge, Enze Xie, Yue Wu, Lewei Yao, Xiaozhe Ren, Zhongdao Wang, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixart-o: Weak-to-strong training of diffusion transformer for 4k text-to-image generation, 2024.
Rewon Child et al. Generating long sequences with sparse transformers. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Krzysztof Choromanski et al. Rethinking attention with performers. In International Conference Learning Representations (ICLR), 2021.
Xiangxiang Chu, Zhi Tian, Bo Zhang, Xinlong Wang, and Chunhua Shen. Conditional positional encodings for vision transformers. arXiv preprint arXiv:2102.10882 2021.
Tri Dao. F Faster attention with better parallelism and work partitioning. In Inter- national Conference on Learning Representations (ICLR), 2024.
Tri Dao and Albert Gu. Transformers SSMs: Generalized models and efficient algorithms through structured state space duality. In International Conference on Machine Learning (ICML), 2024.

---

Published as a conference paper at ICLR 2026
Tri Dao al. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. leee, 2009.
Yichuan Deng, Zhao Song, and Tianyi Zhou. Superiority softmax: Unveiling the performance edge over linear attention. arXiv preprint arXiv:2310. 16602, 2023.
Jacob Devlin, Ming-Wei Weig Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep
bidirectional transformers for language understanding. arXiv preprint arXiv:/810.0488 2018.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. NAACL, 2019.
Alexey Dosovitskiy et al. An image is worth 16x16 words: Transformers for image recognition at scale. International Conference on Learning Representations (ICLR), 2021.
Patrick Esser, Robin Rombach, and Björn Ommer. Taming transformers for high-resolution image synthesis. IEEE/CVF Conference on Computer Vision and Puttern Recognition (CVPR), 2021.
Qihang Fan, Huaibo Huang, Xiaoqiang Zhou, and Ran He. Lightweight vision transformer with bidirectional interaction. In Proceedings of the 37th International Conference on Neural Infor- mation Processing Systems, NIPS '23, Red Hook, NY, USA, 2023. Curran Associates Inc.
Qihang Fan, Huaibo Huang, Yuang Ai, and Ran He. Rectifying magnitude neglect in linear attention. In ICCV, 2025a.
Qihang Fan, Huaibo Huang, and Ran He. Breaking the low-rank dilemma of linear attention. In CVPR, 2025b.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.0077 2023.
Dongchen Han. Xuran Pan, Yizeng Han, Shiji Song, and Gao Huang. Flatten transformer: Vi- sion transformer using focused linear attention. Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 5961-5971, October 2023.
Dongchen Han, Yifan Pu, Zhuofan Xia, Yizeng Han, Xuran Pan, Xiu Li, Jiwen Lu, Shiji Song, and Gao Huang. Bridging the divide: Reconsidering softmax and linear attention. In NeurIPS, 2024.
Ali Hassani, Steven Walton, Jiachen Li. Shen Li, and Humphrey Shi. Neighborhood attention trans- former. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 6185-6194, June 2023.
Ali Hatamizadeh and Jan Kautz. Mambavision: hybrid mamba-transformer vision backbone. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 25261-25270, 2025.
Qibin Hou, Daquan Zhou, and Jiashi Feng. Coordinate attention for efficient mobile network design. In Proceedings of the IEFF/CVF conference on computer vision and pattern recognition, pp- 13713-13722, 2021.
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In Inter- national conference machine learning, pp. 9099-9117. PMLR, 2022.
Angelos Katharopoulos al. Transformers rans: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning (ICML), 2020.
Weijie Kong, Qi Tian, Zijian Zhang, Rox Min, Zuozhuo Dai, Jin Zhou, Jiangfeng Xiong, Xin Li, Bo Wu, Jianwei Zhang, et al. Hunyuanvideo: A systematic framework for large video generative models. arXiv preprint arXiv:2412.03600 2024.
12

---

Published as a conference paper at ICLR 2026
Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang. Qixiang Ye, and Yunfan Liu. Vmamba: Visual state space model. arXiv preprint arXiv:2401.1016, 2024.
Ze Liu al. Swin transformer: Hierarchical vision transformer using shifted windows, Interna- tional Conference on Computer Vision (ICCV), 2021.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference computer vision, pp. 4195-4205. 2023.
Guilherme Penedo, Hynek Kydliček, Loubna Ben allal, Anton Lozhkov. Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale, 2024. URL https: /ttp/aaavvv//////ee/eeeeeeeeeela org/ abs 2tt.//240...5.7777 17557.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung. Matteo Grella, al. Rwkv: Reinventing mas for the transformer era. arXiv preprint arXiv:2305. 2023.
Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Xingjian Du, Teddy Ferdinan, Haowen Hou, et al. Eagle and finch: Rwkv with matrix- valued states and dynamic recurrence. arXiv preprint arXiv:2404. 05892, 2024.
Hao Peng et al. Random feature attention. In International Conference on Learning Representations (ICLR), 2021.
Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. arXiv preprint arXiv:22 10. 10.00400, 2022.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models unsupervised multitask learners. 2019.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Chitwan Saharia et al. Photorealistic text-to-image diffusion models with deep language understand- ing. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Zhiqiang Shen, Tianhua Tao, Liqun Ma Willie Neiswanger, Zhengzhong Liu, Hongyi Wang, Bowen Tan, Joel Hestness, Natalia Vassilieva, Daria Soboleva, and Eric Xing. Slimpajama-de: Under- standing data combinations for Ilm training. 2024. URL 10818.
Zhuoran Shen al. Efficient attention: Attention with linear complexities. In WACV, 2021.
Yuheng Shi. Minjing Dong, and Chang Xu. Multi-scale vmamba: Hierarchy in hierarchy visual state space model. Advances in Neural Information Processing Systems, 37:25687-25708, 2024.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma. Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to Transformer for large language models. ArXiv, abs/2307.08623 08621, 2023.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. Training data-efficient image transformers amp; distillation through attention. In International Conference Machine Learning, volume 139, pp. 10347-10357, July 2021.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 1371,,2025
Ashish Vaswani et al. Attention is all you need. In Advances Neural Information Processing Systems (NeurIPS), 2017.
13

---

Published as a conference paper at ICLR 2026
Jiahao Wang, Ning Kang, Lewei Yao, Mengzhao Chen, Chengyue Wu, Songyang Zhang, Shuchen Xue, Yong Liu, Taiqiang Wu, Xihui Liu, Kaipeng Zhang, Shifeng Zhang, Wenqi Shao, Zhenguo Li and Ping Luo. Lit: Delving into simple linear diffusion transformer for image generation, 2025. URL httpt:tt////c. /wwxiv.iivv///////////aa///aaa org/ abs 22501 10011.12900 12976.
Sinong Wang et al. Linformer: Self-attention with linear complexity. In Advances in Neural Infor- mation Processing Systems (NeurIPS), 2020.
Enze Xie, Junsong Chen, Junyu Chen, Han Cai, Haotian Tang, Yujun Lin, Zhekai Zhang, Muyang Li, Ligeng Zhu, Yao Lu, and Song Han. Sana: Efficient high-resolution image synthesis with linear diffusion transformers, 2024. URL httph: /aaxiivaaiiiv/a//aa httgg//b//210000 10629.
Yunyang Xiong et al. Nyströmformer: A nyström-based algorithm for approximating self-attention. In AAAI Conference on Artificial Intelligence, 2021.
Songlin Yang, Bailin Wang. Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear atten- tion transformers with hardware-efficient training, 2024. URL attpss://arxii org/ /bs/ 2312. 06635.
Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. In The Thirteenth International Conference on Learning Representations, 2025. URL https:://gg http:s://ppeeeeieee wwt.iwwwwwe.eeee..eee.eeeeeeeeee forum? iidiidd-e wdddd-878P2w2222
Manzil Zaheer et al. Big bird: Transformers for longer sequences. In Advances in Neural Informa-
fion Processing Systems (NeurlPS), 2020.
Zhisong Zhang. Yan Wang. Xinting Huang, Tianqing Fang, Hongming Zhang, Chenlong Deng, Shuaiyi Li, and Dong Yu. Attention entropy is key factor: An analysis of parallel context encoding with full-attention-bart pre-trained language models, 2025. URL htRp:t/ts://arrr 2442 1645.
Daquan Zhou, Bingyi Kang, Xiaojie Jin, Linjie Yang, Xiaochen Lian, Zihang Jiang, Qibin Hou, and Jiashi Feng. Deepvit: Towards deeper vision transformer. arXiv preprint arXiv:2103.11888 2021.
Daquan Zhou, Weimin Wang, Hanshu Yan, Weiwei Lv. Yizhe Zhu, and Jiashi Feng. Magicvideo: Efficient video generation with latent diffusion models. arXiv preprint arXiv:2211.11016 2022.
Yupeng Zhou, Daquan Zhou, Ming-Ming Cheng, Jiashi Feng, and Qibin Hou. Storydiffusion: Con- sistent self-attention for long-range image and video generation. Advances in Neural Information Processing Systems, 37:110315-11040 2024.
Lianghui Zhu, Zilong Huang, Bencheng Liao, Jun Hao Liew, Hanshu Yan, Jiashi Feng, and Xing- gang Wang. Dig: Scalable and efficient diffusion models with gated linear attention. In Pro- ceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 7664-7674, June 2025.
14

---

Published as a conference paper at ICLR 2026
A FULL RELATED WORKS
Transformer. Since the introduction of the Transformer architecture (Vaswani et al., 2017), self- attention has become the dominant mechanism across a wide range of domains, including natural language processing (Devlin et al., 2019; Brown et al., 2020), computer vision (Dosovitskiy et al., 2021; Liu et al., 2021; Hou et al., 2021; Zhou et al., 2021), and generative modeling (Esser et al., 2021; Saharia et al., 2022). The expressive power of self-attention stems from its ability to model pairwise interactions among all tokens, but this comes at a quadratic cost in both computation and memory. This limitation becomes particularly pronounced in large-scale or real-time applications, motivating the exploration of more efficient attention mechanisms. A broad spectrum of strategies has been proposed, such sparse attention (Child et al., 2019; Beltagy et al., 2020; Zaheer et al., 2020), low-rank approximations (Wang et al., 2020; Xiong et al., 2021), and hardware hardware-optimizze optimized variants such as FlashAttention (Dao al., 2022: Dao, 2024). Despite these advances, designing ef- ficient attention mechanisms that maintain both scalability and accuracy remains an open challenge.
Linear Attention. Linear attention has emerged as a prominent direction for addressing the quadratic complexity of standard self-attention. Early works reformulated the softmax operation with kernel-based feature mappings, enabling linear-time complexity in both training and infer- ence (Katharopoulos et al., 2020; Choromanski et al., 2021; Peng et al., 2023; 2024; Yang et al., 2024). While these approaches make Transformers scalable to long sequences, they often suffer from reduced representational power compared to full softmax attention, leading to accuracy drops in challenging tasks such vision and generative modeling. To bridge this gap, subsequent research has incorporated additional modules to enrich the expressiveness of linear attention. For example, convolutional layers have been introduced to capture local context (Peng et al., 2021; Shen et al., 2021; Han et al., 2023; Fan et al., 2025b), gating mechanisms have been proposed to better con- trol information flow. More recently, state space models such as Mamba (Gu & Dao, 2023; Dao & Gu, 2024) and its variants (Shi et al., 2024; Liu et al., 2024) have also been explored as effi- cient alternatives to linear attention, showing strong scalability on long sequences and competitive accuracy. However, these methods still face two fundamental limitations: (1) when applied in unidirectional form to tasks requiring bidirectional attention, they exhibit substantial performance degradation; and (2) when augmented with extra modules (c.g., convolutional layers or additional self-attention blocks), they inevitably incur higher computational overhead and remain vulnerable to global context collapse (see Sec. 3.2), where the global summary loses representational diversity
Sparse Attention. In addition to linear attention, sparse attention mechanisms have been another major approach to addressing the computational bottleneck in Transformers. Methods such as Long- former (Beltagy et al., 2020) and BigBird (Zaheer et al., 2020) introduce sparse attention patterns, where each token only attends to subset of the other tokens, reducing the overall number of atten- tion operations. These methods exploit structural sparsity (e.g., local or global attention patterns) to maintain efficiency while still capturing global context in long sequences. Other techniques, such as the Performer (Choromanski et al., 2021), propose using kernel approximations to achieve sparse attention while preserving the model's expressive power. Although sparse attention mechanisms im- prove scalability, they often introduce trade-offs in terms of accuracy, especially in tasks requiring full token interactions.
Applications of Linear and Sparse Attention. Linear and sparse attention mechanisms have been successfully applied across various domains. including NLP, CV, and generative modeling. In NLP, linear attention has been used to scale models like BERT (Devlin et al., 2018) and GPT (Radford et al., 2019) to longer sequences, enabling better handling of long documents and improving ef- ficiency in language models (Devlin et al., 2019; Brown et al., 2020). In computer vision, linear attention methods have been applied to vision transformers to improve efficiency when processing large images, as seen in works like Swin Transformer (Liu et al., 2021) and DeiT (Touvron et al., 2021). These applications demonstrate the broad utility of linear and sparse attention mechanisms, but also highlight the need for continued development to balance efficiency with the expressive power required by complex tasks like image generation and video understanding.
15

---

Published as a conference paper at ICLR 2026
B QUERY-CONDITIONED SELECTIVITY IN SOFTMAX ATTENTION
A key advantage of softmax sell-attention is its query-conditioned selectivity. Recall the standard attention formulation:
exp(g; k.)) t=1 exp(g, kz)
(ttn( (O,KK, = Σ ag
Two properties are crucial: (i) Query-conditioned weighting: each query qi produces its own distribution {= so the relative importance of token Kj is fully dependent qi (īi) Per- token weighting: the weights act directly on each 6j, without collapsing into global summary. Together, these properties give softmax attention the ability produce highly adaptive, sharply concentrated context vectors.
By contrast, global linear attention aggregates all tokens into single summary matrix Sglobal Σ K2V, shared by all queries, yielding
S
9:
Attnin V)K
q. (E,=
where the per-token contributions are no longer explicitly separable by a result, different queries obtain nearly identical context vectors, losing query-conditioned selectivity.
MHLA restores query-conditioned selectivity. MHLA bridges this gap by introducing a learn- able coefficient matrix Mr that forms query-flock-ppecific mixtures of local summaries:
M
3
=
S.
AttoMHLA (Q, = S.
Because m4.6 varies with the query block i, MHLA assigns different effective weights to the same token depending on the querying block. Expanding S, into its token-level definition gives
g S,
Σ ma (5 K.
ii) block-level selection m2,(()) that is query- conditioned, followed by (ii) within-block token reweighting via the kernel inner product % K,- This design reintroduces gunry-conditioned selectivity and per-token weighting while preserving the linear-time complexity of kernelized attention.
revealing a two-stage weighting mechanism:
C MHLA FOR AUTOREGRESSIVE MODELING
autoregressive modeling, the causal mask prevents each token from attending to future tokens. While linear attention normally achieves ( complexity by reusing global key-value sum- mary, under causal masking, the summary must be recomputed or updated for each prefix, which naively results in d( cost over the full sequence. To avoid this quadratic overhead, widely adopted solution for linear attention is chunkwise parallel training (Sun et al., 2023), which splits the sequence into blocks of size C and processes them in parallel avoid the quadratic cost of recomputing attention over all past tokens. For block b, local key-value summary is computed as S, = S,=Σ,eo,,V) K E ddddd and the global summary is updated recursively:
Seloball = sclobal Si,
H, = Q (Q.K.) Vv.
Here, the first term propagates context from preceding blocks via the prefix summary Selobal while the second term captures intra-block attention. This chunkwise scheme preserves causality and allows block-parallel training with per-block complexity O(Cd² C²dd), leading an overall cost (Cd C2d)) for sequence of length T..
16

---

Published as a conference paper at ICLR 2026
MHLA with chunkwise parallel training. MHLA extends this scheme by replacing the single global summary with query-conditioned mixtures of local summaries. Specifically, for block we form mixed summary
S
Σ m. bS
H, 1.=.:5i+m m,,,Q V..
where mi,b are the learnable mixing coefficients from the causal coefficient matrix Mesusal (upper- triangular entries masked to enforce causality). Queries in block i then interact only with S, yielding speciic,,,,,.c, query-adaptiv context representations rather than a shared global one. Because the mixing is performed once block and reused for all tokens in that block, the asymptotic complexity matches that of chunkwise linear attention.
Causal inference. At inference time, we maintain the past local summaries .. and incrementally update the current block summary S, as new tokens arrive. When a block is complete, its contribution to future mixtures fixed and cached. For new token in block we simply update S, f S; + K2V," and recompute the block's mixed summary S by applying m, to the incremental update. This avoids recomputation over previous blocks and keeps per-token complexity (dP).
D DATASET
To assess the effectiivnees of our approach, conduct extensive experiments four tasks: image classification, class-to-image (C2I) generation, text-to-image (T21) generation, and natural language processing. Following prior works (Fan al., 2025a;b; Han et al., 2023), train classification and C21 models on ImageNet-Ie IK (Deng et al., 2009) and evaluate them on the standard validation set. For T2I generation, we finetune pretrained model using relative small collection of 31,292k images gotten from the internet. For natural language processing, we train models with a subset of SlimPajama (Shen al., 2024) with 5B tokens.
E EXTRA IMPLEMENTATION DETAILS
Image Classification. For training of DeiT, we replace the class token with average pooling and train all baselines under identical settings ensure fair comparison. We additionally add CPE (Chu et al., 2021) with kernel size of 3, following previous works for fair comparison. For VLT, we strictly follow the setup (Fan al., 2025b). All models are trained for 300 epochs with batch size of 1024 and peak learning rate of le-3. For models with an input size of 224, we pad the input size to 256 for better splitting of heads. The head number is set to for DeiT modes. For VLT models, the sequence length for the two linear attention layers is (3136, 784). So we set the head number M to (49, 16} for the two layers respectively.
F COMPLETE EXPERIMENTAL RESULTS
F.1 IMAGE GENERATION
We illustrate the complete results on DiT and DiG models in Tab. 10 and Tab. 9, We provide more generation results of SANA-MHLA in Fig. 7.
Table 9: Fast adaptation results on DiT-XL/2 with MHLA, with and without guidance.
Model
Attention Type
Resolution
FID
IS t
sFID
Precision
Recall f
Self Attention
256
9.62
121.50
6.85
0.67
0.67
DIT-XL/2
MHLA (Ours)
256
8.34
121.27
5.52
0.69
0.65
Self Attention
256
2.27
278.24
4.60
0.83
0.57
DIT-XL/2/G)
MHLA (Ours)
256
2.54
252.07
4.67
0.83
0.56
17

---

Published as a conference paper at ICLR 2026
Table 10: Comparison of different attention types across models.
Model
Attention Type
Resolution
FID
IS 1
sFID
Precision 1
Recall f
Self Attention
256
68.40
Linear Attention
256
89.72
15.24
21.87
0.28
0.41
MHLA (Ours)
256
59.80
23.49
10.16
0.39
0.56
DIT-S/2
Self Attention
512
84.54
15.53
17.02
Linear Attention
512
125.33
33.11
11.64
0.22
MHI.A (Ours)
512
78.63
13.11
18.50
0.40
0.49
GLA (Yang al., 2024)
256
62.06
DiG-S/2
GLA
512
99.04
MHLA (Ours)
256
59.49
24.04
11.51
0.40
0.57
Self Attention
DIT-B/2
Linear Attention
256
60.47
24.27
13.69
0.57
MHLA (Ours)
256
37.47
38.79
7.35
0.63
Self Attention
2,5
23.33
Linear Attention
256
32.35
45.57
8.55
0.54
0.62
DIT-L/2
MHLA (Ours, w/None)
256
25.37
54.38
6.06
0.59
0.61
MHLA (Ours, w/ CPE)
256
24.21
57.62
6.12
0.59
0.62
MHLA (Ours, w/ CPE+Gating)
256
21.37
63.47
5.80
0.61
0.62
Self Attention
256
19.47
Linear Attention
256
28.63
51.15
8.23
0.57
0.62
DIT-XL/2
MHLA (Ours, w/ None)
256
20.32
65.95
6.01
0.61
0,62
MHLA (Ours, w/ CPE)
256
22.79
61.80
5.53
0.60
0.62
MHLA (Ours, CPE+Gating)
256
19.17
68.97
5.70
0.63
0.62
Table 1 I:: Comparison with LiT. We report the FID scores (mean + std) over three independent runs for MHIA demonstrate result stability.
Model
FID (mean std)
LiT-S/2
63.21
DIT-S/2 w/ MHLA
59.74 = 0.10
LiT-B/2
40.86
DIT-B/2 w/ MHLA
37.52 + 0.04
LiT-L/2
24.04
DIT-L/2 w/ MHLA
21.43 + 0.05
LiT-XL/2
20.66
DiT-XL/2 w/ MHLA
19.16 + 0.03
We additionally provide more comprehensive comparisons against other recent linear attention methods on image generation tasks (Wang et al., 2025), and report the mean and standard deviation of MHLA over three independent runs demonstrate the stability of our results. The corresponding results are summarized in Tab. 11.
F.2 ABLATION OF CPE AND OUTPUT GATING.
Table 12: Ablation study of MHLA with CPE and output gating.
Setting
FID
Linear Attention
89.7
MHLA w/ None
76.4
MHLA w/ CPE
64.0
MHLA w/ Gating
68.5
MHLA w/ CPE+Gating
59.8
We conducted detailed analysis of the effects of CPE and Output Gating when combined with MHLA in the DiT-S model, as shown in Tab. 12. Our findings show that, in smaller models, CPE and Output Gating serve orthogonal optimizations of MHILA, effectively enhancing the expressive
18

---

Published as a conference paper at ICLR 2026
Table 14: Training throughput (tokens/s) H100 under different sequence length configurations.
batch size
Type
Model
26x16
xKx8
8Kx4
16Kx2
MHLA
83249
83677
81762
81427
PyTorch w/o optimized kernel
GDN
28763
19217
12306
7262
GLA
57815
31875
18092
17375
GDN
81437 (2.8x)
81596
80739
78049
PyTorch w/ optimized kernel
GLA
141386 x2.5x)
139735
139871
137635
Mamba2
96871
96787
96911
96450
ability when the model size is insufficient. However, our experiments Tab. 3a indicate that the performance gains from CPE and Output Gating diminish as the model size increases. In the DiT-XI. model, adding CPE alone actually leads a performance decrease. In contrast, MHLA consistently provides significant improvements in expressivity, regardless of model size.
F.3 CLASSIFICATION RESULTS ON HIGHER RESOLUTIONS
We further conducted additional experiments at resolu- tions of 384x384 and 512x512. using the DeiT-T model to verify the effectiveness of MHLA on high-resolution classification tasks. Results are shown Tab. 13.
Table 13: High-resolution classification accuracy of DeiT-T with and without MHLA.
Model
Reselution
F.4 SCALING ANAYLSIS
Deff-T
x84
74.4
DaiT-T MHLA
384
77.5
In this section, we conduct empirical studies evalu- ate the throughput of MHLA across different tasks under varying sequence lengths N and token-level head num- bers M. The results in Tab. 16 and Tab. 17 show that when M < N is satisfied, MHLA introduces only negligible overhead, whereas larger M leads to more noticeable overhead. However, our abla- tion studies in Tab. 8b have already demonstrated that choosing M such that M < N is sufficient to achieve strong performance.
Deil-T MHLA
512x512
F.5 STREAMING AND CAUSAL INFERENCE EFFICIENCY
To empirically validate MHLA's compatibility with Shunkwise-paraale training and streaming/s- tateful execution under causal masking, we benchmark both training throughput and autoregressive inference latency on NVIDIA H100 GPUs.
Training Throughput. We report tokens per second (tokens/s) under different "sequence length batch size" configurations in Table 14. Under the standard PyTorch implementation without optimized kernels, MHLA consistently achieves the highest throughput across all settings (2K x 16 to 16Kx2), significantly outperforming GLA and GDN.
Although we have not yet implemented fused customized CUDA kernlss for MHLA due to time constraints, prior results on GLA and GDN demonstrate that kernel-level optimization can provide approximately 2-xx acceleration. Since MHLA is built upon the same GEMM-based primi- tives, comparable speedup is expected after similar optimization.
Inference Latency. To further evaluate streaming inference efficiency, measure causal autore- gressive latency at sequence length 4096. As shown Table 15, under the unoptimized Py PoTch implementation, MHLA achieves the lowest latency (204.3s), outperforming both GLA and GDN, and even surpassing the heavily optimized Mamba2 baseline.
These results demonstrate that MHLA not only maintains theoretical compatibility with chunkwise- parallel training and streaming execution, but also delivers strong empirical efficiency under causal masking.
19

---

Published as a conference paper at ICLR 2026
Table 15: Autoregressive inference latency (seconds) H100 with sequence length 4096. MHLA uses head number M = 32.
Type
Modcl
Latency
MHLA
204.3
PyTorch w/o optimized kernel
GLA
269.4
GDN
292.3
GDN
126.8
PyTorch optimized kernel
GLA
144.4
Mambu2
311.0
Table 16: Profiling results of MHLA under varying sequence length N and token-level head number M. Left: DIT-S/2. Right: DeiT-S/16.
M
MIN
256
1024
42ms 3.7G
52ms 7.1G
147ms 20.8G
129
imgals 3.4G
124 imgs/s 8.9G
40ms
51ms
145ms 21.0G
imgs/s
imgs/s
52ms
61ms
12.00
64
150 imgs/s 5.7G
104 imgs/s 1.00
1024
219ms 10.0G
256
89 imgs/s 18.00
G CLARIFICATION ON TERMINOLOGY AND COMPUTATIONAL CONCEPTS
In this section, we provide formal definitions for the terminology used in our method. These terms describe novel computational behaviors in MHLA that lack direct analogues in prior linear attention formulations.
G.1 CONCEPT query-conditioned
The phrase "query-conditioned" describes mechanism where the aggregation of contextual infor- mation is dynamic and specific to each query instance, distinct from the fixed recurrence found in standard linear attention.
Specifically, the process operates as follows:
Each query token associated with unique vector of mixing coefficients.
These coefficients are used to weight and aggregate all local KV summaries independently for every query position.
Consequently, the adaptation occurs per query, rather than globally or via a shared recursive rule.
G.2 CONCEPT 2: KV Summary vs. Hidden States
We introduce the term KV Summary tos strictly distinguish our approach from the Hidden State found in traditional linear attention papers. While the KV summary may seemingly resemble Hidden States in notation, the underlying computation and dependency graphs are structurally different in two key aspects:
Unlike the strict recursive chain in traditional linear attention where h relies on h4-1. MHLA computes each Global KV Summary (Sy) independently, eliminating state propagation across positions.
While traditional states are derived via one-to-one update from the previous step, MHLA follows a many-to-on aggregation pattern, where each Sg computed from all local summaries using specific mixing coefficients.
By avoiding the rigid inheritance of history inherent to hidden states, MHLA'S KV summaries achieve greater expressivity and flexibility.
20

---

Published as a conference paper at ICLR 2026
Table 17: Scaling analysis results for the language model. We report the theoretically derived values, aligned with empirical measurements. We believe that further engineering optimizations will follow the same trend and achieve stronger performance.
M/N
2048
4096
8192
16
11100 tokens's 4.83G
22000 tokens/s 7.46G
40500 tokens/s 13.04G
11000 tokens's
64
100 tokens/s 7.89G
39400 tokcnss/s 13 13.06G
29600 tokensss 13.79G
Figure 7: More generation results from our fine-tuned SANA-MHLA model.
H LLM USAGE.
We used large language models (LLMs) solely writing aid polish the clarity and readability of the manuscript. Specifically, we employed [L.M-based tools to (i) refine grammar and phras- ing for academic style consistency, (ii) improve logical flow between sections, and (iii) condense overly verbose passages. No new research ideas, experimental designs, or results were produced by the LLM; all scientific contributions, methodology development, and experimental analyses were conceived and executed by the authors.
21


Reviews:
[{"review_id": "v0bgbYvp5F", "reviewer": "Reviewer_WAkW", "summary": "This paper raises a failure mode in linear attention, which the authors term \"global context collapse.\" They argue that compressing the entire sequence into a single, shared key-value summary leads to a severe bottleneck in representational capacity, which they quantify through rank and entropy analysis of the attention matrix. To address this issue, this paper proposes Multi-Head Linear Attention (MHLA), a new mechanism that partitions the input sequence along the token dimension into multiple blocks. MHLA computes local summaries for each block. For each query block, it then constructs a query-conditioned context by computing a learnable, weighted mixture of all local summaries. The authors validate MHLA across diverse domains, including image classification, image generation, and natural language processing, demonstrating competitive performance gains over baselines.", "questions": "1. How does MHLA compare, conceptually and empirically, to other block-wise attention strategies? For instance, how does its \"soft mixing\" of all block summaries contrast with the approach in FLASH, which uses a hybrid of quadratic intra-block and linear inter-block attention?\n2. The authors report the “Throughput of DiT-S/2 with 4096 resolution” in Figure 1(d), but no experimental generation results at 4096×4096 resolution are provided in the paper. Presenting throughput for a resolution that was not actually evaluated risks implying overclaimed contribution and should be supported with real experiments.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The diagnosis of \"global context collapse,\" supported by a concise analysis of rank deficiency and entropy elevation, provides intuitive motivation for the work. It clearly articulates why previous linear attention models often underperform.\n2. The paper presents extensive experiments across multiple domains (computer vision and NLP).", "weaknesses": "1. The choice of the term \"Multi-Head\" is confusing and conflicts with the well-established definition from original Transformer, which refers to splitting the channel dimension. In this paper, \"heads\" are defined along the token/spatial dimension. This non-standard usage could lead to significant confusion in the community. \n2. The paper misses some important comparisons with highly relevant prior work such as FLASH[1] and VOLO[2], which also employs a block-wise strategy combining quadratic and linear attention. For example, VOLO-D1 (comparable to DeiT-S) achieves 84.2 on ImageNet, which substantially higher than the MHLA's 81.0. It is worth noting that VOLO-D1 was proposed four years ago.\n\n\n[1] Transformer Quality in Linear Time, ICML 2022\n\n[2] VOLO: Vision Outlooker for Visual Recognition TPAMI", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "OoNFH1nojT", "reviewer": "Reviewer_uTES", "summary": "The paper identifies a failure mode of many kernelized/linear-attention variants namely “global context collapse,” where compressing all keys/values into a single global summary limits rank and drives attention toward near-uniform distributions as sequence length grow . \n\nAutores proposes Multi-Head Linear Attention (MHLA), which partitions tokens into non-overlapping blocks and learns a block-wise mixing of local KV summaries via a nonnegative, normalized coefficient matrix, restoring query-conditioned diversity while keeping linear-time scaling . The authors derive rank bounds showing MHLA can achieve substantially higher attainable rank than global linear attention and empirically observe lower entropy (sharper sparsity) of attention maps .", "questions": "How are the mixing coefficients parameterized to ensure nonnegativity and row normalization in practice (softmax per row, or other), and how sensitive are results to this choice and to the locality-biased initialization? What values of M were used per model and task, and how does wall-clock time and peak memory scale with N and M on H100 across DeiT/DiT/NLP? Please include profiles showing the point where O(M²d²) becomes non-negligible .\n\n\nAlso, can you provide multi-seed runs (≥3) with mean±std for ImageNet accuracy and DiT FID/IS/sFID, including sample counts and evaluation details?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Clear diagnosis of “global context collapse” with complementary theoretical indicators, e.g. rank upper bound ≤d for global linear attention; MHLA’s additive blockwise rank potential;  and empirical entropy analyses .\n\n\n- Simple, hardware-friendly construction: blockwise summaries + learned nonnegative mixing; retains linear-time leading term and is compatible with chunkwise/streaming execution .\n\n\n- Strong cross-domain results: (i) ImageNet improvements over linear attention baselines and competitive with transformer/mamba-style models, (ii) sizable FID gains across DiT scales (including XL) with near-linear-attention throughput, and (iii) modest yet positive NLP benchmarks without positional embeddings .\n\n\n- Useful ablations: initialization vs learnable mixing and head-number sensitivity give practical guidance (e.g., M=16 works well) .", "weaknesses": "- Results are mostly single-number comparisons; there is no reporting of multiple seeds, confidence intervals, or significance tests. For diffusion FID/IS/sFID and ImageNet top-1, please consider report mean±std over ≥3 seeds and specify sample counts and evaluation protocols used for FID/IS (e.g., 50k samples, classifier, resize method) to support the claims of consistent improvements .\n\n\n- Language modeling evaluation is narrow. The 0.3–0.34B model trained on 5B tokens is assessed on a small suite (ARC-c, Wino, CoPA). Please add perplexity on standard corpora and long-context stress tests to directly evaluate the purported benefits at long sequence lengths, and compare to strong linear/SSM baselines under matched budgets .\n\n\n- Complexity & memory overhead clarity. The method introduces an M×M mixing step with O(M²d²) cost. While argued negligible when M²≤N, the paper would benefit from explicit profiling of wall-clock time and peak memory versus N and M across hardware and tasks, and from reporting the realized M values used in each experiment (classification/generation/NLP) .\n\n\n- Streaming/causal claims are not empirically demonstrated. The text asserts compatibility with chunkwise parallel training and streaming/stateful execution; please provide experiments demonstrating long-context streaming inference throughput/latency versus linear attention baselines under causal masking .\n\n\n- Comparisons breadth in T2I/C2I. The DiT/DiG comparisons are strong, but additional comparisons to other recent linear or hybrid attention mechanisms in diffusion backbones would strengthen the case, especially where prior methods also raise rank or sparsify attention (the paper cites some but not all are evaluated directly) .\n\nAdditionally, there are minor typos (e.g., “diversty”, “Iamge”).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "gtTUCQPH6w", "reviewer": "Reviewer_2bkY", "summary": "This paper addresses a key limitation in linear attention models, termed global context collapse, where the model’s representational diversity diminishes as sequence length increases.\nTo overcome this, the authors propose Multi-Head Linear Attention (MHLA), which partitions the token sequence into multiple blocks and computes local key–value summaries within each.\nBy introducing a multi-head mixing mechanism that adaptively combines these local summaries, MHLA restores query-conditioned diversity while maintaining linear computational complexity.\nExtensive experiments demonstrate that MHLA achieves significant improvements over existing linear attention baselines in image classification (ImageNet-1K), image generation (DiT and SANA models), and NLP benchmarks, with minimal computational overhead.", "questions": "See weakness", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Insightful analysis of linear attention – The authors conduct a rigorous examination of why linear attention underperforms, identifying global context collapse as the root cause. Their analysis of rank deficiency and entropy in attention maps provides a strong theoretical foundation.\n\n2. Innovative multi-head design – MHLA introduces a token-level multi-head mechanism that mixes local key–value summaries through learnable coefficients. This design cleverly restores query-dependent diversity and sparsity while keeping linear complexity.\n\n3. Strong empirical results – On ImageNet classification, MHLA improves accuracy by 3.6% over standard linear attention, achieving state-of-the-art results when integrated into DeiT and VLT architectures. On image generation, it delivers up to 12.6% gains over baselines while matching or surpassing softmax attention performance in large-scale DiT models. On NLP tasks, MHLA yields 2.1% improvements over existing linear attention methods, even without additional positional embeddings.", "weaknesses": "### 1. Limited and Inconclusive Experimental Evaluation\nThe major limitation of this paper lies in its lack of comprehensive experiments to support its central claim—improving long-term modeling capability.\nIn the vision domain, experiments are restricted to image classification on 224×224 inputs, which are too short and fail to reflect long-range dependencies.\nIn the NLP domain, the evaluation covers only three small-scale reasoning benchmarks (ARC-c, WinoGrande, CoPA), which do not assess long-context understanding, open-ended text generation, or large-scale language modeling.\nConsequently, it remains unclear whether MHLA truly benefits long-sequence scaling.\nMoreover, the NLP model used (340 M parameters trained on 5 B tokens) is relatively small; while suitable for controlled experiments, it is insufficient to demonstrate MHLA’s effectiveness in large-scale language models, where long-context efficiency becomes critical.\nTo better validate the method’s claims, I recommend incorporating:\n\nVision tasks such as image captioning, VQA, or video captioning, which require multi-modal and temporal reasoning.\n\nNLP datasets such as LongBench, PG19, or BookSum, which directly evaluate scalability and diversity preservation in long-sequence scenarios.\n\n### 2. Incomplete Discussion of the Computational Trade-off\nThe proposed method essentially trades extra computation for improved representational rank.\nThis trade-off is reasonable when the additional $O(M^2d^2)$ cost remains small, but it does not fundamentally solve the scalability problem for very long input sequences. Linear attention mechanisms are expected to handle arbitrarily long contexts, yet MHLA’s computational complexity scales with block size $M$, limiting its practicality for extremely long sequences.\nThe paper should more clearly analyze this scaling limitation and discuss how MHLA behaves as sequence length increases.\n\n### 3. No Throughput or Speed Analysis\nThe paper does not report throughput, latency, or efficiency comparisons.\nAlthough MHLA maintains linear complexity in theory, the method introduces an additional $O(M^2d^2)$ term due to multi-head block mixing.\nThis overhead may be negligible in small-scale settings like image classification or ARC-c, but it could become significant in truly long-sequence scenarios.\nWithout explicit runtime or memory benchmarks, the paper fails to demonstrate that MHLA preserves an effective balance between accuracy and computational cost under long-context conditions.\n\n### 4. Weak Literature Review and Terminology Issues\nThe paper’s literature review is notably shallow.\nThe introduction and related work sections lack a comprehensive discussion of the extensive prior research on linear attention, a highly active field with numerous variants and analyses.\nThe authors should consult and reference existing survey papers on linear attention to position MHLA more accurately within the broader landscape.\nAdditionally, the writing style introduces several non-standard or self-coined terms—such as “adapt to each query individually” and “KV summary”—which are not widely adopted in the community.\nThe authors are encouraged to revise these descriptions using standard terminology and provide clearer, community-recognized definitions to improve readability and academic rigor.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)