Interleaved Head Attention

TL;DR

Interleaved Head Attention (IHA) enhances multi-head attention by enabling cross-head communication, improving reasoning capabilities and efficiency in large language models, with demonstrated benefits on synthetic and real-world tasks.

Contribution

The paper introduces Interleaved Head Attention, a novel mechanism allowing cross-head mixing in multi-head attention to improve reasoning and parameter efficiency.

Findings

IHA outperforms standard MHA on synthetic polynomial and order-sensitive tasks.

IHA improves retrieval accuracy on RULER by 10-20%.

IHA enhances reasoning performance on GSM8K and MATH-500 datasets.

Abstract

Multi-Head Attention (MHA) is the core computational primitive underlying modern Large Language Models (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 $\Theta(\sqrt{k}n^2)$ parameters vs. $\Theta(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.