LightThinker : Penser étape par étape la compression

LightThinker: Thinking Step-by-Step Compression

February 21, 2025

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

cs.AI

Résumé

Les grands modèles de langage (LLMs) ont démontré des performances remarquables dans les tâches de raisonnement complexe, mais leur efficacité est entravée par les coûts substantiels en mémoire et en calcul associés à la génération de tokens longs. Dans cet article, nous proposons LightThinker, une méthode novatrice qui permet aux LLMs de compresser dynamiquement les pensées intermédiaires lors du raisonnement. Inspiré par les processus cognitifs humains, LightThinker compresse les étapes de pensée verbeuses en représentations compactes et élimine les chaînes de raisonnement originales, réduisant ainsi significativement le nombre de tokens stockés dans la fenêtre contextuelle. Cela est réalisé en entraînant le modèle sur quand et comment effectuer la compression grâce à la construction de données, en mappant les états cachés à des tokens résumés, et en créant des masques d'attention spécialisés. De plus, nous introduisons la métrique de Dépendance (Dep) pour quantifier le degré de compression en mesurant la dépendance aux tokens historiques lors de la génération. Des expériences approfondies sur quatre jeux de données et deux modèles montrent que LightThinker réduit l'utilisation de mémoire maximale et le temps d'inférence, tout en maintenant une précision compétitive. Notre travail ouvre une nouvelle voie pour améliorer l'efficacité des LLMs dans les tâches de raisonnement complexe sans sacrifier les performances. Le code sera disponible à l'adresse 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.