지도 학습 기반 사고 상태를 활용한 잠재적 추론

Latent Reasoning with Supervised Thinking States

체인-오브-쏘트(CoT)를 활용한 추론은 대규모 언어 모델(LLM)이 복잡한 문제를 해결하는 데 도움을 주지만, 긴 추론 과정을 생성하기 때문에 상당한 계산 비용이 발생합니다. 본 논문에서는 입력 처리 중에 추론을 수행하는 방법인 '사고 상태(Thinking States)'를 제안합니다. 구체적으로, 사고 상태는 입력 토큰 몇 개마다 '생각 토큰' 시퀀스를 생성하고, 이 토큰들을 임베딩 공간으로 변환하여 다음 입력 토큰에 추가합니다. 이 방법은 다음과 같은 두 가지 주요 이점을 제공합니다. 첫째, CoT의 반복적인 특성을 효과적으로 반영하면서도, '생각 토큰'을 입력 처리 중에 생성합니다. 둘째, '생각'이 토큰으로 표현되기 때문에, 자연어 지도 학습을 통해 학습할 수 있으며, 교사 강제(teacher-forcing)를 사용하여 병렬 처리가 가능합니다. 실험 결과, 사고 상태는 여러 추론 작업에서 다른 잠재적 추론 방법보다 우수한 성능을 보였으며, 수학 문제에서는 CoT와의 성능 격차를 줄이고, 2-Hop QA에서는 CoT와 유사한 성능을 보이면서도 지연 시간(latency)을 개선했습니다. 또한, 상태 추적 작업에서 사고 상태는 CoT보다 더 강력한 추론 능력을 보여주며, 훈련 중에 보지 못한 더 긴 시퀀스에 대한 추론이 가능했습니다.

Reasoning with a chain-of-thought (CoT) enables Large Language Models (LLMs) to solve complex tasks but incurs significant inference costs due to the generation of long rationales. We propose Thinking States, a method that performs reasoning {\em while} the input is processing. Specifically, Thinking States generates sequences of thinking tokens every few input tokens, transforms the thoughts back into embedding space, and adds them to the following input tokens. This has two key advantages. First, it captures the recurrent nature of CoT, but where the thought tokens are generated as input is processing. Second, since the thoughts are represented as tokens, they can be learned from natural language supervision, and using teacher-forcing, which is parallelizable. Empirically, Thinking States outperforms other latent reasoning methods on multiple reasoning tasks, narrowing the gap to CoT on math problems, and matching its performance on 2-Hop QA with improved latency. On state-tracking tasks, we show Thinking States leads to stronger reasoning behavior than CoT, successfully extrapolating to longer sequences than seen during training.