Multi-Head Linear Attention (MHLA)

• Multi-Head Linear Attention (MHLA) is an attention mechanism that combines linear efficiency with multi-head diversity to preserve full context expressivity.
• It partitions tokens and employs specialized mixing strategies or expert routing to overcome global context collapse without quadratic computational costs.
• MHLA demonstrates improved accuracy across vision and language tasks while maintaining linear complexity, making it scalable for large-scale deployments.

Multi-Head Linear Attention (MHLA) is a class of attention mechanisms that combine the scalability advantages of linear attention with the representational diversity of multi-head designs, seeking to overcome global context collapse while enabling efficient deployment in large-scale deep learning models across vision, language, and multimodal domains. MHLA generalizes linear attention by allocating multiple independent or interacting “heads” to partition, summarize, or mix token-level or expert-level information, thus restoring much of the expressivity and functional modularity of standard multi-head softmax attention at strictly linear time and memory complexity.

1. Motivation and Conceptual Foundations

Traditional softmax-based multi-head self-attention (MHSA) incurs $O(N^2)$ complexity in sequence length $N$, limiting its practicality for high-resolution vision, long-context language modeling, and real-time on-device inference. Linear attention variants, in which the softmax kernel is replaced by a positive feature map $\phi(\cdot)$, reduce computational burden to $O(N)$ by collapsing global key–value information into a fixed-rank ($d_\phi$) summary. However, this approach introduces a failure mode—global context collapse—leading to loss of representational diversity, reduced actionable sparsity, and degraded downstream accuracy. To remedy this, Multi-Head Linear Attention partitions the attention operation into parallel heads—defined spatially, channel-wise, or via specialized routing—enabling restoration of higher-rank structure and local/global context diversity without incurring quadratic costs (Zhang et al., 12 Jan 2026, Tuli et al., 30 Oct 2025, Kang et al., 2024, Setyawan et al., 12 Jun 2025).

2. Formalizations and Algorithmic Structures

MHLA admits multiple constructions, each providing a mechanism for multi-head diversity in a linear complexity regime:

Blockwise Token Partitioning and Query-Conditioned Mixing

In “MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head” (Zhang et al., 12 Jan 2026), the input sequence is divided into $M$ disjoint blocks (“heads”) along the token axis. For each block:

• Keys/values are summarized separately: $S_b = \sum_{j\in b}\widetilde K_j V_j^\top$, $z_b = \sum_{j\in b}\widetilde K_j$.
• For a query in block $i$, a learnable coefficient matrix $\mathcal{M}\in\mathbb{R}^{M\times M}$ yields a convex combination: $\widetilde S_i = \sum_{b=1}^M m_{i,b} S_b$, $\widetilde z_i = \sum_{b=1}^M m_{i,b} z_b$.
• Final output for token $t$ in block $i$ is $$\mathrm{MHLA}(Q,K,V)_t = \frac{\widetilde q_t^\top \widetilde S_i}{\widetilde q_t^\top \widetilde z_i}$$.

This two-stage weighting scheme restores token- or block-level specificity absent in global linear attention, with complexity $O(N d_\phi d) + O(M^2 d_\phi d)$ ($M^2\ll N$ in practice) (Zhang et al., 12 Jan 2026).

Mixture-of-Experts SSM and Recurrence-Based Heads

MossNet (Tuli et al., 30 Oct 2025) realizes MHLA via a mixture-of-experts (MoE) over state-space model (SSM) kernels:

• Input tokens $u_t\in\mathbb{R}^d$ pass through top-$k$ MoE projections for input, gating, and output channels, with expert probabilities $p_i(u_t)$ generated by a router network.
• For each expert $i$, state-space parameters $(\bar B^i, C^i, \Delta^i)$ are fused per token via $p_i(u_t)$-weighted sums.
• The ensemble output $y_t$ is a mixture: $y_t = \sum_{i=1}^H p_i(u_t) y_t^{(i)}$.
• Theoretical analysis (Theorem 1) shows a recoverable equivalence to linear MHA:

$$y_t = \sum_{m,n=1}^H \sum_{s=1}^t \langle q_t^m, k_s^n\rangle v_s$$

Each SSM expert thus emulates a separate attention head, with only $k\ll H$ heads selected per token.

Decomposed Cross-Head Interaction

Interactive Multi-Head Self-Attention (iMHSA) (Kang et al., 2024) decomposes the attention matrix for each head via spatial “landmark” pooling, followed by cross-head fully connected interactions. This insertion of $H\times H$ FC layers through similarity channels enables communication among head features in linear $O(NLH+NH^2)$ time, maintaining the benefits of multi-head diversity.

Channel-Wise Token Mixing

In FaceLiVT (Setyawan et al., 12 Jun 2025), MHLA is instantiated as grouped linear token mixers across channel-split heads, without explicit query/key/value formation. Each head applies two linear token-mixing projections with nonlinearity, approximating multi-head token communication in $O(BCN^2)$ for $B$ batch size and $C$ channels, while eschewing softmax operations.

3. Theoretical Properties and Expressivity

MHLA’s fundamental distinction lies in its capacity to restore higher effective attention matrix rank and localized sparsity absent from global linear attention:

• The effective rank of standard linear attention is bounded by $d_\phi$, while MHLA with $M$ heads can scale its rank as $\min(N, d_\phi M)$ (Zhang et al., 12 Jan 2026).
• The query-conditioned mixing matrix $\mathcal{M}$ introduces block-sparsity and diversity, enabling selective information routing analogous to softmax attention’s context preservation.
• MoE-SSM MHLA (MossNet) enables both increased rank and diversified query/key projections for enhanced representational power at linear cost (Tuli et al., 30 Oct 2025).
• Cross-head interaction in iMHSA empirically leads to higher inter-head feature variance and reduced collapse, confirmed via ablation studies (Kang et al., 2024).

4. Computational Complexity and Scaling

A core advantage of MHLA is linear scaling with sequence or spatial resolution:

• Standard MHSA: $O(N^2d)$ flops, $O(N^2)$ memory.
• MHLA approaches: $O(Nd_\phi d)$ (blockwise, kernel-based), $O(kLd+Hd^2)$ (MoE-SSM), $O(NLH+NH^2)$ (iMHSA).
• Head count $M$, landmark count $L$, and number of experts $H$ are hyperparameters that balance accuracy and runtime, subject to $M^2,L,H\ll N$ constraints for retaining linearity (Tuli et al., 30 Oct 2025, Kang et al., 2024).

FaceLiVT achieves further acceleration via structural reparameterization, merging convolutional branches at inference for single-step computation and minimized latency (Setyawan et al., 12 Jun 2025).

5. Empirical Results and Domain Applications

Comprehensive benchmarks demonstrate consistent MHLA superiority under fixed compute budgets:

• On ImageNet-1k with DeiT-T, MHLA achieves 75.8% top-1 accuracy, outperforming both softmax attn (72.2%) and global linear attention (69.8%) at identical flops (Zhang et al., 12 Jan 2026).
• In text/image/video generation and retrieval benchmarks, MHLA variants reduce error rates or improve FID/objective metrics by 3–41% depending on domain and task, sometimes surpassing softmax baselines at twofold speedup (Zhang et al., 12 Jan 2026, Tuli et al., 30 Oct 2025).
• MossNet-8×200M+, trained on 2.8T tokens, outperforms Mamba-790M and Qwen-0.5B across ARC-challenge and MMLU by ≥7% absolute, while offering faster inference and reduced memory on both Nvidia A100 and Samsung Galaxy S24 Ultra (Tuli et al., 30 Oct 2025).
• FaceLiVT with MHLA in later transformer stages achieves 8.6×–21.2× lower inference latency than comparable hybrid or pure ViT models on mobile, with negligible drop in challenging identification metrics (Setyawan et al., 12 Jun 2025).

Model / Domain	Latency / Speed	Accuracy / Metric	Reference
MHLA (DeiT-T, ImageNet)	1.1G FLOPs	75.8% (top-1)	(Zhang et al., 12 Jan 2026)
MossNet-8×8M (PPL, QA)	-	PPL=13.1, ↑QA	(Tuli et al., 30 Oct 2025)
FaceLiVT-M-(LA) (IJB-C)	0.67 ms	94.1–95.7% (↑21.2×)	(Setyawan et al., 12 Jun 2025)
iMHSA (ViT-Tiny, ImageNet)	4.7G FLOPs	79.1% (top-1)	(Kang et al., 2024)

6. Comparisons to Prior and Alternative Methods

MHLA contrasts with:

• Global linear attention (Performer, kernelized Softmax): suffers collapsed context, low effective rank.
• Augmented linear attention (e.g., Focused LA, GLA): supplements with depthwise convolution or gating, incurring extra FLOPs and partial rank restoration, but typically at non-negligible cost (Zhang et al., 12 Jan 2026).
• Pure softmax MHSA: retains full expressivity but at prohibitive $O(N^2)$ cost.

MHLA, in its various instantiations, achieves comparable or superior accuracy to augmented alternatives with no extra convolutional kernels and pure $O(N)$ scaling.

7. Ablations, Limitations, and Model Design Choices

Ablation studies indicate:

• Learned mixing coefficients in the blockwise MHLA provide a consistent accuracy gain over fixed or frozen block mixing (Zhang et al., 12 Jan 2026).
• Block/head count $M$ increases effective attention rank, with marginal throughput cost until $M$ nears $N$.
• Feature map choice (e.g., ELU+1) impacts metric stability in vision settings.
• Addition of CPE or gating is orthogonal—useful for small models but unnecessary at scale (Zhang et al., 12 Jan 2026).

A plausible implication is that head-level granularity and query-conditioned mixing are necessary for fully restoring multi-head expressivity in linear attention, but the best-performing integration strategy is domain- and architecture-dependent.

8. Outlook and Open Challenges

MHLA has catalyzed a paradigm shift toward token-diverse, scalable attention with empirical reach across image, language, and edge domains. Remaining challenges include further improvements in head interaction strategies, automated selection of partitioning or expert routing, and formal characterization of mixability with depthwise/local convolutions or SSM kernels in complex architectures.

MHLA’s demonstrated ability to close the linear attention expressivity gap without quadratic penalties indicates sustained relevance for efficient, real-time AI applications and next-generation large model deployments (Zhang et al., 12 Jan 2026, Tuli et al., 30 Oct 2025, Kang et al., 2024, Setyawan et al., 12 Jun 2025).