穿越瓶颈：多头潜在注意力如何在语言模型中分离内容与位置

# 多头潜在注意力如何在语言模型中分离内容与位置
来源：https://arxiv.org/html/2607.23054
Dhruvil S,1Fenil Sojitra11footnotemark:1,2 Ravirajsinh Chauhan2 1印度理工学院马德拉斯分校 2P P Savani大学 23f2004889@ds\.study\.iitm\.ac\.in22se02it061@ppsu\.ac\.in raviraj\.chauhan@ppsu\.ac\.in同等贡献\. 代码和数据可在https://github.com/Dhruvil-sr24/Small-Language-Model-From-Scratch/tree/main/interp获取

###### 摘要

多头潜在注意力（MLA）由DeepSeek-V2引入，通过共享的低秩瓶颈（cKVc\_\{\\text\{KV\}\}）压缩键值对，在推理过程中实现了81%的KV缓存减少。尽管它已被用于拥有数千亿参数的生产模型中，但先前没有工作研究过*这个瓶颈保留或丢弃了什么信息*，也没有研究过它如何重塑Transformer的内部回路。我们开展了MLA的首次全面机械可解释性研究，训练了一个114M参数的MLA Transformer（在网页/代码/数学混合数据上预训练，在TinyStories上微调），并通过四个互补实验分析其内部表示：压缩矩阵的奇异值分解、注意力头分类、线性探针和cKVc\_\{\\text\{KV\}\}扰动归因分析。我们的主要发现是：(1) cKVc\_\{\\text\{KV\}\}瓶颈学习了一种纯内容表示，保留了实体身份（98%保留率），同时丢弃了位置信息（接近随机水平的准确率），验证了MLA通过RoPE在架构上分离内容与位置的设计；(2) 归纳头集中在单个层（第12层），这与它们在标准MHA中分散形成不同，表明共享瓶颈约束了回路拓扑；(3) 单个“语义枢纽”层（第15层）同时表现出最高的SVD有效秩和最强的扰动归因分数；(4) 该瓶颈全局过度配置，平均仅使用其128维容量的46%，这对异构秩分配有具体影响。这些发现来自一个在窄领域语料库上训练的114M参数模型，表明MLA不仅仅是被动地压缩注意力；它似乎重塑了模型组织内容、位置和回路结构的方式。我们将其视为初始数据点，而非对MLA在大规模下的定论，并在第5节（https://arxiv.org/html/2607.23054#S5）中详细说明了范围限制。

## 1 引言

多头潜在注意力（MLA）(DeepSeek-AI, 2024a, (https://arxiv.org/html/2607.23054#bib.bib8))已迅速成为高效大型语言模型中的关键架构创新。通过将键值（KV）表示压缩到一个共享的低秩瓶颈中，然后将其投影到每个头的键和值，MLA实现了KV缓存内存的大幅减少，相比标准多头注意力（MHA）最高可达81%，且不影响生成质量。该设计已被用于生产规模模型，包括DeepSeek-V2（236B参数）和DeepSeek-V3（671B参数）。

尽管 MLA 在实践中取得了成功，且机械可解释性领域不断发展（Elhage 等人，2021 (https://arxiv.org/html/2607.23054#bib.bib11)；Olsson 等人，2022 (https://arxiv.org/html/2607.23054#bib.bib17)；Conmy 等人，2023 (https://arxiv.org/html/2607.23054#bib.bib6)；Nanda 等人，2022 (https://arxiv.org/html/2607.23054#bib.bib16)），但先前没有工作研究过 MLA 瓶颈内部发生了什么。具体来说：

- •模型选择在压缩后的 cKVc\_\{\text\{KV\}\} 表示中*保留*了哪些语言信息？
- •它*丢弃*了哪些信息，为什么？
- •共享瓶颈如何影响已知回路（例如归纳头）的形成？
- •哪一层的瓶颈对模型的预测具有*因果*重要性？

我们通过对一个 114M 参数的 MLA Transformer 进行全面的机械研究来回答这些问题，该模型在网页文本、代码和数学文本的混合数据上预训练，然后在 TinyStories (Eldan & Li, 2023 (https://arxiv.org/html/2607.23054#bib.bib10)) 上微调（完整数据混合见第 3 节 (https://arxiv.org/html/2607.23054#S3)）。我们的四个实验（SVD 分析、注意力头分类、线性探针和 cKVc\_\{\text\{KV\}\} 扰动归因）汇聚成一幅关于 MLA 如何重塑 Transformer 信息处理的连贯图景。

#### 贡献。

1.  1.我们展示了 cKVc\_\{\text\{KV\}\} 学习了一种*仅内容*的表示：实体信息以 98% 的保真度保留，而位置信息被丢弃（探针准确率接近随机水平）。这验证了 MLA 分离内容（cKVc\_\{\text\{KV\}\}）与位置（RoPEkrop e\_\{\text\{rope\}\}）的设计。
2.  2.我们发现归纳头*集中在单个层*（在第 12 层检测到所有 5 个归纳头），这与它们在标准 MHA (Olsson 等人, 2022 (https://arxiv.org/html/2607.23054#bib.bib17)) 中分散形成形成对比。
3.  3.我们识别出第 15 层是一个*语义枢纽*，它同时具有最高的 SVD 有效秩（88/128）和所有层中最高的扰动归因分数。
4.  4.我们证明了该瓶颈是*全局过度配置*的：平均仅使用 128 维容量的 46%。我们提出异构秩分配作为具体的架构改进。

## 2 背景

### 2.1 多头潜在注意力

标准 MHA 从残差流 x∈Rd\\mathbf\{x\}\\in\\mathbb\{R\}^\{d\} 为每个头独立计算查询、键和值：

qh=WQhx, kh=WKhx, vh=WVhx\\mathbf\{q\}\_\{h\}=W\_\{Q\}^\{h\}\\mathbf\{x\},\\quad\\mathbf\{k\}\_\{h\}=W\_\{K\}^\{h\}\\mathbf\{x\},\\quad\\mathbf\{v\}\_\{h\}=W\_\{V\}^\{h\}\\mathbf\{x\} (1)
这需要在自回归生成期间为所有头缓存 kh\\mathbf\{k\}\_\{h\} 和 vh\\mathbf\{v\}\_\{h\}，每个 token 每层代价为 O(nheads·dhead)O\(n\_\{\\text\{heads\}\}\\cdot d\_\{\\text\{head\}\}\) 内存。

MLA 引入了一个*共享的 KV 压缩瓶颈*：

cKV\\displaystyle\\mathbf\{c\}\_\{\\text\{KV\}\}=RMSNorm(WDKV·x)∈RrKV\\displaystyle=\\text\{RMSNorm\}\(W\_\{\\text\{DKV\}\}\\cdot\\mathbf\{x\}\)\\in\\mathbb\{R\}^\{r\_\{\\text\{KV\}\}\} (2) khnope\\displaystyle\\mathbf\{k\}\_\{h\}^\{\\text\{nope\}\}=WUKh·cKV, vh=WUVh·cKV\\displaystyle=W\_\{UK\}^\{h\}\\cdot\\mathbf\{c\}\_\{\\text\{KV\}\},\\quad\\mathbf\{v\}\_\{h\}=W\_\{UV\}^\{h\}\\cdot\\mathbf\{c\}\_\{\\text\{KV\}\} (3) khrope\\displaystyle\\mathbf\{k\}\_\{h\}^\{\\text\{rope\}\}=RoPE(WKRh·x)\\displaystyle=\\text\{RoPE\}\(W\_\{KR\}^\{h\}\\cdot\\mathbf\{x\}\) (4) kh\\displaystyle\\mathbf\{k\}\_\{h\}=\[khnope; khrope\]\\displaystyle=\[\\mathbf\{k\}\_\{h\}^\{\\text\{nope\}\};\\mathbf\{k\}\_\{h\}^\{\\text\{rope\}\}\] (5)
在推理时，只需要缓存 cKV∈RrKV\\mathbf\{c\}\_\{\\text\{KV\}\}\\in\\mathbb\{R\}^\{r\_\{\\text\{KV\}\}\}（而不是每个头的 K/V），实现压缩比 rKV/(nheads·dhead)r\_\{\\text\{KV\}\}/\(n\_\{\\text\{heads\}\}\\cdot d\_\{\\text\{head\}\}\)。

关键的是，公式 2 (https://arxiv.org/html/2607.23054#S2.E2) 和 4 (https://arxiv.org/html/2607.23054#S2.E4) 揭示了一种架构分离：*内容信息*通过 cKV\\mathbf\{c\}\_\{\\text\{KV\}\}（跨头共享）流动，而*位置信息*通过独立的 RoPE 路径 khrope\\mathbf\{k\}\_\{h\}^\{\\text\{rope\}\}（直接从 x\\mathbf\{x\} 出发，绕过瓶颈）流动。模型是否实际学会利用这种分离是一个经验问题，我们在第 4.3 节 (https://arxiv.org/html/2607.23054#S4.SS3) 中回答。

### 2.2 机械可解释性

我们建立在已有的机械可解释性技术之上：

- •**线性探针** (Belinkov 等人, 2017 (https://arxiv.org/html/2607.23054#bib.bib3); Alain & Bengio, 2016 (https://arxiv.org/html/2607.23054#bib.bib2))：在中间表示上训练线性分类器，测试它们编码了什么信息：如果线性探针达到高准确率，则该信息在表示中是线性可访问的。
- •**激活补丁** (Meng 等人, 2022 (https://arxiv.org/html/2607.23054#bib.bib15); Wang 等人, 2023 (https://arxiv.org/html/2607.23054#bib.bib21))：这种技术的经典版本会破坏输入，然后在外科手术般精确的（层，位置）对上恢复干净激活，测量输出恢复了多少，从而建立真正的因果（而不仅仅是相关）重要性。在第 4.4 节 (https://arxiv.org/html/2607.23054#S4.SS4) 中，我们使用一种更轻量级的变体，即 cKVc\_\{\\text\{KV\}\} 扰动归因，它测量扰动*改变*每个（层，位置）处 cKVc\_\{\\text\{KV\}\} 的程度，并根据由此产生的输出概率下降对其进行重新加权，而不执行恢复步骤。这计算起来更便宜，但建立的是相关重要性，而非严格的因果重要性；我们会在重要之处明确说明这一区别（第 5.4 节 (https://arxiv.org/html/2607.23054#S5.SS4)）。
- •**归纳头分析** (Olsson 等人, 2022 (https://arxiv.org/html/2607.23054#bib.bib17))：识别实现“归纳”回路（复制之前跟随匹配前缀的 token）的注意力头，这是一种基本的上下文学习机制。

## 3 实验设置

### 3.1 模型架构

我们的模型是一个 24 层 MLA Transformer，配置如下：

表 1：模型架构细节。
### 3.2 训练

该模型在混合语料上预训练了 16,500 步：60% FineWeb-Edu (Penedo 等人, 2024 (https://arxiv.org/html/2607.23054#bib.bib19))（通用网页文本，用于基础语言和语法），25% The Stack (Kocetkov 等人, 2022 (https://arxiv.org/html/2607.23054#bib.bib14))（源代码，用于编程逻辑和结构化推理），以及 15% OpenWebMath (Paster 等人, 2023 (https://arxiv.org/html/2607.23054#bib.bib18))（数学网页文本，用于入门级数学语言）。然后在 TinyStories (Eldan & Li, 2023 (https://arxiv.org/html/2607.23054#bib.bib10)) 上微调了 7,200 步，在 TinyStories 验证集上达到了 4.3 的困惑度。本文中的所有可解释性分析均在*微调后*的检查点上运行，除非另有说明，否则使用 TinyStories 分布的段落作为分析输入。因此，虽然模型的表示是由比仅 TinyStories 更广泛、更多样化的预训练混合数据（网页文本、代码和数学）塑造的，但探针、注意力分类和扰动归因结果反映了模型在专门处理更简单的 TinyStories 领域时的行为。我们使用词汇量为 32,768 的 BPE 分词器。

### 3.3 实验概述

我们进行四个实验，每个实验探索 MLA 瓶颈的不同方面：

1.  **SVD 分析**（第 4.1 (https://arxiv.org/html/2607.23054#S4.SS1) 节）：分解压缩矩阵 WDKVW\_\{\\text\{DKV\}\} 和 WDQW\_\{\\text\{DQ\}\} 以测量有效维度。
2.  **注意力模式**（第 4.2 (https://arxiv.org/html/2607.23054#S4.SS2) 节）：将所有 24×8=19224\\times 8=192 个头按功能分类（前一个 token、归纳、BOS 沉没、局部窗口）。
3.  **线性探针**（第 4.3 (https://arxiv.org/html/2607.23054#S4.SS3) 节）：比较从 cKV\\mathbf\{c\}\_\{\\text\{KV\}\}（128 维）与残差流（512 维）中线性可解码的语言特征。
4.  **cKVc\_\{\\text\{KV\}\} 扰动归因**（第 4.4 (https://arxiv.org/html/2607.23054#S4.SS4) 节）：测量每个层扰动引起的 cKV\\mathbf\{c\}\_\{\\text\{KV\}\} 变化与预测退化之间的相关强度，作为因果重要性的轻量级代理。

所有实验均为推理阶段，在单个 NVIDIA T4 GPU 上运行。

## 4 结果

### 4.1 SVD 分析：瓶颈过度配置

我们在 24 个层中的每一层计算 WDKV∈R128×512W\_\{\\text\{DKV\}\}\\in\\mathbb\{R\}^\{128\\times 512\} 的 SVD，并测量*有效秩*：捕获 99% 总能量（∑σi2\\sum\\sigma\_\{i\}^\{2\}）所需的奇异值数量。

表 2：KV 和 Q 压缩矩阵的有效秩。#### 关键发现：

平均而言，该模型仅使用其 128 维 KV 瓶颈容量的 **46%**。在 90% 能量阈值下，平均只需要 17 个维度——表明存在显著的过度配置。

#### 层异构性：

三个层表现出明显更高的有效秩：第 7 层（85）、第 11 层（79）和第 15 层（88）。这些“秩尖峰”对应于架构阶段的转变，正如我们在后续实验中所展示的。

#### KV 和 Q 秩的反相关：

当 WDKVW\_\{\\text\{KV\}\} 使用更多维度时（例如，第 15 层：秩 88），WDQW\_\{\\text\{DQ\}\} 使用更少维度（第 15 层：秩 41），反之亦然。这表明模型学习了一种*非对称容量分配*：当需要更多关于位置*被*注意（KV）的信息时，关于位置*发出*注意（Q）的信息需求就更少。

参照图注图 1：左图：按层划分的 WDKVW\_\{\\text\{DKV\}\} 的 log 奇异值幅度，显示在每层大约前 20-30 个分量之后快速衰减。右图：每层在 90/95/99% 能量阈值下的有效秩，均远低于完整的 128 维容量（虚线）。参照图注图 2：不同层上 WDKVW\_\{\\text\{DKV\}\}（KV 压缩）与 WDQW\_\{\\text\{DQ\}\}（Q 压缩）的有效秩（99% 能量），展示了反相关模式。

### 4.2 注意力模式：归纳头共定位

我们使用在 50 个 TinyStories 段落（序列长度 64）上计算的四个指标对所有 192 个注意力头进行分类：(1) *前一个 token 分数*：注意力在位置 t−1t-1 上的比例；(2) *归纳分数*：注意力在匹配前缀后的位置上的比例；(3) *BOS 分数*：注意力在序列起始 token 上的比例；(4) *局部窗口分数*：注意力在 5-token 窗口内的比例。

表 3：顶部归纳头和前一个 token 头。#### 关键发现：

所有五个检测到的归纳头均*仅在第 12 层*。在标准 MHA 中，Olsson 等人 (2022 (https://arxiv.org/html/2607.23054#bib.bib17)) 发现在 2 层及以上的模型中，归纳头分布在多个层中。我们假设 MLA 的共享 cKV\\mathbf\{c\}\_\{\\text\{KV\}\} 瓶颈迫使归纳头共定位：由于归纳回路需要协调的 K-V 计算，并且 K 和 V 都来自相同的压缩表示，让所有归纳头在相同层的 cKV\\mathbf\{c\}\_\{\\text\{KV\}\} 上运行更为高效。

前一个 token 头实现了更简单的回路（注意位置 t−1t-1），如预期般在早期层（L1-L5）形成。这创建了一个清晰的*前一个 token* → *归纳* 管道，与两步归纳机制 (Olsson 等人, 2022 (https://arxiv.org/html/2607.23054#bib.bib17)) 一致，但压缩到了更少的层中。

参照图注图 3：每个层和头在所有四个头分类分数（前一个 token、归纳、BOS-注意力、局部窗口）上的表现。归纳分数面板显示在第 12 层有一个尖锐、孤立的带，其他任何层都没有可比的信号。参照图注图 4：所有 192 个头的前一个 token 分数与归纳分数对比，按层深着色。蓝绿色范围内的高归纳分数点完全对应于第 12 层的头。

### 4.3 线性探针：内容-位置分离

我们在 cKV\\mathbf\{c\}\_\{\\text\{KV\}\}（128 维）和残差流（512 维）上训练线性探针，以在 8 个层（0, 4, 8, 11, 12, 15, 20, 23）中的每一层分类五种语言属性。探针在 12,800 tok