跳到主要内容

自注意力 Q/K/V

为何需要三个独立投影矩阵、scaled dot-product 怎么防止 softmax 梯度消失

核心要点：

  • 三投影解掉简化版的相似性 / 对称性 / 角色混淆三个限制
  • 三个矩阵不是数学必要，是表达力最优解耦
  • Anthropic Circuits: head = QK circuit （注意力模式） + OV circuit （内容传递）
  • $\sqrt{d}$ 缩放：防 softmax 输入方差爆炸导致梯度消失
  • Llama / Qwen 全保留 $\sqrt{d}$，与 embedding 层 $\sqrt{h}$ 不同来源不同命运

名词定义​

本篇共享名词在 4.1 总览 已定义 (Scaled dot-product attention)。本篇新引入两个 mechanistic interp 术语：

名词定义
QK circuit单个 attention head 中，$W_Q^\top W_K$ 矩阵作为整体决定"位置间如何相互注意"的子电路 （Anthropic Transformer Circuits 框架）
OV circuit单个 attention head 中，$W_O W_V$ 矩阵作为整体决定"被注意后传递什么内容"的子电路；与 QK circuit 近似独立
Induction head实现 in-context learning 的关键 head: Q 查询"当前 token 的身份", K 用跨层组合编码"某位置前一个 token", V 编码"该位置自身内容" (Olsson 2022)

@tbl-qkv-glossary 本篇新引入名词

为什么必须三个独立的 Q/K/V?​

核心问题：简化自注意力 (02 篇) 已经能跑，引入 $W_Q, W_K, W_V$ 三个投影矩阵真的必要？一个矩阵或两个不行吗？

三个矩阵不是数学绝对必要——已有论文证明两个甚至一个矩阵也能跑——但三个是表达力的最优解耦。Vaswani 2017 §3.2 用 "compatibility function" 描述 Q/K 的关系，用 "weighted sum" 描述 V 的载荷，但没有解释为什么要三个独立矩阵。下文从语义、电路、实证三个角度回答这个问题。

数据库类比：三个角色不同​

Q / K / V 的命名直接来自数据库 / 字典 lookup：

  • Query (Q)：检索词，表达"我在找什么样的信息"
  • Key (K)：索引键，表达"我提供什么样的索引方便被查"
  • Value (V)：实际内容，表达"被查到后我要传递什么"

绑定后一个向量无法同时担三个角色。简化自注意力把这三件事都压在 $\mathbf{x}_i$ 一个向量里，模型只能学一个混合表示，没有自由度让"查询模式" 和"内容传递" 各自最优。

Anthropic Transformer Circuits：两条独立路径​

更深的解释来自 Anthropic mechanistic interpretability 工作。Elhage et al. 2021[1] 把每个 attention head 分解成两个近似独立的子电路：

  • QK circuit: $W_Q^\top W_K$ 作为整体，决定"位置 $i$ 注意位置 $j$ 多大权重"，即 attention pattern
  • OV circuit: $W_O W_V$ 作为整体 ($W_O$ 是多头拼接后的输出投影，见 05-多头注意力)，决定"被注意后传递什么内容"

两条路径近似独立：改 QK circuit 影响"看哪里"，改 OV circuit 影响"看到后做什么"，互不干扰。如果强制 K = V （两个矩阵），QK 和 OV 就耦合，模型失去这种独立优化的自由度。

Induction head 实测：三个角色携带不同信息​

Olsson et al. 2022[2] 在 in-context learning 关键的 induction head 上做实测：

  • Q 编码 "当前 token 的 identity"，用于跨上下文匹配
  • K 编码 "某位置的前一个 token 是什么" （通过跨层 K-composition 传入）
  • V 编码 "该位置自身内容"，用于实际复制

三者携带的信息完全不同，强制 Q=K 会破坏 induction head——这是大模型 in-context learning 能力的电路基础。

实证：两个或一个矩阵真的够吗？​

学术上有人测过用更少矩阵：

配置实证结果来源
Q = K （共享）ppl 4.32 → 4.34 （差 0.02），性能几乎持平，但 attention 变对称，方向性丢Shteingart 2024
单矩阵 + 对角缩放参数减 66%, GLUE 差 0.05%，噪声下反而更鲁棒arXiv 2412.00359[3]
K/V 更核心，Q 可简化Q 独立投影可压缩arXiv 2510.23912[4]

@tbl-qkv-reduced 减少 Q/K/V 矩阵数的实证结果

这些工作说明：数学上 Q/K/V 不一定要全分开，但工业大模型仍然全分开，因为表达力和优化自由度的代价 （即使只有 ppl 0.02 的差） 不值得为了省 1/3 参数承受。

DeepSeek MLA：不是合并而是低秩共享​

DeepSeek-V2/V3 的 Multi-head Latent Attention (MLA) 不是 K=V 合并，而是把 K 和 V 各自从一个共享的 low-rank latent （维度 512） 重建[5]。这样 KV cache 压缩到原来 1.4% （节省 98.6%），同时保留 K 和 V 各自独立的语义表示。这是 KV 架构压缩方向的代表，详见 knowledge/03-长上下文/05-kv-cache架构压缩。

Scaled dot-product attention 长什么样？​

核心问题：引入 $W_Q, W_K, W_V$ 后，完整的 attention 公式是什么？为什么要在 softmax 前除以 $\sqrt{d}$?

Vaswani 2017 的 scaled dot-product attention 完整公式：

$$\begin{equation} \mathrm{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \mathrm{softmax}\left(\frac{\mathbf{Q} \mathbf{K}^\top}{\sqrt{d_k}}\right) \mathbf{V} \label{eq:qkv-scaled-attn} \end{equation}$$

其中 $\mathbf{Q}, \mathbf{K}, \mathbf{V}$ 是输入 $\mathbf{X}$ 经投影得到：

$$\begin{equation} \mathbf{Q} = \mathbf{X} W_Q, \quad \mathbf{K} = \mathbf{X} W_K, \quad \mathbf{V} = \mathbf{X} W_V \label{eq:qkv-projection} \end{equation}$$

形状：$\mathbf{X} \in \mathbb{R}^{s \times h}$, $W_Q, W_K \in \mathbb{R}^{h \times d_k}$, $W_V \in \mathbb{R}^{h \times d_v}$，输出 $\mathbb{R}^{s \times d_v}$。通常 $d_k = d_v = d_{\text{head}}$，多头时整体 $h = n_{\text{head}} \cdot d_{\text{head}}$ (05-多头 展开)。

$\sqrt{d}$ 缩放的数学动机​

Vaswani §3.2.1 原话 （脚注 4 给出推导）：

"We suspect that for large values of $d_k$, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients."

完整推导：假设 $\mathbf{q}, \mathbf{k}$ 各分量独立同分布 $\mathcal{N}(0, 1)$，则

$$\begin{equation} \mathrm{Var}(\mathbf{q} \cdot \mathbf{k}) = \mathrm{Var}\left(\sum_{i=1}^{d_k} q_i k_i\right) = \sum_{i=1}^{d_k} \mathrm{Var}(q_i k_i) = d_k \label{eq:qkv-variance} \end{equation}$$

标准差 $\sqrt{d_k}$，除以 $\sqrt{d_k}$ 把方差归一回 1，与 $d_k$ 无关。

不缩放的具体后果​

$d_k = 64$，不缩放时 softmax 输入标准差是 8，即典型值落在 $[-8, 8]$ 范围。

softmax 在这种输入下接近 one-hot 分布：最大那一项接近 1，其他接近 0。梯度：softmax 输出 $p_i$ 对输入 $\omega_i$ 的偏导是 $p_i (1 - p_i)$，当 $p_i$ 接近 0 或 1 时这个值接近 0——梯度消失，训练崩。

除以 $\sqrt{d_k}$ 后输入方差归 1, softmax 在合理范围内，梯度健康。

现代 LLM 全保留 $\sqrt{d}$​

模型Forward 是否除 $\sqrt{d_{\text{head}}}$
nanoGPT (Karpathy)是：att = (q @ k.T) * (1.0 / math.sqrt(k.size(-1)))
Llama 3是：scores = torch.matmul(xq, keys.transpose(2,3)) / math.sqrt(self.head_dim)
Qwen2是：self.scaling = self.head_dim**-0.5; attn_weights = matmul(q,k) * scaling
所有主流 LLM是

@tbl-qkv-scaling 现代 LLM 均保留 $\sqrt{d}$ 缩放 (与 embedding 层 $\sqrt{h}$ 缩放命运相反)

与 embedding 层 $\sqrt{h}$ 缩放的本质区别 (详见 03-token-embedding):

  • Attention $\sqrt{d_k}$：有严格数学依据 （方差归一），Vaswani 给出推导，现代全保留
  • Embedding $\sqrt{h}$: Vaswani §3.4 一句话无解释，推测是匹配 PE 量级；RoPE 取代 absolute PE 后动机失效，Llama / Qwen 全线不做

这是"老论文里的两行代码命运不同"的有趣对照：一行有数学依据被保留，另一行无理由被悄悄淘汰。

Q/K/V 投影的工程实现​

核心问题：实际 PyTorch / Llama 代码里 $W_Q, W_K, W_V$ 怎么实现？单个矩阵还是 fused?

主流实现是三个独立 nn.Linear，输出维度对应 hidden_size 或 $n_{\text{head}} \cdot d_{\text{head}}$:

# Llama 风格 (multi-head, 含 GQA 支持)
self.wq = nn.Linear(h, n_head * d_head, bias=False)
self.wk = nn.Linear(h, n_kv_head * d_head, bias=False)   # GQA: n_kv_head <= n_head
self.wv = nn.Linear(h, n_kv_head * d_head, bias=False)
self.wo = nn.Linear(n_head * d_head, h, bias=False)       # 输出投影 (多头拼接后)

def forward(x):
    bsz, seqlen, _ = x.shape
    q = self.wq(x).view(bsz, seqlen, n_head, d_head).transpose(1, 2)
    k = self.wk(x).view(bsz, seqlen, n_kv_head, d_head).transpose(1, 2)
    v = self.wv(x).view(bsz, seqlen, n_kv_head, d_head).transpose(1, 2)
    # 后续 RoPE + causal mask + scaled dot-product
    ...

几个工程细节：

  • bias=False：现代 LLM 通常去掉 attention 投影的 bias，训练更稳，推理略快
  • 单矩阵 split：有些实现把 $W_Q, W_K, W_V$ 融合成 $W_{QKV} \in \mathbb{R}^{h \times 3h}$ 一次矩阵乘后 split 三块，节省 kernel launch
  • GQA 下 $W_K, W_V$ 输出维度小于 $W_Q$: Llama 3 8B 配置 n_head=32, n_kv_head=8，即 $W_Q$ 输出 $32 \cdot 128 = 4096$, $W_K, W_V$ 输出 $8 \cdot 128 = 1024$。GQA 细节归 05-多头与长上下文章

投影维度与多头的关系预告​

单头时 $d_k = d_v = h$，多头时 $d_{\text{head}} = h / n_{\text{head}}$。Vaswani 2017 选 $n_{\text{head}} = 8, h = 512$ 得 $d_{\text{head}} = 64$。Llama 3 8B 选 $n_{\text{head}} = 32, h = 4096$ 得 $d_{\text{head}} = 128$。

这一切如何分到多个 head，为什么多头比单头表达力更强，留给 05-多头注意力 展开。

从简化版到 Q/K/V 解掉了什么？​

核心问题：教学逻辑上，引入 Q/K/V 投影具体解掉了简化版的哪几个限制？

回顾 02-简化自注意力 列出的三个限制，Q/K/V 投影同时解掉:

限制 （简化版）Q/K/V 投影如何解掉
相似才关注 ($\omega_{ij} = \mathbf{x}_i \cdot \mathbf{x}_j$)$\omega_{ij} = (\mathbf{x}_i W_Q) \cdot (\mathbf{x}_j W_K)$, $W_Q$ 和 $W_K$ 可学不同投影，反义词或句法相关词可以学到高 score
对称性 ($\omega_{ij} = \omega_{ji}$)$\omega_{ij} = \mathbf{q}_i \cdot \mathbf{k}_j \neq \mathbf{q}_j \cdot \mathbf{k}_i$ (除非 $W_Q = W_K$，但实际是分开训练的不同矩阵)，方向性自然涌现
内容 / 查询 / 索引混淆$\mathbf{q}, \mathbf{k}, \mathbf{v}$ 各自一个空间，模型可以分别学"查询模式"、"被查索引"、"实际内容"

@tbl-qkv-fix 简化版三个限制如何被 Q/K/V 投影同时解掉

留下的下一步：因果掩码​

Q/K/V 自注意力是双向的——位置 $i$ 看所有 $j$ (包括未来 $j > i$)。这跟 BERT 这类 encoder-only 编码器适合，但与 GPT 的 next-token prediction 自回归矛盾：训练时如果第 $i$ 个 token 能看到第 $i+1$ 个，模型就直接"作弊"读答案。

下一篇 04-因果掩码 引入上三角 $-\infty$ mask 解掉这个问题，让 decoder-only 训练并行成立。

Takeaway​

知识点核心结论
三个矩阵不是数学必要Q=K 共享 ppl 差 0.02，单矩阵+对角缩放 GLUE 差 0.05%，但工业仍全分开
数据库类比Q = 检索词 / K = 索引键 / V = 实际内容
Anthropic Circuitshead = QK circuit （注意力模式） + OV circuit （内容传递），近似独立
Induction headQ 学 token identity, K 学 "某位置前一个 token", V 学 该位置内容
Scaled dot-product 公式$\mathrm{softmax}(QK^\top / \sqrt{d_k}) V$
$\sqrt{d}$ 数学动机$\mathrm{Var}(q \cdot k) = d_k$，除以 $\sqrt{d_k}$ 归一防 softmax 输入方差爆炸
不缩放后果$d_k = 64$ 时 std=8, softmax 几乎 one-hot，梯度 $p(1-p) \approx 0$，训练崩
现代 LLM 状态$\sqrt{d_k}$ 全保留 (与 embedding $\sqrt{h}$ 命运相反，后者已淘汰)
工程实现三个 nn.Linear(h, d_out), bias=False，有时 fused 成 $W_{QKV}$ 节省 kernel
解掉简化版三个限制学反义/句法 / 方向性 / 角色解耦
留下的问题仍是双向 attention，不兼容自回归——下一篇因果掩码

开放问题​

  • K=V 或更少矩阵能否在大模型上重复：现有实证都在小模型 (PTB / GLUE)，大模型 (8B+) 上 ppl 差距和下游能力如何，未充分研究
  • MLA 这类低秩共享是否会推广：DeepSeek MLA 把 K/V 压缩到 512 dim latent, KV cache 节省 98.6%，是否成为新主流仍待观察
  • Induction head 是否所有模型都涌现：Anthropic 在小到中等模型上观察到 induction head，大模型是否仍依赖同样机制，还是涌现了其他电路，是开放问题
  • $\sqrt{d}$ 的最优是否真的是平方根：数学上确实方差归一，但有论文测过 $d^{0.5 \pm \epsilon}$ 看是否其他指数更好，业界共识仍是平方根

延伸阅读​

参考资料​

  1. Elhage et al. A Mathematical Framework for Transformer Circuits. Anthropic, 2021. https://transformer-circuits.pub/2021/framework/index.html
  2. Olsson et al. In-context Learning and Induction Heads. Anthropic, 2022. https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html
  3. Does Self-Attention Need Separate Weights in Transformers?. 2024. https://arxiv.org/abs/2412.00359
  4. Key and Value Weights Are Probably All You Need. 2025. https://arxiv.org/abs/2510.23912
  5. DeepSeek-AI. DeepSeek-V3 Technical Report. 2024. https://arxiv.org/abs/2412.19437