Research map / Probabilistic and scientific modeling / Causal discovery

A Causal World Model Underlying Next Token Prediction: Exploring GPT in a Controlled Environment

ICML 2025 · Probabilistic and scientific modeling › Causal discovery · approach: Causal representation identifiability · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Derives a causal (SCM) interpretation of GPT masked attention, uses it for zero-shot causal structure learning per sequence with a p-value entropy confidence score, and tests on out-of-distribution random legal Othello and Chess move sequences.

Problem

Whether GPT models trained only on next-token prediction implicitly learn a causal world model, and whether attention reflects causal structure that correlates with generating legal moves on out-of-distribution sequences.

Approach

Interprets the lower-triangular masked attention matrix (after diagonal normalization) as the transitive closure (I-G)^-1 of a linear-Gaussian SCM, enabling recursive constraint-based causal discovery over each sequence and a confidence score R = H_ind - H_dep from CI-test p-values.

Key results

• Othello GPT averages about 95% legal next-move accuracy on random legal sequences, with lower accuracy for lengths 10-30 (e.g., 88% at length 15).
• Structural confidence R is significantly higher for legal than illegal generations for sequence lengths in [15,22] when using all CI tests (case d), and positive at all tested lengths.
• Pruning heads with low structural confidence preserves legal-move accuracy better than pruning high-confidence heads (larger area under curve).

Measured effects reported in the paper

Surprising findings

• Raw attention values alone (pairwise, empty conditioning set) do not give a statistically significant legal vs illegal difference except at one length, while the full CI-based score does.

Limitations

• Causal interpretation assumes linear-Gaussian SCM and CI tests via partial correlation with Fisher z, with Markov and faithfulness assumptions.
• Supports only rung-1 (associational) inference per the authors.
• Evaluated only on two board games with two pretrained models; test distribution is synthetic random play.

Experimental setup

Datasets: Othello game sequences (~132K) Chess games (~2.9M) 1,000 random legal-move synthetic test sequences per game (Chess up to 10,000 for short lengths)

Benchmarks: Othello legal-move generation on random legal-move sequences Chess legal-move generation on random legal-move sequences

Metrics: legal move accuracy structural confidence R difference (legal vs illegal) area under pruning accuracy curve

Baselines: raw attention pairwise correlations BERT-based ABCD (unmasked attention C=AA^T)

Builds on: ABCD iterative causal discovery (ICD) BERT attention causal interpretation

Closest papers by meaning

• Selective Induction Heads: How Transformers Select Causal Structures in Context (ICLR 2025 · similarity 0.88)
• Towards Understanding the Universality of Transformers for Next-Token Prediction (ICLR 2025 · similarity 0.88)
• How Transformers Learn Causal Structure with Gradient Descent (ICML 2024 · similarity 0.86)
• Chess-World-Model: A 10M-Game Benchmark for Exact State Tracking from Chess Move Sequences (NeurIPS 2026 · similarity 0.86)
• SFT Memorizes, RL Generalizes: A Comparative Study of Foundation Model Post-training (ICML 2025 · similarity 0.85)
• The Pitfalls of Next-Token Prediction (ICML 2024 · similarity 0.85)