一个token就能控制模型快些解答或慢点思考。

OpenAIο1模型的发布掀起了人们对AI推理过程的关注，甚至让现在的AI行业开始放弃卷越来越大的模型，而是开始针对推理过程进行优化了。今天我们介绍的这项来自MetaFAIR田渊栋团队的研究也是如此，其从人类认知理论中获得了灵感，提出了一种新型Transformer架构：Dualformer。

根据人类认知理论，人类的思考受到两个系统控制：

System1：系统1，速度快，基于直觉。System2：系统2，速度更慢，更加深思熟虑。

近期有研究表明，如果将系统2过程整合进Transformer和大型语言模型中，就能显著提升它们的推理能力。尽管如此，如果模型只是模仿系统2式的思考过程，那就需要远远更高的计算成本才能完成，同时响应速度也会大幅减慢。

在研究这一难题时，田渊栋团队得到了一项惊人发现：在解决推理任务时，一种简单的数据方案就足以实现即时动态的系统1和系统2配置。

基于此发现，他们提出了Dualformer。这是一种可以轻松配置的Transformer——用户可以指定在推理过程中使用快速或慢速模式，在未指定时模型也可以自行决定。

论文标题：Dualformer:ControllableFastandSlowThinkingbyLearningwithRandomizedReasoningTraces论文地址：https://arxiv.org/pdf/2410.09918

具体而言，为了模仿系统2推理过程，他们让Transformer在包含推理轨迹和最终解答的数据上进行训练。利用推理步骤的结构，他们设计了特定的轨迹丢弃策略，使得生成的轨迹类似于系统1在思考过程中采取的捷径。在极端情况下，会丢弃整个轨迹并鼓励Transformer绕过所有中间步骤，直接输出最终解答。在训练时，他们的策略是随机选择这些结构化的轨迹丢弃策略。

前提准备

他们的这项研究基于田渊栋团队之前的另一项研究《BeyondA*:Betterplanningwithtransformersviasearchdynamicsbootstrapping》，参阅机器之心报道《补齐Transformer规划短板，田渊栋团队的Searchformer火了》。为了执行规划，他们要训练一个Transformer来建模一个token序列，而该序列则是以顺序方式来表示该规划任务、A*算法的计算、由A*搜索得到的最优解。

图3.1展示了其token化方法，其中示例是一个3×3迷宫的导航任务，目标是找到从起点到目标单元格的最短路径。

A*算法已经成功找到了最佳规划。这里使用一个token序列来表示该任务和迷宫结果，其也被用作Dualformer的提示词。该解答由使用坐标描述路径的规划token序列描述。A*算法生成一个搜索轨迹序列，记录执行的搜索动态，如图4.1所示。

回想一下，A*算法是一种在加权图上的寻路算法。create子句将节点（由后续坐标表示）添加到搜索边界中，close子句将节点添加到该闭集。每个子句（create或close）后面都跟着tokenx、y、c0和c1——分别表示节点的坐标、自开始以来的成本值和启发值。

结构化轨迹丢弃和随机训练

田渊栋团队之前提出的Searchformer已被证明可以有效解决多种复杂的决策任务。但是，它仍有两个不足。

1.模型仅能以慢速模式运行并会输出很长的推理链，这会极大延长推理时间。尽管可通过bootstrapping（一种迭代优化技术，包含rollout循环和之后的微调过程）来提速，但这样的过程会对计算资源产生显著的额外需求。

2.Searchformer很难生成多样化的解答，因为其经常会采样相同的rollout。举个例子，在他们测试过的1000个30×30迷宫问题中，Searchformer的推理链平均包含1500多个token，而只能在64个响应中找到7.6条各不一样的可行路径。

为了解决这些挑战，他们提出了一个利用随机化推理轨迹的训练框架。该方法的灵感来自两个研究方向：

该团队注意到，即便Searchformer是在完整的A*搜索轨迹上训练的，但它也会生成更短的勾勒搜索过程的轨迹。研究表明，人类在做决策时往往依赖捷径和模式，这一概念被称为系统1思维。

这些观察再加上dropout技术（在训练时随机丢弃神经网络中的一些单元）的成功，促使该团队研究了随机化推理轨迹的作用，并且他们还希望通过利用结构化元素并选择性地丢弃每个训练示例的某些部分来简化A*搜索轨迹。该方法的细节如下。

如图4.1所示，A*搜索轨迹包含create和close子句，每个子句都包括节点的坐标及其到达起始位置和目标位置的（估计）成本。为了推导得到Dualformer，他们利用了搜索轨迹的结构，并为每个训练示例丢弃轨迹中的某些部分。其有三种自然的丢弃类型：

D1：丢弃一个close子句；D2：丢弃一个子句中的成本token；D3：丢弃一个create子句。

基于此，他们开发出了四个层级逐层递进的丢弃策略：

Level1：去除搜索轨迹中所有close子句。Level2：更进一步，额外丢弃所有成本token。Level3：更加激进，进一步随机丢弃30%的create子句。Level4：丢弃整条搜索轨迹。

图4.1基于上述迷宫任务演示了这些策略。后面我们会看到，这些策略可有效地引导Dualformer学习更简洁、更高效的搜索和推理过程。

为了提升训练数据的多样性，他们没有将丢弃作为一个数据预处理步骤。而是在推理时间，对于一个数据批次中的每个训练样本，都从一个分类分布Cat(p_0,p_1,p_2,p_3,p_4)中随机抽取丢弃策略，其中p_1,...,p_4是执行Level1-4丢弃的概率，p_0是保持完整轨迹的概率。这种训练框架可使Dualformer学习多个经过约简的轨迹，即使对于单个训练示例也是如此，因为同一个示例可能出现在多个批次中。

可控式生成

Dualformer具有一个非常吸引人的特性：在推理时，可以轻松地通过提示词指定以快速或慢速生成模式运行。

该控制机制非常简单：在标准提示词之后添加一个bos和一个控制token，其中控制token是plan或create中的一个。

如果使用plan，则Dualformer将以快速模式运行，绕过推理步骤并直接输出规划。另一方面，如果在bos之后注入create，则Dualformer将以慢速模式工作并生成推理轨迹和最终规划。下面基于迷宫任务展示了这两种模式的示意图。

而如果仅使用标准提示词，则Dualformer将模仿人类决策的双重过程——根据情况，它会选择一种分别对应于系统1和系统2的推理类型进行响应。

实验

实验的目标是解答以下三个问题：

1.Dualformer在快速、慢速和自动模式下的表现是否优于相应的基线？

2.在慢速模式下，Dualformer是否能实现更快的推理，即输出更短的轨迹？

3.结构化的轨迹丢弃技术是否适用于在自然语言数据集上训练的LLM？

为了解答问题1和2，该团队训练了求解迷宫导航任务和紧密相关的推箱子（Sokoban）任务的Transformer。为了解答问题3，他们微调了LLama-3.1-8B和Mistral-7B模型来解答数学问题。

导航任务：迷宫和推箱子

迷宫和推箱子任务使用的数据集与Searchformer研究的一样。这里就不再赘述，我们直接来看结论。

研究表明，Dualformer可以根据控制指令选择快速或慢速的运行模式。在快速模式下，它仅输出最终规划；在慢速模式下，它还会生成推理轨迹。该团队在不同的模式下让Dualformer对比了不同的基线。使用的指标包括生成规划的正确性、最优性和多样性、推理轨迹的长度等。

快速模式

可以看到，在生成正确和最优规划方面，Dualformer在1-Solved-64和1-Optimal-64指标上中都明显优于基线。它在3-Solved-64和3-Optimal-64指标上也明显超过了基线，这证明了Dualformer在规划生成方面的稳健性。

尤其需要注意，随着任务难度提升，Dualformer的优势也会增大。对于最大的30×30迷宫，Dualformer的1-Optimal-64成功率是仅解答模型的2.8倍，在3-Optimal-64上是2.97倍。

Dualformer的SWC分数也比基线高得多——在每个环境中都高于0.9。这表明Dualformer生成的每个单独规划的质量都很高，其成本非常接近最佳规划。

在实验考虑的所有问题上，Dualformer还能稳定地生成更多样化的规划。比如在下面这个迷宫示例中，随着迷宫规模的增加，Dualformer的多样性得分（即64个响应中不同但正确的规划的平均数量）会增加。

一般来说，随着迷宫规模增大，到达单个目标位置的可能路线也越来越多。这表明Dualformer学习了迷宫结构，而仅解答模型可能是记住了最佳规划，因为其多样性得分在所有迷宫规模下都接近1。

慢速模式

相应的基线是Complete-Trace模型，它使用相同的架构并在具有完整A*搜索轨迹的数据上进行了训练。除了之前报告的指标之外，该研究还报告了在所有1000个评估任务中汇总的64个响应的推理轨迹平均长度。结果表明，Dualformer实现了更好的规划能力和推理速度。它在所有正确性和最优性指标方面都优于Complete-Trace模型：包括解决率、最优率和SWC。

此外，Dualformer产生的推理轨迹明显短于基线模型。平均而言，Dualformer在五个任务中将轨迹长度减少了49.4%。与以前一样，与基线相比，Dualformer还生成了更多不同的规划。

与搜索动态引导的比较

Complete-Trace模型是田渊栋团队的基本Searchformer模型。该方法还提出了一种搜索动态引导方法来提高其在推箱子任务上的性能，类似于Anthony等人（2017）；Zelikman等人（2022）的研究。

在训练Searchformer模型后，作者在新创建的自引导数据集上对其进行微调。对于原始数据集中的每个推箱子竞赛，此处生成32个答案，并将最短的最佳答案纳入新数据集。我们可以多次重复此过程。

通过这种方式，Searchformer学会了生成更短的答案。表5.4将Dualformer与最多微调3步的Searchformer模型进行了比较。Dualformer在大多数指标上与引导模型相当或更好，同时仅使用不到45.1%的推理步骤。

该团队发现，每个引导步骤需要推出3.2×10^6个总响应和10^4次迭代的额外微调。这意味着包括8×10^5次预训练迭代。Searchformer步骤3总共需要8.3×10^5次训练迭代和9.6×10^6次rollout，计算成本很高。相比之下，Dualformer只需要一个由8×10^5次迭代组成的训练阶段，没有额外的rollout需求。

自动模式

不仅能通过在bos之后注入控制token的方式来控制Dualformer的推理模式，还可以直接执行采样，使其自由确定操作模式，类似于人类决策的双重过程。这种Dualformer被称为自动模式。表5.3报告了结果。对于这里考虑的所有任务，自动模式Dualformer也优于Complete-Trace和Solution-Only模型。

大模型训练中的应用：数学推理

作者展示了结构化轨迹丢弃技术在训练大规模LLM解决数学问题方面的有效性。具体来说，作者使用了包含各种数学问题和答案的数据集对Llama-3-8B和Mistral-7B模型进行微调，其中包含详细的推理步骤。其中使用了一种轨迹丢弃技术，该技术也利用了数学问题的推理轨迹的特定结构。

最后，作者再对生成的模型与直接在数据集上微调的相应基础模型进行基准测试。

结果见表5.6。作者共测试了p的四个值：0.1、0.2、0.3和0.4。结果表明，新研究所提出的训练策略使这两个LLM更加有效和高效。

首先来看Mistral-7B模型的结果。对于慢速模式推理，使用轨迹丢弃和随机训练对模型进行微调可以改进直接在Aug-MATH数据集上微调的基线模型。当p=0.1时，绝对Greedy@1指标提高了1.7%（相当于10%的相对性能提升），当p=0.2和0.3时提高了0.9%，当p=0.4时提高了0.1%。当p=0.1、0.2和0.3时，新模型也优于Pass@20指标的基线模型，其中绝对正确率增加到61.9%。在两种评估方案下，推理轨迹的平均长度随着p的增加而下降。

同样，对于快速模式下的推理，新模型也实现了更高的正确率。Llama-3-8B模型也具有类似的性能改进趋势。最后，为了供读者参考，作者还列出了在原始MATH数据集上微调的Mistral-7B和Llama-3-8B模型的结果。

未经允许不得转载：头条资讯网_今日热点_娱乐才是你关心的时事 » 补齐Transformer规划短板又不放弃快速思考，Dualformer双重优势