Verification of the Implicit World Model in a Generative Model via Adversarial Sequences

AI summary

Overview

Research area: Machine learning interpretability and formal-language verification of generative sequence models, studied in the domain of chess.

Technical level: Advanced. The paper assumes familiarity with transformer language models, formal-language and world-model definitions, linear probing, and adversarial evaluation methodology.

Scope: The paper introduces adversarial sequence generation as a practical tool for falsifying the soundness of implicit world models, and applies it to 24 chess-playing GPT-2 models trained with different datasets and training objectives.

What This Paper Is About

Generative sequence models are trained on example sequences, and it is unclear whether the implicit "world model" they learn matches the true rule system of the domain. Prior theory suggests that at best we can hope for soundness — generating only valid sequences, though not necessarily all of them — so the practical question is whether we can build tools that detect when a model breaks the rules. The authors use chess as a domain with enough complexity but a simple, fully known rule-based world model, and train adversaries that produce only legal moves while trying to force the sequence model to output an illegal one.

Key Contributions

1. A novel adversarial framework for measuring soundness. The authors formalize soundness of a generative model relative to a ground-truth world model, and propose adversaries that always extend sequences with valid actions chosen to maximize an auxiliary objective while forcing the model into predicting an invalid next action. The framework does not require an ad hoc probability threshold to define the generated language, unlike the sequence-level distinctions approach of Vafa et al. (2024).

2. A large-scale empirical study across 24 chess models. The authors train models on three curated datasets (Millionbase-500k, Stockfish-8M, Lichess-16M) and three random datasets (500K, 2M, 10M valid random games), each with four training objectives: next-token prediction (NT), probability distribution prediction (PD), next-token prediction with a joint board state probe (NT+JP), and probability distribution prediction with a joint probe (PD+JP).

3. Five adversary and baseline strategies with fine-grained failure analysis. Illegal Move Oracle (IMO), Board State Oracle (BSO), Adversarial Detours (AD) from Vafa et al. (2024), Random Move (RM), and Sequence Model Move (SMM), evaluated with per-move attack dynamics as well as final success rates.

4. Evidence on the causal role of board state probes. Through gradient alignment analysis, BSO attack statistics, and joint-probe training, the authors argue that extracted board states have little causal connection to next-token prediction in most models.

Main Findings

• No model is sound. For most models at least one adversary achieves close to a 100% success rate, indicating severe inconsistency between the implicit and the true world model. The paper reports that models trained on the smaller datasets (Random-500k and Millionbase-500k) reach legal move ratios between 94.65% and 96.71% on their test sets, while models trained on the large datasets (Random-10M, Stockfish-8M, Lichess-16M) reach between 99.75% and 99.98%.

• IMO is the strongest adversary, usually by a wide margin. For example, on Random-500k the four models score 0.996, 0.999, 0.996, and 1.000, and on Random-2M they score 0.999, 1.000, 0.997, and 0.999. On Lichess-16M with the NT objective, IMO still achieves 0.387, far above every other adversary for that model.

• AD and BSO are surprisingly weak. AD consistently achieves success rates similar to Random Move, which the authors attribute to most low-probability moves being effectively random from the model's perspective. BSO is sometimes weaker than even the benign SMM baseline, implying a weak causal link between probed board state correctness and move legality.

• Dataset size reliably improves robustness. Increasing dataset size increases resistance to attacks independently of dataset type and training objective.

• Dataset type and objective interact. Curated datasets produce very robust models under the NT objective, especially at large scale, while models trained on random datasets are less robust under NT compared to PD. However, the apparent soundness on curated data is shown to be misleading.

• Joint probes do not help. Adding a joint board state probe to training has a negligible effect on soundness.

• Models overfit sequence length. Many models, especially those trained on large datasets with the PD objective, become extremely unreliable once generation exceeds the dataset's sequence lengths (up to 150), suggesting they do not use sequence-length-independent abstract board state representations.

• Board state probes are functionally independent of generation. The mean sample-wise cosine distance between the gradients of the next-token head and the board state probe ranges from 0.966 to 1.000 across models, meaning the two tasks rely on nearly orthogonal subspaces. Under the BSO attack, the ratio of illegal moves that relate to the probed board state is very low for large datasets, so successful BSO attacks are not caused by misleading the board state predictor.

• Apparent soundness on Lichess-16M is an artifact of in-distribution trajectories. With out-of-distribution warmup sequences the attack success rates rise sharply. For the Lichess-16M NT model, IMO goes from 0.387 to 0.667 with 20 OOD moves (+0.28), 0.785 with 40 (+0.40), and 0.864 with 60 (+0.48); for NT+JP it goes from 0.349 to 0.659, 0.735, and 0.797 respectively.

• Results are robust to decoding policy. With top-k sampling at k=4, all attacks achieve a higher ASR than under greedy decoding while keeping similar relative performance; for example Lichess-16M NT reaches 0.654 under IMO with top-k versus 0.387 with greedy decoding. The authors read this as evidence that IMO uncovers vulnerable state regions rather than one-off errors.

Methodology in Plain English

The authors start from a formal setup. A set of all possible actions defines a language, and a ground-truth world model W tells you which actions are legal continuations of any given action sequence. A sequence is valid if every action in it belongs to the legal continuations of its prefix. A generative model M outputs a probability distribution over actions given a prefix, and a decoding policy (for example greedy argmax, or top-k sampling) turns that into an actual move. A model is sound if, for every valid sequence with at least one legal continuation, its decoded move is legal.

To test this, they build an adversary that plays against the model. The adversary always plays a legal move, chosen from the legal continuations by maximizing an auxiliary scoring function f that tries to make the model fail later. Several variants of f are tested: the Illegal Move Oracle picks the legal move that maximizes the model's probability of an illegal reply; the Board State Oracle picks the legal move that maximizes the error of an attached board state probe; Adversarial Detours picks the lowest-probability legal move, steering toward out-of-distribution regions; Random Move picks uniformly; and Sequence Model Move picks the model's own most likely legal move, acting as a benevolent oracle. Since chess is a two-player game, the attacker plays white and the model plays black, and the attack succeeds if the model ever produces an illegal move.

To populate the study, they train 24 GPT-2 models with 12 hidden layers, 768 hidden dimensions, 12 attention heads, and 86M parameters, for 3 epochs on six datasets using four objectives. The probability distribution objective trains the model to match uniform distributions over all legal move-starting and move-ending squares rather than a single target token, which the authors describe as an explicit way of learning the transition rules. The joint probe objective adds a linear board state probe — a 13-class classification per square (six piece types for each side plus empty) — to the next-token or distribution loss. Evaluation uses 1000 unique 10-move prefixes from each model's training data as warmup sequences, after which the adversary attacks. A separate experiment replaces in-distribution warmup prefixes with random valid out-of-distribution prefixes of 20, 40, and 60 moves per player.

Why This Matters

Impact on research. The paper provides a falsification-based alternative to threshold-sensitive verification methods, and it suggests that probing results can be misleading as evidence about what a model actually uses. It also offers a set of concrete training recipes and failure-mode analyses that other researchers can reuse or contradict.

Real-world applications:

• AI safety and reliability auditing for generative systems deployed in domains with hard rules, such as code generation, where a "sound" model must never emit an invalid construct.
• Chess and board-game engine evaluation, where rule violations in generated move sequences are directly detectable and costly.
• Interpretability tooling, by showing that probe accuracy alone is not evidence of causal use, motivating gradient- and attack-based causal tests.
• Data curation decisions, since the results show that large curated gameplay data can look deceptively safe under in-distribution evaluation but collapse out of distribution.

Industry relevance. Teams that train sequence models on rule-governed data (games, planners, formal languages, structured generation) can adopt adversarial sequence generation as a red-teaming procedure. The finding that multi-task auxiliary heads do not improve rule adherence is directly relevant to practitioners who use auxiliary losses as a shortcut to better internal representations.

Future Directions

• Testing other architectures. The authors state their main limitation is relying on one generative sequence model architecture, and that the effect of different architectures remains to be analyzed.
• Improving adversaries further. Given that IMO dominates and that BSO and AD are weak, a natural next step is designing adversaries that are both stronger and cheaper than explicit lookahead over illegal moves.
• Explaining the sequence-length overfitting. The sudden unreliability beyond training sequence lengths (up to 150) points to an open question about why models fail to learn length-independent state abstractions.
• Resolving the gap between probing and behavior. Since gradient alignment is near-orthogonal and BSO attacks succeed without misleading the probe, more work is needed on why probes can decode board state so well while apparently having no causal role in generation.

Target Audience

Researchers in machine learning interpretability, formal-language learning, and AI safety evaluation, especially those working on emergent capabilities, probing methods, and world models. Practitioners who train sequence models on rule-governed data and need practical red-teaming tools will also benefit, as will readers following the debate about whether board-game language models truly acquire consistent world models. Some background in transformers and formal-language definitions is required. The paper states that its code, models, and datasets are available at https://github.com/szegedai/world-model-verification.

Authors’ abstract

Generative sequence models are typically trained on sample sequences from natural or formal languages. It is a crucial question whether -- or to what extent -- sample-based training is able to capture the true structure of these languages, often referred to as the ``world model''. Theoretical results indicate that we can hope for soundness at best, that is, generating valid sequences, but not necessarily all of them. However, it is still important to have practical tools that are able to verify whether a given sequence model is sound. In this study, we focus on chess, as it is a domain that provides enough complexity while having a simple rule-based world model. We propose adversarial sequence generation for verifying the soundness of the sequence model. Our adversaries generate valid sequences so as to force the sequence model to generate an invalid next move prediction. Apart from the falsification of soundness, this method is also suitable for a more fine-grained analysis of the failure modes and the effects of different choices during training. To demonstrate this, we propose a number of methods for adversarial sequence generation and evaluate the approach on a large set of chess models. We train models on random as well as high-quality chess games, using several training recipes. We find that none of the models are sound, but some training techniques and dataset choices are able to improve soundness remarkably. We also investigate the potential application of board state probes in both our training and attack methods. Our findings indicate that the extracted board states have no causal role in next token prediction in most of the models.