Source-linked AI summary
Training Large Language Models to Reason in a Continuous Latent Space
Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston, Yuandong Tian
TL;DR
Coconut addresses limitations of language-space reasoning, where most tokens contribute little to reasoning while critical tokens require complex planning. It feeds continuous thoughts directly as embeddings, enabling multiple alternative reasoning paths, and improves performance on reasoning tasks while offering a better accuracy–efficiency trade-off.
Problem
Language-space reasoning allocates similar computation to tokens despite varying reasoning demands, while many tokens mainly support fluency and critical tokens require difficult planning.
Method
Coconut feeds the LLM’s last hidden state back as the next input embedding, using continuous thoughts to represent reasoning states and alternative next steps.
Results
Coconut achieves 34.1% accuracy on GSM8k with six continuous thoughts, versus 16.5% for No-CoT, and improves the accuracy–efficiency trade-off relative to traditional CoT.
Takeaways & Limitations
Continuous latent reasoning enables breadth-first search over possible reasoning paths instead of premature commitment to one deterministic CoT trajectory.
Takeaways & Limitations
Coconut’s multi-stage training is effective, but better and more general latent-reasoning strategies are needed, especially without language-chain supervision.
Abstract
from arXiv · showhide
Large language models (LLMs) are typically constrained to reason in the language space, where they express the reasoning process through a chain-of-thought (CoT) to solve complex problems. However, the language space may not always be optimal for reasoning. Most word tokens primarily ensure textual coherence and are not essential for reasoning, while some critical tokens require complex planning and pose challenges to LLMs. To explore the potential of reasoning beyond language, we introduce a new paradigm called Coconut (Chain of Continuous Thought). Coconut utilizes the last hidden state of the LLM as a representation of the reasoning state, termed "continuous thought." Instead of decoding this state into words, we feed it back to the model as the next input embedding directly in the continuous space. This latent reasoning paradigm enables an advanced reasoning pattern, where continuous thoughts can encode multiple alternative next steps, allowing the model to perform a breadth-first search (BFS) rather than committing prematurely to a single deterministic path as in CoT. Coconut outperforms CoT on logical reasoning tasks that require substantial search during planning and achieves a better trade-off between accuracy and efficiency.
1 Introduction
Coconut addresses limitations of language-based reasoning by feeding continuous hidden-state representations back as inputs, enabling latent-space reasoning and search. It improves reasoning accuracy on logical tasks while generating fewer tokens.
- Language-based reasoning allocates similar computation to every token, although many tokens mainly support fluency while critical tokens require complex planning.
- Coconut replaces word-token reasoning with continuous thoughts formed from the last hidden state and reused as the next input embedding.
- Continuous thoughts can encode multiple potential next steps, allowing breadth-first search that progressively eliminates incorrect reasoning paths.
- On logical reasoning tasks including ProntoQA and ProsQA, Coconut and variants surpass language-based CoT while generating significantly fewer inference tokens.
2 Related Work
Related work establishes CoT as an effective language-based reasoning method but highlights its difficulty with planning and search. Prior latent-reasoning approaches manipulate hidden computation or add tokens, with limits on general expressivity.
- CoT generates intermediate language reasoning and can increase transformer effective depth by feeding generated outputs back into the model.
- Autoregressive CoT generation is challenging to extend to complex problems requiring planning and search.
- Prior latent-reasoning studies recover or intervene on hidden representations and identify parallel latent reasoning paths.
- Pause and filler-token methods improve some tasks, but filler tokens may work mainly for highly parallelizable problems and may not extend model expressivity like CoT.
3 Coconut: Chain of Continuous Thought
Coconut alternates between language mode and latent mode, using normalized last hidden states as continuous-thought embeddings. A multi-stage curriculum replaces language reasoning steps with continuous thoughts and trains on later tokens.
- Method Overview: In language mode, Coconut autoregressively generates tokens; in latent mode, it feeds the last hidden state back as the next input embedding.
- Method Overview: During latent mode, the last hidden state represents the current reasoning state and is termed a continuous thought.
- Training Procedure: Training begins with regular CoT instances, then replaces the first k reasoning steps with k × c continuous thoughts at stage k.
- Training Procedure: The training objective masks questions and latent thoughts while optimizing negative log-likelihood on future reasoning and answer tokens.
- Training Details: Continuous thoughts are fully differentiable, but sequential forward passes challenge parallel training efficiency.
- Inference Process: The inference process must determine when to switch between latent and language modes, with <bot> marking latent-mode onset.
4 Continuous Space Enables Latent Tree Search
On ProsQA, Coconut supports latent reasoning that explores multiple candidate paths before committing, improving planning-intensive reasoning over language-space CoT. Its latent search broadens early exploration, then focuses decisions as paths approach terminal states.
- Experimental Setup: Coconut outperforms language-space CoT on ProsQA, a logical reasoning dataset whose DAGs require finding valid paths through complex exploration paths.ProsQA is designed to demand sophisticated planning and search strategies.
- Overall Results: Increasing the number of continuous thoughts improves answer accuracy and raises correct reasoning outcomes while reducing hallucinations and wrong targets.These errors typically arise when the model makes mistakes early in the reasoning process.
- Overall Results: A ProsQA case study shows CoT hallucinating an unsupported edge, Coconut with k=1 ending at an irrelevant node, and Coconut with k=2 solving correctly.The comparison illustrates how additional latent reasoning steps can avoid an early dead end.
- Interpreting the Latent Reasoning as Tree Search: Coconut’s latent search assigns probabilities to candidate concepts as an implicit value function estimating their potential to reach the correct target.In the example, “lempus” initially has value 0.33, while “rorpus” later receives 0.87 after expanding another branch.
- Interpreting the Latent Reasoning as Tree Search: The model maintains diverse reasoning paths during the first latent thought, then narrows the candidate set during the second thought.This pattern indicates a transition from parallel exploration to more focused reasoning.
- Why is Latent Space Better for Planning?: Nodes with lower heights receive more accurate and definitive value estimates, whereas greater heights produce more ambiguous evaluations.Node height is defined as distance to a leaf node, linking proximity to terminal states with confidence in path evaluation.
5 Empirical Results with Coconut
Experiments evaluate Coconut across math and logical reasoning tasks, showing that continuous thoughts can improve accuracy, support test-time scaling, and use fewer generated tokens than language reasoning in logical tasks. The results also identify multi-stage training as important and indicate that latent reasoning still needs guidance.
- 34.1% accuracy on GSM8k with 6 continuous thoughts versus 16.5% for No-CoT.
- Chaining continuous thoughts enhances reasoning, with Coconut outperforming comparable variants and the iCoT baseline on GSM8k.The paper presents chaining more continuous thoughts as an inference-time scaling strategy.
- Performance steadily improves as c increases from 0 to 1 to 2, where c controls latent thoughts per language reasoning step.
- Coconut offers a better accuracy-efficiency trade-off than CoT on GSM8k and maintains higher accuracy when fewer tokens are generated.On ProntoQA and ProsQA, it generates fewer tokens while achieving higher accuracy; decoding a first continuous thought often yields possible intermediate variables.
- Models trained without the multi-stage curriculum do not perform as well, indicating that latent reasoning requires training guidance.The authors state that better and more general latent-reasoning training strategies remain necessary, especially without language-chain supervision.
6 Conclusion
The conclusion presents Coconut as a continuous-latent-space reasoning paradigm whose continuous thoughts can represent alternative next steps and support breadth-first search. It reports improved reasoning performance across tasks while identifying broader pretraining and generalization as directions for further research.
- Coconut enables continuous thoughts to represent multiple alternative next steps and perform BFS over possible reasoning paths.
- Experiments demonstrate that Coconut enhances LLM performance across a variety of reasoning tasks.
- Further research is needed to refine and scale latent reasoning to pretraining for broader generalization across reasoning challenges.
A Datasets
The appendix provides example questions and reasoning solutions for the datasets used in the experiments, including arithmetic chains and logical entailment questions.
- The examples include a grade-school arithmetic problem represented as four intermediate calculation steps ending in answer 300.
- A logical example asks whether Stella is not floral and derives the opposite conclusion through a chain from zumpus to gorpus, rompus, yumpus, lorpus, and floral.
- The appendix also includes additional synthetic logical statements involving entities such as Alex, Tom, and Davis.
A.2 Construction of ProsQA
ProsQA is constructed as a binary-choice logical reasoning dataset whose DAGs guarantee a path from the entity to the correct concept but not the alternative. The construction separates node families, controls graph connections, and assigns names to formulate moderately complex reasoning paths.
- ProsQA frames each instance as asking whether an entity is concept A or concept B, with a DAG path to A but not B.
- The DAG is incrementally built by adding nodes and randomly connecting them with edges.
- New nodes are constrained from becoming descendants of both root nodes, preserving distinct node families and preventing model shortcuts.
- After construction, parentless nodes receive entity names, while other nodes receive concept names; node 0 is the entity, and labeled leaves 1 and 2 are concepts A and B.
- Sampling weights prioritize deeper candidate nodes, encouraging longer solution reasoning chains.
A.3 Statistics
The section reports dataset sizes in Table 3.
- Table 3 reports the sizes of all datasets.
B Clock-Time Reasoning Efficiency Metric
Reasoning efficiency is evaluated using average inference time per test case, measured with batch size 1 on an Nvidia A100 GPU. The reported clock time generally tracks the number of newly generated tokens.
- Inference time is reported as average seconds per test case with batch size 1 on an Nvidia A100 GPU.
- Clock time is generally proportional to the number of newly generated tokens across the evaluated methods.
C.1 Using More Continuous Thoughts
Using three continuous thoughts at once causes a slight GSM8k performance drop, increased variance, and training instability. The authors suggest finer-grained schedules or combining language and latent reasoning as future directions.
- Adding three continuous thoughts at once causes a slight GSM8k performance drop, increased variance, and unstable training.
C.2 Coconut with Larger Models
On larger Llama models, Coconut consistently improves over the no-CoT baseline, but the gains are less pronounced than those previously observed with GPT-2. The authors frame latent-space reasoning as promising while noting that universally surpassing language-based CoT remains an open research challenge.
- Coconut yields consistent performance gains over the no-CoT baseline on both Llama 3.2-3B and Llama 3-8B.
- The improvements on larger Llama models are less pronounced than those previously demonstrated with GPT-2.
- The authors identify latent-space reasoning as a promising direction but state that universally surpassing language-based CoT requires substantial further research.