Understanding Transformer Reasoning Capabilities via Graph Algorithms

TL;DR

This paper explores the algorithmic reasoning capabilities of transformers, identifying the scaling regimes needed for different problems and demonstrating their effectiveness on graph reasoning tasks through theory and experiments.

Contribution

It introduces a representational hierarchy classifying problems by their solvability with transformers under various parameter regimes and provides both theoretical proofs and empirical validation.

Findings

Logarithmic depth is necessary and sufficient for graph connectivity.

Single-layer transformers can solve contextual retrieval tasks.

Transformers outperform some specialized graph neural networks on graph reasoning.

Abstract

Which transformer scaling regimes are able to perfectly solve different classes of algorithmic problems? While tremendous empirical advances have been attained by transformer-based neural networks, a theoretical understanding of their algorithmic reasoning capabilities in realistic parameter regimes is lacking. We investigate this question in terms of the network's depth, width, and number of extra tokens for algorithm execution. Our novel representational hierarchy separates 9 algorithmic reasoning problems into classes solvable by transformers in different realistic parameter scaling regimes. We prove that logarithmic depth is necessary and sufficient for tasks like graph connectivity, while single-layer transformers with small embedding dimensions can solve contextual retrieval tasks. We also support our theoretical analysis with ample empirical evidence using the GraphQA benchmark. These results show that transformers excel at many graph reasoning tasks, even outperforming specialized graph neural networks.