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 xt and encoded convolutionally; non-spatial conditioning is embedded by an MLP. Together they form the Thinker’s input c , over which the Thinker recurses on (y,z) within a single denoising step. Only the spatial tokens of 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
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 7
Figure 6: Qualitative comparison on mazes. When the path is directed into a dead-end, only PaTh is able to correctly revise it.
Figure 7: Difference in retention rate between indirectly and directly available mistakes on extreme MNIST Sudoku (left) and Mazes (right). The diffusion model retains the indirect mistakes far more often early in sampling, while PaTh treats the two types of errors comparably.
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 zi decoded by a linear probe with violated constraints in orange. The Thinker introduces violations absent from xt and then resolves them, reaching a valid grid before anything is committed to xt−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
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
Parameter
MNIST Sudoku
CLEVR
Mazes
N-Queens
Painter
VAE
-
stabilityai/sdxl-vae
-
-
Block types
Block2D ×3
CrossAttBlock ×3
Block2D ×3
Block2D ×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
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
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×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 ↑
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 ℓ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 nsup 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.
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.
Mariia Drozdova, Aidan Sirbu, Pietro Miotti +4
1Google, Paradigms of Intelligence Team · School of Computer Science, McGill University · 4Mila - Quebec AI Institute +3
Diffusion models excel at image synthesis, but they remain limited in their ability to reliably satisfy structured spatial reasoning constraints. In conditional data distribution modeling tasks with implicit logical structure, such as puzzles defined by visible clues paired with consistent solutions, state-of-the-art generative models tend to approximate pixel-space distributions without learning the underlying logical rules required for inference. To address this limitation, we present a novel framework for spatial reasoning with diffusion models that leverages unsupervised object discovery and abstractions of object relations. We show that the relational knowledge derived from object-centric representations enriches diffusion models with structural primitives, allowing them to effectively guide the generative representation space during both training and inference, and enabling conditional image generation that satisfies reasoning constraints. Additionally, we introduce a large-scale generative spatial reasoning benchmark with four datasets inspired by human-solvable puzzles. Our results show that relational abstractions significantly improve reasoning capabilities of diffusion models on a variety of complex reasoning tasks, while enabling robust generalization in out-of-distribution settings.
Ana Ezquerro, Ozan Özdenizci
Institute of Machine Learning and Neural Computation Graz University of Technology, Austria
Inference-time scaling has emerged as a major approach for improving reasoning capabilities, and has been increasingly applied to diffusion models. However, existing inference-time scaling methods for diffusion models typically rely on external verifiers or reward models to rank and select samples, limiting their scalability to settings where such evaluators are available and reliable. Moreover, while recent diffusion models perform sequential inference with region-wise, mixed-noise conditioning, inference-time scaling tailored to this setting remains relatively underexplored. We propose Iterative Partial Refinement (IPR), an inference-time scaling method for sequential diffusion that requires no external verifier. Starting from an already-generated sample, IPR re-noises a subset of regions and regenerates them conditioned on the remaining regions, enabling the model to revise earlier decisions under a richer context than was available during the initial generation. This iterative partial refinement produces more globally consistent samples without external verification. On reasoning tasks requiring global constraint satisfaction, IPR consistently improves performance: on MNIST Sudoku, the valid solution rate increases from 55.8% to 75.0%. These results show that iterative partial refinement alone can serve as an effective inference-time scaling strategy for diffusion models in sequential, mixed-noise settings. Code is available at: https://github.com/ahn-ml/IPR