Causal Masking on Spatial Data: An Information-Theoretic Case for Learning Spatial Datasets with Unimodal Language Models

AI summary

Overview

• Research area: Large language models, attention masking strategies, and spatial vs. sequential data representation, studied through the domain of chess.
• Technical level: Intermediate. The paper uses reinforcement-learning and information-theoretic vocabulary (Markov policies, causal masking, cross-entropy), but explains each idea in plain terms and grounds everything in a concrete chess setup.
• Scope: A methodological case study comparing causal-masked language models trained on spatial (FEN board-state) chess data against models trained on sequential (PGN move-list) chess data.

What This Paper Is About

Language models are built around causal masking—each token can only see the tokens before it—which fits text but seems wrong for data that is spatial rather than sequential, such as a chess board. The paper asks whether it is acceptable to simply apply causal masking to spatial data anyway, accepting the resulting information loss, instead of converting spatial data into a sequential format so a standard causal model can consume it. The authors use chess because it supports both a spatial encoding (FEN) and a sequential encoding (PGN) of the same underlying positions, letting them test the two approaches directly.

Key Contributions

1. A direct masking-vs-sequentialization comparison. The paper isolates the choice of representation by comparing a causal-masked model on spatial FEN data against a causal-masked model on sequential PGN data, holding the setting as constant as possible.
2. An information-theoretic argument. The authors argue that a PGN-trained model must internally reconstruct a latent spatial board representation before choosing a move, forming a composition of functions (PGN to FEN-like latent space to move), while a FEN-trained model maps directly from board state to move. They present this as intuition rather than a formal proof.
3. Empirical validation across three model setups. Two identically sized character-level models (one on PGN with causal masking, one on FEN with bidirectional attention) plus Meta AI's 1.3B-parameter Llama 3.1 fine-tuned on FEN with causal masking.
4. A strong open-source chess baseline with practical lessons. The authors describe the Llama model as, to their knowledge, the first open-source fine-tuned LLM to reach performance in the ~2600 Elo range when calibrated against Stockfish agents, and they document the tokenization and prompting choices that made training converge.

Main Findings

• FEN beats PGN on loss and accuracy. In a preliminary experiment, two identical NanoGPT-style models (based on Karvonen's 50M-parameter decoder-only setup) were trained for 200,000 steps on equivalently sized PGN and FEN datasets with equal batch sizes. The PGN model showed substantially higher cross-entropy loss and lower best-move classification than the FEN model.
• Causal masking on spatial data still beat sequentialization. The causal-masked Llama 3.1 fine-tuned on FEN reached an estimated Elo of 2630 when calibrated against Stockfish agents—roughly 500 points higher than the causal model trained on PGN, and only narrowly worse than the bidirectional model trained on FEN.
• Fine-tuning produced large gains on a held-out set. On ~12,800 held-out positions, the Llama model went from a zero-shot best-move rate of ~0.6% to ≈58% after supervised fine-tuning (roughly a 100× increase). Syntactically valid move rate was 99.94% and legal move rate was 99.91%.
• A clear gradient in best-move quality. All three models produced valid and legal moves with nearly 100% accuracy, but the PGN-trained model chose the best Stockfish move only 40% of the time, the causal FEN Llama model 58%, and the bidirectional FEN model over 60%.
• Tokenization was decisive. Both the Pythia and Llama models failed to converge under their default tokenizers, which merge character sequences such as "pk" (pawn–king) into a single token. Forcing strict character-level tokenization (so "pk" became ['p','k'] → [79, 74]) and flattening run-length encodings in FEN strings stabilized training and produced strong results.
• Exposure bias was modest. Teacher-forced and autoregressive decoding differed only slightly in best-move prediction (Llama: ~3% drop; Pythia: <1% drop). Interestingly, both models were more likely to generate legal and valid SAN moves in autoregressive mode, suggesting they internalized structural constraints of chess notation.
• Pretraining quality mattered in a side comparison. A 140M-parameter Pythia model was consistently outperformed by Llama even at equivalent batch sizes and identical training steps, though both are overparameterized for chess given that 50M-parameter models have reached high skill levels in prior work.

Methodology in Plain English

The researchers picked a domain where the same information exists in two forms. PGN is a list of moves played in order, so it can be fed to a standard next-token-prediction language model. FEN is a snapshot of a board layout plus side-to-move, castling rights, counters, and so on, so it is spatial and has no temporal dimension.

They trained three models. Two were small character-level models trained from scratch: one on PGN with a standard lower-triangular causal attention mask, and one on FEN with no mask at all (bidirectional attention), since the FEN data has no natural ordering to respect. The third was a 1.3B-parameter Llama 3.1 fine-tuned on FEN data while keeping causal masking on, which is the configuration the paper is really about.

For the FEN setup, each board state was paired with a single ground-truth best move from Stockfish, turning the task into classification over legal moves rather than next-token prediction. Prompts embedded the FEN string, the list of legal SAN moves, and the best move. The loss was masked so only the tokens belonging to the best-move label contributed a penalty—a fully teacher-forced objective—and padding tokens were also masked.

Training used 200,000 steps on two Nvidia A100 GPUs (80GB each) with cosine learning-rate decay, 2,000 warmup steps, gradient clipping at ±1.0, mixed precision, and gradient accumulation, taking roughly three weeks per run on their hardware.

Evaluation happened two ways. First, on ~12,800 held-out board positions, measuring the percentage of syntactically valid moves, legal moves, and best Stockfish moves. Second, by playing each model in a simulated tournament against attenuated Stockfish agents at calibration Levels 0–10, with 1,000 games per level (500 as White, 500 as Black), an opening book, and temperature-based sampling; a game was forfeited as a loss if a model produced more than 5 consecutive illegal moves (in practice, no model forfeited this way). Elo updates used K = 16.

Data came from two sources: the ChessBench dataset from Ruoss et al., containing 15 billion board positions annotated with the top move according to Stockfish, and a PGN dataset of approximately 1 billion PGN strings generated from games between a maximum-strength Stockfish agent as White and attenuated engines rated between 1200 and 3100 Elo as Black. The PGN-trained model only played as White in evaluation, while both FEN-trained models played both colors.

Why This Matters

The paper's claim is methodological rather than chess-specific: if causal masking on spatial data is viable—and in this domain preferable to sequentialization—then researchers may not need natively multimodal architectures with their cross-modal attention, paired-dataset curation burdens, and modality-specific tokenization to get useful behavior out of structured spatial inputs. The authors also note that it is possible to learn sophisticated domain behavior with small language models and modest compute.

Real-world applications that follow from this line of reasoning:

• Board-game and strategy agents built on existing pretrained language models rather than bespoke architectures.
• Structured symbolic or tabular domains where the underlying relations are spatial or relational but the input format is naturally symbolic—the paper's argument is that preserving that structure is preferable to linearizing it.
• Planning and world-model research where an agent must reason about a state rather than a history, and where the authors argue causal masking can remain a viable design choice for structured reasoning tasks.
• Interpretability research, since the authors describe abstract games as a small-vocabulary, visually intuitive, quantitatively scorable sandbox for probing how transformers form world models.

Industry relevance: the paper targets the practical tradeoff between training complexity and capability, arguing that simpler unimodal training pipelines can suffice for spatially structured data, which matters for teams with constrained compute who want to adapt pretrained models to a specialized domain.

Future Directions

• Disentangling pretraining quality from parameter count. The authors ask whether the Llama-vs-Pythia gap comes from the quality and duration of pretraining, or from the number of non-embedding parameters—even when both models are overparameterized for the task. They leave both possibilities open.
• Generalizing beyond chess. The empirical results are confined to chess; whether causal masking on spatial data helps in other spatially structured domains remains future work.
• Handling the threefold repetition rule. FEN-trained models are blind to past board states and therefore cannot avoid certain draws without additional mechanisms, which the authors flag as a limitation of the spatial representation.
• Interpretability and scaling-law studies. The paper suggests abstract strategy games as a sandbox for developing interpretability methods, and proposes further study of scaling laws for this kind of domain adaptation.

Target Audience

Researchers and engineers working on attention masking, multimodal or spatial reasoning, and domain adaptation of pretrained language models; chess AI practitioners interested in an open-source fine-tuned baseline; and interpretability researchers looking for a small-vocabulary, quantitatively scorable testbed. Readers need basic familiarity with transformer attention and Elo-style evaluation, but the paper explains its chess encodings and objectives from the ground up.

Authors’ abstract

Language models are traditionally designed around causal masking. In domains with spatial or relational structure, causal masking is often viewed as inappropriate, and sequential linearizations are instead used. Yet the question of whether it is viable to accept the information loss introduced by causal masking on nonsequential data has received little direct study, in part because few domains offer both spatial and sequential representations of the same dataset. In this work, we investigate this issue in the domain of chess, which naturally supports both representations. We train language models with bidirectional and causal self-attention mechanisms on both spatial (board-based) and sequential (move-based) data. Our results show that models trained on spatial board states - \textit{even with causal masking} - consistently achieve stronger playing strength than models trained on sequential data. While our experiments are conducted on chess, our results are methodological and may have broader implications: applying causal masking to spatial data is a viable procedure for training unimodal LLMs on spatial data, and in some domains is even preferable to sequentialization.