Conditional Generation of Creative Chess Puzzles with Diffusion Models

Aatu Selkee; Severi Rissanen; Xidong Feng; Tom Zahavy; Eric Malmi

cs.AISep 29, 2026

Conditional Generation of Creative Chess Puzzles with Diffusion Models

Authors: Aatu Selkee, Severi Rissanen, Xidong Feng, Tom Zahavy, Eric Malmi

Organizations: Aalto University · Google DeepMind

Abstract

While modern language models demonstrate impressive generative capabilities, they often struggle with constrained, counter-intuitive creative tasks. To address this limitation, we explore chess puzzle generation as a rigorous testbed for computational creativity and reasoning, a domain where altering a single piece can invalidate an entire solution. We propose a novel approach for conditional generation of creative chess puzzles using masked diffusion models. Unlike previous methods, our non-directional diffusion approach allows for conditioning on specific tactical themes and partial board positions. We introduce a novel auxiliary task of simultaneous best-move prediction, which improves solution uniqueness by 11.6% and theme-conditioning accuracy by 2.5%. To further optimize solution uniqueness and theme conditioning, we establish a reinforcement learning framework adapted from Denoising Diffusion Policy Optimization (DDPO). This RL training increases the yield of unique and theme-matching positions by 89.1%. Finally, we release the first open-weights models (Appendix B) for chess puzzle generation, offering a new pathway for controllable, creative generation.

Figures & tables

Figure 1 : Example chess puzzles generated by the proposed model. The themes are provided as input to the model. Solutions: (a) Rb5 cxb5 b7 (b) Nf1+ Rxf1 exf1=N+ Kg1 Nxd2

Training	Predict best move	Legal (%)	Unique (%)	Counter-Intuitive (%)	Puzzle (%)	Unique ∗ & Theme (%)	Theme \| Unique ∗ (%)
Supervised	No	98.97 $\pm$ 0.05	39.16 $\pm$ 0.23	1.00 $\pm$ 0.08	0.39 $\pm$ 0.03	19.10 $\pm$ 0.19	82.04 $\pm$ 0.39
Supervised	Block	98.99 $\pm$ 0.05	41.25 $\pm$ 0.24	0.99 $\pm$ 0.07	0.41 $\pm$ 0.03	20.77 $\pm$ 0.20	83.64 $\pm$ 0.36
Supervised	Yes	98.96 $\pm$ 0.05	43.69 $\pm$ 0.24	0.93 $\pm$ 0.07	0.41 $\pm$ 0.03	22.37 $\pm$ 0.20	84.05 $\pm$ 0.34
RL v1	Yes	99.76 $\pm$ 0.04	38.02 $\pm$ 0.43	2.23 $\pm$ 0.20	0.85 $\pm$ 0.08	16.40 $\pm$ 0.32	73.80 $\pm$ 0.82
RL v2	Yes	99.94 $\pm$ 0.01	66.15 $\pm$ 0.20	0.36 $\pm$ 0.03	0.23 $\pm$ 0.02	42.32 $\pm$ 0.22	87.21 $\pm$ 0.21
Supervised ( Feng et al., 2026 )		99.72	30.89	1.11	0.34	—	—

Table 1 : Puzzle quality metrics for different model configurations. All models were sampled with a temperature of 1.0 and 256 denoising steps. Errors represent 95% confidence intervals. Counter-intuitiveness is computed over positions with unique solutions.

Training	Predict best move	Board (Self)	Board (Lichess)	PV (Self)	PV (Lichess)
Supervised	No	11.63 $\pm$ 0.01	11.16 $\pm$ 0.01	0.78 $\pm$ 0.002	0.93 $\pm$ 0.001
Supervised	Block	11.67 $\pm$ 0.01	11.18 $\pm$ 0.01	0.80 $\pm$ 0.002	0.93 $\pm$ 0.001
Supervised	Yes	11.70 $\pm$ 0.01	11.19 $\pm$ 0.01	0.81 $\pm$ 0.002	0.93 $\pm$ 0.001
RL v1	Yes	15.15 $\pm$ 0.01	14.34 $\pm$ 0.01	0.76 $\pm$ 0.002	0.97 $\pm$ 0.001
RL v2	Yes	12.12 $\pm$ 0.01	12.10 $\pm$ 0.01	0.88 $\pm$ 0.002	0.97 $\pm$ 0.001
Supervised ( Feng et al., 2026 )		10.859	8.528	(0.837)	(0.637)

Table 2 : Self-distance and Lichess-distance metrics after supervised training and RL. Metrics are computed over positions passing the uniqueness filter. We align the number of positions with Feng et al. (2026) and use 100k samples for Lichess-distances and 40k samples for the self-distances. Despite this, only board distances are comparable to Feng et al. (2026) due to differing implementations of the PV distance.

Panel of Figure 2 — Figure 2 : Training progress of run v2 in Table 1 during reinforcement learning.

Figure 3 : The distribution of themes for the Lichess puzzles dataset, our supervised model and the v1 RL model. All positions are first filtered for positions with unique solutions and the percentages are computed across the filtered dataset.

Figure 4 : Three positions generated from a common starting point. The only theme sampling was conditioned with is sacrifice . The generation was also conditioned on the bishop sacrifice Bxh7+ and that it is white’s move next.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Parameter	value
Model
Parameter count	268M
Transformer block count	16
Attention head count	8
Embedding dimension	1024
Activation	SwiGLU (widening factor 2.66)

Appendix

Table A1 : All hyperparameters of our training setup.

Category	Themes
State-of-game	Opening, Middlegame, Endgame
Type-of-endgame	Pawn endgame, Bishop endgame, Knight endgame, Rook endgame, Queen endgame, Queen rook endgame
Type-of-checkmate	Mate, Back rank mate, Boden mate, Smothered mate, Hook mate, Double bishop mate, Arabian mate, Dovetail mate, Anastasia mate
Length-of-checkmate	Mate in 1, Mate in 2, Mate in 3, Mate in 4, Mate in 5
Length-of-puzzle	One move, Short, Long, Very long
Winning	Crushing, Advantage

Appendix

Table A2 : The complete list of themes supported by the Lichess puzzler pipeline, grouped by their respective categories.

Figure A1 : The impurity based importance of themes and rating on the puzzle rate after supervised training.

Panel of Figure 10 — Figure A2 : Training progress of run v1 in Table 1 during reinforcement learning.

Explore similar work

Generalization or Memorization? Brittleness Testing for Chess-Trained Language Models
May 17, 2026Ethan TangChess Model Checkpoints
School of Computing and Augmented Intelligence · Arizona State University · Tempe, AZ 85281
RePAIR: Predictive Self-Supervised Representation Learning in Chess
Jun 10, 2026Christoph Koller, Johannes Fürnkranz, Timo BertramChess Self-Supervised Representations
JKU Linz, Austria
Benchmarking Prompt Optimization of Large Language Models With Chess
Sep 30, 2026Timothée Lesort, Alejandra López de Aberasturi Gómez, Tristan Karch +4Chess Large Language Model Benchmarks
imec AI.labs

May 17, 2026, cs.AIJ/K move · Enter open · S save

Generalization or Memorization? Brittleness Testing for Chess-Trained Language Models

Ethan Tang

School of Computing and Augmented Intelligence · Arizona State University · Tempe, AZ 85281

Chess Model Checkpoints Memorization Fragility Synthetic Data

Open paper AskPDF arXiv

Recent work has fine-tuned language models on chess data and reported high benchmark scores as evidence that the resulting models can understand the rules of chess, play full chess games at a professional level, or generate human-readable explanations grounded in expert knowledge. We train KinGPT, a 25M-parameter character-level language model trained only on (position, best-move) pairs, who exceeds 3B-parameter ChessGPT on a 600-puzzle mate-in-N suite and 4B-parameter C1-4B over a 20-theme puzzle benchmark. We examine several claims made in existing literature regarding chess-trained language models and assert that their impressive benchmark performance is largely explained by pattern-matching. We also demonstrate how LLM-Modulo, a verifier-in-the-loop framework, raises RedPajama 3B's best move accuracy from 1.2% to 21.2% and move generation validity from 19.3% to 95.3% on mate-in-N chess puzzles, comparable to gains achieved from ChessGPT's fine-tuning on chess-specific web corpora at a fraction of the cost. Our results illustrate how pairing a general LLM with an external verifier offers a more flexible alternative to directly training on synthetic data for well-defined domains. We open source all training/evaluation code, datasets, puzzle samples, and KinGPT model checkpoints for reproducibility.

Abstract

Explore similar work

Generalization or Memorization? Brittleness Testing for Chess-Trained Language Models

RePAIR: Predictive Self-Supervised Representation Learning in Chess

Benchmarking Prompt Optimization of Large Language Models With Chess