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 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 the number of 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.

翻译：搜索是许多重要任务的基础能力，近期研究表明大型语言模型（LLM）难以稳健地执行搜索。尚不清楚这种能力缺失是由于数据不足、模型参数不够，还是Transformer架构的根本局限。本研究以基础图连通性问题作为测试平台，生成本质上无限的高覆盖数据来训练小型Transformer，并检验其能否学会执行搜索。我们发现，当给定合适的训练分布时，Transformer能够学会搜索。通过一种新颖的机制可解释性技术——该技术使我们能够从训练好的模型中提取计算图——我们分析了Transformer学习到的算法。研究发现，对于输入图中的每个顶点，Transformer会计算从该顶点可达的顶点集合。随后每一层逐步扩展这些集合，使得模型能够在层数指数级增长的顶点范围内进行搜索。然而，随着输入图规模的增大，Transformer学习该任务的难度显著增加。即使增加参数数量也无法解决这一困难，这表明扩大模型规模不会带来稳健的搜索能力。我们还发现，在上下文情境中执行搜索（即思维链）同样无法解决模型在更大规模图上学习搜索的缺陷。