---
title: Multi-Head Mixed Attention (MHMA)
url: https://www.emergentmind.com/topics/multi-head-mixed-attention-mhma
type: topic
---

# Multi-Head Mixed Attention (MHMA)

Multi-Head Mixed Attention (MHMA) generalizes multi-head attention by introducing learned, dynamic, or interaction-rich mechanisms for mixing among attention heads or even entire attention schemes. The MHMA framework encompasses dynamic head selection (mixture-of-head attention), attention expert routing, cross-head feature mixing, and selection among whole attention schemes. These advances improve capacity scaling, training dynamics, inference efficiency, and the representational expressivity of Transformer-based and recurrent architectures.

## 1. Core Concepts and Definitions

Standard Multi-Head Attention (MHA) processes an input $X \in \mathbb{R}^{L \times d}$ via $H$ heads, each with independent projections $W_i^Q, W_i^K, W_i^V$. Classical aggregation concatenates all head outputs and applies a linear projection:
\[
\text{MHA}(X) = \text{Concat}(\text{Attn}_1, \ldots, \text{Attn}_H) W^O
\]
However, this treats head outputs independently until the final projection and does not adaptively exploit head importance, specialization, or non-trivial interactions.

**Multi-Head Mixed Attention** (MHMA) introduces mechanisms that facilitate:
- Mixtures of head outputs with per-token or per-context dynamic gating.
- Routing-based selection of a sparse or contextually determined subset of heads or expert blocks.
- Explicit cross-head or cross-scheme feature mixing prior to output aggregation.
- Extension to recurrent/state-space layers as mixture-of-experts multi-head modules.

MHMA includes prominent instantiations such as Mixture-of-Head Attention (MoH) [2410.11842], Mixture of Attention Heads (MoA) [2210.05144], Knocking-Heads Attention (KHA) [2510.23052], and mixtures of whole attention schemes (MoAS) [2512.20650].

## 2. MHMA Mechanisms: Gating, Routing, and Mixture Formulations

### Mixture-of-Heads (MoH) and MoA

Both MoH and MoA replace uniform head aggregation with a token-dependent weighted mixture. Let $V_i$ denote the $i$th head's output (post-output projection if present), and $g_i(x_t)$ a learned gate for input $x_t$. The generic MHMA formula is:
\[
y_t = \sum_{i=1}^H g_i(x_t) V_i(x_t)
\]
MoH augments standard MHA with a router that splits heads into “shared” (always active) and “routed” (dynamically top-$K$ selected) heads. The router processes per-token input through lightweight MLPs producing selection and weighting scores, enforcing load balancing via auxiliary losses [2410.11842].

MoA further increases scalability by defining a large pool of $E$ attention experts and evaluating only a sparse top-$k$ subset per token [2210.05144], using noisy-top-$k$ MoE gating. Key/value projections may be shared across experts for efficiency.

### Knocking-Heads Attention

KHA enables cross-head feature interaction prior to attention score computation by inserting a shared, diagonally-initialized projection matrix $T^\cdot$ after the per-head linear projections:
\[
Q'_i = Q_i T^Q,\quad K'_i = K_i T^K,\quad V'_i = V_i T^V
\]
\[
\text{Attn}'_i = \text{Softmax}(Q'_i {K'_i}^T / \sqrt{d_k}) V'_i
\]
The shared $T^\cdot$ matrices are initialized as identity to preserve head specialization but learn off-diagonal cross-head mixing during training. This approach is agnostic to the precise head mixing strategy, supports minimal parameter/FLOP overhead, and can be retrofitted to existing attention variants (MHA, GQA, GTA) [2510.23052].

### Mixture of Attention Schemes (MoAS)

MoAS generalizes MHMA over attention *schemes* (e.g., full MHA, GQA, MQA) using a learned per-token router. For input $x_t$, a softmax MLP router produces mixture weights for the three parallel “scheme experts”:
\[
g_t = \text{softmax}(W_2 \cdot \text{GELU}(W_1 x_t))
\]
\[
y_t = \alpha_t^{\text{MHA}} y_t^{\text{MHA}} + \alpha_t^{\text{GQA}} y_t^{\text{GQA}} + \alpha_t^{\text{MQA}} y_t^{\text{MQA}}
\]
A load-balancing regularization term encourages meaningful usage of all schemes [2512.20650].

### MHMA in State-Space Models

MossNet lifts single-head state-space models (SSMs) into H-head linear attention via token-wise mixture-of-experts on the SSM’s time-mixing kernels $\bar{B}_t,\ C_t$ and the channel-mixing MLP. A softmax router generates per-token mixture coefficients selecting the top-k from $E$ SSM experts, creating an ensemble of time-mixing “heads” [2510.26182].

## 3. Algorithmic Designs: Routing Networks and Aggregation

MHMA architectures feature specialized routers and aggregation rules to achieve sparse, context-sensitive, or joint mixing. Key design patterns include:

- **Noisy-Top-$k$ Routing**: Used in MoA [2210.05144], assigns tokens to the top-$k$ scoring experts, injecting noise for improved training stability and load balance.
- **Two-Stage Gate Assignment**: Used in MoH [2410.11842], splits between “shared” always-on heads and token-routed heads, with independent softmax normalization for each group and a softmax mixing for balance.
- **Shared vs. Per-Head Mixing**: KHA [2510.23052] deploys shared projection matrices for cross-head interaction, initialized to identity, ensuring that specialization is preserved early but rich interactions can emerge with training.
- **Per-Scheme Routing**: MoAS [2512.20650] performs softmax-based routing among parallel full attention mechanisms.
- **Iterative Routing-by-Agreement**: Capsule-based aggregation [1904.03100] iteratively refines slot assignments between head outputs and final representation slots, achieving adaptive, content-based head aggregation.

| MHMA Mechanism | Routing Network      | Head/Scheme Selection      | Aggregation Paradigm         |
| -------------- | ------------------- | ------------------------- | --------------------------- |
| MoH/MoA        | Shallow MLP, Noisy-$k$ | Top-$k$ or dynamic, per-token | Weighted sum (gated/expert) |
| KHA            | None (fixed, shared)   | All heads, feature-level mixing | Pre-attention mixing        |
| MoAS           | 2-layer MLP            | Softmax over attention types | Scheme-level mixture        |
| MossNet        | Softmax MLP             | Top-$k$ SSM experts          | MoE over time/channel       |

## 4. Computational Complexity and Scalability

MHMA methods introduce conditional computation, thus decoupling model capacity from per-token inference cost:

- **Efficiency**: MoH/MoA route each token through $k \ll E$ heads, enabling scaling to hundreds or thousands of experts while keeping runtime cost fixed.
- **Parameter Overhead**: KHA adds $3d_k^2$ parameters per layer ($d_k = d/H$), which is $<1\%$ of base MHA parameter count [2510.23052]. MoH routers add $(h_s+(H-h_s)+2)d_{in}$ extra parameters per layer [2410.11842].
- **FLOPs**: KHA-Linear incurs $<1\%$ of baseline MHA FLOPs; MoH only computes attention for active heads, reducing runtime in inference.
- **KV Cache and Memory**: MoAS enables per-token choice of attention scheme, trading off between quality (MHA) and cache efficiency (MQA/GQA) with dynamic memory profiles [2512.20650].
- **SSM Linear Scaling**: MossNet maintains $O(Ld^2)$ complexity due to SSM recurrences, and scales number of experts/heads independently of per-token compute [2510.26182].

## 5. Empirical Performance and Benefits

MHMA instantiations consistently demonstrate improvements in model expressivity, downstream performance, and efficiency:

- **MoA/MoH**: Outperform standard MHA at fixed or lower FLOPs, e.g., in machine translation 29.4 BLEU (MHMA-large) vs. 28.4 (Transformer-big), and in image classification MoH achieves equivalent or better accuracy with only 50–75% of heads active [2410.11842, 2210.05144].
- **KHA**: Reduces loss spikes during pretraining, increases downstream task scores (e.g., +4.32 on RACE, +3.90 on HumanEval-Plus), and yields average +1.26 points across benchmarks at minimal cost [2510.23052].
- **MossNet**: Achieves better perplexity than non-MHMA recurrent SSMs and dense Transformers, with real-device memory and speed advantages for long contexts [2510.26182].
- **MoAS**: Achieves competitive validation loss (2.3074 vs. 2.2940 for pure MHA) but with flexibility to trade memory for compute at runtime via routing [2512.20650].
- **Capsule Routing Aggregation**: Iterative routing yields superior linguistic structure capture and up to +1.16 BLEU on translation compared to linear concat+projection aggregation [1904.03100].

| Architecture | Key Empirical Results                                    |
| ------------ | ------------------------------------------------------- |
| MoH          | +2.4pp accuracy (LLaMA3-8B, 75% heads) [2410.11842]     |
| MoA          | +1.1 BLEU vs. Transformer-base [2210.05144]             |
| KHA          | +4.32 (RACE), +3.90 (HumanEval-Plus), −0.015 loss [2510.23052] |
| MossNet      | PPL=13.1 (Cosmopedia), +5.8% accuracy vs Qwen2.5 [2510.26182]|

## 6. Model Specialization and Interpretability

MHMA’s token-specific routing and adaptive mixing imbue attention heads or experts with distinct functional roles:

- **Specialization**: PMI analysis in MoA shows experts focusing on semantically or syntactically coherent token groups (locations, technological terms, adverbs) [2210.05144].
- **Dynamic Patterns**: MoH’s head-load visualization reveals non-uniform, context-dependent specialization, contrasting with uniform summation in standard MHA [2410.11842].
- **Interpretability Metrics**: Balanced expert assignment (moA and MoH) is achieved using auxiliary load-balance and z-losses, distributing token assignments broadly and preventing collapse.

A plausible implication is that MHMA architectures favor the emergence of modular, interpretable sub-functions within their routing domains.

## 7. Extensions, Limitations, and Theoretical Connections

MHMA is extensible along several axes:

- **Expert Granularity**: Extends from heads (MoA, MoH), to full attention schemes (MoAS), to SSM blocks (MossNet), and can generalize to non-attention experts (convolutions, local, or sparse attention).
- **Mixing Strategies**: Permits integration of richer head-head interactions (KHA), per-layer or per-head adaptation, or dynamic gating coupled with quantization or sparsity.
- **Limitations**: KHA is directly applicable only to softmax-based attention; MoAS incurs extra compute at training due to parallelism, and capsule routing can introduce latency due to iterative refinement [2510.23052, 2512.20650, 1904.03100].

A plausible implication is that, as MHMA mechanisms proliferate and mature, future research will systematize hybrid mixture-of-experts architectures spanning heads, blocks, and even full attention paradigms, yielding increasingly modular and scalable sequence models.

Source: https://www.emergentmind.com/topics/multi-head-mixed-attention-mhma