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

초록

대규모 언어 모델(LLM)은 복잡한 추론 작업에서 뛰어난 성능을 보여주지만, 긴 토큰을 생성하는 데 따른 상당한 메모리 및 계산 비용으로 인해 효율성이 저해되고 있습니다. 본 논문에서는 LLM이 추론 과정에서 중간 사고를 동적으로 압축할 수 있는 새로운 방법인 LightThinker를 제안합니다. 인간의 인지 과정에서 영감을 받은 LightThinker는 장황한 사고 단계를 간결한 표현으로 압축하고 원래의 추론 체인을 폐기함으로써, 컨텍스트 창에 저장되는 토큰 수를 크게 줄입니다. 이는 데이터 구성을 통해 모델이 언제 어떻게 압축을 수행할지 학습하고, 은닉 상태를 요약된 요지 토큰으로 매핑하며, 특수한 어텐션 마스크를 생성함으로써 달성됩니다. 또한, 생성 과정에서 역사적 토큰에 대한 의존도를 측정하여 압축 정도를 정량화하는 Dependency(Dep) 메트릭을 도입합니다. 네 가지 데이터셋과 두 가지 모델에 대한 광범위한 실험을 통해 LightThinker가 경쟁력 있는 정확도를 유지하면서도 최대 메모리 사용량과 추론 시간을 줄이는 것을 확인했습니다. 본 연구는 성능 저하 없이 복잡한 추론 작업에서 LLM의 효율성을 개선하는 새로운 방향을 제시합니다. 코드는 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.