小规模Transformer可学会并行搜索,但面对大图时仍难奏效。
Transformers Struggle to Learn to Search
- 用图连通性问题生成无限数据,训练小模型学习搜索。
- 模型在每顶点并行计算可达集,层数越多搜索范围指数增长。
- 增大模型规模或使用思维链也无法解决大图搜索难题。
搜索是众多重要任务的基础能力,近期研究发现大型语言模型(LLMs)难以稳健执行搜索。这种能力缺失是否源于数据不足、参数量不够,或变压器架构的根本局限尚不明确。本文以基础的图连通性问题为测试平台,生成覆盖广泛的无限数据,训练小型变压器并检验其学习搜索的能力。结果表明,在合适的训练分布下,变压器能够学会搜索。我们通过一种新型机制可解释性技术,从训练好的模型中提取计算图,分析其内部算法。发现变压器在每个顶点上并行执行搜索:对输入图中的每个顶点,计算从该顶点可达的所有顶点集合;每一层逐步扩展这些集合,使模型能以层数 $n_{ ext{layers}}$ 的指数级速度搜索顶点。然而,随着输入图规模增大,模型学习难度显著上升,即使增加参数量也未能缓解,表明扩大模型规模无法带来稳健的搜索能力。此外,在上下文内进行搜索(即思维链)也无法解决在大规模图上的学习困境。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。