PaperlayerView exact revision on arXiv

Source-linked AI summary

Think Before You Paint: Recursive Latent Reasoning for Diffusion Models

Paweł Skierś, Małgorzata Grzanka, Wojciech Masarczyk, Jan-Willem van de Meent, Kamil Deja

arXiv:2610.09876v1cs.AIcs.LG

TL;DR

Diffusion models can depict rule-governed puzzles but struggle to satisfy their constraints directly in pixels. PaTh uses a recursive Thinker to steer a frozen diffusion Painter, improving results across rule-governed tasks, especially on harder instances.

  • Problem

    Diffusion models can struggle to satisfy rules over pixels, even when symbolic representations let diffusion or recursive methods solve hard puzzles.

    Source: §1 · pp. 1–2

  • Method

    PaTh couples a recursive Thinker that refines a latent state within denoising steps to a frozen diffusion Painter through ControlNet adapters, using reconstruction loss without symbolic targets.

    Source: §1 · p. 2, §3.2 · p. 4

  • Results

    Across four rule-governed data families, PaTh improves over diffusion and image-editing baselines; it exceeds the previous best by over 67 p.p. on MNIST Sudoku EXTREME.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results show that reasoning mechanisms developed for symbolic data can be integrated into pixel-space diffusion without symbolic supervision.

    Source: §5 · p. 9

    PaTh requires spatially aligned conditioning and does not recover spatial structure from an unanchored symbolic specification.

    Source: L LIMITATIONS · p. 22

Abstract

from arXiv · show

Diffusion models generate realistic images but often fail on visual reasoning tasks, such as filling in a Sudoku or drawing the path through a maze. When a discrete symbolic representation is available, recursive methods such as the Tiny Recursive Model (TRM) solve even hard instances of these puzzles. We ask how such reasoning can be carried over to pixels, where no symbolic representation is available. We propose Painter-Thinker (PaTh): a small recursive network (the Thinker) reasons over a grid of learned tokens that encode the noisy image and the conditioning, refines a latent state within every denoising step, and steers a frozen diffusion model (the Painter) through ControlNet adapters. The Thinker is trained with the standard reconstruction loss alone, without symbolic targets, a solver, or a verifier. PaTh solves 92.5% of hard MNIST Sudoku puzzles (prior best 75%) and 71.2% of extreme ones (prior best 4.1%), with 10M parameters against 82M for a standard diffusion model. It also improves on mazes, Queens, and CLEVR scenes with specified spatial relations, and its advantage grows with problem size. Diagnostic experiments show that PaTh recovers from injected mistakes that the diffusion model cannot repair, especially when many cells are wrong. Together, these results show that reasoning mechanisms developed for symbolic data can be integrated into pixel-space diffusion without symbolic supervision, opening a path toward generating data under increasingly complex constraints.

1 INTRODUCTION

Diffusion models can satisfy rules over symbolic representations, but visual puzzles remain difficult when those rules must hold over pixels. PaTh separates reasoning from image generation, using a recursive Thinker to steer a frozen Painter and improving results across rule-governed tasks.

1 INTRODUCTION

  • Pixel-space rules are harder to satisfy than rules over symbols because violations span continuous signals depicting units rather than directly manipulated cells, atoms, or bonds.The same difficulty appears in colliding or misplaced vehicles, non-divergence-free velocity fields, and jointly invalid puzzles.Source: §1 · p. 1
  • Vanilla diffusion solves only 8.3% of hard MNIST Sudoku instances despite producing sharp, diverse digits, and adding capacity alone does not improve reasoning.A recursive architecture used as the denoiser yields only a marginal improvement, motivating separation of reasoning and rendering.Source: §1 · pp. 1–2
  • PaTh couples a small recursive Thinker over learned tokens to a frozen diffusion Painter through ControlNet adapters, training both modules with reconstruction loss alone.No solver, verifier, or symbolic targets are used.Source: §1 · p. 2
  • PaTh improves over the previous best by over 67 p.p. on MNIST Sudoku EXTREME and surpasses diffusion models with 700× its capacity on AMAZE.On complex 3D reasoning, a 20M-parameter Thinker directs a 100M-parameter Painter.Source: §1 · p. 2
  • Diagnostics show that PaTh recovers from partially generated samples with many wrong cells, whereas diffusion fails to indicate what to trust when too much is wrong.The Thinker’s latent state supports iterative refinement before the image is committed.Source: §1 · p. 2
  • Across four families of rule-governed data, PaTh improves over standard diffusion, sampling-adaptive diffusion, and large image-editing models, with margins widening as instances grow harder.The diagnostics also find that a Thinker buys more than simply adding parameters.Source: §1 · p. 2

2 RELATED WORK

Related work addresses rule-governed generation through trajectory changes, inference-time search or guidance, symbolic diffusion, and recursive computation. PaTh differs by placing recurrence in a latent state before the result enters the image.

Rule-governed generation with diffusion

  • Pixel-space diffusion methods restructure denoising or supply external reasoning, while SRM reallocates steps that still bear commitments along the trajectory.Other approaches optimize learned energies, use pretrained vision-language models, or delegate plans to multimodal language models.Source: §2 · p. 3

Inference-time compute and guidance

  • Inference-time search and guidance revise samples in data space, requiring either an external verifier or a differentiable constraint; PaTh revises before anything enters the sample.Search approaches include noise search, sequential Monte Carlo, evolutionary and tree search, and resampling.Source: §2 · p. 3

Discrete diffusion for reasoning

  • Discrete diffusion reasons over tokens where masked entries are explicitly undecided, unlike continuous diffusion states that are complete partial samples with no pending marker.The paper identifies this representational difference as the gap its method addresses.Source: §2 · p. 3

Intermediate computation and recursive reasoners

  • Recursive reasoners use weight-tied depth to extend serial computation, and prior work reports strong symbolic-puzzle results and gains from recurrence in tokenized flows.PaTh targets exactly checkable rules in pixel space rather than symbolic puzzles.Source: §2 · p. 3

3 METHOD

The method combines diffusion’s noisy-sample trajectory with a recursive reasoner whose latent state can be revised before an image update. The Thinker’s spatial output conditions a separately trained, fixed Painter, with an efficient variant reducing repeated computation later in sampling.

3.1 PRELIMINARIES

  • The diffusion model predicts the clean sample directly from the noisy state, timestep, and embedded conditioning, and is trained with a reconstruction loss.Each trajectory state is a partially denoised sample.Source: §3.1 · p. 3
  • A recursive reasoner repeatedly updates an answer y and latent state z against fixed conditioning, with effective depth Nsup · H · (L + 1) applications of one shared network.Only y is decoded; z is not itself a solution and can hold information that supports the next update.Source: §3.1 · p. 4

3.2 PAINTER AND THINKER

  • Within each denoising step, the Thinker refines a latent candidate from xt and the conditioning, then returns a signal that conditions the Painter’s prediction of xt−1.Spatial conditioning is encoded convolutionally with xt, while non-spatial conditioning is embedded by an MLP and appended to the input.Source: §3.2 · p. 4, Figure 2 · p. 4
  • Only the spatial tokens of y are decoded and passed through ControlNet adapters into the Painter’s UNet skip connections; non-spatial tokens inform recursion but are discarded at readout.The Painter is trained separately and held fixed.Source: §3.2 · p. 4, Figure 2 · p. 4
  • The Thinker’s latent states are refined within a denoising step and reset before the next, unlike committed partial-image states; one step reaches effective depth 672 in the reported configuration.Intermediate latent states need not be plausible images or improve monotonically over their predecessors.Source: §3.2 · p. 4
  • An efficient variant carries Thinker state across denoising steps, concentrates recursion early, then stops calling the Thinker halfway and reuses its last output.The method uses a UNet Painter, ControlNet adapters, and separately trained Painter and Thinker; it also adapts to DiT.Source: §3.2 · pp. 4–5

4 EXPERIMENTS

Across Sudoku, AMAZE, and CLEVR, PaTh outperforms diffusion baselines on rule-governed generation, with advantages that persist or grow as tasks become more complex. Diagnostics show that its latent reasoning can revise errors before committing to an image, while recursive computation—not extra parameters or denoising steps alone—drives its gains.

4.1 MAIN RESULTS

  • PaTh reaches 92.5% puzzle accuracy on HARD and 71.2% on EXTREME MNIST Sudoku, exceeding every baseline on both setups.EXTREME puzzles have a unique solution, and one misplaced digit invalidates the grid.Source: §4.1 · p. 5, Table 1 · p. 5
  • PaTh solves 75.1% of mazes exactly, versus 43.4% for fine-tuned 14B Bagel and 19.0% for the DM, and solves every Queens board.On Queens, the DM solves 90.9% and fine-tuned Bagel 13.7%.Source: §4.1 · pp. 5–6, Table 2 · p. 6
  • On AMAZE, PaTh's advantage over the diffusion baseline grows with puzzle size; on circular mazes, both models degrade sharply and PaTh's edge lasts only through size 11.On hexagon mazes, the baseline falls below 50% after size 11 while PaTh stays above 50% at the largest size; the pattern also occurs on square and triangle mazes and Queens.Source: Figure 3 · p. 5, §4.1 · p. 6
  • CLEVR spatial accuracy rises from 62.3% for the standard DM to 89.4% for PaTh, while a larger 238M DM reaches 68.8%.PaTh's advantage grows with scene complexity, reaching 85.2% against 62.0% at ten objects.Source: §4.1 · p. 6, Table 3 · p. 6

4.2 DIAGNOSING THE MECHANISM

  • For Sudoku, both models begin near 27 violated constraints, but the DM plateaus around 10.6 while PaTh falls below 1.5.PaTh reaches 6.3 violations within the first 5% of sampling.Source: §4.2 · p. 7, Figure 5 · p. 7
  • After injected errors, PaTh maintains high recovery across early corruption levels, whereas the DM works only in narrow sampling windows and degrades as mistakes accumulate.The DM retains indirect mistakes 15–20+ points more often than direct ones early in sampling; PaTh's difference is 3.2 points on Sudoku and zero or negative on mazes.Source: §4.2 · pp. 7–8, Figure 8 · p. 8
  • A probe shows that in up to a fifth of samples the Thinker considers candidates with substantially more constraint violations before later improving them.Figure 11 illustrates successive recursions introducing violations and resolving them before anything is committed to the next sample.Source: §4.2 · pp. 8–9, Figure 11 · p. 9

4.3 INFERENCE COMPUTE

  • Across inference-FLOP budgets, diffusion models remain below 5% Sudoku EXTREME accuracy while PaTh exceeds 60%; an efficient variant reaches 47% at 10^12 FLOPs.A trained PaTh evaluated with a single recursion is near zero, indicating that recursion—not simply adding the module—drives the gain.Source: §4.3 · p. 9, Figure 10 · p. 9

5 CONCLUSION

PaTh integrates latent recursion with a frozen diffusion model to improve rule-governed generation from pixels, with gains across four families and the hardest instances. The appendix specifies the evaluated datasets and shows conditioning–solution examples.

5 CONCLUSION

  • PaTh couples a recursive reasoner to a frozen diffusion model and improves generation across four families, especially on the hardest instances, using reconstruction loss alone.The conclusion attributes the gains to the Thinker’s latent state and recursion.Source: §5 · p. 9

A DATASETS

  • MNIST Sudoku asks the model to fill missing cells in a 9×9 digit grid, with HARD puzzles typically allowing multiple solutions and EXTREME puzzles having roughly 18 clues and a unique solution.Puzzle accuracy is the reported metric.Source: A DATASETS · p. 13
  • AMAZE tests drawing maze paths and placing Queens subject to row, column, region, and adjacency constraints across boards of varying sizes.The reported measures include coverage, violation, PASS, and EXACT.Source: A DATASETS · p. 13

A DATASETS / A.1 DATASETS - VISUAL EXAMPLES

  • CLEVR scene composition tests whether generated scenes satisfy specified object locations, identities, and pairwise spatial relations that must compose transitively.The examples show conditioning and ground-truth solutions for each dataset.Source: A DATASETS · p. 13, §A.1 · p. 13, Figure 12 · p. 13

B SCENE COMPOSITION EVALUATION METHOD

The CLEVR evaluation compares generated images with their conditioning scene graphs using a shared detection-and-attribute pipeline. The appendix details its calibration, matching, metrics, and model-training setup.

B SCENE COMPOSITION EVALUATION METHOD

  • Generated images are evaluated against their conditioning scene graphs with the same open-vocabulary detect-and-attribute pipeline for every model.The pipeline calibrates pixel-space boxes to 3-D coordinates, detects objects, classifies attributes, and matches detections to ground truth.Source: B SCENE COMPOSITION EVALUATION METHOD · p. 14
  • The evaluation reports attribute agreement and binding, relational accuracy, and precision and recall as measures of conditioning adherence.Counts are accumulated over the evaluation set before ratios are calculated; Table 4 reports performance on CLEVR ground-truth training data.Source: B SCENE COMPOSITION EVALUATION METHOD · p. 15

C IMPLEMENTATION DETAILS

  • The Painter is trained first and frozen; the Thinker and condition encoder are then trained with reconstruction loss alone.The Painter operates in CLEVR’s pretrained SDXL VAE latent space and in pixel space elsewhere; parameter counts are listed by dataset.Source: C IMPLEMENTATION DETAILS · p. 15, Table 5 · p. 15

D ADDITIONAL RESULTS - GEOSOLVERS DATASETS

On the largest Simple Polygon instances, PaTh improves both validity and optimality rate; it is comparable to the diffusion model on Steiner Tree and Inscribed Square. The reported pattern differs between discrete constraint satisfaction and continuous optimization tasks.

D ADDITIONAL RESULTS - GEOSOLVERS DATASETS

E STABILITY OF THE METHOD

The method is stable across seeds and does not require the Thinker’s token grid to align with Sudoku cells. The appendix also reports transfer to a DiT Painter and reduced training cost from halting recursion.

E STABILITY OF THE METHOD

  • EXTREME Sudoku accuracy averaged 70.6 ± 0.7 across three seeds; token grids 10 × 10 and 12 × 12 scored 66.8 and 69.3 versus 71.2 for aligned 9 × 9.The results indicate no required correspondence between Thinker-grid resolution and puzzle structure.Source: E STABILITY OF THE METHOD · p. 17, Table 10 · p. 17

F ADOPTION TO DIT

  • With a DiT Painter, PaTh reaches 0.542 puzzle accuracy against 0.029 for the diffusion baseline, below the 0.712 obtained with a UNet Painter.The authors attribute the gap to the coupling and report that they did not tune it further.Source: F ADOPTION TO DIT · p. 17, Table 11 · p. 18

G HALTING HEAD IMPACT

  • Halting cuts Thinker training from 72h to 26h and total training from 85h to 39h by skipping predicted-unhelpful recursive supervision, not by increasing iteration throughput.The iteration rate is 0.85 iterations/s without halting and 0.81 with halting.Source: G HALTING HEAD IMPACT · p. 18, Table 12 · p. 18

H JOINT VS SEPARATE TRAINING

Separate Painter–Thinker training takes less time and is slightly more accurate than joint training, while producing more diverse digits. The ControlNet interface substantially outperforms concatenating the Thinker output with the diffusion state.

H JOINT VS SEPARATE TRAINING

I CONTROLNET ALTERNATIVE

  • 37.5% versus 71.2% puzzle accuracy on MNIST Sudoku EXTREME: concatenation performs considerably worse than ControlNet coupling.The concatenation variant uses the Thinker output as an additional input channel alongside the diffusion state.Source: I CONTROLNET ALTERNATIVE · p. 19

J RECOVERING MISTAKES: FULL METRICS

The recovery analysis distinguishes restoring corrupted cells from restoring them while keeping the rest of the grid correct. It reports cell recovery and kept rates across Sudoku corruption types and constrained-scene object swaps.

Sudoku corruption metrics

Recovery criteria

  • 58.8% corrupted-cell recovery but only 7.0% valid grids for diffusion, versus 97.7% recovery and 74.7% valid grids for PaTh.These results are for 32 locally visible corruptions injected at 30% of sampling; recovery alone does not require the rest of the grid to remain correct.Source: J RECOVERING MISTAKES: FULL METRICS · pp. 19–20, Table 14 · p. 20

Constrained-scene metrics

K THINKER PROBE TRAINING DETAILS

The probe is trained to read the Thinker’s per-cell latent-derived answer at intermediate reasoning points. Its labels come from a standard MNIST classifier’s reading of each trajectory’s final image.

Probe design

  • The linear probe maps a single per-cell y vector to a nine-way digit distribution, avoiding early image blurring by the Painter.The probe is defined on y because z does not itself decode to a solution.Source: K THINKER PROBE TRAINING DETAILS · p. 21

Training

  • Cross-entropy trains the probe on y recorded at every recursion, using a standard MNIST classifier’s cell-by-cell reading of that trajectory’s final image as the label.The same final-image label is used regardless of when the state was captured.Source: K THINKER PROBE TRAINING DETAILS · p. 21

Validation

  • From denoising step t=5 onward, probe accuracy exceeds 95% and remains essentially unchanged with further reasoning.At t=0, accuracy rises from 66.2% to 89.1% with reasoning depth; accuracy is also lower at t=19, when little reasoning remains.Source: Table 20 · p. 21, K THINKER PROBE TRAINING DETAILS · p. 21

L LIMITATIONS

PaTh depends on spatially aligned conditioning, while its diagnostics show why irreversible local commitments can fail on unique-solution puzzles. The paper also reports compute reductions, full-resolution results, and exploration traces in a fine-tuned diffusion model.

L LIMITATIONS

  • Unaligned conditioning causes Sudoku accuracy to collapse to chance, and unanchored symbolic relations do not supply the spatial structure PaTh needs.Randomly translating the target while leaving the conditioning unchanged yields 11.1% cell accuracy; applying the same transform to conditioning restores partial performance, while increasing tokens from 81 to 144 does not help. In constrained scenes, only conditioning that localizes objects raises both spatial accuracy and attribute binding.Source: Table 21 · p. 22, L LIMITATIONS · p. 22, Table 22 · p. 22

M SRM ON EXTREME

  • Below 50% of samples remain consistent with the unique solution after SRM places only its first cell, although board violations stay near zero.The early commitments can be locally consistent yet globally wrong, and the sampler cannot revisit a commitment once made.Source: M SRM ON EXTREME · pp. 22–23, Figure 17 · p. 23

N EFFICIENT ALLOCATION OF THE RECURSION BUDGET

  • Carrying the Thinker state across denoising steps helps at small reasoning budgets, and concentrating reasoning early can reduce compute without lowering accuracy.At one recursion, carried state gives 56.1% accuracy versus 32.8% when reset every step. Reasoning for the first half reaches 72.5% versus 71.2% for reasoning throughout, at half the cost; late allocation can cost up to 11.3 points.Source: N EFFICIENT ALLOCATION OF THE… · p. 23, Figure 18 · p. 23, Figure 20 · p. 24

O FULL RESOLUTION MNIST SUDOKU RESULTS

P EXPLORATION TRACES IN FINE-TUNED BAGEL

Related papers

Loading 2610.09876v1…