[论文解读] Explicit Multi-head Attention for Inter-head Interaction in Large Language Models

Runyu Peng, Yunhua Zhou+5 more|arXiv (Cornell University)|Jan 27, 2026

Multimodal Machine Learning Applications被引用 0

一句话总结

paper 介绍了多头显式注意力（MEA），通过头级线性组合（HLC）和分组归一化显式建模跨头交互，提高预训练收敛性，并在知识/科学任务上以最小性能损失实现 50% 的 KV 缓存内存缩减。它还统一了几种注意力变体并通过低秩重建支持 KV 缓存压缩。

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.

研究动机与目标

• 说明在 Transformer 中通过头际通信改善注意力性能的动机。
• 提出含头级线性组合的 MEA，以实现显式跨头交互。
• 通过分组归一化稳定训练并将 MEA 与 DFA、THA 等现有变体联系起来。
• 利用缩放定律引导选择更大学习率，实现更快的收敛。
• 展示通过低秩重构实现 KV 缓存的压缩而不造成重大性能损失。

提出的方法

• 定义头级线性组合（HLC），以在头之间混合信息。
• 通过用 HLC 混合版本替换 K 和 V 并对头输出应用分组归一化，提出多头显式注意力（MEA）。
• 提供统一视角，展示 DFA 和 THA 如何作为 MEA 的特例存在。
• 利用缩放定律高效选择学习率，并进行从头预训练比较。
• 提出通过低秩近似实现 KV 缓存压缩以将内存降低 50%。

实验结果

研究问题

• RQ1 与标准 Transformer 和其他跨头变体相比，MEA 是否在优化和最终性能方面具有提升？
• RQ2 分组归一化如何影响 MEA 的训练稳定性和表示多样性？
• RQ3 MEA 是否能在知识/科学任务上实现内存高效的 KV 缓存而不显著损失？
• RQ4 DFA 与 Talking-Heads 变体在统一理论视角下如何与 MEA 相关联？
• RQ5 在继续预训练后的复杂推理基准上，KV 缓存压缩的影响如何？

主要发现

Dataset	PIQA	OBQA	WinoGrande	HellaSwag	ARC-e	ARC-c	Avg.
Transformer	71.93	21.00	56.04	40.62	59.51	26.19	45.88
+GroupNorm	71.38	21.00	56.12	40.59	59.13	25.77	45.67
+DFA	71.76	22.20	54.38	41.29	60.69	27.82	46.36
Ours	73.18	19.80	54.14	42.02	61.57	27.65	46.39

• 在评估的变体中，含分组归一化的 MEA 达到最佳的平均下游性能。
• MEA 相较于基线在预训练阶段实现了更大的稳定学习率和更快的收敛。
• 在知识密集型和科学任务上，KV 缓存内存可降低 50%，且性能损失可以忽略；在完整压缩且可恢复的数学竞赛类基准上约损失 3.59%。
• DFA 和 THA 可以在 MEA 框架中解释，其中未使用分组归一化的 DFA 会退化回标准注意力。
• 分组归一化有助于维持跨头交互并稳定优化，使 MEA 优于缺乏归一化的变体。