LightThinker: 段階的思考による圧縮

LightThinker: Thinking Step-by-Step Compression

February 21, 2025

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

cs.AI

要旨

大規模言語モデル（LLM）は複雑な推論タスクにおいて顕著な性能を示していますが、長いトークンを生成する際の膨大なメモリと計算コストがその効率を阻害しています。本論文では、LLMが推論中に中間的な思考を動的に圧縮することを可能にする新しい手法、LightThinkerを提案します。人間の認知プロセスに着想を得たLightThinkerは、冗長な思考ステップをコンパクトな表現に圧縮し、元の推論チェーンを破棄することで、コンテキストウィンドウに保存されるトークン数を大幅に削減します。これは、データ構築を通じてモデルに圧縮のタイミングと方法を学習させ、隠れ状態を凝縮された要約トークンにマッピングし、特殊なアテンションマスクを作成することで実現されます。さらに、生成中に過去のトークンへの依存度を測定することで圧縮の程度を定量化するDependency（Dep）メトリックを導入します。4つのデータセットと2つのモデルを用いた広範な実験により、LightThinkerがピークメモリ使用量と推論時間を削減しつつ、競争力のある精度を維持することが示されました。本研究は、性能を犠牲にすることなく、複雑な推論タスクにおけるLLMの効率を向上させる新たな方向性を提供します。コードはhttps://github.com/zjunlp/LightThinkerで公開されます。

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 will be released at https://github.com/zjunlp/LightThinker.