Structured Representations in Maze-Solving Transformers
Michael I. Ivanitskiy1, Alex F. Spies23, Tilman Räuker
Guillaume Corlouer, Chris Mathwin, Lucia Quirke, Can Rager, Rusheb Shah, Dan Valentine, Cecilia Diniz Behn1, Katsumi Inoue3, Samy Wu Fung1
1

Colorado School of Mines 2Imperial College London 3 National Institute of Informatics

Why Study Maze-solving Models?

Models Learn Structured Representations

When do World Models Form?

● Small models are easier to analyze and can help us todevelop techniques for investigating large models● Spatial structure: visually intuitive to interpret● Internal search: solving mazes is a good test to see ifmodels can plan, which has implications for AI safety

Linear “World Models” are Learned● Training to follow paths (hallway) rather than solvemazes (jirpy) doesn’t lead to World Models (WMs)● Formation of World Models coincides with ourmodels learning to solve the mazes

In our jirpy model we can faithfully decode the full mazefrom a single token’s* residual stream state at Layer 2!

Jirpy Model

Hallway Model
Optional caption for images, charts, and graphs

Our Setup● Train transformers to solve
mazes, or follow paths● Inputs are tokenized mazed,including the shortest path
● Setup is the same as a normal
next-token prediction task

Average accuracy of the 4-wallprobes. *Using <PATH_START>

Some examples of reconstructed mazes (not cherry-picked).
Here orange → missed and pink→ hallucinatedresiduals from jirpy at layer 2

World Model accuracy (green) against model accuracy throughout training,

Adjacency Heads Encode Valid Neighbors

Example of a maze. The model
must output the shortest pathfrom the start to the goalTakeaways

Looking at per-head logit contributions we identify heads inthe hallway model which capture valid neighbor structure
Attention Weight

A truncated example of the input provided to a model at
train time (the path is omitted during inference)

The Models

Logit Contributions of the “Adjacency Head” on various mazes. Here the dashed line from the redcircle to the red cross indicates the true path, and the current position is the green cube.

We focus our analyses on two performant toy models:● “jirpy” - trained to solve mazes, high data variety○ 9.6M params, 12 layers, 16 heads● “hallway”- trained to follow “forkless” paths○ 1.2M params, 6 layers, 4 heads

Models Capture Neighbor Structure Early● Linear probing works best at early layers● TunedLens provides further evidence that by layer 2models already encode the maze’s topology
hallway

Next Steps● Do the structures we ﬁnd play a causal role? If so, what
does the circuit they are part of look like?
● Can we intervene on the model in order to make itbetter at satisfying constraints we like?
● Can we retarget the model to act sensibly with respect
to different goals?
Author Notes

jirpy
Trained models were consistently able to solve the tasks.
The best models had >95% accuracy on the validation set● Linear WMs form at early layers● Specialized Attention heads respect maze topology● WMs don’t form in transformers which only need tofollow paths (rather than solve mazes)● During training, WMs become more accurate whenTransformers get better at solving mazes.

Accuracy of the linear world model probes
peaks at layer 2. This held across similar models.

We use TunedLens to decode the residualstream, ﬁnding that models already encode
neighbor and valid-neighbor info in early layers➔ The code and models required to reproduce these experiments is available on
our github (linked above) as well as the libraries we created for Maze Datasets
and Maze Transformers➔ Not all transformers form linear world models (but the best ones do)➔ We have since trained models which leverage relative position encodings andgeneralize much more strongly to longer and shorter mazes than those seen
during training.