如何降低推理模型的内存使用
🔳 我们探索 LightThinker 和多头潜在注意力如何结合以减少内存并提升性能
AI 模型已从快速思考（快速给出答案）转变为通过将问题分解为更小的步骤来更仔细地思考。类似 O1 的思维与链式思考方法的实施，使得大型推理模型，如 OpenAI 的 O1、O3 和 DeepSeek-R1，能够回溯、重试并完善其推理，使其在解决棘手问题方面表现得更好。我们已在之前的一篇文章中讨论了扩展测试时间计算的所有重要方面和优势。然而，存在一个大问题：这种推理会产生大量文本（tokens），占用内存并降低速度，从而增加处理成本。这在 Transformer 模型中尤为明显——它们生成的文本越多，所需的内存和计算能力就越多。随着大型推理模型越来越普及，我们必须找到方法来缓解它们的弱点，同时充分发掘它们的改进潜力。
今天，我们将重点关注由于内存使用增加而导致的过长处理时间问题。如果能解决内存效率低下问题，模型在保持高准确性的同时，可以变得更加平衡和高效。目前已经提出了两种显著的方法来减少推理模型中的内存使用：1) LightThinker，它帮助模型学习如何总结自己的“想法”并根据这些简短而有意义的摘要来解决任务；2) 多头潜在注意力（MLA），这是 DeepSeek 提出的解决方案，早在发布 DeepSeek-V2 时就已提出，后来在 DeepSeek-V3 和 DeepSeek-R1 中实现。
今天我们邀请您与我们一起深入探讨这些概念，并思考将它们融合在一起的潜在益处。
📨 点击关注！如果你想直接在收件箱中接收我们的文章，请在此订阅
在本期节目中，我们将讨论
什么是 LightThinker？
LightThinker 背后的理念
正如我们已经提到的，我们需要优化方法，使高质量的推理模型更快、更高效，避免高昂的内存成本。
其中一种方法是由浙江大学-蚂蚁集团知识图谱联合实验室开发的 LightThinker。LightThinker 不仅仅是手动删除单词或内存，它还教导模型在解决问题时“总结”自己的“想法”。可以将其想象成人们如何记录要点而不是写下每个细节。让我们详细了解它是如何工作的。
图片来源：LightThinker 原始论文
LightThinker 如何工作？
总的来说，LightThinker 不会保留冗长、详细的推理步骤，而是将其压缩成更短、更精炼的摘要，然后根据这些摘要继续推理。
LightThinker 做了两件事，这一点很重要：
- 决定何时压缩推理步骤。
- 决定如何压缩它们。
以下是用于执行这些操作的技术。
- 何时压缩？ 有两种方法可以决定何时总结模型的“想法”
- 令牌级压缩： 模型在固定数量的单词后压缩“想法”。这很简单，但可能会不自然地截断想法。
- 思想级压缩： 模型在完成一个完整概念（如句子或段落）后压缩“想法”。这使得想法更具条理，但需要额外的处理来判断何时一个想法已完成。研究人员倾向于在 LightThinker 中使用这种压缩方法，因为它能更好地保留含义。
- 如何压缩？ 还有两种主要的信息总结方式
- 文本压缩： 模型用更短的摘要替换冗长的想法。然而，这需要使用额外的编码模型进行额外处理，并会降低速度。
- 隐藏状态压缩： 模型不重写文本，而是将关键细节存储在特殊令牌中。这些令牌就像 AI 模型以后可以使用的思维笔记一样。研究人员选择隐藏状态压缩，因为它速度更快，并且不需要额外的模型。
以下是 LightThinker 逐步实现压缩的方式
- 在总结之前，LightThinker 将文本分解成更小的部分。它在各部分之间插入特殊令牌来标记摘要
- (可选) → 一个标记，告诉模型压缩先前的思想。
- C (缓存令牌) → 表示存储关键点的先前思想的摘要令牌。
- [o] (输出令牌) → 一个标记，表示“使用此摘要继续推理。”它帮助模型根据已压缩的内容生成新内容。
因此，三个文本部分的序列可能如下所示：→ 第一节 → → 摘要 (C) → [o] → 第二节 → → 摘要 (C) → [o] → 第三节。这个阶段被称为“数据重建”。
- 学习压缩和使用摘要： 一旦数据结构化，模型就会学习何时以及如何压缩信息。基于思想的注意力掩码控制模型在每个步骤中可以和不能“查看”的内容
- 在压缩过程中，模型只能处理原始输入、之前压缩的内容（C）和当前想法。
- 在生成输出令牌时，模型只能“查看”输入问题和之前压缩的内容。这确保模型在推理时只查看摘要而不是原始数据。以下是 LightThinker 在三步推理中的注意力掩码图示
图片来源：LightThinker 原始论文
- 最后，模型被训练成只使用摘要 (C) 来预测下一个令牌。它不允许通过查看完整文本来“作弊”，并且必须学习如何高效地存储信息并逐步检索它。
总的来说，LightThinker 将冗长的想法压缩成简短的内存令牌，并仅用重要的细节继续推理。重要特性是 LightThinker 旨在将内存压缩与推理分离，以实现更高的准确性和效率。
让我们看看 LightThinker 的实际结果，以探索为什么它是一种更好的方法。
LightThinker 的实际表现
LightThinker 使 AI 模型在实际任务中更实用、更高效。它的性能结果令人印象深刻
- 内存减少： LightThinker 将峰值令牌使用量减少了 70%，这意味着它存储了更少的不必要信息。
- 处理速度更快： 它将 Qwen2.5-7B 的推理时间加快了 26%，Llama3.1-8B 的推理时间加快了 41%（它在这两个模型上进行了测试）。推理时间的减少也取决于响应的长度。例如，LightThinker 在生成长响应（32K 令牌）时将推理时间减少了 44%。对于较短的文本（1K–4K 令牌），它仍然节省了 1%–4% 的推理时间。
- 最小的准确性损失： 准确性在 Qwen 上仅下降 1%，在 Llama 上下降 6%，这是效率提升的合理权衡。
- 生成更少的令牌： LightThinker 将模型生成的令牌数量在 Qwen 上减少了 15%，在 Llama 上减少了 13%。
- 对于数学问题等简单任务，它压缩得更彻底，而对于 GPQA（研究生级别的 Google 验证问答基准）等复杂任务，它会更谨慎地压缩以保留重要信息。
- 由于 LightThinker 的设计将压缩和推理步骤分离，准确性提高了 2%，并且当结合其注意力掩码策略时，性能进一步提升了 7%。
- 正确的缓存大小很重要： 较小的缓存大小（LightThinker 存储在内存中的令牌数量）会更频繁地压缩，而较大的缓存大小则压缩较少但保留更多信息。这就是为什么增加缓存大小可以提高准确性并减少推理时间的原因。
- 然而，LightThinker 在数学方面表现不佳： 数学问题有时对它来说很棘手，因为数值可能会被错误压缩。
图片来源：LightThinker 原始论文
由于 LightThinker 是一种摘要方法，让我们简要概述其优点和局限性，以便我们可以保留这些有意义的“笔记”以供文章的结论。:)
为什么 LightThinker 是一种更好的方法？
- 它将思考和总结分离，以提高效率。
- 更好地追踪过去的思想。
- 在保持准确性的同时节省内存。
- 运行速度提高 44%。
也并非没有局限性
- 在数学任务上表现不佳。
- 它不是 Llama 模型的最佳方法。
- 它的动态压缩过程有时会导致内存使用量突然达到峰值。
- LightThinker 在训练时使用固定数量的缓存令牌，但它是否能适应实际任务中不断变化的令牌需求尚不确定。
- 目前尚不清楚，如果 LightThinker 在更大的数据集上训练，它是否会变得更好。
- LightThinker 使用全参数微调进行训练，这需要大量资源，而 LoRA 和 QLoRA 等更高效的微调方法（使用更少参数）尚未经过测试。
这些限制为 LightThinker 提供了改进和增强其灵活性的空间。然而，它已经展示出有益的内存使用减少和更快的推理速度。
现在，是时候转向 DeepSeek 的进步了，它也有效地减少了内存使用，但工作方式完全不同。
什么是多头潜在注意力（MLA）？
为什么需要 MLA？
备受瞩目的 DeepSeek-R1 推理模型融合了两项技术创新，这两项创新在 DeepSeek-V2 发布时就已经提出（我们在 AI101 的关于 DeepSeek 的文章中曾简要介绍过）。它们包括
- 特殊的专家混合系统，由多个专家子模型（每个专家又分为更小的专业部分）和始终活跃处理常见知识的共享专家组成。
- 一种特殊的注意力机制，称为多头潜在注意力（MLA）。
图片来源：DeepSeek-V2 原始论文
为什么 DeepSeek 决定为他们最好的 AI 模型修改传统的 Multi-Head Attention (MHA) 机制？
在 Transformer 中，MHA 机制有助于专注于输入中最相关的部分来处理和生成文本。然而，MHA 在推理过程中存储了大量的键值 (KV) 对，这会占用大量内存并降低模型的工作速度。
DeepSeek 的多头潜在注意力 (MLA) 是一种经过修改的注意力机制，可将 KV 缓存压缩成更小的形式。 MLA 就像一个智能存储系统，它能有效地压缩过去的信息，同时保持其可用于未来使用。它通过一种称为低秩键值联合压缩的技术来实现这一点。这使得模型能够更快、以更少内存处理信息，而不会损失准确性。MLA 究竟如何压缩 KV 缓存？
MLA 如何工作？
MLA 通过低秩键值联合压缩技术减少了 KV 存储，同时保持了强大的性能
图片来源：DeepSeek-V2 原始论文
- 键值 (KV) 对的压缩： MLA 不会为每个令牌单独存储所有完整大小的 KV 对，而是在保存它们之前将其压缩成更小、更低维度的表示。这意味着 MLA 通过数学变换将它们投影到更小的潜在空间中。
- 注意力计算的解压缩： 当模型需要使用存储的 KV 对时，它通过将压缩数据恢复到原始大小来重建它们。压缩的 KV 数据在推理过程中会重新扩展，使模型能够像拥有完整大小的 KV 对一样工作。这很有用，因为模型仍然可以从以前的令牌中获取重要信息，但内存占用仍然很小，从而加快了推理速度。
- 解耦旋转位置嵌入（RoPE） RoPE 修改 K 和 Q 以包含位置数据，因此模型也可以记住词序。然而，当压缩 KV 对时，位置信息也会与压缩形式纠缠在一起，使其难以正确重建。MLA 也解决了这个问题。它将位置信息与主要的 KV 压缩分离或解耦。 它通过使用额外的多头查询和共享键来单独处理位置信息。这种方法具有以下优点
- 它防止位置编码干扰 KV 压缩过程。
- 模型无需重新计算所有过去的键以进行位置调整。
- 模型可以快速检索过去的令牌信息，而无需额外处理。
MLA 的主要优点是什么？
- 由于键和值的低秩压缩，它大大减少了内存使用。
- 加速文本生成： 更少的 KV 存储意味着更快的检索和推理。
- 性能与传统 MHA 一样好： 与 Multi-Query Attention (MQA) 和 Grouped-Query Attention (GQA) 等其他方法不同，MLA 不会牺牲质量。
- 处理位置编码： 解耦的 RoPE 确保词序得以保留。
然而，像往常一样，我们需要讨论 MLA 的局限性，因为没有任何技术是理想的。
也并非没有局限性
- 由于压缩导致一些信息丢失： MLA 可能会轻微削弱长程依赖关系。
- 压缩和解压缩的额外计算： 这可能会减慢训练速度，并略微影响推理。
- RoPE 兼容性问题： MLA 需要解耦 RoPE，增加了实现复杂性。
- 压缩比权衡： 需要仔细调整以平衡内存节省和性能。
- 在不同模型上的基准测试有限： MLA 未在各种 AI 架构上进行足够的测试。
- 与其他注意力机制集成面临挑战： 对于使用自定义注意力方法的模型，需要额外调整。
尽管如此，MLA 是 DeepSeek 最佳 AI 模型的一部分，它通过将键值 (KV) 缓存压缩到潜在向量中，显着降低了键值 (KV) 缓存的大小，从而确保了高效的推理。我们已经看到了 DeepSeek-R1 惊人的性能结果，但为了获得完整的了解，让我们回顾一下它们，并从之前的 DeepSeek 模型中添加一些更多细节。
MLA 结果
这里我们可以看到 MLA 如何为 DeepSeek-V2 的改进做出了贡献。
MLA 将键值（KV）缓存压缩到潜在向量中，减少了 93.3% 的内存需求。由于需要存储和访问的数据更少，每个推理步骤都变得更快。通过降低 KV 内存使用量，可以一次处理更多序列，从而与没有 MLA 的模型相比，生成吞吐量提高了 5.76 倍。当然，DeepSeek 的两项创新——MLA 和 DeepSeekMoE 架构都对改进有所贡献，但内存使用压缩策略来自 MLA。
图片来源：DeepSeek-V2 原始论文
DeepSeek-R1 的惊人结果表明，缓慢的逐步推理不会因使用 MLA 等压缩技术而受到影响。相反，该模型实现了与 OpenAI 高级模型（如 o1-1217）相当的推理性能。
图片来源：DeepSeek-R1 原始论文
❓这引出了一个问题❓
能否在保持 DeepSeek-R1 等模型的准确性并提高其速度和效率的同时，进一步降低其内存使用量？
如果我们结合 LightThinker 和 MLA 的概念会怎样？
LightThinker 和 MLA 都旨在通过减少内存使用和实现更快的推理来提高 LLM 效率。但它们针对模型操作的不同方面。LightThinker 旨在压缩推理，而 MLA 则通过压缩键值 (KV) 缓存来优化注意力内存使用，以减少存储并加快推理速度。
简单来说，MLA 就像压缩和解压缩存储的数据以节省内存，而 LightThinker 就像总结一个冗长的对话，这样 AI 模型就不必记住每个细节。
现在想象一下，如果像 DeepSeek-R1 这样的模型将 LightThinker 技术与 MLA 结合起来。 这种组合可能会带来更高效、更强大的推理模型。以下是一些关于这种合作可能性的建议
- 首先，LightThinker 压缩推理步骤，只将关键信息存储在摘要中，同时丢弃冗余。
- 然后，MLA 可以确保选择性地关注这些压缩的摘要，优先处理最相关的潜在细节，而不是同等对待所有令牌。
这可以使 MLA 充当自适应检索机制，改善 LightThinker 压缩“思想”的召回率。MLA 还可以改进 LightThinker 压缩步骤的检索和扩展方式，在紧凑性和推理深度之间取得平衡。因此，这两种概念的结合可以通过在多个潜在维度上选择性地加权注意力来加速处理，这些维度可以预先有效压缩。
此外，LightThinker 和 MLA 可以相互弥补对方的局限性。例如
- MLA 可以帮助 LightThinker 恢复在激进总结过程中丢失的关键细节，甚至防止关键细节丢失。
- 由于 LightThinker 在数学任务上表现不佳，因为数值可能被错误压缩，MLA 可以通过优先结构化检索来加强压缩思想中的数值一致性。
- MLA 选择性关注 LightThinker 压缩缓存的能力可以使其压缩策略在实际任务中更具灵活性。
- MLA 可能会略微削弱长程依赖关系，在这种情况下，LightThinker 的逐步思想压缩可以减少 MLA 需要处理的初始信息。
总的来说，探索 MLA 如何在 LightThinker 的“笔记”（压缩的“思想”）之间分配注意力，并仅在需要时进行压缩和解压缩以节省内存，是很有趣的。它们共同可以提高推理效率，保持推理深度，并进一步优化内存使用。
我们鼓励您尝试。
结论
内存压缩是一种重要的优化技术，特别是对于需要大量资源进行扩展推理过程的大型推理模型。LightThinker 和 MLA 有效地减少了内存使用并加快了推理处理速度。然而，随着 AI 模型不断扩展其推理步骤和推理时间以实现更高的准确性——以及社区超越 O1 和 DeepSeek-R1 等模型当前成就的推动——我们不能停止开发更有效的压缩技术，因为我们的资源非常宝贵。
这就是为什么像 MLA + LightThinker 这样的混合方法可能是更高效内存使用和更快推理的关键。开发人员和研究人员，你们对此有何看法？
作者：Alyona Vert 想法和编辑：Ksenia Se
额外资源，深入了解
- LightThinker: 逐步思考压缩，作者：Jintian Zhang, Shuofei Qiao, Huajun Chen, Ningyu Zhang 等。
- DeepSeek-V2：一个强大、经济、高效的专家混合语言模型 作者：Bo Liu, Damai Dai 等。
- DeepSeek-R1：通过强化学习激励 LLM 的推理能力 作者：Daya Guo, Qihao Zhu 等。
- 大型语言模型的模型压缩和高效推理：一项调查
- 高效 Transformer：一项调查
- Keyformer：通过关键令牌选择实现高效生成推理的 KV 缓存减少
- FastCache：通过轻量级 KV-缓存压缩框架优化多模态 LLM 服务
- 超越 RAG：面向综合知识推理的任务感知型 KV 缓存压缩 作者：Orion Weller
图灵邮报资源
📨 如果您想直接在收件箱中收到我们的文章，请在此订阅
