Iha Attention

Interleaved Head Attention (IHA).

Implements Interleaved Head Attention, arXiv:2602.21371 (Duvvuri et al. 2026). Standard Multi-Head Attention (MHA) is limited by a fundamental linear scaling constraint: H heads produce exactly H independent attention matrices with no communication between heads during the attention computation. This is 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.

IHA 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 O(H**2 * P). The paper proves IHA uses Theta(sqrt(k) * n**2) parameters vs. Theta(k * n**2) for MHA on the synthetic Polynomial task, and ceil(sqrt(N_max)) heads vs. N_max for MHA on the order-sensitive CPM-3 task. Empirically, IHA improves Multi-Key retrieval on RULER by 10-20% (4k-16k context) and improves GSM8K by 5.8% and MATH-500 by 2.8% (Majority Vote) over full attention after reasoning fine-tuning.

Reference:

Duvvuri, S. S., Ekbote, C., Bansal, R., Tiwari, R., Khatri, D., Brandfonbrener, D., Liang, P., Dhillon, I., & Zaheer, M. (2026). “Interleaved Head Attention”. arXiv:2602.21371.

class src.model.attention.latent.iha_attn.IHAAttention(*args: Any, **kwargs: Any)[source]

Bases: Module

Interleaved Head Attention with learned cross-head mixing.

Parameters:

config – Configuration object. Relevant attributes: hidden_size, num_heads, dropout, use_bitnet, mode, and optional iha_num_pseudo_heads (default num_heads).

hidden_size

Input dimensionality.

num_heads

Number of original attention heads H.

head_dim

Per-head dimensionality.

num_pseudo_heads

Number of pseudo-heads per head P.

q_proj, k_proj, v_proj

Projections producing the H original Q/K/V tensors.

q_mix, k_mix, v_mix

Per-mode mixing matrices H*P x H producing the pseudo-head Q/K/V.

out_proj

Output projection (after averaging pseudo-head outputs).

dropout, mode

As in the rest of the family.

Raises:

ValueError – If hidden_size not divisible by num_heads.

forward(x: torch.Tensor, logical_layer_idx: int | None = None) → torch.Tensor[source]

Compute interleaved-head attention with cross-head mixing.

Parameters:

  • x – Input tensor of shape (batch_size, seq_len, hidden_size).

  • logical_layer_idx – Unused; accepted for interface compatibility.

Returns:

Output tensor of shape (batch_size, seq_len, hidden_size).