LightThinker: Stapsgewijs Denken Compressie

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

Samenvatting

Grote taalmodellen (LLMs) hebben opmerkelijke prestaties getoond in complexe redeneertaken, maar hun efficiëntie wordt belemmerd door de aanzienlijke geheugen- en rekenkosten die gepaard gaan met het genereren van lange tokens. In dit artikel stellen we LightThinker voor, een nieuwe methode die LLMs in staat stelt om tijdens het redeneren dynamisch tussenliggende gedachten te comprimeren. Geïnspireerd door menselijke cognitieve processen, comprimeert LightThinker uitgebreide denkstappen tot compacte representaties en verwijdert de oorspronkelijke redeneerketens, waardoor het aantal tokens dat in het contextvenster wordt opgeslagen aanzienlijk wordt verminderd. Dit wordt bereikt door het model te trainen op wanneer en hoe compressie moet worden uitgevoerd via dataconstructie, het toewijzen van verborgen toestanden aan gecondenseerde kern-tokens, en het creëren van gespecialiseerde aandachtmaskers. Daarnaast introduceren we de Dependency (Dep) metriek om de mate van compressie te kwantificeren door de afhankelijkheid van historische tokens tijdens het genereren te meten. Uitgebreide experimenten op vier datasets en twee modellen tonen aan dat LightThinker het piekgeheugengebruik en de inferentietijd vermindert, terwijl het competitieve nauwkeurigheid behoudt. Ons werk biedt een nieuwe richting voor het verbeteren van de efficiëntie van LLMs in complexe redeneertaken zonder in te leveren op prestaties. Code zal worden vrijgegeven op 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.