Structured World Representations in Maze-Solving Transformers

TL;DR

This paper investigates small transformer models solving mazes, revealing structured internal representations of maze topology, path information, and attention mechanisms, which enhances understanding of how transformers process spatial tasks.

Contribution

It demonstrates that maze topology and paths are internally represented in a structured manner within small transformers, including linearly decodable maze reconstructions and specialized attention heads.

Findings

Residual stream encodes entire maze topology

Token embeddings exhibit spatial structure

Attention heads are involved in path-following

Abstract

Transformer models underpin many recent advances in practical machine learning applications, yet understanding their internal behavior continues to elude researchers. Given the size and complexity of these models, forming a comprehensive picture of their inner workings remains a significant challenge. To this end, we set out to understand small transformer models in a more tractable setting: that of solving mazes. In this work, we focus on the abstractions formed by these models and find evidence for the consistent emergence of structured internal representations of maze topology and valid paths. We demonstrate this by showing that the residual stream of only a single token can be linearly decoded to faithfully reconstruct the entire maze. We also find that the learned embeddings of individual tokens have spatial structure. Furthermore, we take steps towards deciphering the circuity of path-following by identifying attention heads (dubbed $\textit{adjacency heads}$), which are implicated in finding valid subsequent tokens.