Source-linked AI summary
LightThinker: Thinking Step-by-Step Compression
Jintian Zhang, Yuqi Zhu, Mengshu Sun, Yujie Luo, Shuofei Qiao, Lun Du, Da Zheng, Huajun Chen, Ningyu Zhang
TL;DR
Lengthy reasoning improves complex problem solving but creates severe memory and computation costs as attention context and KV caches grow. LightThinker dynamically compresses intermediate thoughts into compact gist-token representations and trains attention patterns for continued generation from compressed history. Across four datasets and two models, it reduces resource use while retaining competitive accuracy, including a 70% peak-token reduction and 26% lower inference time on Qwen.
Problem
Lengthy reasoning improves complex problem solving but sharply increases Transformer attention and KV-cache costs as context length grows.
Method
LightThinker trains LLMs to segment thoughts, compress them into gist-token cache representations, and generate from compressed history using specialized attention masks.
Results
70% lower peak token usage and 26% lower inference time on Qwen accompany a 1% accuracy drop versus Vanilla across four datasets and two models.
Takeaways & Limitations
LightThinker offers a supported balance between reasoning efficiency and accuracy by reducing memory overhead and inference time during complex tasks.
Takeaways & Limitations
Cache-token count is fixed between training and inference, and the ability of representations trained with one count to generalize to another remains uncertain.
Abstract
from arXiv · showhide
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.
1 Introduction
LightThinker addresses the memory and computation costs of lengthy reasoning by dynamically compressing intermediate thoughts while preserving information needed for subsequent inference. The method is motivated by separating fluency-oriented tokens from critical reasoning content and is evaluated with competitive accuracy and lower resource use.
- Motivation: Quadratic attention cost and linear KV-cache growth make long reasoning increasingly expensive in memory and computation.For Qwen32B, a context length of 104 makes the KV Cache comparable in size to the model itself.
- LightThinker: LightThinker compresses each lengthy thought into a compact representation, discards the original chain, and continues reasoning from the compressed content.The approach aims to reduce stored context tokens, memory overhead, and computational costs.
- Training: The method trains LLMs to learn when and how to compress through reconstructed data and specialized attention masks.Its compression process maps thought information into compact representations used for continued generation.
- Evaluation: The Dependency metric measures how many historical tokens generated tokens rely on, with lower values indicating stronger compression.Dep quantifies reliance on the original long context during reasoning.
- Results: With Qwen, LightThinker reduces peak token usage by 70% and inference time by 26% versus Vanilla, with a 1% accuracy drop.These results come from experiments across four datasets and two models.
2 Background
LLM reasoning has progressed from short, fast responses toward longer step-by-step and iterative thinking, improving complex problem solving but increasing token consumption. This growth creates substantial inference-time memory and computational challenges for Transformer models.
- Slow Thinking: Chain-of-Thought decomposes complex problems into sequential subproblems, while o1-like thinking adds trial-and-error, backtracking, correction, and iteration.These approaches represent a shift from fast-thinking toward slow-thinking reasoning.
- Metrics: Figure 3 compares context length with generated-token counts across methods using Peak Token as the curve maximum and Dep as its area under the curve.Peak Token reflects a momentary maximum, whereas Dep summarizes the full inference process.
- LightThinker Overview: Figure 2 illustrates Vanilla and LightThinker attention masks across training and inference in a three-step reasoning example.The figure also indicates that w denotes the matrix size.
- Slow Thinking: Token consumption increases progressively from System 1 to Chain-of-Thought and o1-like thinking modes.Supervised fine-tuning on carefully constructed data can instill slow-thinking behavior in models.
- Inference Challenges: As generated context grows, Transformer attention faces increasing memory overhead because each token’s Key and Value are cached at every layer.The background identifies memory growth as a central inference challenge for long reasoning.
3 Methodology
LightThinker trains an LLM to segment reasoning into thoughts, compress each thought into a small set of cache tokens, and generate subsequent thoughts using compressed history. Its data construction and attention masks support compression during both training and inference.
- Data Reconstruction: Thoughts can be segmented at token level or thought level, with thought-level segmentation aligning compression to semantic boundaries.Token-level segmentation is simpler but may ignore semantic boundaries.
- Compression Design: Hidden-state compression stores a thought in the hidden states of a few special gist tokens without requiring an additional encoding model.The method focuses on retaining information essential for subsequent reasoning rather than preserving every detail.
- Data Reconstruction: LightThinker inserts compression triggers, cache tokens, and output tokens between adjacent thought subsequences.Cache tokens are special gist tokens, while the output token enables continued generation from compressed content.
- Attention Masks: During compression, cache tokens attend to the question, previous compressed content, and the current thought.This attention pattern teaches the model to encode the current thought using available reasoning context.
- Attention Masks: During generation, output tokens attend only to the question and previous compressed content, allowing reasoning to continue without the original long thought.The generation operation therefore conditions on compressed history rather than the discarded chain.
- Training: Training maximizes the model’s next-token probability while masking inputs and special tokens from prediction.The process uses an attention mask to teach both compression and comprehension of compressed content.
4 Experiments
Experiments evaluate LightThinker across four datasets and two models, comparing accuracy and efficiency against multiple baselines. Results show substantial reductions in memory, dependency, token generation, and inference time, with trade-offs shaped by cache size and segmentation.
- Experimental Settings: Experiments use Qwen2.5-7B and Llama3.1-8B on GSM8K, MMLU, GPQA, and BBH, measuring accuracy, time, peak tokens, and dependency.Vanilla, H2O, SepLLM, AnLLM, and CoT provide comparison settings.
- Main Results: LightThinker reduces inference time by an average of 52% on Qwen and 41% on Llama compared with H2O at similar compression rates.It also retains higher accuracy and faster inference than AnLLM.
- Main Results: LightThinker sacrifices 1% accuracy on Qwen and 6% on Llama while reducing peak tokens by 70% on both models.The reported average dependency reductions are 78% on Qwen and 74% on Llama.
- Ablation: Thought-level segmentation improves accuracy over token-level segmentation by 6.2% on Qwen and 5.6% on Llama.The authors link token-level segmentation to loss of semantic boundaries.
- Efficiency: Under the same memory budget, LightThinker reduces Qwen inference time by an average of 2.5× compared with Vanilla across four datasets.The memory reduction permits larger batch sizes under the same budget, improving throughput.
- Efficiency: LightThinker is the only compared method reducing generated tokens relative to Vanilla, by 15% on Qwen and 13% on Llama.The analysis attributes part of its faster inference to this lower token generation.
5 Related Work
Related work accelerates LLM inference by reducing model precision, generating fewer tokens, or reducing the KV cache. KV-cache methods either prune tokens or merge information, with different training and inference-intervention requirements.
- Inference acceleration methods include model quantization, fewer-token generation, and KV-cache reduction.
- The table caption identifies T as AnLLM’s attention mask and F as LightThinker’s attention mask in Qwen ablations.
- KV-cache reduction uses pruning-based selection in discrete space or merging-based compression in continuous space.
- Pruning strategies retain important tokens with eviction policies, whereas merging strategies compress historical information into anchor tokens.
- Pruning is training-free but applies eviction at every generated token; merging is training-based and lets the LLM decide when to apply eviction.
6 Conclusion
LightThinker improves LLM efficiency in complex reasoning by dynamically compressing intermediate thoughts during generation. The method reduces memory and computational costs while maintaining competitive accuracy, and Dep quantifies compression across accelerating methods.
- LightThinker dynamically compresses intermediate thoughts during generation to improve LLM efficiency in complex reasoning tasks.
- The method trains LLMs when and how to compress verbose thought steps into compact representations.
- LightThinker reduces memory overhead and computational costs while maintaining competitive accuracy.
- The Dependency (Dep) metric quantifies compression across different accelerating methods.
- Extensive experiments demonstrate that LightThinker balances efficiency and performance.
Limitations
The paper identifies limitations in fixed cache-token representations, rule-based segmentation, and unassessed task domains.
- The number of cache tokens is fixed during training and must remain consistent during inference.The generalization of representations trained with one token count to inference scenarios requiring more tokens is uncertain.
- LightThinker’s segmentation function relies on relatively simplistic rule-based methods.The paper suggests investigating more advanced segmentation strategies.
- Performance remains unassessed for novel generation, code generation, and multi-turn dialogue.
A Metric: Dependency
Dependency measures the information dependencies accumulated during generation, enabling fairer comparisons between dynamic and traditional KV-cache compression methods than peak token count alone.
- Peak token count is insufficient when compression is dynamic, because memory usage may be brief rather than sustained.LightThinker can occasionally exceed H2O in peak token count even though its peak memory usage occurs only momentarily.
- Output compression prevents a globally preset compression ratio because the number of generated tokens is unknown.Earlier KV-cache methods often compress only prompt content and assume a fixed prompt length.
- Dependency quantifies the total amount of historical information each generated token relies on during generation.It is also described geometrically as the area under the generated-token/context-length curve.
- Dependency is computed from a definition or geometric interpretation using prompt length, output length, and cache context length.For Vanilla, the curve forms a right trapezoid; H2O combines a trapezoid and rectangle; LightThinker and AnLLM require iterative computation.
A.3 Application
Dependency interprets higher values as greater information usage and lower values as stronger effective compression, while the compression operation stores selected thought information in compact cache representations.
- A.3 Application: A higher Dependency means more tokens are considered during generation, whereas a lower value indicates a higher effective compression ratio.
- B Mathematical Description of Compression: During compression, the context is divided into uncompressed preceding content, the thought sequence being compressed, and the sequence storing compressed content.
- B Mathematical Description of Compression: The attention operation produces compact hidden states for the compressed sequence by weighting values from preceding, thought, and compressed-content sequences.Training teaches the model to selectively retain useful information from the current thought sequence before storing the result in the KV cache.
- A.3 Application: H2O and SepLLM are training-free methods with configured retention parameters, whereas AnLLM and LightThinker use training-based dynamic compression.AnLLM stores historical content in anchor tokens, while LightThinker shares a similar overall training-based approach.
C.3 Training Details
The experiments train Vanilla, AnLLM, and LightThinker under closely specified optimization and hardware settings, while using shared prompts across comparison methods.
- C.3 Training Details: Vanilla and AnLLM train for 5 epochs, while LightThinker trains for 6 epochs on B17K with R1-Distill.All use maximum length 4096 and cosine warmup with warmup_ratio 0.05.
- C.3 Training Details: Training runs on 4 A800 GPUs with DeepSpeed ZeRO3 offload, per-GPU batch size 5, and gradient accumulation of 4.These settings produce a global batch size of 80.
- C.3 Training Details: Vanilla, H2O, SepLLM, AnLLM, and LightThinker use the same prompts, with task-specific prompts for downstream evaluations.MMLU and GPQA answer options are randomized.
C.5 Additional Results
The paper situates LightThinker among token-reduction and KV-cache methods and reports additional comparisons across generated-token counts, compressed lengths, and attention-mask designs.
- C.5 Additional Results: Figure 8 compares generated-token counts across two models and multiple datasets, while Figure 10 shows LightThinker’s compressed-length distributions.
- C.5 Additional Results: Figure 7 presents the attention masks used by the baselines in Table 4.
- D Related Work: Inference-acceleration research is organized around model quantization, generating fewer tokens, and reducing the KV cache.The paper focuses its related-work discussion on the latter two categories.
- D Related Work: Generating-fewer-token methods include prompt engineering, instruction finetuning, and reinforcement learning that guide models toward shorter reasoning traces.
- D Related Work: KV-cache reduction includes pruning important discrete tokens and merging historical information into continuous anchor tokens.Training-free pruning applies eviction at every generated token, whereas training-based merging lets the model decide when to apply it.
E.1 Difference between LightThinker and AnLLM
LightThinker differs from AnLLM by separating compression from subsequent generation and by giving compression access to broader context. Unlike manually designed token eviction, it merges prior content into learned continuous representations while retaining discrete reasoning tokens.
- Decoupling Generation and Compression: LightThinker decouples compression and generation: [ci] compresses historical information, while [o] reasons from the compressed content.In AnLLM, [ci] performs both functions, coupling compression with generation.
- Context Visibility during Compression: During compression, LightThinker accesses the input, historical compressed content, and the current thought, whereas AnLLM accesses only the current thought.The broader visibility is intended to enhance contextual understanding, and ablations report significant performance improvements from these designs.
- KV Cache Compression: LightThinker merges previous tokens in continuous space into new representations, rather than selecting tokens through a manually designed eviction policy like H2O.The content and manner of merging are autonomously determined by the language model.
- Cache-Size Behavior: Smaller cache sizes make LightThinker generate more tokens, which the authors attribute to information loss that prompts regeneration of earlier content.This behavior identifies a practical boundary: aggressive compression can increase token generation through repetition.
- Interpretability and Generalization: LightThinker combines continuous compression with discrete-token reasoning, unlike implicit CoT methods that reason entirely in continuous space.The paper associates discrete tokens with better interpretability and reports promising out-of-domain generalization in experiments.
- Reasoning Acceleration and Training: Compared with implicit CoT approaches, LightThinker reduces historical-token dependence while using standard supervised fine-tuning with modified attention masks.The comparison states that LightThinker requires neither specialized training data nor architectural changes.