田渊栋团队论文火了!连续思维链优于思维链,打开LLM推理新范式
来源:机器之心微信公众号
2024-12-17 11:05:05
在认知科学领域,关于语言是用于思考还是用于交流的辩论一直持续。
随着 LLM (Large Language Model, 大语言模型)和 CoT (Chain of Thought,思维链)的兴起,语言已经成为机器推理的默认媒介 —— 但它真的是最佳方法吗?
一般而言,LLM 被限制在语言空间(language space)内进行推理,并通过CoT来表达推理过程,从而解决复杂的推理问题。
然而,语言空间可能并不总是最适合推理的。例如,很多单词 token 主要用于文本连贯性,而不是推理本身,而一些关键 token 则需要复杂的规划,这种差异给 LLM 带来巨大的挑战。
为了探索 LLM 在不受限制潜在空间中的推理潜力,而非使用自然语言,来自 Meta、加州大学圣地亚哥分校的研究者提出了一种新的范式 ——Coconut(连续思维链,Chain of Continuous Thought),来探索 LLM 在潜在空间中的推理。
论文标题:Training Large Language Models to Reason in a Continuous Latent Space
论文地址:https://arxiv.org/pdf/2412.06769
Coconut 涉及对传统 CoT 过程的简单修改:Coconut 不再通过语言模型头(language model head)和嵌入层将隐藏状态与语言 token 进行映射,而是直接将最后的隐藏状态(即连续思维)作为下一个 token 的输入嵌入(如图 1 所示)。
这种修改将推理从语言空间中解放出来,并且由于连续思维是完全可微的,因此可以通过梯度下降对系统进行端到端优化。为了增强潜在推理的训练,本文采用了多阶段训练策略,该策略有效地利用语言推理链来指导训练过程。
这种范式带来了高效的推理模式,与基于语言的推理不同,Coconut 中的连续思维可以同时编码多个潜在下一步,从而实现类似于 BFS(breadth-first search)的推理过程。尽管模型在初始阶段可能做出不正确的决策,但它可以在连续思维中保持许多可能的选项,并通过推理逐步排除错误路径,这一过程由一些隐含的价值函数引导。这种高级的推理机制超越了传统的 CoT,即使模型并没有显式地接受训练或指示以这种方式操作。
实验表明,Coconut 成功增强了 LLM 的推理能力。对于数学推理(GSM8k),使用连续思维被证明有利于提高推理准确率,这与语言推理链的效果相似。通过链接更多连续思维,可以扩展和解决日益具有挑战性的问题。
在逻辑推理方面,包括 ProntoQA 和本文新提出的 ProsQA,这需要更强的规划能力,Coconut 及其一些变体甚至超越了基于语言的 CoT 方法,同时在推理过程中生成的 token 明显更少。
这项研究在 X 上的讨论量非常高,其中单人转发的浏览量就高达 20 多万。
连续思维链:Coconut
方法概述。在 Coconut 方法中,LLM 在语言模式和潜在模式之间切换(图 1):
在语言模式下,该模型作为标准语言模型运行,自回归生成下一个 token。
在潜在模式下,它直接利用最后一个隐藏状态作为下一个输入嵌入。这个最后的隐藏状态代表当前的推理状态,称为连续思维。特殊 token < bot >、< eot > 分别用于标记潜在思维模式的开始和结束。
训练。本文专注于问题 - 解决设置,其中模型接收问题作为输入,并通过推理过程生成答案。作者利用语言 CoT 数据来监督连续思维。如图 2 所示,在初始阶段,模型在常规 CoT 实例上进行训练。在后续阶段,即第 k 阶段,CoT 中的前 k 个推理步骤被替换为 k × c 个连续思维,其中 c 是一个超参数,用于控制替换单个语言推理步骤的潜在思维的数量。
推理过程。Coconut 的推理过程类似于标准的语言模型解码过程,不同之处在于,在潜在模式下,本文直接将最后一个隐藏状态作为下一个输入嵌入。这样做面临的挑战是确定何时在潜在模式和语言模式之间切换。当专注于问题 - 解决设置时,本文会在问题 token 后立即插入一个 < bot >token。对于 < eot >,作者考虑两种潜在策略:a) 在潜在思维上训练二元分类器,使模型能够自主决定何时终止潜在推理,或 b) 始终将潜在思维填充到恒定长度。本文发现这两种方法效果都相当好。除非另有说明,本文在实验中使用第二种选项以简化操作。
实验
研究团队通过三个数据集验证了大语言模型在连续潜空间中进行推理的可行性。实验主要评估模型生成答案的准确性和推理效率。
实验涉及两类主要任务:数学推理和逻辑推理。数学推理使用 GSM8k 数据集。逻辑推理则采用了两个数据集:5-hop ProntoQA 与该团队自行开发的 ProsQA。
ProntoQA 给出一个层级分类的知识结构,要求模型判断不同类别之间的从属关系是否正确。而 ProsQA 中是更具挑战性的推理任务,包含许多随机生成的有向无环图,要求模型进行大量规划和搜索。
实验设置
在实验设置方面,研究采用预训练的 GPT-2 模型,学习率为 1×10^−4,批量大小为 128。
对于数学推理任务,每个推理步骤使用 2 个潜在思维向量表示,整个训练过程分为 4 个渐进式阶段。
在逻辑推理任务中,每步使用 1 个潜在思维向量,训练分为 7 个渐进式阶段,逐步增加难度。所有实验均在标准训练流程后继续训练至第 50 轮,并通过在验证集上评估准确率来选择性能最佳的模型检查点用于最终测试。
基线方法和各种版本的 Coconut
为了全面评估方法效果,研究团队设置了以下基线方法进行对比:
1. 传统的 CoT:使用完整的思维链进行训练,让模型生成每一步的推理过程
2. No-CoT:模型直接生成最终答案,不要求中间推理步骤
3. iCoT:采用渐进式策略,逐步移除推理链中的步骤
4. Pause token:在问题和答案之间插入特殊的暂停 token
同时,他们还评估了 Coconut 的三个变体版本:
1. 无课程学习版本:跳过渐进训练,直接采用最终阶段的训练方式
2. 无思维版本:移除连续思维表示,仅保留分阶段训练机制
3. 思维替换版本:用特殊 token 替代连续思维的表示方式
结果与讨论 表 1 展示了所有数据集的整体结果。连续思维有效增强了大语言模型的推理能力,这从其相比无 CoT 基线的一致性提升可以看出。在 ProntoQA 和 ProsQA 上,其表现甚至超过了 CoT。
研究团队从实验中得出了以下几个关键结论:
连续思维的「链式」组合增强了推理能力。
在传统 CoT 中,输出 token 会作为下一步的输入,这被既有的研究证明可以增加模型的有效深度和表达能力。
该团队进一步探索了这一特性是否也适用于潜空间推理,因为这意味着这种方法可以通过链接多个潜在思维来解决更复杂的问题。
在 GSM8k 数据集的实验中,Coconut 的表现优于其他采用类似策略训练的架构,特别是超过了最新的 iCoT 基线,也显著优于同样能增加计算能力的 Coconut(pause as thought)变体。
虽然此前的研究已经证明特殊的 token 可以解决高度并行化的问题,该研究团队的结果显示 Coconut 架构在一般问题上更有效,比如数学应用题这种后续步骤高度依赖前序步骤的问题。
此外,当调整控制每个语言推理步骤对应潜在思维数量的参数 c 时(见图 3),随着 c 从 0 增加到 1 再到 2,模型性能稳步提升。这表明类似 CoT 的链式效应在潜空间中也存在。
潜在空间推理在规划密集型任务中优于语言推理。复杂推理往往要求模型「前瞻」并评估每一步的合理性。在研究团队的数据集中,GSM8k 和 ProntoQA 由于问题结构直观且分支有限,相对容易预测下一步。相比之下,ProsQA 的随机生成 DAG 结构显著挑战了模型的规划能力。
如表 1 所示,CoT 相比 No-CoT 并无明显改进。然而,Coconut 及其变体和 iCoT 在 ProsQA 上大幅提升了推理能力,表明潜空间推理在需要大量规划的任务中具有明显优势。
模型仍需指导来学习潜在空间推理
理想情况下,模型应该能通过问答数据的梯度下降自动学习最有效的连续思维(即无课程学习版本的 Coconut)。然而,实验结果显示这种训练方式的表现并不优于 no-CoT。将训练分解按照目标划分为多阶段课程,Coconut 在各种任务中都取得了最佳性能。
连续思维是推理的高效表示
虽然连续思维最初并不是为了转换成具体的语言文字而设计的,但该团队发现它可以用来直观地解释推理过程。
图 4 展示了一个由 Coconut(c=1)解决的数学应用题案例研究。第一个连续思维可以解码为「180」、「180」(带空格)和「9」等 token。这个问题的推理过程应该是 3×3×60=9×60=540,或 3×3×60=3×180=540。
这恰好对应了解题过程中的第一步中间计算结果(3×3×60 可以先算出 9 或 180)。更重要的是,连续思维能够同时包含多种不同的解题思路,这种特性使它在需要复杂规划的推理任务中表现出色。 理解 Coconut 中的潜在推理机制
接下来,作者使用 Coconut 的一个变体对潜在推理过程进行了分析。
模型:Coconut 允许通过在推理期间手动设置 < eot > 的位置来控制潜在思维的数量。当强迫 Coconut 使用 k 个连续思维时,该模型预计将从第 k + 1 步开始,用语言输出剩余的推理链。实验采用 k∈{0,1,2,3,4,5,6} 在 ProsQA 上测试 Coconut 的变体。
图 5 展示了在 ProsQA 上对不同推理方法的对比分析。随着更多的推理在连续思维中进行(k 值增加),最终答案的准确率(图 5 左)以及正确推理过程的比例(图 5 右中的 Correct Label 和 Correct Path)都得到了提高。此外,幻觉(Hallucination)和错误目标(Wrong Target)的发生率也下降。这也表明,当更多的推理发生在潜在空间中时,模型的规划能力得到了提升。 图 6 为一个案例研究,其中 CoT 产生了幻觉,Coconut(k = 1)导致错误的目标,但 Coconut(k = 2)成功解决了问题。在此示例中,模型无法准确确定在早期步骤中选择哪条边。但是,由于潜在推理可以避免在前期做出艰难的选择,因此模型可以在后续步骤中逐步消除不正确的选项,并在推理结束时实现更高的准确率。
潜在搜索树的解释
由于连续思维可以编码多个潜在的下一步,潜在推理可以被解释为一个搜索树,而不仅仅是推理「链」。以图 6 为例,第一步可以选择 Alex 的任一子节点:{lempus, sterpus, zhorpus, grimpus}。 图 7 左侧展示了所有可能的分支。同样,第二步的前沿节点是 Alex 的孙节点(图 7 右侧)。
与标准广度优先搜索不同,模型展现出优先探索有希望的节点同时剪枝不相关节点的能力。通过分析模型在语言空间的后续输出,研究团队发现了模型的偏好。例如,当模型在一个潜在思维后切换回语言空间(k=1),它会以结构化格式预测下一步。通过检查概率分布,研究团队得到了根节点 Alex 的子节点的数值(图 7 左)。同样,当 k=2 时,也获得了所有前沿节点的预测概率(图 7 右)。
图 8 展示了模型如何在潜在思维空间中进行推理。在第一个潜在思维阶段,模型会同时考虑多个可能的推理方向,保持思维的多样性。到了第二个潜在思维阶段,模型会逐步缩小范围,将注意力集中在最可能正确的推理路径上。这种从发散到收敛的推理过程,体现了模型在潜在空间中的推理能力。
为什么潜在空间更适合规划?
在这一节中,研究团队探讨了潜在推理在规划中的优势。例如,图 6 中的「sterpus」是叶节点,无法通向目标节点「bompus」,容易被识别为错误选项。相比之下,其他节点有更多后续的节点需要探索,推理难度更大。
研究团队通过测量节点在树中的高度(到叶节点的最短距离)来量化探索潜力。他们发现高度较低的节点更容易评估,因为探索潜力有限。在图 6 中,模型对高度为 2 的「grimpus」和「lempus」节点表现出更大的不确定性。
为了更严格地验证这个假设,研究团队分析了测试集中第一步和第二步潜在推理过程中模型预测概率与节点高度之间的相关性。图 9 揭示了一个规律:当节点高度较低时,模型会为错误节点分配较低值,为正确节点分配较高值。 然而,随着节点高度增加,这种区分变得不那么明显,表明评估难度增大。总之,这些发现突出了利用潜在空间进行规划的优势。模型通过延迟做出决策,并在潜在推理过程中不断探索,最终将搜索推向树的终端状态,从而更容易区分出正确和错误的节点。 更多研究细节,请参阅原文。
参考链接: https://arxiv.org/pdf/2412.06769
责任编辑: 张震