Training Large Language Models to Reason in a Continuous Latent Space

Thank you for recognizing our clear presentation, targeted analyses, and the gains delivered by continuous latent reasoning on ProsQA. Below we address each of your questions.

Convoluted training procedure

Our primary goal in this paper is to expose the benefits of latent reasoning—parallel search and higher reasoning efficiency. We would also like to mention that the current three‑stage curriculum is already simpler than some of the earlier work. For example, iCoT [1] removes CoT tokens one‑by‑one, leading to far longer training, and requires additional tricks like “removal smoothing,”. We expect future pre‑training in latent space to mitigate the need for dataset‑specific multi-stage training altogether.

Quantifying BFS

Fig. 3 reports aggregate accuracy on the full ProsQA test set as we increase the number of continuous thoughts $(k → k + 1)$. The upward trend implies that, for many instances, adding one latent step flips an incorrect outcome to a correct one. For adjacent pair $(k, k + 1)$, especially when $k=1, 2$, there is a subset of testset items where correctness flips. By definition, every flipped case must follow the same logic:

• With $k$ latent steps the model switches to language too early, commits to the “currently most promising” node, and later fails.
• With $k + 1$ latent steps it keeps several branches alive for one more step and then follows the correct path.

Figs. 4–5 present a randomly chosen instance from that subset to visualise this mechanism. The behaviour shown there recurs across all flipped cases, so the example is illustrative—not exceptional.

Furthermore, Figure 6 quantifies the link between a node’s predicted value and its height (distance to a leaf), clarifying why extra continuous thoughts improve BFS performance. Figure 8 plots the exploration distribution of continuous thoughts and shows uncertainty shrinking with each search step. Both analyses use the entire ProsQA test set.

BFS-like behavior on other datasets

We believe that many multi‑step reasoning tasks can be framed as path‑finding on a graph, and that Coconut naturally supports parallel exploration of such paths. Recent theoretical work [3, 4] also shows that the continuous thought architecture inherently enables this search behaviour.

In our experiment, we chose ProsQA for analysis because its explicit DAG allows us to measure search properties such as height and frontier distribution without the confounds of natural‑language understanding. For open-domain tasks like GSM8K, we still see evidence that continuous thoughts encode multiple candidate paths (Fig. 9), but a clean search‑tree metric would be harder to define.

Greedy decoding

We apply greedy decoding for all experiments in Figure 3. We use this setting because previous work [2] shows it typically performs well when using low temperature or greedy decoding for reasoning tasks.

Thanks again for your valuable feedback!

Reference

[1] Deng et al., 2024, “From explicit cot to implicit cot: Learning to internalize cot step by step”

[2] Wei et al., 2022, “Chain-of-Thought Prompting Elicits Reasoning in Large Language Models”

[3] Gozeten et al., 2025, “Continuous Chain of Thought Enables Parallel Exploration and Reasoning”

[4] Zhu et al., 2025, “Reasoning by Superposition: A Theoretical Perspective on Chain of Continuous Thought”

Thank you for the thoughtful review and for highlighting our method’s potential impact, thorough analysis, and better empirical results. Below we address each of your questions.

More general reasoning tasks Regarding GSM8K performance, please see our general response for additional discussion. We agree that fully unlocking the potential of latent reasoning will require further research, and we are actively pursuing larger-scale latent space pretraining to close the gap.

Only GPT‑2 is used

We experimented with COCONUT using Llama 3.2-3B and Llama 3-8B with $c=1$, training for 3 epochs for Llama 3.2-3B and 2 epochs in Stage 0, followed by 1 epoch per subsequent stage. The results are as follows:

Both backbones show clear gains over the no‑CoT baseline, though the boost is smaller than with GPT‑2. One hypothesis is that larger LLMs encode more knowledge for math reasoning problems learned from language-space pre-training, and therefore latent space pre-training might be crucial for coconut to achieve comparable or better performance to language space reasoning.

Scalability

Though the current training method is sequential, we are looking into two potential paths to address this:

• System‑level optimization: For example, FlashRNN [1] shows that careful kernel design (Triton/CUDA, register‑level tuning) can slash the overhead of recurrent computation.
• Post‑training application. Recent work on latent space pretraining [2] applies a fixed sentence embedding, which enables parallel training at sequence level. Coconut could play the role of post‑training stage after such large-scale pre-training, in order to optimize the latent space for reasoning. This mirrors recent progress in online RL for LLMs—despite sequential rollouts, methods like OpenAI’s O1 [3] and DeepSeek’s R1 [4] deliver meaningful gains on long‑horizon reasoning.

Thanks again for your valuable feedback!

References

[1] Pöppel et al., 2025, FlashRNN: I/O‑Aware Optimisation of Traditional RNNs on Modern Hardware

[2] LCM Team, 2024, Large Concept Models: Language Modelling in a Sentence Representation Space

[3] OpenAI, 2024, O1 – https://openai.com/o1/

[4] Guo et al., 2025, DeepSeek‑R1: Incentivising Reasoning Capability in LLMs via Reinforcement Learning

Thank you for the thoughtful review and for highlighting our method’s novelty, thorough experimentation, and clear writing. Below we address each of your questions.

General reasoning tasks

Our main goal is to reveal and analyze the promising properties of latent reasoning, specifically its parallel search behaviour and superior compute efficiency. As discussed in the general response, we expect future works to scale up pretraining in latent space and improve the generality of the latent reasoning model.

Decompose into reasoning steps

This is a great question and an interesting direction to explore. A natural first step is to segment the text into sentences. For example, Large Concept Models [1] apply Segment-any-Text (SaT) [2] for this purpose. A more flexible approach would be to dynamically decide thought boundaries during training, similar to the idea of Byte-Latent Transformer [3], which uses next-byte entropy to mark patch starts.

Dataset specific training / a more general recipe

Currently, Coconut requires dataset-specific training. However, just as language space reasoning relied on task-specific training [4] before very strong LLMs came out, we view this as an intermediate stage. Our paper focuses on uncovering key properties of latent-space reasoning, and future latent-space pre-training might help eliminate the need for task-specific fine-tuning altogether.

KV cache reuse

Consider a sequence of length $L$ containing $k$ continuous thoughts (indices $i$ … $i + k − 1$). Computing its loss requires $k + 1$ forward passes:

• Pass 1 on tokens $[0, i)$ produces the hidden state at $i − 1$ (first continuous thought).
• Pass 2 on $[0, i + 1)$ reuses the KV cache from pass 1 for $[0, i)$ and processes only the new position $i$.
• Subsequent passes extend similarly, each reusing all prior KV entries.

By reusing the KV cache from the previous forward pass, we could avoid repetitive computation.

Changing optimizer state

Following Deng et al. [6], we reset AdamW’s moments to avoid stale momentum when gradients change direction at each stage switch. We did not explore more options in the optimizer setting.

In the absence of CoT data

At the current stage, using CoT data to guide latent reasoning is still necessary. As shown in the ablation setting “Coconut w/o curriculum”, if we directly train the model with the final answer, without the guidance from CoT, the performance is significantly worse. We expect latent space pretraining to mitigate this dependency in future work.

Thanks again for your valuable feedback!

Reference

[1] LCM team, 2024, “Large Concept Models: Language Modeling in a Sentence Representation Space”

[2] Frohmann, 2024, “Segment Any Text: A Universal Approach for Robust, Efficient and Adaptable Sentence Segmentation”

[3] Pagnoni, 2024, “Byte latent transformer: Patches scale better than tokens”

[4] Ling et al., 2017, “Program Induction by Rationale Generation: Learning to Solve and Explain Algebraic Word Problems”

[5] Kobbe et al., 2021, “Training Verifiers to Solve Math Word Problems”

[6] Deng et al., 2024, “From explicit cot to implicit cot: Learning to internalize cot step by step”

Thank you for the thoughtful review and for describing Coconut as impactful and transformative!

Regarding GSM8K performance, please see our general response for additional discussion. We agree that fully unlocking the potential of latent reasoning will require further research, and we are actively pursuing larger-scale latent space pretraining to close the gap.