Transformers struggle to learn search

小 transformer 能否在去掉捷径的数据上学会 DAG 搜索，以及大图上的困难能否靠加参数或中间 token 消除。

• 场景：作者认为搜索是推理、规划与导航的基础。已有工作显示基于 transformer 的 LLM 不能稳健做证明搜索，但不清楚是缺数据、参数不够，还是架构限制；规模继续增大后能否搜索与规划仍不清楚。
• 思维链仍可能不够：思维链把搜索拆成反复预测下一步。作者认为最坏情况下，为避免走错，模型仍须在一次前向传播内完成搜索；走错后会出现错误或幻觉。
• 分布会决定学到什么：朴素随机 DAG 的 lookahead 通常为 1 或 2，更大 lookahead 的概率指数下降，启发式足以解题。作者因此要求训练分布使捷径无用。
• 本文做法：以“给定 DAG、起点和可达目标，预测路径上下一顶点”生成实质上无限的理想数据，训练小 transformer，估计学会搜索的上限；再重建计算图，并考察图规模、参数规模，以及 DFS 与 selection-inference 这类中间 token。

论文第 2 节把相关工作分成搜索能力、机制解释和缩放定律，并在搜索能力中对照可表达性结果。

图搜索与推理

• Fan 等（2024）、Fu 等（2024）、Sanford 等（2024a;b）研究图推理、重复函数复合和 k-hop induction。论文称这些基准相对本文考虑的图仍然偏小。
• Ruoss 等（2024）、Gandhi 等（2024）、Shah 等（2024）报告 transformer 可以近似或模拟搜索算法，但有性能差距，且未检验加大规模或训练能否缩小差距。
• 论文转述：Wang 等（2023）、Bachmann 与 Nagarajan（2024）显示 LLM 能做一些图推理，但会被虚假相关误导；Zhang 等（2023）发现 transformer 强烈偏好启发式，因而学不会证明搜索。Zhang 等（2024）、Borazjanizadeh 等（2024）报告真实图推理与多样搜索题上的困难。作者的对照是：若额外去掉启发式，transformer 可以学会搜索。

可表达性

• Merrill 与 Sabharwal（2024）报告有思维链时 transformer 可模拟任意图灵机。论文称这并不表示能否训练成功；本文在允许中间 token 时仍难以在更大图上学会搜索。
• Sanford 等（2024a;b）给出图连通需要对数层。作者认为这与本文识别出的逐层扩展可达集相符。

机制解释与缩放

• Hou 等（2023）、Kim 等（2024）在激活和注意力中寻找事先指定的电路。本文不预设算法，而从激活与注意力重建计算图，并使用 activation patching（Vig 等，2020；Geiger 等，2021）。
• Brinkmann 等（2024）、Stolfo 等（2023）等把机制解释用于推理；Yang 等（2024）、Jenner 等（2024）给出 LLM 在前向传播中做浅层搜索的证据；Ivanitskiy 等（2023）在迷宫任务上看到地图的内部表示。
• Kaplan 等（2020）、Hoffmann 等（2022）等讨论规模、数据与算力。Caballero 等（2023）的任务不含搜索、推理或规划。Wei 等（2022a）提出能力可能随规模涌现；Schaeffer 等（2023）认为这种涌现可由指标选择解释；Du 等（2024）主张用训练损失衡量能力。

对照设置：准确率实验不使用外部搜索算法作基线，而是比较 naïve、star、balanced 三种自生成分布；机制分析用随机初始化模型作对照。作者把本文定位为：表达能力不保证可训练，去掉启发式后可以学会搜索，但加大规模或中间 token 并未消除更大图上的学习困难。

balanced distribution 上从头训练小 transformer，预测 DAG 路径的下一顶点；再用激活修补重建计算图，检验学到的是逐层合并可达集还是捷径。

任务与三种分布

• 每个样本含 DAG、起点和保证可达的目标。输出是起点到目标路径上的下一顶点。难度用 lookahead：设 P 为该路径，S_i 为从起点出发、其余部分与 P 不交的路径，则 L=min{|P|,max_i|S_i|}，即至少要向前看的步数。
• naïve：顶点先按拓扑序排列，边只向右，入度上限为 4，以免图过密、lookahead 过小；再随机打乱顶点编号。附录 A.1.1 给出父顶点个数的采样。Figure 8 用 41 顶点的 10M 张图显示大 lookahead 很少；作者称大于 6 虽可能，但几乎不会被采样到。
• balanced：给定 L，先建一条 L 条边的真路径，再加长度约为 L 的分支和指向起点的入链，并用偏高入度或出度的附加边去掉“起点是唯一源点、其余入度恰为 1”等捷径。训练时 L 均匀。细节在附录 A.1.2。
• star：中心是起点，k 条长度为 L 的链向外，目标在其中一条末端；k 与 L 均匀。论文称形状对应 Bachmann 与 Nagarajan（2024）的图。

模型、损失与判定

• 架构同 GPT-2，激活改为 ReLU。token 与绝对位置都用 1-hot，拼接而不是相加，以便做机制解释。每层 1 个注意力头，前馈维等于模型维。边顺序随机，注意力不用 causal mask。分布实验为 6 层、隐藏维 16、输入 128 token。
• 流式采样，不固定训练集。最初若干 batch 留作测试，其后用精确匹配去掉测试中出现过的样本。优化器用 Sophia，学习率 10^{-5}，权重衰减 0.1，无 dropout；除另有说明外 batch 为 1024。最小化交叉熵。
• 存在多条路径时，训练均匀随机选一条作标签。评测时，预测顶点落在任意一条通往目标的路径上即算对。

计算图重建

对单个输入做五步。先普通前向。再把最后一层、最后 token 的注意力权重改成 0，或改成该行最大值后重新归一化；预测 logit 下降超过 α 则标为重要。若这一步没有重要操作，则退而标记使 logit 下降最大的操作。对其余层做同样扰动：若最后一层重要操作的 log 注意力变化超过 κ1/√d（d 为模型维），或预测 logit 下降超过 α，也标为重要。然后做 activation patching：把某个顶点编号换成输入中未出现的编号，或把某位置的位置嵌入置零；冻结当前层之前的注意力，并保持原先被 ReLU 置零的位置仍为零，用 query 与 key 的点积变化判断源、目标嵌入中哪些输入特征因果上影响该注意力。最后只保留从输入沿可解释注意力走到输出预测的操作。论文给出该过程需要大量前向传播，因此用于较小模型。对测试准确率很高的模型取 α=0.4、κ1=20、κ2=10。

路径合并的判定

• 作者假设：每个顶点的嵌入保存若干步内可达的顶点集合，或反向保存哪些顶点能到达当前顶点。每一层从可达集边缘拷贝并取并集，集合大小理论上可逐层加倍，从而可搜索的顶点数随层数指数增长。Figure 4 是示意图。搜索可以向前，也可以从目标向回。
• 注意力被分成路径合并，或从起点、目标可达顶点的拷贝；其余丢弃。路径合并可以按可达集重叠做 token 匹配，也可以按位置找一步之遥的顶点。若起点到目标路径上每个顶点都有不断的路径合并，把第一层对应 token 最终拷到最后一层的最后 token，则称该输入被路径合并算法解释。

允许中间 token 的两种任务

• DFS：输入图的边，以及访问到当前顶点为止的序列。右侧 token 留给访问序列，中间填充，使边在各样本中的位置稳定。预测 DFS 轨迹的下一顶点。训练图的 backtrack 距离均匀，对应前面的 balanced。实验固定 3 层。
• selection-inference：每步拆成两个子任务：从已访问且仍有未访问子顶点的顶点中选一个；再预测其一个未访问子顶点。用 frontier 大小 F 和当前顶点的分支数 B，使 (F, B) 均匀。输入是图和已访问的边，同样用填充稳住边的位置。实验固定 4 层。生成细节在附录 A.7 与 A.10。

实验设置

• 模型：从头训练的 GPT-2 架构 transformer，ReLU、1-hot 拼接位置编码、每层 1 头、无 causal mask。分布实验为 6 层、隐藏维 16。规模实验另固定 8 层、隐藏维 16，或改变非嵌入参数量。附录 A.5 为 decoder-only、可学习 token 嵌入与 RoPE。
• 数据：naïve、star、balanced 三种 DAG 分布。另有自然语言证明搜索：每条边写成条件句，词和标点是 token，任务改成从起始命题找证明的下一步。还有 DFS 与 selection-inference。Figure 2 的模型最大 41 顶点、输入 128 token、训练 883M 例。规模实验因显存把 batch 改为 256。
• 指标与判定：准确率的规则是预测顶点落在任一通往目标的路径上。训练目标是交叉熵。机制分析报告被路径合并解释的比例，以及 maximal 路径合并占比。规模实验把训练准确率大于 0.995 记为收敛。无 LLM 裁判。
• 测试与重复：测试样本来自预留 batch，并经精确匹配过滤。Figure 6 每个最大图规模 14 个种子；Figure 7 每个模型规模 4 个种子；Figure 9 上排每个图规模 5 个种子；Figure 14 每个点是 15 个种子的最小损失；Figure 11、12 是 2 个种子的最小损失。机制分析正文写共 2000 个输入、每个 lookahead 100 个；Figure 5 每格 100 例。

主结果

Figure 2：最大 41 顶点、883M 例、输入 128 token。L 列是 balanced 分布上该 lookahead 的留出测试，不是 naïve 图内部的 lookahead。

Swipe horizontally to read the table

训练分布	naive测试	balanced L=5	balanced L=14	balanced L=20
naive	0.99	0.26	0.51	0.49
star	0.65	1.00	0.99	1.00
balanced，L≤20	0.99	1.00	0.99	0.97
balanced，L≤12	0.99	1.00	0.96	0.49

• 作者称 naïve 训练不能稳健处理更大 lookahead，即使对训练中出现过的 lookahead（例如 L=5）准确率也低；作者把原因归于该分布生成大 lookahead 的概率指数下降。
• 作者称 star 训练在 balanced 测试上尚可，在 naïve 测试上较差；完整 balanced 训练在全部测试设置上都能搜索，并能泛化到未见过的 lookahead 数量。
• 只训到 L≤12 时，作者称能泛化到 lookahead 13 和 14，但不能泛化到更大的 lookahead。表外同图同列：L=13 为 0.99，L=15 为 0.84。

机制解释

• 测了什么：对 balanced 上训练、测试准确率很高的模型，在 naïve 与 balanced 留出样本上重建计算图，阈值 α=0.4、κ1=20、κ2=10。对照是随机初始化模型。
• 结果：作者称该方法特异性高，能在训练模型中识别路径合并，随机模型则不能。Figure 5 上半，随机模型可见值为 0.05 或 0.00；完整 balanced 训练行的可见值最低为 0.92（每格 100 例）。naïve 训练行有低至 0.43 的格。下半：作者称 maximal 路径合并比例明显小于 1。
• 作者的解读：模型没有理由学 maximal 合并，因为训练 lookahead 不对齐 2 的幂，且层数多于最大 lookahead 所需。作者用这一点解释 L≤12 模型不能完全泛化到更大 lookahead。附录 A.12 用 lookahead 为 9 的 100 例看方向：最大 lookahead 20 的模型前几层更像从目标向回搜；L≤12 的模型有的层像向前、有的层像向回。作者认为方向可能依赖初始化种子，并称还要更多实验才能描述完整。

规模

• 固定 8 层、隐藏维 16，balanced 上训练 236M 例，测试为 naïve 留出样本，14 个种子。作者称最大图规模增大时，达到训练准确率 ≥ 0.995 的种子比例变得非常小；已收敛模型上，14 个种子的最小测试损失以越来越快的速率上升（Figure 6）。正文未给出各图规模的具体损失值。
• 固定最大图规模 31，非嵌入参数为 0.9M、1.3M、6.0M、60.4M（Figure 7，每规模 4 个种子）。作者称更大模型更快到达局部最小值，但模型规模与找到全局最小值所需训练量之间没有可辨识的模式。
• 附录 A.5 改为可学习 token 嵌入与 RoPE 相乘的 decoder-only，d_model 与第 5 节对应实验相同。作者称缩放行为没有改变（Figure 11、12，每点为 2 个种子的最小损失）。

中间步骤

• DFS：固定 3 层并改变图规模时，作者称模型在所测图规模上都能学会训练分布，因而只需常数层；但图变大后仍然困难（Figure 14 上；测试为 backtrack 8 的 balanced 留出样本，15 个种子的最小损失）。固定最大 35 顶点后加大模型，作者称并不能更容易学会。更大模型用更少样本就能学会，但按 FLOPs 这一好处消失，更大模型需要多得多的 FLOPs。
• selection-inference：固定 4 层时，作者称图变大后仍然困难（Figure 16）。固定最大 45 顶点后加大规模也没有帮助（Figure 17；非嵌入参数含 3.2M、3.8M、13.4M、14.8M、40.3M）。附录 A.10 称 transformer 在 selection 子任务上做得好，在大图的 inference 子任务上做得差。论文未给出这两句对应的准确率。
• padding：Figure 15，7 层、输入 128、训练 746M 例。标准 padding 在 backtrack 0 为 1.00、backtrack 15 为 0.15；随机 padding 在 backtrack 0 至 4 为 1.00，backtrack 15 为 0.87。作者称随机 padding 足以让模型回溯更远，但不能帮助在更大图上学会搜索；规模实验的训练分布已均匀包含全部 backtrack 距离。

其他消融与分析

• 自然语言证明搜索（第 3.1.2 节，Figure 9）：训练行为与图搜索定性相似；图越大越难学，按 FLOPs 更明显。图例中的最大图规模包括 6、9、14、19、25、30；上排每个规模 5 个种子。正文未给该设置的准确率表。
• Figure 8：naïve、41 顶点、10M 图；作者称 lookahead 大于 6 虽可能，但几乎不会被采样到。
• Figure 10：与 Figure 6 同一设置的训练动态，14 个种子；Figure 6 是 236M 例处的切片。
• 若干损失曲线按图注做平滑：窗口 81，每 2^{18}=262K 例记录一点。论文未报告这些曲线上各点的具体损失数字。

作者指出的局限与后续方向

论文没有名为 Limitations 的小节。下列都出自结论或附录中作者明确写出的后续方向。

• 结论：放到大得多的模型规模，有可能导致搜索能力涌现。这是作者写出的可能性，不是本文已经观察到的结果。
• 结论：别的训练方法可能让 transformer 更容易学会搜索，例如课程学习；别的架构也可能有帮助，例如 looped transformers。
• 结论：路径合并能解释已训练模型的几乎所有样本，但某些样本上模型可能同时使用其他算法或启发式。
• 结论：欢迎进一步工作，以提高该分析对更大模型的可扩展性；小模型上的分析可以单独拿到大模型里检验。
• 附录 A.12：搜索方向可能依赖随机初始化种子，需要进一步实验才能得到更完整的描述。
• 结论同时判断：标准训练下，未来基于 transformer 的模型不会解决搜索任务，可能需要别的训练方法。

实验覆盖范围

• 实验对象是从头训练的小 transformer：主设置 6 层、隐藏维 16、1 头、ReLU、无 causal mask；规模设置固定 8 层、隐藏维 16，或固定 31 顶点、非嵌入参数取 0.9M、1.3M、6.0M、60.4M。另有 RoPE decoder-only（附录 A.5，Figure 11 每点 2 个种子）。
• 图分布为 naïve（入度上限 4）、star 与 balanced。任务还包括自然语言条件句证明、DFS（实验固定 3 层；变模型时最大 35 顶点）和 selection-inference（固定 4 层；变模型时最大 45 顶点）。
• 分布实验为输入 128 token、最大 41 顶点、883M 例、batch 1024、Sophia、学习率 10^{-5}。规模实验为 236M 例、batch 256、14 个种子；收敛线是训练准确率大于 0.995，测试分布是 naïve。
• 机制分析使用 α=0.4、κ1=20、κ2=10。正文写共 2000 个输入、每个 lookahead 100 个；Figure 5 每格 100 例，并以随机初始化模型为对照。
• 成功判定是规则：预测顶点在任一条到目标的路径上即正确。论文未报告预训练 LLM 上的实验，也未报告该规则与人工判定的一致性。

作者用 DAG 连通作测试台，检查小 transformer 学不会搜索是因为数据、参数，还是架构；理想化的无限数据被用来估计学会搜索的上限。

• 任务是给定有向无环图、起点和可达目标，预测路径上的下一顶点。难度由 lookahead 衡量。naïve 与 star 会留下捷径；balanced 让 lookahead 均匀，并去掉“起点是唯一源点”一类启发式。
• 模型是从头训练的 GPT-2 式 transformer，主实验 6 层、隐藏维 16、无 causal mask。作者用激活修补重建计算图，并假设学到的是在所有顶点上并行做指数路径合并：嵌入保存可达集，每层扩展该集合。分析找到了路径合并，但 maximal 合并比例小于 1；搜索方向在不同种子上可以向前，也可以向回。
• Figure 2（最大 41 顶点、883M 例）：naive 训练在 balanced 测试、lookahead=5 的准确率为 0.26；完整 balanced 训练在 naive 测试上为 0.99，lookahead=20 为 0.97。只训到 L≤12 时，lookahead=14 为 0.96，lookahead=20 为 0.49。作者称完整 balanced 训练后能在所测设置上搜索，L≤12 只能泛化到 13 和 14。
• 图规模增大后，8 层、隐藏维 16 的模型在 236M 例、14 个种子上越来越难达到训练准确率 0.995（Figure 6）。非嵌入参数增到 60.4M，作者称也看不出更容易找到全局最小值（Figure 7）。DFS 固定 3 层、selection-inference 固定 4 层时，更大图仍然困难，加规模不能消除。
• 作者的结论是：标准训练下，未来 transformer 不会只靠加大规模解决搜索，可能需要课程学习或 looped transformer 等其他做法。作者同时认为，放到大得多的规模仍可能导致搜索能力涌现；路径合并之外，模型在部分样本上可能还使用别的算法。