Think Before You Paint: Recursive Latent Reasoning for Diffusion Models

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

cs.AIOct 7, 2026

Think Before You Paint: Recursive Latent Reasoning for Diffusion Models

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

Organizations: Warsaw University of Technology · IDEAS Research Institute · University of Amsterdam

Abstract

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.

Figures & tables

Figure 1: Diffusion models are not limited by scale but by where they spend compute. Top: a standard diffusion model (DM) spends its entire inference budget on denoising ( ) and commits to errors early, leaving an invalid sudoku ( ✗ ). Bottom: PaTh (Painter–Thinker) pairs a smaller denoiser, the Painter , with a recursive Thinker ( , loop) that refines a draft of the solution before each denoising step, with most thinking at high noise, and solves the puzzle ( ✓ ). Both models have the same total inference budget. On MNIST Sudoku extreme , PaTh solves 71.2% of puzzles versus 2.9% for the vanilla DM (Sec. 4 ).

Figure 2: PaTh architecture. Spatial conditioning is concatenated with

\mathbf{x}_{t}

and encoded convolutionally; non-spatial conditioning is embedded by an MLP. Together they form the Thinker’s input

\mathbf{c}

, over which the Thinker recurses on

(\mathbf{y},\mathbf{z})

within a single denoising step. Only the spatial tokens of

\mathbf{y}

are decoded and passed to the Painter through ControlNet adapters; the non-spatial tokens (gray squares) shape the recursion but are not read out. The Painter is trained separately and held fixed.

Model	Hard	Extreme
DM	0.083	0.029
TRM	-	0.083
SRM	0.516	0.032
IPR	0.750	0.041
PaTh	0.925	0.712

Table 1: Puzzle accuracy on MNIST Sudoku. While hard puzzles usually admit many solutions extreme has exactly one. PaTh significantly outperforms other methods on both setups.

Figure 3: On Amaze benchmark, PaTh has an increasing advantage over standard DM as the complexity of the puzzles increases - x-axis: puzzle size, y-axis: pass@1; for both datasets

proprietary models
Model	Params	Continuous (Maze) Task						Discrete (Queen) Task
Model	Params	Violation $\downarrow$	Coverage $\uparrow$	Pass@1 $\uparrow$	Pass@5 $\uparrow$	Exact@1 $\uparrow$	Exact@5 $\uparrow$	Violation $\downarrow$	Coverage $\uparrow$	Pass@1 $\uparrow$	Pass@5 $\uparrow$	Exact@1 $\uparrow$	Exact@5 $\uparrow$
GPT-image-1	N/A	62.88	58.97	5.40	6.06	–	–	62.91	37.09	0.00	2.28	–	–
NanoBanana-Pro	N/A	47.76	64.21	4.82	9.28	–	–	32.56	67.43	30.35	35.58	–	–
Seedream-4.5	N/A	16.90	25.67	2.14	3.21	–	–	76.86	23.14	2.86	2.86	–	–
open-source models (w/o chain-of-thought reasoning)

Table 2: Results on the AMAZE benchmark’s Maze and Queens tasks. PaTh is able to outperform both standard DMs trained on the same data and the image editing models evaluated on the benchmark, including their fine-tuned (FT) versions and proprietary systems.

Method	Recall	Prec.	Attr.	Spatial
DM (101M)	96.5	97.5	54.6	62.3
DM (124M)	96.5	97.5	56.1	64.3
DM (236M)	96.9	97.5	56.3	68.8
PaTh	95.9	97.5	53.8	89.4

Table 3: Constrained scene composition on CLEVR (%). PaTh substantially improves relational consistency while maintaining object and attribute metrics.

Figure 6: Qualitative comparison on mazes. When the path is directed into a dead-end, only PaTh is able to correctly revise it.

Panel of Figure 8 — Figure 8: Recovery rate after corrupting a valid solution and re-noising it to a given point of the trajectory, on extreme MNIST Sudoku (left) and constrained scenes (right). The diffusion model recovers only within a narrow band, at the start of sampling on CLEVR and at $50\%$ on Sudoku, and degrades sharply as mistakes accumulate. PaTh recovers throughout the region where change remains possible, declining only as less trajectory is left.

Figure 9: Rethink rate (fraction of samples) and magnitude (violations) during sampling.

Figure 10: Puzzle accuracy on Sudoku extreme against inference FLOPs. DM stay below

5\%

at every budget, whether the compute is spent on denoising steps or parameters, while PaTh exceeds

60\%

Figure 11: Latent reasoning within one denoising step, each recursion’s state $\mathbf{z}_{i}$ decoded by a linear probe with violated constraints in orange. The Thinker introduces violations absent from $\mathbf{x}_{t}$ and then resolves them, reaching a valid grid before anything is committed to $\mathbf{x}_{t-1}$ .

Appendix figures & tables30 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 12: MNIST Sudoku puzzle examples.

Figure 13: Mazes puzzle examples for all shapes.

Figure 14: N-Queens puzzle examples for different sizes.

Figure 15: Constrained scene generation puzzle examples. Object tokens visualized for readability.

General adherence		Attribute matching					Spatial relations accuracy
recall	precision	color	shape	material	size	all	Spatial relations accuracy
99.9	99.8	95.2	96.9	96.4	100	89.5	97.2

Appendix

Table 4: Performance of the evaluation method on CLEVR ground truth training data.

Dataset	Painter	Thinker	PaTh	DM
MNIST Sudoku	1.5M	8.3M	10.3M	81.6M
Mazes	1.5M	8.3M	10.3M	62.7M
Queens	1.5M	8.3M	10.3M	62.7M
Constrained scenes	101.9M	19.7M	121.6M	236.0M

Appendix

Table 5: Parameter counts by dataset. PaTh is the sum of its Painter, Thinker, conditioning encoders, and ControlNet adapters; DM is the diffusion baseline

Painter
Parameter	MNIST Sudoku	CLEVR	Mazes	N-Queens
VAE	-	stabilityai/sdxl-vae	-	-
Block types	Block2D $\times 3$	CrossAttBlock $\times 3$	Block2D $\times 3$	Block2D $\times 3$
Block out channels	32, 64, 64	128, 256, 512	32, 64, 64	32, 64, 64
Cross att dim	-	512	-	-
Layers per block	2	2	2	2

Appendix

Table 6: Hyperparameters across Components and Tasks

Model	7-12			13-15
Model	valid	ratio	opt rate	valid	ratio	opt rate
PaTh	0.985	0.998	0.783	0.869	1.000	0.264
visual geosolvers	0.953	0.988	0.574	0.620	0.962	0.062

Appendix

Table 7: Simple Polygon, by number of vertices. valid is the fraction of outputs forming a simple polygon, ratio their quality relative to the reference, and opt rate the fraction attaining the optimum.

Model	10-20		21-30		31-40		41-50
Model	valid	ratio	valid	ratio	valid	ratio	valid	ratio
PaTh	0.997	1.0008	0.983	1.0018	0.796	1.0040	0.340	1.0081
visual geosolvers	0.996	1.0008	0.986	1.0018	0.834	1.0044	0.334	1.0092

Appendix

Table 8: Steiner Tree, by number of terminals. Lower ratio is better. The two models are comparable throughout; neither solves the largest instances.

Model	Squareness	Alignment
PaTh	0.874	-0.84
visual geosolvers	0.891	-0.90

Appendix

Table 9: Inscribed Square. Neither model has a consistent advantage.

TRM num tokens	8x8	9x9	10x10	12x12
puzzle acc	59.8	71.2	66.8	69.3

Appendix

Table 10: Puzzle accuracy on extreme MNIST Sudoku for different Thinker token grids. Only

9\times 9

aligns with the Sudoku cells.

Painter architecture	Puzzle accuracy
UNet	0.712
DiT	0.542
DM	0.029

Appendix

Table 11: Puzzle accuracy on extreme MNIST Sudoku with different Painter architectures. DM is a diffusion model without a Thinker.

Configuration	iter/s	Training time
Separate, Painter	5.8	13h
Separate, Thinker	0.85	72h
Separate, total	–	85h
Separate, Thinker with halting	0.81	26h
Separate, total with halting	–	39h

Appendix

Table 12: Training cost on MNIST Sudoku. Total training times sum the Painter and Thinker stages; iteration rates are not defined for the combined two-stage configurations.

Training configuration	Training time	Puzzle accuracy	Digit diversity $\uparrow$
Joint	113h	69.9%	0.721
Separate	85h	71.2%	0.855

Appendix

Table 13: Comparison of joint and separate Painter–Thinker training on MNIST Sudoku. Separate training is less costly, slightly more accurate and produces more diverse digits in this setting. Digit diversity is the within-class mean pairwise

\ell_{2}

distance between unit-normalized samples, averaged over digit classes.

Figure 16: Examples generated by jointly and separately trained PaTh models. Separately trained models exhibit greater diversity in the generated MNIST digits.

	PaTh							DM
% sampling	$n{=}1$	$n{=}2$	$n{=}4$	$n{=}8$	$n{=}16$	$n{=}32$	$n{=}64$	$n{=}1$	$n{=}2$	$n{=}4$	$n{=}8$	$n{=}16$	$n{=}32$	$n{=}64$
10	94.5	96.5	96.1	94.3	94.5	94.2	86.8	49.2	50.0	49.6	52.2	50.5	49.0	44.7
30	96.9	96.1	96.5	98.1	96.0	97.7	64.8	85.9	88.3	86.9	85.0	77.6	58.8	42.4
50	89.1	95.3	91.2	90.1	90.4	90.6	45.9	97.7	96.9	96.5	95.8	92.0	74.2	36.8
70	42.2	46.5	43.9	40.9	39.2	35.4	15.9	77.3	72.7	77.9	70.6	57.7	32.3	12.1
90	0.0	0.4	0.4	0.3	0.3	0.3	0.3	10.9	7.0	8.2	6.9	6.3	3.6	1.9

Appendix

Table 14: Recovery rate (%) for locally visible corruptions on extreme MNIST Sudoku, by percentage of sampling completed when the mistakes are injected and number

n

of corrupted cells. This does not require the rest of the grid to remain correct.

	PaTh							DM
% sampling	$n{=}1$	$n{=}2$	$n{=}4$	$n{=}8$	$n{=}16$	$n{=}32$	$n{=}64$	$n{=}1$	$n{=}2$	$n{=}4$	$n{=}8$	$n{=}16$	$n{=}32$	$n{=}64$
10	94.5	95.2	93.3	93.7	92.4	91.5	76.7	48.4	51.2	55.9	50.6	49.1	45.9	40.9
30	97.7	96.8	96.0	95.1	95.0	90.8	40.7	89.0	90.5	82.2	80.9	67.2	42.8	30.6
50	92.2	87.1	90.0	90.5	90.7	69.4	26.5	95.3	95.3	93.7	90.3	78.9	39.6	19.9
70	43.2	41.3	37.2	42.5	36.0	20.2	8.8	76.8	74.1	70.1	61.4	43.0	17.8	8.0
90	0.0	0.0	0.2	0.3	0.6	0.2	0.3	7.1	12.7	5.7	7.1	4.7	2.2	1.1

Appendix

Table 15: Recovery rate (%) for globally visible corruptions.

	PaTh							DM
% sampling	$n{=}1$	$n{=}2$	$n{=}4$	$n{=}8$	$n{=}16$	$n{=}32$	$n{=}64$	$n{=}1$	$n{=}2$	$n{=}4$	$n{=}8$	$n{=}16$	$n{=}32$	$n{=}64$
10	0.8	0.0	0.0	0.2	0.1	0.1	0.2	0.0	0.0	0.2	0.2	0.3	0.2	0.3
30	0.0	0.4	0.2	0.1	0.0	0.2	0.3	0.0	0.0	1.4	0.0	0.4	0.5	0.5
50	4.7	1.6	3.3	4.2	4.3	4.0	6.6	1.6	0.0	1.2	0.9	1.0	1.5	2.7
70	37.5	40.6	42.8	42.7	45.9	48.8	52.2	14.1	22.7	15.0	20.9	28.4	38.1	45.4
90	94.5	98.8	97.7	97.5	97.4	97.6	97.5	86.7	88.3	87.1	89.4	88.4	90.3	90.2

Appendix

Table 16: Kept rate (%) for locally visible corruptions.

	PaTh							DM
% sampling	$n{=}1$	$n{=}2$	$n{=}4$	$n{=}8$	$n{=}16$	$n{=}32$	$n{=}64$	$n{=}1$	$n{=}2$	$n{=}4$	$n{=}8$	$n{=}16$	$n{=}32$	$n{=}64$
10	3.9	1.2	2.8	2.6	3.0	3.1	8.7	15.1	16.5	19.6	18.5	18.4	20.6	21.7
30	0.8	0.4	0.6	2.1	2.0	4.5	26.6	5.5	4.3	6.9	9.1	13.8	24.7	31.4
50	3.9	6.6	4.7	3.7	4.6	16.5	40.7	2.4	2.8	3.9	6.0	12.6	33.9	43.9
70	43.2	46.5	46.3	44.0	49.4	65.5	72.5	17.6	19.6	23.4	28.3	43.1	60.2	67.2
90	97.7	98.0	98.6	97.8	97.4	97.7	97.8	88.2	82.1	90.2	87.1	91.0	93.6	94.4

Appendix

Table 17: Kept rate (%) for globally visible corruptions. The difference between this table and Tab. 16 is plotted in Fig. 7 .

	PaTh			DM
% sampling	$n{=}1$	$n{=}2$	$n{=}3$	$n{=}1$	$n{=}2$	$n{=}3$
10	53.2	55.7	53.0	40.7	41.4	41.1
30	51.2	52.1	51.2	25.5	26.4	26.3
50	39.8	41.4	39.7	14.2	14.7	15.5
70	21.0	24.6	23.4	9.2	9.8	9.1
90	13.0	12.2	12.6	6.0	7.6	6.8

Appendix

Table 18: Per-object recovery rate (%) on constrained scenes.

	PaTh			DM
% sampling	$n{=}1$	$n{=}2$	$n{=}3$	$n{=}1$	$n{=}2$	$n{=}3$
10	43.8	39.4	41.7	53.2	55.5	56.0
30	44.0	41.7	43.5	71.0	70.1	70.8
50	54.2	52.2	53.6	82.5	81.6	80.7
70	74.0	71.1	71.6	87.7	87.5	87.5
90	80.0	82.5	82.3	92.0	88.9	89.4

Appendix

Table 19: Per-object kept rate (%) on constrained scenes.

t \ nsup	1	4	8	12	16
0	66.2	81.9	86.4	87.9	89.1
5	96.2	95.7	95.8	95.9	96.0
10	98.3	98.2	98.2	98.2	98.4
15	98.8	98.8	98.8	98.7	98.8
19	55.3	62.3	62.1	62.3	62.1

Appendix

Table 20: Probe accuracy (%) by denoising step

t

and completed supervision steps, without resetting the recursive state between columns.

Conditioning	Puzzles	TRM num tokens	Puzzle acc.	Cell acc.
puzzle	translated+scaled	81	0.0	11.1
puzzle	translated	81	0.0	11.1
puzzle	translated	144	0.0	11.1
aligned puzzle	translated+scaled	81	3.4	44.3

Appendix

Table 21: MNIST Sudoku with the solution rendered on a larger canvas under a random translation, or translation and scaling. Aligned applies the same transform to the conditioning. Cell accuracy of

11.1\%

is chance.

Conditioning	Attr. binding (%)	Spatial acc. (%)
tokens	7.8	51.2
tokens + part of solution	8.3	50.4
tokens + visual attributions	21.2	51.6
tokens + centroids	53.8	89.4
centroids with attributes	54.8	91.2

Appendix

Table 22: Constrained scene composition under different conditioning. Only conditioning that localises objects in the image improves spatial accuracy. The final row, in which the mask carries one channel per attribute and so states exactly where each object belongs, is a sanity check rather than a proposed setting.

Figure 17: As cells are denoised one by one, the fraction of samples still consistent with the unique solution (blue) falls below

50\%

after the first cell, while the number of rules violated on the board (orange) stays near zero. SRM’s commitments conflict with nothing already present and are therefore invisible to it, yet most have already ruled out the only valid completion.

Figure 18: Puzzle accuracy against reasoning budget

n_{\text{sup}}

for different reset intervals. The Thinker’s state carries across denoising steps , improving accuracy by

23

points at low budgets.

Figure 19: Change in puzzle accuracy when a fixed recursion budget is concentrated at different points of the trajectory (columns) with different spread (rows), relative to uniform allocation. Allocating early helps and allocating late costs up to $11.3$ points.

Figure 20: Puzzle accuracy when the Thinker stops reasoning after a given fraction of sampling and its last output is reused thereafter. Reasoning for the first half suffices, slightly exceeding reasoning at every step (

72.5\%

against

71.2\%

) at half the cost.

Resolution	hard	extreme
144x144	0.925	0.712
252x252 (full-resolution)	0.926	0.682

Appendix

Table 23: Puzzle accuracy comparison of PaTh on different resolutions of MNIST Sudoku.

Figure 21: Examples of solutions to mazes and queens puzzles generated by Bagel.

Explore similar work

Diffusion as a Training Curriculum for Timestep-Free Iterative Reasoning
Sep 1, 2026Mariia Drozdova, Aidan Sirbu, Pietro Miotti +4Functional Code Solvers Denoising Diffusion Probabilistic Model
1Google, Paradigms of Intelligence Team · School of Computer Science, McGill University · 4Mila - Quebec AI Institute +3
Relational Abstractions for Spatial Reasoning with Diffusion Models
Oct 7, 2026Ana Ezquerro, Ozan ÖzdenizciSpatial Reasoning
Institute of Machine Learning and Neural Computation Graz University of Technology, Austria
Inference-Time Scaling in Diffusion Models through Iterative Partial Refinement
May 19, 2026Taegu Kang, Jaesik Yoon, Sungjin AhnDiffusion Models Inference-Time
KAIST · KAIST & SAP · KAIST & NYU

Sep 1, 2026, cs.LGJ/K move · Enter open · S save

Diffusion as a Training Curriculum for Timestep-Free Iterative Reasoning

Mariia Drozdova, Aidan Sirbu, Pietro Miotti, Robert Obryk+3

1Google, Paradigms of Intelligence Team · School of Computer Science, McGill University · 4Mila - Quebec AI Institute+3

Functional Code Solvers Denoising Diffusion Probabilistic Model Inference-Time Recursive Reasoner Timestep

Open paper AskPDF arXiv

Diffusion models and recursive reasoners are both iterative, but they carry information across iterations differently. We add a persistent hidden state to a diffusion denoiser and remove its timestep conditioning, leaving a single shared update that can be run to arbitrary depth. The result is an anytime solver: accuracy keeps improving with inference depth far beyond the rollout lengths and backpropagation window used in training, reaching 99.90% exact solve on Sudoku-Extreme. We also obtain 98.93% solve rate on Maze-Unique. Surprisingly, progressive denoising is unnecessary at inference: holding corruption at its maximum by replacing every non-clue variable with fresh Gaussian noise at each step retains near-perfect solving and converges to stable solutions. This simple noise-injection mechanism enables a single trajectory to efficiently explore the solution space and settle on the correct answer without parallel rollouts, candidate selection, or external verifiers required by prior reasoning models. Nonetheless, ordered annealed corruption remains critical during training, which suggests that diffusion's primary contribution to our anytime solver is not a sampling procedure at inference, but a denoising training curriculum.

Abstract

Explore similar work

Diffusion as a Training Curriculum for Timestep-Free Iterative Reasoning

Relational Abstractions for Spatial Reasoning with Diffusion Models

Inference-Time Scaling in Diffusion Models through Iterative Partial Refinement