Explicit Multi-head Attention for Inter-head Interaction in Large Language Models

TL;DR

This work introduces Multi-head Explicit Attention (MEA), an attention variant that explicitly models cross-head interactions via a Head-level Linear Composition (HLC) and stabilizes training with GroupNorm. MEA provides a unified view that subsumes and contextualizes prior designs like Differential Transformer and Talking-Heads Attention, showing how reordering linear modules and applying normalization affects optimization and expressiveness. Through from-scratch pretraining on a LLaMA3.2-1B-like architecture, MEA enables larger, more stable learning rates and faster convergence, outperforming baselines on knowledge, science, and reasoning tasks. Additionally, MEA supports a practical 50% KV-cache memory reduction via low-rank, head-wise reconstruction, preserving performance on most tasks while highlighting sensitivity in mathematical reasoning to compression and data quality effects. Overall, MEA advances inter-head communication in Transformers and offers tangible efficiency benefits for large-scale pretraining and inference.

Abstract

In large language models built upon the Transformer architecture, recent studies have shown that inter-head interaction can enhance attention performance. Motivated by this, we propose Multi-head Explicit Attention (MEA), a simple yet effective attention variant that explicitly models cross-head interaction. MEA consists of two key components: a Head-level Linear Composition (HLC) module that separately applies learnable linear combinations to the key and value vectors across heads, thereby enabling rich inter-head communication; and a head-level Group Normalization layer that aligns the statistical properties of the recombined heads. MEA shows strong robustness in pretraining, which allows the use of larger learning rates that lead to faster convergence, ultimately resulting in lower validation loss and improved performance across a range of tasks. Furthermore, we explore the parameter efficiency of MEA by reducing the number of attention heads and leveraging HLC to reconstruct them using low-rank "virtual heads". This enables a practical key-value cache compression strategy that reduces KV-cache memory usage by 50% with negligible performance loss on knowledge-intensive and scientific reasoning tasks, and only a 3.59% accuracy drop for Olympiad-level mathematical benchmarks.

Figures (8)

• Figure 1: Comparison of how different attention variants compute pre-softmax attention scores, using MHA-based variants as an example. In the modified THA, the learnable linear combination is moved earlier in the computation flow, making it equivalent to MEA's pre-softmax operation.
• Figure 2: Test loss curves fitted using scaling laws in the learning rate selection experiment, under fixed learning rates $1 \times 10^{-3}$ for different attention variants.
• Figure 3: Test loss curves fitted using scaling laws in the full-scale pretraining setting, using cosine decay schedules initialized from the optimal fixed learning rates in Figure \ref{['fig:lr-selection']}, comparing different model variants.
• Figure 4: Test loss curves fitted using scaling laws in the learning rate selection experiment, under different fixed learning rates for Transformer.
• Figure 5: Test loss curves fitted using scaling laws in the learning rate selection experiment, under different fixed learning rates for Transformer+GroupNorm.