Training Large Language Models to Reason in a Continuous Latent Space | Summary

Contents

Summary
1 Introduction
2 Related Work
3 Coconut: Chain of Continuous Thought
4 Continuous Space Enables Latent Tree Search
5 Empirical Results with Coconut
6 Conclusion
Appendix
Brief Thoughts

This article explains the key points of Training Large Language Models to Reason in a Continuous Latent Space.

2024-12-09 (arXiv)
Hao, Shibo, Sukhbaatar, Sainbayar, Su, DiJia, Li, Xian, Hu, Zhiting, Weston, Jason, Tian, Yuandong.
FAIR at Meta, UC San Diego
Paper

Summary

Large language models (LLMs) usually express intermediate reasoning as text, which can spend computation on linguistic fluency and commit too early to an incorrect reasoning path. Coconut, short for Chain of Continuous Thought, instead feeds internal state vectors directly into subsequent reasoning steps and achieves 97.0±0.3% accuracy on ProsQA, a logical reasoning benchmark requiring search, versus 77.5±1.9% for chain-of-thought (CoT), which generates intermediate reasoning in language, under the paper’s GPT-2 model setup.
Coconut uses a training curriculum that progressively replaces textual reasoning steps with continuous thoughts: the model’s final hidden states, or internal state vectors, reused as input embeddings, the vectors supplied to the model at successive positions. On GSM8k, a grade-school mathematics benchmark, it reaches 34.1±1.5% accuracy versus 16.5±0.5% without intermediate reasoning, while remaining below textual CoT at 42.9±0.2%.
Probes that decode intermediate states support an interpretation resembling breadth-first search (BFS), which explores alternatives at the same depth before proceeding deeper. The evidence remains task dependent: curriculum-based baselines also perform strongly on the synthetic logic tasks, training without the curriculum fails to improve on direct answers, and larger-model experiments show smaller gains.

1 Introduction

Training Large Language Models to Reason in a Continuous Latent Space questions whether every intermediate reasoning state needs to become a discrete token, a unit of generated text. The authors argue that nearly uniform computation per generated token poorly matches the uneven difficulty of reasoning: many tokens maintain fluent prose, whereas a few require substantial planning.

Coconut bypasses the conversion from hidden states to discrete tokens during intermediate reasoning, while retaining language for the question and answer.
The cited neuroscience motivates separating language from reasoning, but the paper’s experimental evidence comes from machine reasoning benchmarks.
The supplied version is arXiv:2412.06769v4, last updated August 25, 2026. Code: https://github.com/facebookresearch/coconut.

Textual reasoning and continuous hidden-state feedback between the question and answer.

The paper connects Coconut to work showing that feeding generated reasoning back into a transformer, the underlying language-model architecture, increases its effective computational depth. It also reviews explicit search, reasoning through hidden states, additional computation tokens, and methods that gradually remove textual reasoning during training.

Pause tokens are additional input tokens that provide computation before the answer without expressing a reasoning step.
iCoT is a curriculum-based method that internalizes CoT by gradually removing initial reasoning tokens until the model directly predicts the answer.
Coconut combines a reasoning-removal curriculum with repeated use of continuous hidden states. The supplied revision also discusses subsequent theoretical studies of superposition, meaning multiple reasoning paths represented together.

3 Coconut: Chain of Continuous Thought

Coconut alternates between ordinary language generation and latent mode, where the previous position’s final hidden state becomes the next input embedding. The special tokens <bot> and <eot> mark the beginning and end of latent reasoning; the reused hidden states have already passed through the model’s final normalization layer, which controls their magnitude.

In ordinary language mode, the model applies its output head, the mapping from hidden states to token probabilities, to select the next token. In latent mode, it reuses the hidden vector without selecting a token.
The output head can still decode a latent state for analysis, although producing language from that state is not part of the latent inference step.
Because the recurrent states are continuous, gradients can propagate through the sequence of latent computations during training.

At training stage k, the first k textual reasoning steps are replaced with \(k\times c\) continuous thoughts, where c specifies how many latent steps replace one language step. The optimizer state is reset at stage transitions, and the prediction loss excludes questions and latent thoughts while supervising the remaining reasoning text and answer.

This objective trains latent states to support future predictions; it does not explicitly require reconstructing or compressing the removed reasoning text. If a chain contains fewer than k textual steps, all of them are removed.
Training with n latent thoughts requires n+1 forward passes. A key-value (KV) cache, which stores previously computed attention keys and values, avoids repeated computation, but the passes remain sequential.
Inference normally uses a manually fixed latent length. The authors also tried a binary classifier that decides whether to end latent reasoning and report comparable performance, without a detailed quantitative comparison.

Curriculum that progressively replaces language reasoning steps with continuous thoughts.

4 Continuous Space Enables Latent Tree Search

The latent tree-search investigation uses ProsQA, short for Proof with Search Question-Answering, as a proof of concept for planning across distracting logical branches. By changing when the same trained model exits latent mode, the authors compare fully textual reasoning, mixed reasoning, and fully latent reasoning without changing model weights.

4.1 Experimental Setup

ProsQA presents logical statements describing a directed acyclic graph (DAG), a directed graph with no cycles, and asks which of two concepts is reachable from an entity. The GPT-2 experiments use a learning rate of \(1\times10^{-4}\), effective batch size 128, six replacement stages after the initial stage, and five epochs per stage before continuing the final stage until 50 total epochs.

Inference varies the number of continuous thoughts over \(k\in\{0,1,2,3,4,5,6\}\); after those thoughts, the model generates any remaining reasoning in language.
Final-answer accuracy is evaluated separately from reasoning quality. Correct Path denotes a shortest valid path to the answer; Longer Path denotes a valid but longer route; Hallucination denotes nonexistent edges or a disconnected path; Wrong Target denotes a valid route to an irrelevant destination. Answer-only outputs receive Correct Label or Incorrect Label.
For partially verbalized paths, Correct Path means that a valid explanation can complete the observed suffix. This tests compatibility with a valid path rather than directly observing the hidden reasoning.

4.2 Overall Results

Figure 3 shows higher ProsQA answer accuracy once Coconut uses several continuous thoughts, together with fewer invalid edges and wrong-target outputs. The case study in Figure 4 illustrates the proposed explanation: textual CoT commits to a dead end and invents an edge, one latent thought still leads to an irrelevant target, and two latent thoughts produce a correct continuation.

The accuracy pattern is not strictly monotonic at every adjacent value of k: the clearest improvement occurs when moving from short latent prefixes to several latent steps.
At larger k, many outputs contain only the answer. Their Correct Label classification does not verify an explicit proof.

ProsQA final-answer accuracy and reasoning-output categories across latent lengths and baselines.

A ProsQA graph and example outputs from textual CoT and Coconut with one or two latent reasoning steps.

4.3 Interpreting the Latent Reasoning as Tree Search

The authors probe intermediate latent states by forcing a return to language and measuring probabilities of candidate next concepts. In the illustrated example, “lempus” initially has the highest probability at 0.33, but the next latent step favors “rorpus” at 0.87, although it descends from the alternative “grimpus” branch.

A concept’s score is the product of the conditional probabilities of its constituent tokens. The paper interprets this score as an implicit value function, a score estimating the concept’s usefulness for reaching the target.
Across the test set, the first latent step distributes appreciable probability across several candidates; the second step concentrates more strongly on a leading candidate.
These observations support a BFS-like interpretation of deferred commitment. They do not directly enumerate the complete internal search process.

Candidate-concept probabilities decoded after the first and second continuous thoughts.

Test-set percentile curves for cumulative candidate values in the first two latent steps.

4.4 Why is Latent Space Better for Planning?

The planning hypothesis is that early candidate nodes are harder to evaluate because they remain far from terminal outcomes. Figure 7 relates candidate scores to node height, a measure of distance to a leaf, and shows more decisive separation between correct and incorrect candidates near leaves, supporting the proposed benefit of delaying a discrete choice.

This is correlational evidence for the proposed explanation rather than an intervention isolating deferred commitment as the cause of improved accuracy.
The source defines height inconsistently: Figure 5 uses the longest distance to a leaf, whereas Section 4.4 uses the shortest distance to a leaf. The text does not resolve which definition underlies the analysis.

Predicted candidate values versus node height for correct and incorrect nodes.

5 Empirical Results with Coconut

The broader evaluation tests continuous reasoning on mathematics and two synthetic logical reasoning datasets. It examines final-answer accuracy, inference length, curriculum dependence, and test-time scaling, meaning the effect of allocating more computation through additional continuous steps.

5.1 Experimental Setup

For GSM8k, the model trains on synthetic mathematics data with c=2, using six initial-stage epochs and three epochs per subsequent stage. Three stages progressively replace reasoning, followed by an additional stage that removes any remaining textual reasoning while retaining six continuous thoughts.

ProntoQA, a synthetic logical deduction benchmark, and ProsQA use c=1, six replacement stages after the initial stage, and five epochs per stage.
All schedules continue in the final stage until 50 total epochs. Checkpoints are selected by validation accuracy, inference uses the final training stage’s latent length, and all experiments use greedy decoding, selecting the most probable next token.
The main evaluation uses GPT-2; separate larger-model results appear in Appendix C.2.

5.2 Baselines and Variants of Coconut

The baselines are textual CoT, No-CoT trained only on question-answer pairs, iCoT with progressive reasoning-token removal, and Pause Token trained with additional computation tokens but no reasoning-chain supervision. Three Coconut variants distinguish the contribution of the curriculum from that of continuous states.

w/o curriculum trains directly with fully latent reasoning, removing intermediate curriculum stages.
w/o thought retains Coconut’s curriculum but adds no continuous thoughts. Its schedule follows Coconut rather than iCoT.
pause as thought replaces continuous states with <pause> tokens while retaining the same curriculum. This differs from the Pause Token baseline, which lacks the reasoning-removal curriculum.

5.3 Results and Discussion

Table 1 reports Coconut at 34.1±1.5% accuracy and 8.2 inference tokens on GSM8k, compared with CoT at 42.9±0.2% and 25.0 tokens. Coconut reaches 99.8±0.2% with 9.0 tokens on ProntoQA versus CoT’s 98.8±0.8% with 92.5 tokens, and 97.0±0.3% with 14.2 tokens on ProsQA versus 77.5±1.9% with 49.4 tokens.

Coconut’s reported inference lengths include continuous-thought positions; they are not counts of generated text alone. It improves on No-CoT’s 16.5±0.5% GSM8k accuracy and exceeds the cited iCoT mathematics result of 30.0%, but textual CoT remains the most accurate mathematics method in the table.
On ProsQA, iCoT achieves 98.2±0.3% with 8.2 tokens, exceeding Coconut in both accuracy and brevity. The logical reasoning results therefore do not establish Coconut as the strongest evaluated method.
The reported ± quantities are preserved as printed; the surrounding experimental description does not specify their statistical definition.

Accuracy and inference-token comparisons for three reasoning datasets, baselines, and Coconut ablations.

The curriculum is essential under the tested conditions: removing it lowers Coconut to 14.4±0.8% on GSM8k, 52.4±0.4% on ProntoQA, and 76.1±0.2% on ProsQA. Increasing c from 0 to 1 to 2 improves mathematics performance, while Figure 8 shows that replacing textual steps with continuous thoughts retains accuracy better than simply removing those textual steps.

On GSM8k, w/o thought obtains 21.6±0.5% and pause as thought obtains 24.1±0.7%, both below full Coconut. This comparison supports a contribution from continuous feedback beyond the curriculum alone.
The same variants perform strongly on the synthetic logic tasks, so those tasks provide weaker evidence that continuous feedback is necessary.
Figure 9 decodes an intermediate mathematics state into plausible numerical quantities, including “180” and “9”. This is an illustrative probe rather than a general verification of latent-state meaning.

GSM8k accuracy versus inference length and continuous thoughts per replaced reasoning step.

A mathematics problem with numerical tokens decoded from an intermediate continuous thought.

6 Conclusion

The conclusion presents continuous hidden-state feedback as a way to reason beyond explicit language tokens and interprets the observed alternative-path behavior as latent search. Scaling this approach into pretraining and improving generalization are proposed research directions; the experiments do not establish those outcomes.

Appendix

A Datasets describes arithmetic and logical problem formats, the synthetic ProsQA graph-generation procedure, and dataset statistics and splits. Answer choices are randomly reordered to reduce positional bias.
B Clock-Time Reasoning Efficiency Metric reports average latency with batch size 1 on an Nvidia A100 GPU. Coconut is faster than CoT under the tested hardware and implementation; this is a scoped timing comparison.
C More Discussion reports that adding more continuous thoughts can reduce GSM8k performance and destabilize training. Larger-model experiments show modest gains and omit a textual CoT comparison; language-focused pretraining is offered as a possible explanation for difficult latent adaptation.

Brief Thoughts

The main contribution is continuous hidden-state feedback paired with probes showing that intermediate computation can retain alternatives beyond the initially leading linguistic branch. The mathematics experiments support a useful accuracy–length trade-off under the tested setup; strong iCoT results, curriculum dependence, sequential training cost, and limited larger-model comparisons constrain broader claims of superiority.