Training Large Language Models to Reason in a Continuous Latent Space — part 2/3

Cleaned-text code points [15912, 31499). Bibliography omitted here; source.md includes it.

In this section, we provide a proof of concept of the advantage of continuous latent space reasoning. On ProsQA, a new dataset that requires extensive planning ability, Coconut outperforms language space CoT reasoning. Interestingly, our analysis indicates that the continuous representation of reasoning can encode multiple alternative next reasoning steps. This allows the model to perform a breadth-first search (BFS) to solve the problem, instead of prematurely committing to a single deterministic path like language CoT.

We start by introducing the experimental setup (Section 4.1 ). By leveraging Coconut ’s ability to switch between language and latent space reasoning, we are able to control the model to interpolate between fully latent reasoning and fully language reasoning and test their performance (Section 4.2 ). This also enables us to interpret the latent reasoning process as tree search (Section 4.3 ). Based on this perspective, we explain why latent reasoning can help LLMs make better decisions (Section 4.4 ).

### 4.1 Experimental Setup

Dataset. We introduce ProsQA ( Pro of with S earch Q uestion- A nswering), a new logical reasoning dataset. A visualized example is shown in Figure 5 . Each instance in ProsQA consists of a directed acyclic graph (DAG) of logical relationships between concepts, presented as natural language statements. The task requires models to determine logical relationships by finding valid paths through this graph, demanding sophisticated planning and search strategies. Unlike previous logical reasoning datasets like ProntoQA ( 35 ) , ProsQA’s DAG structure introduces complex exploration paths, making it particularly challenging for models to identify the correct reasoning chain. More comprehensive details about the dataset construction and characteristics can be found in Appendix A .

Setup. We use a pre-trained GPT-2 model as the base model for all experiments. The learning rate is set to 1 × 10 − 4 1\times 10^{-4} while the effective batch size is 128. We train a Coconut model following the training procedure in Section 3 . Since the maximum reasoning steps in ProsQA is 6, we set the number of training stages to N = 6 N=6 in the training procedure. In each stage, we train the model for 5 epochs, and stay in the last stage until the 50 epochs. The checkpoint with the best accuracy in the last stage is used for evaluation. As reference, we report the performance of (1) CoT : the model is trained with CoT data, and during inference, the model will generate a complete reasoning chain to solve the problem. (2) no-CoT : the model is trained with only the question and answer pairs, without any reasoning steps. During inference, the model will output the final answer directly.

To understand the properties of latent and language reasoning space, we manipulate the model to switch between fully latent reasoning and fully language reasoning , by manually setting the position of the <eot> token during inference. When we enforce Coconut to use k k continuous thoughts, the model is expected to output the remaining reasoning chain in language, starting from the k + 1 k+1 step. In our experiments, we test variants of Coconut on ProsQA with k ∈ { 0 , 1 , 2 , 3 , 4 , 5 , 6 } k\in\{0,1,2,3,4,5,6\} . Note that all these variants only differ in inference time while sharing the same model weights.

Metrics. We apply two sets of evaluation metrics. One of them is based on the correctness of the final answer , regardless of the reasoning process. It is also the main metric used in the later sections (Section 5.3 ). To enable fine-grained analysis on ProsQA, we define another metric on the reasoning process . We classify a reasoning chain into (1) Correct Path : The output is one of the shortest paths to the correct answer. (2) Longer Path : A valid path that correctly answers the question but is longer than the shortest path. (3) Hallucination : The path includes nonexistent edges or is disconnected. (4) Wrong Target : A valid path in the graph, but the destination node is not the one being asked. These four categories naturally apply to the output from Coconut ( k = 0 k=0 ) and CoT , which generate the full path. For Coconut with k > 0 k>0 that outputs only partial paths in language (with the initial steps in continuous reasoning), we classify the reasoning as a Correct Path if a valid explanation can complete it . Also, we define Longer Path and Wrong Target for partial paths similarly. If no valid explanation completes the path, it’s classified as Hallucination. In no-CoT and Coconut with larger k k , the model may only output the final answer without any partial path, and it falls into (5) Correct Label or (6) Incorrect Label . These six categories cover all cases without overlap.

### 4.2 Overall Results

Figure 3 presents a comparative analysis of various reasoning methods evaluated on ProsQA. The model trained using CoT frequently hallucinates non-existent edges or outputs paths leading to incorrect targets, resulting in lower answer accuracy. In contrast, Coconut , which leverages continuous space reasoning, demonstrates improved accuracy as it utilizes an increasing number of continuous thoughts. Additionally, the rate of correct reasoning processes (indicated by “Correct Label” and “Correct Path”) significantly increases. At the same time, there is a notable reduction in instances of “Hallucination” and “Wrong Target,” issues that typically emerge when the model makes mistakes early in the reasoning process.

An intuitive demonstration of the limitations of reasoning in language space is provided by the case study depicted in Figure 5 . As shown, models operating in language space often fail to plan ahead or backtrack. Once they commit to an incorrect path, they either hallucinate unsupported edges or terminate with irrelevant conclusions. In contrast, latent reasoning avoids such premature commitments by enabling the model to iteratively refine its decisions across multiple reasoning steps. This flexibility allows the model to progressively eliminate incorrect options and converge on the correct answer, ultimately resulting in higher accuracy.

### 4.3 Interpreting the Latent Reasoning as Tree Search

To better understand Coconut , we probe the latent reasoning process by forcing the model to explicitly generate language reasoning steps following intermediate continuous thoughts (Figure 5 ). Using the example presented in Figure 5 , at the initial reasoning step, the model must select which immediate child node of “Alex” to consider next, specifically from the set {“lempus”, “sterpus”, “zhorpus”, “grimpus”}. The distribution over these candidate next steps is visualized in Figure 5 , left. In the subsequent reasoning step, these nodes expand further into an extended set of potential paths, including all grandchildren of “Alex” (Figure 5 , right).

We define the predicted probability of a concept following continuous thoughts as a value function (Figure 5 ), estimating each node’s potential for reaching the correct target. Interestingly, the reasoning strategy employed by Coconut is not greedy search: while “lempus” initially has the highest value (0.33) at the first reasoning step (Figure 5 , left), the model subsequently assigns the highest value (0.87) to “rorpus,” a child of “grimpus,” rather than following “lempus” (Figure 5 , right). This characteristic resembles a breadth-first search (BFS) approach, contrasting sharply with the greedy decoding typical of traditional CoT methods. The inherent capability of continuous representations to encode multiple candidate paths enables the model to avoid making immediate deterministic decisions. Importantly, this tree search pattern is not limited to the illustrated example, but constitutes a fundamental mechanism underlying the consistent improvement observed with larger values of k k in Coconut .

Figure 6 presents an analysis of the parallelism in the model’s latent reasoning across the first and second thoughts. For the first thoughts (left panel), the cumulative values of the top-1, top-2, and top-3 candidate nodes are computed and plotted against their respective percentiles across the test set. The noticeable gaps between the three lines indicate that the model maintains significant diversity in its reasoning paths at this stage, suggesting a broad exploration of alternative possibilities. In contrast, the second thoughts (right panel) show a narrowing of these gaps. This trend suggests that the model transitions from parallel exploration to more focused reasoning in the second latent reasoning step, likely as it gains more certainty about the most promising paths.

### 4.4 Why is Latent Space Better for Planning?

Building upon the tree search perspective, we further examine why latent reasoning benefits planning tasks—specifically, why maintaining multiple candidate paths and postponing deterministic decisions enhances reasoning performance. Our hypothesis is that nodes explored in the early reasoning stages are inherently more challenging to evaluate accurately because they are farther from the final target nodes. In contrast, nodes positioned closer to potential targets, having fewer subsequent exploration possibilities, can be assessed accurately with higher confidence.

To systematically test this, we define the height of a node as its shortest distance to any leaf node and analyze the relationship between node height and the model’s estimated value. Ideally, a correct node—one that can lead to the target node—should receive a high estimated value, whereas an incorrect node—one that cannot lead to the target node—should receive a low value. Empirical results across the test set (Figure 7 ) support our hypothesis: nodes with lower heights consistently receive more accurate and definitive probability evaluations. Conversely, nodes with greater heights exhibit more ambiguous evaluations, reflecting increased uncertainty.

These findings underscore the advantage of latent space reasoning. By delaying deterministic decisions and allowing exploration to proceed toward terminal states, latent reasoning significantly enhances the model’s ability to differentiate correct paths from incorrect ones, thereby improving performance on complex, planning-intensive tasks compared to traditional greedy methods.

## 5 Empirical Results with Coconut

After analyzing the promising parallel search pattern of Coconut , we validate the feasibility of LLM reasoning in a continuous latent space through more comprehensive experiments, highlighting its better reasoning efficiency over language space, as well as its potential to enhance the model’s expressivity with test-time scaling.

### 5.1 Experimental Setup

Math Reasoning. We use GSM8k ( 6 ) as the dataset for math reasoning. It consists of grade school-level math problems. To train the model, we use a synthetic dataset generated by 8 . We use two continuous thoughts for each reasoning step (i.e., c = 2 c=2 ). The model goes through 3 stages besides the initial stage. We then include an additional stage where still 3 × c 3\times c continuous thoughts are used as in the previous stage, but with all the remaining language reasoning chain removed. This handles the long-tail distribution of reasoning chains longer than 3 steps. We train the model for 6 epochs in the initial stage, and 3 epochs in each remaining stage.

Logical Reasoning. Logical reasoning involves the proper application of known conditions to prove or disprove a conclusion using logical rules. We use the ProntoQA ( 35 ) dataset, and our newly proposed ProsQA dataset, which is more challenging due to more distracting branches. We use one continuous thought for every reasoning step (i.e., c = 1 c=1 ). The model goes through 6 training stages in addition to the initial stage, because the maximum number of reasoning steps is 6 in these two datasets. The model then fully reasons with continuous thoughts to solve the problems in the last stage. We train the model for 5 epochs per stage.

For all datasets, after the standard schedule, the model stays in the final training stage, until reaching 50 epochs. We select the checkpoint based on the accuracy on the validation set. For inference, we manually set the number of continuous thoughts to be consistent with their final training stage. We use greedy decoding for all experiments.

### 5.2 Baselines and Variants of Coconut

We consider the following baselines: (1) CoT , and (2) No-CoT , which were introduced in Section 4 . (3) iCoT ( 7 ) : The model is trained with language reasoning chains and follows a carefully designed schedule that “internalizes” CoT. As the training goes on, tokens at the beginning of the reasoning chain are gradually removed until only the answer remains. During inference, the model directly predicts the answer. (4) Pause token ( 18 ) : The model is trained using only the question and answer without a reasoning chain. However, different from No-CoT , special <pause> tokens are inserted between the question and answer, which provides the model with additional computational capacity to derive the answer. The number of <pause> tokens is set the same as continuous thoughts in Coconut .

We also evaluate some variants of Coconut : (1) w/o curriculum , which directly trains the model in the last stage. The model uses continuous thoughts to solve the whole problem. (2) w/o thought : We keep the multi-stage training, but don’t add any continuous latent thoughts. While this is similar to iCoT in the high-level idea, the exact training schedule is set to be consistent with Coconut , instead of iCoT , for a strict comparison. (3) Pause as thought : We use special <pause> tokens to replace the continuous thoughts, and apply the same multi-stage training curriculum as Coconut .

### 5.3 Results and Discussion

We show the overall results in Table 1 . Using continuous thoughts effectively enhances LLM reasoning over the No-CoT baseline. For example, by using 6 continuous thoughts, Coconut achieves 34.1% accuracy on GSM8k, which significantly outperforms No-CoT (16.5%). We highlight several key findings below.

“Chaining” continuous thoughts enhances reasoning. Language CoT proves to increase the effective depth of LLMs and enhance their expressiveness ( 13 ) . Thus, generating more tokens serves as a way to inference-time scaling for reasoning ( 19 ; 38 ) . This desirable property holds naturally for Coconut too. On GSM8k, Coconut outperformed other architectures trained with similar strategies, including Coconut ( pause as thought ) and Coconut ( w/o thought ). Particularly, it surpasses the latest baseline iCoT ( 7 ) , which requires a more carefully designed training schedule.

Additionally, we experimented with adjusting the hyperparameter c c , which controls the number of latent thoughts corresponding to one language reasoning step (Figure 8 , II). As we increased c c from 0 to 1 to 2, the model’s performance steadily improved. 2 2 2 We discuss the case of larger c c in Appendix C.1 . This further validates the potential of continuous thoughts to scale up to harder problems. In two other synthetic tasks, we found that the variants of Coconut ( w/o thoughts or pause as thought ), and the iCoT baseline also achieve impressive accuracy. This indicates that the model’s computational capacity may not be the bottleneck in these tasks. In contrast, GSM8k involves more complex contextual understanding and modeling, placing higher demands on computational capability.