LightThinker: Пошаговое сжатие мышления

LightThinker: Thinking Step-by-Step Compression

February 21, 2025

Авторы: Jintian Zhang, Yuqi Zhu, Mengshu Sun, Yujie Luo, Shuofei Qiao, Lun Du, Da Zheng, Huajun Chen, Ningyu Zhang

cs.AI

Аннотация

Крупные языковые модели (LLMs) демонстрируют выдающиеся результаты в сложных задачах рассуждения, однако их эффективность ограничивается значительными затратами памяти и вычислительных ресурсов, связанных с генерацией длинных токенов. В данной статье мы предлагаем LightThinker — новый метод, который позволяет LLMs динамически сжимать промежуточные мысли в процессе рассуждения. Вдохновленный когнитивными процессами человека, LightThinker сжимает многословные шаги рассуждения в компактные представления и отбрасывает исходные цепочки рассуждений, тем самым значительно сокращая количество токенов, хранящихся в контекстном окне. Это достигается за счет обучения модели на том, когда и как выполнять сжатие, путем построения данных, отображения скрытых состояний в сжатые токены-суть и создания специализированных масок внимания. Кроме того, мы вводим метрику Dependency (Dep) для количественной оценки степени сжатия, измеряя зависимость от исторических токенов в процессе генерации. Многочисленные эксперименты на четырех наборах данных и двух моделях показывают, что LightThinker снижает пиковое использование памяти и время вывода, сохраняя при этом конкурентоспособную точность. Наша работа открывает новое направление для повышения эффективности LLMs в сложных задачах рассуждения без ущерба для производительности. Код будет доступен по адресу 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.