PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Kewei Zhang, Ye Huang, Yufan Deng, Jincheng Yu, Junsong Chen, Huan Ling, Enze Xie, Daquan Zhou

arXiv:2601.07832v2cs.CVcs.AI

TL;DR

Transformers are difficult to scale because self-attention has quadratic time and memory costs, while linear attention can lose query-conditioned expressiveness. MHLA groups tokens and mixes key–value summaries according to query blocks, preserving token-wise diversity with linear complexity and no additional modules. Across classification, generation, NLP, and video tasks, it reports improvements ranging from 3.6% to 41% under the stated comparisons.

  • Problem

    Quadratic self-attention limits long-sequence scalability, while linear attention’s shared global key–value summary can cause rank-deficient, low-diversity representations.

  • Method

    MHLA partitions tokens into groups and uses query-conditioned mixing of key–value summaries to restore token-level diversity while retaining O(N) complexity without additional modules.

  • Results

    MHLA achieves state-of-the-art performance across discriminative and generative tasks, including 3.6% higher ImageNet accuracy, 6.3% NLP improvement, 12.6% image-generation improvement, and 41% video-generation improvement over vanilla linear attention.

  • Takeaways & Limitations

    MHLA preserves token-wise diversity and achieves performance comparable to or surpassing self-attention without depthwise convolutions or hybrid self-attention layers.

  • Takeaways & Limitations

    The single-global-summary formulation is strictly rank-bounded by d, and the stated upper-bound analysis assumes full row rank for each block product and suitable row-space conditions.

Abstract

from arXiv · show

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

1 Introduction

Transformers face scalability limits from quadratic self-attention, while linear attention improves efficiency but can lose query-specific expressiveness. MHLA restores token-level diversity through token-dimension grouping and query-conditioned mixing while retaining linear complexity and negligible overhead.

  • Quadratic self-attention time and memory complexity limits scalability for long-sequence tasks such as high-resolution image and video generation.
  • Linear attention reduces complexity to linear by compressing keys and values into a global summary, but this removes query-specific adaptation and often degrades accuracy.
  • Global key–value sharing limits representational capacity, producing rank deficiency, more uniform attention weights, and reduced diversity as sequences lengthen.
  • MHLA partitions tokens into non-overlapping heads, computes local key–value summaries, and mixes them with query-conditioned weights to retrieve tailored contexts.Within selected blocks, a query-dependent reweighting module further refines token contributions.
  • MHLA uses standard GEMMs, maintains O(N) complexity, supports streaming/stateful execution, and adds negligible computational overhead.
  • 3.6% accuracy gain over self-attention on ImageNet, 12.6% improvement for image generation, 6.3% improvement on NLP, and 41% improvement over vanilla linear attention for video generation.

2 Related Works

Related work replaces softmax attention with kernel-based linear mechanisms to improve efficiency, but these methods often lose expressiveness and rely on added components to mitigate performance drops.

  • Linear attention replaces softmax with kernel-based methods to achieve linear time complexity, but often suffers performance drops when capturing complex token interactions.

3 Analysis of Linear Attention

Linear attention achieves linear-time scaling by reordering feature-map computations around a shared global key–value summary, but this fixed-size summary limits rank and query-specific token weighting as sequences grow. The resulting global context collapse appears as reduced diversity, higher entropy, and weaker concentration than softmax attention.

  • Linear Attention Formulation: Softmax attention computes and normalizes all pairwise query–key similarities, resulting in O(N^2) complexity.Its per-query normalization enables query-specific attention distributions.
  • Linear Attention Formulation: Linear attention replaces the exponential softmax kernel with a positive feature map and precomputes global summaries, reducing complexity from O(N^2) to O(Ndϕ).The summaries are reused across queries to obtain linear-time scaling with sequence length.
  • Global Context Collapse: The fixed-size global key–value summary creates an information bottleneck, causing performance saturation as sequence length increases; this is termed global context collapse.The information to represent can exceed the capacity of the d × d summary matrix.
  • Global Context Collapse: The rank of linear-attention matrices is bounded by d regardless of sequence length, producing rank-deficient approximations when n ≫ d.Experiments show the rank is capped by the head dimension, typically dh ≤72, while relative expressivity declines as sequence length grows.
  • Global Context Collapse: Shared global summaries prevent query-specific reweighting of individual keys, making attention distributions less sparse and increasingly uniform as sequence length grows.Higher entropy indicates weaker concentration; linear attention exhibits significantly higher entropy than softmax-based attention.
  • Global Context Collapse: Figure 3 visualizes the resulting attention-map differences and reports that MHLA produces richer, more focused attention than the baselines.The comparison uses attention-map visualization together with average rank and entropy for DeiT-T.

4 Multi-Head Linear Attention

MHLA splits tokens into multiple heads, forms local key–value summaries, and mixes them according to query-block-specific coefficients to restore token-level diversity while retaining linear complexity. Its blockwise structure increases representational capacity and yields higher-rank, sparser attention than global linear attention.

  • Architecture: MHLA partitions the sequence into M non-overlapping token blocks and computes local key–value summaries for each block.Vision models define blocks on spatial or spatiotemporal grids rather than only flattened sequences.
  • Multi-Head Mixing: Each query block forms a distinct mixture of local summaries using a learned coefficient row, producing query-specific context instead of one shared global summary.The coefficient matrix Mc stores affinities between query blocks and local summaries, while each row specifies the mixture for one query block.
  • Token-Level Adaptivity: MHLA restores token-level weighting in two stages: block selection through learned coefficients followed by intra-block reweighting through kernel query–key similarities.This combines block-dependent scaling with token differentiation inside each selected block.
  • Efficiency: MHLA retains linear complexity because block mixtures are computed once per block and reused across all queries in that block.With M^2 ≤ N, the O(Nd^2 + M^2d^2) cost is dominated by O(Nd^2).
  • Rank Analysis: Under mild generic conditions, MHLA’s attention rank can reach min(N, Σ_b r_b), whereas global linear attention remains limited by d.Even when block row spaces are not fully independent, the rank grows roughly additively with the number of blocks M.
  • Sparsity Analysis: MHLA yields lower attention entropy than linear-attention baselines and softmax attention, indicating more concentrated, query-conditioned token selection.The coefficient matrix emphasizes relevant blocks, while kernel inner products further sharpen token contributions within those blocks.

5 Experiments

Experiments evaluate MHLA across classification, image and video generation, NLP, long-context understanding, and ablations. Across these settings, MHLA generally improves performance while retaining low computational overhead and linear-attention efficiency.

  • Image Classification: MHLA achieves the best accuracy across all evaluated DeiT model sizes and state-of-the-art performance in VLT with consistent baseline improvements.It introduces the fewest extra parameters among the compared linear-attention baselines.
  • Image Generation: MHLA maintains throughput nearly identical to linear attention while matching or surpassing self-attention performance in class-to-image generation.Across model sizes, it achieves the best reported performance without extra modules at L and XL scales.
  • Image Generation: MHLA rapidly adapts during SANA fine-tuning, matching the pretrained checkpoint within 2k steps and later converging to a lower loss.Replacing linear attention with MHLA improves multiple evaluation metrics over the baseline SANA model and PixArt series.
  • Video Generation: At sequence length 31,500, MHLA preserves linear-time complexity, matches the original FlashAttention-based model comparably, and achieves a 2.1× inference speedup.The hybrid model achieves a 1.6× speedup with better overall performance, while vanilla linear attention suffers severe degradation.
  • Video Generation: MHLA rapidly adapts during ultra-long video fine-tuning, whereas vanilla linear attention fails to train effectively and plateaus at a high loss.The loss behavior supports the paper’s analysis of global context collapse in ultra-long visual sequences.
  • Natural Language Processing: In NLP, MHLA performs comparably to Transformer++ and leading linear models, surpasses all baselines on the aggregated benchmark, and attains the highest average LongBench score.Its advantages on LongBench are especially pronounced for Multi-Doc QA, summarization, and code tasks.
  • Ablation Study: Locality-biased initialization provides a strong prior, while learnable mixing coefficients adapt further to the dataset distribution and improve performance.MHLA reaches excellent FID at M=16 while maintaining the highest throughput, indicating low overhead with a relatively small head count.

6 Conclusion

The paper introduces Multi-Head Linear Attention as a linear attention mechanism that preserves token-wise diversity without auxiliary depthwise convolutions or hybrid self-attention layers. It reports performance comparable to or surpassing self-attention-based models across downstream applications.

  • MHLA partitions tokens into multiple groups to preserve token-wise diversity in linear attention.
  • MHLA achieves performance comparable to or surpassing self-attention-based models without additional depthwise convolutions or hybrid self-attention layers.

Multi-Head

Prior work develops efficient alternatives to quadratic self-attention, but linear and sparse approaches make trade-offs in expressiveness, context modeling, or scalability. Additional convolutions, gating, and state-space components address some limitations while adding complexity or leaving fundamental constraints.

  • Linear Attention: Linear attention replaces quadratic softmax attention with kernel-based feature mappings, enabling linear-time training and inference for long sequences.
  • Linear Attention: Linear attention often reduces representational expressiveness, motivating added convolutional and gating modules to enrich local context and information control.
  • Linear Attention: State space models such as Mamba and its variants provide efficient alternatives with strong scalability on long sequences and competitive accuracy.
  • Sparse Attention: Sparse attention reduces computation by restricting each token to a subset of other tokens through local or global structural patterns.
  • Applications: Linear and sparse attention mechanisms have been applied across NLP, computer vision, and generative modeling to improve efficiency and handle longer sequences.

B Query-Conditioned Selectivity in Softmax Attention

Softmax attention provides query-conditioned and per-token weighting, whereas global linear attention shares one key–value summary across queries and loses selectivity. MHLA restores query-adaptive context through mixtures of local summaries while retaining chunkwise linear-attention complexity.

  • Query-Conditioned Selectivity: Softmax attention assigns each query its own token-weight distribution, allowing sharply concentrated and adaptive context vectors.
  • Global Linear Attention: Global linear attention compresses all keys and values into one summary shared by every query, making different queries receive nearly identical contexts.
  • MHLA: MHLA uses query-block-specific mixtures of local summaries so different querying blocks assign different effective weights to the same token.
  • MHLA: The mechanism combines query-conditioned block selection with within-block token reweighting through the kernel inner product.
  • Causal MHLA: In causal chunkwise training, MHLA forms block-specific mixed summaries and reuses each block’s mixture, preserving causality while matching chunkwise linear-attention asymptotic complexity.

D Dataset

The paper evaluates MHLA across image classification, image generation, and natural language processing, using standard benchmarks and a fine-tuning setup for text-to-image generation.

  • Experiments cover image classification, C2I generation, T2I generation, and natural language processing.Classification and C2I use ImageNet-1K, while T2I uses a pretrained model fine-tuned on 31,292k internet images.
  • Classification and C2I models are trained on ImageNet-1K and evaluated on its standard validation set.
  • T2I evaluation fine-tunes a pretrained model using a relatively small collection of 31,292k internet images.

E Extra Implementation Details

Additional implementation details describe fair classification comparisons, expanded generation reporting, and stability checks across independent runs.

  • Classification baselines use identical training settings, with average pooling replacing DeiT’s class token.The setup also adds CPE with a kernel size of 3 for fair comparison.
  • 300 epochs, batch size 1024, and peak learning rate 1e-3 define the reported classification training setup.
  • Generation results are expanded through complete DiT and DiG results, additional SANA-MHLA samples, and comparisons with recent linear-attention methods.
  • MHLA generation results report mean and standard deviation over three independent runs to demonstrate stability.The corresponding comparisons are summarized in Table 11, including FID scores for MHLA against LiT.

F.2 Ablation of CPE and output gating.

The ablations examine auxiliary modules, resolution, computational overhead, and MHLA terminology, showing that MHLA’s expressivity gains remain consistent across model sizes while overhead depends on token-level head count.

  • F.2 Ablation of CPE and output gating.: MHLA consistently improves expressivity across model sizes, whereas CPE and output-gating gains diminish as models become larger.In DiT-XL, CPE alone decreases performance, while MHLA remains beneficial.
  • High-resolution classification experiments evaluate DeiT-T at 384×384 and 512×512 resolutions.
  • MHLA adds negligible overhead when M^2 < N, but larger token-level head counts produce more noticeable overhead.Ablations indicate that M^2 < N is sufficient for strong performance.
  • The study reports DiT and DiG comparisons, LiT FID stability measurements, and profiling across sequence lengths and token-level head numbers.
  • Query-conditioned aggregation assigns each query token mixing coefficients to combine local KV summaries independently at every query position.This makes adaptation per-query rather than global or governed by a shared recursive rule.
  • MHLA computes global KV summaries independently through many-to-one aggregation instead of the strict recurrent hidden-state chain used in traditional linear attention.The paper argues that avoiding rigid history inheritance increases KV-summary expressivity and flexibility.
Loading 2601.07832v2…