技术分析：DeepSeek 如何改进 Transformer 架构？

DeepSeek 最近发布了 DeepSeek v3，这是目前在开放权重模型中基准性能表现最好的模型，同时还发布了一份技术报告，详细描述了该模型的训练过程。令人印象深刻的是，他们仅使用了 280 万个 H800 小时的硬件训练时间就实现了这一 SOTA 性能——如果我们假设 40% MFU，这相当于大约 4e24 FLOP。这比性能类似的 Llama 3.1 405B 少了大约 10 倍的训练计算量。

在本期中，我将介绍 DeepSeek 在其报告中强调的一些重要架构改进，以及为什么我们应该期望它们比普通 Transformer 带来更好的性能。完整的技术报告还包含大量非架构细节，如果您想更好地了解在安排中等规模的训练运行时必须解决的工程问题，我强烈建议您阅读它。

图 1：DeepSeek v3 架构及其两项最重要的改进：DeepSeekMoE 和多头潜在注意力 (MLA)。未显示多标记预测。来自 DeepSeek v3 技术报告。

多头潜在注意力（MLA）

多头潜在注意力（简称 MLA）是 DeepSeek 长上下文推理模型中最重要的架构创新。该技术首次在 DeepSeek v2 中引入，与分组查询和多查询注意力等传统方法相比，它是减少 KV 缓存大小的更佳方法。

我先简单介绍一下 KV 缓存。如果你对此很熟悉，可以直接跳到下一小节。

什么是 KV 缓存以及它为什么重要？

当使用 Transformer 在推理过程中按顺序生成 token 时，它需要查看所有过去 token 的上下文，以决定接下来输出哪个 token。最简单的方法是，每次我们想要生成新 token 时，简单地进行一次包括所有过去 token 的前向传递，但这样做效率低下，因为那些过去的 token 之前已经被处理过了。我们只是重新计算之前已经获得并丢弃的结果。

为了避免这种重新计算，一种有效的做法是缓存所有过去 token 的 Transformer 相关内部状态，然后在需要未来 token 时从此缓存中检索结果。因为过去 token 对未来 token 产生影响的唯一方式是通过注意力机制中的键和值向量，所以缓存这些向量就足够了。这就是键值缓存（简称 KV 缓存）名称的由来。

当上下文长度较短时，这种方法效果很好，但当上下文长度变长时，这种方法就会变得很昂贵。这是因为缓存读取不是免费的：我们需要将所有这些向量保存在 GPU 高带宽内存 (HBM) 中，然后在需要将它们参与计算时将它们加载到张量核心中。如果每个 token 都需要知道其所有过去的上下文，这意味着对于我们生成的每个 token，我们必须从 HBM 读取整个过去的 KV 缓存。

在使用标准多头注意力机制的 vanilla Transformer 中，每个过去 token 的 KV 缓存参数数量可以表示为：

2 * 注意力头尺寸 * 注意力头数量 * Transformer 块数量

例如，GPT-3 有 96 个注意力头，每个注意力头有 128 个维度和 96 个块，因此对于每个标记，我们需要 2.36M 个参数的 KV 缓存，或者 4.7 MB，每个 KV 缓存参数的精度为 2 字节。

GPT-3 不支持长上下文窗口，但如果我们暂时假设它支持，那么在 100K 上下文长度下生成的每个额外token将需要 470 GB 的内存读取，或者在 H100 的 HBM 带宽为 3.3 TB/s 的情况下，大约需要 140 毫秒的 H100 时间。以每 H100 每小时 2 美元的价格计算，每百万个token的价格将为 80 美元，大约是 Claude 3.5 Sonnet 向客户收取的价格的 5 倍（这可能大大高于 Anthropic 本身的成本）。这个简单的成本可以通过推测抽样等方式降低，但它给出了一个不错的粗略估计。

这个粗略的计算说明了为什么当我们处理 100K 或以上的上下文长度时，找到减少 KV 缓存大小的方法至关重要。到目前为止，开源模型中最流行的方法是分组查询注意。在这种架构设置中，我们为每对键和值头分配多个查询头，有效地将查询头分组在一起 - 因此得名该方法。这会将 KV 缓存的大小减少一个因子，该因子等于我们选择的组大小。在 Llama 3.3 70B 和 Mistral Large 2 等模型中，分组查询注意将 KV 缓存大小减少了大约一个数量级。

击败分组查询注意力

分组查询注意或 KV 缓存量化等方法的根本问题是，为了减少 KV 缓存的大小，它们会牺牲模型质量。相反，DeepSeek 找到了一种在不影响质量的情况下减少 KV 缓存大小的方法，至少在他们的内部实验中是这样。

他们通过将残差流中键和值向量的计算转变为两步过程来实现这一点。在 vanilla Transformer 中，键和值向量是通过将残差流向量直接乘以形状为

（头部数量·头部尺寸）x（模型尺寸）

DeepSeek 的方法本质上是强制该矩阵为低秩：它们选择一个潜在维度并将其表示为两个矩阵的乘积，一个矩阵的维度为潜在维度乘以模型，另一个矩阵的维度为（头部数量 · 头部维度）乘以潜在维度。然后，在推理过程中，我们仅缓存潜在向量，而不缓存完整的键和值。然后，我们可以通过减小潜在维度来缩小 KV 缓存的大小。

图 2：DeepSeek v2 技术报告中对多头潜在注意力的说明。

简单来说，这不应该解决我们的问题，因为每次我们需要生成新标记时，我们都必须重新计算实际的键和值。毕竟，我们需要完整的向量才能使注意力发挥作用，而不是它们的潜在向量。多头潜在注意力基于一个聪明的观察，即这实际上并不正确，因为我们可以将矩阵乘法合并，这些矩阵乘法将从潜在向量计算升级的键和值向量，分别与查询和后注意力投影合并。

低秩压缩之所以如此有效，是因为不同注意力头需要了解的信息之间存在大量重叠。如果我们对单个头的键和值向量使用低秩压缩，而不是对所有头堆叠在一起的所有键和值使用低秩压缩，那么该方法就相当于一开始就使用较小的头维度，我们不会获得任何好处。利用不同头需要访问相同信息这一事实对于多头潜在注意力机制至关重要。

分组查询注意等方法利用了相同重叠的可能性，但它们通过强制分组在一起的注意头对查询做出类似的响应而效率低下。换句话说，信息共享与在某种限制意义上具有相同的行为相结合，这显然是一种不受欢迎的特性。另一方面，低秩压缩允许不同的注意头以非常不同的方式使用相同的信息。从理论上讲，这甚至可能对训练产生有益的正则化效果，DeepSeek 在其技术报告中报告了这种效果。

我认为这是那些事后看来显而易见的创新之一，但需要很好地理解注意力头实际上在做什么才能想出这个创新。一旦你看到这种方法，你就会立即发现它不会比分组查询注意力更差，而且很可能要好得多。然而，想出尝试这个想法又是另一回事。

专家混合创新

对 vanilla Transformer 最受欢迎的改进之一是引入了混合专家 (MoE) 模型。这些模型将 Transformer 的前馈块划分为多个不同的专家，并添加一种路由机制，该机制以上下文相关的方式将每个 token 发送给少数专家。这意味着模型可以拥有比它为每个特定 token 激活的参数更多的参数，从某种意义上说，将模型的知识量与处理单个 token 的算术成本分离开来。目前已知的最具影响力的 MoE 模型可能是原始的 GPT-4。

专家路由算法的工作原理如下：一旦我们退出任何层的注意块，我们就会有一个残差流向量作为输出。每个专家都有一个相同维度的对应专家向量，我们通过查看哪些专家与当前残差流具有最高的内积来决定哪些专家将被激活。

问题在于，它引入了一个行为相当不良的不连续函数，模型的核心是离散图像，这与实现连续输入输出关系的 vanilla Transformers 形成鲜明对比。这导致梯度下降优化方法在 MoE 训练中表现不佳，经常导致“路由崩溃”，即模型陷入始终为每个 token 激活相同的少数专家的困境，而不是将其知识和计算传播到所有可用的专家。

为了直观地了解路由崩溃，可以尝试训练一个模型，例如 GPT-4，其中总共有 16 位专家，每个 token 有 2 位活跃专家。现在，假设由于随机初始化的原因，其中两位专家恰好是开始时表现最好的专家。梯度下降将强化选择这些专家的倾向。这意味着这些专家将在更新期间获得几乎所有的梯度信号并变得更好，而其他专家则落后，因此其他专家将继续不被选中，从而产生一个正反馈循环，导致其他专家永远不会被选中或训练。

根本问题是梯度下降只会朝着局部最优的方向前进。这通常在神经网络训练中遇到的非常高维度的优化问题中效果很好。然而，当我们的神经网络行为如此不连续时，即使问题空间的高维度也可能无法让我们免于失败。

解决这些训练难题并非易事。DeepSeek v3 通过结合多种不同的创新来实现这一目标，我将逐一讨论这些创新。

无辅助损耗负载平衡

避免路由崩溃的一种常用方法是强制“平衡路由”，即每个专家在足够大的批次中被激活的次数大致相等，方法是在训练损失中添加一个术语来衡量专家路由在特定批次中的不平衡程度。这个术语被称为“辅助损失”，直观地说，引入它会推动模型走向平衡路由。然而，DeepSeek v3 技术报告指出，即使这种辅助损失确保了平衡路由，也会损害模型性能。

他们的替代方案是将特定于专家的偏差项添加到路由机制中，并将其添加到专家亲和力中。这些偏差项不会通过梯度下降进行更新，而是在整个训练过程中进行调整以确保负载平衡：如果某个专家没有获得我们认为应该获得的命中次数，那么我们可以在每个梯度步骤中将其偏差项稍微增加一个固定的小量，直到达到预期。技术报告指出，这比依赖辅助损失可以获得更好的性能，同时仍能确保适当的负载平衡。

共享专家

上述解决路由崩溃问题的方法存在一个严重问题，即它毫无根据地假设经过最佳训练的 MoE 会实现平衡路由。然而，这是一个可疑的假设。

要了解原因，请考虑一下，任何大型语言模型可能都拥有少量经常使用的信息，而拥有大量不经常使用的信息。例如，几乎所有对法学硕士提出的英语要求都要求模型知道如何说英语，但几乎没有对法学硕士提出的要求要求它知道 1510 年的法国国王是谁。因此，最佳 MoE 应该有一些经常访问并存储“通用信息”的专家，而其他访问较少并存储“专业信息”的专家则非常合理。

如果我们强制平衡路由，我们将失去实现此类路由设置的能力，并且必须在不同的专家之间冗余地复制信息。但是，如果我们不强制平衡路由，我们将面临路由崩溃的风险。为了避免这种困境，DeepSeek 将专家分为两种类型：共享专家和路由专家。无论如何，共享专家始终会被路由：他们被排除在专家亲和力计算和任何可能的路由不平衡损失项之外。我们只关心确保路由专家的平衡路由。

这里的关键观察是“路由崩溃”是一种极端情况，其中每个专家被选中的可能性为 1 或 0。简单的负载平衡通过尝试推动分布均匀来解决此问题，即每个专家应该有相同的被选中的机会。但是，如果我们唯一关心的是避免路由崩溃，那么我们没有理由专门针对均匀分布。DeepSeek v3 的目标是每个专家要么肯定会被选中（概率为 1），要么以每个标记的某个固定概率 p > 0 被选中。

我认为这种分布可能不是最优的，而更好的分布选择将产生更好的 MoE 模型，但这已经比强制均匀分布有了显著的改进。

多标记预测

DeepSeek v3 对 vanilla Transformer 做出的最后一个改变是，它能够在模型的每次前向传递中预测出多个 token。这使得他们可以在训练期间使用多 token 预测目标，而不是严格的下一个 token 预测，并且他们在消融实验中展示了这一变化带来的性能提升。

基本思路如下：我们首先进行普通的前向传递，以预测下一个标记。与在 vanilla Transformer 中一样，我们使用最终残差流向量通过解嵌入和 softmax 生成下一个标记概率。但是，与 vanilla Transformer 不同，我们还将此向量输入到后续 Transformer 块中，并使用该块的输出对第二个下一个标记进行预测。我们可以随意迭代此过程，但 DeepSeek v3 在训练期间仅预测出两个标记。

图 3：摘自 DeepSeek v3 技术报告的多标记预测设置图示。

他们将这些关于更远的 token 的预测纳入训练目标中，方法是向训练损失添加一个额外的交叉熵项，该项的权重可以作为超参数调高或调低。这不仅为他们在训练期间提供了一个额外的信号目标，还允许模型用于推测性地解码自身。我们可以在每次前向传递中生成一些 token，然后将它们展示给模型，以决定我们需要从哪个点开始拒绝提议的延续。

DeepSeek v3 仅使用多token预测到第二个token，技术报告引用的第二个token预测的接受率在 85% 到 90% 之间。如果我们使用上述推测解码设置，这非常令人印象深刻，并且应该可以以固定的每个token价格将推理速度提高近一倍（以每个用户每秒token为单位）。它看起来并不比使用 Llama 3 70B 解码 Llama 3 405B 时获得的接受概率更差，甚至可能更好。

我很好奇，如果他们预测的比第二个标记更远，他们会得到什么。例如，如果每个后续标记给我们带来 15% 的相对接受度减少，那么通过预测更多的标记，可能可以从这种推测性解码设置中获得更多收益。

结论

我认为 DeepSeek 所做的许多改进“事后看来显而易见”：如果有人事先问我，我会说这些创新是好主意。然而，正如我之前所说，这并不意味着一开始就很容易想出这些想法。

我听到很多人表达了这样的观点：DeepSeek 团队在研究方面“品味不错”。仅从这些架构改进来看，我认为这种评价是正确的。这些改进似乎都不是通过对可能的想法进行强力搜索而发现的。相反，它们看起来像是由了解 Transformer 的工作原理以及如何解决其各种架构缺陷的研究人员精心设计的。

如果我不得不猜测下一步可能在哪里找到类似的改进，那么计算优先级可能是个不错的选择。目前，Transformer 在每个 token 上花费的计算量相同，无论它正在处理或预测哪个 token。这在直觉上似乎效率低下：如果模型做出更难的预测，它应该考虑更多，如果做出更容易的预测，它应该考虑更少。在某种程度上，这可以通过可变测试时间计算扩展纳入推理设置，但我认为也应该有一种方法将其直接纳入基础模型的架构中。