---
title: Knocking-Heads Attention (KHA)
url: https://www.emergentmind.com/topics/knocking-heads-attention-kha
type: topic
---

# Knocking-Heads Attention (KHA)

Knocking-Heads Attention (KHA) is an enhancement for multi-head attention (MHA) architectures, enabling lightweight and direct cross-head feature interactions prior to the scaling dot-product operation. KHA addresses the representational trade-offs inherent in classical MHA as the number of attention heads increases—specifically, the loss of per-head expressiveness due to the reduction in dimensionality. By introducing shared, diagonally-initialized projection matrices across all heads, KHA preserves head specificity at initialization but allows feature-level communication, leading to improved expressiveness and training stability with minimal computational overhead [2510.23052].

## 1. Motivation and Context

Classical MHA architectures, foundational to modern large language models, decompose their input $X \in \mathbb{R}^{L \times d}$ into $n$ attention heads of dimension $d_k = d / n$. An increased number of heads is believed to offer finer-grained relational modeling, yet each head consequently operates in a lower-rank subspace, causing a loss in representational power—a phenomenon highlighted in the literature as a low-rank bottleneck. Furthermore, traditional MHA and prominent variants such as Grouped-Query Attention (GQA) and Grouped-Tied Attention (GTA) handle each head in isolation, concatenating outputs without inter-head feature mixing. Methods like Talking-Heads Attention introduce mixing at the attention logits or post-softmax level but at the expense of quadratic overhead and limited scalability. CollabHead applies large shared projections but at the cost of specialization and increased FLOPs.

Knocking-Heads Attention circumvents these drawbacks by integrating a shared, lightweight projection in the per-head feature pipeline, effecting $O(d^2/n)$ additional computation. This arrangement maintains head specialization while enabling beneficial cross-head communication [2510.23052].

## 2. Formal Description and Architectural Details

Given $X \in \mathbb{R}^{L\times d}$ and $n$ heads, KHA first applies standard head-specific projections:

\[
Q_i = X W_i^Q, \quad K_i = X W_i^K, \quad V_i = X W_i^V \qquad (W_i^* \in \mathbb{R}^{d \times d_k})
\]

In KHA-Linear, three shared matrices $T^Q$, $T^K \in \mathbb{R}^{d_k \times d_k}$, and $T^V \in \mathbb{R}^{d_v \times d_v}$ are introduced and applied post projection but before attention computation:

\[
\tilde{Q}_i = Q_i T^Q, \quad \tilde{K}_i = K_i T^K, \quad \tilde{V}_i = V_i T^V
\]

Attention proceeds as in standard MHA but on these transformed representations, and the same $T^*$ matrices are reused across all heads, enforcing feature-level coupling. For inference, $T^*$ can be fused into $W_i^*$, incurring no additional inference cost.

In KHA-MLP, the shared $T^V$ is replaced with a small, gated MLP:

\[
\tilde{V}_i = \mathrm{MLP}(V_i) = 2(V_i W^{\mathrm{up}} \odot \sigma(V_i W^{\mathrm{gate}})) W^{\mathrm{down}}
\]
where $W^{\mathrm{up}}$, $W^{\mathrm{gate}}$, $W^{\mathrm{down}} \in \mathbb{R}^{d_v \times d_v}$ are shared and $\sigma(\cdot)$ denotes the sigmoid.

## 3. Initialization and Training Dynamics

KHA employs diagonal initialization to preserve head-specific inductive bias. For KHA-Linear, $T^Q$, $T^K$, $T^V$ are initialized as identity matrices plus small positive diagonal perturbations; for KHA-MLP, $W^{\mathrm{up}}$ and $W^{\mathrm{down}}$ are initialized as diagonal, and $W^{\mathrm{gate}}$ as zeros, ensuring the initial forward pass is equivalent to unmodified MHA. As training progresses, gradients activate off-diagonal elements and $W^{\mathrm{gate}}$, progressively enabling cross-head feature blending without abrupt function drift [2510.23052].

## 4. Integration with Attention Variants

KHA’s head-sharing projection is inserted directly after per-head linear maps for $Q$, $K$, $V$, making it compatible with any softmax-based multi-head framework including MHA, GQA, GTA, multi-query attention (MQA), and multi-linear attention (MLA). For example, within GQA, $T^V$ is applied to each group’s $V$ projection; within GTA, to the shared $V$ projection. No architectural changes are required beyond the introduction of the shared projection.

A concise pseudocode for KHA-Linear:

```python
def KHA_Attention(X):
    for each head i in 1...n:
        Q_i = X @ W_i^Q
        K_i = X @ W_i^K
        V_i = X @ W_i^V

        tilde_Q_i = Q_i @ T^Q
        tilde_K_i = K_i @ T^K
        tilde_V_i = V_i @ T^V

        O_i = softmax(tilde_Q_i @ tilde_K_i^T / sqrt(d_k)) @ tilde_V_i

    return concat(O_1...O_n) @ W^O
```
[2510.23052]

## 5. Computational and Parameter Overhead

KHA introduces negligible resource requirements relative to standard MHA. For KHA-Linear, the additional training computation is $6 L d^2 / n$ FLOPs per layer, corresponding to only 0.55% of full layer compute and 1.17% of MHA’s FLOPs for typical configurations ($L=2048$, $d=1024$, $n=32$). Parameter overhead is $O(d^2/n)$—under 1% of attention-block parameters. At inference, projections may be statically fused, resulting in zero additional runtime cost [2510.23052].

## 6. Empirical Evaluation and Results

KHA demonstrated substantial improvements in large-scale training regimes. A 6.1B parameter Mixture-of-Experts (MoE) model—1.01B parameters active per token, with GQA (32 query heads, 4 KV groups) and head size 128—was trained on 1T high-quality tokens using the FSDP-based Adam optimizer. KHA, notably the MLP variant applied to $V$, reduced the occurrence and severity of loss spikes, delivering a consistent 0.015 lower loss in the stabilized regime.

Results summary:

| Setting                           | KHA Effect (Loss / Downstream)                                     |
|:---------------------------------- |:------------------------------------------------------------------ |
| More KV heads (1→4, Table 2)      | KHA-MLP lowers loss up to 0.024; KHA-Linear smaller benefit        |
| Projection ablation (Table 3)      | $V$ projection is most effective; MLP > linear; gating alone harmful|
| Attention variant compatibility    | Loss reductions in MHA, GQA, GTA, MQA, MLA ($-0.010$ to $-0.020$)  |
| Downstream task metrics (Table 5)  | RACE +4.32, Code +3.90, Math +1.62, Overall +1.26 pts (across 30+) |

Empirical results indicate KHA is most effective when cross-head transformations are applied to $V$ via a shared MLP, though $Q$/$K$ projections yield marginal gains [2510.23052].

## 7. Interpretation, Limitations, and Prospective Extensions

KHA’s improvements arise from two factors: (1) increased expressiveness via statistical strength-sharing among heads, mitigating low-rank constraints; (2) implicit regularization, as diagonal initialization stabilizes early training and reduces catastrophic gradient spikes.

Limitations include variable gains across tasks, with minor impact on general-knowledge benchmarks. Further research could address task-specific gating or sparsity, richer nonlinear sharing architectures (such as intra-MLP cross-head attention), and adaptive per-layer or per-token sharing to optimize specialization versus collaboration. Extending the diagonal-initialized head-sharing principle to other neural components, including convolutional, recurrent, or retrieval-augmented modules, remains an open avenue [2510.23052].

Source: https://www.emergentmind.com/topics/knocking-heads-attention-kha