Training Large Language Model to Reason in a Continuous Latent Space

We thank the reviewer for their thoughtful comments and appreciate the opportunity to address the concerns raised. Below, we respond to each point in detail and provide clarifications regarding our work.

Scope of Work and Performance in Comparison with CoT

Please kindly refer to our general response regarding the scope of this work and its comparison with CoT. Our primary objective is not to position COCONUT as a direct replacement for CoT but to explore the core attributes of latent space reasoning and analyze its potential as a novel paradigm. It’s also worth noting that COCONUT achieves a relative improvement of 13.7% over the latest implicit CoT reasoning method [1] on GSM8k, demonstrating progress in this emerging area.

Conceptual Connection to CoT

Thank you for your insightful comments. We agree that the differences between our method and CoT can be understood along two key dimensions:

• Learning Objective: Imitation learning on the reasoning process vs. optimization towards future outcomes (e.g., the likelihood of future steps and the final answer)
• Reasoning Representation: Discrete vs. continuous representations for intermediate reasoning steps.

Breaking Down the Combinations:

• Imitation learning + discrete representation: This corresponds to CoT, where reasoning steps are explicitly represented as language tokens and supervised through imitation learning.
• No imitation learning + continuous representation: This is the approach introduced in our work, where reasoning is performed in a continuous latent space, and the objective is to optimize the likelihood of future steps and the final answer.
• Imitation learning + continuous representation: This approach is not feasible due to the absence of ground truth labels for continuous reasoning representations. While previous work [1] attempted to train smaller models to imitate the hidden states of larger models for reasoning via knowledge distillation, their performance remains limited. In our study, we have compared COCONUT against its follow-up work [2] as a stronger baseline.
• No imitation learning + discrete representation: This combination is particularly challenging due to the non-differentiable nature of discrete tokens. Traditional methods aim to solve this with techniques such as Gumbel-softmax [3] to enable backpropagation, or reinforcement learning (RL) [4, 5] to optimize discrete objectives. However, these approaches have only become practical with the advent of powerful pre-trained models. Moreover, most existing RL methods require supervised fine-tuning (which is CoT training in our setting) before applying RL [6, 7]. At the initial stages of this project, we experimented with RL-based methods for reasoning using discrete tokens. However, we observed poor performance on GSM8k with GPT-2, often resulting in near-zero accuracy. This motivated our shift toward exploring continuous representations for reasoning, as they are inherently differentiable and offer a promising alternative for effective training.

Different Model in Section 5

We appreciate your concern and would like to clarify our intent. Our primary goal is to explore and analyze latent space reasoning versus language space reasoning, rather than advocating for any specific training schedule. As such, while the model used in Section 5 was not trained under exactly the same settings as the one in Section 4, we believe the results and analysis remain consistent with our objectives and claims.

We would also like to highlight that the accuracy of the model used in Section 5 is comparable to that of the model evaluated in Section 4 (95.8% and 97.0% on ProsQA, as shown in Figure 5 and Table 1), while both are significantly higher than CoT (77.5%, as reported in Table 1). This consistency indicates that the reasoning patterns observed in Section 5 are likely robust and universal across these two different but closely related training methods. We acknowledge your feedback and will revise the paper to explicitly clarify this point at the start of Section 5, ensuring there is no ambiguity regarding the relationship between the models and analyses across different sections. Thank you for bringing this to our attention!

We thank the reviewer for their detailed feedback and suggestions. Below, we address the raised concerns and clarify the contributions of our work.

The scope of paper and general-purpose reasoning method

As noted in our general response, the primary focus of this paper is to explore the core attributes of latent space reasoning and analyze its potential as a novel paradigm. While our method does require specific training datasets, this is a necessary step to investigate the fundamental feasibility of reasoning in a latent space. Importantly, the CoT method reported in our paper is also fine-tuned on specific training datasets. Though our approach may "diverge from the current research trend" of general-purpose LLM solutions, we view this as an opportunity to address a more fundamental problem, rather than building upon existing language space reasoning paradigms. This aligns with the long-term goal of enhancing machine reasoning capabilities.

Efficiency Metric

Our original motivation is to minimize external factors that could influence the comparison. For instance, existing libraries may have been specifically optimized for text generation with latest ML system techniques, whereas our implementation of continuous thought generation may not yet be fully optimized. Fundamentally, the computational cost of generating a continuous thought is inherently lower than that of generating a token. As such, we use the number of tokens as a conservative estimate of its efficiency.

Below, we supplement the clock-time comparison. The reported numbers are the average time of inference on one test case (in second), with batch size = 1, on an A100 GPU. We use the standard generate method in the transformers library for no CoT and CoT.

The clock time is generally proportional to the number of new generated tokens reported in Table 1. We will add the new results to our paper.

Impact of Prompt Variations

The GSM8k dataset includes diverse math problems expressed in free text, and therefore we believe the improved accuracy of our model can indicate its robustness to certain prompt variations. This experimental setting also closely follows previous works [1, 2]. Could you clarify the specific types of prompt changes you are interested in?

More thoughts on GSM8k

Please kindly refer to our supplemented experiments and discussions in the general response.

Divergent Benchmark Results

• Real-World vs. Synthetic Domains: GSM8k represents a real-world, open-domain question-answering task. Unlike the synthetic datasets used in our study, it demands more complex contextual understanding and modeling, which can place greater demands on computational capabilities. This hypothesis is supported by the observation that Coconut outperforms all other latent reasoning methods, and its accuracy steadily improves as the number of thoughts per step ($c$) increases from 0 to 2. Additionally, GSM8k requires diverse commonsense and world knowledge. This may give CoT an advantage, as it aligns closely with the pretraining objectives of the underlying language model, enabling it to better leverage its knowledge compared to Coconut.
• Planning Requirements: Complex reasoning tasks often require the model to "look ahead" to determine whether a particular step is optimal (also known as planning). Among the datasets in our experiments, GSM8k involves grade-school-level math word problems that allow for intuitive judgment of the next reasoning step. Similarly, ProntoQA includes distracting branches of limited size, making it relatively straightforward to identify the correct next step. In contrast, ProsQA, based on a randomly generated Directed Acyclic Graph (DAG) structure, presents a significant challenge to the model's planning abilities. Our experimental results suggest that tasks requiring extensive planning benefit more from latent space reasoning (including Coconut, some of its variants, and iCoT) than from reasoning using language tokens (CoT). Further analysis of this phenomenon is provided in Section 5.

These discussions can be found in lines 328–332 and 334–343 of the paper.

We hope this response clarifies the scope and contributions of our work, and we appreciate the valuable insights provided by the reviewers. We plan to add all the new results and discussions to the paper.

Reference

[1] Deng et al., 2023, “Implicit Chain of Thought Reasoning via Knowledge Distillation”

[2] Deng et al., 2024, “From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step”

Dear Reviewer,

We appreciate the time you took to review our paper and provide valuable feedback. We have carefully addressed your concerns in our rebuttal and would be grateful if you could take a moment to read it.

In particular, we would like to draw your attention to Section 5 (updated for better clarity), which reveals an interesting tree search pattern in latent space. It's worth noting that other reviewers have highlighted this section as one of the most valuable contributions of our paper. We would greatly appreciate hearing your thoughts on this section and any feedback you may have.

Thank you for your time and consideration.

We thank the reviewer for their detailed feedback and suggestions. Below, we address the raised concerns and clarify the contributions of our work.

The scope of this paper and comparison to CoT

While the performance of Coconut is worse than CoT on GSM8k, we’d like to emphasize that it outperforms the latest implicit CoT reasoning method [1] by a relative improvement of 13.7%. We’d like to emphasize that the scope of this paper is to focus on analyzing the core attributes of latent space reasoning (please refer to our general response), rather than proposing an immediate replacement of CoT.

Generation efficiency

To address concerns about the efficiency-performance trade-off, we report the number of generated tokens as a conservative metric for generation efficiency (Table 1), given that the computation of a continuous thought is strictly less than that of a token. Additionally, we have supplemented this with clock-time comparisons. Please kindly refer to our response to G21p.

More continuous thoughts in Figure 3; Results on larger models

Please kindly refer to our supplemented experiments and discussions in the general response.

Latent tree search analysis

Latent reasoning offers unique advantages by enabling the model to encode multiple potential next steps simultaneously in its continuous thought representations. This contrasts with verbal CoT, which requires deterministic decisions at every step. We illustrate this distinction with a concrete example in Figure 6:

• In verbal CoT, the model prematurely selects “lempus” as the next step, which is incorrect. In contrast, Coconut maintains a latent distribution representing multiple options, assigning similar scores to “lempus” (0.33) and “grimpus” (0.32).
• By the second latent reasoning step, Coconut progressively eliminates incorrect options, assigning a high score (0.87) to “rorpus,” which is on the correct reasoning path.

Therefore, we respectfully disagree with the statement that “a similar analysis can be done to verbal CoT”. While one could also construct a search tree for verbal CoT, the decoding process is essentially a preliminary greedy search, which fundamentally differs from latent reasoning's ability to perform BFS.

We have reorganized Section 5 for improved clarity and logical flow to address this feedback. The new version has been uploaded.

Trainable linear layer

We appreciate the reviewer’s suggestion to include a trainable linear projection layer to address potential input-output mismatches. We conducted experiments with this design during the initial stages of our work, but the results showed minimal differences. Our hypothesis is that since we train all parameters of the LLM, other parameters adapt effectively to produce and process continuous thoughts, mitigating the mismatch.

We hope this response clarifies the scope and contributions of our work, and we appreciate the valuable insights provided by the reviewers. We plan to add all the new results and discussions to the paper.

[1] Deng et al., 2023, “From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step”

Dear Reviewer,

We appreciate the time you took to review our paper and provide valuable feedback. We have carefully addressed your concerns in our rebuttal and would be grateful if you could take a moment to read it.

In particular, we would like to draw your attention to Section 5 (updated for better clarity), which reveals an interesting tree search pattern in latent space. It's worth noting that other reviewers have highlighted this section as one of the most valuable contributions of our paper. We would greatly appreciate hearing your thoughts on this section and any feedback you may have.

Thank you for your time and consideration.

To address the reviewer’s concerns:

The tree-search analysis only uses one dataset (ProsQA)

There is a trade-off between the depth and the generality of an analysis. We used ProsQA for this analysis because its graph structure allows us to clearly define a search tree and apply metrics like “height” to quantify the exploration potential of nodes. These are not applicable to open-domain datasets like GSM8k. This dataset enables us to conduct a focused, in-depth analysis of core planning abilities while minimizing confounding factors such as language understanding.

It’s worth noting that many “analysis papers” on LLM reasoning also focus on one synthetic datasets [5, 6, 7], like simple two-hop questions, but they still provide valuable insights into LLM reasoning by deeply exploring specific aspects of model behavior.

… and mainly focuses on the first two reasoning steps.

We focus on the first two reasoning steps because latent reasoning tends to be most effective during the early stages. This is supported by two key observations:

• The most significant improvement occurs when transitioning from Coconut (k=1) to Coconut (k=3) (Figure 5).
• As discussed in Section 5.4, nodes in the early steps are more challenging to evaluate, making parallel search particularly beneficial at this stage.

Focusing on these initial steps does not contradict our claims. On the contrary, it’s intended to enhance the clarity of our tree-search analysis and highlights the core strengths of latent reasoning.

The whole tree search is only analyzed with one specific example.

This is factually incorrect. The majority of our analysis is conducted on the full test set, as shown in Figures 5, 7, 9. Figures 6, 8 present an example, which are for illustrative purposes and to better explain some key concepts, such as value function and height, as a preparation for the quantitative analysis on the entire test set.

We also want to highlight that the examples in Figure 6 and 8 are not exceptional cases. They are just typical scenarios where increased latent reasoning enables problem-solving, while fewer latent steps do not. These examples do not compromise the generality of our findings, as they align with the broader statistical trends observed across the entire test set, e.g., the improvements from Coconut (k=1) to Coconut (k=3) shown in Figure 5.

In summary, the better planning performance that might be achieved by the proposed latent CoT is only hypothetical, with one synthetic dataset.

We respectfully disagree with the characterization of the improved planning performance as "hypothetical." In addition to better overall accuracy on ProsQA, we have included extensive analysis (the whole section 5) to substantiate our claims, as summarized above.

We would also like to emphasize that the ProsQA dataset is particularly well-suited for evaluating planning ability, as it is specifically designed to enlarge distractive branches, requiring the model to plan ahead before making decisions. In contrast, ProntoQA features similar question formats but with fewer branches. As shown in Table 1, the gap between latent reasoning and CoT is significantly larger on ProsQA compared to ProntoQA, further highlighting its planning advantages.

Reference

[1] Goyal et al. "Think before you speak: Training language models with pause tokens.", ICLR 2024.

[2] Pfau et al., “Let's Think Dot by Dot: Hidden Computation in Transformer Language Models”, COLM 2024.

[3] Deng et al., “Implicit Chain of Thought Reasoning via Knowledge Distillation”, arXiv preprint arXiv:2311.01460 (2023).

[4] Deng et al., “From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step” arXiv preprint arXiv:2405.14838 (2024).

[5] Yang, et al. "Do Large Language Models Latently Perform Multi-Hop Reasoning?." ACL 2024

[6] Biran, et al. "Hopping Too Late: Exploring the Limitations of Large Language Models on Multi-Hop Queries." EMNLP 2024

[7] Wang et al., “Grokked Transformers are Implicit Reasoners: A Mechanistic Journey to the Edge of Generalization” NeurIPS 2024

Thanks for acknowledging the tree search analysis is interesting and raising the score. We respectfully disagree with the dichotomy between method and analysis papers. By proposing a new training method, we believe this paper has provided both strong empirical results and an insightful analysis on latent space reasoning.

Empirical Results

We contend that the value of a method should not be assessed solely on whether it outperforms existing methods across all datasets. We believe our empirical results presented in this work are significant for several reasons:

• Reasoning without generating a language reasoning chain is an extremely challenging task, especially on GSM8k. As a reference:

  • The pause-token paper [1] does not compare to CoT throughout the paper. The proposed method achieves 8.5% accuracy on GSM8k, slightly improving over the no-CoT baseline, which yields 7.5% accuracy.
  • The filler token paper [2] focuses on synthetic datasets, and the performance is still worse than CoT.
  • The iCoT papers [3, 4] propose methods that match CoT performance on certain synthetic datasets but remain behind CoT on GSM8k. Furthermore, [4] highlights that even GPT-4, when prompted to reason without generating explicit language reasoning chains, achieves only a 44% accuracy on GSM8k, underscoring the difficulty of this task.

• To the best of our knowledge, our paper is the first to demonstrate that latent reasoning can outperform explicit CoT on any individual dataset. On ProsQA—a dataset requiring strong planning abilities—COCONUT exceeds CoT by 25.2%. This provides the first evidence of latent reasoning's advantages over CoT beyond efficiency. Moreover, our analysis explores the underlying mechanisms driving this enhanced planning capability, offering a clear explanation of why latent reasoning excels in these scenarios.

• On the challenging GSM8k dataset, our method clearly outperforms no-CoT baselines on GSM8K, and achieves a 13% relative improvement over the latest baseline iCoT [4]. Even though our performance is below that of models trained with explicit CoT, achieving 34% accuracy on GSM8k using a GPT-2 small model without generating explicit language reasoning chains is a non-trivial accomplishment. We also discussed the differences between datasets, explaining why latent reasoning faces greater challenges in outperforming language reasoning on GSM8k (Line 323).

• In addition to overall performance, our empirical results reveal several interesting findings:

  • Enhanced Performance with Increased Latent Computing: As demonstrated in Figure 3, employing more latent computation steps improves reasoning performance.
  • Interpretability of Latent Representations: The latent representations are easily interpretable, and they are shown to encode multiple potential intermediate variables simultaneously (Figure 4).

Analysis

We believe our analysis is thorough and well-supported. To provide further clarification:

• Interpolating between latent and language reasoning (Section 5.2): By adjusting the number of continuous thoughts, we explore a spectrum from fully latent reasoning to fully language-based reasoning. The results show that increasing latent reasoning improves both final answer accuracy and reasoning path accuracy (Figure 5). This indicates that sometimes while the model initially struggles with decisions in early steps, it’s often able to make a correct decision later when earlier steps are conducted in latent space. Figure 6 provides an example to illustrate this behavior.
• Tree Search Interpretation (Section 5.3): To better understand this process, we interpret latent reasoning as a tree search within the problem graph of ProsQA, and also define the model’s implicit value function to evaluate each node.
• Advantages of Latent Reasoning (Section 5.4): Based on the search tree perspective, we further explain why latent reasoning can be advantageous. When a node is closer to the leaf nodes (indicating smaller room for exploration), the LLM can accurately estimate its value (Figure 7). As the latent search tree expands (through using more latent reasoning steps), the search frontier is pushed closer to the leaf nodes, thus making the decision easier for LLMs.

Thank you for the thoughtful review and valuable feedback. Below, we address the key concerns raised and clarify our contributions.

Scope of Work and comparison to CoT

Please kindly refer to our general response. Our focus is on exploring latent space reasoning, especially its interesting attributes compared to language space reasoning, rather than framing this work as "Coconut vs. CoT." In case it wasn't immediately apparent, we would like to clarify that in this work, the CoT method refers specifically to training LLMs to generate natural language reasoning chains, rather than CoT prompting as proposed by Wei et al. [1].

The Number of Training Epochs and Usability

While the statement that our method “requires a large number of epochs” is true, it needs to be understood in the following contexts:

• GPT-2 is a relatively small language model and needs more training to converge. As a reference, in our experiments on GSM8k, it requires about 10 epochs for CoT training to converge, while the baseline iCoT method [2] takes ~50 epochs. Coconut converges in 25–30 epochs in total.

• Coconut’s learning objective does not align with existing pretraining paradigms, putting it at a disadvantage for continual training. However, as discussed in the general response, future work will explore pretraining paradigms optimized for latent space reasoning to mitigate this constraint and improve usability.

Evaluation Domains and Generalization

We acknowledge the importance of evaluating reasoning methods across diverse domains and out-of-distribution (OOD) tasks. However, we’d also like to point out that recent analysis [3] has shown that CoT can hardly help with any tasks other than math or symbolic reasoning. E.g. It cannot improve the performance on ARC, and on MMLU it only helps with math problems. [4] shows that CoT doesn’t help with Game-of-24 and Mini-crossword either.

Given these observations, many recent works on LLM reasoning [5, 6] also focus on math reasoning exclusively, as it’s a representative form of reasoning requiring diverse skills and presents significant challenges. While extending evaluations to more tasks (including those requiring longer-context CoT chains) or under the OOD settings are important future directions, it is beyond the immediate scope of this paper, which focuses on analyzing the core attributes of latent space reasoning.

We hope this response clarifies the scope and contributions of our work, and we appreciate the valuable insights provided by the reviewers. We will add all the new results and discussions to the paper.

Reference

[1] Chain-of-Thought Prompting Elicits Reasoning in Large Language Models

[2] From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step

[3] To CoT or Not to CoT? Chain-of-thought helps mainly on math and symbolic reasoning

[4] Tree of Thoughts: Deliberate Problem Solving with Large Language Models

[5] Teaching Large Language Models to Reason with Reinforcement Learning

[6] Physics of Language Models: Part 2.1, Grade-School Math and the Hidden Reasoning Process

Scale up to large models

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:

We still observe performance gains across both Llama 3.2-3B and Llama 3-8B models over the no-CoT baseline, though the gains are not as significant as those with GPT-2. A potential reason is that smaller models may particularly benefit from continuous thoughts to increase their effective network depth. Future work could explore improved training schedules, such as those proposed in iCoT, to enhance performance of larger models.

We also want to highlight that our goal is not to replace CoT, but to explore the potential of latent reasoning. For instance, our results indicate that this approach might enable smaller language models to serve as a reasoning backbone, balancing performance and computational constraints by dynamically leveraging continuous thoughts. Moreover, the method's inherent advantages in planning and search (shown in Section 5) may inspire novel approaches to solving complex problems.

---

We thank the reviewers once again for their thoughtful feedback and hope that the additional results and clarifications address their concerns.

[1] Goyal et al., 2023, “Training Language Models With Pause Tokens”

[2] Gu et al., 2023, “Linear-Time Sequence Modeling with Selective State Spaces”

[3] Deng et al., 2023, “From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step”