Out-of-distribution Tests Reveal Compositionality in Chess Transformers

TL;DR

The paper investigates whether chess Transformers demonstrate genuine compositional generalization by evaluating a ~270M-parameter policy on diverse out-of-distribution positions and chess variants. Using behavior cloning on ChessBench and Stockfish as oracle, it measures legal-move adherence, topK move quality, puzzle-sequence accuracy, and Elo across Standard, Chess960, and Horde. It finds strong rule extrapolation with near-perfect legality in ID and robust legality in many OOD cases, plus high-quality OOD puzzle play, but limited strategy adaptation in the Horde and highly divergent starts. The results indicate that purely data-driven Transformers encode substantial compositional structure of chess, though there remains a gap to explicit symbolic search, with implications for understanding reasoning and generalization in AI, as well as for robust deployment in varied game settings.

Abstract

Chess is a canonical example of a task that requires rigorous reasoning and long-term planning. Modern decision Transformers - trained similarly to LLMs - are able to learn competent gameplay, but it is unclear to what extent they truly capture the rules of chess. To investigate this, we train a 270M parameter chess Transformer and test it on out-of-distribution scenarios, designed to reveal failures of systematic generalization. Our analysis shows that Transformers exhibit compositional generalization, as evidenced by strong rule extrapolation: they adhere to fundamental syntactic rules of the game by consistently choosing valid moves even in situations very different from the training data. Moreover, they also generate high-quality moves for OOD puzzles. In a more challenging test, we evaluate the models on variants including Chess960 (Fischer Random Chess) - a variant of chess where starting positions of pieces are randomized. We found that while the model exhibits basic strategy adaptation, they are inferior to symbolic AI algorithms that perform explicit search, but gap is smaller when playing against users on Lichess. Moreover, the training dynamics revealed that the model initially learns to move only its own pieces, suggesting an emergent compositional understanding of the game.

Figures (6)

• Figure 1: OOD board types: We consider four out-of-distribution chess scenarios. The first board shows a position with more pieces of a certain type than allowed (the exception is via pawn promotion), 3 white queens in this case; while the second one illustrates 2 white bishops on squares with the same color. We also study two chess variants: Chess960, where in the starting position pieces on the first rank are randomly reordered and this is mirrored in the 8th rank (third board); and Horde, where White only has pawns and Black has a different objective (fourth board). The next moves of our model are highlighted with purple.
• Figure 2: Knights&Rooks: extreme cases of boards having more pieces (knights and rooks) than normally allowed.
• Figure 3: The illegal moves on OOD puzzles arise from the model trying to move pinned pieces. E.g. on the fourth board, moving the black queen d5g2 would mate the white king, but it allows to the white queen on b3 to put the black king into check.
• Figure 4: Training dynamics of piece selection: Heatmap of an ID Puzzle board showing summed probabilities of next moves starting from each square, depicted at initialization (left), during training (middle), and after training (right). On each board, the probabilities are normalized to $[0;1]$ by the maximum value, dark blue indicates 0, yellow indicates 1.
• Figure 5: Training dynamics of move legality: the left and middle plots illustrate the legal next move accuracy during training on the ID (blue) and the OOD Puzzles (orange); and the sum of all legal moves' probabilities from all possible moves for the ID (green) and the OOD Puzzles (red). Averages are taken over 1000 puzzles. The two plots are scaled differently to better see the beginning of the training. The right plot shows the relative legal probability of the pieces, calculated as $p(\text{legal moves of a given piece})/p(\text{all moves with a given piece})$. The notation of the pieces is the following: N - knight, R - rook, B - bishop, Q - queen, P - pawn, K - king. These values were calculated on very simple boards containing one black, one white king and only the type of piece whose legal move was examined.