Technical Summary: REPAIR – Predictive Self-Supervised Representation Learning in Chess

Problem Statement

Current methods for learning chess representations predominantly rely on reinforcement learning (RL) with self-play (e.g., AlphaZero, Leela Chess Zero) or supervised learning using handcrafted heuristics and domain expertise. These approaches often require costly training setups and explicit move labels. While self-supervised learning has achieved significant success in image and natural language domains (e.g., MAE, BERT, JEPA), its application to sequential state data like chess games remains underexplored. The challenge is to learn a semantically rich, low-dimensional representation of chess states without relying on explicit move annotations, player ratings, or expensive RL, while enabling the model to reason about piece movements and reconstruct missing states in a sequence.

Methodology: The REPAIR Architecture

The authors propose REPAIR (Representation Prediction via Autoencoding using Iterative Refinement), a self-supervised architecture that synthesizes concepts from Masked Autoencoders (MAE), Joint Embedding Predictive Architectures (JEPA), and BERT. The framework operates on sequences of states (e.g., consecutive chess positions) through the following pipeline:

1. Encoder (Independent State Mapping):

   • Each chess state $s_i$ (represented as a $19 \times 8 \times 8$ bitmap) is mapped independently to a latent vector $z_i$ by a convolutional encoder.
   • The encoder utilizes Squeeze-and-Excitation blocks to learn channel weights, emphasizing important features.
   • Crucially, the encoder processes states independently; no temporal information flows between states at this stage.

2. Masking:

   • A sequence of latent vectors $Z$ is generated. Random portions of the sequence (excluding the first and last elements) are masked with a learnable token $\bar{z}_i$ with probability $p$.
   • The paper employs high masking ratios (up to 90%), leveraging the fact that chess sequences allow for large information inference from small contexts, unlike natural language.

3. Predictor (Iterative Refinement):

   • The Predictor is the only component allowing information exchange between states. It uses a single-layer, encoder-only Transformer with four attention heads.
   • It operates iteratively: starting with the masked sequence, it updates the latent representations $\hat{z}_i$ by attending to the entire sequence. This process repeats $n$ times, aiming to converge the masked latent states to the unmasked ground truth ($z_i \approx \hat{z}_i^{(n)}$).
   • This iterative refinement allows the model to "repair" gaps in the latent sequence.

4. Decoder (State Reconstruction):

   • The decoder is an inverted encoder that maps the refined latent vectors back to the input space ($s_i \approx \hat{s}_i$).
   • The goal is to reconstruct the original bitmap, effectively recovering missing board states.

5. Loss Functions:

   • JEPA Loss: Minimizes the distance between the Predictor's output and the unmasked ground truth latent representation.
   • Decoder Loss (Long): Computes the reconstruction error between the original state and the state decoded from the Predictor's output. This prevents representation collapse.
   • Decoder Loss (Short): Computes reconstruction error between the original state and the state decoded directly from the Encoder's output (without Predictor refinement).
   • The paper finds that combining the Long Decoder Loss with the Short Decoder Loss yields the best performance, effectively training the Predictor to refine representations without needing the JEPA latent alignment loss.

Key Contributions

• Novel Architecture: Introduction of REPAIR, which integrates MAE-style masking, JEPA-style latent prediction, and BERT-style sequence processing into a unified self-supervised framework for sequential state data.
• Self-Supervised Chess Learning: Demonstration that meaningful chess representations can be learned without reinforcement learning, handcrafted heuristics, or labeled move data, relying solely on the statistical structure of game sequences.
• Iterative Latent Repair: A mechanism where a lightweight Predictor iteratively refines masked latent states, enabling the model to reason about piece movements and reconstruct large gaps in game sequences.
• Semantic Clustering: Evidence that the learned latent space naturally clusters chess concepts, including game phases (opening, middlegame, endgame), specific opening families (ECO codes), and puzzle motifs.

Experimental Results

The model was trained on 80,000 Lichess games (June 2025) with a 50% masking ratio, followed by a phase with 90% masking.

• Reconstruction Accuracy:

  • Models trained with the Long Decoder Loss achieved a top-1 reconstruction accuracy of approximately 93% on masked squares, even with 80% of the sequence masked.
  • Combinations lacking the Long Decoder Loss resulted in significantly lower accuracy or model collapse.
  • The model demonstrated the ability to reconstruct legal move patterns and exclude non-sensible moves (e.g., ignoring moves that lead to immediate piece capture).

• Attention Analysis:

  • Visualization of the Predictor's attention weights showed that information initially spreads from masked states to neighboring unmasked states and then propagates further.
  • The attention mechanism stabilizes after approximately five iterations, corresponding to the point where state changes stagnate.

• Representation Space Semantics:

  • Game Phases: PCA projections of full game trajectories revealed distinct clusters corresponding to openings, castling events (separate for White and Black), middlegames, and endgames.
  • Openings: t-SNE projections of ECO opening positions showed clear clustering by ECO code (e.g., King-pawn vs. Queen-pawn openings), with the learned embeddings showing better separation than raw bitmaps.
  • Puzzles: The embedding space successfully separated chess puzzles based on semantic motifs (e.g., "advanced pawn," "back-rank mate," "queen endgame"), confirming the model captures high-level strategic concepts.

Significance and Claims

The paper claims that REPAIR provides a semantically rich representation space for chess that emerges purely from self-supervised training. Unlike previous approaches where embeddings are a by-product of a policy trained for winning, REPAIR focuses explicitly on learning meaningful embeddings for the states themselves.

The authors assert that the model demonstrates an "intuition" of chess rules and strategies, evidenced by:

1. The ability to reconstruct missing board states with high accuracy.
2. The generation of probability distributions that respect legal move constraints and tactical consequences (e.g., avoiding moves that lose pieces).
3. The natural clustering of game trajectories and puzzle motifs in the latent space, allowing for intuitive dissection of games by observing their paths.

The work suggests that such self-supervised architectures can be applied to other domains involving successive transformations, offering a pathway to explainable processes without the need for costly reinforcement learning or explicit labeling.