LightThinker++: 추론 압축에서 메모리 관리까지

LLM이 복잡한 추론을 수행할 때 발생하는 방대한 중간 사고 과정(Thought Traces)은 메모리와 연산 비용을 급격히 증가시킵니다. 이 논문은 중간 사고를 압축된 벡터 형태로 변환하거나 명시적인 메모리 관리 기능을 부여하여, 긴 추론 과정에서도 성능 저하 없이 효율성을 극대화하는 방법을 제시합니다.

Transformer 아키텍처는 입력 시퀀스가 길어질수록 Attention 연산량이 제곱으로 늘어나고 KV Cache 메모리가 선형적으로 증가하는 한계를 가진다. 특히 복잡한 문제를 풀기 위해 Chain-of-Thought(CoT)를 길게 생성할수록 모델의 효율성은 급격히 떨어진다. 기존의 KV Cache 최적화 방식은 중요도가 낮은 토큰을 단순히 삭제하여 복구 불가능한 정보 손실을 초래하는 경우가 많았다.

이 논문은 인간의 작업 기억(Working Memory) 원리를 차용한다. 인간은 복잡한 문제를 풀 때 모든 단어를 기억하지 않고 핵심 결론만 머릿속에 요약해 두었다가, 논리적 병목 지점에 도달했을 때만 세부 사항을 다시 떠올린다. LightThinker는 긴 사고 과정을 소수의 'Gist Tokens'라는 압축된 임베딩 벡터로 요약하여 정보의 병목(Information Bottleneck)을 형성함으로써 메모리 효율을 극대화한다.

LightThinker++는 여기서 한 발 더 나아가 모델에게 '메모리 제어권'을 부여한다. 모델은 현재 정보가 충분히 요약 가능하다고 판단되면 'commit' 명령으로 압축하고, 더 자세한 근거가 필요하면 'expand' 명령으로 원래의 상세 내용을 복원한다. 이를 통해 정보 밀도가 높은 구간에서는 상세하게, 그렇지 않은 구간에서는 압축적으로 사고하여 긴 시퀀스에서도 일관된 추론 능력을 유지한다.

LightThinker는 사고 과정을 k개의 세그먼트로 분할하고 각 세그먼트 사이에 , C, [o]와 같은 특수 토큰을 삽입한다. C는 Gist Tokens 집합으로, 이전 세그먼트의 정보를 압축하여 담는 그릇 역할을 한다. [o] 토큰은 압축된 표현으로부터 다음 추론을 재개하는 신호로 사용된다.

압축 연산 Cmp()는 현재 사고 S_i와 이전 압축 이력을 입력으로 받아 Softmax(mask(QK^T/√d)) 연산을 수행한다. 이때 mask는 현재 사고가 이전 압축 토큰과 질문에만 집중하도록 설계되어, 모델이 핵심 정보를 C_i라는 소수의 벡터로 증류(Distill)하도록 유도한다. 결과값 H = A × V를 통해 정보가 KV Cache에 저장된다.

LightThinker++는 행동 수준의 메모리 관리를 위해 세 가지 기본 연산(Primitives)을 도입한다. commit(R, Z)은 가공되지 않은 사고 R을 요약 Z로 변환하여 아카이브에 저장한다. expand(k)는 k번째 아카이브된 요약을 다시 상세한 R_k로 복원하며, fold(k)는 복원된 상세 내용을 다시 요약본으로 되돌려 컨텍스트를 정화한다.

학습을 위해 Online Thought Synthesis 프레임워크를 사용한다. 강력한 교사 모델이 추론 과정 중에 메모리 관리 명령을 내리도록 유도하고, 환경이 이에 맞춰 컨텍스트를 동적으로 수정(상세 내용 숨기기 또는 노출)하여 모델이 압축된 상태에서도 논리적 일관성을 유지하도록 SFT(Supervised Fine-Tuning)를 수행한다.

Qwen2.5-7B 모델 기준, LightThinker는 Vanilla 모델 대비 피크 토큰 사용량을 70% 절감하고 추론 시간을 26% 단축하면서도 정확도 손실은 1% 내외로 억제했다. 특히 32K 토큰 이상의 긴 시퀀스에서는 추론 시간을 최대 44%까지 줄이는 효과를 보였다.

LightThinker++는 Throughput 설정에서 피크 메모리 사용량을 69.9% 절감하면서 베이스라인 수준의 정확도를 유지했다. 성능을 우선시하는 Budget 설정에서는 피크 메모리를 45.0% 줄이면서도 평균 정확도가 +2.42% 향상되는 결과를 얻었다. 이는 불필요한 노이즈를 제거한 고신호 컨텍스트가 복잡한 추론에 더 유리함을 입증한다.

장기 에이전트 작업(Long-horizon Agentic Tasks) 벤치마크인 xBench-DeepSearch 등에서 LightThinker++는 Vanilla 에이전트 대비 평균 14.8%의 성능 향상을 기록했다. 특히 60라운드 이상의 장기 상호작용에서도 컨텍스트 크기를 30k~40k 수준으로 안정적으로 유지하며 '컨텍스트 폭발' 현상을 효과적으로 방지했다.

LightThinker는 표현 계층(Representation-level)에서 Attention Mask를 조작하여 정보 병목을 강제하는 구조다. 반면 LightThinker++는 행동 계층(Behavioral-level)에서 모델이 스스로 가시성 상태(Visibility State) σ_k ∈ {archive, active}를 결정하게 한다. σ_k가 archive일 때는 요약 Z_k가, active일 때는 원본 R_k가 컨텍스트에 노출된다.

학습 데이터 정제 시 'Memory Lifecycle' 제약 조건을 적용한다. fold 연산은 반드시 이전에 수행된 expand 연산에 대응되어야 하며, 동일한 메모리 조작이 연속으로 발생하는 것을 방지하는 Anti-Jitter Heuristics를 통해 모델이 목적성 있는 메모리 관리를 수행하도록 보장한다.

정보 밀도 분석 결과, LightThinker++가 생성한 사고 세그먼트는 일반적인 CoT보다 평균 길이가 길고(38자 → 51자) 논리적 앵커의 밀도가 높다. 이는 모델이 압축을 염두에 두고 더 의미론적으로 완결된 단위로 사고하도록 학습되었음을 시사한다.

GSM8K와 같은 수치 데이터가 중요한 작업에서 두 번째 압축 단계가 필수적인 숫자 정보를 누락시켜 계산 오류를 범하는 사례가 발견됐다. 현재의 압축 메커니즘이 미세한 수치 정보에 대해 충분히 민감하지 못할 수 있음을 보여준다.