Transformers Struggle to Learn to Search
作者: Abulhair Saparov, Srushti Pawar, Shreyas Pimpalgaonkar, Nitish Joshi, Richard Yuanzhe Pang, Vishakh Padmakumar, Seyed Mehran Kazemi, Najoung Kim, He He
分类: cs.CL, cs.AI, cs.LG
发布日期: 2024-12-06 (更新: 2025-03-16)
备注: Published as a conference paper at ICLR 2025
💡 一句话要点
研究表明Transformer在学习搜索能力上存在困难，并提出一种新颖的可解释性分析方法。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: Transformer 搜索算法 图连通性 可解释性 机制分析 语言模型 图神经网络
📋 核心要点
- 大型语言模型在搜索任务中表现不佳，原因尚不明确，可能是数据不足或架构限制。
- 论文利用图连通性问题生成无限数据，训练小型Transformer，研究其学习搜索的能力。
- 实验发现，Transformer在特定训练分布下能学习搜索，但随着图增大，学习难度增加。
📝 摘要（中文）
搜索是许多重要任务的基础能力，但最近的研究表明，大型语言模型（LLM）在执行搜索时表现不佳。目前尚不清楚这种无力是由于缺乏数据、模型参数不足还是Transformer架构的根本限制。本文使用基础图连通性问题作为测试平台，生成有效无限的高覆盖率数据来训练小型Transformer，并测试它们是否可以学习执行搜索。研究发现，当给定正确的训练分布时，Transformer能够学习搜索。通过一种新颖的机制可解释性技术，从训练后的模型中提取计算图，分析了Transformer学习到的算法。结果表明，Transformer并行地在每个顶点执行搜索：对于输入图中的每个顶点，Transformer计算从该顶点可到达的顶点集合。然后，每一层逐步扩展这些集合，允许模型搜索的顶点数量以$n_{ ext{layers}}$指数增长。然而，研究发现，随着输入图大小的增加，Transformer在学习该任务时遇到更大的困难。即使增加参数数量也无法解决这种困难，这表明增加模型规模不会带来强大的搜索能力。研究还发现，上下文搜索（即，思维链）并不能解决在较大图上学习搜索的问题。
🔬 方法详解
问题定义：现有的大型语言模型在执行搜索任务时表现出明显的不足，尤其是在面对复杂或大规模的搜索空间时。这种不足可能是由于训练数据覆盖率不足、模型参数规模不够，或者是Transformer架构本身的局限性所导致。论文旨在探究Transformer架构是否具备学习搜索算法的潜力，以及在何种条件下能够有效地学习。
核心思路：论文的核心思路是将搜索问题简化为图连通性问题，并利用图连通性问题可以无限生成数据的特性，为Transformer提供充足的训练数据。通过观察Transformer在学习图连通性问题上的表现，来推断其学习搜索算法的能力。此外，论文还提出了一种新颖的机制可解释性方法，用于分析Transformer学习到的具体算法。
技术框架：论文的技术框架主要包括以下几个部分：1) 数据生成：利用图连通性问题生成大规模训练数据，保证数据覆盖率；2) 模型训练：使用小型Transformer模型进行训练，观察其在图连通性问题上的学习效果；3) 算法提取：利用提出的机制可解释性方法，从训练后的模型中提取计算图，分析模型学习到的具体算法；4) 性能评估：评估模型在不同规模图上的搜索性能，并分析模型规模和上下文学习对性能的影响。
关键创新：论文的关键创新在于：1) 将搜索问题简化为图连通性问题，并利用其数据生成优势；2) 提出了一种新颖的机制可解释性方法，能够从训练后的Transformer模型中提取计算图，从而理解模型学习到的具体算法；3) 实验结果表明，即使在充足的数据下，Transformer在学习大规模图上的搜索任务时仍然存在困难，这暗示了Transformer架构可能存在固有的局限性。
关键设计：论文的关键设计包括：1) 使用小型Transformer模型，以便于进行可解释性分析；2) 设计了特定的训练分布，以引导Transformer学习有效的搜索算法；3) 提出了基于计算图提取的机制可解释性方法，该方法能够有效地揭示Transformer内部的计算过程；4) 实验中，对比了不同模型规模和上下文学习策略对搜索性能的影响，从而更全面地评估了Transformer的学习能力。
🖼️ 关键图片
📊 实验亮点
实验结果表明，在合适的训练分布下，Transformer能够学习执行搜索任务。然而，随着输入图规模的增大，Transformer的学习难度显著增加，即使增加模型参数或采用上下文学习策略也无法有效解决。这暗示了Transformer架构在处理大规模搜索问题时可能存在固有的瓶颈。
🎯 应用场景
该研究成果有助于理解Transformer架构在解决复杂搜索问题上的局限性，并为未来设计更有效的搜索算法和模型架构提供指导。潜在应用领域包括知识图谱推理、路径规划、推荐系统等，有助于提升这些领域中算法的效率和准确性。
📄 摘要（原文）
Search is an ability foundational in many important tasks, and recent studies have shown that large language models (LLMs) struggle to perform search robustly. It is unknown whether this inability is due to a lack of data, insufficient model parameters, or fundamental limitations of the transformer architecture. In this work, we use the foundational graph connectivity problem as a testbed to generate effectively limitless high-coverage data to train small transformers and test whether they can learn to perform search. We find that, when given the right training distribution, the transformer is able to learn to search. We analyze the algorithm that the transformer has learned through a novel mechanistic interpretability technique that enables us to extract the computation graph from the trained model. We find that transformers perform search at every vertex in parallel: For each vertex in the input graph, transformers compute the set of vertices reachable from that vertex. Each layer then progressively expands these sets, allowing the model to search over a number of vertices exponential in $n_{\text{layers}}$. However, we find that as the input graph size increases, the transformer has greater difficulty in learning the task. This difficulty is not resolved even as the number of parameters is increased, suggesting that increasing model scale will not lead to robust search abilities. We also find that performing search in-context (i.e., chain-of-thought) does not resolve this inability to learn to search on larger graphs.