Spectral Journey: How Transformers Predict the Shortest Path
This work addresses the interpretability of transformer models for AI researchers, though it is incremental as it focuses on a controlled, simplified setting.
The researchers tackled the problem of understanding whether decoder-only transformers can plan or reason by training them to predict shortest paths on simple graphs, finding that two-layer models can handle graphs up to 10 nodes and learn embeddings correlated with spectral decomposition.
Decoder-only transformers lead to a step-change in capability of large language models. However, opinions are mixed as to whether they are really planning or reasoning. A path to making progress in this direction is to study the model's behavior in a setting with carefully controlled data. Then interpret the learned representations and reverse-engineer the computation performed internally. We study decoder-only transformer language models trained from scratch to predict shortest paths on simple, connected and undirected graphs. In this setting, the representations and the dynamics learned by the model are interpretable. We present three major results: (1) Two-layer decoder-only language models can learn to predict shortest paths on simple, connected graphs containing up to 10 nodes. (2) Models learn a graph embedding that is correlated with the spectral decomposition of the line graph. (3) Following the insights, we discover a novel approximate path-finding algorithm Spectral Line Navigator (SLN) that finds shortest path by greedily selecting nodes in the space of spectral embedding of the line graph.