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

# Transformer 真的需要三个投影矩阵吗？QKV 变体的系统性研究
来源：https://arxiv.org/html/2606.04032

###### 摘要

Transformer 已成为各类 AI 任务的标准解决方案，其中查询、键和值（QKV）注意力机制居于核心地位。然而，这三个投影矩阵各自的贡献，以及省略其中某些矩阵所带来的影响，目前仍缺乏深入理解。我们系统评估了三种投影共享约束：a）Q-K=V（键值共享），b）Q=K-V（查询键共享），以及 c）Q=K=V（单一投影）。后两种变体会产生对称注意力图；为解决这一问题，我们还通过 2D 位置编码探索了非对称注意力机制。通过涵盖合成任务、视觉任务（MNIST、CIFAR、TinyImageNet、异常检测）以及语言建模（在 100 亿 token 上训练的 3 亿和 12 亿参数模型）的实验，我们发现所提出的 Transformer 变体性能与标准 QKV Transformer 相当，甚至偶尔更优。在语言建模任务中，Q-K=V 投影共享可实现 50% 的 KV 缓存减少，而困惑度仅上升 3.1%。关键在于，投影共享与头共享（GQA/MQA）具有互补性：将 Q-K=V 与 GQA-4 结合可实现 87.5% 的缓存减少，而 Q-K=V + MQA 可达到 96.9% 的缓存减少——从而支持实际的端侧推理。我们的研究表明，Q-K=V 能够保持质量，是因为键和值可以占据相似的表示空间，且注意力在低秩条件下运行；而 Q=K-V 则会破坏注意力的方向性。我们的结果将投影共享系统性地刻画为注意力机制中一种尚未充分探索的权重绑定形式，具有直接、可量化的推理内存收益——对边缘部署尤为有价值。代码已公开发布于 https://github.com/Brainchip-Inc/Do-Transformers-Need-3-Projections。

机器学习，ICML

## 1 引言

自诞生以来，Transformer（Vaswani 等人，2017 (https://arxiv.org/html/2606.04032#bib.bib3)）已从专用于语言的工具演变为多模态 AI 的骨干架构（Yin 等人，2024 (https://arxiv.org/html/2606.04032#bib.bib64)；Han 等人，2022 (https://arxiv.org/html/2606.04032#bib.bib35)）。然而，随着上下文窗口不断扩大以及对实时推理需求的增长，研究界逐渐将目光转向架构效率。从 Performer、Linformer 等线性复杂度模型，到 Ring Attention、分块计算等现代实现，各类高效变体纷纷涌现，旨在缓解自注意力的平方瓶颈（Tay 等人，2022 (https://arxiv.org/html/2606.04032#bib.bib4)）。

尽管如此，一个根本性的结构问题依然悬而未决：三分式的（Query, Key, Value）投影是否真正必要？卷积神经网络（CNN）（LeCun 等人，1995 (https://arxiv.org/html/2606.04032#bib.bib2)）和当代状态空间模型（SSM）（Gu and Dao，2023 (https://arxiv.org/html/2606.04032#bib.bib36)）通常采用更为统一的内部表示，而 Transformer 却在其投影矩阵中持续保留着冗余。为探究这一问题，我们提出并评估三种**投影共享**架构：

- **Q=K-V**：统一 Q 和 K；V 独立。
- **Q-K=V**：Q 独立；统一 K 和 V。
- **Q=K=V**：三者共用单一投影。

我们的研究发现，减少投影矩阵的数量能够显著降低参数量和计算开销，同时对下游性能的影响微乎其微。我们还观察到，这些简化方案的有效性具有任务依赖性：例如，**对称注意力**（Q=K）对图像分类等非时序任务高度有效，而时序任务则受益于保持一定程度的非对称性。

### 1.1 投影共享 vs. 头共享

我们的方法所针对的效率维度不同于当前业界标准，如 Ainslie 等人（2023 (https://arxiv.org/html/2606.04032#bib.bib65)）提出的**分组查询注意力（GQA）**和 Shazeer（2019 (https://arxiv.org/html/2606.04032#bib.bib88)）提出的**多查询注意力（MQA）**。GQA 和 MQA 通过在层内共享**头**来减小 KV 缓存规模，而我们的方法则共享**投影矩阵**本身。这两类策略彼此正交：将投影共享与头共享结合使用，可以在内存效率和吞吐量上实现复合收益。

### 1.2 主要贡献

- **系统性评估**：我们在 12 个多样化任务上对投影共享策略进行基准测试，涵盖合成推理、计算机视觉和大语言模型（LLM）预训练。
- **缓存优化**：我们证明 **Q-K=V** 配置可将 KV 缓存占用减少 **50%**，而对于 3 亿参数模型，困惑度仅增加 **3.1%**。
- **规模验证**：我们在 12 亿参数规模（约 100 亿 token）上验证了研究结论，确认相对质量排名在不同模型规模下保持稳定。在更大规模下，MQA 与 QKV 基本持平（困惑度仅上升 1.06%），同时提供 97% 的缓存减少。
- **架构协同**：我们证明投影共享与头共享严格互补。组合的 **Q-GQA-4** 配置可实现 **87.5%** 的缓存减少，而 **Q-MQA** 可达 **96.9%** 的减少。
- **深入洞察**：我们提供了架构层面的洞察，解释了为何 Q-K=V 有效（共享表示空间），而 Q=K-V 失败（破坏注意力方向性）。此外，我们证明在 QKV 折叠下，核化注意力可以化为纯递归形式，其中注意力状态通过外积更新演化，并由当前输入读取，使线性注意力成为具有自适应观测的状态空间模型的特例（附录 A.1 (https://arxiv.org/html/2606.04032#A1.SS1)）。

## 2 相关工作

### 2.1 背景：标准注意力机制

Transformer 架构（Vaswani 等人，2017 (https://arxiv.org/html/2606.04032#bib.bib3)）已成为现代深度学习在多个领域的基础，涵盖自然语言处理（Brown 等人，2020 (https://arxiv.org/html/2606.04032#bib.bib24)）、计算机视觉（Dosovitskiy 等人，2021 (https://arxiv.org/html/2606.04032#bib.bib7)）等众多方向。Transformer 块的核心由多个相互关联的组件构成：多头自注意力、逐位置前馈网络、层归一化（Ba 等人，2016 (https://arxiv.org/html/2606.04032#bib.bib26)）、残差连接（He 等人，2016 (https://arxiv.org/html/2606.04032#bib.bib6)）以及位置编码。

自注意力机制——亦称为内部注意力——是 Transformer 的核心创新。该机制使序列中的每个位置能够有选择性地从所有其他位置聚合信息，从而计算依赖上下文的表示。自注意力已在多种任务中展现出卓越的有效性，包括机器翻译、抽象摘要（Gupta and Gupta，2019 (https://arxiv.org/html/2606.04032#bib.bib67)）、视觉问答（Wu 等人，2017 (https://arxiv.org/html/2606.04032#bib.bib66)）、多模态理解（Radford 等人，2021 (https://arxiv.org/html/2606.04032#bib.bib15)）以及目标识别（Dosovitskiy 等人，2021 (https://arxiv.org/html/2606.04032#bib.bib7)）。

形式上，对于在输入 $X \in \mathbb{R}^{n \times d}$ 上运行的单个注意力头，注意力机制计算如下：

$$A_h = \text{Softmax}(\alpha Q_h K_h^T) V_h, \tag{1}$$

其中 $Q_h = XW_q$，$K_h = XW_k$，$V_h = XW_v$ 表示学习的线性投影，权重矩阵 $W_q, W_k, W_v \in \mathbb{R}^{d \times d_k}$。缩放因子 $\alpha = 1/\sqrt{d_k}$ 用于稳定训练过程中的梯度，其中 $d_k = d/H$，$H$ 表示注意力头的数量。softmax 操作按行应用以产生注意力权重。在多头注意力中，$H$ 个头并行计算注意力 $A_1, \ldots, A_H$，这些输出拼接后经过最终的线性变换。注意力分数 $QK^T$ 编码了词元对之间的亲和性，查询-键点积决定了哪些值对每个位置是相关的。

### 2.2 三个独立投影的必要性

尽管 QKV 公式已成为标准，其必要性仍是一个开放问题。与 CNN（LeCun 等人，1998 (https://arxiv.org/html/2606.04032#bib.bib28)）、RNN 或状态空间模型（Gu and Dao，2023 (https://arxiv.org/html/2606.04032#bib.bib36)）更为简约的表示不同，Transformer 为每个 token 维护三个不同的表示。近期研究已开始质疑这一设计：线性注意力（Katharopoulos 等人，2020 (https://arxiv.org/html/2606.04032#bib.bib69)）、核方法注意力（Choromanski 等人，2021 (https://arxiv.org/html/2606.04032#bib.bib82)）以及无注意力模型（Zhai 等人，2021 (https://arxiv.org/html/2606.04032#bib.bib70)）等方法表明，更简单的机制或许已经足够。然而，这些方法往往以牺牲标准注意力的灵活性为代价。

我们的工作采取了一种互补的思路：与其完全替换注意力机制，不如研究三个投影矩阵是否可以统一，同时保留核心注意力机制。我们最初在 Borji（2023 (https://arxiv.org/html/2606.04032#bib.bib94)）¹中提出了这一想法。随后，Kowsher 等人（2025 (https://arxiv.org/html/2606.04032#bib.bib95)）提出了类似的方案。另有若干工作与此相关（Fusco 等人，2022 (https://arxiv.org/html/2606.04032#bib.bib92)；Mai 等人，2023 (https://arxiv.org/html/2606.04032#bib.bib91)）。

DeepSeek-V2 的多头潜在注意力（MLA）（Liu 等人，2024 (https://arxiv.org/html/2606.04032#bib.bib98)）通过将 K 和 V 压缩为共享的潜在向量（在推理时缓存并展开）来减少 KV 缓存。与 Q-K=V 不同，K 和 V 在展开后在功能上仍保持独立——MLA 以增加投影参数为代价换取更丰富的压缩表示，而 Q-K=V 则通过简单的硬等式约束实现缓存减少。

## 3 我们的方法

图 1：我们提出的投影共享注意力变体。带有 2D 位置编码的注意力机制标记为 (X)+。

### 3.1 提出的投影共享注意力变体

我们系统地考察三种投影共享约束，这些约束逐步减少学习变换的数量（图 1 (https://arxiv.org/html/2606.04032#S3.F1.2)）。

**变体 1：Q=K-V。** 我们消除独立的查询投影，令 $Q = K$：

$$A = \text{Softmax}(\alpha KK^T) V. \tag{2}$$

该公式产生对称的注意力矩阵 $KK^T$。对称注意力已在图神经网络（Veličković 等人，2018 (https://arxiv.org/html/2606.04032#bib.bib71)）和关系推理（Santoro 等人，2017 (https://arxiv.org/html/2606.04032#bib.bib81)）的先前工作中有所探索，在这些场景中缺乏方向偏置是有益的。然而，对于需要因果依赖的时序任务，对称性可能构成限制。

为解决这一问题，我们引入 **(Q=K-V)+**，通过 2D 位置编码注入非对称性。我们首先构造固定的 2D 正弦位置编码 $P \in \mathbb{R}^{n \times n \times m}$（Vaswani 等人，2017 (https://arxiv.org/html/2606.04032#bib.bib3)）。将 $n \times n$ 的注意力图沿通道维度广播并与 $P$ 相加。为将结果张量映射回 2D 注意力矩阵，我们应用 $1 \times 1$ 卷积（等价于跨通道的线性投影）。该设计受到相对位置编码（Shaw 等人，2018 (https://arxiv.org/html/2606.04032#bib.bib83)；Huang 等人，2020 (https://arxiv.org/html/2606.04032#bib.bib12)）和视觉 Transformer 中 2D 位置嵌入（Dosovitskiy 等人，2021 (https://arxiv.org/html/2606.04032#bib.bib7)）的启发。完整构造见附录 A.2 (https://arxiv.org/html/2606.04032#A1.SS2)。

**变体 2：Q-K=V。** 我们统一键和值的投影，令 $V = K$：

$$A = \text{Softmax}(\alpha QK^T) K. \tag{3}$$

由于 $Q$ 和 $K$ 保持独立，该公式保留了非对称注意力图。键和值共享表示的约束可视为一种权重绑定（Press and Wolf，2017 (https://arxiv.org/html/2606.04032#bib.bib11)），这已被证明在语言建模中是有效的。

**变体 3：Q=K=V。** 最激进的简化方案对三个角色使用单一投影：

$$A = \text{Softmax}(\alpha KK^T) K. \tag{4}$$

该方案结合了变体一的对称注意力和变体二的表示瓶颈。我们还评估了 **(Q=K=V)+**，与第一个变体类似，通过添加 2D 位置编码来缓解对称性约束。

##### (X)+ 变体的适用范围

(X)+ 变体中的 2D 位置编码针对非因果场景（视觉任务、合成任务），在这些场景中 $Q=K$ 导致的对称注意力是主要限制。因果语言建模已通过因果掩码强制实现非对称性，因此 (X)+ 解决的问题在该场景中并不实质存在；我们因此仅在非因果任务上评估 (X)+（表 2 (https://arxiv.org/html/2606.04032#S4.T2) 和表 3 (https://arxiv.org/html/2606.04032#S4.T3)），并将其视为任务特定的启发式方法，而非通用增强手段。

### 3.2 将投影共享与头共享相结合

我们的投影共享方法在不同于近期头共享方法的维度上运作，从而支持复合优化。

**头共享机制。** 分组查询注意力（GQA）（Ainslie 等人，2023 (https://arxiv.org/html/2606.04032#bib.bib65)）和多查询注意力（MQA）（Shazeer，2019 (https://arxiv.org/html/2606.04032#bib.bib88)）通过在多个查询头之间共享键值头来减少内存占用。在 GQA-g 中，$H$ 个查询头关注 $g$（$j > i$）个头。

**与注意力的集成：** 给定原始注意力分数 $A = QK^\top \in \mathbb{R}^{n \times n}$，我们将 $A$ 沿通道维度广播，加入位置编码，然后应用 $1 \times 1$ 卷积（线性投影）将 $A' \in \mathbb{R}^{n \times n \times m}$ 映射回 $\mathbb{R}^{n \times n}$。

**直觉：** 这一操作使注意力能够将基于内容的相似性与位置/方向偏置相结合，打破投影共享带来的对称性，从而支持对顺序敏感的行为。

### A.3 额外的合成任务和视觉任务结果

#### A.3.1 合成任务结果

图 3 (https://arxiv.org/html/2606.04032#A1.F3) 展示了合成任务中随时间变化的损失曲线。图 4 (https://arxiv.org/html/2606.04032#A1.F4) 显示了注意力图样本。值得注意的是，KV（Q=K-V）Transformer 的注意力图在直线 $y = x$ 附近呈现对称性。注意力图中还可以观察到一些值得关注的模式。例如，在反转任务中，QKV 模型已学会关注位于自身翻转索引处的 token。然而，它也会将部分注意力分配给翻转索引附近的值。这种行为的出现是因为模型并不需要精确、严格的注意力来解决该问题，而是受益于近似的、带噪声的注意力图。图 5 (https://arxiv.org/html/2606.04032#A1.F5) 展示了用于计算和归一化自注意力图的代码，以及注意力图的可视化结果。

图 3：QKV、Q=K-V 和 (Q=K-V)+ 在合成任务中随时间变化的损失曲线。