Transformers Struggle to Learn to Search

搜索是许多重要任务的基础能力，最近的研究表明，大型语言模型（LLMs）在执行搜索时表现得不够稳健。目前尚不清楚这种无能是由于数据不足、模型参数不足，还是变压器架构的根本限制所致。在这项工作中，我们使用基础图连通性问题作为测试平台，生成实际上无限量的高覆盖率数据来训练小型变压器，并测试它们是否能够学会进行搜索。我们发现，当提供正确的训练分布时，变压器能够学会搜索。 我们通过一种新颖的机制可解释性技术分析了变压器所学的算法，该技术使我们能够从训练好的模型中提取计算图。我们发现，对于输入图中的每个顶点，变压器都会计算从该顶点可达的顶点集。然后每一层逐步扩展这些集合，使模型能够在层数的指数级顶点数量上进行搜索。 然而，我们发现随着输入图大小的增加，变压器在学习任务时遇到了更大的困难。即使增加参数数量，这一困难也无法解决，这表明增加模型规模不会导致稳健的搜索能力。我们还发现，在上下文中执行搜索（即，链式思维）也无法解决在更大图上学习搜索的能力不足问题。

该论文旨在探讨大型语言模型（LLMs）在执行搜索任务时的表现不佳的原因，并通过图连通性问题作为测试平台来训练小型Transformer模型，以验证其学习搜索的能力。这是一个尝试理解并解决现有模型局限性的问题。

论文的关键思路是利用图连通性问题生成无限量的高覆盖率数据来训练小型Transformer模型，从而探索模型是否能够学习到搜索算法。与现有研究不同，该论文不仅关注模型参数和数据量的影响，还通过机制可解释性技术分析了模型内部的计算过程。

论文通过机制可解释性技术揭示了Transformer如何逐步扩展每个节点的可达集合，从而实现搜索。然而，当输入图的规模增大时，模型的学习难度增加，即使增加模型参数也无法解决这一问题。此外，论文指出在上下文中进行搜索（即链式思维）同样无法提高模型在较大图上的表现。这些发现为未来的研究提供了重要的方向。

近期在这个领域的一些相关研究包括： 1. 'Probing Neural Network Comprehension of Natural Language' - 探讨神经网络对自然语言的理解能力。 2. 'Transformers as Soft Decision Trees: An Empirical Study' - 研究Transformer模型作为软决策树的行为。 3. 'On the Limitations of Representational Similarity for Understanding Neural Networks' - 讨论表征相似性在理解神经网络中的局限性。 4. 'Emergent Properties in Large-Scale Language Models' - 探索大规模语言模型中的涌现属性。