Transformers Provably Learn Algorithmic Solutions for Graph Connectivity, But Only with the Right Data

TL;DR

This paper demonstrates that Transformers can learn algorithmic solutions for graph connectivity when trained on data within their capacity, combining theoretical proofs and empirical evidence.

Contribution

It provides a theoretical analysis of how Transformers learn graph algorithms and shows that data within the model's capacity is crucial for learning exact solutions.

Findings

Transformers can compute graph connectivity using matrix powers.

Training data within model capacity leads to exact algorithm learning.

Beyond-capacity data results in heuristic-based solutions.

Abstract

Transformers often fail to learn generalizable algorithms, instead relying on brittle heuristics. Using graph connectivity as a testbed, we explain this phenomenon both theoretically and empirically. We consider a simplified Transformer architecture, the Disentangled Transformer, and prove that an $L$-layer model can compute connectivity in graphs with diameters up to $3^L$, implementing an algorithm equivalent to computing powers of the adjacency matrix. By analyzing training dynamics, we prove that whether the model learns this strategy hinges on whether most training instances are within this model capacity. Within-capacity graphs (diameter $\leq 3^L$) drive the learning of the algorithmic solution while beyond-capacity graphs drive the learning of a simple heuristic based on node degrees. Finally, we empirically show that restricting training data to stay within a model's capacity makes both standard and Disentangled Transformers learn the exact algorithm.