#模型时代# DeepSeek是如何改善Transformer架构的?
epoch专家撰写了一篇技术博客(链接:http://t.cn/A6usSgx5),对DeepSeek v3的技术架构做了一番分析。结论是:“我看到DeepSeek所做的许多改进在“事后看来是显而易见的”:它们是那种如果有人事先问我,我会说是好主意的创新。然而,正如我之前所说,这并不意味着最初想出这些想法是容易的”。
***
DeepSeek 最近发布了 DeepSeek v3,该版本在开放权重模型中在基准性能方面处于最先进水平,并附有一份技术报告,详细描述了该模型的 training。令人印象深刻的是,他们仅使用了 280 万 H800 小时的 training 硬件时间就达到了这一 SOTA 性能——如果假设 40% 的 MFU,这相当于大约 4e24 FLOP。这比同样表现的 Llama 3.1 405B 所需的计算量少了大约十倍 training。
在本期中,我将介绍DeepSeek在其报告中强调的一些重要架构改进,以及为什么我们应该期待它们在性能上优于普通的Transformer。完整的技术报告还包含大量非架构细节,如果您想更好地了解在协调中等规模的training运行时需要解决的工程问题,我强烈建议您阅读它。
- 多头潜在注意力(MLA)
多头潜在注意力(简称为 MLA)是 DeepSeek 模型在长上下文推理中的最重要架构创新。该技术首次在 DeepSeek v2 中引入,是一种比传统方法(如分组查询和多查询注意力)更优越的减少 KV 缓存大小的方法。
我将简要解释一下 KV 缓存的内容。如果您对此很熟悉,可以直接跳到下一小节。
- KV 缓存是什么,为什么它很重要?
当Transformer在推理过程中用于顺序生成tokens时,它需要查看所有过去tokens的上下文,以决定下一个输出哪个token。这样做的简单方法是每次想要生成新的token时,进行一次包含所有过去tokens的前向传播,但这效率低下,因为那些过去的tokens已经被处理过。我们只是在重新计算之前已经获得并丢弃的结果。
为了避免这种重新计算,缓存所有过去 tokens 的相关内部状态 Transformer 是高效的,然后在我们需要它们用于未来 tokens 时从这个缓存中检索结果。因为过去 tokens 对未来 tokens 的影响仅通过它们在注意力机制中的键和值向量,因此只需缓存这些向量。这就是键值缓存或简称 KV 缓存名称的来源。
当上下文长度较短时,这种方法效果很好,但当长度变长时可能会变得昂贵。这是因为缓存读取并不是免费的:我们需要将所有这些向量保存在 GPU 高带宽内存(HBM)中,然后在需要将它们参与计算时将它们加载到张量核心中。如果每个token都需要知道其所有过去的上下文,这意味着对于我们生成的每个token,我们必须从 HBM 中读取整个过去的 KV 缓存。
在一个标准的多头注意力机制的基础上,过去每个 token 的 KV 缓存参数数量可以表示为:
2 * 注意力头维度 * 注意力头数量 * Transformer 块数量
例如,GPT-3 有 96 个注意力头,每个头有 128 个维度和 96 个块,因此对于每个 token,我们需要一个 2.36M 参数的 KV 缓存,或者在每个 KV 缓存参数精度为 2 字节的情况下为 4.7 MB。
GPT-3 不支持长上下文窗口,但如果暂时假设它支持,那么在 100K 上下文长度下生成的每个额外 token 将需要 470 GB 的内存读取,或者大约 140 毫秒的 H100 时间,考虑到 H100 的 HBM 带宽为 3.3 TB/s。每百万个 tokens 的价格为每小时每个 H100 2 美元,那么价格将是 80 美元,约为 Claude 3.5 Sonnet 客户价格的 5 倍(这可能远高于其自身的 Anthropic 成本)。这种简单的成本可以通过例如投机采样来降低,但它提供了一个合理的粗略估计。
这个粗略的计算显示了在处理 100K 或以上的上下文长度时,减少 KV 缓存大小的重要性。目前开源模型中最流行的方法是分组查询注意力。在这种架构设置中,我们为每对键和值头分配多个查询头,有效地将查询头分组在一起——因此得名。这将 KV 缓存的大小减少到我们选择的组大小的一个倍数。在如Llama 3.3 70B 和 Mistral Large 2 等模型中,分组查询注意力将 KV 缓存大小减少了大约一个数量级。
- 击败分组查询注意力
诸如分组查询注意力或 KV 缓存量化等方法的根本问题在于,它们涉及在降低 KV 缓存大小的同时妥协模型质量。与此不同,DeepSeek找到了一种在不妥协质量的情况下减少 KV 缓存大小的方法,至少在他们的内部实验中是这样。
他们通过将残差流中键和值向量的计算转变为一个两步过程来实现这一点。在一个普通的Transformer中,键和值向量是通过直接将残差流向量与形状为的矩阵相乘来计算的。
(头数 · 头部尺寸)x (模型尺寸)
DeepSeek的方法本质上强制这个矩阵为低秩:他们选择一个潜在维度,并将其表示为两个矩阵的乘积,一个矩阵的维度为潜在维度乘以模型维度,另一个矩阵的维度为(头的数量·头的维度)乘以潜在维度。然后,在推理过程中,我们只缓存潜在向量,而不是完整的键和值。通过减小潜在维度,我们可以缩小 KV 缓存的大小。
天真地说,这不应该解决我们的问题,因为每次需要生成新的 token 时,我们都必须重新计算实际的键和值。毕竟,我们需要完整的向量才能使注意力机制正常工作,而不是它们的潜在表示。多头潜在注意力基于一个巧妙的观察,即这实际上并不正确,因为我们可以将计算从潜在表示中上采样的键和值向量的矩阵乘法与查询和后注意力投影合并。
低秩压缩之所以如此有效,是因为不同注意力头需要了解的信息之间存在大量重叠。如果我们对单个头的键和值向量使用低秩压缩,而不是对所有头的所有键和值进行堆叠,那么该方法实际上等同于一开始就使用更小的头维度,我们将不会获得任何收益。利用不同头需要访问相同信息的事实对于多头潜在注意力机制至关重要。
诸如分组查询注意力的方法利用了相同重叠的可能性,但它们通过强制分组在一起的注意力头对查询做出相似的响应而效果不佳。换句话说,信息共享变得与在某种限制意义上具有相同的行为相耦合,这显然是一种不理想的特性。另一方面,低秩压缩允许不同的注意力头以非常不同的方式使用相同的信息。从理论上讲,这甚至可能对training产生有益的正则化效果,而DeepSeek在其技术报告中发现了这种效果。
我将其视为那些在事后看来显而易见的创新,但实际上需要对注意力头的实际作用有很好的理解才能提出这个想法。一旦你看到这种方法,立刻就会明白它不可能比分组查询注意力更糟,而且很可能会显著更好。然而,想出尝试这个想法是另一回事。
- 专家混合创新
最受欢迎的改进之一是引入了混合专家(MoE)模型。这些模型将Transformer的前馈块划分为多个不同的专家,并添加了一种路由机制,以上下文相关的方式将每个token发送到少数这些专家。这意味着模型可以拥有比每个特定token激活的更多参数,从某种意义上说,将模型的知识量与处理单个tokens的算术成本解耦。目前已知的最具影响力的 MoE 模型可能是原始的 GPT-4。
专家路由算法的工作原理如下:一旦我们退出任何层的注意力块,我们就会得到一个残差流向量作为输出。每个专家都有一个相应的相同维度的专家向量,我们通过查看与当前残差流的内积最高的专家来决定哪些专家将被激活。
这个问题在于它在模型的核心引入了一个表现不佳的间断函数,具有离散的图像,这与实现连续输入输出关系的普通Transformers形成鲜明对比。这导致梯度下降优化方法在 MoE training中表现不佳,常常导致“路由崩溃”,模型总是激活同几个专家,而不是将其知识和计算分散到所有可用的专家中。
为了直观理解路由崩溃,考虑尝试训练一个模型,例如总共有 16 个专家的 GPT-4,每token激活 2 个专家。现在,假设由于随机初始化的原因,这两个专家恰好是开始时表现最好的专家。梯度下降将加强选择这些专家的倾向。这意味着在更新过程中,这些专家几乎会获得所有的梯度信号并变得更好,而其他专家则滞后,因此其他专家将继续不被选择,产生一个正反馈循环,导致其他专家永远不会被选择或训练。
根本问题在于梯度下降只是朝着局部最优的方向前进。这通常在神经网络 training 中遇到的高维优化问题中效果良好。然而,当我们的神经网络在其行为上如此不连续时,即使问题空间的高维度也可能无法拯救我们免于失败。
解决这些training困难并非易事。DeepSeek v3 通过结合几种不同的创新来实现这一点,我将逐一讨论每一种。
- 辅助无损失负载均衡
避免路由崩溃的一个流行方法是强制“平衡路由”,即每个专家在足够大的批次中大致被激活相同次数的属性,通过向training损失添加一个测量特定批次中专家路由不平衡程度的项。这个项被称为“辅助损失”,引入它使模型朝向平衡路由的方向发展是合乎直觉的。然而,DeepSeek v3 技术报告指出,即使它确保了平衡路由,这种辅助损失也会损害模型性能。
他们的替代方案是向路由机制添加专家特定的偏置项,这些偏置项会被添加到专家亲和力中。这些偏置项不是通过梯度下降更新的,而是在training期间进行调整,以确保负载平衡:如果某个特定专家的访问量没有达到我们认为的应有水平,那么我们可以在每个梯度步骤中将其偏置项略微增加一个固定的小量,直到达到预期水平。技术报告指出,这种方法比依赖辅助损失实现了更好的性能,同时仍然确保了适当的负载平衡。
- 共享专家
上述解决路由崩溃的方法存在一个严重问题,即它在没有任何理由的情况下假设一个经过最佳训练的 MoE 将具有平衡的路由。然而,这一假设是值得怀疑的。
要理解原因,可以考虑任何大型语言模型可能有少量信息被频繁使用,而有大量信息则很少使用。例如,几乎任何对LLM的英语请求都要求模型知道如何说英语,但几乎没有对LLM的请求会要求它知道 1510 年法国国王是谁。因此,最优的 MoE 很可能应该有一些专家被频繁访问并存储“常见信息”,同时还有一些专家被稀疏访问并存储“专业信息”。
如果我们强制平衡路由,我们将失去实施这种路由设置的能力,并且不得不在不同的专家之间冗余地复制信息。然而,如果我们不强制平衡路由,我们面临路由崩溃的风险。为了摆脱这个困境,DeepSeek将专家分为两种类型:共享专家和路由专家。共享专家无论如何总是被路由:它们在专家亲和力计算和任何可能的路由不平衡损失项中都被排除在外。我们只关注确保路由专家的平衡路由。
这里的关键观察是,“路由崩溃”是一种极端情况,在这种情况下,每个单独专家被选择的可能性要么是 1,要么是 0。简单的负载均衡通过试图将分布推向均匀来解决这个问题,即每个专家应该有相同的被选择机会。然而,如果我们唯一关心的是避免路由崩溃,那么就没有理由特别针对均匀分布。DeepSeek v3 则针对一种分布,其中每个专家要么肯定被选择(概率为 1),要么以某个固定概率 p > 0 被选择。token
我认为即使这种分布也可能不是最优的,选择更好的分布将会产生更好的 MoE 模型,但这已经比仅仅强制使用均匀分布有了显著改善。
- 多重token预测
DeepSeek v3 对原始 Transformer 的最终更改是能够在模型的每次前向传播中预测多个 tokens。这使得他们能够在 training 期间使用多 token 预测目标,而不是严格的下一个 token 预测,并且他们在消融实验中展示了这一变化带来的性能提升。
基本思想如下:我们首先进行一次普通的前向传播以进行下一个token预测。与普通的Transformer一样,我们使用最终的残差流向量通过去嵌入和 softmax 生成下一个token的概率。然而,与普通的Transformer不同,我们还将这个向量输入到后续的Transformer块中,并使用该块的输出对第二个下一个token进行预测。我们可以根据需要进行多次迭代,尽管DeepSeek v3 仅在training期间预测两个tokens。
他们将关于更远期tokens的预测纳入training目标,通过向training损失中添加一个额外的交叉熵项,权重可以作为超参数进行调节。这不仅为他们在training期间提供了一个额外的目标以获取信号,还允许模型用于推测性地自我解码。我们可以在每次前向传播中生成一些tokens,然后将它们展示给模型,以决定从哪个点开始拒绝提议的延续。
DeepSeek v3 仅使用多token预测到第二个下一个token,技术报告中引用的第二个token预测的接受率在 85%到 90%之间。这相当令人印象深刻,并且如果我们使用上述的推测解码设置,应该可以在每个token的固定价格下实现几乎翻倍的推理速度(以每秒每个用户tokens为单位)。它看起来并不比解码Llama 3 405B 时获得的接受概率差,甚至可能更好。
我很好奇如果他们预测的时间超过下一个token会得到什么。如果例如每个后续的token都能使接受率相对降低 15%,那么通过预测更多的tokens,可能可以从这个投机解码设置中挤出更多的收益。
