Training Large Language Models to Reason in a Continuous Latent Space

📄 arXiv: 2412.06769v3 📥 PDF

作者: Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston, Yuandong Tian

分类: cs.CL

发布日期: 2024-12-09 (更新: 2025-11-03)

备注: Accepted to COLM 2025


💡 一句话要点

提出Coconut:利用连续隐空间进行LLM推理,提升逻辑推理能力。

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 连续隐空间 推理 思维链 广度优先搜索

📋 核心要点

  1. 现有LLM推理受限于语言空间,存在冗余信息和规划困难的问题。
  2. Coconut将LLM的隐藏状态作为连续思维,在隐空间中迭代推理,避免过早确定路径。
  3. 实验表明,Coconut在逻辑推理任务上优于CoT,实现了更好的精度和效率平衡。

📝 摘要(中文)

大型语言模型(LLM)通常被限制在语言空间中进行推理,通过思维链(CoT)表达推理过程以解决复杂问题。然而,语言空间可能并非推理的最佳选择。大多数词语token主要确保文本连贯性,对推理并非至关重要,而一些关键token需要复杂的规划,对LLM构成挑战。为了探索超越语言的推理潜力,我们引入了一种名为Coconut(Chain of Continuous Thought,连续思维链)的新范式。Coconut利用LLM的最后一个隐藏状态作为推理状态的表示,称之为“连续思维”。我们不将此状态解码为单词,而是将其作为下一个输入嵌入直接反馈到连续空间中的模型。这种潜在的推理范式能够实现高级推理模式,其中连续思维可以编码多个替代的下一步骤,使模型能够执行广度优先搜索(BFS),而不是像CoT那样过早地提交到单个确定性路径。在需要大量搜索的逻辑推理任务中,Coconut优于CoT,并在准确性和效率之间实现了更好的权衡。

🔬 方法详解

问题定义:现有的大型语言模型在进行复杂推理时,通常采用思维链(Chain-of-Thought, CoT)的方法,即通过生成中间步骤的自然语言描述来辅助推理。然而,这种方法存在局限性:一是语言空间本身可能不是最优的推理空间,因为自然语言中包含大量与推理无关的信息;二是CoT方法本质上是一种深度优先搜索,容易陷入局部最优解,难以进行充分的探索和回溯。因此,如何让LLM在更有效的空间中进行推理,并具备更强的搜索能力,是本文要解决的问题。

核心思路:本文的核心思路是让LLM在连续的隐空间中进行推理,而不是在离散的语言空间中。具体来说,就是将LLM的最后一个隐藏状态(即“连续思维”)作为推理状态的表示,并将其直接反馈到模型作为下一个输入,从而实现迭代推理。这种方法避免了将推理状态解码为自然语言,减少了冗余信息的干扰,并允许模型探索多个可能的推理路径。

技术框架:Coconut的技术框架主要包括以下几个步骤:1. 使用LLM对输入问题进行编码,得到初始的隐藏状态(continuous thought)。2. 将该隐藏状态作为输入,再次输入到LLM中,得到新的隐藏状态。3. 重复步骤2,进行多次迭代,每次迭代都相当于在连续空间中进行一次推理。4. 在推理结束后,将最终的隐藏状态解码为自然语言答案。整个过程类似于在连续空间中进行广度优先搜索,模型可以同时探索多个可能的推理路径,并在最终选择最优的答案。

关键创新:Coconut最重要的技术创新点在于它将LLM的推理过程从离散的语言空间转移到了连续的隐空间。这种转变使得模型可以更有效地利用其内部的知识表示,避免了自然语言的冗余信息和歧义性。此外,Coconut还通过迭代推理的方式,实现了类似于广度优先搜索的推理策略,提高了模型探索复杂问题的能力。

关键设计:Coconut的关键设计包括:1. 使用LLM的最后一个隐藏状态作为连续思维的表示,这保证了推理状态能够充分利用LLM的知识。2. 直接将隐藏状态反馈到模型作为下一个输入,避免了解码和编码的过程,提高了效率。3. 通过控制迭代次数,可以调节模型的搜索深度和广度。4. 在训练过程中,可以使用强化学习等方法来优化模型的推理策略。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,Coconut在需要大量搜索的逻辑推理任务上显著优于CoT。例如,在某些任务上,Coconut的准确率比CoT提高了10%以上,并且在达到相同准确率的情况下,Coconut所需的计算资源更少,效率更高。这些结果表明,Coconut是一种有效的LLM推理方法,具有很大的应用潜力。

🎯 应用场景

Coconut的潜在应用领域包括逻辑推理、数学问题求解、规划和决策等。该研究的实际价值在于提升LLM在复杂推理任务上的性能,使其能够更好地解决实际问题。未来,Coconut可以与其他技术相结合,例如知识图谱、符号推理等,进一步增强LLM的推理能力。

📄 摘要(原文)

Large language models (LLMs) are typically constrained to reason in the language space, where they express the reasoning process through a chain-of-thought (CoT) to solve complex problems. However, the language space may not always be optimal for reasoning. Most word tokens primarily ensure textual coherence and are not essential for reasoning, while some critical tokens require complex planning and pose challenges to LLMs. To explore the potential of reasoning beyond language, we introduce a new paradigm called Coconut (Chain of Continuous Thought). Coconut utilizes the last hidden state of the LLM as a representation of the reasoning state, termed "continuous thought." Instead of decoding this state into words, we feed it back to the model as the next input embedding directly in the continuous space. This latent reasoning paradigm enables an advanced reasoning pattern, where continuous thoughts can encode multiple alternative next steps, allowing the model to perform a breadth-first search (BFS) rather than committing prematurely to a single deterministic path as in CoT. Coconut outperforms CoT on logical reasoning tasks that require substantial search during planning and achieves a better trade-off between accuracy and efficiency.