Mixture of Sparse Attention: Content-Based Learnable Sparse Attention via Expert-Choice Routing

TL;DR

MoSA tackles the quadratic cost of self-attention by introducing per-head, content-based sparse token selection via Expert-Choice routing, reducing per-head complexity to O(k^2 + T) and enabling more specialized heads under the same compute budget. The method jointly learns token selections and attention, achieving perplexity gains up to 27% over dense baselines under iso-FLOP budgeting while also reducing KV-cache and improving wall-clock time with a pure PyTorch implementation. Hybrid models that blend MoSA heads with a small set of dense heads are essential for stability and peak performance, outperforming fixed sparse and Routing Transformer baselines across multiple model scales. The approach demonstrates strong results on standard and long-sequence language modeling tasks, with practical implications for efficient deployment and potential extension to longer inputs and cross-modal settings.

Abstract

Recent advances in large language models highlighted the excessive quadratic cost of self-attention. Despite the significant research efforts, subquadratic attention methods still suffer from inferior performance in practice. We hypothesize that dynamic, learned content-based sparsity can lead to more efficient attention mechanisms. We present Mixture of Sparse Attention (MoSA), a novel approach inspired by Mixture of Experts (MoE) with expert choice routing. MoSA dynamically selects tokens for each attention head, allowing arbitrary sparse attention patterns. By selecting $k$ tokens from a sequence of length $T$, MoSA reduces the computational complexity of each attention head from $O(T^2)$ to $O(k^2 + T)$. This enables using more heads within the same computational budget, allowing higher specialization. We show that among the tested sparse attention variants, MoSA is the only one that can outperform the dense baseline, sometimes with up to 27% better perplexity for an identical compute budget. MoSA can also reduce the resource usage compared to dense self-attention. Despite using torch implementation without an optimized kernel, perplexity-matched MoSA models are simultaneously faster in wall-clock time, require less memory for training, and drastically reduce the size of the KV-cache compared to the dense transformer baselines.

Figures (7)

• Figure 1: MoSA layer compared to the dense attention layer. MoSA replaces each dense head with multiple heads with a learnable sparsity pattern. Each head selects its own $k$ tokens to process. MoSA calculates query, key, and value projections only for the selected token and computes the attention only between them. It drops the rest of the tokens, leading to more efficient compute utilization. This reduces the computational and memory complexity on a sequence of length $T$ from $O(T^2)$ to $O(k^2+T)$. The saved compute budget can be used to scale up the number of heads.
• Figure 2: Attention variants visualized. In the plot, the colors indicate different heads. Sparse attention methods are roughly FLOP-matched and have sparsity $\rho=2$. One Routing Attention head corresponds in FLOP-cost to $\rho$ Fixed/MoSA heads. Fixed sparse attention uses only $k=\frac{T}{\rho}$ tokens in specific positions, with regular stride. The Routing Attention clusters tokens within each head into $\rho$ clusters of size $k$ based on their representations. MoSA selects $k$ tokens for each attention head independently based on their representations.
• Figure 3: Perplexity ($\downarrow$) of FLOP matched models under different sparsities. Each plot corresponds to a specified FLOP budget per step. The number in parenthesis is the number of parameters of the dense baseline. Sparsity 1 represents the dense baseline. As sparsity increases, MoSA's perplexity improves monotonically until reaching a saturation point around sparsity 32-64, beyond which performance deteriorates. This is likely because at very high sparsity levels, each attention head selects only a few tokens, which is insufficient to capture the complex relations. On the other hand, other sparse methods fail to reach the perplexity of the dense baseline in the IsoFLOP setting. We explore fewer sparsity levels for larger models due to excessive memory requirements.
• Figure 4: Perplexity of sparse‑attention methods (MoSA, Fixed, and Routing) as sequence length increases. Each method has a fixed size window size (cluster size for the Routing Transformer, number of tokens selected for each head in MoSA and Fixed) regardless of total sequence length. MoSA matches the computational cost of the fixed sparsity baseline while requiring fewer FLOPs than the Routing Attention and consistently achieves the lowest perplexity.
• Figure 5: Perplexity of IsoFLOP matching models under pure MoSA setting. Each curve corresponds to a given FLOP budget. For a given sparsity, we replace all dense heads with a FLOP equivalent number of MoSA heads. In contrast to Fig \ref{['fig:main_results']}, sparse models fail to outperform the baseline (apart from the Large model). This demonstrates the symbiotic relation between dense heads and MoSA heads in the hybrid model.