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

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

# Training Large Language Models to Reason in a Continuous Latent Space

###### Abstract

Large language models (LLMs) are restricted to reason in the “language space”, where they typically express the reasoning process with a chain-of-thought (CoT) to solve a complex reasoning problem. However, we argue that language space may not always be optimal for reasoning. For example, most word tokens primarily ensure textual coherence and are not essential for reasoning, while some critical tokens require complex planning and pose huge challenges to LLMs. To explore the potential of LLM reasoning in an unrestricted latent space instead of using natural language, we introduce a new paradigm Coconut ( C hain o f Con tin u ous T hought). We utilize the last hidden state of the LLM as a representation of the reasoning state (termed “continuous thought”). Rather than decoding this into a word token, we feed it back to the LLM as the subsequent input embedding directly in the continuous space. This latent reasoning paradigm leads to the emergence of an advanced reasoning pattern: the continuous thought can encode multiple alternative next reasoning steps, allowing the model to perform a breadth-first search (BFS) to solve the problem, rather than prematurely committing to a single deterministic path like CoT. Coconut outperforms CoT on certain logical reasoning tasks that require substantial search during planning, and shows a better trade-off between accuracy and efficiency.

## 1 Introduction

Large language models (LLMs) have demonstrated remarkable reasoning abilities, emerging from extensive pretraining on human languages ( 9 ; 1 ) . While next token prediction is an effective training objective, it imposes a fundamental constraint on the LLM as a reasoning machine: the explicit reasoning process of LLMs must be generated in word tokens. For example, a prevalent approach, known as chain-of-thought (CoT) reasoning ( 44 ) , involves prompting or training LLMs to generate solutions step-by-step using natural language. However, this is in stark contrast to certain human cognition results. Neuroimaging studies have consistently shown that the language network – a set of brain regions responsible for language comprehension and production – remains largely inactive during various reasoning tasks ( 2 ; 31 ; 29 ; 30 ; 11 ) . Further evidence indicates that human language is optimized for communication rather than reasoning ( 12 ) .

A significant issue arises when LLMs use language for reasoning: the amount of reasoning required for each particular token varies greatly, yet current LLM architectures allocate nearly the same computing budget for predicting every token. Most tokens in a reasoning chain are generated solely for fluency, contributing little to the actual reasoning process. By contrast, some critical tokens require complex planning and pose huge challenges to LLMs. While previous work has attempted to fix these problems by prompting LLMs to generate succinct reasoning chains ( 27 ) , or performing additional reasoning before generating some critical tokens ( 52 ) , these solutions remain constrained within the language space and do not solve the fundamental problems. On the contrary, it would be ideal for LLMs to have the freedom to reason without any language constraints, and then translate their findings into language only when necessary.

In this work we instead explore LLM reasoning in a latent space by introducing a novel paradigm, Coconut (Chain of Continuous Thought). It involves a simple modification to the traditional CoT process: instead of mapping between hidden states and language tokens using the language model head and embedding layer, Coconut directly feeds the last hidden state (a continuous thought) as the input embedding for the next token (Figure 1 ). This modification frees the reasoning from being within the language space, and the system can be optimized end-to-end by gradient descent, as continuous thoughts are fully differentiable. To enhance the training of latent reasoning, we employ a multi-stage training strategy inspired by 7 , which effectively utilizes language reasoning chains to guide the training process.

Interestingly, our proposed paradigm leads to an efficient reasoning pattern. Unlike language-based reasoning, continuous thoughts in Coconut can encode multiple potential next steps simultaneously, allowing for a reasoning process akin to breadth-first search (BFS). While the model may not initially make the correct decision, it can maintain many possible options within the continuous thoughts and progressively eliminate incorrect paths through reasoning, guided by some implicit value functions. This advanced reasoning mechanism surpasses traditional CoT, even though the model is not explicitly trained or instructed to operate in this manner, as seen in previous works ( 47 ; 21 ) .

Experimentally, Coconut successfully enhances the reasoning capabilities of LLMs. For math reasoning (GSM8k, 6 ), using continuous thoughts is shown to be beneficial to reasoning accuracy, mirroring the effects of language reasoning chains. This indicates the potential to scale and solve increasingly challenging problems by chaining more continuous thoughts. On logical reasoning including ProntoQA ( 35 ) , and our newly proposed ProsQA (Section 4 ) which requires stronger planning ability, Coconut and some of its variants even surpasses language-based CoT methods, while generating significantly fewer tokens during inference. We believe that these findings underscore the potential of latent reasoning and could provide valuable insights for future research.

## 2 Related Work

Chain-of-thought (CoT) reasoning. We use the term chain-of-thought broadly to refer to methods that generate an intermediate reasoning process in language before outputting the final answer. This includes prompting LLMs ( 44 ; 23 ; 53 ) , or training LLMs to generate reasoning chains, either with supervised finetuning ( 51 ; 49 ) or reinforcement learning ( 42 ; 22 ; 37 ; 48 ) . 27 classified the tokens in CoT into symbols, patterns, and text, and proposed to guide the LLM to generate concise CoT based on analysis of their roles. Recent theoretical analyses have demonstrated the usefulness of CoT from the perspective of model expressivity ( 13 ; 28 ; 26 ) . By employing CoT, the effective depth of the transformer increases because the generated outputs are looped back to the input ( 13 ) . These analyses, combined with the established effectiveness of CoT, motivated our design of feeding the continuous thoughts back into the LLM as input embeddings. While CoT has proven effective for certain tasks, its autoregressive generation nature makes it challenging to mimic human reasoning on more complex problems ( 24 ; 21 ) , which typically require planning and search. There are works that equip LLMs with explicit tree search algorithms ( 45 ; 47 ; 20 ) , or train the LLM on search dynamics and trajectories ( 25 ; 14 ; 39 ) . In our analysis, we find that after removing the constraint of a language space, a new reasoning pattern similar to BFS emerges, even though the model is not explicitly trained in this way.

Latent reasoning in LLMs. Previous works mostly define latent reasoning in LLMs as the hidden computation in transformers ( 46 ; 4 ) . 46 constructed a dataset of two-hop reasoning problems and discovered that it is possible to recover the intermediate variable from the hidden representations. 4 further proposed to intervene the latent reasoning by “back-patching” the hidden representation. 36 discovered parallel latent reasoning paths in LLMs. Another line of work has discovered that, even if the model generates a CoT to reason, the model may actually utilize a different latent reasoning process. This phenomenon is known as the unfaithfulness of CoT reasoning ( 41 ; 40 ) . To enhance the latent reasoning of LLMs, previous research proposed to augment it with additional tokens. 18 pretrained the model by randomly inserting a learnable <pause> tokens to the training corpus. This improves LLM’s performance on a variety of tasks, especially when followed by supervised finetuning with <pause> tokens. On the other hand, 32 further explored the usage of filler tokens, e.g., “ ... ”, and concluded that they work well for highly parallelizable problems. However, 32 mentioned these methods do not extend the expressivity of the LLM like CoT; hence, they may not scale to more general and complex reasoning problems. 43 proposed to predict a planning token as a discrete latent variable before generating the next reasoning step. Recently, it has also been found that one can “internalize” the CoT reasoning into latent reasoning in the transformer with knowledge distillation ( 8 ) or a special training curriculum which gradually shortens CoT ( 7 ) . 50 also proposed to distill a model that can reason latently from data generated with complex reasoning algorithms. These training methods can be combined to our framework, and specifically, we find that breaking down the learning of continuous thoughts into multiple stages, inspired by iCoT ( 7 ) , is very beneficial for the training. Other work explores alternative architectures for latent reasoning, including looped transformers ( 16 ; 10 ) , diffusion models in sentence embedding space ( 3 ) . Recurrent memory transformers similarly pass information across segments through continuous memory embeddings updated at each recurrent step rather than through language tokens ( 5 ) , with associative memory mechanisms later introduced to scale this recurrence to longer contexts ( 34 ) . Different from these works, we focus on general multi-step reasoning tasks and aim to investigate the unique properties of latent reasoning in comparison to language space. In addition to reasoning tasks, 33 also explored using continuous space for multi-agent communication. Building on Coconut , 55 developed a theoretical framework demonstrating that continuous CoT can be more efficient than discrete CoT on certain tasks by encoding multiple reasoning paths in superposition states. Subsequently, 54 analyzed the training dynamics to explain how such superposition emerges under the Coconut training objective.

## 3 Coconut : Chain of Continuous Thought

In this section, we introduce our new paradigm Coconut (Chain of Continuous Thought) for reasoning in an unconstrained latent space. We begin by introducing the background and notation we use for language models. For an input sequence x = ( x 1 , … , x T ) x=(x_{1},...,x_{T}) , the standard large language model ℳ \mathcal{M} can be described as:

H t = Transformer ​ ( E t ) H_{t}=\text{Transformer}(E_{t}) ℳ ⁡ ( x t + 1 ∣ x ≤ t ) = softmax ​ ( W ​ h t ) \mathcal{M}(x_{t+1}\mid x_{\leq t})=\text{softmax}(Wh_{t})

where E t = [ e ⁡ ( x 1 ) , e ⁡ ( x 2 ) , … , e ⁡ ( x t ) ] E_{t}=[e(x_{1}),e(x_{2}),...,e(x_{t})] is the sequence of token embeddings up to position t t ; H t ∈ ℝ t × d H_{t}\in\mathbb{R}^{t\times d} is the matrix of the last hidden states for all tokens up to position t t ; h t h_{t} is the last hidden state of position t t , i.e., h t = H t [ t , : ] h_{t}=H_{t}[t,:] ; e ⁡ ( ⋅ ) e(\cdot) is the token embedding function; W W is the parameter of the language model head.

Method Overview. In the proposed Coconut method, the LLM switches between the “language mode” and “latent mode” (Figure 1 ). In language mode, the model operates as a standard language model, autoregressively generating the next token. In latent mode, it directly utilizes the last hidden state as the next input embedding. This last hidden state represents the current reasoning state, termed as a “continuous thought”.

Special tokens <bot> and <eot> are employed to mark the beginning and end of the latent thought mode, respectively. As an example, we assume latent reasoning occurs between positions i i and j j , i.e., x i = x_{i}= <bot> and x j = x_{j}= <eot> . When the model is in the latent mode ( i < t < j i<t<j ), we use the last hidden state from the previous token to replace the input embedding, i.e., E t = [ e ⁡ ( x 1 ) , e ⁡ ( x 2 ) , … , e ⁡ ( x i ) , h i , h i + 1 , … , h t − 1 ] E_{t}=[e(x_{1}),e(x_{2}),...,e(x_{i}),h_{i},h_{i+1},...,h_{t-1}] . After the latent mode finishes ( t ≥ j t\geq j ), the input reverts to using the token embedding, i.e., E t = [ e ⁡ ( x 1 ) , e ⁡ ( x 2 ) , … , e ⁡ ( x i ) , h i , h i + 1 , … , h j − 1 , e ⁡ ( x j ) , … , e ⁡ ( x t ) ] E_{t}=[e(x_{1}),e(x_{2}),...,e(x_{i}),h_{i},h_{i+1},...,h_{j-1},e(x_{j}),...,e(x_{t})] . It is worth noting that the last hidden states have been processed by the final normalization layer, so they are not too large in magnitude. ℳ ⁡ ( x t + 1 ∣ x ≤ t ) \mathcal{M}(x_{t+1}\mid x_{\leq t}) is not defined when i < t < j i<t<j , since the latent thought is not intended to be mapped back to language space. However, softmax ⁡ ( W ​ h t ) \mathrm{softmax}(Wh_{t}) can still be calculated for probing purposes (see Section 5 ).

Training Procedure. In this work, we focus on a problem-solving setting where the model receives a question as input and is expected to generate an answer through a reasoning process. We leverage language CoT data to supervise continuous thought by implementing a multi-stage training curriculum inspired by 7 . As shown in Figure 2 , in the initial stage, the model is trained on regular CoT instances. In the subsequent stages, at the k k -th stage, the first k k reasoning steps in the CoT are replaced with k × c k\times c continuous thoughts 1 1 1 If a language reasoning chain is shorter than k k steps, then all the language thoughts will be removed. , where c c is a hyperparameter controlling the number of latent thoughts replacing a single language reasoning step. Following 7 , we also reset the optimizer state when training stages switch. We insert <bot> and <eot> tokens (which are not counted towards c c ) to encapsulate the continuous thoughts.

During the training process, we optimize the normal negative log-likelihood loss, but mask the loss on questions and latent thoughts. It is important to note that the objective does not encourage the continuous thought to compress the removed language thought , but rather to facilitate the prediction of future reasoning . Therefore, it’s possible for the LLM to learn more effective representations of reasoning steps compared to human language.

Training Details. Our proposed continuous thoughts are fully differentiable and allow for back-propagation. We perform n + 1 n+1 forward passes when n n latent thoughts are scheduled in the current training stage, computing a new latent thought with each pass and finally conducting an additional forward pass to obtain a loss on the remaining text sequence. While we can save any repetitive computing by using a KV cache, the sequential nature of the multiple forward passes poses challenges for parallelism. Further optimizing the training efficiency of Coconut remains an important direction for future research.

Inference Process. The inference process for Coconut is analogous to standard language model decoding, except that in latent mode, we directly feed the last hidden state as the next input embedding. A challenge lies in determining when to switch between latent and language modes. As we focus on the problem-solving setting, we insert a <bot> token immediately following the question tokens. For <eot> , we consider two potential strategies: a) train a binary classifier on latent thoughts to enable the model to autonomously decide when to terminate the latent reasoning, or b) always pad the latent thoughts to a constant length. We found that both approaches work comparably well. Therefore, we use the second option in our experiment for simplicity, unless specified otherwise.

## 4 Continuous Space Enables Latent Tree Search