Deepseek推理太慢，蚂蚁集团&浙大提出LightThinker：动态压缩推理链，LLM效率飙升

随着Deepseek R1范式的爆发，越来越多人相信给定足够的思考时间，模型能更好地回答问题，但是当前的大模型多采用基于Transformer架构的模型，其注意力机制的计算复杂度随上下文长度呈二次增长，而KV缓存的存储开销随上下文长度线性增长。这导致在长文本生成和复杂推理任务中，LLMs的效率受到严重限制。

小马不会过河

1157人浏览 · 2025-02-28 23:00:28

小马不会过河 · 2025-02-28 23:00:28 发布

为了解决这一问题，由浙江大学、蚂蚁集团以及浙江大学-蚂蚁集团知识图谱联合实验室的研究人员共同撰写的论文《LightThinker: Thinking Step-by-Step Compression》，提出了一种名为 LightThinker 的新方法，旨在通过动态压缩推理过程中的中间思考步骤来提高大型语言模型（LLMs）在复杂推理任务中的效率。

LightThinker方法

LightThinker 的核心思想是模仿人类的认知过程，将冗长的思考步骤压缩成紧凑的表示形式，并丢弃原始的推理链，从而显著减少存储在上下文窗口中的token数量。具体而言，该方法通过以下步骤实现：

数据构造：通过构造数据集，训练模型学习何时以及如何进行压缩。
隐藏状态映射：将需要压缩的思考步骤的隐藏状态映射到少量特殊token（即要点token，gist tokens）对应的隐藏状态上。
注意力掩码设计：通过精心设计的注意力掩码，让LLMs学习如何压缩以及如何基于压缩后的内容继续生成文本。

此外，作者引入了一个新的指标——依赖度（Dependency, Dep），用于量化推理过程中对历史token的依赖程度，从而衡量压缩的程度。较低的Dep值表示对原始长上下文的依赖度降低，压缩效果更显著。

数据构造

数据构造的核心目标是：

教会模型何时压缩：模型需要学会在推理过程中识别出哪些思考步骤可以被压缩。
教会模型如何压缩：模型需要学会将冗长的思考步骤压缩成紧凑的表示形式（即要点token，gist tokens），同时保留关键信息以支持后续推理。

数据构造的具体步骤

输出分割（Output Segmentation）

首先，对于每个训练样本中的输出（即推理过程和答案），使用分割函数 Seg() 将其划分为多个子序列（subsequences）。分割函数可以根据不同的策略实现，例如：

基于token数量的分割：每隔固定数量的token进行一次分割。
基于语义的分割：以句子或段落为单位进行分割，更适合保留语义完整性。

例如，假设一个推理过程包含多个步骤，分割函数可以将每个步骤作为一个子序列。

插入特殊token

在分割后的子序列之间插入一组特殊token，这些特殊token用于触发压缩操作并存储压缩后的内容。具体插入的特殊token包括：

压缩触发token ：用于指示模型需要对前面的思考步骤进行压缩。如果分割是基于token数量的，这个token可以省略。
要点token [c]：用于存储压缩后的内容。这些token是压缩后的关键信息，模型需要学会将冗长的思考步骤映射到这些要点token上。
输出token [o]：用于基于压缩后的内容继续生成后续推理内容。

例如，假设原始输出 Y 被分割为 S1, S2, …, Sk，则增强后的输出 Ŷ 为：

Ŷ = {S1, , [c], [o], S2, , [c], [o], …, Sk}

构造增强数据集

通过上述步骤，原始数据集 D 被增强为新的数据集 ˆD，其中每个样本的输出部分被替换为增强后的输出 Ŷ。增强后的数据集 ˆD 用于训练模型，使其学会在推理过程中动态压缩思考步骤。

通过上述数据构造方法，模型在训练过程中会学习到以下内容：

何时压缩：模型通过token识别出需要压缩的位置。
如何压缩：模型将冗长的思考步骤 Si 压缩到要点token [c] 中，并基于压缩后的内容 [o] 继续生成后续推理内容。
语义保留：通过分割函数和特殊token的设计，模型能够保留关键信息，同时丢弃冗余内容，从而实现高效的推理。

实验结果

作者在四个数据集（GSM8K、MMLU、GPQA、BBH）上对两种模型（Qwen2.5-7B系列和Llama3.1-8B系列）进行了广泛的实验。实验结果表明，LightThinker在减少峰值内存使用和推理时间的同时，保持了与原始模型相当的准确率。具体来说：

在Qwen模型上，LightThinker将峰值token使用量减少了70%，推理时间减少了26%，同时准确率仅下降了1% ，Dep. 降低了 78%（即达到了 16.6/3.7=4.5 倍的压缩比）。
在Llama模型上，LightThinker将峰值token使用量减少了70%，推理时间减少了1%，准确率下降了6%，Dep. 降低了 74%（即达到了 10.5/2.7=3.9 倍的压缩比）。
Thought级别的分段方法比 Token 级别更为有效，在 Qwen 上，准确度提高了 6.2%，在 Llama 上提高了 5.6%。这表明 Token 级别的分段方法可能会丧失语义边界。

LightThinker的效率

左图展示了 H2O、AnLLM、LightThinker 和 Vanilla 在四个数据集上的平均生成 Token 数量，右图展示了在 Qwen 上进行 GPQA 时压缩 Token 数量的分布
LightThinker 是唯一在生成 Token 数量上减少的模型，相较于 Vanilla，Qwen 上减少了 15%，Llama 上减少了 13%。
压缩后的 Token 数量分布呈现长尾模式，这意味着大多数压缩操作发生在少数几个 Token 上，而其他 Token 则相对较少被压缩。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述