Article Overview: Training Large Language Models to Reason in a Continuous Latent Space

Link: https://arxiv.org/abs/2412.06769

This paper introduces COCONUT (Chain of Continuous Thought), a new paradigm exploring the potential for Large Language Models (LLMs) to reason within an unrestricted latent space rather than being confined to natural language. The core idea challenges the optimality of language space for complex reasoning, noting that many tokens serve textual coherence rather than reasoning, while critical steps can be challenging for LLMs. This contrasts with neuroimaging studies suggesting human reasoning doesn’t always heavily rely on language centers.

Key Concepts

• Latent Reasoning: COCONUT utilizes the LLM’s last hidden state as a “continuous thought” representing the reasoning state. Instead of decoding this state into a token, it’s fed directly back into the model as the next input embedding, allowing reasoning to proceed in the continuous latent space (as illustrated in Figure 1).
• Mode Switching: The model switches between a standard “language mode” and the new “latent mode”. Special tokens, <bot> and <eot>, mark the beginning and end of the latent reasoning phase.
• Training: A multi-stage training curriculum is employed, inspired by prior work (iCoT). Starting with standard language Chain-of-Thought (CoT) data, language reasoning steps are gradually replaced by continuous thoughts across training stages (visualized in Figure 2). The process is end-to-end differentiable.
• Inference: During inference, the main challenge is determining when to switch out of latent mode. The paper primarily uses a fixed length for latent thought sequences for simplicity.

Evaluation & Findings

Datasets

The approach was tested on mathematical reasoning (GSM8k) and logical reasoning tasks requiring planning (ProntoQA and a new, more challenging dataset, ProsQA).

Baselines

COCONUT was compared against standard CoT, No-CoT (direct answer generation), iCoT (internalized CoT), and Pause Token methods.

Performance

• COCONUT consistently improved reasoning over No-CoT methods and outperformed standard CoT on logical reasoning tasks like ProntoQA and ProsQA, often using significantly fewer tokens (see Table 1).
• On GSM8k, COCONUT showed benefits, and performance increased with more continuous thoughts per step (c=1, c=2), suggesting a “chaining” effect similar to CoT (Figure 3). Performance slightly dropped for c=3, potentially due to training instability.
• The multi-stage training curriculum proved essential; training directly on the final objective performed poorly (COCONUT w/o curriculum in Table 1).

Emergent Behaviors

• COCONUT demonstrated an advanced reasoning pattern akin to Breadth-First Search (BFS). Continuous thoughts can encode multiple potential next steps, allowing the model to explore options in parallel and prune incorrect paths later, rather than committing to a single path early like CoT.
• This latent planning advantage was particularly evident on the ProsQA dataset, where standard CoT showed little improvement over No-CoT (Table 1).

Analysis & Interpretation

• Latent thoughts can sometimes be decoded into meaningful intermediate steps (Figure 4).
• The latent reasoning process can be interpreted as a search tree (Figures 6, 7). The model implicitly assigns values (probabilities) to different paths, pruning less promising ones.
• Analysis suggests greater parallelism (exploration of alternatives) in earlier continuous thoughts compared to later ones (Figure 8).
• Latent reasoning aids planning by delaying commitment; it’s easier for the model to evaluate nodes closer to the leaves of the search tree (Figure 9).
• Interestingly, even when forced to generate a full language chain (k=0), COCONUT trained with the mixed-stage curriculum outperformed standard CoT, suggesting the training improves planning (Figure 5).

Conclusion

COCONUT demonstrates the potential of LLM reasoning in a continuous latent space, offering advantages in efficiency and planning capabilities, especially for tasks requiring exploration and backtracking. The emergence of BFS-like search patterns highlights the flexibility gained by removing language constraints. Future work includes refining training methods, potentially pretraining with continuous thoughts, and exploring hybrid language/latent reasoning approaches.