24年2月来自Meta 田博的团队论文“ Beyond A∗: Better Planning with Transformers via Search Dynamics Bootstrapping”。
尽管 Transformer 在各种应用环境中取得了巨大进步，但此类架构在解决复杂决策任务方面仍然落后于传统的符号规划器。这项工作是有关如何训练 Transformer 来解决复杂的规划任务。其训练编码器-解码器 Transformer 模型来预测 A
搜索算法的搜索动态。对该模型进行微调，可获得 Searchformer，这是一个 Transformer 模型，可以在 93.7% 的时间内以最佳方式解决以前未见过的 Sokoban puzzle，同时使用的搜索步骤比最初用于训练的 A
实现少 26.8%。在训练方法中，A
的搜索动态表示为一个token序列，概述了在符号规划期间何时将任务状态添加到搜索树中以及从搜索树中删除任务状态。Searchformer 的表现明显优于直接预测最佳规划的基线，其模型要小 5-10 倍，训练数据集小 10 倍。
基于 Transformer 的架构 (Vaswani，2017) 在不同任务中表现出吸引人的性能，包括在人类层面进行对话 (Shuster，2022；OpenAI，2022、2023；Touvron，2023)、高质量图像理解 (Caron，2021；Oquab，2024；Assran，2023) 和视频生成 (Singer，2023)、多模态生成 (Girdhar，2023；Radford，2021) 和代码完成 (Roziere，2023；OpenAI，2021)。通过在互联网规模的数据集上训练这些架构，生成的模型（例如大语言模型）可以在实际用例中很好地推广。
尽管取得了这些成功，但基于 Transformer 的架构和 LLM 在解决规划和推理任务时仍然举步维艰。先前的研究表明，LLM 在多步规划任务（Valmeekam，2023a、b）或执行高阶推理（Momennejad，2023；Fan，2020）中表现不佳。
近年来，人们提出了各种方法来提高 Transformer 在这些环境中的性能。一种方法是模拟人类的思维过程并在输出响应之前产生中间“想法”。思维链 (CoT) 提示 (Wei，2022) 和思维树 (ToT) 方法 (Yao，2023) 鼓励模型一步一步地“思考”。虽然这些技术通常很有效，但它们也可能导致性能下降，例如由于自我强化 (Huang，2023)。此外，由于所涉及的推理类型变化（例如，空间推理与数学推理），在一个数据集上有效的技术可能不适用于其他数据集。如何使 Transformers 和 LLM 能够规划、解决多步决策任务并执行推理仍然难以捉摸，也是一个活跃的研究领域。
虽然现有研究（Trinh，2024；Ruoss，2024）利用合成数据集来学习推理策略，但本文研究在这方面有根本的不同，专注于提高 Transformer 权重中嵌入的推理能力。AlphaZero（Silver，2018）、MuZero（Schrittwieser，2020）和 AlphaGeometry（Trinh，2024）等现有算法，使用现有符号规划算法的输出来优化神经网络，这些算法的内部状态未被使用（即被视为黑盒）。例如，（Silver，2017）使用 MCTS 作为策略改进运算符来更新神经网络的权重。相比之下，所提出的搜索动态自举方法使用 Transformer 模型来推广到更高效的搜索模式，并改进模型本身。规划算法（连同其内部搜索动态）用于初始训练 Transformer 模型。
先前的工作重点是训练神经网络执行推理任务的轨迹（Nye，2021）或训练神经网络预测最佳动作（Yang，2022；Pallagani，2022；Ruoss，2024）。相比之下，本文专注于训练 Transformer 以在计算最佳规划时生成 A∗ 的整个搜索过程。提出的模型不是只预测单个动作，而是预测解决任务的整个多步规划。
在强化学习 (RL) 环境中，先前的研究已经研究了使用 Transformer 架构来解决复杂的顺序决策任务 (Chen，2021；Janner，2021；Laskin，2023)。然而，先前的研究提出了不同的方法来建模试错交互的轨迹，并侧重于预测下一个动作、状态或奖励或它们的组合。相比之下，本文演示如何使用 Transformer 来建模计算最佳多步规划所涉及的搜索步骤。MCTSNet (Guez，2018) 也试图学习搜索过程本身，但仍然将 MCTS 搜索过程 (Coulom，2006) 硬编码到神经网络中，这会导致二次反向传播开销，并且只能处理最多 500 步展开，而本文方法可以处理更长的搜索执行轨迹。Transformers 不仅可以模仿符号规划算法，还可以通过微调发现更有效的启发式方法。
如图概述合成数据集的生成过程。考虑两个领域：迷宫导航（图 (a)）和 Sokoban puzzle（图（d））。在迷宫导航中，目标是找到穿过 n×n 迷宫的最短路径。在 Sokoban 中，一个工人可以向上、向下、向左或向右移动，并且必须将每个箱子推到码头上才能解决难题。错误的移动可能会立即导致死胡同，因此需要跨多个时间步进行推理才能解决难题。难题中的每个状态都由箱子和工人位置的组合组成，这使得推箱子在计算上比迷宫导航更难解决。
（d）Sokoban puzzle
在图 (a) 中的迷宫示例中，每个节点都对应一个空的（非墙）网格单元。对于每个节点，算法都会计算一个启发式值和起始值的一个成本。在任何给定的迭代中，下一个搜索的节点由边界和封闭集的内容以及启发式值和起始值成本决定（图 (c)，左图）。A
的执行轨迹通过跟踪所有插入边界和封闭集的操作以及启发式值和起始值成本来收集（图 (c)，右图）。图 (c) 中的右图，显示图 (b) 中所示的迷宫示例结果轨迹。每一行对应于将节点插入边界（由create token表示）或将节点移动到封闭集（由close token表示）。每个节点都由其在迷宫中的 (x, y) 位置以及两个成本tokens表示。然后将结果规划附加到此轨迹。此轨迹的构造使得给定任何前缀，都可以正确预测下一个token。对于迷宫数据集，A* 用距目标位置的曼哈顿距离作为启发式方法。在 Sokoban 中，A* 首先将每个箱子与最近的码头匹配，然后计算每个箱子和码头对之间所有曼哈顿距离的总和。
对于每个实验，生成两个 token 序列变型，如上图所示：
由于每个模型都是从头开始训练的，因此生成的模型经过专门训练，仅预测概述一组不同规划任务的最佳规划序列。训练后，将解析并评估模型的输出是否包含最佳或可行的解决方案规划。
在生成 token 序列数据集时，每个任务都是唯一的，并且测试集的构造方式使其不包含任何训练集的重复项。通过这种实验设计，深入了解如何使用 Transformers 来解决规划任务并推广到以前从未见过的测试任务。
通过包含中间计算步骤，Transformer 模型经过训练可以有效地模仿 A* 算法执行的计算。与过程克隆 (Yang et al., 2022) 不同，在过程克隆中，神经网络被学习来预测最佳状态/动作序列（在我们的例子中是任务提示和最佳计划），本文的 Transformer 模型还学习预测整个思考过程，包括尝试但失败的路径，从而得出最佳规划。
对于每个实验，都使用编码器-解码器 T5 架构 (Raffel et al., 2020) 的改版，该架构集成旋转位置嵌入 (RoPE) (Su et al., 2023)。编码器处理训练序列的
部分，解码器处理
格式的序列（搜索增强模型）或仅处理
格式的序列（仅解决方案模型）。根据模型变型，每个网络都经过训练以最大化解码器生成分布与从训练数据集中采样相应序列的分布之间的交叉熵。
为了减少搜索增强模型在推理过程中生成的tokens数量，实现一种方法来迁移解码器生成执行轨迹的分布。首先，训练搜索增强模型来模仿 A 搜索的搜索动态。为了使用此搜索增强模型发现新的搜索动态并探索执行轨迹空间，搜索增强模型必须为同一任务提示生成不同的序列。在训练数据中引入非确定性来实现这一点，并使用非确定性的 A 实现，该实现随机打破成本关系并随机化子节点的扩展顺序。这种方法不会降低 A 搜索本身的效率，而只改变搜索不同节点的顺序，同时仍然尊重 A* 的启发式和成本计算。然后，生成的搜索增强模型将近似生成训练序列的概率分布。
