LightThinker: Thinking Step-by-Step Compression
Jintian Zhang, Yuqi Zhu, Mengshu Sun, Yujie Luo, Shuofei Qiao, Lun Du, Da Zheng, Huajun Chen, Ningyu Zhang
2025-02-24
Summary
This paper talks about LightThinker, a new method that helps large language models (LLMs) think more efficiently by compressing their thoughts as they reason, similar to how humans focus on key ideas when solving problems.
What's the problem?
Large language models are great at complex reasoning tasks, but they use up a lot of memory and computing power because they generate and store long chains of thoughts. This makes them slow and expensive to run, especially for complicated problems that require a lot of thinking steps.
What's the solution?
The researchers created LightThinker, which teaches LLMs to summarize their thoughts into shorter, more compact versions as they go along. It does this by training the model to know when and how to compress its thoughts, turning detailed reasoning steps into condensed 'gist tokens' and using special attention masks to focus on the most important information. They also came up with a new way to measure how well the compression is working, called the Dependency metric.
Why it matters?
This matters because it makes LLMs much more efficient without making them less accurate. In tests, LightThinker used 70% less memory at its peak and was 26% faster, while still being just as good at solving problems. This could make advanced AI systems cheaper to run and more practical to use in real-world applications, potentially leading to breakthroughs in fields that require complex reasoning.
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.