Interleaved Head Attention

TL;DR

Interleaved Head Attention (IHA) breaks Multi-Head Attention’s “one attention pattern per head” bottleneck by constructing $p$ pseudo-queries, pseudo-keys, and pseudo-values per head as learned linear combinations across all heads, then running standard scaled dot-product attention on the resulting length- $pT$ interleaved sequence. This induces up to $p^2$ attention patterns per head ( $p^2 h$ overall) while preserving the standard attention operator (FlashAttention- and RoPE-compatible). The authors prove a quadratic head-count savings on two compositional-reasoning proxies (Polynomial Filters and the new CPM-3 task) and, in FLOP-matched 2.4B / 240B-token pretraining, report 10–20% relative gains on RULER Multi-Key Retrieval at 4k–16k context and +5.8 / +2.8 majority-vote points on GSM8K / MATH-500 after OpenThoughts SFT.

Key claims

Standard MHA has a “linear scaling” expressivity bottleneck — $h$ heads produce exactly $h$ attention matrices, so representing $k$ distinct relational/step patterns in one layer requires $h \ge k$ heads or extra depth [§1, §3.1].
On the Polynomial Filter problem (a controlled $k$ -hop reasoning proxy), MHA needs $\Theta(k)$ heads and $\Theta(k d^2)$ parameters, while IHA with $p = \sqrt{k}$ pseudo-heads needs only $\Theta(\sqrt{k})$ heads and $\Theta(\sqrt{k}\, d^2)$ parameters — a quadratic-in-task-complexity reduction [§4.2, Thm./proof sketch].
On the new Count Permutation Match-3 (CPM-3) task (order-sensitive 2-step composition with counting), one-layer MHA constructions require $h = \Theta(L)$ heads to build the per-position “workspace” of all cyclic shifts, while IHA realizes the same workspace with $h = \Theta(\sqrt{L})$ heads via factored shift indices [§4.3].
IHA strictly generalizes MHA: there exists a parameter-zero-overhead choice of pseudo-mixing tensors that recovers any MHA module exactly, and on the repeated-token subspace MHA is provably linear while IHA can be nonlinear — so the inclusion $\mathcal{F}_{\mathrm{MHA}} \subsetneq \mathcal{F}_{\mathrm{IHA}}$ is strict for every $p \ge 2$ [§4.1].
IHA performs head mixing before attention (mixing Q/K/V projections) rather than mixing attention logits or weights (Talking-Heads, Knocking, Differential Attention), so the post-mix operation is plain scaled dot-product attention and remains compatible with FlashAttention and RoPE (each pseudo-token gets a distinct RoPE phase via interleaving into a length- $pT$ sequence) [§4].
On RULER long-context evaluation after fine-tuning a 2.4B model from 8k → 64k context, IHA outperforms Global Attention on Multi-Key Retrieval by +27% (4k), +32% (8k), and +112% (16k), and achieves the best RULER average EM (44.0% vs. 35.0% Global, 40.6% Global+Local, 37.2% Diff Transformer) [§5.2, Fig. 2].
After OpenThoughts SFT (8B tokens, temp 0.6, 16 generations), IHA leads all reasoning metrics under FLOP-matched comparison: GSM8K Maj@16 54.2% (+5.8 over Global), MATH-500 Maj@16 18.4% (+2.8), and second-best on both MBPP P@1 / P@10; best Avg. Rank 1.5 across the 6 reported metrics [§5.3, Table 1].
Pretraining-only (5-shot, no SFT) results show IHA already leads on GSM8K EM (+2.73), GSM8K Maj@5 (+2.81), and MATH-500 EM (+0.66) over Global Attention, suggesting the gains come from added architectural expressivity rather than fine-tuning interaction [§5.3, Table 2].

Method

IHA replaces the standard $h$ -head attention block as follows. For each head $i \in [h]$ and each token position $t$ , it builds $p$ pseudo-queries $\tilde{q}_{i,1}, \ldots, \tilde{q}_{i,p}$ as learned linear combinations of all $h$ original queries at that position (and analogously for pseudo-keys and pseudo-values), parameterized by mixing tensors of shape $h \times p \times h$ . The $p$ pseudo-tokens per original token are then interleaved into the sequence dimension, producing a length- $pT$ sequence on which a single standard scaled dot-product attention runs, after which a collapse map / output projection rejoins the per-position pseudo-blocks back to $T$ outputs. The extra parameter cost is $\mathcal{O}(p \cdot h^2 \cdot d_{\text{head}})$ from the mixing tensors — modest because $p \cdot h \ll d$ in practice.

Two design choices matter for compatibility: (1) interleaving rather than stacking the pseudo-head dimension gives each pseudo-token a distinct RoPE phase (it sits at its own position index in the length- $pT$ sequence), and (2) because the post-mix kernel is plain attention, FlashAttention applies directly — the cost is the natural $p^2$ factor from working at sequence length $pT$ instead of $T$ . For experimental FLOP-matching against Global Attention, the authors use a hybrid schedule of four local sliding-window IHA layers (window 512) followed by one global IHA layer in a 4:1 ratio, with $p$ chosen so the average per-layer cost matches global MHA.

The theory rests on two synthetic tasks: Polynomial Filter (representing $\hat{A}, \hat{A}^2, \ldots, \hat{A}^k$ simultaneously in one layer) and CPM-3 (counting ordered pairs satisfying a modular predicate $b \equiv c \cdot a^2 \pmod{p}$ , an arithmetic analogue of two-fact-chained bAbI). The Polynomial Filter result factors $\hat{A}^{ij}$ into $\hat{A}^i \cdot \hat{A}^j$ to assign query/key matrices to blocks of powers; the CPM-3 result factors the cyclic-shift index $s$ as $s = i + j\sqrt{L}$ .

Results

Setup: 2.4B decoder-only Transformer (2560 hidden, 26 layers, 20 heads × 128, FFN 10240, Llama-3 tokenizer, RoPE base 50000), pretraining 240B tokens / 240K steps / 8k context, AdamW + cosine to 10× lower, BF16 + FSDP over 128 H200s. Five attention variants compared FLOP-matched: Global, Global+Local (4:1, window 512), Talking-Heads, Diff Transformer, IHA.
Long context (RULER, fine-tuned 8k → 64k). IHA Multi-Key Retrieval: +27% / +32% / +112% over Global at 4k / 8k / 16k. RULER average EM: IHA 44.0%, Global+Local 40.6%, Diff Transformer 37.2%, Global 35.0% [§5.2, Fig. 2].
Reasoning, SFT (Table 1). GSM8K Maj@16: IHA 54.2% vs Global 48.4% (+5.8). MATH-500 Maj@16: 18.4% vs 15.6% (+2.8). GSM8K P@1: 34.3% vs 29.5% (+4.8). MATH-500 P@1: 10.0% vs 8.8% (+1.2). MBPP P@1: 15.5% (Talking-Heads leads at 15.9%). MBPP P@10: 41.6% (Talking-Heads leads at 43.1%). HumanEval P@1: 0.4% gain. Avg. Rank: IHA 1.5, Talking-Heads 2.5, Diff Transformer 2.8, Global 3.8, Global+Local 4.3.
Reasoning, pretrained-only 5-shot (Table 2). GSM8K EM 8.34% vs Global 5.61% (+2.73); GSM8K Maj@5 8.42% vs 5.61% (+2.81); MATH-500 EM 3.54% vs 2.88% (+0.66); MBPP P@1 24.5% (+1.1); HumanEval P@1 17.1% (−0.1); Avg. Rank 1.4.
The authors emphasize the FLOP-match: IHA’s per-head cost scales as $\mathcal{O}(p^2 T^2 d)$ vs $\mathcal{O}(T^2 d)$ for global MHA, so the 4:1 local/global hybrid schedule is required for parity. Pure global IHA is a $p^2$ slower configuration.

Why it’s interesting

IHA is the most direct competitor on the wiki to Exclusive Self Attention (XSA) — both are surgical, two-equation modifications of vanilla attention that try to free up capacity within a head, and both ship strong long-context / reasoning results from FLOP-matched 1–3B-scale runs. The difference is which limitation each diagnoses: XSA argues each head wastes capacity by re-emitting its own value vector (cosine( $y_i$ , $v_i$ ) is large), so it projects that component out; IHA argues each head wastes capacity by producing only one attention pattern, so it factorizes Q/K/V across heads to induce $p^2$ patterns per head. The two recipes are mechanically orthogonal — XSA edits the output of attention, IHA edits its inputs — and the “stack XSA + IHA” experiment is the obvious follow-up neither paper has run. Both papers also share IHA’s pitch that the modification stays FlashAttention- and RoPE-compatible, which is what makes them deployable at scale (compare to multi-token / simplicial attention, which doesn’t).

The architectural lineage is the same family already mapped by Training stability at scale: a wave of attention-side micro-architecture surgery (XSA, Gated Attention in A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training, Mamba-3’s gating in Mamba-3: Improved Sequence Modeling using State Space Principles, Differential Attention) that share the recipe “find a load-bearing geometric artifact in MHA, patch it without breaking efficient kernels.” IHA is the one that targets expressivity at fixed depth rather than stability or scaling-law slope. For Luma the concrete hook is the same as XSA’s: gains grow with sequence length / context, the regime Luma’s video DiTs spend their compute on; and the multi-key retrieval result (+112% at 16k) is the kind of finding that would motivate trying IHA in a DiT’s cross-attention to long text/image context. Worth noting that IHA’s empirical wins (5.8 points on GSM8K Maj@16) are reasoning-tasks; the paper does not test on diffusion / generative video.