NLP——MHA2MLA(Partial-RoPE)

注：本文包含 AI 辅助创作

参考链接：
- 原始论文：Towards Economical Inference: Enabling DeepSeek’s Multi-Head Latent Attention in Any Transformer-based LLMs, Fudan & ECNU & Shanghai Al Lab
- GitHub 开源链接：github.com/JT-Ushio/MHA2MLA

Paper Summary

整体总结：
- MHA2MLA 包含贡献感知的部分 RoPE 移除和 SVD 驱动的低秩投影
  - 论文实现了 KV 缓存的近乎无损压缩
- 本文创新：通过针对性的参数重用和数据高效的微调来实现大语言模型架构迁移的可行性
问题提出：
- 基于 MHA 及 MHA 的变体（GQA 等）的标准 LLM 在成本上存在显著劣势
DeepSeek 的解法：MLA
- Multi-head Latent Attention（MLA，多头潜在注意力）是 DeepSeek 提出的一种创新架构
- MLA 通过将键值（Key-Value, KV）缓存显著压缩为一个潜在向量，从而实现高效且经济的推理
本文核心解决的现实问题及挑战：
- 使训练良好的大语言模型（例如 Llama）能够快速适配 MLA 而无需从头预训练，既具有重要意义又充满挑战
论文提出了首个数据高效的精调方法，用于从 MHA 迁移至 MLA（MHA2MLA），该方法包含两个关键组件：
- 对于 部分旋转位置编码（partial-RoPE） ，论文从对注意力分数贡献较小的查询和键的维度中移除 RoPE
- 对于 低秩近似（low-rank approximation） ，论文基于预训练的键和值参数引入联合奇异值分解（SVD）近似
MHA2MLA 仅需使用一小部分（3% 到 6%）数据即可恢复性能，显著降低了推理成本，同时可与 KV 缓存量化等压缩技术无缝集成
- 举例：Llama2-7B 的 KV 缓存大小减少了 92.19%，而在 LongBench 上的性能仅下降 0.5%

Introduction and Discussion

LLM 的快速发展显著加速了通往通用人工智能（Artificial General Intelligence, AGI）的进程，模型能力随参数数量增加而呈现可预测的扩展 (2020)
- 然而，这些收益伴随着高昂的代价：训练的计算需求不断攀升，推理吞吐量下降，导致巨大的能源消耗和碳排放 (2019)
随着下游任务日益复杂，长上下文处理和高计算量的推理已成为大语言模型应用的核心 (2024)
- 一个关键瓶颈在于 MHA (2017) 机制固有的键值（Key-Value, KV）缓存的内存占用，其随序列长度和模型大小线性增长
- 为缓解此问题，研究者探索了诸如分组查询注意力（Grouped-Query Attention, GQA）(2023) 和多查询注意力（Multi-Query Attention, MQA）(2019) 等变体
- 然而，这些方法不仅减少了 KV 缓存大小，也减少了注意力中的参数数量，导致性能下降
DeepSeek 引入了 MLA (2024)
- MLA 是一种配备低秩键值联合压缩的注意力机制
- 经验上，MLA 实现了优于 MHA 的性能，同时显著减少了推理期间的 KV 缓存，从而提升了推理效率
一个关键但尚未探索的问题随之产生：原本为 MHA 良好训练的大语言模型能否适配 MLA 以进行推理？
- MHA 与 MLA 之间固有的架构差异使得零样本迁移不切实际，而从头预训练的惊人成本使得这种转变在技术上具有挑战性且在现有研究中探索不足
为填补这一空白，论文提出了首个精心设计的 MHA2MLA 框架，该框架最大化地复用了预训练 MHA 网络的参数 ，同时将 KV 缓存存储和推理过程与 MLA 的范式对齐（图 1）
论文的框架具有两项关键的技术创新：
- 部分旋转位置编码（partial rotary position embedding, partial RoPE）
- 低秩近似（low-rank approximation）
MHA2MLA 的主要目标是实现数据高效的性能恢复，即使用最少的精调数据来恢复由架构变更引起的能力下降
MLA 的推理加速机制与 RoPE 之间固有的不兼容性 necessitates 架构上的折衷
DeepSeek 的解决方案是在有限维度中保留位置编码（PEs）同时压缩其他维度，这需要在 MHA 中策略性地移除 RoPE 维度（将其转换为 NoPE）以实现与 MLA 的对齐
虽然更高的移除比率提升了压缩效率，但也加剧了性能下降，形成了效率与能力之间的权衡
通过系统性地探索 RoPE 移除策略，论文发现基于贡献的维度选择（保留按注意力分数影响排序的前 k 个维度）能最优地平衡这些竞争目标
尽管先前的研究已经探索了从头训练部分 RoPE 大语言模型 (2024; 2021)，但论文的工作开创了在大语言模型中进行全 RoPE 到部分 RoPE（Partial RoPE）转换的数据高效精调方法
MLA 通过将键和值投影到低秩潜在表示空间（存储在 KV 缓存中）来减少内存占用
MHA2MLA 也可以对剥离了 RoPE 的值和键（NoPE 维度）应用低秩近似
通过对对应于 NoPE 子空间的预训练参数矩阵 $\boldsymbol{W}_{v}$ 和 $\boldsymbol{W}_{k}$ 执行奇异值分解（Singular Value Decomposition, SVD），论文将这些组件压缩到一个潜在空间中，同时最大限度地保留原始模型学到的知识
论文的主要贡献是：
- 论文提出了 MHA2MLA，这是首个参数高效的精调框架，能够仅使用 $3%$ 到 $6%$ 的训练数据将预训练的基于 MHA 的大语言模型适配到 MLA 架构，而无需从头训练
- 论文证明了 MHA2MLA 架构可以与 KV 缓存量化技术集成，以实现更经济的推理（最高减少 96.87%）
- 论文在四种模型规模（从 135M 到 7B，涵盖 MHA 和 GQA）上进行了实验，并进行了详细的消融研究，为 MHA2MLA 提供了指导和见解

Preliminary

多头注意力机制 (Multi-Head Attention, MHA)

给定一个输入序列 $\{\boldsymbol{x}_{1},\ldots,\boldsymbol{x}_{l}\} \in \mathbb{R}^{l \times d}$，标准的 MHA (2017) 将每个 token $\boldsymbol{x}_{i}$ 投影为查询向量 $\boldsymbol{q}_{i}^{(h)} = \boldsymbol{x}_{i}\boldsymbol{W}_{q}^{(h)}$、键向量 $\boldsymbol{k}_{i}^{(h)} = \boldsymbol{x}_{i}\boldsymbol{W}_{k}^{(h)}$ 和值向量 $\boldsymbol{v}_{i}^{(h)} = \boldsymbol{x}_{i}\boldsymbol{W}_{v}^{(h)}$，其中对于每个头 $h \in \{1,\ldots,n_{h}\}$，有 $\boldsymbol{W}_{q}^{(h)}, \boldsymbol{W}_{k}^{(h)}, \boldsymbol{W}_{v}^{(h)} \in \mathbb{R}^{d \times d_{h} }$。旋转位置编码（Rotary Positional Encoding, RoPE）(2021) 被应用于查询和键（例如，$\boldsymbol{q}_{i,\text{rope} }^{(h)} = \text{RoPE}(\boldsymbol{q}_{i}^{(h)})$），随后进行缩放点积注意力计算：
$$
\boldsymbol{o}_{i}^{(h)} = \text{Softmax}\left( \boldsymbol{q}_{i,\text{rope} }^{(h)} \boldsymbol{k}_{\leq i,\text{rope} }^{(h)\top} \right) \boldsymbol{v}_{\leq i}^{(h)}, \\
\text{MHA}(\boldsymbol{x}_{i}) = \left[ \boldsymbol{o}_{i}^{(1)}, \ldots, \boldsymbol{o}_{i}^{(n_{h})} \right] \boldsymbol{W}_{o},
$$
- 其中 $\boldsymbol{W}_{o} \in \mathbb{R}^{(n_{h}d_{h}) \times d}$
- $[\cdot,\cdot]$ 表示向量拼接
- 注：为简化符号，此处忽略了 $\frac{1}{\sqrt{d} }$ 缩放因子
在自回归推理过程中，MHA 需要存储大小为 $O(2ln_{h}d_{h})$ 的键值（KV）缓存 $\{\boldsymbol{k}_{\text{rope} }^{(h)}, \boldsymbol{v}^{(h)}\}_{h=1}^{n_{h} }$，该大小随序列长度 $l$ 线性增长，造成了内存瓶颈

分组查询注意力 (Grouped-Query Attention, GQA)

GQA (2023) 通过在 $n_{g}$ 个组（$n_{g} \ll n_{h}$）之间共享键/值来减少 KV 缓存。对于每个头 $h$，它映射到组 $g = \lfloor h / n_{g} \rfloor$：
$$
\boldsymbol{o}_{i}^{(h)} = \text{Softmax}\left( \boldsymbol{q}_{i,\text{rope} }^{(h)} \boldsymbol{k}_{\leq i,\text{rope} }^{(g)\top} \right) \boldsymbol{v}_{\leq i}^{(g)}, \\
\text{GQA}(\boldsymbol{x}_{i}) = \begin{bmatrix} \boldsymbol{o}_{i}^{(1)}, \ldots, \boldsymbol{o}_{i}^{(n_{h})} \end{bmatrix} \boldsymbol{W}_{o}. \tag{2}
$$

多头查询注意力 (Multi-Query Attention, MQA)

MQA (2019) 是 GQA 的一个特例，其中 $n_{g} = 1$，即所有头共享一个全局的键/值
虽然 GQA 和 MQA 方法将 KV 缓存减少到 $O(2ln_{g}d_{h})$，但由于参数剪枝，它们会导致性能下降

多头潜在注意力 (Multi-Head Latent Attention, MLA)

MLA (DeepSeek-AI, 2024) 引入了一种混合架构，将位置编码（PE）与潜在 KV 压缩解耦
对于每个头 $h$，输入 $\boldsymbol{x}_{i}$ 被投影为两个互补的分量：
位置感知分量 (Position-Aware Component) ：一部分维度保留 PE 以保持位置敏感性：
$$
\boldsymbol{q}_{i,\text{rope} }^{(h)}, \boldsymbol{k}_{i,\text{rope} } = \text{RoPE}\left( \boldsymbol{x}_{i}\boldsymbol{W}_{dq}\boldsymbol{W}_{qr}^{(h)}, \boldsymbol{x}_{i}\boldsymbol{W}_{kr} \right),
$$
- 其中 $\boldsymbol{W}_{dq} \in \mathbb{R}^{d \times d_{q} }$，$\boldsymbol{W}_{qr}^{(h)} \in \mathbb{R}^{d_{q} \times d_{r} }$，$\boldsymbol{W}_{kr} \in \mathbb{R}^{d \times d_{r} }$ 将查询/键投影到保留 RoPE 的 $d_{r}$ 维分量中
位置无关分量 (Position-Agnostic Component) ：剩余的 $d_{c}$ 个维度被移除 PE（即 NoPE），并将 $\boldsymbol{k}_{i,\text{nope} }^{(h)}$ 和 $\boldsymbol{v}_{i}^{(h)}$ 压缩成一个共享的潜在向量 $\boldsymbol{c}_{i,kv}^{(h)}$：
$$
\boldsymbol{q}_{i,\text{nope} }^{(h)} = \boldsymbol{x}_{i}\boldsymbol{W}_{dq}\boldsymbol{W}_{qc}^{(h)}, \\
\boldsymbol{c}_{i,kv} = \boldsymbol{x}_{i}\boldsymbol{W}_{dkv}, \\
\boldsymbol{k}_{i,\text{nope} }^{(h)}, \boldsymbol{v}_{i}^{(h)} = \boldsymbol{c}_{i,kv}\boldsymbol{W}_{uk}^{(h)}, \boldsymbol{c}_{i,kv}\boldsymbol{W}_{uv}^{(h)},
$$
- 其中 $\boldsymbol{W}_{qc}^{(h)} \in \mathbb{R}^{d_{q} \times d_{c} }$，$\boldsymbol{W}_{dkv} \in \mathbb{R}^{d \times d_{kv} }$，$\boldsymbol{W}_{uk}^{(h)} \in \mathbb{R}^{d_{kv} \times d_{c} }$，$\boldsymbol{W}_{uv}^{(h)} \in \mathbb{R}^{d_{kv} \times d_{h} }$
注意 $d_{r} + d_{c} = d_{h}$。MLA 的注意力输出结合了两个分量：
$$
\boldsymbol{o}_{i}^{(h)} = \text{Softmax}\left( \boldsymbol{q}_{i,\text{rope} }^{(h)} \boldsymbol{k}_{\leq i,\text{rope} }^{(h)\top} + \boldsymbol{q}_{i,\text{nope} } \boldsymbol{k}_{\leq i,\text{nope} }^{(h)\top} \right) \cdot \boldsymbol{v}_{\leq i}^{(h)} \\
\text{MLA}(\boldsymbol{x}_{i}) = \begin{bmatrix} \boldsymbol{o}_{i}^{(1)}, \ldots, \boldsymbol{o}_{i}^{(n_{h})} \end{bmatrix} \cdot \boldsymbol{W}_{o}. \tag{3}
$$
与 MHA 及其变体不同，MLA 存储潜在向量 $\boldsymbol{c}_{kv}$ 和 $\boldsymbol{k}_{i,\text{rope} }^{(h)}$（$\mathcal{O}(ld_{r} + ld_{kv})$）而不是全秩的 $\boldsymbol{k}_{i}$, $\boldsymbol{v}_{i}$（$\mathcal{O}(2ln_{h}d_{h})$），其中 $(d_{r} + d_{kv}) \ll 2n_{h}d_{h}$
为什么 MLA 需要分离 RoPE 和 NoPE？
- MLA 在推理过程中对 NoPE 部分引入了矩阵合并技术，有效减少了内存使用
- 对于点积操作 $\boldsymbol{q}_{i,\text{nope} }^{(h)} \boldsymbol{k}_{j,\text{nope} }^{(h)\top}$，可以应用以下恒等变换：
  $$
  \boldsymbol{q}_{i,\text{nope} } \boldsymbol{k}_{j,\text{nope} }^{\top} = (\boldsymbol{x}_{i}\boldsymbol{W}_{dq}\boldsymbol{W}_{qc}) (\boldsymbol{c}_{j,kv}\boldsymbol{W}_{uk})^{\top} = \boldsymbol{x}_{i} (\boldsymbol{W}_{dq}\boldsymbol{W}_{qc}\boldsymbol{W}_{uk}^{\top}) \boldsymbol{c}_{j,kv}^{\top}
  $$
  - 注：为简化符号，论文省略了上标 ${}^{(h)}$。矩阵 $\boldsymbol{W}_{uv}$ 和 $\boldsymbol{W}_{o}$ 也可以合并，请参阅 DeepSeek-AI 等人 (2024) 的附录 C
  - 其中 $(\boldsymbol{W}_{dq}\boldsymbol{W}_{qc}\boldsymbol{W}_{uk}^{\top})$ 可以预先合并为单个矩阵，而 $\boldsymbol{c}_{j,kv}$ 已经存储在 KV 缓存中。对于 RoPE 部分，RoPE() 函数将输入向量乘以旋转矩阵（例如，$\text{RoPE}(\boldsymbol{q}_{i}) = \boldsymbol{q}_{i}\boldsymbol{R}_{i}$，$\boldsymbol{R}_{i}$ 的具体形式将在第 3.1 节介绍）
因此，恒等变换变为：
$$
\boldsymbol{q}_{i,\text{rope} } \boldsymbol{k}_{j,\text{rope} }^{\top} = (\boldsymbol{x}_{i}\boldsymbol{W}_{dq}\boldsymbol{W}_{qr} \boldsymbol{R}_{i}) (\boldsymbol{x}_{j}\boldsymbol{W}_{kr} \boldsymbol{R}_{j})^{\top} = \boldsymbol{x}_{i} (\boldsymbol{W}_{dq}\boldsymbol{W}_{qr} \boldsymbol{R}_{j-i} \boldsymbol{W}_{kr}^{\top}) \boldsymbol{x}_{j}^{\top}
$$
由于 $(\boldsymbol{W}_{dq}\boldsymbol{W}_{qr} \boldsymbol{R}_{j-i} \boldsymbol{W}_{kr}^{\top})$ 与相对位置 $j-i$ 相关，它不能被合并成一个固定矩阵。考虑到 LLM 中的相对距离可能非常长（例如 128K），RoPE 部分更适合使用原始形式进行计算

MHA2MLA

部分旋转位置编码（Partial-RoPE）

为实现从标准 MHA 到 MLA 的迁移，论文提出了部分旋转位置编码微调（partial-RoPE finetuning）策略，该策略从目标比例维度中移除 RoPE，并将其转换为 NoPE
关键的是，尽管先前的工作已经探索了从头开始训练具有部分 RoPE 的 LLM（实现了比全 RoPE 略好的困惑度 (2021; 2024)），但现有方法均未解决如何高效地将预训练的全 RoPE 模型（例如 Llama）适配到部分 RoPE，而无需昂贵的重新训练
论文的工作通过系统评估部分 RoPE 的变体，以确定最数据高效的微调方案来恢复模型在适配后的性能，从而弥补了这一空白
RoPE 的数学形式化表示如下：
- 对于维度为 $d_h$ 的查询或键向量，RoPE 将向量划分为 $\frac{d_h}{2}$ 个子空间，每个子空间包含两个连续维度（例如，第 $k$ 个子空间包含维度 $2k$ 和 $2k+1$）
- 每个子空间以不同的旋转速度（频率）旋转，其中第 $k$ 个子空间的旋转角为 $\theta_k = b^{-2k/d_h}$，$b$ 是预定义的基数（例如，Llama 使用 $b=10000$）
- 因此，对查询和键应用 RoPE 变为：
  $$
  \boldsymbol{q}_{i,rope} =\left[\boldsymbol{R}_{i}^{[2k,2k+1]}(\theta_{k})\boldsymbol{q}_{i }^{[2k,2k+1]}\right]_{0\leq k < \frac{d_{h} }{2} }, \\
  \boldsymbol{k}_{i,rope} =\left[\boldsymbol{R}_{i}^{[2k,2k+1]}(\theta_{k})\boldsymbol{k}_{i }^{[2k,2k+1]}\right]_{0\leq k < \frac{d_{h} }{2} }.
  $$

全 RoPE 到部分 RoPE 的策略（Full-RoPE to Partial-RoPE Strategies）

给定保留的旋转子空间数量 $r$（$r=\frac{d_r}{2} \ll$ 总子空间数 $\frac{d_h}{2}$），论文提出了四种策略（如图 2 所示）来选择哪些 $r$ 个子空间保留 RoPE 编码
高频保留（High-Frequency Preservation） 保留 $r$ 个旋转最快（高频）的子空间：
$$
\mathcal{S}_{\text{high} }=\left\{k,|,0\leq k<r\right\}.
$$
- 这与 Barbero 等人 (2024) 提出的 p-RoPE 方法一致，他们探索了 $r$ 占总子空间数 25%、50% 和 75% 的设置，并观察到在从头训练的 LLM 中比全 RoPE 略有优势
低频保留（Low-Frequency Preservation） 保留 $r$ 个旋转最慢（低频）的子空间：
$$
\mathcal{S}_{\text{low} }=\left\{k,\Big{|},\frac{d_{h} }{2}-r\leq k<\frac{d_{h} } {2}\right\}.
$$
- 选择该策略作为高频策略的对照实验
均匀采样（Uniform Sampling） 以等间隔选择 $r$ 个子空间：
$$
\mathcal{S}_{\text{uniform} }=\left\{\left.\left\lfloor k\frac{d_{h} }{2r}\right \rfloor,\right|0\leq k<r\right\}
$$
- 这通过几何间距平衡了高频和低频分量。实践中，$2r$ 通常能整除 $d_h$。这与 GPT-Neo (2021) 中使用的部分 RoPE 类似
头部感知 2-范数贡献度（Head-wise 2-norm Contribution） Barbero 等人 (2024) 首次提出了 2-范数贡献度来研究这些频率是否被使用以及它们如何发挥作用。该方法基于这样的观察：根据柯西-施瓦茨不等式，第 $k$ 个频率子空间对注意力对数几率（logits）的影响受相应查询和键分量的 2-范数上界限制，即 $\left|\left\langle\mathbf{q}_{i}^{[2k,2k+1]},\mathbf{k}_{j}^{[2k,2k+1]}\right \rangle\right|\leq\left|\mathbf{q}_{i}^{[2k,2k+1]}\right|\left|\mathbf{k}_{j }^{[2k,2k+1]}\right|$。对于每个头 $h$，论文在长序列上计算 LLM 中每个子空间的平均 2-范数得分 $^4$。然后，论文提出按它们的 2-范数得分对所有子空间进行排序，并选择前 $r$ 个：
$$
\mathcal{S}_{2\text{-norm} }=\operatorname*{top}\nolimits_{r} \left(\left|\mathsf{q}^{[2k,2k+1]}_{*}\right|\left|\mathsf{k}^{[2k,2k+1]}_{*}\right|\right).
$$
- 这种头部特定的选择自适应地保留了旋转关键的子空间
- 图 3 可视化了 Llama2-7B 四个头部的 2-范数
论文将在第 4.3 节分析这四种策略的有效性，并在附录 D 中对关键超参数 $r$ 进行消融研究
对于所有策略，未选择的子空间（$k \notin \mathcal{S}$）变为 NoPE 维度，从而实现与 MLA 潜在压缩的无缝集成

Low-rank Approximation, Low-rank Approximation

在从全 RoPE 转换为部分 RoPE 后，论文得到了 MLA 中 KV 缓存的第一个分量，表示为：
$$ \boldsymbol{k}_{i,rope}=\left[\boldsymbol{R}^{[2k,2k+1]}_{i}(\theta_{k})\boldsymbol{k}^{[2k,2k+1]}_{i} \right]_{k\in\mathcal{S} } $$
论文的下一个目标是推导第二个分量
$$ \boldsymbol{c}_{i,kv} \in \mathbb{R}^{d_{kv} } $$
- 它作为 $\boldsymbol{k}_{i,\text{nope} }$ 和 $\boldsymbol{v}_{i}$ 的低秩表示
给定 MHA 中的键 $\boldsymbol{k}_{i}=\boldsymbol{x}_{i}\boldsymbol{W}_{k}$ 和值 $\boldsymbol{v}_{i}=\boldsymbol{x}_{i}\boldsymbol{W}_{v}$，论文首先提取 $\boldsymbol{W}_{k}$ 中对应于 $\boldsymbol{k}_{i,\text{nope} }$ 的子空间，即未包含在 $\mathcal{S}$ 中的维度，得到：
$$ \boldsymbol{k}_{i,\text{nope} }=\boldsymbol{x}_{i}\boldsymbol{W}_{k,\text{nope} } $$
论文提出了两种基于奇异值分解（SVD）的策略（如图 4 所示）来在实现降秩的同时保留预训练知识：
解耦 SVD（Decoupled SVD, SVD${}_{\text{split} }$） 分别将 $\boldsymbol{W}_{k,\text{nope} }$ 和 $\boldsymbol{W}_{v}$ 分解为截断 SVD，各分配 $d_{kv}/2$ 个维度：
$$
\boldsymbol{W}_{k,\text{nope} }=\boldsymbol{U}_{k}\boldsymbol{\Sigma}_{k}\boldsymbol{V}^{\top}_{k}, \quad \boldsymbol{W}_{v}=\boldsymbol{U}_{v}\boldsymbol{\Sigma}_{v}\boldsymbol{V}^{\top}_{v},
$$
其中 $\boldsymbol{U}_{k},\boldsymbol{U}_{v},\boldsymbol{V}_{k},\boldsymbol{V}_{v} \in \mathbb{R}^{d_{h} \times \frac{d_{kv} }{2} }$，$\boldsymbol{\Sigma}_{k},\boldsymbol{\Sigma}_{v} \in \mathbb{R}^{\frac{d_{kv} }{2} \times \frac{d_{kv} }{2} }$。下投影矩阵 $\boldsymbol{W}_{d\cdot}$ 和上投影矩阵 $\boldsymbol{W}_{u\cdot}$ 变为：
$$
\boldsymbol{W}_{dk} =\boldsymbol{U}_{k}\boldsymbol{\Sigma}^{1/2}_{k}, \quad \boldsymbol{W}_{uk} =\boldsymbol{\Sigma}^{1/2}_{k}\boldsymbol{V}^{\top}_{k},
$$
$$
\boldsymbol{W}_{dv} =\boldsymbol{U}_{v}\boldsymbol{\Sigma}^{1/2}_{v}, \quad \boldsymbol{W}_{uv} =\boldsymbol{\Sigma}^{1/2}_{v}\boldsymbol{V}^{\top}_{v}.
$$
低秩表示 $\boldsymbol{c}_{i,kv}$ 可以使用 $\boldsymbol{c}_{i,kv}=[\boldsymbol{x}_{i}\boldsymbol{W}_{dk}, \boldsymbol{x}_{i}\boldsymbol{W}_{dv}]$ 构建
联合 SVD（Joint SVD, SVD${}_{\text{joint} }$） 为保留 $\boldsymbol{K}_{\text{nope} }$ 和 $\boldsymbol{V}$ 之间的相互作用，论文联合分解拼接后的矩阵：
$$
[\boldsymbol{W}_{k,\text{nope} }, \boldsymbol{W}_{v}] = \boldsymbol{U}_{kv}\boldsymbol{\Sigma}_{kv}\boldsymbol{V}^{\top}_{kv},
$$
其中 $\boldsymbol{U}_{kv}, \boldsymbol{V}_{kv} \in \mathbb{R}^{d_{h} \times d_{kv} }$，$\boldsymbol{\Sigma}_{kv} \in \mathbb{R}^{d_{kv} \times d_{kv} }$。潜在投影则为：
$$
\boldsymbol{W}_{dkv} = \boldsymbol{U}_{kv}\boldsymbol{\Sigma}^{1/2}_{kv},
$$
$$
\boldsymbol{W}_{uk} = \boldsymbol{\Sigma}^{1/2}_{kv}\boldsymbol{V}_{kv}[:, :-d_{v}], \quad \boldsymbol{W}_{uv} = \boldsymbol{\Sigma}^{1/2}_{kv}\boldsymbol{V}_{kv}[:, d_{v}:].
$$
这联合优化了键和值的潜在空间，即 $\boldsymbol{c}_{i,kv} = \boldsymbol{x}_{i}\boldsymbol{W}_{dkv}$，保留了对自回归生成至关重要的跨参数依赖性 $^5$。第 4.3 节显示 SVD${}_{\text{joint} }$ 优于 SVD${}_{\text{split} }$ ，验证了联合分解能更好地保留预训练知识

Experiment

论文在不同规模（SmoILM-135M/360M/1B7, Llama2-7B）且使用 MHA 或 GQA 预训练的 LLM 上评估了论文的方法
- 选择 SmoILM 系列是因为其预训练数据和框架都是开源的，这可以最大程度地减少微调数据和过程上的差异
- 选择 Llama2-7B 是因为它是广泛使用的开源大语言模型之一（但其预训练数据未开源，微调数据可能存在潜在差异）
论文分别使用 MHA2MLA 和 GQA2MLA 来表示架构迁移
- 两者均采用数据高效的全参数微调（data-efficient full-parameter fine-tuning）
- 默认使用基于头部的 2-范数贡献度选择（$\mathcal{S}_{2\text{-norm} }$，$r=\frac{d_{h} }{16}$）作为部分旋转位置编码（Partial-RoPE）策略
- 联合奇异值分解（SVD${}_{\text{joint} }$）作为低秩近似策略
论文的实验旨在回答三个关键问题：
- 1）MHA2MLA 如何最小化由架构转变引起的准确性下降？
- 2）MHA2MLA 在 KV 缓存减少比率方面取得了什么成果？
- 3）MHA2MLA 能否与 KV 缓存量化技术结合以实现复合收益？

Commonsense Reasoning Tasks

Main Results

如表 1 所示，论文的方法在四种模型规模（135M 到 7B）和不同的 KV 缓存压缩比（通过潜在维度 $d_{kv}$ 控制）下均实现了高效的架构迁移
当比较论文的微调方法与原始大语言模型的性能时
- 论文观察到四个基础模型的性能仅有微小变化：
  - 135M 模型下降 -0.25%
  - 360M 模型上升 +0.03%
  - 1B7 模型上升 +0.03%
  - 7B 模型上升 +0.37%
- 这表明微调数据并未显著降低或提高原始模型的性能，为 MHA2MLA 框架提供了一个合适的实验环境
随着 $d_{kv}$ 减小（例如从 32 到 16 再到 8），KV 缓存减少量增加（即从 -68.75% 到 -81.25% 再到 -87.5%），但通过微调恢复性能损失变得更具挑战性
- 图 5 显示了 135M（代表 GQA）和 7B（代表 MHA）在不同压缩比下的微调损失曲线
- 随着压缩比增加，与基线的损失差异变大
- 论文还观察到损失曲线的波动趋势几乎一致，这表明论文的架构迁移并未显著损害模型的内部知识
更大的模型在迁移到 MLA 架构时经历的性能下降更小
- 例如，压缩至 18.75% 时，性能下降分别为：
  - 135M 下降 2.41%
  - 360M 下降 2.69%
  - 1B7 下降 1.28%
  - 7B 下降 0.61%
- 这揭示了 MHA2MLA 的潜在缩放定律 (potential scaling law of MHA2MLA)
最后，从 135M 模型到 7B 模型，微调所需的 token 数量仅占预训练 token 的约 0.3% 到 0.6%，证明了论文方法的数据效率
总体而言，无论是使用 GQA2MLA 还是 MHA2MLA，架构迁移都以极小的成本实现，从而带来高效且经济（economical）的推理
表 1： 使用 MHA2MLA 或 GQA2MLA 的四个大语言模型的常识推理能力
- 六个基准测试包括 MMLU (2021)、ARC 简单和挑战集 (ARC, 2018)、PIQA (2020)、Hellaswag (HS, 2019)、OpenBookQA (OBQA, 2018)、Winogrande (WG, 2021)

Long Context Tasks

Settings

为评估模型的生成能力，论文采用 LongBench (2024) 作为生成性能的基准
所有模型均使用贪心解码策略进行测试
上下文窗口大小根据模型微调时使用的序列长度确定
使用 HQQ ( 2023) 和 Quanto 以不同精度级别设置缓存，以评估原始模型的性能作为基线
由于论文的方法与 KV 缓存量化兼容，论文还进行了额外实验来评估两种方法结合的效果

Main Results

如表 2 所示，在 LongBench 上，与训练后量化方法相比，MHA2MLA 实现了具有竞争力或更优的效率-准确性曲线
- 原生的 4 位量化在可比压缩比下仅带来适度的性能下降（-0.2% 到 -0.4%）
- 原生的 2 位量化实现了 87.5% 的 KV 缓存减少，但出现了严重的性能崩溃（-6.2% 到 -9%）
- 相比之下
  - MHA2MLA 在达到 87.5% 压缩（$d_{kv}=16$）时仅造成 3% 的准确性损失
  - 进一步与 4 位量化协同作用，实现了 92.19%/96.87% 的压缩（$d_{kv}=64/16$+Int4HQQ），同时将性能下降限制在 -0.5%/-3.2%，优于所有 2 位基线
- 这突显了 MHA2MLA 的潜在空间设计与数值精度降低是正交的，从而能够实现复合效率增益 (compound efficiency gains) 而不会产生破坏性干扰
表 2： Llama2-7B 和 MHA2MLA 在 LongBench 上的评估结果。粗体表示压缩比大于或等于 Int2 量化，同时性能也高于 Int2
图 5： 不同 KV 缓存存储比率下的微调损失曲线（颜色从浅到深代表 12.5%, 18.75%, 31.25%, 和 100%）

Ablation Study

四种部分旋转位置编码策略：$\mathcal{S}_{\text{high} }$, $\mathcal{S}_{\text{low} }$, $\mathcal{S}_{\text{uniform} }$, $\mathcal{S}_{\text{2-norm} }$

表 3 展示了四种将完整旋转位置编码（full-RoPE）转换为部分旋转位置编码（partial-RoPE）的策略结果
- 当将这四种策略与完整旋转位置编码进行比较时
  - 低频保留策略 $\mathcal{S}_{\text{low} }$ 遭受了最大的性能损失（135M 减少 -6.49%，1B7 减少 -1.21%）
  - 高频保留策略 $\mathcal{S}_{\text{high} }$ 的性能下降显著较小（135M 减少 -0.85%，1B7 减少 -0.76%）
  - 强调了高频子空间的重要性
- $\mathcal{S}_{\text{uniform} }$ 和 $\mathcal{S}_{\text{2-norm} }$ 都产生了更好的性能，$\mathcal{S}_{\text{uniform} }$ 保留了跨频率谱的子空间
- $\mathcal{S}_{\text{2-norm} }$ 则根据子空间对注意力分数的贡献来保留子空间
- 论文选择 $\mathcal{S}_{\text{2-norm} }$ 作为默认配置，因为被移除的子空间（即 NoPE）更适合（基于 SVD 的）低秩近似

两种基于 SVD 的低秩近似：$\text{SVD}_\text{split}$, $\text{SVD}_\text{joint}$

表 3 中每个组的最后两行比较了两种 SVD 方法的效果
- 在两个大语言模型上，$\text{SVD}_\text{joint}$ 方法 consistently 优于 $\text{SVD}_\text{split}$，在 135M 模型上平均性能提升 0.92%，在 1B7 模型上平均提升 0.74%
- 这表明 $\text{SVD}_\text{joint}$ 成为明确的默认选择

Efficient Attention Architectures

标准的多头注意力机制（Multi-Head Attention, MHA）(2017) 在上下文长度上具有二次复杂度，这促使了众多效率创新
MHA 变体，如多头查询注意力（Multi-Query Attention, MQA）和分组查询注意力（Grouped-Query Attention, GQA）(2023)
- 通过在不同头之间共享键/值来减少内存开销
- 但这是以参数剪枝和性能下降为代价的
其他并行的工作，如线性 Transformer (2019; 2020; 2021)、RWKV (2023) 和 Mamba (2023)
- 用线性循环或状态空间模型替代了 softmax 注意力，但在自回归生成中难以匹配标准注意力的表达能力
多头潜在注意力（Multi-Head Latent Attention, MLA）(2024) 通过将 KV 缓存压缩为低秩潜在向量而无需剪枝注意力参数，从而脱颖而出
论文的工作将 MLA 与主流架构（MHA/GQA）连接起来，通过数据高效的微调实现无缝迁移
许多线性注意力变体放弃了 softmax 查询-键交互（例如，通过核近似），但保留查询-键点积结构（即使是分解形式）的架构仍然与论文的 MHA2MLA 框架兼容

Economical Key-Value Cache（经济的键值缓存）

KV 缓存的内存占用已成为长上下文推理的关键瓶颈。最近的进展分为三类：
- 创新架构方法 ，如 MLA (DeepSeek-2024)、MiniCache (2024a) 和 MLKV (2024)，跨层或头共享或压缩 KV 表示
  - 虽然有效，但跨层共享可能混淆不同的注意力模式，可能损害特定任务的性能
  - 只有 MLA 在 DeepSeek 的 LLM 中得到了成功验证
- 量化技术 ，如 GPTQ (2022)、FlexGen (2023) 和 KIVI (2024b)
  - 以低比特格式（例如 2 比特）存储 KV 缓存，以精度损失为代价实现内存节省
- 动态剪枝方法
  - A2SF (2024) 和 SnapKV (2024) 从 KV 缓存中剪枝“不太重要”的 Token
    - 但 Token 剪枝可能丢弃关键的长距离依赖
  - 头剪枝（例如 SliceGPT (2024)、Sheared (2024) 和 Simple Pruning (2024)）则不可逆地降低了模型容量
论文的 MHA2MLA 方法实现了标准基于 Transformer 的大语言模型向更经济的 MLA 架构的迁移，并已证明其能够与 KV 量化技术集成以实现约 97% 的缓存节省
- 它在理论上也与其他方法（如剪枝）兼容

NLP——MHA2MLA(Partial-RoPE)

Introduction and Discussion

Preliminary

多头注意力机制 (Multi-Head Attention, MHA)

分组查询注意力 (Grouped-Query Attention, GQA)

多头查询注意力 (Multi-Query Attention, MQA)

多头潜在注意力 (Multi-Head Latent Attention, MLA)

MHA2MLA

部分旋转位置编码（Partial-RoPE）

全 RoPE 到部分 RoPE 的策略（Full-RoPE to Partial-RoPE Strategies）

Low-rank Approximation, Low-rank Approximation

Experiment

Commonsense Reasoning Tasks

Main Results

Long Context Tasks

Settings

Main Results

Ablation Study

四种部分旋转位置编码策略：\(\mathcal{S}_{\text{high} }\), \(\mathcal{S}_{\text{low} }\), \(\mathcal{S}_{\text{uniform} }\), \(\mathcal{S}_{\text{2-norm} }\)

两种基于 SVD 的低秩近似：\(\text{SVD}_\text{split}\), \(\text{SVD}_\text{joint}\)

Efficient Attention Architectures

Economical Key-Value Cache（经济的键值缓存）

Introduction and Discussion

Preliminary

多头注意力机制 (Multi-Head Attention, MHA)

分组查询注意力 (Grouped-Query Attention, GQA)

多头查询注意力 (Multi-Query Attention, MQA)

多头潜在注意力 (Multi-Head Latent Attention, MLA)

MHA2MLA

部分旋转位置编码（Partial-RoPE）

全 RoPE 到部分 RoPE 的策略（Full-RoPE to Partial-RoPE Strategies）

Low-rank Approximation, Low-rank Approximation

Experiment

Commonsense Reasoning Tasks

Main Results

Long Context Tasks

Settings

Main Results

Ablation Study

四种部分旋转位置编码策略：\(\mathcal{S}_{\text{high} }\), \(\mathcal{S}_{\text{low} }\), \(\mathcal{S}_{\text{uniform} }\), \(\mathcal{S}_{\text{2-norm} }\)

两种基于 SVD 的低秩近似：\(\text{SVD}_\text{split}\), \(\text{SVD}_\text{joint}\)

Related Work

Efficient Attention Architectures

Economical Key-Value Cache（经济的键值缓存）