Source-linked AI summary
Chess as a Testbed for Language Model State Tracking
Shubham Toshniwal, Sam Wiseman, Karen Livescu, Kevin Gimpel
TL;DR
The paper addresses the difficulty of measuring how accurately language models track the world state underlying text. It uses chess notation to probe board-state tracking and reports that results depend on training data and access to game history.
Problem
The complexity of natural language makes it challenging to measure how accurately language models represent the world state underlying text.
Source: §1 · p. 1Method
Results
Takeaways & Limitations
Chess is proposed as a benchmark for developing and analyzing language models’ world-state tracking capabilities.
Source: §1 · p. 2, §7 · p. 7Reformer results use transformers-library version 4.2.2 while its implementation was still a work in progress.
Source: §5.1 · p. 5
Abstract
from arXiv · showhide
Transformer language models have made tremendous strides in natural language understanding tasks. However, the complexity of natural language makes it challenging to ascertain how accurately these models are tracking the world state underlying the text. Motivated by this issue, we consider the task of language modeling for the game of chess. Unlike natural language, chess notations describe a simple, constrained, and deterministic domain. Moreover, we observe that the appropriate choice of chess notation allows for directly probing the world state, without requiring any additional probing-related machinery. We find that: (a) With enough training data, transformer language models can learn to track pieces and predict legal moves with high accuracy when trained solely on move sequences. (b) For small training sets providing access to board state information during training can yield significant improvements. (c) The success of transformer language models is dependent on access to the entire game history i.e. "full attention". Approximating this full attention results in a significant performance drop. We propose this testbed as a benchmark for future work on the development and analysis of transformer language models.
1 Introduction
Chess offers a simple, deterministic testbed for probing whether language models track the world state underlying text. The paper reports that performance depends on training data and access to game history.
1 Introduction
- Chess is proposed as a testbed for evaluating language models’ world-state tracking, using notation that permits simple prompts without new classifiers or changes to the language-modeling objective.The paper also proposes the testbed for model development and analysis.Source: §1 · pp. 1–2
- Chess’s deterministic transitions support fine-grained evaluation, including whether a predicted move is legal and counterfactual queries.Illegal moves can also be analyzed automatically.Source: §1 · p. 1
- Generating a valid next move requires tracking both sides’ pieces, following chess rules, and predicting the actual move.The passage frames these as distinct challenges for a language model.Source: §1 · p. 1
- With enough data, transformers track piece locations and predict legal moves accurately; with small datasets, board-state tokens improve piece tracking.The models also show a performance drop when access to the full game history is limited.Source: §1 · p. 1
- Models’ performance also drops when they cannot access the whole game history, despite chess being Markovian.The paper also reports robustness to changes in input distribution when board-state tokens are added only during training.Source: §1 · p. 1
2 Chess Preliminaries
The paper represents chess moves in UCI notation and tokenizes board squares as individual tokens. Its vocabulary contains square, piece-type, and special-symbol tokens.
2 Chess Preliminaries
- UCI notation represents a move with its starting and destination squares, as in f1b5.The paper chooses UCI over standard algebraic notation for its move representation.Source: §2 · p. 2
- The tokenizer treats each board square as one token, producing a vocabulary of 77 token types.The vocabulary includes the 64 squares, piece-type symbols, and other special symbols.Source: §2 · p. 2
- Removing the move-boundary delimiter did not degrade performance and made training faster by reducing sequence length.The paper reports this in an implementation note.Source: §2 · p. 2
3 Language Model Prompts as Board State Probes
The section uses game-prefix prompts and next-token predictions to probe board-state tracking, with separate tasks for predicting move destinations and piece locations. It defines RAP and an always-annotated variant to enable direct piece-location probes.
3 Language Model Prompts as Board State Probes
- A game prefix followed by a prompt token lets next-token predictions reveal what the model understands about the implied board state.For a bishop at f1, an impossible destination such as g1 may indicate piece-type confusion, while g2 may indicate failure to track an occupied square.Source: §3 · p. 2
3.1 Randomly Annotated Piece type (RAP)
- RAP randomly adds piece-type tokens to moves during training, enabling prompts that ask where a piece currently is rather than where it might move.At test time, the prompt supplies the piece type without annotating moves in the history.Source: §3.1 · p. 3
- The oracle UCI + AP variant includes piece types in move notation during both training and testing.AP stands for “always piece type.”Source: §3.1 · p. 3
- RAP changes training sequences but not inference sequences, creating a distribution mismatch that the authors hypothesize may provide useful piece-tracking supervision.The proposed supervision may also improve language-modeling performance.Source: §3.1 · p. 3, Table 2 · p. 3
3.3 Ending Square Tasks
- Ending-square prompts ask for a move’s starting square, then use the predicted next token as its destination; End-Actual and End-Other differ in which piece is prompted.Table 3 gives examples; End-Actual uses the piece moved next in the game, while End-Other allows any legally movable piece.Source: §3.3 · p. 3, §3.2 · p. 3, Table 3 · p. 4
- End-Actual is evaluated by exact-move and legal-move accuracy, while End-Other permits only legal-move evaluation.Legal-move evaluation also uses R-Precision; the authors note that move-destination predictions partly conflate board-state knowledge with chess strategy, whose capabilities are tied to training-game quality.Source: §3.3 · pp. 3–4
3.4 Starting Square Tasks
- Starting-square prompts supply a piece type and test whether the model tracks where that piece is on the board; only models trained with piece types can be tested this way.Start-Actual and Start-Other mirror the actual-piece versus any-legally-movable-piece distinction, with exact-move and legal-move evaluations for Start-Actual.Source: §3.4 · p. 4
4 Experimental Setup
The experiments use filtered chess-game datasets, separate language-modeling and board-state-probing evaluations, and GPT2-small trained from scratch. The setup also compares RAP rates and alternative attention models.
4 Experimental Setup
- The filtered Millionbase collection contains about 2.5 million games, with 200K for training, 15K each for development and test, and 50K reserved for probing.Training subsets of 15K, 50K, and 200K games are called Train-S, Train-M, and Train-L.Source: §4 · p. 4
- The board-state evaluation uses 1,000 unseen-prefix instances from non-pawn pieces, with prefixes 51–100 moves long.The 50K games reserved for probing supply these examples.Source: §4 · p. 4
- Models use GPT2-small’s 12-layer, 12-head architecture with a 512-token context and are trained from scratch.The context covers the longest game in the training set.Source: §4 · p. 4
- RAP probability is tuned over 5%, 15%, 25%, 50%, 75%, and 100% using validation perplexity; 25% is best for Train-S and Train-M, and 15% for Train-L.Higher probabilities increase train–inference mismatch, while lower probabilities may provide insufficient training signal.Source: §4 · p. 4
- The experiments include Reformer and Performer as approximate-attention transformers, alongside LSTM models with and without RAP.The reported comparison also includes canonical validation and test perplexity, counting one move as one token.Source: §4 · pp. 4–5, Table 4 · p. 5
- Training runs for up to 10 epochs with batch size 60, Adam optimization, and validation-based early stopping.The setup uses a 5×10^-4 learning rate with warmup and linear decay, plus 0.01 L2 weight decay.Source: §4 · p. 5
5 Results
Across the results, transformers learn board-state tracking and legal-move prediction with enough data, while state tokens can help in smaller-data settings. Restricting access to game history reduces performance, and approximate-attention models do not generally match full-attention GPT2.
5 Results
- The results show high-accuracy piece tracking and legal-move prediction with enough data, while limiting access to game history reduces performance.RAP and AP lower perplexity, particularly with small training sets; models with restricted or approximate attention perform worse than the full-attention baseline.Source: §5 · p. 5, §5.1 · p. 5, §5.3 · p. 6
5.1 Language Modeling
- RAP and AP reduce perplexity across training sizes; RAP probabilities up to 50% help smaller sets, while UCI + AP achieves the lowest perplexities on larger sets.For larger sets, lower RAP probabilities are preferred; piece-type logits are masked when calculating perplexity for UCI + RAP models.Source: §5.1 · p. 5, Figure 2 · p. 5
5.2 Board State Tracking
- UCI + RAP predicts legal starting positions with perfect accuracy and R-Precision on Train-L, but accuracy falls to 91.3% on Train-S.For other pieces’ locations, accuracy drops only slightly, to 99.6%; the difference between exact-move and piece-location accuracy exceeds 8 percentage points and reflects strategy rather than state tracking.Source: §5.2 · p. 5, Table 5 · p. 6
- UCI and UCI + RAP exceed 97% accuracy for top legal ending-square predictions, but R-Precision for predicting all legal moves is about 85%.The models train on actual games emphasizing meaningful moves, and performance also drops when predicting endings for a different legal starting square.Source: §5.2 · pp. 5–6, Table 6 · p. 6
5.3 Compressing the Game History
- A 50-token GPT2 attention window performs worse than full attention, especially on small and medium datasets, while the LSTM drops further in all settings.LSTM + RAP improves performance only on larger sets; transformer models generally improve with RAP data.Source: §5.3 · p. 6, Table 7 · p. 6
Approximate Attention Transformers
- Performer generally outperforms Reformer except on the small dataset, but neither substantially beats GPT2 with a 50-token window except Performer on medium data.Both approximate-attention architectures typically outperform LSTMs, illustrating the challenge of modeling chess with approximate attention.Source: Approximate Attention Transformers · p. 7
6 Related Work
Prior simulated-world and chess work provides access to world state but generally targets other tasks, while natural-language cloze and probing evaluations can require curation or added models. This work’s distinction is that chess notation supports direct state probing under a language-modeling objective.
Simulated Worlds
- Simulated-world frameworks make true world state available, but this work uses chess notation to probe state tracking in a model trained with a language-modeling objective.Prior examples include bAbI, TextWorld, and 3D world simulations.Source: §6 · p. 7
Cloze Tasks for Natural Language Models
- Natural-language cloze tasks cover text understanding, commonsense, narrative, and factual knowledge, but often require human curation and exact-match evaluation.The proposed chess tasks are precisely automated without human curation and permit fine-grained evaluation.Source: §6 · p. 7
Probing and Deep Learning for Chess
- Prior probing commonly trains an additional model, which is time-consuming and can make conclusions depend on probe complexity; chess notation instead enables simple-prompt probing.Chess-playing work focuses on move quality given the true board state, unlike this paper’s focus on language-model board-state tracking.Source: §6 · p. 7
7 Conclusion
The conclusion proposes chess as a testbed for probing language models’ world-state tracking, and reports that board-state training information can help when data is limited. The accompanying sections explain why UCI is used instead of SAN, describe vocabulary choices, and show how model-size effects depend on training-set size.
7 Conclusion
- Chess provides a testbed for probing language models’ world-state tracking; with enough data they track board state, while board-state information during training can improve results with limited data.The conclusion also notes that chess’s precise dynamics allow varying explicit state during training and evaluating predictions at a fine-grained level.Source: §7 · p. 7
A SAN Notation
- SAN represents a move by the moved piece type and destination square, but omits the starting square, limiting piece-specific prompts.The authors therefore do not use SAN in their experiments.Source: A SAN Notation · p. 10
B Model Vocabulary
- The vocabulary includes tokens for promoted pawn piece types; for example, e7e8q represents a pawn moving from e7 to e8 and becoming a queen.The section refers readers to Table 1 for the vocabulary.Source: B Model Vocabulary · p. 10
C Effect of Model Size
- With Train-S, GPT2-small outperforms GPT2-intermediate and GPT2-medium on almost all evaluations, but the larger models outperform GPT2-small on most evaluations as training data increases.The authors attribute the pattern to larger training sets alleviating overfitting and note that tuning regularization such as dropout could further improve larger models on small datasets.Source: C Effect of Model Size · p. 10
D Data Statistics
The dataset statistics describe game lengths and the prefixes used in board-state prompts, including how prompts are constructed for actual and alternative moves.
D Data Statistics
- The average game length is around 75 moves across all dataset splits.Figure 3 shows the distribution of tokenized UCI game lengths for Train-L.Source: D Data Statistics · p. 10
- Board-state prompts share the same game prefix, but starting-square prediction prompts use the actual move’s piece type while ending-square prompts use its starting square.End-Actual and Start-Actual therefore share the same piece type for each prompt.Source: D Data Statistics · p. 10
- For alternative-move prompts, End-Other samples a different legal starting square, while Start-Other samples a different legally movable piece type.Figure 4 shows the distribution of game-prefix lengths used in board-state prompts.Source: D Data Statistics · p. 10
E Model Hyperparameters and Training time
The experiments use separate hyperparameter searches for several training configurations, with special handling for Reformer. The section also reports approximate training times by model family and dataset size.
E Model Hyperparameters and Training time
- For baselines other than GPT2-small, hyperparameters are searched separately for Train-S and Train-M, with Train-M settings used for Train-L except for Reformer.GPT2-small uses default hyperparameters.Source: E Model Hyperparameters and Training… · p. 10
- Reformer combines local and LSH-based attention with 50-token chunks; its Train-L loss rose after a few epochs, prompting separate tuning and an inverse-square-root decay schedule.The schedule produced slightly better final results; the other experiments use Train-M hyperparameters and the learning-rate schedule described in Section 4.Source: E Model Hyperparameters and Training… · pp. 10–11
- Transformers take around 4 hours for Train-S, under 10 hours for Train-M, and under 24 hours for Train-L on one GeForce RTX 2080 Ti.LSTMs take under 2, 4, and 8 hours for Train-S, Train-M, and Train-L, respectively, on the same GPU.Source: E Model Hyperparameters and Training… · p. 11
F Error Analysis
The error analysis separates ending-square mistakes into categories tied to different aspects of board-state tracking. Errors decline with more training data, but the patterns differ by error type and prediction setting.
F Error Analysis
- The analysis distinguishes unreachable, syntax, path-obstruction, and pseudo-legal errors; unreachable errors are rare, with fewer than five for every model.Syntax errors indicate failure to track the piece type at the starting square, while path obstruction reflects missed pieces or failure to account for blocked paths.Source: F Error Analysis · p. 11
- UCI + RAP consistently outperforms UCI on syntax errors, with the largest differences on small training sets, but shows no consistent trend across other error categories.The Performer generally makes more errors than transformers, especially in syntax errors; the authors suggest its partial attention may limit access to relevant prior positions.Source: F Error Analysis · p. 12, Table 12 · p. 13
- Actual-move ending-square predictions are easier than randomly selected legal moves, which have more path-obstruction and pseudo-legal errors despite comparable syntax errors.The proposed explanations include training on meaningful actual games and differences in the piece-type distribution, including a higher share of king moves among End-Other prompts.Source: F Error Analysis · p. 12
- Syntax errors decrease by roughly an order of magnitude from Train-S to Train-M, while path-obstruction and pseudo-legal errors decline more gradually.Across error categories, counts decrease with more training data; the authors relate the slower reductions to computations involving multiple piece locations.Source: F Error Analysis · pp. 11–12, Table 12 · p. 13
F.1 Detailed Error Analysis
The detailed error analysis examines Pseudo Legal and Path Obstruction errors in the UCI model trained with Train-L for ending-square prediction.
F.1 Detailed Error Analysis
- The analysis focuses on Pseudo Legal and Path Obstruction errors in the UCI model trained with Train-L for ending-square prediction.This section introduces a more in-depth analysis of those two categories.Source: §F.1 · p. 13
- Table 14 reports counts by piece type for the Path Obstruction error category.Its total column removes instances with errors from other categories.Source: Table 14 · p. 13
- The Path Obstruction totals therefore exclude instances classified as other error categories.Source: Table 14 · p. 13
F.2 Pseudo Legal Errors
Pseudo Legal errors are categorized by the king’s check status and whether the king moves; Path Obstruction errors vary by piece type and prompt split. The analysis also compares move-path lengths for obstructed and legal predictions.
F.2 Pseudo Legal Errors
- Pseudo Legal errors are categorized by whether the king is in check and whether the current move moves the king, yielding four combinations.Figure 6 presents one instance for each combination.Source: §F.2 · p. 13, Figure 6 · p. 14
- “Check + Other” errors are high in both splits, particularly End-Other, where some prompts offer few meaningful legal moves.For example, when the king is in check, the queen’s only legal move may be c7, while the model predicts c6.Source: §F.2 · p. 13, §F.3 · p. 14
- “No Check + King” is another common error category, with higher End-Other counts possibly reflecting prompt frequency and out-of-distribution prompts.Source: §F.2 · p. 13
F.3 Path Obstruction
- Rooks, bishops, and queens have more Path Obstruction errors than knights, with a larger frequency increase for End-Other than End-Actual.Predicted moves with Path Obstruction errors have significantly longer average paths than legal predictions for both prompt types; knight and king path lengths are constant.Source: §F.3 · p. 13, Figure 8 · p. 16
- Path Obstruction examples include a rook overlooking its own king and a bishop blocked by another piece.Source: §F.3 · p. 15