TLDR: A research paper demonstrates that training language models with causal masking on spatial chess board data (FEN) leads to stronger performance than training on sequential move data (PGN). This suggests that accepting information loss from causal masking on spatial inputs can be preferable to linearizing spatial data, offering a viable method for training unimodal LLMs on spatial datasets.
A recent research paper explores an intriguing question in the field of artificial intelligence: can language models, traditionally designed for sequential data like text, effectively learn from spatial data using their inherent “causal masking” mechanism? This mechanism, which enforces a left-to-right prediction, is often considered unsuitable for data with spatial or relational structures, where information is not strictly ordered.
The study, titled “Causal Masking on Spatial Data: An Information-Theoretic Case for Learning Spatial Datasets with Unimodal Language Models,” by Jared Junkin and Samuel Nathanson from Johns Hopkins University, delves into this issue using the game of chess as a unique testbed. Chess is ideal because it offers two distinct ways to represent a game: PGN (Portable Game Notation), which is a sequence of moves, and FEN (Forsyth-Edwards Notation), which describes the spatial arrangement of pieces on the board at a given moment.
The researchers trained several language models, including Meta AI’s 1.3 Billion Parameter Llama3.1, on both types of chess data. A key aspect of their investigation was comparing models trained with causal masking (the standard for language models) against those using bidirectional attention (which can look at all parts of the data simultaneously). Surprisingly, their findings revealed that models trained on spatial board states (FEN), even when constrained by causal masking, consistently achieved stronger playing strength than models trained on sequential move data (PGN).
This suggests a significant methodological insight: for certain domains, applying causal masking directly to spatial data might be a viable, and even preferable, approach compared to first converting spatial data into a sequential format. The paper posits that this is because PGN-based models must internally reconstruct a spatial understanding of the chessboard from the sequence of moves, adding a layer of functional complexity. FEN-based models, on the other hand, can directly process the spatial structure, making the learning task more efficient.
The Llama model, fine-tuned on the FEN dataset with causal masking, demonstrated grandmaster-level performance, achieving an estimated ELO rating of 2630. This was notably higher than a causal model trained on PGN data, which reached an ELO of 2000. A bidirectional model trained on FEN performed slightly better at 2680 ELO, but the strong performance of the causal FEN model highlights the potential of this approach.
Crucial to these results were careful tokenization and prompting strategies. The team enforced character-level tokenization for FEN strings and used templated prompts that embedded the FEN, legal moves, and the best Stockfish move. These steps were vital for stable training and effective learning.
Also Read:
- GeoFM: A New Approach to Boost Geometric Reasoning in AI Models
- LLMCOMP: Advancing Scientific Data Compression with Large Language Models
While the study’s empirical results are confined to chess, the authors believe their findings have broader implications for training unimodal large language models on other spatially structured datasets. Chess, often called “the Drosophila of artificial intelligence,” continues to be a fertile ground for exploring fundamental questions in AI research, including representation learning and masking strategies. For more details, you can read the full research paper here.
