arXiv:2502.15589cs.CLcs.AI2025-02EMNLP被引 116

让大模型自动压缩推理过程,省内存还提速。

LightThinker: Thinking Step-by-Step Compression

论文配图:LightThinker: Thinking Step-by-Step Compression
图 1 · 摘自论文原文
  • 动态压缩推理中间步骤,减少上下文存储
  • 峰值内存下降42%,推理时间减少38%仍保持准确
  • 适合需要高效推理的复杂任务场景

大型语言模型在复杂推理任务中表现优异,但生成长序列令牌带来的内存和计算开销限制了其效率。本文提出LightThinker,一种使LLM在推理过程中动态压缩中间思考的方法。受人类认知启发,LightThinker将冗长的思维步骤压缩为紧凑表示,并丢弃原始推理链,显著减少上下文窗口中的令牌数量。通过数据构建、隐藏状态映射到浓缩概要令牌以及专用注意力掩码训练模型判断何时何地进行压缩。此外,引入依赖度(Dep)指标量化压缩程度,衡量生成时对历史令牌的依赖性。在四个数据集和两个模型上的实验表明,LightThinker有效降低峰值内存使用和推理时间,同时保持竞争力的准确率。本工作为提升LLM在复杂推理中的效率提供了新方向。代码已开源:https://github.com/zjunlp/LightThinker。

原文摘要 · Abstract (English)

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.

大模型压缩推理优化上下文管理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。