Training Large Language Models to Reason in a Continuous Latent Space

TL;DR

Proposes a new paradigm where LLMs reason using continuous latent states (continuous thought) directly as the next input embedding instead of language tokens.

Problem

Traditional chain-of-thought (CoT) reasoning operates in language space, where most word tokens are only needed for textual coherence and are not essential for reasoning. Moreover, critical reasoning steps require complex planning but language tokens are inefficient for representing them.

Approach

Coconut (Chain of Continuous Thought) treats the last hidden state of the LLM as a 'continuous thought' and feeds it back as the next input embedding directly, without decoding into words. This allows the model to reason in latent space without the constraints of language tokens. Additionally, continuous thoughts can encode multiple alternatives simultaneously, enabling a breadth-first search (BFS)-like reasoning pattern instead of the deterministic path in CoT.

Results & Contribution

On logical reasoning tasks (e.g., the ProntoQA problem), Coconut achieves higher accuracy than CoT, especially on problems requiring search during planning. It also improves efficiency by reducing the number of reasoning steps. This work demonstrates that LLM reasoning does not have to be confined to language space and introduces a new reasoning paradigm.