LightThinker: Thinking Step-by-Step Compression

2025-02-21

AI Summary

Large language models (LLMs) suffer from high memory consumption and computational latency during complex reasoning due to excessively long chains of thought. To address this, we propose a **Dynamic Thought Compression Framework**, the first to incorporate a human-cognition-inspired, progressive compression mechanism into LLM inference. Our approach introduces a *Dependency* metric to quantify the importance of historical reasoning steps, enabling selective discarding of low-dependency intermediate states. Compressed *gist tokens* are then generated via hidden-state mapping, integrated with customized attention masking and controllable compression timing during training to achieve efficient context reduction. Extensive experiments across four reasoning benchmarks and two mainstream LLMs demonstrate that our method reduces peak memory usage by 38% on average and inference latency by 31%, while preserving accuracy comparable to baseline models. This work establishes a novel paradigm for efficient long-chain reasoning in LLMs.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Cognitive Modeling & Cognitive Systems: Conceptual Inference and Reasoning

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

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.

Problem

• Compress intermediate thoughts in LLMs
• Reduce memory and computational costs
• Maintain accuracy in complex reasoning

Innovation

• Dynamic thought compression
• Gist token mapping
• Dependency metric compression