LightThinker: Compresión mediante pensamiento paso a paso

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

Resumen

Los modelos de lenguaje de gran escala (LLMs, por sus siglas en inglés) han demostrado un rendimiento notable en tareas de razonamiento complejo, pero su eficiencia se ve limitada por los costos significativos de memoria y computación asociados con la generación de tokens extensos. En este artículo, proponemos LightThinker, un método novedoso que permite a los LLMs comprimir dinámicamente pensamientos intermedios durante el razonamiento. Inspirado en los procesos cognitivos humanos, LightThinker comprime pasos de pensamiento verbosos en representaciones compactas y descarta las cadenas de razonamiento originales, reduciendo así significativamente el número de tokens almacenados en la ventana de contexto. Esto se logra entrenando al modelo sobre cuándo y cómo realizar la compresión mediante la construcción de datos, mapeando estados ocultos a tokens de resumen condensados y creando máscaras de atención especializadas. Además, introducimos la métrica de Dependencia (Dep) para cuantificar el grado de compresión midiendo la dependencia de tokens históricos durante la generación. Experimentos exhaustivos en cuatro conjuntos de datos y dos modelos muestran que LightThinker reduce el uso máximo de memoria y el tiempo de inferencia, manteniendo una precisión competitiva. Nuestro trabajo proporciona una nueva dirección para mejorar la eficiencia de los LLMs en tareas de razonamiento complejo sin sacrificar el rendimiento. El código se publicará en 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.