Transformers Provably Learn Graph Search: Training Dynamics and the Exponential Power of Depth
Xutao Ma ⋅ Somayeh Sojoudi
Abstract
Large Language Models (LLMs) have achieved remarkable progress on complex reasoning tasks, yet the mechanisms by which they acquire reasoning capabilities and the efficiency of their reasoning remain poorly understood. In this paper, we investigate the path-finding problem over directed graphs—a symbolic abstraction of multi-step reasoning. We first characterize the architectural requirements for this task, proving that one-layer transformers require $\Omega(N^2)$ size on dense graphs to implement the key DFS child-selection primitive. While a two-layer transformer implements full DFS with $O(N\log(N))$ size. Thus, a second layer is necessary and sufficient for near-linear DFS reasoning. We then extend this expressivity analysis to deeper models, proving the existence of an $L$-layer transformer that performs DFS with a look-ahead horizon of $2^{L-3}$, establishing that increasing depth yields an exponential gain in search efficiency. Finally, going beyond the expressivity, given appropriately curated training data, we show that two-layer transformers can provably learn to execute DFS via gradient flow and generalize to unseen graphs. Together, our results provide a mechanistic account of how transformers implement graph search, how depth improves reasoning efficiency, and how such reasoning algorithms can emerge through training.
Chat is not available.
Successful Page Load