🤖 machine learning

Transformers Linearly Represent Highly Structured World Models

By training a transformer on Sudoku solving traces, the study reveals that the model internally constructs a structured world model aligned with the domain's constraint algebra rather than surface features, utilizing sparse, interpretable circuits like a "naked-single" detector to solve the task.

Original authors: Roman Kniazev, Nathanaël Fijalkow

Published 2026-05-20

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Roman Kniazev, Nathanaël Fijalkow

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you hire a brilliant but mysterious apprentice to learn how to solve Sudoku puzzles. You don't teach them the rules directly; instead, you just show them thousands of examples of someone else solving puzzles step-by-step. After a while, the apprentice gets really good at solving them.

But here's the big question: How does the apprentice's brain actually work? Do they see the puzzle as 81 individual empty squares, or do they see it as a set of interconnected rules?

This paper investigates exactly that by "opening up" a computer model (a Transformer) trained on Sudoku. The researchers found three surprising things about how this model thinks.

1. The "Team Captain" vs. The "Individual Worker"

Most people assume that if you show a model a Sudoku grid, it will learn to track every single cell (the little boxes) individually, like a worker checking off 81 separate items on a list.

The Paper's Finding: The model doesn't do that. Instead, it organizes its thinking around groups.
Think of a Sudoku puzzle not as 81 boxes, but as three types of teams:

The Row Team (9 horizontal lines)
The Column Team (9 vertical lines)
The Box Team (9 small 3x3 squares)

The model builds a "world model" where it tracks what numbers are present in these teams, not just in individual boxes. It's as if the apprentice stopped looking at individual bricks and started looking at the walls, floors, and ceilings. The researchers found that the model is perfect at answering, "Is the number 5 in the top row?" but only about 80% accurate at guessing, "What number is in this specific box?"

Why? Because in Sudoku, a number's validity depends on the group it belongs to, not the box itself. The model learned to think in terms of the rules (the algebra), not the surface appearance (the grid).

2. The "Traffic Controllers" in the Middle

The model has 8 layers of "thinking" steps. The researchers looked at the middle layers to see how information flows.

The Paper's Finding: They discovered specific "traffic controllers" (called attention heads) that act like specialized managers for each team.

One manager only looks at Row 4.
Another only looks at Box 5.
Another only looks at Column 2.

These managers have a very simple job: The "No-Go" Signal.
If a number (say, a 7) is already in Row 4, this manager immediately puts a "Do Not Use" sign on the number 7 for every empty spot in that row. They don't just ignore it; they actively suppress the idea of using it. This happens simultaneously for all rows, columns, and boxes, creating a clean list of what can be used.

3. The "Lightbulb Moment" Neurons

Finally, the researchers looked at the very last step before the model makes a move. This is where the model decides, "Okay, I'm going to put a 7 here."

The Paper's Finding: They found a tiny, specific group of neurons (only 91 of them) that act like Lightbulbs.

These lightbulbs only turn on when a specific cell has exactly one possible number left (a situation called a "Naked Single").
When the lightbulb turns on, it doesn't just say "7 is good." It shouts, "Hey, this whole cell is solved!" and boosts the confidence for all numbers in that cell equally.
Because the middle layers had already done the hard work of eliminating the wrong numbers, the "correct" number was already the top choice. The lightbulb just gives it a final, massive push to ensure the model commits to that answer.

The Big Picture

The paper concludes that this AI didn't just memorize patterns; it built an internal map that perfectly matches the logic of the puzzle.

It doesn't see a grid; it sees constraints.
It doesn't guess randomly; it uses specialized managers to enforce rules.
It doesn't hesitate; it has dedicated switches to lock in the answer when the logic is clear.

The researchers proved this by "breaking" the model: when they turned off the "Row Managers," the model started suggesting numbers that were already in that row. When they turned off the "Lightbulb Neurons," the model got confused about which number to pick, even though it knew the answer was there.

In short, the AI learned to solve Sudoku the way a human expert does: by understanding the rules of the game, not just by looking at the empty squares.

Technical Summary: Transformers Linearly Represent Highly Structured World Models

Problem Statement

The paper investigates whether transformers trained on sequential reasoning traces develop internal "world models" that reflect the underlying algebraic structure of the task domain, rather than merely its surface presentation. While prior work (e.g., Li et al. [2023], Nanda et al. [2023]) demonstrated that models trained on Othello spontaneously learn linear representations of board occupancy (one feature per cell), it remains unclear if this holds for tasks with different structural constraints. Specifically, the authors ask: in a combinatorial constraint-satisfaction task like Sudoku, where validity is determined by overlapping substructures (rows, columns, and 3×3 boxes) rather than isolated cells, does the model represent the state cell-by-cell or substructure-by-substructure?

Methodology

The authors trained an 8-layer causal decoder-only transformer (8 attention heads, $d_{model}=576$ ) on 2.7 million Sudoku solving traces generated by a Norvig-style solver with backtracking. The model achieved 98.4% per-cell and 97.5% per-grid accuracy on a held-out test set.

To analyze the internal computation, the authors employed a "probe-and-patch" methodology:

Linear Probing: Logistic regression classifiers were trained on residual stream activations at the [clues_end] token (marking the end of initial clues) to decode specific concepts. They tested three families of probes:
- Cell-state: Predicting the digit in a specific cell (81 probes).
- Cell-candidates: Predicting valid candidate digits for a specific cell (729 probes).
- Substructure-state: Predicting if a digit is present in a specific substructure (row, column, or box) (243 probes).
Causal Intervention (Activation Patching): To verify if decoded representations are causally used, the authors subtracted specific probe directions from the residual stream and measured the resulting drop in target token logits.
Attention Analysis: The authors decomposed attention heads into Query-Key (QK) and Output-Value (OV) circuits to determine what information heads attend to and how they modify the residual stream.
Logit Lens & Direct Logit Attribution (DLA): These tools were used to track how predictions evolve across layers and to attribute logit changes to specific components (heads or neurons).
Neuron Ablation: Specific neurons in the final MLP layer were identified and ablated to test their causal role in decision-making.

Key Contributions and Results

1. Emergence of a Substructure World Model

Contrary to the "cell-by-cell" representation found in Othello models, the Sudoku transformer builds a world model organized around the puzzle's constraint algebra.

Linear Separability: Linear probes for "is digit $d$ present in substructure $S$ ?" achieve perfect exact-match accuracy (1.0) in mid-layers (Layers 4–7) across all 243 (substructure, digit) pairs.
Cell-Level Limitations: In contrast, probes predicting the state of individual cells plateau at ~80% accuracy. Probes for cell candidates reach 100% accuracy but exhibit structural alignment: candidate vectors for cells sharing a substructure are highly correlated, suggesting they are linear combinations of substructure-level features.
Stability: This substructure representation emerges at Layer 4 and remains geometrically stable throughout the solving trace, as confirmed by cross-layer and cross-position transfer experiments.
Causal Usage: Activation patching confirms these directions are causally necessary; overwriting them to indicate a digit is present in a substructure causes the model's prediction for that placement to collapse.

2. Substructure-Aligned Attention Mechanisms

Mid-layer attention heads (Layers 4–6) are specialized to operate on specific substructures.

Routing: Attention heads concentrate their focus on specific rows, columns, boxes, bands, or stacks.
Constraint Enforcement: The OV circuits of these heads implement a symmetric constraint-checking mechanism. If a digit is present in the attended substructure, the head suppresses the logits for that digit in the relevant cells; if absent, it promotes the digit.
Localization: Causal ablation of specific heads (e.g., a head attending to columns 4–6) increases the logits of illegal placements only within that specific substructure, leaving other substructures unaffected.

3. The Naked-Single Circuit

The authors identified a sparse, monosemantic circuit in the final MLP layer responsible for committing to "naked single" placements (cells with exactly one valid candidate).

Pre-MLP State: Before the final MLP, the residual stream already ranks the correct digit first in 98.77% of unique naked-single cases.
Dedicated Neurons: The final MLP contains 91 dedicated neurons that fire exclusively when their associated cell has exactly one candidate. These neurons are monosemantic and do not activate for other candidate counts.
Amplification Mechanism: When a naked-single neuron fires, it applies a uniform, massive boost (+3.00 logits) to all placement tokens for that specific cell. Since the correct digit was already top-ranked by mid-layers, this uniform amplification preserves the lead and concentrates the softmax probability on the correct digit.
Causal Confirmation: Ablating these neurons drops the probability of the target naked single by ~0.59 without affecting other simultaneous naked singles on the same board, confirming their role as the final commitment gate.

Significance and Claims

The paper claims to demonstrate that the geometry of emergent world models is shaped by the constraint algebra of the domain rather than its surface decomposition. In Sudoku, where validity is defined by the intersection of rows, columns, and boxes, the model learns to represent the state in terms of these substructures.

Furthermore, the study provides a fully mechanistic, end-to-end algorithmic account of how a transformer solves a combinatorial reasoning task:

Representation: Mid-layers construct a linear, substructure-aligned world model.
Constraint Propagation: Specialized attention heads enforce constraints by suppressing invalid candidates within those substructures.
Decision: A sparse set of monosemantic neurons in the final layer detects specific logical states (naked singles) and amplifies the pre-computed correct answer.

The authors conclude that mechanistic interpretability tools can successfully reverse-engineer these complex reasoning processes, revealing that the resulting decision circuits are sparse, monosemantic, and fully interpretable.

Limitations

The authors note that these findings are specific to a single model architecture (8 layers, 8 heads) trained on a single puzzle type (Sudoku). They do not claim that substructure world models or naked-single circuits will necessarily appear in models of different sizes or those trained on non-backtracking traces. Additionally, the analysis of how attention patterns are selected (QK circuits) is not fully detailed, and the causal evidence relies on a specific patching protocol.