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

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 our insights transfer to standard Transformers: restricting training data to stay within a model's capacity makes both standard and Disentangled Transformers learn the exact algorithm.

Today's AI systems often lean on brittle shortcuts rather than truly learning to solve a problem, so they work on familiar examples but fail on anything new. Using a simple task of figuring out which points in a network are connected, we found that a model can only handle problems up to a certain size, and it learns the real, reliable method only when its training examples stay within that limit. Otherwise it falls back on guessing. By matching the training examples to what the model can actually handle, we got it to learn the correct general method and succeed even on situations it had never seen before.