Relational Abstractions for Spatial Reasoning with Diffusion Models

Ana Ezquerro; Ozan Özdenizci

cs.CVOct 7, 2026

Relational Abstractions for Spatial Reasoning with Diffusion Models

Authors: Ana Ezquerro, Ozan Özdenizci

Organizations: Institute of Machine Learning and Neural Computation Graz University of Technology, Austria

Abstract

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.

Figures & tables

Figure 1: Illustration of RDM with Discoverer and the Reasoner blocks. During training, the clues ( $\mathbf{x}_{c}$ ) and solutions $(\mathbf{x}_{0})$ are individually fed to the pretrained frozen ( ) slot encoder to obtain the corresponding feature ( $\mathbf{f}^{\prime}_{1:m}$ ) and positional ( $\mathbf{p}^{\prime}_{1:m}$ ) embeddings. Slot Attention yields the slots ( $\mathbf{s}_{1:k}^{\prime}$ ) for reconstruction by the slot decoder, and object-specific features and positions, $\smash{\mathbf{z}^{\prime}=(\tilde{\mathbf{f}}^{\prime}_{1:k},\tilde{\mathbf{p}}^{\prime}_{1:k})}$ for relational abstractions. The relational abstraction module consists of $L$ abstractor layers that optimize the relational embeddings ( ), which are introduced into the U-Net bottleneck through a cross-attention layer. Since the solution is unavailable at inference time, its disentangled representations ( $\tilde{\mathbf{f}}^{0}_{1:k},\tilde{\mathbf{p}}^{0}_{1:k}$ ) are suppressed by a default learnable embedding ( $\varnothing$ ) with probability $p_{\text{drop}}$ .

Panel of Figure 2 — Figure 2: Akari as a grid-based image completion puzzle.

Panel of Figure 3 — Figure 3: Coldoku puzzle, which combines Sudoku with a color arrangement that also satisfies Sudoku rules. Tableau 10 palette is used, which is designed to have a large perceptual difference between colors.

Panel of Figure 4 — Figure 4: Tangram example. A valid solution uses all pieces matching the colors and shapes, given the (a) and (b) as clues .

Panel of Figure 5 — Figure 5: LogicFace example with 4 features: smiling ( OR ), male ( AND ), glasses ( XOR ), young ( IMPLIES ). Features from $A$ and $B$ are combined in $S$ via the corresponding logic operators (e.g., glasses : $A=0$ , $B=1$ , $S=1$ , satisfying XOR ).

	Akari						Coldoku
	in-dist.		low		high		Easy (50-60)			Medium (40-50)			Hard (30-40)
	blb	acc	blb	acc	blb	acc	dgt	clr	acc	dgt	clr	acc	dgt	clr	acc
Base	86.02	80.04	43.62	39.48	92.22	88.92	93.53	94.73	89.00	61.81	78.56	49.06	11.38	25.31	3.88
SRM ( Wewer et al., 2025 )
parallel	35.82	18.80	1.60	0.94	54.18	33.34	94.67	96.74	91.06	61.02	78.53	51.01	12.17	25.91	4.00
sequential	00.00	00.00	–	–	–	–	93.05	96.70	92.37	82.21	83.61	68.35	40.95	36.66	20.78

Table 1: Experimental results on Akari and Coldoku datasets.

	Tangram							LogicFace
		$\square$	$\triangle^{S}$	$\triangle^{M}$	$\triangle^{L}$	shape	acc	smiling ( OR )	male ( AND )	glasses ( XOR )	young ( IMP. )	avg	acc
Base	23.06	24.22	42.82	41.50	55.60	57.72	15.46	54.60	76.52	67.34	38.90	57.84	11.24
SRM ( Wewer et al., 2025 ) 3 3 3 Diffusion forcing led to highly unstable training in Tangram , where both SRM parallel and sequential were unable to outperform Base (similar to LogicFace ). These findings also align with Wewer et al. (2025) .
parallel	–	–	–	–	–	–	–	37.00	38.16	49.26	39.44	40.97	2.60
sequential	–	–	–	–	–	–	–	25.04	58.38	50.58	73.86	51.94	5.70
RDM [S]

Table 2: Experimental results on Tangram and LogicFace datasets.

Panel of Figure 6 — Figure 6: Performance of RDM [ L ] in Tildoku .

Appendix figures & tables21 assets

Supplementary material from the paper’s appendix.

Appendix

	Akari	Coldoku	Tangram	LogicFace
$\texttt{H}\times\texttt{W}$	$198\times 198$	$180\times 180$	$128\times 128$	$128\times 128$
$\texttt{h}\times\texttt{w}$	$11\times 11$	$9\times 9$	$8\times 8$	$8\times 8$
$k$	121 / 16	81 / 16	7	6
# samples	1M	1M	1,013	30,000
neural evaluator	✗	✓	✗	✓

Appendix

Table 3: General statistics: image (

\texttt{H}\times\texttt{W}

) and bottleneck (

\texttt{h}\times\texttt{w}

) resolution; number of slots (

k

) and unique samples (# samples); and neural evaluation.

Panel of Figure 7 — Figure 7: Akari samples from the training set.

Panel of Figure 8 — Figure 8: Visualization of Coldoku samples from the training set (top: solution , bottom: clue ).

Panel of Figure 9 — Figure 9: Visualization of Tildoku solution samples with different tilting ratios.

Panel of Figure 10 — Figure 10: Tangram layouts ordered by reasoning difficulty. From easy configurations (left) containing salient geometric cues that simplify piece identification to complex silhouettes (right) characterized by high structural ambiguity, where multiple valid piece arrangements may exist for the same layout.

Panel of Figure 11 — Figure 11: Visualization of Tangram samples from the training set.

OR			AND			XOR			IMPLIES
$A$	$B$	$S$	$A$	$B$	$S$	$A$	$B$	$S$	$A$	$B$	$S$
0	0	0	0	0	0	0	0	0	0	0	1
0	1	1	0	1	0	0	1	1	0	1	1
1	0	1	1	0	0	1	0	1	1	0	0
1	1	1	1	1	1	1	1	0	1	1	1

Appendix

Table 4: Logical operators between

A

and

B

, yielding

S

. Since there are four operator and each operator has four valid statements, LogicFace has

256

possible configurations.

Panel of Figure 13 — Figure 13: LogicFace examples from the training set with the four binary features: smiling ( OR ), male ( AND ), glasses ( XOR ), young ( IMPLIES ).

Figure 17

	$k$	$d$	channels	#conv/layer	#params.	#steps	Accuracy
Slot-Attention
Akari	16	128	[16, 32, 64, 128]	1	1.4M	100K	99.58
Coldoku	16	128	[32, 64, 128, 256]	1	5.4M	400K	88.19
Tangram	7	128	[16, 32, 64, 128]	1	1.4M	300K	16.34
LogicFace	6	256	[32, 64, 128, 256]	2	7.6M	200K	72.54
Tildoku	20	128	[16, 32, 64, 128]	2	7.6M	400K	61.15

Appendix

Table 5: Architecture, training and performance details of the Discoverer in configurations S and M .

	$d$	channels	#conv/layer	#params of $\mathcal{D}$	#params of $\mathcal{R}$	#params of U-Net	# steps
RDM [S]	128	[64, 128, 256, 512]	[2, 2, 1, 1]	(Table 5 )	4M	25M	300K
RDM [M]	256	[64, 128, 256, 512]	[2, 2, 2, 2]	(Table 5 )	4M	35M	300K
RDM [L]
Akari	512	[128, 256, 512, 512]	2	40M	15M	120M	40K
Coldoku	512	[128, 256, 512, 512]	2	40M	15M	120M	45K
Tangram	512	[128, 256, 512, 512]	2	40M	15M	120M	80K

Appendix

Table 6: Architecture and training details of the Reasoner .

Panel of Figure 14 — Figure 14: Disentanglement of the tans on Tangram inputs.

Panel of Figure 15 — Figure 15: Disentanglement of the facial attributes on LogicFace inputs.

		ID					OoD
		ID					5-10% (low)					35-40% (high)
		unl	dst	ovp	blb	acc	unl	dst	ovp	blb	acc	unl	dst	ovp	blb	acc
$\tau=10$	Base	0.46	0.10	0.03	74.08	67.35	3.22	0.22	0.48	17.02	13.80	0.17	0.04	0.01	89.78	86.80
	SRM
	parallel	2.38	0.81	0.15	26.52	12.26	15.43	0.63	0.77	1.04	0.54	1.27	0.63	0.06	45.42	27.22
	sequential	7.01	9.81	17.79	0.00	0.00	0.26	4.64	87.15	0.00	0.00	13.87	13.47	6.28	0.00	0.00

Appendix

Table 7: Detailed results on Akari .

acc	Easy (50-60)			Medium (40-50)			Hard (20-30)
acc	dgt	clr	acc	dgt	clr	acc	dgt	clr	acc
Base	59.20	66.84	40.28	22.62	45.92	12.48	1.50	9.42	0.24
SRM ( seq. )	44.87	25.87	13.80	24.31	21.31	7.08	3.62	5.62	0.62
RDM [L]
clue	62.52	79.74	50.28	33.02	52.84	19.18	3.78	10.12	0.64
solution	65.72	80.92	53.26	43.70	60.12	28.40	8.34	14.48	1.92

Appendix

Table 8: Detailed results on the Tildoku dataset.

		Configuration [S]							Configuration [L]
$(^{*}\uparrow)$			$\square$	$\triangle^{S}$	$\triangle^{M}$	$\triangle^{L}$	shape	acc		$\square$	$\triangle^{S}$	$\triangle^{M}$	$\triangle^{L}$	shape	acc
$\tau=10$	Base	17.96	17.82	29.78	34.98	39.33	40.09	12.37	16.08	15.24	21.78	24.64	28.08	28.64	11.46
	RDM
	clue	24.84	25.37	37.55	40.71	41.90	42.41	22.20	34.72	34.26	41.80	43.34	44.88	44.98	31.96
	solution	31.43	31.80	36.73	38.69	40.06	40.51	28.45	33.34	32.66	39.90	41.64	43.16	43.52	29.94
	full	34.49	34.69	38.65	42.67	43.82	4.31	30.18	22.20	21.38	33.26	36.28	40.48	41.44	16.74

Appendix

Table 9: Detailed results on Tangram .

$(^{*}\uparrow)$		Configuration [S]							Configuration [L]
$(^{*}\uparrow)$		smiling ( OR )	male ( AND )	glasses ( XOR )	young ( IMP. )	avg	acc	fid	smiling ( OR )	male ( AND )	glasses ( XOR )	young ( IMP. )	avg	acc	fid
$\tau=10$	Base	71.28	79.39	60.77	73.50	71.24	26.03	44.95	55.08	55.42	39.24	70.28	54.51	8.62	35.92
	RDM
	clue	87.74	96.66	93.61	81.67	89.92	66.08	37.66	83.34	94.52	49.60	70.02	74.37	27.98	37.18
	solution	68.53	78.34	49.81	72.10	67.19	18.98	50.21	65.41	66.44	49.41	66.05	61.83	14.95	39.71
	full	84.97	94.74	90.81	78.55	87.27	59.57	36.37	78.38	91.74	50.68	65.66	71.61	25.08	38.04

Appendix

Table 10: Detailed results on LogicFace .

Panel of Figure 18 — Figure 18: Examples of generated Tangram solutions.

Panel of Figure 20 — Figure 20: Examples of generated LogicFace solutions shown in $S$ , when conditioned on $A$ and $B$ .

acc	Coldoku ( $\tau=10$ )			Tangram
acc	Easy	Medium	Hard	( $\tau=100$ )
Base	90.84	48.72	3.43	15.46
Slot-based	93.76	54.94	4.74	27.09
RDM ∗ ( clue )	96.66	65.62	7.76	34.37

Appendix

Table 11: RDM vs slot-based diffusion.

	MNIST-Sudoku		Counting Polygons
	Medium	Hard	Original	Unbiased
Base	54.18 ( 53.6 )	1.9 ( 0.8 )	13.2	9.8
SRM
parallel	31.04 ( 59.0 )	0.00 ( 1.0 )	16.6	10.1
sequential	40.02 ( 75.4 )	0.01 ( 51.6 )	18.6	0 5.4
Non-object-centric features

Appendix

Table 12: Results on previous benchmarks.

Explore similar work

Think Before You Paint: Recursive Latent Reasoning for Diffusion Models
Oct 7, 2026Paweł Skierś, Małgorzata Grzanka, Wojciech Masarczyk +2Diffusion Models
Warsaw University of Technology · IDEAS Research Institute · University of Amsterdam
The Art of Interrogation: Consistency Amplifies Factuality in Spatial Reasoning
Jun 10, 2026Theo Uscidda, Marta Tintore Gazulla, Maks Ovsjanikov +2Spatial Reasoning LLM Reasoning Strategies
CREST, ENSAE, Institut Polytechnique de Paris · Google Zurich · Google DeepMind
The Thinking Pixel: Recursive Sparse Reasoning in Multimodal Diffusion Latents
Apr 28, 2026Yuwei Sun, Yuxuan Yao, Hui Li +1Multimodal Diffusion Models Diffusion Models
Shanghai Academy of AI for Science · Fudan University

Oct 7, 2026, cs.AIJ/K move · Enter open · S save

Think Before You Paint: Recursive Latent Reasoning for Diffusion Models

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

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

Diffusion Models

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Think Before You Paint: Recursive Latent Reasoning for Diffusion Models

The Art of Interrogation: Consistency Amplifies Factuality in Spatial Reasoning

The Thinking Pixel: Recursive Sparse Reasoning in Multimodal Diffusion Latents