Blog Post👁️ | Interactive Demo🎮 | Model🤗 | W&B Logs📊
ROOK-CLF-9M is a 9M parameter chess move prediction model using a classification approach, reproducing the methodology from Google DeepMind's "Grandmaster-Level Chess Without Search". Unlike traditional chess engines that rely on search algorithms, this model directly predicts the best move from a position using a transformer architecture.
Key Features:
- 🎯 49% action accuracy (ChessBench 40M; LAION note)
- ✅ 57% accuracy on BIG-bench Checkmate-in-One (LAION note)
- 🌐 In-browser inference via ONNX Runtime (try the interactive demo)
- 🚀 9M parameters - Efficient decoder transformer (Llama architecture)
- 📊 Custom tokenization - FEN positions encoded as 78 tokens (77 chars + [CLS])
- Model Type: LlamaForSequenceClassification (decoder-only transformer)
- Parameters: 9M (8 layers, 8 heads, embedding dim 256, context length 78)
- Task: Text classification over 1968 legal chess moves
- Training: HuggingFace Transformers with custom chess tokenization
- Inference: ONNX export for web deployment
The model uses a custom tokenization scheme critical for proper inference:
Step 1: FEN Processing (77 characters fixed)
# Original FEN
fen = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
# Process FEN to fixed 77-character format:
# 1. Expand numbers to dots (e.g., "8" → "........")
# 2. Remove slashes
# 3. Pad castling to 4 chars, en passant to 2 chars, halfmove to 3 chars, fullmove to 3 chars
def process_fen(fen):
position, turn, castling, en_passant, halfmove, fullmove = fen.split(" ")
# Expand empty squares: "8" → "........"
position = re.sub(r'\d+', lambda m: "." * int(m.group()), position)
position = position.replace("/", "") # Remove row separators
castling = castling.ljust(4, ".") # Pad to 4 chars
en_passant = en_passant.ljust(2, ".") # Pad to 2 chars
halfmove = halfmove.ljust(2, ".") + "." # Pad to 3 chars total
fullmove = fullmove.ljust(3, ".") # Pad to 3 chars
return "".join([position, turn, castling, en_passant, halfmove, fullmove])
# Result: exactly 77 characters
processed = process_fen(fen)
# "rnbqkbnrpppppppp................................PPPPPPPPRNBQKBNRwKQkq-...0..1.."Step 2: Add [CLS] token and convert to token IDs
# Add classification token
final_input = processed + "[CLS]" # 78 characters total
# Convert to token IDs (character-level tokenization)
tokens = [char_to_id[c] for c in final_input] # 78 tokensComplete example:
Input FEN: "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
Processed: "rnbqkbnrpppppppp................................PPPPPPPPRNBQKBNRwKQkq-...0..1.."
With [CLS]: "rnbqkbnrpppppppp................................PPPPPPPPRNBQKBNRwKQkq-...0..1..[CLS]"
Token IDs: [13, 11, 3, 12, 10, 3, 11, 13, 15, 15, 15, 15, 15, 15, 15, 15, ...] # 78 tokens
Key Details:
- Position: 64 chars (numbers expanded to dots, slashes removed)
- Turn: 1 char (w/b)
- Castling: 4 chars (KQkq padded with dots)
- En passant: 2 chars (- padded with dots)
- Halfmove: 3 chars (padded with dots)
- Fullmove: 3 chars (padded with dots)
- Total: 77 chars + [CLS] = 78 tokens
This fixed-length encoding ensures consistent attention patterns and enables efficient batch processing.
- Training Data: Lichess games + Stockfish 16.1 selfplay
- Annotations: Best moves and top-5 candidates with evaluation scores
- Size: Multiple scales tested (20k to 5M positions)
- Preprocessing: FEN standardization, move legality validation
| Dataset Size | Best Move Accuracy | Top-5 Accuracy | Checkmate-in-One |
|---|---|---|---|
| 20k | 0.6% | 1.4% | 0.0% |
| 709k | 8.8% | 28.2% | 7.0% |
| 5M | 13.4% | 39.6% | 11.5% |
# Clone repository
git clone https://github.com/jorahn/rook.git
cd rook
# Install dependencies
pip install -r requirements.txt
# Download model (optional - for local use)
python download.shfrom train import train_model
from data import load_chess_dataset
# Load dataset
dataset = load_chess_dataset("data/chess_positions.csv")
# Train model
model = train_model(
dataset=dataset,
model_name="rook-clf-9m",
num_epochs=3,
batch_size=32
)from eval import evaluate_model
# Evaluate on benchmarks
results = evaluate_model(
model_path="checkpoints/rook-clf-9m",
benchmarks=["checkmate_in_one", "puzzles"]
)
print(f"Checkmate accuracy: {results['checkmate_in_one']}")from src.policy import ChessPolicy
from src.model import make_model, make_tokenizer
# Load model
model = make_model()
tokenizer = make_tokenizer()
policy = ChessPolicy(model, tokenizer)
# Predict move from FEN position
fen = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
move = policy.predict_move(fen)
print(f"Predicted move: {move}")Try ROOK-CLF-9M directly in your browser: https://jorahn.github.io/research/rook-clf-demo/
The demo features:
- Self-play Analysis: Interactive chess board with move predictions
- Benchmark Evaluation: Test against research datasets
- Attention Visualization: Explore model internals
rook/
├── data.py # Dataset loading and preprocessing
├── train.py # Training script
├── eval.py # Evaluation on benchmarks
├── download.sh # Download pretrained models
├── src/
│ ├── model.py # Model architecture definition
│ ├── policy.py # Inference and move generation
│ ├── const.py # Chess constants and move mappings
│ └── utils/ # Data conversion utilities
├── tests/ # Unit tests
└── checkpoints/ # Saved models
# Run all tests
pytest
# Run specific test
pytest tests/test_model.py
# Run with coverage
pytest --cov=src- RookWorld: Language model approach with chain-of-thought reasoning
- LAION Research: Full research blog post and findings
- HuggingFace Model: Pretrained weights and model card
If you use this work, please cite:
@article{rook2024,
title={ROOK: Strategic Reasoning in Chess Without Search},
author={Rahn, Jonathan and Jitsev, Jenia and Sun, Qi},
journal={LAION Research Notes},
year={2024},
url={https://laion.ai/notes/rook/}
}MIT License - see LICENSE file for details
- LAION for compute resources and collaboration
- Google DeepMind for the original research inspiration
- HuggingFace for the transformer infrastructure