PaperlayerView exact revision on arXiv

Source-linked AI summary

Conditional Generation of Creative Chess Puzzles with Diffusion Models

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

arXiv:2609.38577v1cs.AIcs.LG

TL;DR

Modern generative models can struggle with novel, constrained creative tasks, and chess puzzles provide a test that demands both creativity and reasoning. The paper uses masked diffusion with conditioning and auxiliary best-move prediction, then applies reinforcement learning; the yield of unique, theme-matching positions rises from 22.37% to 42.32%.

  • Problem

    Modern generative models can struggle to produce novel, highly constrained, counter-intuitive artifacts, while chess puzzles test creativity and reasoning because a single piece change can alter the puzzle.

    Source: §1 · pp. 1–2

  • Method

    The paper trains a masked diffusion model for conditional chess-puzzle generation and combines board generation with best-move prediction, then applies reinforcement learning.

    Source: §6 · p. 10

  • Results

    The yield of generated positions with a unique solution and correct themes rises from approximately 22.37% after supervised training to 42.32% with reinforcement learning.

    Source: §6 · pp. 10–11

  • Takeaways & Limitations

    The released open-weight models provide conditional chess-puzzle generation to the community.

    Source: §6 · p. 11

    Puzzles with similar themes can still express similar ideas despite having different main lines, and preventing model collapse remains difficult.

    Source: §6 · p. 11

Abstract

from arXiv · show

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.

1. Introduction

Chess puzzle generation tests whether models can create novel, constrained artifacts while balancing creativity and reasoning. The paper uses masked diffusion for controllable generation, adding best-move prediction and reinforcement learning to improve puzzle quality.

1. Introduction

  • Masked diffusion is chosen because it iteratively refines the whole sequence, enabling non-directional conditioning for the puzzle-generation task.The paper contrasts this with autoregressive generation, which builds sequences token by token.Source: §1 · p. 1
  • Chess puzzles test creativity and reasoning under tight constraints, since changing one piece can transform a puzzle entirely.The paper presents the task as a contained framework for studying controllability in creative generation.Source: §1 · pp. 1–2
  • RL raises the yield of positions with a unique solution and correct themes from 22.37% after supervised training to 42.32%.The model is trained on the Lichess puzzles dataset; the paper also reports that the best-move auxiliary task improves uniqueness and theme-conditioning accuracy.Source: §1 · p. 2
  • The model supports conditioning on tactical themes, partial board positions, and the best move, giving users control over puzzle generation.The paper describes practicing specific types of puzzles as a practical use.Source: §1 · p. 2
  • The authors release open-weight models for chess puzzle generation.The examples in Figure 1 were hand-picked from generated puzzles based on solution elegance.Source: §1 · p. 2

2. Related work

Related work situates chess puzzle generation within computational creativity and contrasts search-based and deep-learning approaches. It also reviews discrete diffusion and reinforcement-learning methods, motivating the paper’s use of masked diffusion and DDPO.

Machine intelligence and creativity

  • Computational creativity research addresses generative systems’ limits in producing novel, constrained artifacts, including through evolutionary search and expert evaluation.Related examples include semantic discovery, 3D-form co-evolution, and origami design.Source: §2 · pp. 2–3

Chess puzzle generation

  • Prior chess-puzzle generators use evolutionary or iterative search, while Feng et al. (2026) train a deep-learning model that beats the Lichess baseline on counter-intuitivity.The authors do not claim to surpass that work on puzzle-quality metrics; they emphasize conditional generation and open-source implementation instead.Source: §2 · p. 3

Masked diffusion models

  • Discrete diffusion globally refines sequences with bidirectional attention, supports variable inference-time efficiency, and can suit tasks without a fixed left-to-right order.The cited work also reports higher pretraining costs for comparable text quality.Source: §2 · p. 3

Reinforcement learning with diffusion models

  • Prior reinforcement-learning methods for diffusion models address difficult sequence-level likelihoods through step-wise evaluation or one-pass approximations.The authors report that ESPO did not yield a working configuration for their research, so they used DDPO, which performed better despite being less efficient in theory.Source: §2 · p. 3

3. Methods

The methods combine masked diffusion with supervised best-move prediction and an adapted DDPO objective, then define puzzle quality through legality, uniqueness, and counter-intuitiveness criteria.

3.1. Masked diffusion models

  • Masked diffusion corrupts FEN-tokenized boards into increasingly masked sequences, then a transformer denoiser iteratively unmasks tokens, with cross-attention incorporating theme conditions.The reverse process samples tokens from the network as it transitions to less-masked states.Source: §3.1 · pp. 3–4

3.2. Supervised training

  • Supervised training minimizes an ELBO on Lichess puzzles, estimated with one-point Monte Carlo; experiments compare joint best-move generation, puzzle-only generation, and block diffusion.Block diffusion keeps best-move tokens masked during board generation and samples them afterward; the selected checkpoint has the lowest validation loss at 1,000,000 steps.Source: §3.2 · p. 4

3.3. Denoising Diffusion Policy Optimization

  • DDPO treats reverse-diffusion transitions as a step-wise policy, with partially masked sequences as states and sequence-level rewards propagated through the trajectory.The objective adds entropy and KL terms; group-relative advantages are estimated and normalized group-wise.Source: §3.3 · pp. 4–5

3.4. Detecting chess puzzles

  • Puzzle evaluation requires legality, uniqueness, and counter-intuitiveness; uniqueness compares best- versus second-best-move winning chances, while counter-intuitiveness uses Stockfish discovery depth and captured-piece value.A position is counter-intuitive when its counter-intuitiveness score exceeds 0.1.Source: §3.4 · p. 5
  • The RL reward penalizes illegal positions and rewards legal positions only when uniqueness, diversity, piece-count, theme, and best-move filters pass.Diversity filtering checks similarity within and across batches and differences in opponent responses.Source: §3.4 · p. 5

4. Experimental evaluation

The experiments show that simultaneous best-move prediction improves supervised puzzle metrics, while reinforcement learning raises the yield of unique, theme-matching puzzles. Conditioning also supports targeted themes and partial-board prompts, although RL produces fewer sufficiently difficult puzzles.

4.1. Best move prediction improves puzzle metrics

  • Simultaneous best-move prediction yields an 11.6% higher uniqueness rate and a 2.5% higher theme-match rate among unique solutions than board-only generation.The highest puzzle rate and counter-intuitiveness rate do not identify a clear best supervised configuration because confidence intervals overlap.Source: §4.1 · pp. 5–6, Table 1 · p. 6
  • Supervised configurations closely preserve Lichess board and principal-variation distances, and best-move generation strategy has virtually no impact on these diversity metrics.Against Feng et al. (2026), only board distances are comparable because the principal-variation distance implementations differ.Source: §4.1 · pp. 6–7, Table 2 · p. 6

4.2. Reinforcement learning

  • Conditioning RL on sampled best-move tokens in v2 helps prevent policy collapse, enabling longer training than v1.The v2 run is initialized from the supervised checkpoint with simultaneous best-move prediction.Source: §4.2 · p. 7
  • The share of positions passing both uniqueness and theme criteria rises from 22.37% after supervised training to 42.32% after RL.The RL model produces more unique, theme-matching puzzles, even though fewer meet the counter-intuitiveness criterion.Source: §4.2 · p. 7
  • During v2 training, uniqueness rises from below 45% to above 65%, while the share passing the counter-intuitiveness threshold falls despite a slight increase in mean counter-intuitiveness.Illegal positions decrease from around 1.0% to around 0.1%.Source: §4.2 · p. 7

4.3. Manual control via conditioning

  • The RL model retains the full spectrum of tactical themes, including mate-in-one puzzles that are practically never counter-intuitive.The reward enforces requested-theme matching, and advantages are computed within groups of the same theme.Source: §4.3 · pp. 8–9
  • Masked diffusion supports conditioning on themes, partial board positions, and the best move, enabling targeted practice and varied completions from a shared starting position.In the Figure 4 example, the bishop sacrifice Bxh7+ is conditioned as the best move, while the generated continuations differ.Source: §4.3 · p. 9

5. Hierarchical generation

Hierarchical generation branches a partially denoised chess position into multiple independent continuations. This creates many related positions while allowing shared board elements to occur across the board rather than only in a sequential prefix.

5. Hierarchical generation

  • The method duplicates a partially unmasked state at selected denoising steps, then independently continues each branch while preserving tokens already unmasked.This produces different final piece configurations from a shared partial position.Source: §5 · p. 9
  • Six rounds of splitting and denoising produce 4,096 positions, with most puzzle diversity established early in generation.By step 128, about half the tokens are masked, but many puzzles already share similar principal variations.Source: §5 · p. 10
  • Unlike autoregressive branching, masked diffusion can distribute shared tokens anywhere on the board rather than restricting shared ancestry to a sequential prefix.The paper argues this lets hierarchical variations appear naturally random to human observers.Source: §5 · p. 10

6. Conclusions and discussion

The conclusion presents masked diffusion as a framework for controllable chess-puzzle generation, with best-move prediction and reinforcement learning improving puzzle quality. It also identifies model collapse as an unresolved challenge and suggests the approach may extend to other constrained creative-generation tasks.

Contributions

  • Joint best-move prediction improves generation quality, while masked diffusion supports conditioning on partial board positions and best moves as well as tactical themes.The authors suggest the auxiliary task may encourage richer internal representations of chess concepts.Source: §6 · p. 10

Generation process

  • Figure 5 depicts hierarchical generation for one theme set, splitting the process into four to form a tree of rook-endgame positions.The figure caption describes predetermined splits in the generative process.Source: Figure 5 · p. 11

Results

  • 42.32% of generated positions had a unique solution and matched requested themes after reinforcement learning, compared with 22.37% after supervised training.The authors report that reinforcement learning nearly doubled this yield while producing more unique and positionally diverse puzzles.Source: §6 · p. 11

Contributions

  • The authors release the first open-weight models for conditional chess-puzzle generation.They present this as a resource for the community.Source: §6 · p. 11

Limitations and future work

  • Preventing model collapse remains the most difficult future-work challenge; puzzles with similar themes can share ideas even when their main lines differ.The authors suggest diversity filtering or reward-formulation changes as possible approaches.Source: §6 · p. 11

Broader implications

  • The authors suggest that masked diffusion combined with auxiliary prediction could be a promising design pattern for other generation problems with strict structural constraints and controllable outputs.Argumentation is offered as an example, with rhetorical strategies as themes and a concluding claim as the best move.Source: §6 · p. 11

A. Reproducibility

The reproducibility appendix describes the model architecture, board and move representations, supervised-training setup, and reinforcement-learning configuration.

Model architecture

  • The diffusion denoising network uses a 268M-parameter transformer encoder, with cross-attention conditioning accounting for 68M parameters.The model has 16 transformer blocks and 8 attention heads per layer; the best-move model uses sequence length 81 rather than 76.Source: A. Reproducibility · p. 15

Representations

  • The model conditions on one-hot puzzle themes and represents boards with tokenized FEN, while the best-move model appends a UCI-format move.The tokenization distinguishes concepts such as a black bishop and black-to-move even though both use the character b in FEN.Source: A. Reproducibility · p. 15

Supervised training

  • Supervised training minimizes the negative evidence lower bound with AdamW, using a 1,000-step warmup and cosine annealing through 2,000,000 steps.Training uses a total batch size of 1024 positions on eight Nvidia H200 GPUs.Source: A. Reproducibility · pp. 15–16

Reinforcement learning

  • The reinforcement-learning phase uses AdamW at 3 · 10^-5 and generates trajectories with 64 denoising steps at temperature 1.0.Policy updates use six groups of size 16, with diversity and realism parameters β=0.03 and γ=0.03 and PPO clip parameter ε=0.2.Source: A. Reproducibility · p. 16

Reward and diversity criteria

  • Generated positions count as different when board distance is at least 6 or principal-variation distance is at least 1.The appendix treats even slightly differing best moves as different under principal-variation distance.Source: A. Reproducibility · p. 16

B. Model weights and training

The appendix provides access to the model weights and training code.

Access

  • The model weights are available on Hugging Face, and the training code is available on GitHub.The passage identifies these access locations without specifying a repository link.Source: B. Model weights and training · p. 16

C. Theme computation

The appendix describes how the authors identify objective, board-verifiable themes and determine whether generated puzzles satisfy requested theme conditions.

Theme identification

  • The Lichess puzzler pipeline identifies tactical themes by iterating over a generated position's principal variation.The authors group themes into state of game, endgame type, checkmate type and length, puzzle length, winning states, and other tactics.Source: C. Theme computation · p. 16

Evaluation scope

  • The evaluation uses objective, position-based themes and excludes metadata tags such as “master vs. master” that cannot be verified from the board state alone.Source: C. Theme computation · p. 16

Theme matching

  • Theme matching compares generated themes with requested themes using category-specific checks rather than requiring the generated set to be a strict superset.The complete matching logic is given in Algorithm 1.Source: C. Theme computation · p. 16

Theme analysis

  • Figure A1 concerns impurity-based theme importance and rating on the puzzle rate after supervised training.The caption states the plotted topics but does not report a result.Source: Figure A1 · p. 17

D. Conditioning on difficulty ratings

Difficulty-rating conditioning has little effect on final puzzle rate compared with theme conditioning. The authors attribute this to themes carrying difficulty information that the model can infer.

D. Conditioning on difficulty ratings

E. Reinforcement learning

The RL runs trade off counter-intuitiveness against controllability and diversity, with version two more usable in practice. Their reward and filtering choices differ, and version one’s training curves show collapse.

E. Reinforcement learning

  • Version one’s training curves show counter-intuitiveness rising as uniqueness and theme match generally fall; board distances decline after 3,000 steps and PV distance drops sharply after 4,500.The authors select the 3,500th checkpoint for comparison with version two.Source: Figure A2 · p. 18, E. Reinforcement learning · pp. 18–19
  • Version one omits best-move conditioning and the opponent-best-move diversity filter, while version two omits the low-entropy filter.Source: E. Reinforcement learning · p. 18
  • Version two samples themes from a more diverse set and randomizes some best-move tokens; version one uses larger KL and entropy coefficients.The coefficients are 0.1 and 0.2, respectively, and are not needed in version two due to its diversity-preserving best-move conditioning.Source: E. Reinforcement learning · p. 18
  • Version one has a higher counter-intuitive rate but lower PV distance and diversity, making it less usable than version two under the authors’ self-PV-distance criterion.Version one also has much higher board distances.Source: E. Reinforcement learning · p. 18

Related papers

Loading 2609.38577v1…