연속 잠재 공간에서 추론하는 대규모 언어 모델 훈련

Training Large Language Models to Reason in a Continuous Latent Space

December 9, 2024

저자: Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston, Yuandong Tian

cs.AI

초록

대형 언어 모델(Large language models, LLMs)은 "언어 공간"에서 추론하는 것으로 제한되어 있으며, 일반적으로 복잡한 추론 문제를 해결하기 위해 사고 연쇄(chain-of-thought, CoT)로 추론 과정을 표현합니다. 그러나 우리는 언어 공간이 항상 추론에 최적이 아닐 수 있다고 주장합니다. 예를 들어, 대부분의 단어 토큰은 주로 텍스트 일관성을 위한 것이며 추론에 필수적이지 않습니다. 반면 일부 중요한 토큰은 복잡한 계획이 필요하며 LLM에게 엄청난 도전을 제기합니다. 자연어 대신 제한되지 않은 잠재 공간에서 LLM 추론의 잠재력을 탐색하기 위해 코코넛(Coconut, Chain of Continuous Thought)이라는 새로운 패러다임을 소개합니다. 우리는 LLM의 마지막 숨겨진 상태를 추론 상태의 표현으로 활용하고("연속적 사고"라고 함), 이를 단어 토큰으로 디코딩하는 대신 연속적 공간에서 직접 LLM에게 후속 입력 임베딩으로 다시 공급합니다. 실험 결과, 코코넛은 여러 추론 작업에서 LLM을 효과적으로 보완할 수 있음을 보여줍니다. 이 새로운 잠재 추론 패러다임은 신흥 고급 추론 패턴으로 이어지며, 연속적 사고는 다양한 대안적 다음 추론 단계를 인코딩할 수 있어 모델이 CoT와 같이 단일 결정적 경로로 일찍 커밋하는 대신 문제를 해결하기 위해 너비 우선 탐색(BFS)을 수행할 수 있습니다. 코코넛은 계획 중에 상당한 되감기가 필요한 특정 논리 추론 작업에서 CoT보다 우수한 성능을 보이며 추론 중에 적은 사고 토큰을 필요로 합니다. 이 결과들은 잠재 추론의 가능성을 입증하고 향후 연구에 대한 가치 있는 통찰을 제공합니다.

English

Large language models (LLMs) are restricted to reason in the "language space", where they typically express the reasoning process with a chain-of-thought (CoT) to solve a complex reasoning problem. However, we argue that language space may not always be optimal for reasoning. For example, most word tokens are primarily for textual coherence and not essential for reasoning, while some critical tokens require complex planning and pose huge challenges to LLMs. To explore the potential of LLM reasoning in an unrestricted latent space instead of using natural language, we introduce a new paradigm Coconut (Chain of Continuous Thought). We utilize the last hidden state of the LLM as a representation of the reasoning state (termed "continuous thought"). Rather than decoding this into a word token, we feed it back to the LLM as the subsequent input embedding directly in the continuous space. Experiments show that Coconut can effectively augment the LLM on several reasoning tasks. This novel latent reasoning paradigm leads to emergent advanced reasoning patterns: the continuous thought can encode multiple alternative next reasoning steps, allowing the model to perform a breadth-first search (BFS) to solve the problem, rather than prematurely committing to a single deterministic path like CoT. Coconut outperforms CoT in certain logical reasoning tasks that require substantial backtracking during planning, with fewer thinking tokens during inference. These findings demonstrate the promise of latent reasoning and offer valuable insights for future research.