Interleaved Head Attention

Abstract: Multi-Head Attention (MHA) is the core computational primitive underlying modern LLMs. However, MHA suffers from a fundamental linear scaling limitation: $H$ attention heads produce exactly $H$ independent attention matrices, with no communication between heads during attention computation. This becomes problematic for multi-step reasoning, where correct answers depend on aggregating evidence from multiple parts of the context and composing latent token-to-token relations over a chain of intermediate inferences. To address this, we propose Interleaved Head Attention (IHA), which enables cross-head mixing by constructing $P$ pseudo-heads per head (typically $P=H$), where each pseudo query/key/value is a learned linear combination of all $H$ original queries, keys and values respectively. Interactions between pseudo-query and pseudo-key heads induce up to $P^2$ attention patterns per head with modest parameter overhead $\mathcal{O}(H^2P)$. We provide theory showing improved efficiency in terms of number of parameters on the synthetic Polynomial task (IHA uses $Θ(\sqrt{k}n^2)$ parameters vs. $Θ(kn^2)$ for MHA) and on the synthetic order-sensitive CPM-3 task (IHA uses $\lceil\sqrt{N_{\max}}\rceil$ heads vs. $N_{\max}$ for MHA). On real-world benchmarks, IHA improves Multi-Key retrieval on RULER by 10-20% (4k-16k) and, after fine-tuning for reasoning on OpenThoughts, improves GSM8K by 5.8% and MATH-500 by 2.8% (Majority Vote) over full attention.

• The paper introduces IHA, a novel method that overcomes the compositional bottleneck of standard multi-head attention by interleaving pseudo-heads to allow cross-head interactions.
• It demonstrates a quadratic increase in interaction patterns per head, achieving superior parameter efficiency and enhanced performance on benchmarks like GSM8K and long-context retrieval.
• Empirical and theoretical analyses confirm faster convergence and improved multi-hop reasoning, offering a robust alternative to traditional attention mechanisms.

Interleaved Head Attention: Breaking the Compositional Bottleneck of MHA in LLMs

Introduction and Motivation

LLMs rely fundamentally on Multi-Head Attention (MHA) [vaswani2017attention], but MHA enforces a structural barrier: each of the $H$ heads produces exactly one independent attention matrix, with no direct communication between heads within a single attention computation. This linear scaling restricts MHA's expressivity, especially for multi-step, compositional reasoning tasks that require iterative evidence aggregation and multi-hop relational composition. For instance, many real-world tasks (e.g., multi-hop question answering, arithmetic chain reasoning) necessitate the ability to compose inferences across a sequence of token-to-token relations; a scenario where traditional MHA’s isolated heads are sub-optimal.

Interleaved Head Attention (IHA): Architecture and Mechanism

Interleaved Head Attention (IHA) is introduced as an architectural generalization of MHA, enabling cross-head mixing and enhancing the representational and computational capacity of self-attention modules without sacrificing kernel efficiency and compatibility (e.g., with FlashAttention). IHA constructs $P$ pseudo-queries, pseudo-keys, and pseudo-values for every original head by learned mixing across all $H$ original heads, typically with $P = H$. These pseudo-entities induce up to $P^2$ possible interaction patterns per head, allowing for quadratic scaling in the diversity of attention maps per layer, with additional parameters scaling only as $\mathcal{O}(H^2 P)$. Unlike post-attention mixing (e.g., Talking-Heads [shazeer2020talking]), this process happens before attention, preserving the algebraic structure of attention itself.

The architectural flows are summarized in Figure 1.

Figure 1: Overview of Interleaved Head Attention (IHA). The model mixes original head projections into pseudo-heads, interleaves them in the sequence dimension, and applies standard attention on the expanded sequence, enabling cross-head interaction patterns quadratically within each head.

Theoretical Properties and Expressivity Analysis

IHA is strictly more expressive than traditional MHA for a fixed number of heads. Formally, the parameter set of single-layer $H$-head MHA modules is a strict subset of the set attainable by $H$-head IHA modules with $P \geq 2$ pseudo-heads (plus $4H^2P$ extra parameters). IHA accommodates all MHA behaviors but can express nonlinear operations on repeated-token subspaces that MHA can only represent linearly.

On synthetic polynomial filter tasks (serving as proxies for multi-step reasoning and multi-hop information propagation), MHA needs $\Theta(k)$ heads and $\Theta(k n^2)$ parameters to simultaneously realize all $k$-hop aggregations. In contrast, IHA matches this expressivity with only $\mathcal{O}(\sqrt{k})$ heads and $\Theta(\sqrt{k} n^2)$ parameters, leveraging the combinatorial explosion of $P^2$ interaction patterns per head. This quadratic improvement extends to other expressive tasks, such as CPM-3 (Count Permutation Match-3), where MHA’s minimum parameter complexity scales as $\Theta(n^3)$ while IHA achieves the same with $\Theta(n^2 \sqrt{n})$.

Empirical Validation

Long-Context and Reasoning Benchmark Results

IHA is compared with several state-of-the-art attention mechanisms: Global Attention, Global+Local, Talking-Heads, and Diff Transformer. All evaluations are FLOP-matched to control for computational budget and architectural variance. The benchmark suite includes RULER (for long-context retrieval and generalization), GSM8K and MATH-500 (mathematical reasoning), and multiple code generation datasets (MBPP, HumanEval).

IHA yields strong numerical improvements:

• On RULER after 64k-context fine-tuning, IHA enhances Multi-Key Retrieval accuracy by +10–20% in the 4k–16k regime and yields the highest overall EM (exact match) (see Figure 2).

  Figure 2: IHA delivers strong improvements in Multi-Key Retrieval on long-context RULER benchmarks and raises overall exact match (EM) rates compared to global and hybrid attention mechanisms.

• For reasoning, IHA consistently leads in GSM8K (8.34% EM, +2.73 over Global Attention) and in MATH-500, both before and after supervised fine-tuning (OpenThoughts SFT: +5.8% on GSM8K, +2.8% on MATH-500 Maj@16).
• On code tasks, IHA matches or slightly lags the best baselines, but remains highly robust across all metrics (see the detailed tables in the main text).

Synthetic Compositional Evaluation

To stress-test compositional generalization, IHA’s superiority is further validated in synthetic binary and ternary relation composition benchmarks. Under identical conditions (single-layer, $H=8$, $L=1$), IHA matches or surpasses benchmarked alternatives, including simplicial attention [roy2025fastsimplex], and learns more efficiently.

Figure 3: Final test accuracy comparison for binary and ternary relation composition: IHA achieves systematically higher accuracy than standard MHA and simplicial attention across learning rates, confirming improved sample efficiency for compositional tasks.

Learning dynamics highlight that IHA both converges faster and reaches higher ultimate generalization performance.

Figure 4: Learning curves for binary relation composition showing superior convergence of IHA versus MHA and simplicial attention.

Figure 5: Learning curves for ternary relation composition demonstrating IHA’s advantage in learning three-step compositions.

Architectural and Practical Implications

The major theoretical outcome is that cross-head pseudo-mixing confers an expressivity class and inductive bias unattainable by straightforward linear expansion of MHA. This enables single-layer transformers with IHA to efficiently realize complex compositional and multi-hop relational patterns. Practically, this improves the efficiency of parameter utilization and compute for reasoning tasks, state-tracking, and long-context retrieval use-cases in LLMs. Due to the quadratic scaling in induced pattern diversity, parameters grow only modestly.

One limitation is that IHA, in its global form, increases attention cost quadratically with $P$. This is mitigated in practice by using sliding-window schedules and adapting window sizes to match compute budgets.

Future Directions

Paths for future work include dynamic allocation of pseudo-heads, application to encoder-decoder or vision architectures, and investigation into the efficacy of IHA in more general settings with adaptive pseudo-mixing or learnable composition depths. Additionally, more advanced regularization or routing schemes could leverage IHA's cross-head interaction capacity for hierarchical or multi-modal reasoning.

Conclusion

Interleaved Head Attention constitutes a strategic, theoretically grounded augmentation of attention architectures, breaking the linear combinatorial barrier of MHA with only modest parameter overhead. It demonstrates clear empirical and theoretical benefits for compositional and reasoning-intensive LLM applications, offering a robust path toward more efficient and capable sequence modeling architectures.

---

References

• "Interleaved Head Attention" (2602.21371)
• Vaswani et al., "Attention is All You Need", NeurIPS 2017
• Shazeer et al., "Talking-Heads Attention", (Shazeer et al., 2020)
• Roy et al., "Fast and Simplex: 2-Simplicial Attention in Triton", (Roy et al., 3 Jul 2025)
• Kozachinskiy et al., "Strassen Attention...", (Kozachinskiy et al., 31 Jan 2025)