Head-level Linear Composition (HLC)

• HLC is a module in multi-head attention that linearly recombines key and value vectors to enable explicit inter-head communication.
• It enhances training dynamics and achieves memory savings by enabling robust convergence and compressing the KV-cache efficiently.
• Integration within the MEA framework leverages low-rank factorizations and head-level GroupNorm to optimize parameter efficiency while preserving full multi-head capacity.

Head-level Linear Composition (HLC) is a structural module for multi-head attention mechanisms in LLMs, designed to enable explicit inter-head communication by applying learned linear combinations to the key and value vectors across attention heads. Introduced as a central component of the Multi-head Explicit Attention (MEA) framework, HLC facilitates richer interaction between attention heads, promoting robust training dynamics, enhanced parameter efficiency, and practical advantages for memory usage via virtual head reconstruction and KV-cache compression (Peng et al., 27 Jan 2026).

1. Mathematical Formulation

Let $N$ denote sequence length, $D$ the hidden dimension, $H$ the number of attention heads, and $d_k = d_v = D/H$ the per-head key/value dimensionality. After the standard input projection, the model yields:

$$\mathbf{K}_{\rm comp} = \mathrm{reshape}(\mathbf{X} \mathbf{W}^0_K, N \times H \times d_k), \quad \mathbf{V}_{\rm comp} = \mathrm{reshape}(\mathbf{X} \mathbf{W}^0_V, N \times H \times d_k)$$

These tensors represent the stack of $H$ key and value heads per token.

To achieve inter-head mixing, HLC introduces two learnable matrices $$\mathbf{W}_K, \mathbf{W}_V \in \mathbb{R}^{H \times H'}$$ that linearly recombine the $H$ heads into $H'$ "mixed" heads, where typically $H' = H$. The recombined keys and values are:

$$\begin{aligned} \mathbf{K}_{\rm hlc} &= \mathrm{HLC}(\mathbf{W}_K, \mathbf{K}_{\rm comp}) = \left[\sum_{j=1}^H (\mathbf{W}_K)_{j,i} \mathbf{K}_j \right]_{i=1}^{H'} \in \mathbb{R}^{N \times H' \times d_k}\ \mathbf{V}_{\rm hlc} &= \mathrm{HLC}(\mathbf{W}_V, \mathbf{V}_{\rm comp}) = \left[\sum_{j=1}^H (\mathbf{W}_V)_{j,i} \mathbf{V}_j \right]_{i=1}^{H'} \in \mathbb{R}^{N \times H' \times d_k} \end{aligned}$$

A high-performance implementation uses the Einstein summation:

K_hlc = einsum("n h d, h h' -> n h' d", K_comp, W_K) V_hlc = einsum("n h d, h h' -> n h' d", V_comp, W_V)

This flexible parameterization enables the module to learn arbitrary head-wise mixing before scaled-dot-product attention.

2. Mechanism of Inter-Head Communication

In standard multi-head attention, each head operates independently, generating its own keys and values with no explicit cross-head interaction prior to output aggregation. HLC, by contrast, synthesizes each new head as a learned linear combination of all original heads:

• Pre-softmax, each query head attends to a synthesized key subspace $\phi(\mathbf{K}_{\rm hlc})$ rather than a single, isolated $\phi(\mathbf{K}_j)$.
• Post-softmax, each output head aggregates a learned mixture of original value heads.

The weights $\mathbf{W}_K$ and $\mathbf{W}_V$ are trained to adaptively route information across heads at every layer, enabling the model to construct specialized combinations and promote cross-head feature extraction beyond "blindly" independent attention branches. This design supports richer internal representations within each Transformer layer (Peng et al., 27 Jan 2026).

3. Integration into Multi-head Explicit Attention (MEA)

MEA incorporates HLC by replacing the standard keys and values in the scaled-dot-product attention operation with their HLC recombined counterparts, followed by a head-level Group Normalization (GroupNorm) to align statistical properties and preserve diversity among heads. The MEA layer operates as:

$$\begin{aligned} \mathbf{A} &= \mathrm{softmax}\left( \frac{\phi(\mathbf{Q}) \phi(\mathbf{K}_{\rm hlc})^\top}{\sqrt{d_k}} \right) \ \mathbf{C}_i &= \sum_{t=1}^N A_{i, t} \, \mathbf{V}_{{\rm hlc}, t, i} \end{aligned}$$

The collected outputs $[\mathbf{C}_1, \dots, \mathbf{C}_{H'}]$ undergo GroupNorm along the head dimension:

$$\widetilde{\mathbf{C}} = \mathrm{GroupNorm}(\mathbf{C}_1, \dots, \mathbf{C}_{H'})$$

The reshaped tensor is finally projected by the output linear layer. In practical settings, $H'=H$. The following summarizes the graph:

Q = X W_Q^0 K_comp = X W_K^0 V_comp = X W_V^0 K_hlc = HLC(W_K, K_comp) V_hlc = HLC(W_V, V_comp) A = softmax((φ(Q)·φ(K_hlc).T)/sqrt(d_k)) C = A · V_hlc # shape N×H×d_k Ĉ = GroupNorm_over_heads(C) output = reshape(Ĉ,N×(H d_k)) · W_O

The insertion of HLC and subsequent GroupNorm yields a plug-and-play extension of standard attention mechanisms, preserving computational and architectural simplicity (Peng et al., 27 Jan 2026).

4. Parameter Efficiency and Virtual Head Reconstruction

HLC's head-combination matrices $\mathbf{W}_K$, $\mathbf{W}_V$ permit further parameter and memory optimization. By factoring each as a low-rank product, one can reduce $H$ original heads to $H'<H$ "basis" heads (virtual heads) and reconstruct the full set via matrix multiplication. The low-rank approximation: $$\mathbf{W}_K \approx U_K V_K^\top, \quad U_K, V_K \in \mathbb{R}^{H \times r}, \quad r \ll H$$ permits basis-head compression, only requiring storage of $U_K$, $V_K$. In practice, an SVD on the key-projection matrix $\mathbf{W}^0_K$ enables truncation to $r$ leading singular vectors, yielding a decomposition for compressed inference.

During forward inference, only the $H'$ compressed heads must be cached (KV-cache), with the full set reconstructed as needed. This reduces KV-cache memory from $\mathcal{O}(L H T d_k)$ to $\mathcal{O}(L H' T d_k)$, where $L$ is the layer count and $T$ is the sequence position. Empirically, nearly full multi-head capacity is retained despite large memory savings (Peng et al., 27 Jan 2026).

5. Empirical Performance and Observed Benefits

HLC, through its deployment in MEA, demonstrates robust empirical performance:

• Pretraining robustness: In 1B-parameter models, MEA tolerates peak learning rates up to $3 \times 10^{-3}$, while competing baselines diverge at rates beyond $1 \times 10^{-3}$. Models converge faster and reach lower validation loss by 500B tokens.
• Downstream accuracy: On benchmarks including PIQA, OBQA, WinoGrande, HellaSwag, ARC-e, and ARC-c, MEA with HLC achieves 0.5–1.5 accuracy point gains over Transformer + GroupNorm or DFA baselines.
• KV-cache compression: In Qwen3-30B-A3B, compressing cached heads from $4\to2$ (50% memory reduction) with low-rank HLC reconstruction results in negligible drops of $-0.32\%$ for science, $-1.40\%$ for knowledge, and $-3.59\%$ for challenging math benchmarks after brief recovery (Recover+CPT) training.

Summary properties of HLC:

Property	Description
Head-wise transformation	$$\mathbb{R}^{N \times H \times d_k} \xrightarrow{\;\mathrm{HLC}(W)\;} \mathbb{R}^{N \times H \times d_k}$$
Expressiveness	Learns arbitrary linear mixtures of heads pre-attention
Parameter efficiency	Can be factorized to $\mathcal{O}(rH)$ for virtual head reconstruction
Training stability	Enhanced by head-level GroupNorm
Empirical effect	Faster convergence, lower loss, improved accuracy, 50% KV-cache reduction with minor degradation

This combination of expressiveness, efficiency, and practical benefit establishes HLC as a concise, effective core for inter-head interaction in contemporary attention architectures (Peng et al., 27 Jan 2026).

6. Context and Significance in Attention Architectures

HLC addresses a key limitation in classic multi-head attention: the lack of explicit information exchange among parallel attention branches prior to output aggregation. By synthesizing new heads as learned, adaptively weighted mixtures of originals, HLC allows the network to develop task-adaptive, cross-head representations and routing strategies. Its lightweight design and compatibility with fast attention implementations make it suitable for scaling in LLMs and enhance parameter and memory efficiency through factorized, "virtual" heads. A plausible implication is that future attention modules may generalize this head-level transformation paradigm for broader sequence modeling and resource-constrained deployments, as indicated by the empirical scalability and effectiveness of MEA with HLC (Peng et al., 27 Jan 2026).