第九便利

Transformer 的三个投影真的都需要吗？QKV 变体系统性研究

在 Transformer 架构中，注意力机制的 QKV（Query、Key、Value）三个投影矩阵一直被视为不可或缺的组成部分。但一个根本性的问题很少有人问：这三个投影真的都需要吗？每个投影对模型性能的独立贡献到底有多大？

来自以色列理工学院（Technion）等机构的研究者在一篇系统性的论文中，对这一问题给出了令人惊讶的答案。他们系统地研究了三种投影共享约束变体，揭示了注意力机制中一个长期被忽视的优化空间。

Transformer 中的 QKV（Query、Key、Value）三个投影矩阵是注意力机制的核心。但三个投影是否都必要？每个投影的独立贡献是什么？作者系统地研究了三种投影共享约束变体。

研究者定义了三种投影共享约束模式：

a) Q-K=V（共享 Key-Value）：Key 和 Value 使用同一投影矩阵
b) Q=K-V（共享 Query-Key）：Query 和 Key 使用同一投影矩阵
c) Q=K=V（单投影）：三个全部共享

值得注意的是，变体 b 和 c 会产生对称注意力图（symmetric attention map），因此需要引入 2D 位置编码来恢复注意力图中的方向性信息。

实验给出了清晰且令人意外的结论：

1. Q-K=V 是最佳变体。 仅损失 3.1% 的 perplexity，却实现了 50% 的 KV cache 缩减。这意味着在几乎不影响质量的前提下，推理时的显存占用直接减半。

2. 单投影（Q=K=V）令人惊讶地有效。 在多数任务上，完全共享三个投影的极端变体也能与标准 QKV 持平，说明注意力机制对投影冗余有很强的容忍度。

3. 投影共享与头共享（GQA/MQA）具有互补性。 这一发现极具实用价值：

— Q-K=V + GQA-4：KV cache 缩减达到 87.5%
— Q-K=V + MQA：KV cache 缩减达到惊人的 96.9%，使设备端推理（on-device inference）在内存受限的场景下首次成为可能

4. Q=K-V 表现最差。 共享 Query 和 Key 破坏了注意力的方向性，导致在所有任务中效果最差。

研究者在多种任务和规模上验证了结论的普适性：

— 合成任务：验证注意力的理论性质
— 视觉任务：MNIST、CIFAR、TinyImageNet 等标准图像分类基准
— 语言建模：分别在 300M 和 1.2B 参数规模的模型上，使用 10B tokens 进行训练和评估

实验还发现，Key 和 Value 可以占据相似的表示空间，注意力实际上在低秩空间中运行——这从理论上解释了为什么共享投影不会显著损害性能。

Q-K=V 本质上是注意力机制中被忽视的一种权重绑定（weight tying）实例。与标准 Transformer 中的嵌入层权重绑定类似，它直接带来可量化的推理内存优势。在 LLM 推理成本日益成为瓶颈的今天，这一发现对设备端部署和 KV cache 优化具有直接且实际的价值。

值得注意的是，Q-K=V 与现有的 KV cache 量化、剪枝等技术正交，意味着可以与它们叠加使用，进一步放大优化效果。