MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head

🎯 Pitch

Standard linear attention suffers from a hidden failure mode called 'global context collapse'—its single shared key-value summary cannot provide each token with distinct, useful context. This paper introduces Multi-Head Linear Attention (MHLA), which lets each query dynamically select from a set of local summaries, restoring expressivity to match or surpass softmax attention, as shown by a 3.6% accuracy gain on ImageNet classification and a 12.6% FID improvement on image generation, all while preserving linear complexity.

1. Executive Summary

This paper introduces Multi-Head Linear Attention (MHLA), a new linear attention mechanism that restores the expressivity lost when replacing softmax attention — which suffers from quadratic complexity — with conventional linear attention, which compresses all key–value pairs into a single global summary shared across all queries. Through an empirical analysis of DeiT and DiT models on ImageNet-1K classification and generation, the authors identify a failure mode they term global context collapse (manifested as rank deficiency and elevated attention entropy that erode per-query selectivity), and MHLA addresses it by partitioning tokens into spatial blocks and computing query-conditioned mixtures of local key–value summaries via a learned coefficient matrix — a two-stage mechanism combining block-level selection with intra-block token reweighting. Across four domains, MHLA achieves a 3.6% accuracy gain over self-attention on ImageNet classification, a 12.6% FID improvement on DiT-based class-to-image generation, a 6.3% gain on NLP commonsense reasoning, and a 41% improvement over vanilla linear attention on video generation, while maintaining $O(Nd^2)$ linear complexity — establishing that query-conditioned token diversity can be recovered without auxiliary modules such as depthwise convolutions, and that conventional linear attention degrades catastrophically only when forced to share a single global context across all query positions.

2. Context and Motivation

The Fundamental Tension: Expressivity vs. Efficiency in Attention

The Transformer architecture has become the dominant backbone across computer vision, natural language processing, and generative modeling because of one core operation: self-attention. Self-attention computes pairwise interactions between every token in a sequence, allowing each query to selectively attend to any key regardless of distance. This global, query-conditioned interaction is what gives Transformers their expressive power — a model can look at an image patch and precisely weight its relationship to every other patch, or a language model can retrieve context from thousands of tokens away.

But this expressivity comes at a steep price: the computation and memory cost scales as $O(N^2 d)$ , where $N$ is the sequence length and $d$ is the feature dimension. For a 1024×1024 image processed at patch size 16, $N$ is already 4,096 tokens, and the attention matrix alone consumes millions of entries. For video generation — where sequences easily reach 30,000+ tokens — quadratic attention becomes prohibitively slow, even on modern hardware. This is not a theoretical concern. As the paper states, the quadratic complexity "severely limit[s] scalability to long sequence tasks such as high-resolution image generative and video generation tasks" (Section 1), and this bottleneck is the primary obstacle to deploying Transformers on real-world long-context problems.

The Linear Attention Solution — and Why It Breaks

The standard remedy is linear attention. The key insight, dating back to Katharopoulos et al. (2020) and Choromanski et al. (2021), is to replace the softmax exponential kernel with a decomposable feature map:

$\text{Sim}(Q_i, K_j) \approx \phi(Q_i)\phi(K_j)^\top$

This factorization is powerful because it lets us reverse the order of operations: instead of computing all pairwise similarities first (the $O(N^2)$ step), we first aggregate all keys and values into a global key–value summary $G = \sum_j \phi(K_j)^\top V_j$ and then let each query interact with this single compressed representation. The result is a dramatic reduction from $O(N^2 d)$ to $O(N d^2)$ , which is linear in the sequence length.

However, this computational trick introduces a fundamental representational bottleneck. Every query in the sequence now retrieves context from the same global summary — there is no per-query selectivity. In softmax attention, query $q_1$ might attend heavily to tokens 5–10 while query $q_2$ attends to tokens 100–105, each producing a sharply different attention distribution. In linear attention, both queries multiply the same $G$ matrix, differing only by the query vector itself. As the paper puts it in Section 3.2, "linear attention achieves linear-time complexity by reusing a global key–value summary across all queries," but "this fixed-size design introduces an intrinsic information bottleneck."

The consequences are measurable. The paper quantifies this degradation through two complementary metrics (Section 3.2):

Rank collapse. The attention matrix in linear attention is $A_{\text{lin}} = \widetilde{Q}\,\widetilde{K}^\top$ , and its rank is bounded by $\min(\text{rank}(\widetilde{Q}), \text{rank}(\widetilde{K})) \leq d$ , where $d$ is the feature dimension — typically 72 or less per head in practical architectures. As the sequence length $N$ grows far beyond $d$ , the relative expressivity of the attention map degrades because the matrix can represent at most $d$ independent interaction patterns regardless of how many tokens exist. Figure 3(b) shows this empirically: linear attention's attention-score rank plateaus while softmax attention's rank grows with the sequence.

Loss of sparsity. Softmax attention uses the exponential function to produce sharply peaked distributions — a query can concentrate nearly all its attention mass on a handful of relevant tokens. Linear attention, by collapsing all tokens into a shared summary, cannot reweight individual keys according to query-specific relevance. The paper quantifies this via attention entropy (Figure 3(a)): linear attention exhibits significantly higher entropy, meaning its attention is spread uniformly rather than focused on informative tokens. This is not merely an aesthetic difference; sparse, low-entropy attention distributions are known to benefit optimization and improve generalization (Zhang et al., 2025; Deng et al., 2023).

The paper gives this combined failure mode a name: global context collapse. It is "manifested as both rank deficiency and elevated entropy in the attention map" (Section 3.2), and it means that linear attention loses precisely the property that made self-attention powerful in the first place — the ability for different queries to retrieve qualitatively different context.

Prior Attempts to Fix Linear Attention — and Their Hidden Costs

The research community has recognized this performance gap and produced a variety of remedies, but the paper argues they share a common flaw: they reintroduce computational overhead that partially or fully defeats the purpose of linear attention.

Convolution-based augmentations. Several works, including RALA (2024), Focused Linear Attention (Han et al., 2024), and Inline Attention (Han et al., 2024), add depthwise separable convolutions or similar local operators to linear attention modules. The intuition is that convolutions capture local spatial structure that the global linear summary misses. However, these are additional modules with their own parameter counts and FLOPs — they are not free. The paper notes in Section 1 that "this reliance on external modules introduces additional computational overhead and continues to suffer from performance degradation as sequence length increases." In other words, convolutions patch a symptom (lost local structure) without addressing the root cause (the shared global summary), and they add cost that scales — albeit linearly — with the sequence.

Gating mechanisms. Approaches like Gated Linear Attention (GLA; Yang et al., 2024) and Gated DeltaNet introduce data-dependent gating to control information flow through the recurrent state. These add parameters and computation within the attention module itself. The paper includes GLA as a baseline (Table 3(a)) and shows it helps but does not close the gap to softmax attention.

Hybrid architectures. Another line of work combines linear attention with occasional full self-attention layers or with state-space models like Mamba and Mamba2. These hybrid designs can recover performance but at the cost of retaining some quadratic or near-quadratic operations, making them less suitable for the longest sequences where linear attention is most needed.

State-space models (SSMs) as alternatives. Mamba (Gu and Dao, 2023) and Mamba2 (Dao and Gu, 2024) reformulate sequence modeling through a continuous-time state-space lens rather than attention. They achieve strong results on long sequences and are included as baselines in the paper's NLP experiments (Tables 6 and 8). However, the paper identifies a critical limitation in Appendix A: "when applied in a unidirectional form to tasks requiring bidirectional attention, they exhibit substantial performance degradation." This matters for vision tasks (where images have no inherent left-to-right ordering) and for bidirectional language understanding — precisely the domains where the paper demonstrates MHLA's advantages.

The common thread across all these approaches is that they add auxiliary modules — convolutions, gates, hybrid attention layers — that increase parameter count, FLOPs, or both. Table 2(a) makes this visible: Focused LA adds 0.4M parameters to DeiT-T, Inline Attn adds 0.8M, and MALA adds 0.6M, while MHLA adds zero extra parameters over vanilla self-attention. The paper's core criticism is that "existing fixes typically re-introducing computational overhead through extra modules... that defeat the original purpose" (Section 1). The goal of linear attention was efficiency; patchwork fixes that erode that efficiency are, at best, a compromise.

The Deeper Problem: Why a Global Summary Is Insufficient

To understand why the paper's diagnosis matters, we need to look at exactly what the global summary loses. In softmax attention, each query $q_i$ produces a distribution $\{\alpha_{ij}\}_{j=1}^N$ where $\alpha_{ij} \propto \exp(q_i^\top k_j)$ . Two critical properties emerge (Appendix B):

Query-conditioned weighting: The weight assigned to token $j$ depends on the specific query $q_i$ through the inner product $q_i^\top k_j$ . Different queries produce entirely different distributions — even orthogonal ones — because the exponential amplifies differences in the inner products.
Per-token weighting: The attention weights act directly on individual value vectors $v_j$ . The value matrix $V$ is never collapsed before the attention weighting is applied, so the model preserves fine-grained token-level information.

Linear attention loses both. Because the summary $G = \sum_j \phi(K_j)^\top V_j$ is computed once and shared, per-token contributions are no longer separable by query index $i$ . As Appendix B explains, different queries obtain "nearly identical context vectors" — their differences are limited to a linear projection through the query vector $\widetilde{q}_i$ , which is far less expressive than the full pairwise softmax.

This is not a minor implementation detail. It means linear attention cannot represent attention patterns where different queries focus on entirely different subsets of tokens — the very thing that makes attention useful for tasks like visual object recognition (where different image regions correspond to different objects) or language understanding (where different words in a sentence depend on different antecedents). The paper's rank and entropy measurements in Figure 3 are the empirical manifestation of this structural limitation.

The Paper's Positioning: A Root-Cause Fix Without Auxiliary Modules

The paper positions MHLA as a direct, principled solution to global context collapse rather than an incremental improvement that papers over symptoms with extra modules. The key intellectual move is stated clearly in Section 1:

"Our key insight is that, in conventional linear 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."

The proposed fix is therefore not to augment linear attention with convolutions or gates, but to restructure the summary mechanism itself so that different query positions can retrieve different context — restoring query-conditioned selectivity while preserving the $O(N d^2)$ complexity bound. This restores query-conditioned selectivity while keeping the $O(N d^2)$ complexity that made linear attention attractive in the first place.

The paper explicitly positions this as a fundamental attention mechanism rather than a task-specific trick. The conclusion states the ambition clearly: "We envision this work as establishing a fundamental attention mechanism that can benefit a wide range of downstream applications, such as high-quality image generation, long-horizon video synthesis, and large-scale language modeling." The breadth of the experimental validation — four domains, multiple model scales, both discriminative and generative tasks — is designed to support this claim of generality.

Summary: The Gap MHLA Fills

To synthesize the motivation:

The problem: Quadratic self-attention is too expensive for long sequences, but linear attention loses the query-conditioned selectivity that makes attention powerful — a failure mode the paper names global context collapse (rank deficiency + elevated entropy).
Why it matters: Long-sequence tasks — high-resolution image generation, video synthesis, long-document NLP — are among the most important and computationally demanding applications of Transformers. Efficient attention that preserves expressivity would directly enable larger models, longer contexts, and broader deployment.
Prior approaches fall short: Existing fixes (convolutions, gating, hybrid architectures, SSMs) add auxiliary modules that reintroduce computational overhead, partially defeating the efficiency gains, and/or fail on bidirectional tasks.
The paper's position: MHLA addresses the root cause (the shared global summary) by introducing query-conditioned mixtures of multiple local summaries, restoring token-level diversity without auxiliary modules, maintaining strict $O(N d^2)$ complexity, and being applicable across domains without modification.

3. Technical Approach

3.1 Reader orientation (approachable technical breakdown)

This paper proposes a new attention mechanism called Multi-Head Linear Attention (MHLA) that can be dropped into any Transformer architecture as a direct replacement for the standard self-attention module. The system it builds is not a separate pipeline or training procedure, but rather a redesigned attention operator that computes how each token in a sequence should aggregate information from all other tokens.

The core problem MHLA solves is that conventional linear attention forces every query token to retrieve context from the same compressed global summary, which destroys the query-conditioned selectivity that makes standard softmax attention powerful. The "shape" of the solution is to split the token sequence into spatial blocks, compute a local key–value summary for each block, and then let each query block form its own custom mixture of these block summaries — restoring per-query diversity while keeping the computation linear in the sequence length.

3.2 Big-picture architecture (diagram in words)

The MHLA attention operator has four major components that process an input token sequence $X \in \mathbb{R}^{N \times d}$ :

Linear projections — Standard learned matrices $W_Q, W_K, W_V$ project the input into queries, keys, and values ( $Q, K, V \in \mathbb{R}^{N \times d}$ ). A chosen feature map $\phi(\cdot)$ maps keys and queries into a kernelized space ( $\widetilde{Q} = \phi(Q)$ , $\widetilde{K} = \phi(K)$ ) to enable the associativity trick that gives linear attention its efficiency.
Block partitioning — The token sequence is divided into $M$ non-overlapping spatial blocks (the "MHLA heads") along the token dimension. For vision, blocks are defined on 2D or 3D grids rather than flattened 1D sequences, preserving spatial locality.
Local key–value summary computation — For each block $b$ , a local summary matrix $S_b = \sum_{j \in b} \widetilde{K}_j V_j^\top \in \mathbb{R}^{d \times d}$ and a normalizer $z_b = \sum_{j \in b} \widetilde{K}_j \in \mathbb{R}^d$ are computed. These are the atomic building blocks that replace the single global summary of conventional linear attention.
Multi-Head Mixing module — A learned coefficient matrix $\mathcal{M}_c \in \mathbb{R}^{M \times M}$ specifies, for every query-block $i$ , a nonnegative weight vector $m_i \in \mathbb{R}^M$ that determines how to linearly combine all $M$ local summaries into a query-block-specific global summary $\widetilde{S}_i = \sum_{b=1}^M m_{i,b} S_b$ . Each query in block $i$ then retrieves context from $\widetilde{S}_i$ rather than from a shared global summary.

Information flows as follows: input tokens → linear projections → feature map application → block partitioning → per-block summary computation → Multi-Head Mixing (mixing coefficients combine summaries per query-block) → query-summary interaction (each query multiplies its block's mixed summary) → output tokens. The normalizer term $z_b$ follows a parallel path and is combined with the numerator to produce the final normalized output.

3.3 Roadmap for the deep dive

First, the kernelized linear attention formulation — because MHLA is built on this foundation, understanding the associative property that enables $O(N d^2)$ complexity is essential before seeing what MHLA changes.
Second, the block partitioning and local summary computation — this is the first departure from standard linear attention and establishes the building blocks that enable per-query diversity.
Third, the Multi-Head Mixing coefficient matrix — the intellectual core of MHLA: how the learned coefficients $\mathcal{M}_c$ let each query-block form its own custom global summary, what the initialization strategy is, and why the resulting attention matrix achieves higher rank and lower entropy.
Fourth, the output computation and two-stage weighting — how the mixed summaries interact with individual queries to produce the final attention output, revealing the block-selection + intra-block-reweighting mechanism.
Fifth, the chunkwise parallel form for autoregressive modeling — how MHLA adapts to causal masking during training and inference without losing linear complexity, which matters for language and video generation tasks.

3.4 Detailed, sentence-based technical breakdown

This is primarily a mechanism design paper whose core idea is that query-conditioned selectivity in linear attention can be restored by splitting the token sequence into spatial blocks and using a learned mixing matrix to let each query-block retrieve a custom-weighted combination of block-level key–value summaries, maintaining $O(N d^2)$ complexity without auxiliary modules.

Kernelized Linear Attention Preliminaries

MHLA inherits the associative computation strategy from kernelized linear attention, so understanding this foundation is necessary before seeing what MHLA modifies.

The standard attention formulation. Given an input sequence $X \in \mathbb{R}^{N \times d}$ , the query, key, and value matrices are produced by learned linear projections:

$Q = X W_Q, \quad K = X W_K, \quad V = X W_V$

where $W_Q, W_K, W_V \in \mathbb{R}^{d \times d}$ are learnable weight matrices. For a single query token $i$ , the softmax attention output is:

$Y_i = \frac{\sum_{j=1}^N \text{Sim}(Q_i, K_j) V_j}{\sum_{m=1}^N \text{Sim}(Q_i, K_m)}$

In softmax attention, $\text{Sim}(Q_i, K_j) = \exp(Q_i K_j^\top / \sqrt{d})$ , which requires computing all $N^2$ pairwise similarities and normalizing per query — the source of $O(N^2)$ complexity.

The kernel trick. Linear attention replaces the exponential softmax kernel with a decomposable feature map $\phi(\cdot)$ that satisfies:

$\text{Sim}(Q_i, K_j) \approx \phi(Q_i) \phi(K_j)^\top$

where $\phi(\cdot)$ maps from $\mathbb{R}^d$ to a (potentially higher-dimensional) feature space $\mathbb{R}^{d_\phi}$ . This factorisation is the critical property because it allows the summation order to be swapped — a mathematical trick with profound computational consequences.

Why swapping the summation matters. With the kernelised similarity, the attention output for token $i$ becomes:

$Y_i = \phi(Q_i) \frac{\sum_{j=1}^N \phi(K_j)^\top V_j}{\phi(Q_i) \sum_{m=1}^N \phi(K_m)^\top}$

Denote the numerator's inner sum as $G = \sum_{j=1}^N \phi(K_j)^\top V_j \in \mathbb{R}^{d \times d}$ and the denominator's inner sum as $z = \sum_{m=1}^N \phi(K_m)^\top \in \mathbb{R}^d$ .

The computation now proceeds in two linear-time stages. First, aggregate all key–value pairs into $G$ and the normalizer $z$ by iterating over all $N$ tokens once — this costs $O(N d^2)$ . Second, for each query $i$ , compute $\phi(Q_i) G$ and divide by $\phi(Q_i) z$ — this is $O(d^2)$ per query, so $O(N d^2)$ total. The $N^2$ pairwise operation has been eliminated entirely. The cost is now $O(N d^2)$ , which is linear in $N$ because $d$ is fixed by the model architecture.

The expressivity cost, restated mathematically. The attention matrix implied by linear attention is $A_{\text{lin}} = \widetilde{Q} \widetilde{K}^\top$ where $\widetilde{Q} = \phi(Q)$ and $\widetilde{K} = \phi(K)$ . Linear algebra gives a hard bound:

$\text{rank}(A_{\text{lin}}) \leq \min\{\text{rank}(\widetilde{Q}), \text{rank}(\widetilde{K})\} \leq d$

This means that no matter how long the sequence $N$ grows, the attention matrix can represent at most $d$ linearly independent interaction patterns — typically $d \leq 72$ per head in practical transformers, far smaller than $N$ for long sequences. The model's representational capacity is capped by the feature dimension, not by the sequence length. This is the mathematical root of global context collapse.

The paper's feature map choice. While the paper does not specify the exact $\phi(\cdot)$ function in the main text (the choice is architecture-dependent and inherited from prior linear attention implementations), the standard options include the ReLU feature map $\phi(x) = \max(0, x)$ , the ELU+1 feature map $\phi(x) = \text{ELU}(x) + 1$ used in Katharopoulos et al. (2020), or the Performer's random Fourier features (Choromanski et al., 2021). The key requirement is that $\phi(\cdot)$ produces nonnegative outputs so that the resulting similarities are valid attention weights. The paper uses the notation $\widetilde{Q} = \phi(Q)$ and $\widetilde{K} = \phi(K)$ to keep the formulation general across feature map choices.

What MHLA inherits and what it changes. MHLA inherits the feature map formulation and the two-stage aggregation-then-query computation. What it changes is that instead of computing a single global summary $G$ , it computes $M$ local summaries $S_1, \ldots, S_M$ — one per spatial block — and lets each query-block combine them via learned coefficients. This preserves the $O(N d^2)$ complexity while breaking the $d$ -bounded rank constraint.

Block Partitioning and Local Summary Computation

From 1D sequences to spatial blocks. Standard linear attention treats the input as a flat token sequence. MHLA instead partitions the $N$ tokens into $M$ non-overlapping blocks along the spatial dimensions of the input. For a 2D image with height $H$ and width $W$ (each divided into patches), blocks are contiguous rectangular regions on this 2D grid. For video, blocks extend to 3D spatiotemporal regions. The paper states: "In practice on vision models, blocks are defined on spatial (2D) or spatiotemporal (3D) grids rather than by flattening to 1D" (Section 4.1). This spatial blocking is essential because it preserves locality — tokens that are spatially adjacent remain in the same block, and the mixing coefficients (initialized to favor nearby blocks) can learn meaningful spatial attention patterns.

Formally, block $b$ contains $N_b$ tokens, and the blocks form a partition of the token set: $\sum_{b=1}^M N_b = N$ . For a model with input resolution 256×256 and patch size 16, the token grid is 16×16 = 256 tokens. With $M = 16$ , each block would cover a 4×4 spatial region. For DeiT models at resolution 224 padded to 256 (Section 5.1), the token count depends on the patch size and the number of MHLA heads is set to $M = 16$ for DeiT-T and DeiT-S (Table 2), and $\{49, 16\}$ for the two linear attention layers in VLT respectively.

Computing a local key–value summary. For each block $b$ , MHLA computes:

$S_b = \sum_{j \in b} \widetilde{K}_j V_j^\top \in \mathbb{R}^{d \times d}$

where $\widetilde{K}_j = \phi(K_j) \in \mathbb{R}^d$ is the feature-mapped key vector for token $j$ , and $V_j \in \mathbb{R}^d$ is its value vector.

What this computes operationally: for every token $j$ inside block $b$ , take its feature-mapped key vector $\widetilde{K}_j$ (a $d$ -dimensional row vector), take its value vector $V_j$ (a $d$ -dimensional column vector), compute their outer product $\widetilde{K}_j V_j^\top$ (a $d \times d$ matrix), and sum these outer products over all tokens in the block. The result $S_b$ is a $d \times d$ matrix — the same shape as the global summary $G$ in standard linear attention, but representing only the tokens within block $b$ .

Why outer products: the outer product $\widetilde{K}_j V_j^\top$ encodes the key–value association for a single token as a rank-1 matrix. When a query later interacts with $S_b$ via $\widetilde{q}^\top S_b$ , the operation expands to $\sum_{j \in b} (\widetilde{q}^\top \widetilde{K}_j) V_j^\top$ — each token $j$ in block $b$ contributes to the output in proportion to its kernel similarity with the query, weighted by its value vector. This is precisely the linear attention computation, restricted to the tokens within block $b$ .

The normalizer path. In parallel, each block computes a normalizer vector:

$z_b = \sum_{j \in b} \widetilde{K}_j \in \mathbb{R}^d$

This is a $d$ -dimensional vector that accumulates the sum of all feature-mapped key vectors in block $b$ . It will later be used to normalize the attention output, ensuring that the effective attention weights sum to one (a proper probability distribution). The denominator in the final output is $\widetilde{q}^\top \widetilde{z}_i = \sum_{b=1}^M m_{i,b} (\widetilde{q}^\top z_b)$ .

Why per-block summaries instead of one global summary: the paper's diagnosis is that a single global summary $G$ forces all queries to retrieve identical context. By computing $M$ separate summaries, MHLA creates $M$ atomic context representations that can later be differentially combined. A query in block 5 might want mostly information from blocks 5–7 (nearby tokens) and almost nothing from block 20, while a query in block 20 might want the opposite pattern. The per-block summaries make this differential retrieval possible because the mixing step (next) can assign different weights to each block depending on which query-block is asking.

Computational cost of this stage. Computing $S_b$ for a block with $N_b$ tokens requires $N_b$ outer products of size $d \times d$ , costing $O(N_b d^2)$ per block. Summing over all $M$ blocks gives $O(N d^2)$ total — identical to the cost of computing the single global summary in standard linear attention. The block structure adds no asymptotic overhead.

Multi-Head Mixing: The Learned Coefficient Matrix

The core idea. The intellectual center of MHLA is a learned mixing matrix $\mathcal{M}_c \in \mathbb{R}^{M \times M}$ where $M$ is the number of token blocks. The element at position $(i, j)$ — denoted $m_{i,j}$ — represents the affinity between query-block $i$ and the key–value summary of block $j$ . Row $i$ of $\mathcal{M}_c$ , denoted $m_i \in \mathbb{R}^M$ , is a weight vector that tells query-block $i$ how to combine all $M$ local summaries into a query-specific global summary.

The mixture computation. For query-block $i$ , the mixed summary and mixed normalizer are:

$\widetilde{S}_i = \sum_{b=1}^M m_{i,b} S_b, \qquad \widetilde{z}_i = \sum_{b=1}^M m_{i,b} z_b$

where $S_b \in \mathbb{R}^{d \times d}$ and $z_b \in \mathbb{R}^d$ are the local summaries computed in the previous stage.

What this computes operationally: take the $M$ local summary matrices $S_1, \ldots, S_M$ , each a $d \times d$ matrix, and for query-block $i$ , form a weighted sum of these matrices using the coefficients $m_{i,1}, \ldots, m_{i,M}$ as weights. The result $\widetilde{S}_i$ is again a $d \times d$ matrix — it looks like a standard linear attention global summary, but it is customized for query-block i. Different query-blocks get different mixed summaries because different rows of $\mathcal{M}_c$ have different weight patterns. The normalizer $\widetilde{z}_i$ follows an identical weighted-sum logic but operates on vectors rather than matrices.

Why learned rather than fixed: if the mixing coefficients were fixed (e.g., uniform weights, or purely distance-based), the model could not adapt its attention patterns to the data distribution. A classification model might learn that certain spatial regions are globally informative (e.g., the center of an image) and up-weight those blocks for all query-blocks. A generative model at different layers might learn different mixing patterns — early layers might favor broad, uniform mixtures while later layers focus on highly specific block subsets. Making $\mathcal{M}_c$ learnable lets the model discover these patterns during training.

Constraints and enforcement. The paper enforces two properties on the mixing coefficients:

Nonnegativity: $m_{i,j} \geq 0$ for all $i, j$ . This ensures the mixed summary is a convex-like combination of local summaries — no block gets "subtracted" — which has a clean interpretation as an attention distribution over blocks. Negative weights could cause unstable training because the normalizer denominator might approach zero or become negative.
Normalization: The coefficients in each row are constrained such that they can be interpreted as a distribution, though the paper does not specify whether hard $\sum_j m_{i,j} = 1$ normalization is applied at every step or soft normalization emerges through training. In practice, the paper states the coefficients are "clipped to the interval $(0, 1)$ on every update" (Section 4.2) to ensure stability.

Hardware-efficient implementation via GEMM. The paper emphasizes that the mixture computation can be implemented as a standard matrix-matrix multiplication (GEMM) between the coefficient matrix $\mathcal{M}_c$ and a stacked representation of the local summaries. Because GEMM is one of the most heavily optimized operations on modern GPUs (NVIDIA Tensor Cores are designed specifically for it), this means the mixing overhead is minimal in practice. The paper states: "The process can be done with a highly hardware-efficient GEMM operation between key–value summaries and coefficient matrix" (Section 4.1). No custom CUDA kernels or sparse operations are needed.

Locality-biased initialization. Prior to training, the coefficient matrix is not initialized randomly or uniformly. Instead, the paper uses a locality-biased initialization that encodes the inductive bias that nearby spatial blocks are more relevant than distant ones:

$m_{i,j}^{(0)} \propto 1 - \frac{\text{dist}(i, j)}{\max_k \text{dist}(i, k)}$

where $\text{dist}(i, j)$ measures the Euclidean distance between the spatial centers of blocks $i$ and $j$ , and $\max_k \text{dist}(i, k)$ is the maximum distance from block $i$ to any other block $k$ . The coefficients are then normalized so that $\sum_j m_{i,j}^{(0)} = 1$ .

What this initialization encodes: for a query-block $i$ , the initial mixing weight assigned to block $j$ is largest when $j$ is spatially closest to $i$ and linearly decreases to zero for the most distant block. This means that, at the start of training, each query-block retrieves context primarily from its own spatial neighborhood — which is a sensible default for vision tasks where spatial locality is a strong prior. Figure 4(b) visualizes two rows of this initialized matrix for a model with $M = 25$ blocks, reshaped to 2D for spatial interpretation. Block 1 (a corner) shows highest weights in its own corner region with a smooth falloff. Block 14 (near the center) shows a radially symmetric pattern peaking at its own location.

Why locality-biased rather than uniform: uniform initialization would give every query-block an identical summary (the average of all local summaries), which is essentially standard linear attention — exactly the thing MHLA is trying to escape. Locality bias provides a structured starting point that already breaks the symmetry, giving each query-block a distinct context. The ablation in Table 7(a) confirms that locality-biased initialization alone (with frozen coefficients) achieves 75.4% accuracy on DeiT-T, compared to 75.1% with uniform initialization. Learning further improves this to 75.8%, showing the initialization provides a strong prior that training can refine.

Inference-time behavior. After training, $\mathcal{M}_c$ is a fixed learned matrix. The mixing step is a single matrix multiplication per attention layer, adding $O(M^2 d^2)$ to the total cost. The paper ensures $M^2 \leq N$ (Section 4.3), so the $O(N d^2)$ term from the per-token operations dominates and the mixing cost is subsumed into the linear complexity bound.

Output Computation and the Two-Stage Weighting Mechanism

The query–summary interaction. Given a query vector $\widetilde{q} \in \mathbb{R}^d$ from block $i$ , the output for that token is:

$o = \frac{\widetilde{q}^\top \widetilde{S}_i}{\widetilde{q}^\top \widetilde{z}_i} = \frac{\sum_{b=1}^M m_{i,b} \, \widetilde{q}^\top S_b}{\sum_{b=1}^M m_{i,b} \, \widetilde{q}^\top z_b}$

where $\widetilde{S}_i$ is the mixed summary for query-block $i$ , $\widetilde{z}_i$ is the mixed normalizer, and the right-hand side expands the mixture to show the individual block contributions.

What the numerator computes: the query vector $\widetilde{q}$ is multiplied against the mixed summary $\widetilde{S}_i$ , producing a $d$ -dimensional output vector. Expanding through the mixture: $\widetilde{q}^\top (\sum_b m_{i,b} S_b) = \sum_b m_{i,b} (\widetilde{q}^\top S_b)$ . For each block $b$ , $\widetilde{q}^\top S_b$ is a $d$ -dimensional vector representing the aggregated contribution of all tokens in block $b$ to the query's output, weighted by the mixing coefficient $m_{i,b}$ .

What the denominator computes: the query vector $\widetilde{q}$ is multiplied against the mixed normalizer $\widetilde{z}_i$ , producing a scalar. Since $z_b = \sum_{j \in b} \widetilde{K}_j$ , the denominator expands to $\sum_b m_{i,b} \sum_{j \in b} (\widetilde{q}^\top \widetilde{K}_j) = \sum_b m_{i,b} \sum_{j \in b} \phi(q)^\top \phi(K_j)$ . This scalar is the sum of all kernel similarities (weighted by mixing coefficients) and serves as the normalization factor — dividing by it ensures the effective attention weights form a distribution that sums to one, though exactly what distribution is slightly more subtle than in softmax attention because of the block-level pre-aggregation.

The two-stage weighting mechanism revealed. Expanding the numerator to its full token-level form reveals the key structural innovation:

$\widetilde{q}^\top \widetilde{S}_i = \sum_{b=1}^M m_{i,b} \sum_{j \in b} (\widetilde{q}^\top \widetilde{K}_j) V_j^\top = \sum_{t=1}^N m_{i, b(t)} (\widetilde{q}^\top \widetilde{K}_t) V_t^\top$

where $b(t)$ denotes the block index of token $t$ .

What this equation shows operationally: for every token $t$ in the entire sequence, its contribution to the output is the product of two factors. First, the block-level mixing coefficient $m_{i, b(t)}$ , which depends on which query-block $i$ is asking and which block $b(t)$ the token belongs to. Second, the intra-block kernel similarity $\widetilde{q}^\top \widetilde{K}_t = \phi(q)^\top \phi(K_t)$ , which depends on the specific query $q$ and the specific key token $K_t$ . These two factors multiply together to weight the token's value vector $V_t$ .

Why this is query-conditioned: different query-blocks use different rows of $\mathcal{M}_c$ , so $m_{i, b(t)}$ varies with $i$ . A token in block 5 might get a large coefficient when queried from block 5 (nearby) but a small coefficient when queried from block 20 (distant), even though the token itself and the kernel similarity $\widetilde{q}^\top \widetilde{K}_t$ remain the same. This two-stage design — coarse block selection followed by fine token reweighting — is what restores query-conditioned selectivity without explicitly computing per-query pairwise similarities against all tokens.

Relationship to softmax attention: softmax attention achieves query-conditioned selectivity through the exponential $\exp(q_i^\top k_j)$ acting individually on each token pair. MHLA approximates this selectivity through a factorized form: block-level weights $m_{i,b}$ provide the coarse selectivity (which blocks to attend to), and the kernel inner products provide the fine selectivity (which tokens within those blocks to attend to). The factorization is what makes the computation linear in $N$ rather than quadratic — the block summaries are precomputed, and the mixture weights are shared across all tokens in a query-block.

The normalizer in practice. The paper notes that "In tasks like language modeling and video generation, the normalizer term can be omitted for better training stability when the sequence is getting longer" (Section 4.1). This refers to the denominator $\widetilde{q}^\top \widetilde{z}_i$ . Omitting it means the output is just the unnormalized numerator $\widetilde{q}^\top \widetilde{S}_i$ , which is a common practice in recurrent architectures like Mamba and RWKV where the normalizer can become numerically unstable for very long sequences. Whether the normalizer is used or not is a task-dependent hyperparameter.

The complete MHLA attention output for a query block. For all tokens in query-block $i$ , the computation is identical except for the individual query vectors $\widetilde{q}$ . The mixed summary $\widetilde{S}_i$ is computed once and reused for every token in the block. The output for each token is $o_q = (\widetilde{q}^\top \widetilde{S}_i) / (\widetilde{q}^\top \widetilde{z}_i)$ . This blockwise reuse is what makes the $O(M^2 d^2)$ mixing cost amortize over $N_b = N/M$ tokens, keeping the per-token overhead negligible.

Rank Analysis: Why MHLA Escapes the d-Bound

The paper provides a formal rank analysis to demonstrate mathematically why MHLA achieves higher representational capacity than standard linear attention.

Construction of the MHLA attention matrix. Let the full query matrix be partitioned by blocks: $\widetilde{Q} = [\widetilde{Q}_1^\top, \ldots, \widetilde{Q}_M^\top]^\top$ where $\widetilde{Q}_b \in \mathbb{R}^{N_b \times d}$ . For query-block $i$ , the key sequence it effectively sees — after mixing — is:

$Y_i = [m_{i, b(1)} \widetilde{K}_1, \; m_{i, b(2)} \widetilde{K}_2, \; \ldots, \; m_{i, b(N)} \widetilde{K}_N] \in \mathbb{R}^{d \times N}$

Each key vector is scaled by the mixing coefficient that query-block $i$ assigns to the key's block. The attention submatrix for block $i$ is $A_i = \widetilde{Q}_i Y_i \in \mathbb{R}^{N_b \times N}$ , and the full MHLA attention matrix stacks these submatrices:

$A_{\text{MHLA}} = [A_1 \; A_2 \; \cdots \; A_M]^\top \in \mathbb{R}^{N \times N}$

Rank bound derivation. For any submatrix $A_b = \widetilde{Q}_b Y_b$ , linear algebra gives:

$\text{rank}(A_b) \leq \min\{\text{rank}(\widetilde{Q}_b), \text{rank}(Y_b)\} \leq \min(N_b, d)$

This bound is tight for each block: the query submatrix $\widetilde{Q}_b$ has at most $N_b$ rows and $d$ columns, so its rank cannot exceed $\min(N_b, d)$ . The mixed key matrix $Y_b$ similarly has at most $d$ rows (from the key dimension) and $N$ columns. The product inherits the tighter of these bounds.

The global rank bound. Since the full attention matrix is a vertical stack of the block submatrices, its rank is bounded by the sum of the individual submatrix ranks (ranks are subadditive under stacking):

$\text{rank}(A_{\text{MHLA}}) \leq \min\left(N, \sum_{b=1}^M \min(N_b, d)\right)$

Why this is substantially larger than standard linear attention. In standard linear attention, the rank is bounded by $d$ regardless of $N$ — that is a single min, not a sum of mins. In MHLA, if each block has $N_b \geq d$ tokens, then $\min(N_b, d) = d$ for every block, and the sum is $M \cdot d$ . The rank bound becomes $\min(N, M \cdot d)$ . For $M = 16$ and $d = 72$ , the bound is up to $16 \times 72 = 1152$ , compared to $d = 72$ for standard linear attention. This is a 16× increase in achievable representational capacity.

When the bound is attainable. The paper notes that this upper bound "is attainable under mild, generic conditions: if each block product $\widetilde{Q}_b Y_b$ has full row rank $r_b = \min(N_b, d)$ and the row spaces of $\{\widetilde{Q}_b Y_b\}_{b=1}^M$ are linearly independent" (Section 4.3). In practice, exact linear independence is unlikely, but "the blockwise mixture still expands the diversity of the row spaces, causing $\text{rank}(A_{\text{MHLA}})$ to grow roughly additively with $M$ ." The empirical evidence in Figure 3(b) confirms this: MHLA's attention-score rank is substantially higher than other linear attention variants across all sequence lengths tested.

The key insight of this analysis: standard linear attention's rank cap of $d$ comes from the fact that every query sees the same key matrix. MHLA breaks this by giving different query-blocks differently scaled key matrices (via the mixing coefficients), so each block submatrix can contribute independent rank to the overall attention matrix. The independence is not guaranteed, but the structural possibility exists and is realized in practice.

Sparsity Analysis: How Block Selection Enables Focused Attention

Beyond rank, the paper analyzes attention entropy to show that MHLA produces more focused, sparse attention distributions.

The mechanism for reduced entropy. The learned coefficient matrix $\mathcal{M}_c$ allows each query-block to assign higher weights to a subset of blocks that are more relevant and lower (near-zero) weights to irrelevant blocks. At the block level, this acts as a coarse sparsification: tokens in low-weight blocks have their effective attention weights scaled down by a factor $m_{i, b(t)}$ that can be close to zero. Within the selected (high-weight) blocks, the kernel inner products $\widetilde{q}^\top \widetilde{K}_t$ further differentiate token contributions.

Two-stage sparsification in operation: consider a query in block $i$ that assigns $m_{i,5} = 0.8$ to block 5 (relevant) and $m_{i,20} = 0.001$ to block 20 (irrelevant). All tokens in block 20 have their contributions suppressed by a factor of $800\times$ compared to tokens in block 5, regardless of how similar their keys are to the query. Within block 5, the kernel similarity then determines which specific tokens get the most attention. The result is a distribution that concentrates mass on a small, semantically relevant subset of tokens — sparse, low-entropy attention.

Empirical validation. Figure 3(b) shows that MHLA consistently yields lower attention entropy than other linear-attention baselines and even lower than softmax attention in some configurations. This is a striking result — MHLA is not just catching up to linear attention baselines but exceeding softmax attention on this metric. The paper interprets this as evidence that "MHLA preserves query-conditioned selectivity and achieves substantially higher sparsity, enabling the model to attend to a small, semantically relevant subset of tokens rather than spreading attention uniformly" (Section 4.3).

Why lower entropy than softmax is possible: softmax attention computes similarities for all token pairs, but the resulting distribution's entropy depends on the temperature (controlled by $1/\sqrt{d}$ ). With small $d$ , softmax can be relatively flat. MHLA's block-level selection provides an additional mechanism for concentrating attention — it can completely down-weight entire blocks before the kernel similarity even comes into play — which can produce sharper distributions than softmax alone when the block structure aligns with the task structure.

Efficiency Analysis: The $O(N d^2 + M^2 d^2)$ Complexity

Decomposing the MHLA computation. The forward pass of MHLA consists of three stages, each with its own cost:

Local summary computation — for each of $M$ blocks, sum $N_b$ outer products of size $d \times d$ . Cost: $O(M \cdot N_b \cdot d^2) = O(N d^2)$ .
Multi-Head Mixing — for each of $M$ query-blocks, compute a weighted sum of $M$ summary matrices of size $d \times d$ . Cost: $O(M^2 d^2)$ . This is a dense matrix-matrix operation (GEMM) between the $M \times M$ coefficient matrix and the stacked summaries.
Output computation — for each of $N$ tokens, compute a query-summary multiplication ( $d \times d$ times $d$ -vector) and a query-normalizer dot product. Cost: $O(N d^2)$ .

Total:

$O(M N_b d^2 + M^2 d^2 + M N_b d^2) = O(N d^2 + M^2 d^2)$

The dominance condition. The paper ensures $M^2 \leq N$ (Section 4.3: "the number of blocks $M$ is usually set to satisfy $M^2 \leq N$ "). Under this condition, $M^2 d^2 \leq N d^2$ , so the $O(N d^2)$ term dominates and the total complexity is $O(N d^2)$ — linear in $N$ with the same asymptotic form as standard linear attention.

Practical scaling verification. Appendix Table 14 profiles MHLA throughput under varying $N$ and $M$ values. When $M^2 < N$ (e.g., $M=16$ , $N=1024$ , so $M^2=256 < 1024$ ), MHLA introduces only negligible overhead over standard linear attention — 51ms vs. 52ms on DiT-S/2, or 118 imgs/s vs. 124 imgs/s on DeiT-S/16. When $M$ is larger (e.g., $M=64$ , $N=256$ , so $M^2=4096 > 256$ ), overhead becomes noticeable — 4.8G vs. 3.7G FLOPs. The paper's ablation in Table 7(b) shows that $M=16$ is already sufficient for strong performance on DiT-S/2 at 512px resolution, so the condition $M^2 \leq N$ is not restrictive in practice.

Memory complexity. Standard linear attention stores the single global summary $G \in \mathbb{R}^{d \times d}$ , consuming $O(d^2)$ memory independent of sequence length. MHLA stores $M$ local summaries, each $d \times d$ , for a total of $O(M d^2)$ memory. Under the condition $M^2 \leq N$ , the memory is still independent of $N$ and the additional factor of $M$ is modest (e.g., $M=16$ means $16 \times$ the memory of standard linear attention for the summary storage, which is typically small relative to the activation memory for the tokens themselves).

Comparison to self-attention's memory: self-attention requires storing the full $N \times N$ attention matrix, so $O(N^2)$ memory. For $N=4096$ and $d=72$ , self-attention stores a $4096 \times 4096$ float matrix (∼67 MB per head per layer), while MHLA stores $M=16$ summaries of size $72 \times 72$ (∼0.04 MB). The difference is roughly three orders of magnitude and grows with sequence length.

Throughput advantage at high resolution. Figure 1(b) shows that at resolution 4096 (where $N$ is large), MHLA maintains throughput nearly identical to linear attention while self-attention throughput drops sharply. At 512 resolution on DiT-S/2, MHLA achieves "better FID scores while doubling the throughput of self-attention" (Section 5.2). This throughput advantage comes from the linear complexity — MHLA avoids the quadratic attention matrix computation that bottlenecks self-attention at long sequences.

Chunkwise Parallel Form for Autoregressive Modeling

The causal masking challenge. In autoregressive language modeling and video generation, a causal mask prevents each token from attending to future tokens. Standard linear attention handles this by maintaining a running global summary $G_t = \sum_{j=1}^t \widetilde{K}_j V_j^\top$ that is updated incrementally as each new token arrives. However, naively recomputing $G_t$ for each prefix during training would cost $O(N^2 d)$ over the full sequence — quadratic again, defeating the purpose of linear attention.

The chunkwise solution (inherited from prior work). The standard remedy, used in architectures like RetNet and GLA, is chunkwise parallel training. The sequence is split into blocks (chunks) of size $C$ . For each block $b$ , a local summary $S_b = \sum_{j \in b} \widetilde{K}_j V_j^\top$ is computed. The global summary is updated recursively across blocks:

$S_i^{\text{global}} = S_{i-1}^{\text{global}} + S_i$

The attention output for block $i$ combines context from previous blocks (via $S_{i-1}^{\text{global}}$ ) with intra-block attention (computed within block $i$ using standard pairwise attention). The per-block cost is $O(C d^2 + C^2 d)$ , and with $L/C$ blocks, the total cost is $O(L d^2 + L C d)$ — linear in $L$ when $C$ is constant. This scheme preserves causality exactly while enabling block-parallel training.

How MHLA adapts this. MHLA replaces the single global summary with query-conditioned mixtures of local summaries. For block $i$ , the mixed summary is restricted to only past blocks (to enforce causality):

$\widetilde{S}_i = \sum_{b \leq i} m_{i,b} S_b$

where $m_{i,b}$ comes from a causal coefficient matrix $\mathcal{M}_c^{\text{causal}}$ in which upper-triangular entries (where $b > i$ , representing future blocks) are masked to enforce causality. The attention output for block $i$ is then:

$H_i = Q_i \widetilde{S}_{i-1} + m_{i,i} (Q_i \widetilde{K}_i^\top) V_i$

What this computes operationally: the first term, $Q_i \widetilde{S}_{i-1}$ , propagates context from all blocks before $i$ through the query-specific mixture $\widetilde{S}_{i-1}$ — this is the cross-block context. The second term, $m_{i,i} (Q_i \widetilde{K}_i^\top) V_i$ , captures intra-block attention within block $i$ itself, scaled by the self-mixing coefficient $m_{i,i}$ (which represents how much block $i$ attends to itself). The two terms are summed to produce the block's output.

Why this preserves causality: $\widetilde{S}_{i-1}$ is formed only from blocks $1$ through $i-1$ , so no future information leaks through the cross-block term. The intra-block term uses only tokens within block $i$ , so within-block causality is handled by standard triangular masking of the $Q_i \widetilde{K}_i^\top$ matrix.

Asymptotic training cost. Because the mixing is performed once per block and reused for all tokens in that block, the cost of forming $\widetilde{S}_i$ is $O(M^2 d^2)$ per layer per sequence, which under $M^2 \leq N$ is subsumed by the $O(N d^2)$ token-level operations. The overall training complexity matches chunkwise linear attention: $O(N d^2 + N C d)$ with block size $C$ .

Causal inference with incremental updates. At inference time, when generating tokens one by one, MHLA maintains the set of past local summaries $\{S_1, \ldots, S_{i-1}\}$ . When a new token arrives in block $i$ , its contribution is incrementally added to the block's summary:

$S_i \leftarrow S_i + \widetilde{K}_t V_t^\top$

The mixed summary $\widetilde{S}_i$ is then updated by applying $m_{i,i}$ to this incremental change — there is no need to recompute the entire mixture over all past blocks, because past blocks' summaries are cached and their contributions to $\widetilde{S}_i$ are unchanged. This keeps per-token inference cost at $O(d^2)$ .

The key design insight for autoregressive modeling: MHLA's block structure naturally maps onto the chunks used in chunkwise parallel training — each MHLA head is a chunk. This means no additional infrastructure is needed to support causal masking; the same block partitioning that enables query-conditioned mixtures also enables efficient autoregressive training and inference. The paper emphasizes this in Section 4.2: "MHLA naturally fits this setting: each head can be directly mapped to a chunk."

Summary of Design Choices and Their Justifications

Spatial block partitioning (not 1D flattening): preserves locality, makes the learned coefficients interpretable as spatial attention patterns, and maps naturally to chunkwise training for autoregressive models.
$M$ local summaries instead of one global summary: creates $M$ atomic context representations that can be differentially combined, breaking the rank-d constraint of standard linear attention.
Learned mixing coefficients $\mathcal{M}_c$ with nonnegativity and clipping: enables data-driven adaptation of per-query-block context retrieval while maintaining training stability through bounded coefficients.
Locality-biased initialization (not random or uniform): provides a strong spatial prior that gives each query-block a distinct context from the start of training; the ablation shows it achieves 75.4% accuracy even without learning, versus 75.1% for uniform initialization.
Two-stage weighting (block selection + intra-block kernel similarity): factorizes the attention computation into a coarse spatial selection and a fine token reweighting, approximating full query-conditioned selectivity without $O(N^2)$ cost.
$M^2 \leq N$ constraint: ensures the $O(M^2 d^2)$ mixing cost is subsumed by the $O(N d^2)$ per-token operations, preserving the linear complexity guarantee. Empirical profiling (Table 14) confirms negligible overhead when this condition holds.
Causal mixture matrix for autoregressive tasks: masks future blocks, maps MHLA heads to chunkwise training chunks, and enables $O(d^2)$ per-token incremental inference without recomputing over past summaries.
No auxiliary modules (no depthwise convolutions, no gating, no hybrid self-attention layers): the paper's core design philosophy is to fix the root cause (the shared global summary) within the attention mechanism itself rather than patching symptoms with external modules that add parameters, FLOPs, and engineering complexity.

4. Key Insights and Innovations

Innovation 1: Global Context Collapse — Naming and Quantifying the Core Failure Mode of Linear Attention

The paper's most important conceptual contribution is not a new mechanism but a diagnosis: identifying, naming, and quantitatively characterizing why linear attention degrades performance. Prior work acknowledged that linear attention underperforms softmax attention — everyone knew there was a gap — but the typical response was to add auxiliary modules (convolutions, gates, hybrid layers) that treated symptoms without isolating the root cause. The field lacked a precise language for what was being lost.

The paper coins the term global context collapse (Section 3.2) and decomposes it into two measurable phenomena that are individually diagnostic and together constitute a unified failure mode:

Rank deficiency as a hard information-theoretic bottleneck. The attention matrix $A_{\text{lin}} = \widetilde{Q} \widetilde{K}^\top$ has rank bounded by $d$ (the feature dimension, typically ≤72 per head) regardless of how many tokens $N$ exist in the sequence. The paper shows this mathematically and validates it empirically in Figure 3(b): standard linear attention's attention-score rank plateaus while softmax attention's rank grows with sequence length. This is not merely an implementation quirk — it is a fundamental linear algebra constraint that means linear attention can represent at most $d$ independent interaction patterns, no matter how long the sequence. Prior work on low-rank bottlenecks in attention (Bhojanapalli et al., 2020) had studied rank in the context of multi-head softmax attention, but the paper's insight is that linear attention's rank problem is qualitatively different: in softmax attention, multi-head design raises the effective rank, but in linear attention, the shared global summary imposes a hard cap that additional heads cannot circumvent because the bottleneck is in the key–value aggregation step, not in the query–key interaction.

Loss of sparsity as a behavioral symptom of the same structural cause. Softmax attention produces sharply peaked, low-entropy distributions because each query interacts individually with every key through an exponential kernel that amplifies differences. Linear attention, by collapsing all key–value pairs into a shared summary $G$ , cannot reweight individual keys according to query-specific relevance — every query sees the same aggregated representation. The paper's entropy measurements in Figure 3(a) quantify this: linear attention consistently shows higher entropy (more uniform attention) than softmax attention, and the visualization of attention maps reveals diffuse, unfocused attention patterns that fail to concentrate on semantically relevant tokens.

Why naming this matters conceptually. Before this paper, the degradation of linear attention was variously attributed to "lost expressivity," "missing local structure," or "kernel approximation error" — all vague and non-specific. By providing a precise, two-dimensional quantitative characterization (rank + entropy), the paper transforms the problem from "linear attention isn't as good" to "linear attention suffers from a specific, measurable representational bottleneck with clear mathematical origins." This changes how subsequent research approaches the problem: instead of asking "what module can I add to make linear attention better?", researchers can ask "does my method increase the rank of the attention matrix?" or "does my method reduce attention entropy?" — concrete, falsifiable diagnostic questions.

Diagnostic significance beyond the proposed solution. The paper uses rank and entropy analysis not only to motivate MHLA but also to evaluate it — Figure 3 shows MHLA achieving both higher rank and lower entropy than competing linear attention variants. But the diagnostic framework stands independently of MHLA. Other researchers could apply the same measurements to evaluate novel linear attention mechanisms, making global context collapse a general evaluation criterion rather than a paper-specific observation. This is the hallmark of a genuine conceptual contribution: it provides tools that outlast the specific method that motivated them.

Comparison to prior diagnostic work. Prior analyses of attention rank (e.g., Bhojanapalli et al., 2020; Dong et al., 2021) focused on softmax attention and argued that low rank limits expressivity in standard Transformers. The paper's contribution is extending this analytical lens to linear attention and showing that the problem is far more severe there — not merely a capacity concern but a hard algebraic bound. Additionally, prior work studied attention entropy (Zhang et al., 2025; Deng et al., 2023) as a correlate of model quality, but the paper connects entropy explicitly to the structural property of query-conditioned selectivity: high entropy in linear attention is a direct consequence of sharing a global summary, not a separate phenomenon.

Innovation 2: Query-Conditioned Selectivity as the Organizing Principle for Linear Attention Design

The paper reframes the design space of linear attention around a single organizing concept: query-conditioned selectivity — the ability for different query tokens to retrieve qualitatively different context from the same set of key–value pairs. This reframing is the intellectual bridge between the diagnosis (global context collapse) and the solution (MHLA), and it clarifies what had been a messy landscape of disparate fixes.

What the field was doing before. Prior approaches to improving linear attention were characterized by symptomatic augmentation:

Convolution-based methods (RALA, Focused Linear Attention, Inline Attention) added depthwise convolutions because linear attention "lost local structure." The implicit assumption was that the missing ingredient was spatial locality.
Gating mechanisms (GLA, Gated DeltaNet) added data-dependent gates because linear attention "couldn't control information flow." The implicit assumption was that the problem was insufficient dynamic range in the recurrent state.
Hybrid architectures (MambaVision, VMamba) combined linear attention with occasional full self-attention layers to "recover global interactions." The implicit assumption was that some quadratic operations were inevitable.

Each of these approaches addressed a symptom of the underlying problem without explicitly naming what was fundamentally lost. The paper's reframing reveals that all of these symptoms — lost local structure, uncontrolled information flow, degraded global interactions — are consequences of a single root cause: the absence of per-query adaptation in the context retrieval mechanism.

The conceptual move. The paper identifies query-conditioned selectivity as the property that distinguishes softmax attention from standard linear attention, and argues that restoring this property is the sufficient design objective — if a linear attention mechanism can re-enable per-query diversity in context retrieval, the other symptoms (lost local structure, high entropy, low rank) should resolve automatically without explicit auxiliary modules. This is a shift from "what's missing?" to "why is it missing, and what single change would recover it?"

What makes this reframing novel. Query-conditioned computation is not a new concept — it's implicit in softmax attention and has been discussed in the context of dynamic routing and selective state-space models. The innovation is making it the explicit design principle for linear attention and showing that a single, simple mechanism (learned block-level mixing coefficients) can achieve it without auxiliary modules, without quadratic operations, and without domain-specific assumptions. The paper's Appendix B formalizes this by contrasting the mathematical forms: softmax attention provides query-conditioned per-token weighting; standard linear attention provides query-independent summary-based retrieval; MHLA restores query-conditioned block-level weighting.

Why this matters beyond MHLA. By articulating query-conditioned selectivity as the organizing principle, the paper establishes a lens through which to evaluate all linear attention mechanisms — past, present, and future. A researcher can ask: "Does this method give different queries different effective views of the key–value context?" If the answer is no, the method is fundamentally limited regardless of its other virtues. If yes, the method addresses the root cause. This is a more principled and general design criterion than "add convolutions" or "add gates," and it points toward a research program rather than a single solution.

Evidence of the principle's validity. Figure 3's double dissociation — MHLA simultaneously achieves higher rank (more diverse representations) and lower entropy (more focused attention) while using no auxiliary modules — is evidence that query-conditioned selectivity, when restored, naturally produces both improved expressivity and improved sparsity. The paper does not need separate mechanisms for rank and for entropy; the single design principle addresses both because they are consequences of the same underlying cause.

Innovation 3: Block-Level Mixing as a Factorized Approximation of Full Pairwise Attention — A Provably Linear-Complexity Alternative to the Global Summary

The paper's mechanism-level innovation — the learned coefficient matrix $\mathcal{M}_c$ that enables query-block-specific mixtures of local key–value summaries — represents a structural insight rather than an architectural hack: that full query-conditioned token-level selectivity can be factorized into a coarse spatial selection (block-level mixing) and a fine token reweighting (intra-block kernel similarity) without incurring quadratic cost.

The intellectual counterweight to the dominant assumption. The field has largely operated under an implicit assumption that the choice is binary: either you compute full pairwise attention (softmax, $O(N^2)$ ) or you compress all tokens into a shared summary (linear attention, $O(N d^2)$ ). Prior attempts to find a middle ground — sparse attention patterns (Longformer, BigBird), low-rank approximations (Linformer), kernel methods with random features (Performer) — either retained some quadratic dependence or sacrificed generality for specific sparsity patterns. The paper's insight is that there exists a second decomposition strategy: instead of sharing one summary across all queries or computing all pairwise scores, you can compute multiple local summaries and let each query-group combine them differently using a learned mixing matrix.

Why this factorization works. The two-stage mechanism — $m_{i, b(t)} \cdot (\widetilde{q}^\top \widetilde{K}_t)$ — decomposes the attention weight into a factor that depends on which query-block is asking ( $m_{i, b(t)}$ ) and a factor that depends on which specific token is being attended to ( $\widetilde{q}^\top \widetilde{K}_t$ ). This factorization is not mathematically equivalent to softmax attention's $\exp(q_i^\top k_j)$ , but it captures the same essential structure: different queries can assign different importance to the same token. The factorization is what makes the computation linear: the block-level summaries $S_b$ are precomputed once ( $O(N d^2)$ ), the mixing is a small $M \times M$ operation ( $O(M^2 d^2)$ , subsumed when $M^2 \leq N$ ), and the per-token interaction with the mixed summary is $O(d^2)$ .

The rank analysis as theoretical justification for the factorization. Section 4.3's proof that $\text{rank}(A_{\text{MHLA}})$ can grow additively with $M$ — reaching up to $\min(N, M \cdot d)$ versus the standard linear attention bound of $d$ — is the theoretical demonstration that the factorization successfully escapes the rank bottleneck. This is not an empirical observation but a mathematical consequence of the block structure: different query-blocks see differently scaled key matrices, so their attention submatrices can span independent row spaces. The paper's empirical rank measurements in Figure 3(b) confirm that this theoretical capacity is realized in practice.

Comparison to prior factorization approaches. The closest conceptual relative is probably sparse attention patterns (Beltagy et al., 2020; Zaheer et al., 2020), where each token attends to a fixed or learned subset of other tokens. MHLA's block-level mixing can be seen as a soft, learned, global sparsity pattern — every query-block can attend to every key-block, but with learned, continuous weights rather than hard binary masks. This combines the flexibility of full attention (all blocks are accessible) with the efficiency of sparse attention (the mixing matrix is small relative to the full $N \times N$ attention matrix).

Another comparison is to multi-query attention (Shazeer, 2019) and grouped-query attention (Ainslie et al., 2023), where multiple query heads share key–value projections to reduce memory. These methods reduce the number of key–value representations but still compute full pairwise attention within each shared group. MHLA inverts this logic: it increases the number of key–value summaries (from 1 to $M$ ) but avoids pairwise computation by factorizing through the mixing matrix.

The significance is in the provable efficiency–expressivity tradeoff. Table 1 is the paper's succinct statement of what MHLA achieves: linear time complexity, memory independent of $N$ , a rank bound that scales with $M$ , and restored query-conditioned selectivity — all without auxiliary modules. No prior linear attention mechanism satisfies all four conditions simultaneously. This table, while simple, represents a genuine advance in the theoretical understanding of what's achievable within the linear-complexity constraint.

Innovation 4: Auxiliary-Module-Free Design as a Principled Methodology, Validated by Empirical Scaling Behavior

The paper's design philosophy — that linear attention should be fixed from within rather than augmented from without — is not merely an aesthetic preference. It is validated by a specific empirical phenomenon that the paper documents and that has implications for how the field approaches efficient attention design: auxiliary modules provide diminishing (or negative) returns at larger model scales.

The empirical evidence from DiT scaling. Table 3(a) shows a progression across DiT model sizes that tells a clear story. At DiT-S/2, MHLA with no auxiliary modules achieves FID 76.4; adding CPE (a depthwise convolution) improves to 64.0; adding gating on top of that further improves to 59.8. The auxiliary modules help at small scale. But at DiT-L/2, the pattern shifts: plain MHLA achieves 25.37, adding CPE improves to 24.21 (a smaller relative gain), and adding CPE+gating reaches 21.37. At DiT-XL/2, the pattern reverses: plain MHLA already achieves 20.32 — close to self-attention's 19.47 — and adding CPE degrades performance to 22.79. Only CPE+gating narrowly beats self-attention (19.17).

What this scaling behavior reveals. Auxiliary modules like depthwise convolutions add inductive biases (e.g., locality, translation equivariance) that are beneficial when the model's own capacity is insufficient to learn those patterns from data. At small scales, these biases compensate for limited parameters. But as model capacity grows, the biases become constraints — the model could learn more flexible, data-driven attention patterns if not forced through a fixed convolutional structure. MHLA's learned mixing matrix $\mathcal{M}_c$ , by contrast, is capacity-adaptive: it can learn any attention pattern over blocks, from purely local to purely global to anything in between, and its expressivity grows with the model's ability to learn meaningful coefficients.

The implication for efficient attention research. The paper's finding suggests a methodological principle: when evaluating a new efficient attention mechanism, report performance across a range of model scales, not just at the smallest convenient size. A method that looks promising at the DeiT-T or DiT-S scale because it cleverly adds spatial priors may be a dead end at larger scales where those priors become limiting. The paper's auxiliary-module-free approach is not just cleaner — it scales better, and the scaling behavior is evidence that the method addresses the fundamental bottleneck rather than patching symptoms.

Comparison to prior work's methodology. Most prior linear attention papers (RALA, Focused LA, Inline Attention, MALA) report results at one or two model scales, often small ones where auxiliary modules show the largest relative benefit. The paper's systematic scaling across four DiT sizes (S, B, L, XL) and two DeiT sizes (T, S) — plus the additional scaling analysis in Appendix F.3 (higher resolution classification) and F.4 (throughput profiling) — sets a higher methodological bar. It doesn't just claim that MHLA works; it characterizes where and why alternative approaches break down.

Relation to the "bitter lesson." The paper's approach aligns with the broader principle that learned, general mechanisms tend to outperform hand-designed, domain-specific ones at sufficient scale. The mixing matrix is more general than a fixed convolution; it learns spatial attention patterns rather than having them hard-coded. The fact that it outperforms convolution-augmented variants only at larger scales is consistent with the pattern observed in many ML domains: inductive biases help when data or capacity is limited, but learned flexibility wins at scale.

Caveat on the strength of this innovation. The "no auxiliary modules" principle is a design philosophy, not a theoretical contribution, and it's not entirely novel — some prior linear attention variants (e.g., basic Performer, basic Linear Transformer) were also auxiliary-module-free. What's novel is the combination of this philosophy with (a) a specific mechanism that achieves competitive performance without modules, (b) empirical evidence that the modules become counterproductive at scale, and (c) the reframing of the problem around query-conditioned selectivity, which explains why the modules were never going to be the right fix — they address symptoms (local structure) rather than the root cause (lack of per-query adaptation).

5. Experimental Analysis

Evaluation Methodology

Dataset. The paper evaluates on four domains using standard benchmarks. For image classification: ImageNet-1K (1.28M training images, 50K validation images, 1000 classes), following the setup in prior linear attention work (RALA, MALA, Focused LA). For class-to-image generation: ImageNet-1K at resolutions 256×256 and 512×512, trained for 400K steps. For text-to-image generation: a custom collection of 31,292K images from the internet, used to fine-tune a pretrained SANA-0.6B checkpoint for 40K steps. For NLP: a 5B-token subset of SlimPajama for training from scratch, with evaluation on standard commonsense reasoning benchmarks (MMLU, CSRs, WinoGrande, PIQA, ARC, OpenBookQA, BoolQ) and LongBench for long-context understanding. For video generation: fine-tuning on 81-frame videos at 480×800 resolution (sequence length 31,500 tokens), evaluated on VBench.
Base model(s). The paper integrates MHLA into multiple architectures across scales, chosen to test generality rather than optimize a single model. For image classification: DeiT-T and DeiT-S (Touvron et al., 2021) and VLT-T and VLT-S (RALA, 2024). For class-to-image generation: DiT-S/2, DiT-B/2, DiT-L/2, DiT-XL/2 (Peebles and Xie, 2023) and DiG-S/2 (Zhu et al., 2025). For text-to-image: SANA-0.6B (Xie et al., 2024). For NLP: a 340M-parameter Transformer++ trained from scratch following the GLA recipe (Yang et al., 2024). For video generation: a pretrained Wan2.1-1.3B model (replacing FlashAttention with MHLA or vanilla linear attention). The diversity of architectures (vision transformers, diffusion transformers, autoregressive LMs, video diffusion models) and scales (6M to 1.3B parameters) is deliberate: the paper aims to establish MHLA as a general-purpose attention replacement, not a task-specific trick.
Metrics. Classification uses Top-1 accuracy (%) on the ImageNet-1K validation set. Class-to-image generation uses FID (Fréchet Inception Distance) as the primary metric, with supplementary IS (Inception Score), sFID, Precision, and Recall reported in Appendix F.1 (Table 10). Text-to-image generation reports FID, CLIP score, and GenEval. NLP reports perplexity (ppl) on held-out text, zero-shot accuracy on 7 commonsense reasoning benchmarks, MMLU accuracy, and LongBench task-specific scores (averaged within Multi-Doc QA, Single-Doc QA, Few-shot, Synthetic, Summarization, and Code categories, plus an overall average). Video generation reports VBench Quality, Semantic, and Total scores, plus inference latency (seconds). Throughput is measured in images/second (classification) or milliseconds (generation) on Nvidia H100 GPUs.
Baselines. The paper compares against a comprehensive set of attention mechanisms. For image classification: Self Attention (standard softmax), Linear Attention (vanilla kernelized linear attention without auxiliary modules), Focused Linear Attention (Han et al., 2024), Inline Attention (Han et al., 2024), MALA (2024), and RALA (2024). For class-to-image generation: Self Attention and Linear Attention in DiT, and Gated Linear Attention (GLA) (Yang et al., 2024) in DiG. For NLP: Transformer++ (Touvron et al., 2023, a strong softmax-attention baseline with RoPE, SwiGLU, and RMSNorm), GLA (Y ang et al., 2024), Mamba (Gu and Dao, 2023), Mamba2 (Dao and Gu, 2024), and Gated DeltaNet (GDN) (Yang et al., 2025). For video generation: the original Wan2.1-1.3B with FlashAttention (Wan-FA) and a Wan-LA variant with all attention layers replaced by vanilla linear attention. For text-to-image: PixArt-α, PixArt-Σ, and the official SANA checkpoint.
Generation budget / compute accounting. The paper measures compute through several complementary metrics to enable fair comparisons across attention types. Parameter count is reported for all classification models (Table 2) to verify MHLA adds negligible parameters (0M over self-attention in DeiT, 0M in DiT when used without CPE/gating). FLOPs are reported for classification (Table 2) and profiled across varying N and M in Appendix F.4 (Table 14). Throughput (images/second or milliseconds) is benchmarked on Nvidia H100 GPUs at multiple resolutions (Figure 1b, Figure 1d, Table 7b). Inference latency is reported for video generation (Table 5) in seconds per sample. The paper does not use a unified "generation budget" metric because the tasks are too heterogeneous; instead, it ensures that within each task, the comparison to baselines holds parameters, FLOPs, or throughput constant. For the throughput comparisons in Figure 1, MHLA is measured against self-attention and linear attention on the same hardware at the same resolution, making the speedup claims directly interpretable.
Cross-validation / statistical protocol. The paper reports mean and standard deviation over three independent runs for key image generation results in Appendix F.1 (Table 11), explicitly stating: "we report the mean and standard deviation of MHLA over three independent runs to demonstrate the stability of our results." For classification, the standard ImageNet training protocol (300 epochs, fixed seed practices from DeiT) is followed without explicit cross-validation — this is consistent with the baseline papers (DeiT, RALA, MALA) that MHLA compares against. For NLP, a single training run is reported following the GLA protocol, which is standard practice in the linear attention literature at this scale. The paper does not employ two-fold cross-validation or multiple seeds for the main classification and NLP results, which is a limitation — some of the reported differences (e.g., 75.8% vs. 75.1% on DeiT-T) are within the range where seed variance could matter.

Main Quantitative Results

Image Classification (Section 5.1, Tables 2(a) and 2(b))

Headline result. MHLA achieves state-of-the-art accuracy among linear attention methods while adding zero extra parameters over vanilla self-attention. On DeiT-T, MHLA reaches 75.8% Top-1 accuracy — a +3.6% absolute improvement over self-attention's 72.2% and +6.0% over vanilla linear attention's 69.8% (Table 2(a)). On DeiT-S, MHLA reaches 81.0%, beating self-attention's 79.8% by +1.2% and vanilla linear attention's 77.6% by +3.4%.

Comparison to prior linear attention methods at matched parameters and FLOPs. Table 2(a) reports all models at the same parameter count and FLOPs where possible. On DeiT-T at 5.7M parameters and 1.1G FLOPs: Focused LA achieves 74.1% (+4.3% over vanilla linear attention), Inline Attn reaches 74.5% (+4.7%), MALA reaches 75.1% (+5.3%), and MHLA reaches 75.8% (+6.0%). The trend holds on DeiT-S at 22M parameters: RALA achieves 80.4% (+2.8% over vanilla LA), MALA achieves 80.3% (+2.7%), and MHLA achieves 81.0% (+3.4%). The key differentiator is that Focused LA, Inline Attn, and MALA all add parameters (0.4M–0.8M on DeiT-T, 2M on DeiT-S) and FLOPs to achieve their gains, while MHLA adds no parameters and no FLOPs beyond standard self-attention. The paper emphasizes this: "We reach the best accuracy in linear attention across all model sizes, while introducing the fewest extra parameters compared with baselines" (Section 5.1).

Comparison to state-of-the-art vision architectures. Table 2(b) integrates MHLA into the VLT architecture (RALA, 2024) and compares against a broader set of efficient vision models at matched compute budgets (~2.5G and ~4.5G FLOPs). At ~2.5G FLOPs, MHLA-VLT-T achieves 82.6% accuracy, outperforming FL-PVT-T (77.8%), FL-PVTv2-B1 (79.5%), MSVMamba-M (79.8%), NAT-M (81.8%), RAVLT-T (82.3%), and MAViT-T (82.4%). At ~4.5G FLOPs, MHLA-VLT-S achieves 84.6%, outperforming FAT-B3 (83.6%), Vmamba-T (82.6%), MV-T (82.3%), MSVMamba-T (83.0%), and MAViT-S (84.3%). The margin over the closest competitor is 0.2–0.3% at both scales, which is modest but consistent — and achieved without the hybrid attention or SSM components that some baselines use.

High-resolution classification (Appendix F.3, Table 13). On DeiT-T at 384×384 resolution, MHLA achieves 77.5% vs. self-attention's 74.4% (+3.1%). At 512×512, MHLA achieves 78.3% vs. 75.3% (+3.0%). The improvement is consistent across resolutions, suggesting the mechanism scales with sequence length.

Class-to-Image Generation (Section 5.2, Table 3)

Headline result. MHLA achieves FID scores that match or exceed self-attention on DiT models at all scales, while maintaining throughput roughly 2× higher than self-attention at 512 resolution. On DiT-S/2 at 256 resolution, MHLA achieves FID 59.80 vs. self-attention's 68.40 (a 12.6% relative improvement) and vanilla linear attention's 89.72 (Table 3(a)). On DiT-XL/2 at 256 resolution, plain MHLA (no CPE, no gating) achieves FID 20.32, within 0.85 of self-attention's 19.47, while vanilla linear attention trails at 28.63.

Scaling behavior: auxiliary modules help at small scale, hurt at large scale. The progression across DiT sizes reveals diminishing returns from CPE (depthwise convolution) and gating:

DiT-S/2 (256px): MHLA alone = 76.4; +CPE = 64.0; +Gating = 68.5; +CPE+Gating = 59.8. CPE provides a large gain (+12.4 FID) at this scale.
DiT-B/2 (256px): MHLA alone achieves 37.47 (already beating self-attention's 43.47), with no CPE or gating reported — suggesting plain MHLA is sufficient by this scale.
DiT-L/2 (256px): MHLA alone = 25.37; +CPE = 24.21 (smaller gain of +1.16); +CPE+Gating = 21.37. The marginal benefit of CPE shrinks.
DiT-XL/2 (256px): MHLA alone = 20.32 (near self-attention's 19.47); +CPE = 22.79 (degradation of 2.47 FID); +CPE+Gating = 19.17 (modest gain over plain MHLA, narrowly beating self-attention).

The paper interprets this as evidence that "although modules like DWConv may offer gains at small scales, their benefits do not scale with model size or sequence length" (Section 5.2). At XL scale, the learned mixing matrix alone provides sufficient expressivity, and the fixed spatial prior of convolution becomes a constraint.

Resolution scaling. At 512 resolution on DiT-S/2, MHLA achieves FID 78.63 vs. self-attention's 84.54 and vanilla linear attention's 125.33 (Table 3(a)). The gap between MHLA and self-attention widens at higher resolution (+5.91 FID at 512 vs. +8.60 at 256), consistent with the hypothesis that MHLA's linear complexity provides an advantage over quadratic self-attention as sequence length grows. The paper reports that MHLA "achieves better FID scores while doubling the throughput of self-attention" at 512 resolution (Section 5.2).

Fast adaptation from pretrained checkpoints. Fine-tuning a pretrained DiT-XL/2 for 400K steps with MHLA produces FID 8.34 (without CFG), outperforming the original self-attention model's 9.62 (Table 3(b)). With classifier-free guidance, MHLA achieves FID 2.54 vs. self-attention's 2.27 — a slight degradation but still highly competitive. This demonstrates that MHLA can be retrofitted into existing pretrained models with minimal performance loss and rapid adaptation.

Comparison to GLA (Table 3(a)). On DiG-S/2, which uses Gated Linear Attention (Yang et al., 2024) as the baseline, MHLA achieves FID 59.49 at 256 resolution vs. GLA's 62.06. At 512 resolution, the gap widens: MHLA at 78.63 FID in DiT-S/2 with MHLA (directly comparable) vs. GLA's 99.04 in DiG-S/2 — though this comparison is imperfect because the base architectures differ, the trend suggests MHLA scales better with sequence length.

Text-to-Image Generation (Section 5.2, Table 4)

Headline result. Fine-tuning the pretrained SANA-0.6B model for 40K steps with MHLA replacing the original linear attention layers yields FID 5.90, CLIP score 28.26, and GenEval 0.68, improving over the official SANA checkpoint (FID 6.10, CLIP 28.15, GenEval 0.64) across all three metrics (Table 4). MHLA-SANA also outperforms PixArt-α (FID 6.14, CLIP 27.55, GenEval 0.48) and PixArt-Σ (FID 6.34, CLIP 27.62, GenEval 0.52).

Training dynamics (Figure 5). The loss curves show MHLA-based SANA rapidly adapting: it "matching the pretrained checkpoint within the first 2k steps and subsequently converging to a lower loss" (Section 5.2). This rapid adaptation from a pretrained checkpoint — with only 40K fine-tuning steps — demonstrates that MHLA can be a drop-in replacement for existing linear attention modules in production text-to-image models without extensive retraining.

Video Generation (Section 5.3, Table 5)

Headline result. At an extreme sequence length of 31,500 tokens (81 frames at 480×800), vanilla linear attention catastrophically fails: Wan-LA achieves only 58.24 Total on VBench vs. Wan-FA's 83.31 (a 30% relative degradation), with the Semantic score collapsing to 11.38 vs. 75.65 (Table 5). This is the paper's most dramatic demonstration of global context collapse — at ultra-long sequences, the shared global summary becomes essentially useless. MHLA recovers to 82.62 Total, nearly matching Wan-FA's 83.31 while providing a 2.1× inference speedup (81s vs. 166s latency). The hybrid variant (Wan-MHLA-H, replacing only 2/3 of attention layers) achieves the best overall performance at 83.82 Total, surpassing the original FlashAttention model, with a 1.6× speedup (103s latency).

Training dynamics (Figure 6). The loss curves for video generation provide striking evidence of the failure mode. Wan-FA (the original FlashAttention model) shows a steadily decreasing loss. Wan-MHLA "rapidly adapts during fine-tuning and quickly approaches the pretrained model's loss trajectory." Wan-LA (vanilla linear attention) "effectively fails to train under such long sequences, with its loss plateauing at a high level." This is the paper's strongest empirical validation of the global context collapse diagnosis: vanilla linear attention does not just underperform — it stops learning entirely at extreme sequence lengths, providing a controlled demonstration that the degradation is structural, not just a matter of suboptimal hyperparameters.

Key detail on the hybrid variant. The paper states Wan-MHLA-H replaces only 2/3 of layers with MHLA, keeping the remaining 1/3 as FlashAttention. This hybrid achieves the best VBench Total score (83.82) and a 1.6× speedup, suggesting that even partial adoption of MHLA can provide substantial efficiency gains with minimal performance impact. The paper does not report which specific layers were replaced or whether the choice matters, which is a missing detail.

Natural Language Processing (Section 5.4, Tables 6 and 8)

Headline result (perplexity and commonsense reasoning). On a 340M-parameter model trained from scratch on 10B tokens, MHLA achieves competitive but not dominant performance on perplexity and commonsense reasoning (Table 6). MHLA achieves MMLU 23.7 — the best among all baselines (Transformer++: 22.9, Mamba: 23.5, Mamba2: 23.0, GDN: 23.0, GLA: 22.9). On the commonsense reasoning average, MHLA scores 47.1, essentially tied with Mamba2 (47.0), GDN (46.9), and Transformer++ (46.8). On perplexity, MHLA at 71.64 (WikiLM) trails Transformer++ (60.46) and Mamba2 (58.51) substantially — this is the one metric where MHLA does not match the best baselines. The paper acknowledges this indirectly by reporting the numbers without claiming dominance.

Headline result (long-context understanding). On LongBench, MHLA achieves the highest average score of 7.41, outperforming all recurrent baselines (Table 8). The advantage is most pronounced in Multi-Doc QA (3.58 vs. 3.37 for Mamba), Summarization (18.59 vs. 18.36 for Mamba), and Code (12.72 vs. 12.55 for GLA). The paper attributes this to MHLA's query-conditioned selectivity enabling more precise retrieval across long contexts: "The result demonstrates the superior long context understanding capability of the proposed MHLA" (Section 5.4).

Interpretation of the NLP results. The NLP experiments are arguably the weakest domain for MHLA — it matches but does not clearly exceed strong baselines on commonsense reasoning, and it lags on perplexity. However, the paper's central claim is not that MHLA is universally superior to all sequence models; it is that MHLA restores the expressivity lost by vanilla linear attention while maintaining linear complexity. The NLP results support this narrower claim: MHLA dramatically outperforms GLA (the most comparable linear attention baseline) on LongBench (7.41 vs. 6.53 average) and MMLU (23.7 vs. 22.9), and it is competitive with the strongest baselines (Mamba2, Transformer++) on most metrics. The perplexity gap suggests MHLA may be a better retrieval mechanism than a compression mechanism — it excels at selective attention across long contexts but may not model local token statistics as precisely as state-space models or softmax attention.

Absence of a controlled comparison to vanilla linear attention in NLP. A notable omission: the paper does not report a vanilla linear attention baseline in the NLP experiments (Tables 6 and 8). GLA is the closest, but GLA includes gating mechanisms. This makes it harder to isolate how much of MHLA's NLP improvement comes from restoring query-conditioned selectivity versus how much GLA already recovers. The video generation experiment (Table 5), where vanilla LA is included and collapses, provides the cleanest controlled comparison.

Ablation Studies and Robustness Checks

Multi-Head Mixing initialization strategy (Table 7(a)): On DeiT-T, locality-biased initialization without learning (frozen coefficients) achieves 75.4%, demonstrating that the spatial prior alone provides a strong inductive bias. Uniform initialization with learning achieves 75.1% — worse than frozen locality-biased initialization, suggesting that standard optimization struggles to discover good mixing patterns from a symmetric start. The full model (locality-biased initialization + learnable coefficients) achieves 75.8% , confirming that the initialization provides a foundation that training can refine but cannot easily discover from scratch.

Head number M (Table 7(b)): On DiT-S/2 at 512 resolution (N=1024 tokens), M={4, 16, 64} achieves FID {79.56, 78.63, 79.50} respectively, with throughput {435, 435, 408} imgs/s. The optimal M is 16 — small enough to satisfy M² ≤ N (256 ≤ 1024) and thus incur negligible overhead, large enough to achieve good performance. M=64 (where M²=4096 > 1024=N) drops throughput by ~6% while FID degrades slightly to 79.50, validating the paper's claim that M² ≤ N is both sufficient and necessary for efficiency. M=4 (only 4 blocks for 1024 tokens) shows worse FID (79.56), confirming that too few blocks limit the diversity of the mixed summaries.

CPE and output gating (Appendix F.2, Table 12): On DiT-S/2 at 256 resolution, MHLA alone achieves FID 76.4. Adding CPE improves to 64.0. Adding gating alone improves to 68.5. Adding both improves to 59.8. The ablation shows that CPE and gating are orthogonal optimizations — each provides independent benefit at small scale — but the paper's DiT-XL results (Table 3(a)) demonstrate these benefits diminish or reverse at scale. This is a critical ablation because it distinguishes MHLA's core contribution (block-level mixing) from the auxiliary modules that prior work relied on.

High-resolution classification (Appendix F.3, Table 13): On DeiT-T at 384×384, MHLA improves from 74.4% to 77.5% (+3.1%). At 512×512, from 75.3% to 78.3% (+3.0%). The consistency of the improvement (~3%) across resolutions confirms that MHLA's mechanism scales with sequence length — the rank advantage grows with N, compensating for any increased difficulty.

Stability across runs (Appendix F.1, Table 11): On DiT image generation, MHLA's FID over three independent runs shows standard deviations of 0.03–0.10 FID across model scales. DiT-S/2 with MHLA: 59.744 ± 0.100. DiT-B/2: 37.519 ± 0.039. DiT-L/2: 21.426 ± 0.051. DiT-XL/2: 19.164 ± 0.031. These tight standard deviations suggest the learned mixing matrix converges stably and the reported gains are not due to lucky seeds.

Throughput profiling under varying N and M (Appendix F.4, Table 14): Across N={256, 1024, 4096} and M={4, 16, 64, 256, 1024}, the profiling demonstrates that MHLA's overhead is negligible when M² < N (e.g., M=16, N=1024: 51ms for MHLA vs. 52ms for standard linear attention on DiT-S/2). When M is large relative to N (M=256, N=1024: 61ms, a ~17% increase), the M² d² mixing cost becomes noticeable. This empirically validates the O(N d² + M² d²) complexity analysis.

Critical Assessment

Claim 1: MHLA achieves a 3.6% improvement over self-attention on ImageNet classification. This claim is supported by Table 2(a) (75.8% vs. 72.2% on DeiT-T), but the comparison needs qualification. The 3.6% figure is specifically for DeiT-T, and on DeiT-S the margin shrinks to +1.2% (81.0% vs. 79.8%). The paper does not report MHLA integrated into larger classification models (DeiT-B, ViT-L), so it's unclear whether this trend of diminishing relative gain continues. More importantly, the self-attention baseline at 72.2% on DeiT-T is relatively weak — modern training recipes and architecture improvements (better augmentations, LayerScale, etc.) can push DeiT-T higher, which might narrow the gap. The comparison to state-of-the-art models in Table 2(b) shows MHLA-VLT achieving the highest accuracy, but by margins of 0.2–0.3% — competitive but not transformative. The stronger claim is that MHLA achieves the best accuracy among linear attention mechanisms, not that it universally dominates self-attention.

Claim 2: MHLA improves DiT-based class-to-image generation by 12.6%. This claim is supported for the specific comparison of DiT-S/2 with MHLA (FID 59.80) vs. DiT-S/2 with self-attention (FID 68.40), which represents an 12.6% relative improvement in FID. However, the framing is slightly misleading: the 12.6% figure is the relative improvement at the smallest model scale (DiT-S). At DiT-XL/2, plain MHLA (20.32) is worse than self-attention (19.47) by 0.85 FID, and only beats it when augmented with CPE+gating (19.17 vs. 19.47). The fair summary is: MHLA matches or slightly exceeds self-attention at large scales, and provides large gains over vanilla linear attention at all scales. The 12.6% figure should be understood as the maximum observed improvement, not a universal one.

Claim 3: MHLA achieves a 41% improvement over vanilla linear attention in video generation. This claim is strongly supported by Table 5: Wan-MHLA achieves 82.62 Total vs. Wan-LA's 58.24 (a 41.9% relative improvement, reclaiming most of the 30% gap between Wan-LA and Wan-FA). The video generation experiment is the paper's cleanest and most compelling result because it provides the strongest controlled comparison: same model, same data, same training, same sequence length — only the attention mechanism differs. The near-total collapse of Wan-LA (Semantic score dropping from 75.65 to 11.38) and the near-full recovery by MHLA (Semantic 76.16) is a textbook demonstration that global context collapse is real, severe, and addressable by the proposed mechanism.

Claim 4: MHLA maintains linear complexity while restoring expressivity. The throughput measurements (Figure 1, Tables 7b and 14) and the theoretical analysis (Section 4.3, Table 1) support this claim. MHLA's throughput is nearly identical to linear attention (51ms vs. 52ms at N=1024 on DiT-S/2) and substantially higher than self-attention at long sequences (2× at 512 resolution). The paper demonstrates that the M² d² mixing overhead is negligible under the M² ≤ N condition, and ablations show M=16 is sufficient for strong performance on sequences of length 1024+. However, the paper only profiles throughput, not memory usage during training. Linear attention's main advantage is often memory (O(d²) vs. O(N²) for the attention matrix), not just speed. The paper reports memory complexity in Table 1 (O(M d²) for MHLA) but does not provide empirical memory measurements, which would strengthen this claim for practitioners deciding whether to adopt MHLA.

A critical missing experiment: the normalizer ablation. The paper mentions in Section 4.1 that "In tasks like language modeling and video generation, the normalizer term can be omitted for better training stability when the sequence is getting longer." This is a significant design choice — omitting the normalizer means the attention output is no longer a proper weighted average, which could affect training dynamics. The paper does not report whether the normalizer was omitted for the NLP or video generation experiments, nor does it provide an ablation comparing normalized vs. unnormalized MHLA on any task. This is a genuine gap: the normalizer's presence or absence changes the mathematical properties of the attention mechanism, and readers deserve to know which variant was used for which results and what the sensitivity is.

A second missing experiment: comparison to sparse attention at matched efficiency. The paper positions MHLA as an alternative to both linear attention and self-attention, but doesn't compare against sparse attention mechanisms (Longformer, BigBird) or local+global hybrids at matched throughput. On image classification, Swin Transformer's windowed attention achieves strong results with linear complexity — how does window size in Swin compare to block size M in MHLA? On long-context NLP, sparse attention patterns (sliding window + global tokens) are a practical baseline that the paper doesn't include. This omission is understandable given the paper's focus on linear attention, but it limits the claim that MHLA is a general-purpose drop-in replacement — practitioners choosing between sparse attention and MHLA would need to run their own comparisons.

Dataset and model scale limitations. The NLP experiments use only 10B training tokens — modern language models are trained on trillions. The paper's 340M model may not be representative of behavior at the 7B+ scale where linear attention is most practically relevant. On classification, the largest model tested is VLT-S (~27M parameters) — MHLA's behavior on ViT-L or ViT-H is unknown. The text-to-image experiment uses only 40K fine-tuning steps rather than training from scratch, so it demonstrates compatibility but not that MHLA would be optimal for a from-scratch SANA-scale training run. These are not fatal weaknesses — no paper can test every scale — but they mean the "universal drop-in replacement" framing is aspirational rather than fully demonstrated.

The LongBench result deserves scrutiny. While MHLA achieves the highest LongBench average (7.41), the individual task scores are within typical variance for models of this size. The largest absolute advantage is in Summarization (18.59 vs. 18.36 for Mamba, a 0.23 difference on a scale where scores span 11–18), and in Multi-Doc QA (3.58 vs. 3.37 for Mamba, a 0.21 difference on a scale where scores span 1–5). Without confidence intervals, it's unclear whether these differences are statistically reliable. The paper reports these as point estimates from single training runs, which is standard in the linear attention literature but limits the strength of the "state-of-the-art" claim.

The paper's strongest result is the video generation experiment, and it deserves emphasis as such. This is where the paper's diagnosis (global context collapse), proposed mechanism (MHLA), and empirical demonstration align most cleanly: a controlled A/B test at extreme sequence length where vanilla linear attention collapses, MHLA nearly fully recovers, and the hybrid variant even exceeds the FlashAttention baseline. This single experiment, more than any other, validates the paper's central thesis. The weaker results on NLP perplexity and the diminishing margins on classification suggest that MHLA's advantages are most pronounced in the regime it was designed for: very long sequences where the rank bottleneck of standard linear attention becomes acute.

6. Limitations and Trade-offs

The Inherent Competence Ceiling: MHLA Cannot Recover Capabilities the Base Model Lacks

The paper is transparent about a fundamental boundary condition of its approach, though it does not explore the consequence experimentally. MHLA restores query-conditioned selectivity within linear attention, but it operates entirely on the key–value representations produced by the base model's learned projections. If a task requires reasoning patterns, factual knowledge, or representational structures that the base model fails to encode in its $W_K$ and $W_V$ projections, no amount of mixing-block sophistication can compensate. This is the same competence-ceiling principle that governs all linear attention methods: the expressivity gain is multiplicative (diversifying how queries retrieve context) rather than additive (introducing new semantic capabilities absent from the key–value representations themselves).

The evidence for this limitation is distributed rather than concentrated in a single ablation, but the video generation experiment in Section 5.3 (Table 5) provides the clearest demonstration in negative form. When vanilla linear attention (Wan-LA) collapses to a VBench Total of 58.24 — with the Semantic score cratering to 11.38 — the failure is catastrophic: the model essentially stops producing semantically meaningful outputs. MHLA recovers to 82.62, nearly matching the Wan-FA FlashAttention baseline at 83.31. This recovery demonstrates that MHLA can restore query-conditioned access to representations the base model does possess, which standard linear attention destroys through global context collapse. However, the experiment does not demonstrate — and no experiment in the paper tests — whether MHLA can surpass the FlashAttention baseline by a wide margin on tasks that genuinely challenge the base model's capacity. The hybrid variant (Wan-MHLA-H) edges ahead slightly (83.82 vs. 83.31 Total), suggesting MHLA can provide a small efficiency-adjusted improvement, but the ceiling remains anchored to the Wan2.1-1.3B model's fundamental capabilities.

The NLP experiments (Section 5.4, Table 6) provide complementary evidence. On perplexity — a metric sensitive to precise local token prediction — MHLA at 71.64 ppl (WikiLM) trails Transformer++ at 60.46 and Mamba2 at 58.51 by substantial margins. This gap suggests MHLA may be structurally worse at modeling the fine-grained, high-rank statistics that softmax attention and state-space models capture at this scale. The paper does not explore whether scaling model size closes this gap, whether it reflects a fundamental expressivity ceiling of the factorized block-mixing design, or whether the normalizer omission (mentioned in Section 4.1 but not specified for the NLP experiments) contributes to the degradation. A practitioner deploying MHLA for language modeling would need to assess whether the long-context advantages (LongBench average 7.41 vs. 6.97 for Mamba) justify the perplexity cost in their specific use case — the paper provides no guidance on this tradeoff.

The mitigation status is that the paper explicitly positions MHLA as a restoration of lost expressivity, not an augmentation beyond softmax attention's capabilities. The conclusion states the ambition is to "establish a fundamental attention mechanism," implying that further development could address the perplexity gap. But the mechanism itself — block-level mixing of precomputed key–value summaries — may impose inherent limits on the granularity of token-level interaction that constrain perplexity regardless of scale. The paper neither proves nor disproves this, leaving it as an unresolved question.

The Cost of Difficulty Estimation (Conceptual Analog): The Locality-Biased Initialization Is Not a Universal Prior

The paper designs a locality-biased initialization for the mixing coefficient matrix $\mathcal{M}_c$ that encodes the inductive bias that nearby spatial blocks are more relevant than distant ones (Section 4.2). This is a sensible choice for vision tasks where Euclidean distance correlates with semantic relevance — adjacent image patches are more likely to belong to the same object or texture than distant ones — and the ablation in Table 7(a) confirms it provides a substantial benefit: frozen locality-biased initialization achieves 75.4% on DeiT-T, outperforming uniform initialization with learning at 75.1%.

The limitation is that this initialization strategy implicitly assumes Euclidean spatial adjacency is a good proxy for semantic relevance. For 2D images and 3D video volumes, this assumption holds reasonably well, which explains MHLA's strong performance on vision tasks. For 1D token sequences in NLP — where the token ordering is sequential but semantic relationships are mediated by syntax and discourse structure, not Euclidean distance — the optimal initialization is far less obvious. The paper does not describe what initialization was used for the NLP experiments (Section 5.4), nor does it provide an ablation comparing different NLP-appropriate initialization strategies (e.g., distance based on token position with varying window sizes, or content-based initialization using pretrained embeddings). The LongBench results (Table 8) show MHLA performing well on long-context tasks, but it is unclear whether the initialization adapts to the 1D sequential structure naturally, whether it requires a different prior, or whether the learning process overrides the initialization quickly enough that the initial bias is irrelevant at 10B training tokens.

The consequence for practitioners is that applying MHLA to domains with non-Euclidean token geometries — graphs, point clouds, irregular meshes, tabular data, multi-modal sequences with heterogeneous token types — requires designing a new distance-based or content-based initialization strategy, with no guidance from the paper on what properties matter. The paper's ablation (Table 7(a)) shows initialization matters (frozen locality-bias outperforms uniform learning), suggesting that a poor choice of initialization in a new domain could significantly degrade performance or slow convergence. The paper does not ablate the learning rate or optimization dynamics of $\mathcal{M}_c$ — it is possible that the coefficient matrix learns different patterns at different stages of training, and that the initialization's influence decays at a rate that depends on dataset size, model scale, and task type.

The mitigation status is that the paper provides a general formula for locality-biased initialization ( $m_{i,j}^{(0)} \propto 1 - \text{dist}(i,j) / \max_k \text{dist}(i,k)$ ) and notes that it is "free to adapt during training" (Section 4.2). The implicit assumption is that learning will correct any misalignment between the spatial prior and the task structure. But the DeiT-T ablation shows that pure learning from a uniform start achieves worse performance than the frozen spatial prior (75.1% vs. 75.4%), which complicates this assumption — at least at the 300-epoch training budget typical of ImageNet experiments, optimization alone cannot fully compensate for a poor initialization. Whether longer training (the NLP model sees 10B tokens) or larger models overcome this initialization dependence is unmeasured.

The M² ≤ N Constraint and Unanswered Scaling Questions for Very Long Sequences

MHLA's linear complexity guarantee depends on the condition $M^2 \leq N$ , where $M$ is the number of token blocks and $N$ is the sequence length. Under this constraint, the $O(M^2 d^2)$ mixing cost is subsumed by the $O(N d^2)$ token-level operations, maintaining linear scaling. The paper validates this empirically in Appendix F.4 (Table 14): at $M=16, N=1024$ , MHLA throughput (51ms) is nearly identical to standard linear attention (52ms); at $M=64, N=256$ (where $M^2=4096 > 256=N$ ), throughput drops ~17% relative to standard linear attention.

The practical question this raises is: how should $M$ scale with $N$ for optimal performance, and does the M² ≤ N constraint become binding at extreme sequence lengths? The paper's ablation in Table 7(b) shows that on DiT-S/2 at N=1024, $M=16$ achieves optimal FID (78.63) with no throughput penalty, while $M=64$ degrades FID slightly (79.50) and reduces throughput. This suggests the optimal $M$ is considerably smaller than $\sqrt{N}$ for this task — the constraint is not binding in practice. However, the video generation experiment (Section 5.3) uses $M=105$ for $N=31,500$ , where $\sqrt{N} \approx 177.5$ , so $M=105$ satisfies $M^2=11,025 \ll 31,500$ . The constraint is comfortably satisfied. But the implication is that at $N=31,500$ , the model uses only 105 mixing blocks to cover 31,500 tokens — each block averages ~300 tokens. What happens if performance demands more blocks? If the optimal $M$ grows proportionally to $N$ (e.g., keeping a fixed tokens-per-block ratio), then $M^2$ eventually dominates $N$ and MHLA loses its linear complexity guarantee. If $M$ grows only as $\sqrt{N}$ , then tokens-per-block grows as $\sqrt{N}$ , making intra-block attention coarser. The paper does not study this tradeoff — it is unclear whether there exists a sequence length where the optimal $M$ violates $M^2 \leq N$ , forcing a choice between expressivity and complexity.

The video generation experiment is illuminating but does not resolve this question. At N=31,500 with M=105, the model performs well (VBench 82.62 vs. Wan-FA 83.31), but we do not know whether M=200 (with M²=40,000 > 31,500) would improve quality at the cost of exceeding linear complexity, or whether M=50 (coarser blocks, cheaper mixing) would suffice. The paper's scaling analysis (Appendix F.4) profiles throughput at fixed N and varying M, but does not profile FID or accuracy — it is a computational scaling study, not a quality-vs-M study. A practitioner deploying MHLA on million-token sequences (e.g., hour-long video, full-codebase modeling) would need to choose $M$ without knowing whether the optimal value respects the linear-complexity constraint.

The mitigation status is that the paper acknowledges the constraint implicitly through the complexity analysis in Section 4.3, but does not frame it as a potential limitation or provide guidance on choosing $M$ as a function of $N$ for new tasks. The fact that $M=16$ works well for N values from 256 (classification) to 4096 (generation at 512px) suggests the optimal $M$ may be relatively insensitive to $N$ within typical ranges — a practitioner-favorable finding — but this is not tested systematically.

The Perplexity Gap: NLP Evidence That MHLA Is Not a Universal Replacement

The NLP experiments in Section 5.4 are positioned as demonstrating MHLA's applicability across domains, and the LongBench results (Table 8) genuinely support the claim of strong long-context retrieval. However, the perplexity results in Table 6 reveal a consistent gap that the paper does not directly address: MHLA achieves 71.64 perplexity on WikiLM compared to 60.46 for Transformer++ and 58.51 for Mamba2. On LMB (Lambada), MHLA reaches 38.31 vs. Transformer++'s 34.57 and Mamba2's 35.40. These are not small differences — they suggest MHLA has a structural disadvantage at the fine-grained token prediction that perplexity measures.

The paper's narrative emphasizes the positive results (MMLU 23.7, best among all baselines; LongBench average 7.41, best among recurrent models) without analyzing the perplexity gap. This is a meaningful omission because perplexity is the standard training objective for language models, and gaps in perplexity at the 340M scale often persist or widen at larger scales. The mechanism-level question is whether the two-stage factorization (block-level mixing + intra-block kernel similarity) inherently cannot capture the high-rank, fine-grained token interactions that softmax attention models — and that state-space models like Mamba2 capture through their continuous-time formulation — or whether the gap is specific to the 340M scale and 10B-token training budget.

The NLP experiments also lack a controlled comparison to vanilla linear attention, making it difficult to isolate MHLA's contribution in the language domain. GLA (which includes gating mechanisms) serves as the closest linear-attention baseline, and MHLA outperforms GLA on MMLU (23.7 vs. 22.9) and LongBench (7.41 vs. 6.53). But without a vanilla LA baseline at the same scale, we cannot determine whether the bulk of the improvement over GLA comes from restoring query-conditioned selectivity (MHLA's claimed mechanism) or from other architectural differences (presence/absence of gating, normalizer behavior, feature map choice). The video generation experiment (Table 5), where vanilla LA is included and catastrophically fails, provides the cleanest evidence for MHLA's claimed mechanism. The NLP experiments provide weaker evidence because the baselines are stronger and the controlled comparison is missing.

A practitioner choosing between MHLA and alternatives for language modeling faces an underspecified tradeoff: MHLA offers strong long-context retrieval (LongBench) and MMLU performance, but at the cost of higher perplexity compared to both softmax attention and state-space models, with no analysis of whether the perplexity gap closes at larger scales. The paper's claim that MHLA "consistently outperforms existing linear attention baselines with negligible computational overhead" is supported for vision and video, but for NLP, the comparison is against GLA (a gated linear attention with its own expressivity mechanisms, not vanilla linear attention), and the overhead relative to Transformer++ includes a perplexity penalty whose practical significance is unquantified.

The Normalizer Omission: A Design Choice with Unspecified Consequences Across Experiments

Section 4.1 contains a brief but potentially consequential statement: "In tasks like language modeling and video generation, the normalizer term can be omitted for better training stability when the sequence is getting longer." The normalizer is the denominator $\widetilde{q}^\top \widetilde{z}_i$ in the attention output formula, which ensures the effective attention weights sum to one — a proper convex combination of value vectors. Omitting it means the output is the unnormalized numerator $\widetilde{q}^\top \widetilde{S}_i$ , which can have arbitrary magnitude and removes the probabilistic interpretation of the attention operation.

The paper does not specify:

Whether the normalizer was omitted for the NLP experiments (Section 5.4), the video generation experiments (Section 5.3), or both.
What the quantitative effect of omission is on any reported metric — no normalizer ablation exists anywhere in the paper.
Whether omission affects different metrics differently (e.g., more impact on perplexity than on long-context retrieval).
The mechanism by which the normalizer becomes unstable at long sequences (numerical overflow? vanishing denominator? training dynamics?) and at what sequence length this transition occurs.

This is a substantive limitation because the presence or absence of normalization changes the mathematical form of the attention operator. With the normalizer, MHLA produces a normalized weighted average — each token's contribution is bounded and the output lies in the convex hull of the value vectors, which has known benefits for training stability and gradient flow. Without the normalizer, the output magnitude scales with the query–key similarities and the mixing coefficients, which could interact with layer normalization downstream in unpredictable ways. A practitioner attempting to replicate the NLP or video generation results would not know which variant to implement, and if they tried both, they would not know how much the performance difference stems from the normalizer choice versus other hyperparameters.

The text-to-image generation results (Table 4) and image classification results (Section 5.1) presumably use the normalizer (since the caveat is specific to "language modeling and video generation"), but this is not stated. The image generation experiments (Section 5.2, Table 3) are also ambiguous. This lack of specification weakens the paper's claim of providing a turnkey attention replacement — the normalizer behavior is a task-dependent hyperparameter whose sensitivity is entirely uncharacterized.

The mitigation status is that the paper mentions the omission as a pragmatic choice without examining its consequences. The stated rationale ("better training stability when the sequence is getting longer") is plausible — unnormalized linear attention is used in architectures like RWKV and Mamba where normalizer instability at long sequences is a known issue — but the paper provides no empirical evidence that instability actually occurs with MHLA's normalizer at the tested sequence lengths (2048 for NLP, 31,500 for video). A simple ablation comparing normalized and unnormalized MHLA on a long-sequence language modeling task would clarify whether the omission is necessary or merely precautionary, and what performance tradeoff it entails.

Single-Model-Family Validation in the Context of General-Purpose Claims

The paper evaluates MHLA across an admirably broad set of tasks — image classification, class-to-image generation, text-to-image generation, video generation, and NLP — but the vision architectures all share the same ViT/DiT lineage, and the NLP experiments use a single model scale (340M parameters) trained on a single data distribution (SlimPajama) with a single training recipe (10B tokens). The paper states in the conclusion that it envisions MHLA as "a fundamental attention mechanism that can benefit a wide range of downstream applications," but the evidence base has specific architectural and scale boundaries that are not fully acknowledged.

For vision, all experiments use variants of the DeiT, VLT, DiT, or SANA architectures, which share common design elements: patch-based tokenization, fixed-resolution input, class-token or average-pooling aggregation, and standard training recipes. Whether MHLA transfers to vision architectures with different tokenization strategies (e.g., convolution-based token mixing, deformable attention, hierarchical feature pyramids), to dense prediction tasks (segmentation, detection, depth estimation) where per-pixel rather than per-image accuracy matters, or to multi-scale architectures where the block partitioning would need to adapt across feature levels is unknown. The paper's spatial block partitioning assumes a regular 2D grid — architectures with irregular token geometries (graph neural networks, point cloud transformers) would require a redesign of how blocks are defined and how the distance-based initialization is computed.

For NLP, the 340M-parameter, 10B-token scale is far below the regime where linear attention is most practically relevant. The cost advantage of linear attention over softmax attention grows with sequence length, but the quality gap may also change with scale in ways the paper does not explore. State-space models like Mamba show that the softmax-to-linear quality gap narrows at larger scales and longer training — but MHLA's behavior in that regime is untested. The LongBench results are encouraging but come from a single run at a single scale, without the error bars or multi-seed analysis that would make the claim of superiority robust.

The paper does not evaluate MHLA on any task requiring structured output beyond classification scores or generative samples — there are no experiments on machine translation, summarization quality (beyond LongBench metrics), code execution accuracy, mathematical reasoning, or instruction following. These tasks stress attention mechanisms in different ways (cross-lingual alignment, long-range logical dependencies, precise token-level control), and it is plausible that MHLA's factorized two-stage weighting behaves differently under these demands than it does under the relatively homogeneous requirements of image generation and commonsense reasoning.

The mitigation status is that breadth across four domains is substantially more than most linear attention papers provide — the paper is genuinely more comprehensive in its empirical scope than the typical efficient-attention manuscript. But the framing as a "fundamental attention mechanism" carries a burden of evidence that the current experiments do not fully discharge. A single architecture family in vision, a single scale in NLP, and no structured-output tasks means the generality claim is aspirational — well-motivated by the mechanism's design but not yet empirically established. The paper would be stronger if it explicitly acknowledged these boundaries and suggested specific architectures or tasks where validation would be most informative.

7. Implications and Future Directions

How This Work Changes the Landscape

MHLA is best understood not as a new attention mechanism that competes with existing ones on standard benchmarks — though it does — but as a diagnostic contribution with a mechanism attached. What will persist beyond this specific paper is the reframing of the linear attention problem around two measurable quantities: the rank of the attention matrix and the entropy of the attention distribution. Prior work treated the performance gap between linear and softmax attention as diffuse and multifactorial — lost local structure, kernel approximation error, insufficient nonlinearity — and the field responded with a proliferation of auxiliary modules (convolutions, gates, hybrid layers) that addressed symptoms without converging on a unified understanding of the root cause. The paper's identification and quantitative characterization of global context collapse — the joint phenomenon of rank deficiency and elevated entropy arising from a shared global key–value summary — changes what it means to evaluate a new linear attention method. Instead of asking "does it beat baseline X on benchmark Y?", researchers can now ask "does it increase the rank of the attention matrix?" and "does it reduce attention entropy?" — questions with clear mathematical definitions and diagnostic value that outlast any particular benchmark.

The magnitude of this shift is diagnostic reframing, not paradigm shift. The paper does not overturn the fundamental approach to efficient attention — softmax attention remains the gold standard for tasks where quadratic complexity is tolerable, and state-space models like Mamba continue to show strong results on long sequences through a different computational primitive entirely. What changes is the design philosophy within the linear attention subfield: the paper demonstrates that a principled fix to the root cause (query-conditioned selectivity, restored through block-level mixing) can match or exceed the performance of methods that add auxiliary modules, and it provides evidence that those auxiliary modules become counterproductive at scale (Table 3(a), DiT-XL CPE degradation). This shifts the research agenda from "what module can I add to make linear attention better?" to "how can I design the summary mechanism itself to preserve per-query diversity?" — a more constrained and potentially more productive question.

The paper also resolves a latent tension in the literature between advocates of linear attention (who emphasize its efficiency) and critics (who point to its degraded performance). By showing that the performance degradation has a specific, measurable cause — the shared global summary — rather than being an inherent cost of linear complexity, the paper suggests that the earlier negative results on linear attention were not about linear complexity per se, but about a particular design choice (the single global summary) that is not necessary. This reframes the debate: the question is not whether linear attention works, but whether a given linear attention design avoids global context collapse. The video generation experiment (Section 5.3, Table 5) provides the most dramatic empirical resolution: vanilla linear attention catastrophically fails at 31,500 tokens (Semantic score dropping from 75.65 to 11.38), while MHLA at the same sequence length and the same latency nearly fully recovers. This demonstrates that the earlier negative results were measuring the consequences of the global summary design, not an intrinsic limitation of linear-complexity attention.

The research directions that become more attractive after this work include: principled design of linear attention mechanisms evaluated by rank and entropy metrics rather than solely by downstream task performance; systematic study of query-conditioned selectivity as a design axis orthogonal to kernel choice and feature map design; and exploration of factorized attention patterns that occupy the middle ground between full pairwise attention and fully shared summaries. The directions that become less attractive include: incremental addition of auxiliary modules (convolutions, gates, hybrid layers) without demonstrating that they address the root cause rather than papering over symptoms; and comparisons of linear attention methods that do not control for or report the attention matrix's rank and entropy, since the paper demonstrates these metrics are diagnostic of the underlying failure mode.

Follow-Up Research This Work Enables

Systematic rank-entropy benchmarking of all linear attention variants. The paper introduces rank and entropy as diagnostic tools for global context collapse, but applies them only to a subset of methods (self-attention, vanilla linear attention, Focused LA, Inline Attention, MALA, and MHLA) on a single DeiT-T model (Figure 3). A comprehensive study would measure the attention-matrix rank and entropy for every published linear attention variant — GLA, Mamba2, Gated DeltaNet, RWKV, RetNet, H3, Performer, Linformer — across multiple model scales, sequence lengths, and tasks (vision, language, video). This would reveal which methods genuinely escape the rank bottleneck and which achieve their performance gains through mechanisms orthogonal to query-conditioned selectivity (e.g., better gating, improved feature maps, training dynamics). The paper's diagnostic framework makes such a study newly tractable because it provides the measurement methodology and the theoretical motivation; previously, there was no unified lens through which to compare approaches as diverse as kernel-based linear attention and state-space models.

The normalizer ablation: when and why does the denominator matter? Section 4.1 states that the normalizer can be omitted "for better training stability when the sequence is getting longer," but no ablation quantifies the effect. A controlled experiment would train MHLA models at multiple sequence lengths (1K, 2K, 4K, 8K, 16K tokens) with and without the normalizer, measuring perplexity, downstream accuracy, and training stability (loss spikes, gradient norms). The key question is whether the normalizer degrades gracefully with sequence length (in which case it should be kept for its probabilistic interpretation) or whether there exists a sharp threshold beyond which it causes divergence (justifying omission for long-sequence tasks). This study would also test whether the normalizer interacts with the mixing coefficients — if the denominator becomes unstable because $\widetilde{q}^\top \widetilde{z}_i$ approaches zero for certain query-block–summary-block combinations, this would suggest a design refinement where the normalizer is omitted only for blocks with small denominator values rather than globally. The paper's video experiment (31,500 tokens) and NLP experiment (2,048 tokens) used potentially different normalizer settings without specification, so this ablation is also necessary for replicability.

Scaling behavior of the optimal block count M with sequence length N. The paper establishes that $M^2 \leq N$ is the condition for linear complexity and that $M=16$ is sufficient for strong performance on sequences up to 4,096 tokens (Table 7b, Appendix F.4). However, it does not characterize how the optimal $M$ scales when $N$ grows by orders of magnitude — to 100K tokens (long documents), 1M tokens (hour-long video), or beyond. A systematic study would train MHLA models at fixed model size across a range of $N$ values (e.g., 1K, 4K, 16K, 64K, 256K) while sweeping $M$ (keeping $M^2 \leq N$ , $M^2 \approx N$ , and $M^2 > N$ ) and measuring both quality and throughput. The central question is whether the optimal $M$ grows as $\sqrt{N}$ (maintaining fixed tokens-per-block), stays constant (increasing tokens-per-block and relying on intra-block kernel similarity for fine-grained attention), or follows some intermediate scaling law. If the optimal $M$ grows faster than $\sqrt{N}$ , MHLA eventually loses its linear complexity guarantee, which would define a practical upper bound on the sequence lengths where it is applicable. If the optimal $M$ is insensitive to $N$ , this would be strong evidence that the block structure serves primarily to break the rank bottleneck at modest $M$ , and that infinite sequence length is achievable without increasing $M$ — a highly attractive property for deployment.

Interaction between MHLA's mixing coefficients and the feature map choice. The paper uses a generic kernelized formulation ( $\phi(Q), \phi(K)$ ) without specifying the feature map or exploring how it interacts with the mixing mechanism. Different feature maps (ReLU, ELU+1, Performer's random Fourier features, cos-based features) produce different similarity distributions and may benefit from different mixing-coefficient patterns. A study would train MHLA models with a fixed architecture and vary the feature map, measuring both downstream performance and the learned mixing-coefficient matrix $\mathcal{M}_c$ to see whether the optimal spatial mixing pattern depends on the kernel. For example, a feature map that produces sharper similarities (more selective intra-block attention) might benefit from broader inter-block mixing (higher off-diagonal coefficients), while a feature map that produces flatter similarities might need more localized mixing. This would inform whether the feature map and mixing matrix are independent design choices or whether they should be co-designed.

MHLA on structured-output and cross-attention tasks. The paper evaluates MHLA on classification, generation, and language modeling, all of which use self-attention (tokens attending to tokens from the same sequence). Many practical applications use cross-attention (encoder–decoder models for translation, image–text models for captioning, retrieval-augmented generation where queries attend to retrieved documents). Cross-attention is structurally different: the query sequence and key–value sequence have different lengths and potentially different token geometries, making the spatial-block initialization formula ( $\text{dist}(i, j)$ based on Euclidean distance in a shared grid) inapplicable. A study would adapt MHLA to cross-attention by designing content-based or learned initialization strategies for the mixing matrix, and evaluate on tasks like machine translation (where the query–key alignment is monotonic but with local reordering) and image captioning (where the query–key relationship is between text tokens and image regions with no inherent ordering). This would stress-test the paper's claim that MHLA is a "fundamental attention mechanism" by moving it beyond self-attention on regular grids.

Negative result: does MHLA fail when the block structure misaligns with task structure? The paper's block partitioning assumes spatial contiguity on a 2D or 3D grid, which works well for images and videos where adjacency correlates with relevance. A deliberate stress test would construct a task where relevant tokens are scattered non-locally in the token ordering (e.g., randomly permuted image patches, or a language task where the tokens relevant to answering a question are dispersed throughout a long document with no positional pattern). If MHLA's performance degrades sharply relative to softmax attention on such tasks, this would reveal that the block-level mixing mechanism relies on spatial locality as a prior, not just as a useful initialization — i.e., that the block structure imposes a capacity limit on modeling long-range, non-local dependencies that the learned coefficients cannot fully overcome. This would refine the paper's claim by clarifying that MHLA restores query-conditioned selectivity subject to a block-level locality bias, and that tasks requiring highly non-local attention patterns may need a different partitioning strategy or a mechanism operating at multiple spatial scales.

Practical Applications and Downstream Use Cases

High-resolution image and video generation at scale. This is the paper's most immediately actionable result. Table 3 shows that MHLA achieves FID scores matching or exceeding self-attention on DiT-XL/2 while providing 2× throughput at 512px resolution. Table 5 shows that for video generation at 31,500-token sequences, MHLA achieves a 2.1× inference speedup (81s vs. 166s) over FlashAttention while recovering nearly all of the quality lost by vanilla linear attention. For production video generation systems — where generating a single high-resolution clip currently takes minutes on expensive hardware — this speedup translates directly to cost reduction and latency improvement without meaningful quality degradation. The hybrid variant (replacing only 2/3 of attention layers with MHLA, keeping the rest as FlashAttention) provides an even more attractive deployment profile: 1.6× speedup with quality exceeding the original FlashAttention model (VBench Total 83.82 vs. 83.31). A team deploying Wan, Sora, or similar video diffusion models could adopt the hybrid strategy immediately with minimal engineering effort (drop-in layer replacement) and measurable throughput gains.

Long-context retrieval and multi-document QA systems. The LongBench results in Table 8 show MHLA achieving the highest average score (7.41) among recurrent models, with particular strength in Multi-Doc QA (3.58 vs. 3.37 for Mamba), Summarization (18.59 vs. 18.36 for Mamba), and Code (12.72 vs. 12.55 for GLA). For systems that need to process long contexts — legal document review, scientific literature synthesis, codebase understanding — MHLA provides a mechanism that can selectively attend across long sequences without the quadratic cost of full attention and without the documented degradation of vanilla linear attention. The throughput advantage over softmax attention grows with context length, making MHLA particularly suitable for retrieval-augmented generation pipelines where the retriever provides dozens of documents and the model must synthesize across them. The caveat is the perplexity gap on standard language modeling (71.64 vs. 60.46 for Transformer++ on WikiLM), which suggests practitioners should use MHLA when long-context retrieval is the primary requirement, and consider hybrid architectures (some softmax layers for local fluency, MHLA layers for cross-document attention) when both fluency and retrieval matter.

Fine-tuning pretrained models for long-sequence tasks without full retraining. The SANA fine-tuning experiment (Section 5.2, Table 4, Figure 5) demonstrates that MHLA can be retrofitted into a pretrained linear-attention model with only 40K fine-tuning steps, rapidly matching and then surpassing the original checkpoint's loss. The DiT-XL/2 fast adaptation experiment (Table 3b) shows MHLA fine-tuned from a pretrained self-attention checkpoint matches the original FID after 400K steps. For teams with existing pretrained models facing deployment on longer sequences than originally anticipated (e.g., moving from 256px to 1024px image generation, or from 2K to 32K context windows), MHLA offers a pathway to extend the model's effective sequence length without retraining from scratch and without the catastrophic degradation of naive linear attention. The engineering cost is low — MHLA is a drop-in replacement for self-attention or linear attention layers — and the fine-tuning cost is modest relative to pretraining.

When to Prefer This Method

The paper provides enough comparative data against named alternatives (softmax self-attention, vanilla linear attention, GLA, Mamba, Mamba2, Gated DeltaNet) to support a conditional decision rule grounded in its empirical results, though the rule necessarily reflects the limitations of the evaluated scales and domains.

Prefer MHLA over vanilla linear attention when: the sequence length exceeds a few thousand tokens AND the task requires non-uniform, query-dependent attention patterns. The video generation experiment (Table 5) is the decisive evidence: at 31,500 tokens, vanilla LA's Semantic score collapses to 11.38 vs. 75.65 for FlashAttention, while MHLA recovers to 76.16 with the same 2.1× speedup. At shorter sequences (e.g., 256-token image classification, Table 2a), the gap is smaller (6.0 percentage points on DeiT-T) but still substantial. There is no scenario in the paper's results where vanilla LA matches MHLA, so the preference is unambiguous whenever the engineering cost of implementing MHLA over vanilla LA is acceptable.

Prefer MHLA over softmax self-attention when: throughput at long sequences is the binding constraint AND a small quality tradeoff is acceptable OR the task is within the regime where MHLA matches or exceeds self-attention. On image classification, MHLA (75.8%) exceeds self-attention (72.2%) on DeiT-T but the margin shrinks on DeiT-S (+1.2%) and is untested on larger classification models. On image generation, MHLA with CPE+gating narrowly beats self-attention on DiT-XL/2 (FID 19.17 vs. 19.47) while plain MHLA without auxiliary modules trails slightly (20.32 vs. 19.47). On video, the hybrid MHLA-FlashAttention model exceeds the pure FlashAttention baseline (83.82 vs. 83.31 Total). On NLP, MHLA leads on long-context understanding (LongBench 7.41) but trails on perplexity (71.64 vs. 60.46 for Transformer++). The decision rule is therefore task-dependent: for generation and classification, MHLA is favorable at or above the DiT-XL/DeiT-S scale; for language, MHLA is favorable when long-context retrieval is the primary metric and perplexity is secondary.

Prefer state-space models (Mamba, Mamba2) over MHLA when: language modeling perplexity is the primary evaluation metric AND the sequence length does not exceed the regime where Mamba's bidirectional performance degradation matters. Table 6 shows Mamba2 achieving WikiLM perplexity of 58.51 vs. MHLA's 71.64 — a gap that may or may not close at larger scales but is substantial at 340M parameters. Mamba also slightly edges MHLA on commonsense reasoning average (47.0 vs. 47.1 — essentially tied). However, the paper notes (Appendix A) that Mamba-based models "exhibit substantial performance degradation" on bidirectional tasks, which MHLA handles natively. The tradeoff is therefore between perplexity/efficiency (Mamba) and bidirectional compatibility/LongBench performance (MHLA), with the decision depending on whether the deployment requires unidirectional or bidirectional attention.

The paper does not provide enough data to make strong recommendations against sparse attention mechanisms (Longformer, BigBird, sliding window), against Performer-style random-feature approaches, or against the full range of hybrid architectures — these comparisons would need to be run by practitioners evaluating specific use cases at their target sequence lengths and model scales.