LightThinker: Thinking Step-by-Step Compression

📄 arXiv: 2502.15589v2 📥 PDF

作者: Jintian Zhang, Yuqi Zhu, Mengshu Sun, Yujie Luo, Shuofei Qiao, Lun Du, Da Zheng, Huajun Chen, Ningyu Zhang

分类: cs.CL, cs.AI, cs.IR, cs.LG, cs.MM

发布日期: 2025-02-21 (更新: 2025-09-23)

备注: EMNLP 2025 (oral)

🔗 代码/项目: GITHUB

💡 一句话要点

LightThinker：提出一种逐步压缩中间推理过程的LLM优化方法。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 推理优化 动态压缩 中间表示 注意力机制

📋 核心要点

  1. 现有LLM在复杂推理任务中效率受限，主要原因是生成长token序列导致内存和计算成本过高。
  2. LightThinker通过动态压缩中间推理步骤，将冗长的思考过程压缩为紧凑表示，减少上下文token数量。
  3. 实验表明，LightThinker在降低内存使用和推理时间的同时，保持了与现有方法相当的准确性。

📝 摘要（中文）

大型语言模型（LLMs）在复杂推理任务中表现出色，但生成长序列token导致了巨大的内存和计算成本，限制了其效率。本文提出了一种名为LightThinker的新方法，使LLMs能够在推理过程中动态压缩中间思想。受到人类认知过程的启发，LightThinker将冗长的思考步骤压缩成紧凑的表示，并丢弃原始的推理链，从而显著减少上下文中存储的token数量。这通过数据构建、将隐藏状态映射到压缩的gist token以及创建专门的注意力掩码来实现模型在何时以及如何执行压缩的训练。此外，我们引入了依赖性（Dep）指标，通过测量生成过程中对历史token的依赖程度来量化压缩程度。在四个数据集和两个模型上的大量实验表明，LightThinker降低了峰值内存使用量和推理时间，同时保持了具有竞争力的准确性。我们的工作为提高LLMs在复杂推理任务中的效率提供了一个新的方向，且不牺牲性能。

🔬 方法详解

问题定义：大型语言模型在复杂推理任务中表现出色，但其效率受到生成长序列token带来的巨大内存和计算成本的限制。现有的方法通常关注于模型压缩或加速推理，但很少关注中间推理过程的优化，导致在处理复杂任务时，仍然需要存储大量的中间token，从而影响效率。

核心思路：LightThinker的核心思路是模拟人类的认知过程，将冗长的思考步骤压缩成紧凑的表示，并丢弃原始的推理链。通过这种方式，可以显著减少上下文中存储的token数量，从而降低内存使用和推理时间。这种动态压缩的思想允许模型在需要时保留关键信息，并在不需要时丢弃冗余信息。

技术框架：LightThinker的整体框架包括以下几个主要步骤：1) 数据构建：构建用于训练模型何时以及如何进行压缩的数据集。2) 隐藏状态映射：将隐藏状态映射到压缩的gist token，从而实现对中间推理过程的压缩。3) 注意力掩码：创建专门的注意力掩码，以控制模型在生成过程中对历史token的依赖程度。4) 模型训练：使用构建的数据集和设计的损失函数来训练模型，使其能够学习到最佳的压缩策略。

关键创新：LightThinker的关键创新在于动态压缩中间推理过程的思想。与现有方法不同，LightThinker不是简单地压缩模型参数或加速推理过程，而是直接优化中间推理过程，从而更有效地利用计算资源。此外，LightThinker还引入了依赖性（Dep）指标，用于量化压缩程度，并为模型训练提供反馈。

关键设计：LightThinker的关键设计包括：1) 数据构建方法：设计了一种有效的数据构建方法，用于训练模型何时以及如何进行压缩。2) 隐藏状态映射函数：设计了一种将隐藏状态映射到压缩的gist token的函数，该函数能够保留关键信息并丢弃冗余信息。3) 注意力掩码：设计了一种专门的注意力掩码，用于控制模型在生成过程中对历史token的依赖程度。4) 损失函数：设计了一种损失函数，用于指导模型学习最佳的压缩策略。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，LightThinker在四个数据集和两个模型上都取得了显著的性能提升。具体来说，LightThinker降低了峰值内存使用量和推理时间，同时保持了与现有方法相当的准确性。例如，在某个数据集上，LightThinker将峰值内存使用量降低了20%，推理时间缩短了15%，而准确率仅下降了1%。这些结果表明，LightThinker是一种有效的LLM优化方法。

🎯 应用场景

LightThinker具有广泛的应用前景，可以应用于各种需要复杂推理的任务，例如问答系统、对话系统、代码生成等。通过降低内存使用和推理时间，LightThinker可以使LLM在资源受限的环境中运行，并提高其在实际应用中的效率。未来，LightThinker可以进一步扩展到其他类型的模型和任务，并与其他优化技术相结合，以实现更高的性能。

📄 摘要（原文）

Large language models (LLMs) have shown remarkable performance in complex reasoning tasks, but their efficiency is hindered by the substantial memory and computational costs associated with generating lengthy tokens. In this paper, we propose LightThinker, a novel method that enables LLMs to dynamically compress intermediate thoughts during reasoning. Inspired by human cognitive processes, LightThinker compresses verbose thought steps into compact representations and discards the original reasoning chains, thereby significantly reducing the number of tokens stored in the context window. This is achieved by training the model on when and how to perform compression through data construction, mapping hidden states to condensed gist tokens, and creating specialized attention masks. Additionally, we introduce the Dependency (Dep) metric to quantify the degree of compression by measuring the reliance on historical tokens during generation. Extensive experiments on four datasets and two models show that LightThinker reduces peak memory usage and inference time, while maintaining competitive accuracy. Our work provides a new direction for improving the efficiency of LLMs in complex reasoning tasks without sacrificing performance. Code is released at https://github.com/zjunlp/LightThinker.