LightThinker: Compressão com Raciocínio Passo a Passo

LightThinker: Thinking Step-by-Step Compression

February 21, 2025

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

cs.AI

Resumo

Modelos de linguagem de grande escala (LLMs) demonstraram desempenho notável em tarefas complexas de raciocínio, mas sua eficiência é prejudicada pelos custos substanciais de memória e computação associados à geração de tokens extensos. Neste artigo, propomos o LightThinker, um método inovador que permite que LLMs comprimam dinamicamente pensamentos intermediários durante o raciocínio. Inspirado pelos processos cognitivos humanos, o LightThinker comprime etapas verbosas de pensamento em representações compactas e descarta as cadeias de raciocínio originais, reduzindo significativamente o número de tokens armazenados na janela de contexto. Isso é alcançado ao treinar o modelo sobre quando e como realizar a compressão por meio da construção de dados, mapeamento de estados ocultos para tokens de resumo condensados e criação de máscaras de atenção especializadas. Além disso, introduzimos a métrica de Dependência (Dep) para quantificar o grau de compressão ao medir a dependência de tokens históricos durante a geração. Experimentos extensos em quatro conjuntos de dados e dois modelos mostram que o LightThinker reduz o uso máximo de memória e o tempo de inferência, mantendo uma precisão competitiva. Nosso trabalho fornece uma nova direção para melhorar a eficiência de LLMs em tarefas complexas de raciocínio sem sacrificar o desempenho. O código será disponibilizado em 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.