Transformers Use Causal World Models in Maze-Solving Tasks

arXiv 24.12 2024 37.3 method

TLDR

Transformers trained on maze-solving tasks develop causal world models, revealed by SAEs and attention analysis, with feature intervention and generalization tests.

Reasoning

Strengths: The paper provides a clear methodology using Sparse Autoencoders and attention analysis to identify and intervene on world model features, demonstrating causal roles and generalization properties. Weaknesses: The study is limited to synthetic maze tasks, lacks real-world validation, and some findings (e.g., easier to activate than suppress) may be task-specific.

Read-first score

Read-first score 37.3, weighted from topical fit, citation, graph, method, reproducibility, and recency signals. Original total remains 11.

Recency 8%

75.1

Uses a gentle age decay so recent papers surface without erasing older foundations. 2024

Methodology quality 25%

Screens visible abstract and analysis fields for experiment, dataset, baseline, metric, and limitation evidence. markers=analysis

Reproducibility 25%

Screens links and visible text for paper, code, dataset, artifact, and repository signals. pdf=True; code=False; dataset=False; markers=code

Topical relevance 42%

15.7

Uses existing LLM keyword relevance scores normalized to 0-100. world model,world simulator,generative world model,interactive world model,video world model,world dynamics prediction,model-based reinforcement learning world model

Field roles

Candidate

Rank sensitivity

Stability: volatile; rank range: 98.

Keyword Scores

world model

world dynamics prediction

world simulator

generative world model

interactive world model

video world model

model-based reinforcement learning world model

Deep Analysis

Innovations

Identification of causal world models in transformers trained on maze-solving tasks using Sparse Autoencoders (SAEs) and attention pattern analysis.
Demonstration of consistency between SAE feature-based and circuit-based analyses for world model construction.
Discovery that it is easier to causally activate features than to suppress them via intervention.
Finding that models can generalize to mazes with more simultaneously active features than seen during training, but fail when those mazes are provided via input tokens.
Evidence that positional encoding schemes influence how world models are structured in the residual stream.

Methodology

The study trains transformer models on maze-solving tasks and uses Sparse Autoencoders (SAEs) to extract interpretable features, combined with attention pattern analysis to examine world model construction. Causal interventions are performed on isolated features to confirm their causal roles, and the models are evaluated on mazes with varying numbers of connections, both as internal features and as input tokens.

Key Results

Models can reason about mazes involving more simultaneously active features than encountered during training, but fail when the same mazes (with greater numbers of connections) are provided via input tokens. Additionally, activating features is easier than suppressing them, and positional encoding schemes affect how world models are structured in the residual stream.

Limitations

The study is limited to maze-solving tasks, so generalizability to other domains is not established.
Models fail to handle mazes with many connections when those mazes are presented as input tokens, indicating a limitation in processing input complexity.
Causal interventions reveal asymmetry (easier to activate than suppress), which may constrain the interpretability and control of model behavior.
The analysis relies on specific interpretability tools (SAEs, attention patterns) that may not capture all aspects of world model formation.

Resources

Paper