Spatial Reasoning with Denoising Models
Authors: Christopher Wewer
Bartlomiej Pogodzinski
Bernt Schiele
Jan Eric Lenssen
Venue: ICML 2025 poster (2025)
Keywords: Generative Models
Reasoning
Image Generation
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: A diffusion-based framework for structured reasoning over continuous spatial variables — enabling flexible sampling strategies w.r.t. amount and order of sequentialization. Introduces a benchmark to measure hallucination.

Abstract: We introduce Spatial Reasoning Models (SRMs), a framework to perform  reasoning over sets of continuous variables via denoising generative models. SRMs infer continuous representations on a set of unobserved variables, given observations on observed variables. 
Current generative models on spatial domains, such as diffusion and flow matching models, often collapse to hallucination in case of complex distributions. To measure this, we introduce a set of benchmark tasks that test the quality of complex reasoning in generative models and can quantify hallucination. The SRM framework allows to report key findings about importance of sequentialization in generation, the associated order, as well as the sampling strategies during training. It demonstrates, for the first time, that order of generation can successfully be predicted by the denoising network itself. Using these findings, we can increase the accuracy of specific reasoning tasks from <1% to >50%. Our [project website](https://geometric-rl.mpi-inf.mpg.de/srm/) provides additional videos, code, and the benchmark datasets.

Spatial Reasoning with Denoising Models
Christopher Wewer'
Bart Pogodzinski
Bernt Schiele
Jan Eric Lenssen
Playing MNIST Sudoku with Image Generators....
SRMs (Ouns)
Standard DDPM DDPM
007
27
8
1
H
112
Figure 1: We introduce Spatial Reasoning Models (SRMs), a framework systematically investigate diffusion/flow-basee generative models with respect to their reasoning capabilities multiple variables. We introduce benchmarks that allow quantification of higher-level reasoning capabilities and show that manual and automatic schemes for sequentialization can heavily reduce hallucination. Here, we show the solving process of one of the benchmarks, Sudoku game consisting of MNIST images, which is solved correctly by our SRMs, while standard diffusion models fail.
Abstract
1. Introduction
We introduce Spatial Reasoning Models (SRMs), a framework to perform reasoning over sets of continuous variables via denoising generative models. SRMs infer continuous representations on set of unobserved variables, given observa- tions on observed variables. Current generative models spatial domains, such diffusion and flow matching models, often collapse to hallucina- tion in case of complex distributions. To measure this, we introduce a set of benchmark tasks that test the quality of complex reasoning in generative models and can quantify hallucination. The SRM framework allows report key findings about importance of sequentialization in generation, the associated order, as well as the sampling strate- gies during training. It demonstrates, for the first time, that order of generation can successfully be predicted by the denoising network itself. Using these findings, we can increase the accuracy of specific reasoning tasks from < 1% > 50%. Our project website provides additional videos, code, and the benchmark datasets.
Conditional generative models, such as GPTs (Radford et al., 2018) or denoising models like diffusion//low-baat els (Ho et al., 2020; Song et al., 2021; Lipman et al., 2023; Liu al., 2023), promise significant advances in practical reasoning. They allow to model highly complex and multi- modal distributions by learning to sample from them. While reasoning capabilities of LLMs are extensively explored in many recent works (Huang & Chang, 2023), similar ef- forts in continuous, spatial domains are needed to profit from the semantic structure that can be learned from high- dimensional, continuous data. Our work approaches this goal by providing novel framework benchmark and ad- vance the reasoning capabilities of diffusion//low-b -bseddddd gen- erative models, which we deem of large importance for the further development of large image, video, and physically- grounded world models (Agarwal et al., 2025).
mod-
The process of reasoning can be defined as inferring the states of unobserved variables x;, given distinct set of observed variables y/j (Kwan et al., 2008), with varying de- pendencies between all variables, loosely inspired by tradi- tional fields of probabilistic graphical models and Bayesian networks. Most reasoning tasks are subject to a high degree of inherent uncertainty due to incompleteness of informa- tion given in the observed variables. Thus, it is natural to tackle reasoning problems with probabilistic inference, i.e. modeling p(=1...., Zn V/...., When reasoning across nigher-dimensiones continuous domains, such as images or image patches, these distributions become complex and multi-modal, preventing the application of traditional ap-
'Max Planck Institute for Informatics, Saarland Informat- ics Campus, Germany. Correspondence Christopher Wewer <ewewer@me mpi-ini inf mpf.ddpgee,, Bart Pogodzinski <bpogodzi @ mpi- inf mef.mpp.e>>, Jan Eric Lenssen <jlenssen& @ mpi- ing.mmg.fm..>>.....
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Spatial Reasoning with Denoising Models
3 3/6 3/6/72258
695/27381
#9725 8
77 1792000 601
3
/ 7875 2396
138 1885566 6
8 4
4
87
22 1-
25
4
1243 5
65 1549271
3
9 26583241
55 751364m38
12433
65127 90M8
1
9
9
H
/
87 0739810 144
375
13 2 A 5
2
/ 147 GK
81 1137544 74754629
94 54 1223222
8
1
5 6
23 687
26
2 2 1
44
7
2
79
4
2
7
29
59
4
9
5
4 23 97
5
8 9
4
42 1267888 3
(a) Training Data
(b) Conditioning image and generation output, standard DDPM
Figure 2: MNIST Sudoku. One of our reasoning benchmarks. (a) A dataset of correct Sudokus, consisting of random MNIST images. (b) When training conditional diffusion model correct examples it fails to conditionally generate correct solutions for hard Sudokus (almost 0% accuracy). In contrast, SRMs achieve > 50% accuracy for these cases.
proaches that rely on Gaussian assumptions or predefined families of discrete distributions.
To this end, we introduce Spatial Reasoning Models (SRMs), an architecture-agnostic framework to formulate (soft or hard) sequentialization strategies that can be used for rea- soning on sets of continuous variables without canonical order. Within this framework, create multiple of such strategies and perform exhaustive evaluation on bench- mark tasks as described above.
Generative models for spatial domains typically learn to sample from approximated data distribution p(21,,-,,%.) To perform reasoning, it is desired that such model ab- stracts from the low-level data space and is able to detect high-level patterns that occur in seen data. To obtain an insight into the level of such capabilities of existing diffu- sion/flow models, we introduce a set of benchmarks. One example based on the well-known Sudoku game, where each number from
In summary, our contributions are
a general framework for reasoning over sets of contin- uous variables with denoising generative models,
must occur exactly once in each row, column and 3x3 block, is shown in Fig. 2. A (conditional) diffusion model is trained to generate or complete correct Sudokus consisting of varying MNIST digits. While the gen- erated images look realistic on the first glance, as individual numbers are high- high-quality samples from the MNIST dataset, the generation process fails to sample correct instances. We identify this effect as an instance of hallucination, where the model falls back superficial solutions that satisfy the pixel-wise MSE loss well enough, in case it fails to capture the actual, complex patterns in the underlying distribution.
a novel algorithm for t-sampling during training, which adjusts for multiple variables,
different task-specific and task-agnostic, soft and hard variants of sequentialization within reasoning models,
a benchmark to quantitatively measure hallucination in reasoning over visual domains.
We can summarize the key findings as
diffusion//low-basee generative models capable of simple visual reasoning but fall back to hallucination when the modeled distribution becomes too complex,
In the space of large language models (LLMs), chain-of- thought prompting has achieved successes in fighting such hallucinations by prompting the model to make smaller, sequential steps towards the solution (Wei et al., 2022). In- spired by these advances, develop and investigate strate- gies in our spatial setting that can lead to similar effects. The chain-rule of probabilities lets us decompose
inducing sequentialization can simplify the task by decomposing into simpler distributions, reducing hal- lucination, and improving reasoning.
greedy orders based on predicted uncertainty can sig- nificantly improve reasoning quality further,
the choice of training (-sampling strategies is crucial when adjusting for sequentialization.
p( In) = II pla (f
(I)
Our framework, code, and benchmarks are available on our project website for further investigation and development.
for arbitrary permutations n. While all orders ग model the correct distribution theory, the individual orders lead to chains of varying complexity the individual distributions, determined by hidden dependencies, potentially leading to varying amounts of hallucination during sampling. We be- lieve investigating strategies to find (I) the correct amount of sequentialization and (2) the best order with least amount of hallucination is a promising venture for further development of continuous generative models.
2. Related Work
Probabilistic reasoning is currently receiving a lot of at- tention, mostly driven by recent advancements in large language models (LLMs) (Huang & Chang, 2023; Plaat et al., 2024) that allow to generate tokens from a discrete codebook. To leverage their capabilities, many modalities
2

---

Spatial Reasoning with Denoising Models
have been mapped such discrete spaces, using special- ized tokenizers (van den Oord et al., 2017), allowing to reason about originally continuous domains to certain de- grees (Chen et al., 2023). Vision Language Models (VLMs) have been heavily used to connect spatial domains with language, to perform reasoning about images in the text domain (Chen et al., 2024b). In this work, we move in different direction and investigate reasoning capabilities across continuous domains by generative models that have explicitly been designed to work on these domains, such as DDPM (Ho et al., 2020), DDIM (Song et al., 2021), recti- fied flow (Liu et al., 2023), or flow matching (Lipman et al., 2023). Investigating reasoning capabilities of such models has remained largely unexplored so far.
Belief Propagation The above paradigm is loosely in- spired by belief propagation through Probabilistic Graphical Models (Pearl, 1982), propagating information from ob- served to unobserved variables via sequential probabilistic inference. However, the set of assumptions is drastically re- duced. It does not make any assumptions about distributions of x and does not strictly require the Markov assumption about conditional independence between variables. Instead, the network can (but not has to) be conditioned on all other variables, utilizing compression capabilities of neural net- works. Due to the nature of generative models, however, the formulation does not allow for explicit representation of probabilities but allows only sampling.
3.1.1. SAMPLING IN CONTINUOUS DOMAINS
Closest to our work are recent methods that introduce differ- ing levels of noise in sequential generation with diffusion models (Wu et al., 2023; Zhang et al., 2024; Ruhe et al., 2024; Chen et al., 2024a). AR AR-Diffusion (Wu et al., 2023) performed generation via diffusion on text spaces. Later, the idea was extended to sequential, continuous spaces, such as human poses (Zhang et al., 2024), or videos (Ruhe et al., 2024; Chen et al., 2024a), which also introduced overlap- ping via different noise levels of variables in purely sequen- tial settings, Our work extends this direction to arbitrary spatial settings without canonical orders. Also, we heavily improve on the training sampling strategy used in Diffusion Forcing (Chen et al., 2024a). Another recent related work is MAR (Li et al., 2024), which performs autoregressive generation of image patches via denoising in random order.
We first introduce our formulation for sampling a single continuous random variable. Denoising generative models are the established state-of-the-art for learning continuous data distributions (Dhariwal & Nichol, 2021). While there are multiple different derivations like score matching (Song & Ermon, 2019), DDPM (Ho et al., 2020), or the recently more popular (conditional) flow matching (Lipman et al., 2023), they all share the same idea of learning a step-wise mapping of scalar or higher-dimensional continuous vari- ables x' from a simple and known Gaussian distribution xt ~ (O, 7) to the complex and unknown data distribu- tion x0 ~ q. With small enough step sizes t' -'-E f (Ho et al., 2020), each step itself can be modeled by sampling from a Gaussian distribution (with possibly zero variance)
x² ~ NNool ), Eo(x")) )).
(3)
3. Spatial Reasoning Models
parameterized by a neural network with weights 0 trained to denoise noisy versions of samples from the data distribution. Training examples are constructed by interpolating between samples x² ~ q and Gaussian noise e ~ No 7) as
In this section, we introduce our spatial reasoning mod- els. We begin with formulating our general framework in Sec. 3.1. After training a SRM as outlined in See. 3.2, it can be leveraged with various sampling strategies. We introduce a selection of those in Sec. 3.3.
x² === =a()a=² + b(()e,
(4)
where t E [0, 1] is a continuous level of noise and a, : [0. - [0, 1] define interpolation weights. While DDPM (Ho et al., 2020) limits itself to subset of possible noise schedules, alternatives like rectified flows were re- cently explored as general Gaussian probability paths in the context of conditional flow matching (Lipman et al., 2023).
3.1. General Framework
Our goal is to learn to reason over sets of continuous ran- dom variables. Practical examples of such variables can be patches of images, frames of videos, even multiple views of 3D scene. For our SRM setting, we define reasoning over a set of continuous random variables as sampling
In this work, we construct a specific reverse distribution, i.e., mean and variance in Eq. 3 for arbitrary noise schedules a, b with its marginal distribution satisfying Eq. 4, similar as but more general than DDIM (Song et al., 2021). We provide all details and proofs in Appendix A. As result, our formu- lation can combine non-diffusion noise schedules such as from rectified flows (Liu et al., 2023) (a(t) = 1 1-t, b(t) t) with stochastic sampling as in DDPM and follow-ups such as learning the variances of the reverse distributions (Nichol & Dhariwal, 2021).
at ...,, 2fn q(a"' yin shr
2t).
where oft denotes the variable with noise level 1; during the denoising process and t. < ty. Spatiality is (optionally) encoded via positional encodings on the variables. Choosing noise levels & allows explicit control over amount and order of sequentialization, as discussed further in Sec. 3.1.2.
3

---

Spatial Reasoning with Denoising Models
Parallel
Overlapping
Full Autoregressive
Reasoning Proutss
2
Random
Reasoning Model
Predicted Uncertainty
3
0
J2
Manual Graph-based
2'1
27
2
3
4
2
2
(a) Amount of Sequentialization
(b) Order of Sequentialization
Figure 3: Investigated Degrees of Freedom. (a) The SRM framework allows to define different amounts of sequentialization, Lee. parallel generation, autoregressive generation or mixtures with varying overlap, modeled by differing levels noise on individual variables at the same time. (b) Further, it provides different options to define the order of sequentialization, allowing random ordering. greedy heuristic based on predicted uncertainty and manually-defing graphs.
3.1.2. SAMPLING SPATIAL DOMAINS
where t (f1:-.., (n), x (21,--. Fn)) and x² (zf),..,,<? x..). Additionally, we let the network predict the variance of the reverse process in Eq. 3, optimized for the variational lower bound (Nichol & Dhariwal, 2021). As shown in previous works (Esser et al., 2024), the training of a denoising network very sensitive to the sampling of noise levels. This is especially true for our setting of indi- vidual noise levels for possibly large number of variables.
We model our spatial reasoning problems inferring sets of continuous random variables, unifying previous denoising generative approaches and spatially autoregressive methods.
On the one hand, existing diffusion- and flow-based models like popular image generators (Rombach et al., 2022) reason only along the noise dimension by assuming single, shared noise levels for all variables at any point in time during the denoising process. By setting & = ... tu and tj t, in Eq. 2, this special case is subsumed by our framework. We hypothesize potential limitations of this approach in the presence of complex dependencies between spatially distinct variables, as in the case of MNIST Sudoku (Fig.2).
3.2.1. NOISE LEVEL SAMPLING
Sampling noise levels during training of diffusion models allows the model to learn denoising images of all noise levels t. In SRMs, the image patches also act as conditioning to other patches, so we need to ensure that the whole t distribution similar what happens during inference.
On the other hand, spatially autoregressive approaches like MAR (Li et al., 2024) learn to sample a next variable z; con- ditioned the history of previously sampled variables z j E I and ((i), J, X.) being partitioning of {1, n Our framework covers this specific way of sampling by set- ting 12 12 = 0 for E J and l t = for & E IC. Although the application of the chain-rule can decompose the complex joint distribution of variables into potentially simpler conditional distributions, improvements depend on the order, which is usually simply set to be random or raster scan.
At the beginning of inference, all patches p; start with 4 = 1, so the mean noise level is initially ī = 1. Similarly, at the end of inference, all patches reach t; = 0, giving = 0. For diffusion models with linear time schedules, or our SRM with parallel schedule, the mean noise level decreases by a fixed amount per inference step AT - where N is the number of inference steps. In the autore- gressive generation, similar relation holds At = M-P- where M is the number of steps per patch, and P is the number of patches. In both cases, A/ does not depend on the noise level, which suggests that with sufficient num- ber of inference steps, should ideally follow a uniform distribution. See Fig. appendix for visualization.
argue that potential
These two special instances depict the extreme cases of our framework. Given that both have their own strengths and weaknesses, we aim to explore the space in between. There- fore, we train spatial reasoning models to jointly denoise multiple variables but with individual noise levels.
Diffusion Forcing (Chen et al., 2024a) suggests indepen- dently sample noise level vector t ~ Le([0, 1]"), ensuring uniform marginal distribution, ~ ~((0, 1). However, this leads to the mean i following a Bates distribution, which is highly concentrated around 0.5 - siiniiinffiaattt) different from the distribution encountered during inference (c.f. Fig. 4a), As a result, trained model is undertrained for early (and late) inference steps, where most patches have noise levels close (and 0), reducing performance (c.f. Sec. 4.6).
3.2. Training
We parameterize a SRM as a noise prediction network 6p that is trained to regress the ground truth noise ~ I)
Cp = -x,x,xx,xoo -112

---

Spatial Reasoning with Denoising Models
3.3. Sampling
A trained SRM can be used for various sampling techniques and therefore benchmark them fairly using a single model. We investigate different amounts and orders of sequential- ization for the denoising of spatial variables (see Fig. 3).
Individual C
Regarding the amount of sequentialization, we explore par- allel generation as usual in denoising generative models, the other extreme of fully autoregressive sampling, and mix- ture achieved by overlapping the denoising step intervals of the different variables by variable amount.
Mean
Mean
(a) Uniform i (Diffusion Forcing)
(b) Uniform if (ours)
Figure 4: Noise Level Sampling, We visualize the joint distribution of noise levels i and mean noise levels f for (a) Diffusion Forcing training and (b) our Uniform approach. In Diffusion Forcing, the distribution of t; is correctly mod- eled but f is overly concentrated around f 0.5, which does not reflect the distribution seen during inference. Our Uniform I addresses this issue, though it introduces some bias in the distribution of individual &; which we address by per-patch loss weighting.
Besides the amount of sequentialization, we further com- pare different orders of variables. Leveraging the estimated uncertainty in the noise prediction from Sec. 3.2.2, we pro- pose to adaptively sample the next variable with the lowest uncertainty. This strategy follows the assumption that more uncertain variables can benefit from stronger conditioning of many clean variables later during sampling. For some tasks, we know the dependency structure between spatial variables in advance. Considering our Sudoku example, a cell de- pends directly on the other cells in its row, column, and 3x3 block. Encoding this information in form an adjacency matrix, we propose to propagate certainty based on each variable's noise level to sample the most certain variable next. We include random order as further baseline.
To address this issue, we introduce a two-step sampling strategy called Uniform F. We first sample i ~ (0, 1), and then generate t ~ p(t | f) using our recursive allocation sampling algorithm (Appendix C. The algorithm allows control of the sharpness of the p(t t) the spread of individual f/ around f. This makes it adaptable to different inference scenarios flexible for different inference settings. For example, in the parallel generation, & should stay close to (high sharpness). In contrast, autoregression benefits from greater cross- pooss-pacch noise level variance.
4. Experiments
We evaluate SRMs for reasoning on three new benchmark datasets that we introduce in Sec. 4.1. We consider images split into patches as our sets of continuous random variables and use 2D UNets for denoising as described in Sec. 4.2, Besides our main findings detailed in Sec. 4.3 to 4.5, we provide additional ablations in Sec. 4.6.
Uniform E sampling preserves the uniform mean distribution but oversamples individual patch noise levels close to 4 0 and f; = 1 (c.f. Fig. 4b). This occurs because if f = 0 is sampled, all & must be close to 0 and similarly for F x 1, Vit, = 1. On the other hand, for f z 0.5, the t, samples can take any value within the range. To counteract this bias towards sampling t4
4.1. Benchmark Datasets
We introduce three different datasets to quantify reasoning capabilities. They are aimed at different aspects to be tested. The MNIST Sudoku dataset captures complex (NP-hard) dependencies that need to be understood. The Even Pixels dataset is an easier task that can be solved in greedy fash- ion. Finally, we introduce the Counting Polygons Stars FFHQ dataset, which moves closer to real-world images.
t s 1, we introduce per-patch loss weights W/ p())' where p(t) is empirically estimated.
3.2.2. UNCERTAINTY ESTIMATION
To sample variables in meaningful order, we propose to train the denoising network to predict the uncertainty in its noise prediction. As common in heteroscodastic uncertainty estimation (Seitzer et al., 2022), we model the uncertainty as standard deviation de (xt) and minimize the negative log-likelihood (NLL) of the ground truth noise c
MNIST Sudoku Dataset We create a dataset based on one million correct Sudoku instances by randomly sam- pling MNIST representatives for each cell (see Fig. 2a for examples). The classification of Sudoku images into cor- rect and incorrect ones is then done via an application of pre-trained MNIST digit classifier MLP on each cell and a check whether all Sudoku rules are fulfilled. To avoid classi- fication errors distorting our reasoning results, we limit the set of MNIST digit representatives to the top 1000 training examples per class of the classifier w.r.t, its confidence.
L. = E,,,x ((cc(ccxxx""
(6)
To avoid deteriorating the noise prediction, we only train the uncertainty estimation with this loss and a small weight.
5

---

Spatial Reasoning with Denoising Models
(a) Even Pixels dataset
(b) Counting Polygons/ / Stars FFHQ dataset
Figure 5: Two Further Datasets. In addition to MNIST Sudoku, we introduce two further datasets. (a) The Even Pixels dataset requires the model generate images with an equal number of pixels per color. (b) Counting Polygons Stars FFHQ requires that all images contain two digits: the number of objects and the number of polygon edges star points, respectively. We introduce real-world complexity by placing the synthetic objects over FFHQ samples as backgrounds.
For testing, we use a held-out dataset split of valid Sudokus and apply random masking of cells with the number of masked ones randomly sampled from the intervals [1, 27], [28, 54], and [55, 81], resulting three levels of difficulty easy, medium, and hard, respectively. As metrics, we con- sider accuracy well the sum of Li-distances of row-, column-, and block-wise digit histograms to the all ones vector (zero if correct), averaged over all test examples.
number of polygons / stars, vertices / points, and the digits.
To introduce real- world complexity, we place these synthetic objects over backgrounds sampled from the FFHQ dataset. This forces the model to notice and learn meaningful spatial dependencies in setting much closer to the real world. We train ResNet classifiers to estimate correctness of generated samples and verify their capabiliit on held-out validation splits. More details are given in Appendix E.
The game Sudoku involves complex spatial dependencies, whereas the distribution of individual cells being MNIST numbers is simple to fit by generative model. Note that the general task of solving Sudokus for variable grid size is NP-complete and that depending the level of masking, there can be multiple valid solutions resulting in uncertainty. We further highlight the difficulty of even a discrete Sudoku version for state-of-thh-ar LLMs in Appendix H.
4.2. Experimental Setup
We compare SRMs in pixel space with standard denoising diffusion models that are equivalent in architecture and train- ing up to the sampling of individual spatial variables, which are chosen to be image patches of task-specific sizes for our visual reasoning tasks. Since our MNIST Sudoku bench- mark can be seen as an instance of inpainting, the baseline diffusion model is additionally conditioned on the incom- plete Sudoku and its mask via concatenation in the input. Note that this is not necessary for SRMs, already given variables are identified by the noise level zero. If not stated otherwise, we use stochastic sampling similar to DDPM with a total of 1000 steps (network evaluations) indepen- dent of the method of sequentialization during sampling. Therefore, the computational cost equal for all sampling methods resulting in a fair comparison. This also means that there are fewer denoising steps for individual variables the higher the degree of sequentialization (lower overlap). SRMs are agnostic to both noise schedules and denoiser ar- chitectures. For the main paper, we use rectified flows (Liu et al., 2023) and lightweight versions of widely established 2D UNets with spatial attention in low low-resolution layers to avoid results being dominated by too extreme overparam- eterization. We report additional results with the cosine noise schedule (Nichol & Dhariwal, 2021) and Diffusion Transformers (DiT) (Peebles & Xic, 2023) in Appendix G, describe the exact sampling process in Appendix A, and provide implementation details in Appendix F.
For our graph-based order for sequentialization during sam- pling, we encode the direct dependencies between each cell and its neighbors within row, column, and block as the ad- jaceney matrix. Additionally, we provide results for oracle algorithms that directly solve discrete Sudokus by randomly sampling a number for the next cell from all possible digits avoiding collisions with neighbors if possible. The order of cells is chosen either randomly or greedy fashion, with the next cell being the one with the largest number of already sampled neighbors in the constructed graph.
Even Pixels Dataset We construet a dataset of images where each pixel is assigned one of two opposite colors (c.f. Fig. 5a). The number of pixels of each color is always equal. The task requires a model to learn this implicit con- straint without explicit supervision. We detail our evaluation procedure in Appendix D.
Stars FFHQ Dataset We also in- troduce two datasets in which each image contains set of randomly positioned polygons or stars and numbers (see Fig. 5b), where: The number of synthetic objects, N, is randomly sampled from {1,..., .9). each polygon / star has K E .3...., 7} vertices or K E (2...., 99 points, respec- tively, and the numbers N and K are explicitly present in the image. The model's task is to generate images following this rule and hence to capture the relationship between the
Counting Polygons/
4.3. MNIST Sudoku Results
We provide quantitative comparison of SRM with different sampling strategies, a standard diffusion model baseline, and the task-specific discrete oracle algorithms in Tab. 1.
6

---

Spatial Reasoning with Denoising Models
Easy
Medium
Hard
Model
Sampling
L.1.
Accf
L.1.p
Accf
1.1.
Acct
Diffusion Model
Parallel
0.032
0.994
3.924
0.536
14.120
0,008
Parallel
0.012
0.998
3,312
0.590
19.156
0.010
Predicted Order w/o Overlap
0.012
0.998
1.616
0.754
3.212
0.516
Predicted Order Overlap
0.000
1.000
1.816
0.716
4.340
0.422
SRM (Ours)
Random Order w/o Overlap
0.012
0.998
3.828
0.568
12.612
0.024
Random Order Overlap
0.000
1.000
4.172
0.542
13.896
0.020
Graph-based Order w/o Overlap
0.020
0.996
2.544
0.662
5.816
0.266
Graph-based Order Overlap
0.000
1.000
3.240
0.576
6.036
0.238
Random
0.751
0.059
0.000
Oracle
Greedy
1.000
0.987
0.672
Table I: MNNST Sudoku Quantitative Results. While standard DDPM diffusion models fail to solve the hard Sudoku cases (accuracy near 0%, SRMs achieve an accuracy of >> 50%). Interestingly, the Sudoku task particularly profits from predicted uncertainty, which heavily outperforms random sequentialization order and also the graph- based ordering. We provide the accuracy of combinatorial Sudoku solver using a greedy strategy (without backtracking) as reference (Oracle).
While the (inpainting) diffusion model that does fully paral- lel denoising of all patches is able to solve easy examples with a large number of given cells, its performance dete- riorates drastically with weaker conditioning. SRM with parallel sampling behaves similar, indicating that this sam- pling strategy is inappropriate for the Sudoku task with complex dependencies, independent of the network training. For all spatially autoregressive sampling methods, having an overlap between denoising intervals of patches sampled one after the other, decreases the performance compared to the full autoregressive extreme. This highlights the importance of spatial sequentialization for reasoning.
Masked Input
Sample
Sample
Sample
5 f
67 3
2
5 1 2
o
S
3 y
6
7 6
2.77 4
5737 y
6 5
/2537 3
S F43
2 12447663
3
7.755
11425594
4 (
P
7 7868y3 /
/
7
15158 69
R2
6
7
1623500
66
9773773333 74232
15234700
66 160/g9365
317926 57 A PETT29110
25 88 / 8 2 2
7 3
K K67 S 2
9 7
32 22627
37
3 374 6 29
9 A 7 7 4
Furthermore, our results clearly show that the order of se- quentialization matters. SRM with a random order achieves mediocre performance. As the model always sees the entire (noisy) Sudoku, the network learns that the distribution of the currently denoised patch does not only depend on all pre- viously denoised cells but also on the possible solutions for the random
93 2 7
3 5 S 6
6 1 4
29
7 2 9
S 2 22
5983 62
7 2 2
1362544 2
7 2 7
Figure 6: Sample Diversity. SRMs can sample multiple different, correct solutions given an incomplete observation.
all future ones, explaining the advantage oracle (which fails completely). Using task-specific knowl- edge about spatial dependencies in form of a general graph helps to significantly improve the order of sequentialization and as a result sampling from the correct distribution.
account other cells that are not yet denoised but already determined by applying the rules of Sudoku. We encourage the reader to verify this for themselves. In the end, we have case of multiple valid solutions, from which SRMs, being generative models, able to sample. We further demon- strate the sample diversity for multiple incomplete Sudokus in Fig. 6. We provide videos on our project website that visualize the full sampling process for multiple examples.
However, we achieve the best performance by a large margin with the task-agnostic order based on the predicted uncer- tainty as described in Sec. 3.2.2. An example for such sampling process is visualized in Fig. 7. We depict three pairs of close steps at the start, middle, and end of sampling. By following an exclusion procedure for the digit one in the middle right block using clean conditionings only, we understand why the model correctly predicted the lowest uncertainty for that patch the start. Moving to the middle case, the chosen cell is also completely determined by the conditioning, which only becomes visible if one takes into
Overall, we improve the correctness of samples in the hard setting from 0.8% to 51.6% (cf. Tab. 1) using the same trained model, only varying the sampling strategy.
4.4. Even Pixels Results
Tab. shows the quantitative comparison of our best settings for all sampling methods. We can see clear gap between

---

Spatial Reasoning with Denoising Models
Samplex x
Notse lovelt
Uncertainty sulk))
Single Stop x
Accuracy
0.20
93
0.0
0.2
0.4
JAG
Overlap
1576
Predicted Order
Random Order
6
S V
Figure 8: Even Pixel Accuracy for Different Overlap. Larger overlap is better for generating even number of pixels. However, when approaching the parallel setting at 1.0, the performance heavily drops.
J247 18SJa
1177
17878792
EV3
plify the task. This behavior can be further encouraged by lowering the sharpness of our noise level sampling dur- ing training (cf. Sec. 3.2.1), i.e., slightly biasing the SRM training towards spatially autoregressive generation.
7 1 192636622
1573 72144
13 675666 I
G 11V895794
3 07299/8
Unlike for the MNIST Sudoku experiment, we find over- lapping denoising variables be beneficial and further visualize this relationship in Fig. 8. For both predicted and random order, the accuracy gradually increases with higher overlap until sweet spot 0.95, after which it falls quickly. This shows that depending on the data distribution, mixtures of parallel and sequential generation enabled by SRMs can be advantageous. We can see a small, consistent advantage of the predicted uncertainty order over the random one.
7269 19539
1575 1.95222 22PN
13 138546721
0VV355994
7 77886/374
3 389276616
195J7
&
6 155111
175597 15596/324 7
319274614
4.5. Counting Polygons/ Stars Results
Figure 7: Sequentialization With Predicted Order. We
For the final Counting Polygons / Stars on FFHQ datasets, we provide the quantitative results for sample accuracy (ful- fillment dataset rule) and consistency of objects (same number of polygon vertices / star points) in Tab. 3. SRM outperforms the diffusion model for almost all sampling methods. This indicates potential of our training with indi- vidual noise levels, as this form of spatial disentanglement might be generally beneficial for certain data distributions.
visualize the sampling process with the current sample x noise level t, estimated uncertainty (darker = lower) ds (x(x*) and the single step to = O result x SRMs are able to reason over spatial variables by capturing complex depen- dencies (first two blocks) and handle uncertainty, e.g., in the case of multiple valid solutions in the last block.
Samling
#E
Acct
For the Counting Polygons dataset variant, there no clearly winning sampling strategy, unlike for the other two bench- marks. We attribute this dataset. First, using real images for backgrounds makes the task significantly more complex, as the denoising objective becomes less sensitive to the dependencies between num- bers and polygons and capacity of the denoising network spent for fitting the distribution of FFHQ faces. Secondly, the simultaneous generation of matching numbers and poly- gons can be approached in a coarse-to-fine manner with the numbers being more high-frequency details compared to larger owwffreuency polygons. For MNIST Sudoku, all numbers have the same sizes such that, during sampling with a parallel denoising strategy, the commitment to indi- vidual digits has to happen at similar points in time. Due to the spatial dependencies of Sudoku, this is suboptimal and a
Diffusion Model
1.270
0.250
main differences of the
Ours, Parallel
5,184
0.054
Ours, Predicted Order + Overlap
0.534
0,518
Ours, Random Order + Overlap
0.584
0.476
Table Even Pixels Experiments. SRMs clearly outper- form standard diffusion model in this task. Random order and predicted order perform similar, with slight advantage for the predicted one. See Fig, 8 for overlap analysis.
the diffusion baseline and SRM. While the diffusion model must learn to balance pixel colors evenly on a global level, sequentially eliminating differences in pixel counts can sim-
8

---

Spatial Reasoning with Denoising Models
Dataset
Polygons
Stars
Sampling / t-Sampling
Uniform
Ours
Sampling
Acct
Cont
Acct
Cont
Parallel
0.000
0.008
Diffusion Model
0.132
0.990
0.070
0.544
Predicted Order w/o Overlap
0.018
0.516
Ours, Parallel
0.166
0.956
0.034
0.576
Predicted Order + Overlap
0.008
0.422
Ours, Predicted Order w/o Overlap
0.144
0.994
0.076
0.844
Random Order w/o Overlap
0,000
0.024
Ours, Predicted Order + Overlap
0.142
0.986
0.150
0,888
Random Order + Overlap
0.000
0.020
Ours, Random Orler w/o Overlap
0.164
0.998
0.080
0.872
Graph-based Order w/o Overlap
0.012
0.266
Ours, Random Order Overlap
0.186
0.964
0.104
0.938
Graph-based Order + Overlap
0.010
0.238
Table 3: Counting Polygons/ / Stars FFHQ Experiments. Although this task is generally hard and not solved well, sequentialization is beneficial, especially for consistency.
Table 4: Noise Level Sampling Ablation. Ablation per-
formed on MNIST Sudoku hard. Our strategy for sampling t's during training is crucial for all strategies to work.
spatially autoregressive strategy together with good order can commit to the digits in cell-by-cell fashion. For the Counting Polygons dataset, the numbers are high-frequency details compared to larger polygons. As a result, coarse-to- fine generation with the diffusion baseline can first commit to a number of polygons of a certain vertex count and then generate matching digits.
denoiser architecture, noise schedule, and the sharpness hyperparameter (cf. Sec. 3.2.1) in Appendix G.
5. Conclusion
We introduced Spatial Reasoning Models (SRMs), a frame- work that allows to reason over sets of continuous, spatial variables. Our framework allows parameterization for sev- eral different strategies of sequentialization and generation order. Of particular interest is an automatic order prediction based on uncertainty. The SRM framework also is agnostic to the employed denoising formulation and parameteriza- tion, and allows to use wide variety of different denoising architectures. We introduced three benchmark tasks, which demonstrate that SRMs significantly outperform standard diffusion models in higher level reasoning tasks.
We further validate this hypothesis by replacing the low- frequency polygons with stars that are composed of higher- frequencies, moving the "point of commitment" to numbers and stars closer together in time. For this dataset variant, we observe that the accuracy of the diffusion baseline de- creases significantly compared to the version with polygons (7% vs 13.2%), while sequential sampling with overlap and predicted order maintains the same performance. More noticed hallucina-
interestingly, for parallel generation, tions of samples with stars having inconsistent numbers of points (cf. consistency column). For sequential sampling, the model can replicate stars after the first one has been generated, whereas in parallel sampling, the decisions over the number of points for all stars are again closer in time. We hypothesize that this behavior is not visible for poly- gons because of differences in terms of frequencies, with the diffusion model having more "time to correct itself" for low-frequency polygons, as their generation starts earlier than for the high-frequency stars.
Future Directions While we made significant steps for- ward, our models are still far away from solving the given tasks an optimal fashion. However, we believe that the presented paradigm is able to achieve even better results. particular topic for future research can be more advanced strategies for automatic prediction of generation order. It is an interesting venture to further investigate what makes a distribution difficult to generate for a diffusion model. A po- tential difficulty metric could help optimizing for an optimal order. Also, we speculate that backtracking- that allow to increase noise levels again hold much potential.
strategies
4.6. A Ablations
In Tab. 4, we provide the results of an ablation of the noise level sampling from Sec. 3.2.1 on the hard difficulty of the MNIST Sudoku dataset. Enforcing a uniformly distributed mean during training essential for all sampling methods. While independent uniform sampling of noise levels might be good enough for very small number of variables as in Diffusion Forcing (Chen et al., 2024a), high and low mean noise levels quickly become severely undersampled when we increase it. As these two cases represent start and end points at test time, the input to the network is immediately out of the training distribution during sampling. We pro- vide additional ablations w.r.t stochasticity of inference,
Impact Statement
This paper presents fundamental research with the goal of advancing reasoning capabilities of generative models. There are many potential societal consequences of advances in this direction further down the road, however, none are immediate enough to be specifically highlighted here.
A Acknowledgements
This work was partially supported by the Saarland/Intel Joint Program on the Future of Graphics and Media. We thank Thomas Wimmer for proofreading and helpful discussions.

---

Spatial Reasoning with Denoising Models
References
Lipman, Y., Chen, R. T. Q.. Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In International Conference Learning Representations, 2023.
Agarwal, N., Ali, A., Bala, M.. Balaji, Y., Barker, E., Cai,
T., et al. Cosmos world foundation model platform for physical ai, 2025. URL pettps://arxiee org/ abs/ 2501. 03575.
Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning generate and transfer data with rectified flow. In International Conference Learning Representations, 2023.
Bishop, C. M. and Nasrabadi, N. M. Pattern recognition and machine learning, volume 4. Springer, 2006.
Chen, B., Monso, D. M., Du, Y., Simchowitz, M., Tedrake, R., and Sitzmann, V. Diffusion forcing: Next-token pre- diction meets full-sequence diffusion. Advances in Neural Information Processing Systems, 2024a.
Nichol, A. Q. and Dhariwal, P. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning. 2021.
Chen, B., Xu, Z., Kirmani, S., Ichter, B., Sadigh, D., Guibas, L., and Xia, Spatialvim: Endowing vision-land -aniuag models with spatial reasoning capabilities. In Proceed- ings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp. 14455-14465, June 2024b.
Pearl, J. Reverend bayes on inference engines: distributed hierarchical approach. In Proceedings of the Second AAAI Conference on Artificial Intelligence, AAAI'22, pp. 133-136. AAAI Press, 1982.
Peebles, W. and Xie, S. Scalable diffusion models with transformers. In 2023,
Chen, L., Li, B., Shen, S., Yang, J., Li, C., Keutzer, K., Darrell, T., and Liu, Z. Large language models are visual reasoning coordinators. In Proceedings of the 37th Inter- national Conference on Neural Information Processing Systems, NIPS 233, Red Hook, NY, USA, 2023. Curran Associates Inc.
Plaat, A., Wong, A., Verberne, S., Broekens, J., van Stein, N., and Back, T. Reasoning with large language models, survey, 2024. URL https: /ttp:://axxvvvvoo /bbs/ 1157.11511.
Radford, A., Narasimhan, K., Salimans, T., and Sutskever, I. Improving language understanding by generative pre- training. 2018.
Dhariwal, and Nichol, A. Diffusion models beat gans on image synthesis. In Advances in Neural Information Processing Systems, 2021.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. Computer Vision and Pattern Recog- nition (CVPR), 2022.
Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F., al. Scaling rectified flow transformers for high-resolution image synthesis. In International Conference on Machine Learning, 2024.
Ruhe, D., Heek, J., Salimans, T., and Hoogeboom, E. Rolling diffusion models. In International Conference on Machine Learning, 2024.
Ho. J.. Jain, A., and Abbeel, P. Denoising diffusion prob- abilistic models. In Advances in Neural Information Processing Systems, 2020.
Seitzer, M., Tavakoli, A., Antic, D., and Martius, G. On the pitfalls of heteroscedastic uncertainty estimation with probabilistic neural networks. In International Confer- ence on Learning Representations, 2022.
Huang, J. and Chang, K. C.-C. Towards reasoning in large language models: A survey. Rogers, A., Boyd-Graber, J., and Okazaki, N. (eds.), Findings of the Association for Computational Linguistics: ACL 2023, pp. 1049-1065, Toronto, Canada, July 2023. Association for Computa- tional Linguistics.
Song, J., Meng, C., and Ermon, Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.
Kwan, M., Chow, K.-P., Law, F., and Lai, P. Reasoning about evidence using bayesian networks. In Advances in Digital Forensics IV, pp. 275-289, Boston, MA, 2008. Springer US
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. In Advances Neural Information Processing Systems, 2019.
van den Oord, A., Vinyals, O., and kavukcuoglu, k. Neural discrete representation learning. In Guyon, L., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S.,, and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 30, 2017.
Li, T., Tian, Y., Li, H., Deng, M., and He, K. Autore- gressive image generation without vector quantization. In Advances in Neural Information Processing Systems, 2024.
10

---

Spatial Reasoning with Denoising Models
Wei, J., Wang, X., Schuurmans, D., Bosma, M., Ichter, B., Xia, F,, Chi, E. H., Le, Q. V., and Zhou, D. Chain-of- thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems, 2022.
Wu, T., Fan, Z., Liu, X., Gong, Y., Shen, Y., Jiao, J., Zheng, H.-T., Li, J., Wei, Z., Guo, J., Duan, N., and Chen, W. Ar-diffusion: Auto-regressive diffusion model for text generation. In Advances in Neural Information Process- ing Systems, 2023.
Zhang, Z., Liu, R., Aberman, K., and Hanocka, R. Tedi: Temporally-entange diffusion for long-term motion syn- thesis. In SIGGRAPHI, Technical Papers, 2024. doi: 10.114556641599.3655555

---

Spatial Reasoning with Denoising Models
The appendix is structured follows. First, the detailed unified denoising framework is given in Sec. A. Then, graph-based sampling strategies given in Sec. B. It follows detailed description of the Uniform i sampling in Sec. C and dataset descriptions in Sec. D and Sec. E. The appendix concludes with further details about experimental setup in Sec. F. Please also take look at our project website with, among other things, videos showing Sudoku reasoning over time.
A. Generation Process
Notation Given noise schedule a, [0, + [0, 1] for continuous levels of noise E [o, 1]. we define a, == a(t), := b(t) for short notation. Furthermore, consider single (possibly higher-dimensional) continuous random variable for simplicity, write Ty for the variable at noise level t instead of using superseript as in the main paper.
Reverse Distribution Similarly to DDIM (Song et al., 2021), we define family of inference distributions indexed by the function Ø {(t", .) E [0 t" < t + Rzo with o == att", t) S by- for fixed-sizer schedule of continuous noise levels 4g E [0, 1] with S N. N N22, = 0, iN 1, and 4-11 < 44:
N
g (To. (m0) = to (x (20)
II
do (T [ 2o):
i=2
where 2 (2) (aa bfT and for all 0 tt t, 1:
(xt. 20) vo) N a0. by,
a1, To
V b3
oi
(8)
We show that the choice of the mean ensures the desired marginal distribution go (xe/2o) N Nar 1)
Proof. By construction, the statement holds for t = 1. Let < 1, then we have
go (2z NN(aro b11)
(9)
qo (2x- |2,, mo) N/ (0-10 + - -,00
19. of 2s o..))
(10)
b,
From (Bishop& & Nasrabadi, 2006) (2.115), have that go (ot- (20) is normal distribution N(, E) with:
ag. +
0,30 a,30
00
of
Zo
(11)
b3
2
Σ = o²_, I
r
1 b
by.
(12)
D2
Therefore, go (z) (r)) (a,TO, b holds for all t E 1].
Via simple variable substitution conditionings:
2a- and Zg se-
we obtain the corresponding distributions for alternative
ga (Ie lc, xo) = N (av. + / b3. oi or-,,0 or,aa)
(13)
(1p22 + bV. 13- of. ar-be yc, of oz,
g ( (o (c, = N
(14)
where we use Eq. 13 the posterior distribution approximated by minimizing the variational lower bound and Eq. 14 for the following definition of our generative process by replacing the ground truth epsilon with the prediction of the denoising neural network
(4.- + (v 63. of Or by )eg Mir.
q (ar 1: = =N
(15)
az
where Ex e be either fixed as in DDPM (Ho et al., 2020), e.g., equivalent upper Ex,, = b. Ex-, = of.,' depending on (=x)) being isotropic noise or a delta function, or learned by the network as an interpolation between these two optimized for the VLB, as proposed by (Nichol & Dhariwal, 2021) for DDPM. In all experiments the paper, we choose the last option.
b,- of _ and lower bounds
12

---

Spatial Reasoning with Denoising Models
Choice of For the choice of the posterior standard deviation oe .t. we follow DDIM (Song al., 2021) and consider the following special case. If and only if att,, t) = by vI (arbr, (() for z t* < t S 1, the forward process becomes Markovian, i.., 9o (2e xo) qo (ze (x-).
Proof. Using the abbreviation a for (t", 1), we have
2o ( = N Na
(16)
qr (1 [e, (n) = N (4:200
Ft -z apt0
V 12. 02, 0²,
(17)
bt
From (Bishop Nasrabadi, 2006) (2.116), have go (Ie)
-N 2:) with:
b
a
1)2
-1
170(t (t)
Σ
+
I=
(18)
bi
ttt,,))2
B7.
12
bi. d²
0,30
/
by
0,02
It
by
v
0²
a2
Iu
69
(19)
We can see that j becomes independent of xg for the case:
Vr 0²
(
11
01 0
b102
0
12
0t.
b7
(20)
at
(b)
0²)
ar-
1/- 0² +
at 0
(21)
b.02
0²
ae (b). 0²)
aa-bi b3g - σ² + 0,02 = 0
(22)
- a.bi. a by V by. - 0²
(23)
aebr a2- ar-
- 0 = br /
1
(24)
We define (t", : 17. by / 11111(a10 -aaabb-/(aa-b))) for 1) E [0, 1] as the variance of the posterior distribution. By setting = 0 we obtain deterministic sampling equivalent to DDIM with discrete diffusion noise schedules and equivalent to flow-based methods (Lipman et al., 2023) when solving the generative ODE with the Euler method. However, by setting
7 = 1 we enable stochastic sampling with general noise schedules (Gaussian probability paths in the context of flow matching), which has not yet been explored up to our best knowledge.
B. Graph-Sequential sampling
The graph-based sampling order is motivated by the idea that the dependency structure between patches can be exploited to estimate the level of knowledge about given patch based on the noise levels of patches connected to it. By leveraging these dependencies, the sampling process can be guided in a structured and efficient manner, ensuring that patches with stronger conditioning from their neighbours are prioritized during denoising.
The adjacency matrix is defined as A E (NN-M))(N-NM) where (N, MM are the dimensions of the image patch grid. In the case of Sudoku, this grid is of size (9. 9). The adjacency matrix is binary mask where all elements of the same row, column, or 3 3 subgrid are connected. Other elements are not connected (represented by value of 0 in the adjacency matrix).
When selecting the next patch to evaluate, we take arg maXpateh (1 t) A) C K, where t E [0, M is the current noise level vector, and K E (0, 1} is a mask with for all patches that didn't 't start denoising process and 0 otherwise. So we propagate 1 - t (denoising level) through the graph, and select a patch that has the noise level f; = 1, but has the strongest conditioning from its predecessors.
13

---

Spatial Reasoning with Denoising Models
B.1. Drawbacks
The certainty obtained using this algorithm is solely noise-level based. cannot distinguish whether a patch is being conditioned by the same number across its row, column, and subgrid (weak conditioning) by three different numbers (strong conditioning). As result, the model with predicted uncertainty (as described in Sec. 3.2.2) can surpass the performance of the graph-based approach (Tab. 1).
C. Uniform f sampling
For better intuition of why Uniform f should be one of our goals, we show the distributions of individual t, and f in Fig. 9.
Individua)
Mean
Mean
(a) Parallel
(b) Autoregressive
Figure 9: The distributions of individual patch noise level t, and mean over the image in parallel sampling (left) and autoregressive generation (right). In a parallel generation, all patches have the same noise level throughout inference, so both 4 and I marginal distributions are uniform. Autoregressive generation's 's 4 distribution is very dense around and 1, as during inference only one patch can have value other than 0 the one currently generated. It's important, that regardless of which inference strategy choose, the distribution of i remains uniform.
C.1. Recursive allocation sampling p(t F)
To sample p(t f), where t E [0, 1]d, we generate vector with a specified mean i by recursive allocation algorithm. The key idea is to define the total sum of the sampled vector as d f, then recursively partition into two sum contributions from the first and second halves of the vector. This process is illustrated in Algorithm 1.
Algorithm 1 Recursive Sampling of sum constrained Vector
Require: s (total sum), d (vector dimension)
Ensure: A vector x E [0. 1]d such that Σ rr
1: Function GerSumnConttraineeVeettoo d)
2: if dd= then
3: output [S] (End of recursion for 1 element]
4: end if
5: di [d/2] (Get Split dimension)
6: d2 d- di
7: sqaax + min(s, (s, di) (Define upper bounds of the contributions)
8: smax + min(s, (s, d2)
9: smin + max(0,
) (Define lower bound)
10: Sample r ~ Peppii (r(d), where E [0, 1] (Sample a split point)
81 f mmin + +s - mmnn
$2 --s; 81
13: output GeeSumCCnnttrineeWeccrrrs,,,,, d1) GerSSumConstrainedVVetoorr.
The algorithm first determines the range of possible contributions from the first half of the vector, (s7mnn,s7ar) Then samples split point from a distribution (r)da(r)d)) which has support on [0, 1]. The sampled value then scaled to the appropriate range of split points.
14

---

Spatial Reasoning with Denoising Models
The split distribution Peplie (d) is modelled symmetrical of Beta distribution, empirically unnd match the natural mid-point splitting behaviour of independently sampled uniform noise vectors. Explicitly.
Psplic (rld) = Beta(rla,,
(25)
where
3 (d mod 2)) sharpness.
(26)
uniformly sampled vector ((0, 11 as the number patches d increases, extreme split points become less likely. For example with 100 patches and = 0.5 $1 could be between 0 and 50 but is highly unlikely that all patches in the first half of the vector will be ones and all in the second half will be zeros. Instead, 81 is most likely be close to 25. That is why the split point distribution gets more center-heavy with higher vector dimensionality d d we model this by increasing the a and 3 parameters for higher d as in Eq. 26. For sharpness = 1 and i = 0.5 the distribution of t. E [0, 1j" closely resembles 2[00, I]") regardless of n. Higher sharpness results in distributions more tightly centred around É, while sharpness E (0. 1) leads to oversampling values closer to 0 and 1. The effect of the sharpness parameter on the distribution illustrated Fig. 10. Its influence on the Even Pixels dataset is presented in Tab.
sharpness 0.1 j 0.25
sharpness 0.1 ==
sharpness 0.1 f 0.75
sharpness 0.1 | 0.95
Density
0.25
0.25
0.00
0.0
0.00
0.0
sharpness=::3 0.25
sharpness 0.5 0.5| == 0.5
sharpness 0.5 / 0.75
sharpness 0.5 == 0.95
Densicy
0.2
0.10
0.2
0.50
0.1
0.03
0.1
0.25
0.00
sharpness 1.0 0.25
sharpness 1.0 1 0.5
sharpness
0.75
sharpness 1,0 i
Density
0.10
0.10
0.05
0.02
0.05
0.2
0.00
0.00
0.00
0.0
sharpness = 2.0 0.25
sharpness 2.0 0.5
sharpness 2.0 5 0.75
sharpness 2.0 0.95
Density
0.30
0.50
0.10
0.4
0.05
0.025
0.05
0.2
0.00
0.000
0.00
0.0
sharpness | i 0.25
sharpness = 5.0 0.5
sharpness 5.0
0.75
sharpness 5.0
0.95
Density
0.1
0.05
0.1
0.5
sharpness 10.0 0.25
sharpness 10.0 =
sharpness 10.0
0.75
sharpness 10.0
0.95
Density
0.10
0.1
0.1
0.5
0.0
0.00
0.0
0.0
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.30
0.75
1.00
Individual
Individual &
Individual
Individual e,
Figure 10: Distributions of 4; noise level values given different sharpness (ι-axis) and mean noise level T.
The advantage of this algorithm over, for example, iterative perturbation methods is that it can be easily parallelized across the batch dimension, ensuring efficient utilization of computational resources. Additionally, it exhibits stable execution time, a crucial property for maintaining predictable and consistent performance in the training pipeline.
15

---

Spatial Reasoning with Denoising Models
D. Even Pixels Evaluation
For evaluation, we generate a hue histogram from the model's output and identify the peak hue maxx. We then count the number of pixels closer to max than to complementary hue h (himax + 180%) mod 360%. The error is computed as the absolute difference between this count and half the total number of pixels
width height
error # (pixels closer hmax than H max)
2
The accuracy indicates in which fraction of samples where the pixel hue counts equal (error 0).
E. Counting Polygons/Stars Dataset
E.1. objects color selection
The color of the numbers and polygons/stars chosen HSV space hsv(c, 1.9,,.9), where
C = arg min ((histogram(Homg) * Ga) (h)
/t
and G. is Gaussian kernel with a small In other words, we search for the hue with the lowest count in the smoothed hue histogram of an image. We found this method as an intuitive approach to increase the visibility of the objects on the FFHQ background.
E.2. Evaluation
To assess the generative model's ability to learn the implicit counting constraint, for each of the object types (polygons or stars) we employ a ResNet-50-bbsee classifier with three prediction heads:
Predicting the set of two numbers (N, K)) present in the image.
Predicting the number of polygons/stars.
Predicting the number of object vertices.
For star objects, only count the convex, outward-facing vertices. The classifiers trained on a dataset with similar image compositions, but the consistency constraint between the numbers and the actual object counts is dropped in 50% of cases. For all the prediction tasks mentioned above, both classifiers have achieved over 99.9% accuracy in the validation set.
For evaluation, the classifier detects the numbers present in the generated image and compares them to the true object and vertex/spikes counts. A generated image is classified as correct the set of extracted numbers the same as the set of actual polygon/star objects and vertex counts (regardless of order): correct + {num,, num; = (#objects, #verices)..
The evaluation metric of the generative model, accuracy, is computed as the fraction of images for which the number-to- object-count correspondence correct.
F. Implementation Details
For reasoning in the spatial visual domain, we divide images into patches of a fixed size. While we use the size of a cell, i.e.,
MNIST example 28 x 28 as the patch size for the MNIST Sudoku dataset with image resolution 252 x 252, we choose patch sizes of 4 and 8 for the even pixel and counting polygons/stars datasets with resolutions 32 x 32 and 128 128. respectively.
Although SRMs are agnostic w.r.t. the architecture choice, we choose simple 2D UNets for all experiments, as they are widely established for image generation. We provide all architecture and training hyperparameters in Tab. 5.
For training the SRMs with individual noise levels per patch and the estimation of its uncertainty in the noise prediction, we additionally apply to simple architecture modifications. Instead of conditioning the UNet on a single encoding of the noise level f, we compute map with per-pixel noise levels, where all pixels within a patch share the same This map
16

---

Spatial Reasoning with Denoising Models
MNIST Sudoku
Even Pixels
Counting Polygons/Stars FFHQ
Channels
128
64
128
Depth
2
2
2
Channel multipliers
1, 1, 2, 2, 4,
1, 2, 2, 4
1, 1, 2, 2, 4
Head channels
64
64
64
Attention resolution
16, 8
8, 4
16, 8
Parameters
118M
19.7M
76,8M
Effective batch size
56
2048
56 250k
Iterations
250k
100k
Learning Rate
le-4
8e-4
le-4
Table 5: Hyperparameters used for all experiments.
Deterministic Acct
Stochastic
Sampling
L11
L1V
Acct
Parallel
26.288
0.004
19.156
0.010
Predicted Order w/o Overlap
3.652
0.500
3.212
0.516
Predicted Order Overlap
4.192
0.420
4.340
0.422
Graph-based Order w/o Overlap
6.212
0.240
5,816
0.266
Graph-based Order + Overlap
6.652
0.188
6.036
0.238
Table 6: Ablation over Deterministic vs. Stochastic Sampling on MNIST Sudoku hard. Stochastic sampling almost consistently outperforms deterministic sampling, highlighting the benefits of our generative process combining arbitrary (non-diffusion) noise schedules with stochastic sampling.
Sharpness
0.50
0.75
1.00
Sampling
#E
Accf
#Е
Accf
#H
Acct
Parallel
3.580
0.096
5.184
0.054
4.076
0.074
Predicted Order Overlap
0.840
0.368
0.534
0.518
0.590
0.484
Random Order + Overlap
0.776
0.342
0.584
0.476
0.656
0.430
Table 7: Sharpness Ablation on Even Pixels. By lowering the sharpness hyperparameter of our noise level sampling algorithm, we can train SRMs be more tailored towards sequential sampling methods.
bilinearly interpolated to the spatial resolution the feature maps in each layer, encoded using a shared MLP, and finally used to compute scale and shift maps applied pixel-wise the features. While this modification specific to UNets, similar adjustments can be made for other architectures like transformers, which already consider patches as tokens in the vision domain. In order to predict single uncertainty per patch, our denoising network outputs additional one-dimensional feature map, on which we apply 2D average pooling with the kernel size and stride corresponding to the patch size of our spatial variables. The result is interpreted as log-variance of the predicted noise.
G. Additional Ablations
We provide additional ablation studies considering hyperparameter choices for our experiments. Tab. 6 compares the performance of deterministic and stochastic sampling the hard difficulty of the MNIST Sudoku dataset by setting either 0 or 1 in the generation process described in Sec. A. We can see that stochastic sampling is favorable for spatial reasoning, which we enable in combination with arbitrary noise schedules.
In Tab. 7, we ablate the choice of the sharpness hyperparameter in our noise level sampling algorithm for training SRMs. On the Even Pixels dataset, we can oversample noise level combinations for spatial variables that are more likely in sequential sampling strategies by lowering the sharpness. This results in higher performance depending on the data distribution.
17

---

Spatial Reasoning with Denoising Models
Easy
Medium
Hard
Metric
Sampling
UNet
DiT
UNet
DiT
UNet
DiT
Parallel
0.998
0.992
0.590
0.328
0.010
0.000
Predicted Order w/o Overlap
0.998
0.988
0.754
0.618
0.516
0.248
Predicted Order + Overlap
1.000
0.992
0.716
0.598
0,422
0.166
Accuracy t
Random Order w/o Overlap
0.998
0.992
0.568
0.346
0.024
0.004
Random Order Overlap
1.000
0.974
0.542
0.330
0.020
0.004
Graph-based Order w/o Overlap
0.996
0.978
0.662
0,470
0.266
0.086
Graph-based Order Overlap
1.000
0.992
0.576
0.430
0.238
0.056
Parallel
0,012
0.044
3.312
10.940
19.156
39.048
Predicted Order w/o Overlap
0.012
0.068
1.616
3,960
3.212
8.620
Predicted Order + Overlap
0.000
0.060
1.816
4.216
4.340
11.244
LI Error
Random Order w/o Overlap
0.012
0.052
3.828
9.184
12.612
23.792
Random Order Overlap
0.000
0.172
4.172
10.324
13.896
25.368
Graph-based Order w/o Overlap
0.020
0.180
2.544
5.864
5,816
12.752
Graph-based Order + Overlap
0.000
0.072
3.240
7.100
6.036
14.912
Table 8: Denoising Architecture Ablation on MNIST Sudoku. We compare SRM variantes using UNet and DiT denoising backbones. For each metric, model, and difficulty, the best-performing configuration is highlighted bold. Sequential generation with predicted order yields the strongest performance across both metrics independent of the used architecture.
Easy
Medium
Hard
Metric
Sampling
Linear
Cosine
Linear
Cosine
Linear
Cosine
Parallel
0.998
1.000
0.590
0.622
0.010
0.018
Predicted Order w/o Overlap
0.998
1.000
0.754
0.748
0.516
0.378
Predicted Order Overlap
1.000
0.998
0.716
0.706
0.422
0.298
Accuracy f
Random Order w/o Overlap
0.998
1.000
0.568
0.502
0.024
0.036
Random Order Overlap
1.000
0.994
0.542
0.510
0.020
0.024
Graph-based Order w/o Overlap
0.996
1.000
0.662
0.630
0.266
0.220
Graph-based Order + Overlap
1.000
0.996
0.576
0.620
0.238
0.210
Parallel
0.012
0.000
3.312
3.048
19.156
17.140
Predicted Order w/o Overlap
0,012
0.000
1.616
1,812
3.212
5.140
Predicted Order Overlap Random Order w/o Overlap Random Order Overlap
0.000
0.012
1.816
2.056
4.340
5.920
L1 Error $
0.012
0.000
3.828
4.812
12.612
13.232
0.000
0.032
4.172
5.020
13.896
14.200
Graph-based Order w/o Overlap
0.020
0.000
2.544
2.752
5.816
6.592
Graph-based Order + Overlap
0.000
0.032
3.240
2.924
6.036
6.848
Table 9: Noise Level Schedule Ablation MNIST Sudoku. We compare SRM variants using Linear (Liu al., 2023) and Cosine (Nichol & Dhariwal, 2021) noise level schedules. The winning schedule is highlighted in bold for each row. While the linear schedule clearly outperforms the cosine one for our setting, the benefits of sequentialization with a meaningful order remain independent of the choice of the noise level schedule.
To further demonstrate the architecture--anostic nature of the SRM framework, we also trained Diffusion Transformer (DiT B) (Peebles & Xie, 2023) with patch size of 7 and 130M parameters compared to those of 2D UNet are presented in Tab. 8. Notably, in this highly structured domain, sequential generation with predicted variable order yields the best performance ndependent of the used architecture.
MNIST Sudoku benchmark. The results,
As an alternative to the linear schedule used Rectified Flows (Liu et al., 2023), we also explored the cosine noise level schedule (Nichol & Dhariwal, 2021). The results are presented in Tab. 9. Independent of the noise level schedule, sequentialization with predicted order consistenly outperforms all other sampling strategies. Interestingly, the linear schedule
18

---

Spatial Reasoning with Denoising Models
seems be more suitable than the cosine one for the dataset hand. We suspect that this is because the cosine schedule allocates more time steps low-noise regions, effectively focusing the denoising process the high-frequency details, which are not relevant for the digit identity and thus not captured by our metrics.
H. Sudoku Completion with LLMs
To evaluate the difficulty of generation-based Sudoku solving, we examined the performance of general-purpon large language models (LLMs). Specifically, tested GPT-4o, GPT-4o-mini, and open-souree Phi-4 on text-based Sudoku completion task. Each model was given textual representation of a Sudoku grid with missing elements and was tasked with filling in the blanks complete the puzzle. The prompts included the full context of the Sudoku task: description of the rules, output format, and three fho- examples. An example is shown below:
[Example 1]
Input:
<sudoku>
34_ 179| 258
187 523 | 964
529 648 371
965 832 417
472 11] 835
13 754
629
798 261 _43
6381485| 792
1499 18_
</sudoku>
Output:
<sudoku>
346 179 258
187 523] 964
5291 648 371
965 832
417
472 916
835
813 754
629
798 261 | 543
631 485
792
254 397 186
</sudoku>
If a model's output deviated from the required format, we resampled until a well-formed 9x9 grid of digits suitable evaluation was produced. The model was then asked to provide only the full solution, with no "inner thoughts". This yielded an autoregressive generation variant similar to our image-based generation approach, but in the text domain.
suitable for
Table 10 compares the performance of the image-based Diffusion and SRM models with the few-shot LLMs across different difficulty levels. Although this comparison is not fair whereas LLMs are few-shot prompted via text - it highlights the difficulty of solving Sudoku via generative models.
models are trained specifically for the task, in image space,
Figure 11 further breaks down model performance into validity, consistency, and overall accuracy as the number of masked cells increases. Validity measures whether the output follows the Sudoku rules, consistency ensures that the solution matches the initial grid, and accuracy requires both. Interestingly, while performance drops with increasing difficulty, accuracy improves when most of the board masked, possibly due to resemblance with the few-shot examples in the prompt.
19

---

Spatial Reasoning with Denoising Models
Sudoku Difficulty
Method
Modality
Setup
Easy
Medium
Hard
Diffusion
0,994
0.536
0.008
Image
Trained
SRM (Ours)
0.998
0.754
0.516
Phi-4
0.038
0.000
0.000
GPT-4o-mini
Text
Few-shot
0.205
0.000
0,000
GPT-40
0.556
0.001
0.011
Table 10: SRM LLM. To highlight the complexity of our MNIST Sudoku benchmark, we compare Sudoku solving capabilities of image generators with general-purpose LLMs used in a few-shot setting. Note that this is not a fair comparison, as SRM and Diffusion are trained directly for the task. Despite tackling discrete and therefore simpler version of Sudoku, LLMs perform poorly in this task compared to SRM.
Validity vs. number of masked cells
Consistency vs. number of masked cells
Accuracy vs. number of masked cells
1.0
Model name
1.0
Model name
1.0
Model name
Phi-4
Phi-4
Phi-4
0.8
Prompt consistency
0.8
0.8
Sudoku validity
GPT-4o mini
GPT-4o
0.6
0.6
GPT-4o
Accuracy
0.6
GPT-4o
0.4
0.2
0.2
0.2
0.0
0.0
0.0
20
40
80
20
40
60
80
20
40
60
80
Number masked cells
Number masked cells
Number of masked
(a) Validity: proportion of outputs that proper Sudoku grids (no duplicate digits).
(b) Consistency: solutions that match the (c) Accuracy: solutions that are both valid initially masked grid entries.
and consistent.
Figure 11: LLM Performance in Discrete Sudoku Completion. As the number of masked cells increases corresponding an increasing difficulty level of the Sudoku completion task, both validity and consistency drop. Surprisingly, when almost all elements are masked, the accuracy improves possibly due to similarity to some example grids in the prompt.
We additionally experimented with a chain-of-thought setup, where the model was prompted to iteratively fill one cell at time, selecting its generation order. However, consistency with the input grid was insufficient for a reliable evaluation. We also tested the Deepseek-RI reasoning model, allowing unconstrained internal reasoning before producing a final output. Although the model attempted to verify its own solutions, frequently entered loops of self-correction or made mistakes while evaluating the constraints, resulting in unfinished rollouts invalid Sudokus. Overall, the results suggest that solving Sudoku, even in simplified textual form, remains challenging task for seneral-purpose generative models.
20


Reviews:
[{"review_id": "4VC9rXFUuN", "reviewer": "Reviewer_7X1h", "summary": "This paper studies how diffusion models perform on higher-level reasoning tasks, such as the Sudoku game. The authors introduce a novel SRM framework to integrate several key improvements for semanticalization in generation, the associated order, and the sampling strategies. The experimental results are encouraging to some degree. The key idea of this work is to train a spatial reasoning models to jointly denoise multiple variables but with individual noise levels, which is inspired by the previous diffusion forcing work. The authors futher introduce a two-step sampling strategy to overcome the", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your valuable feedback. To address your concerns, we present a point-to-point response in the following:\n\n> Why not use the latest VLMs [...] to address this symbolic visual reasoning task?\n\nThank you for this important question. Our response is threefold:\n\n1. The goal of our paper is not to propose the best-performing method for individual visual reasoning tasks from all possible approaches including VLMs. We aim for **evaluating and improving reasoning capabilities of denoising generative models for continuous, spatial domains**. Just like for LLMs, hallucinations are an issue, especially in the case of complex spatial dependencies. Our different strategies can reduce them by a significant amount. We will make sure to resolve any lack of clarity in the final version.  \n     \n2. **TLDR: We tested LLMs and VLMs. They do not naively solve Sudoku.**  \n   We agree that “Sudokus can be transformed into either a text-only mathematical reasoning task or an image-based math-VQA task.” To evaluate LLMs for this, we conducted an experiment for Sudoku solving with GPT-4o, GPT-4o-mini, and the open-source Phi-4, given full context about the game of Sudoku including its rules, the required output format, and few-shot completion examples (similar to the text-only evaluation of the mentioned Sudoku-Bench). As we noticed deviations from the required format, we resample until obtaining a 9x9 grid of numbers that can be evaluated. Please check out the following figure with quantitative results for 10 samples per number of masked cells: [https://figshare.com/s/b22f205dd69a815d2b89](https://figshare.com/s/b22f205dd69a815d2b89) . While SOTA LLMs (GPT-4o) are able to correctly complete Sudokus with up to 10 missing numbers, their performance quickly deteriorates for more cells to fill. With increasing difficulty, they further start to violate the completion task or in other words cheat by overriding given cells. The following table compares the accuracy of diffusion, our SRM, and LLMs:  \n   \n| Method | Easy | Medium | Hard |\n| :---- | :---: | :---: | :---: |\n| Diffusion | 0.994 | 0.536 | 0.008 |\n| SRM (Ours) | 0.998 | 0.754 | 0.516 |\n| GPT-4o-mini | 0.205 | 0.000 | 0.000 |\n| GPT-4o | 0.556 | 0.001 | 0.011 |\n| Phi-4 | 0.038 | 0.000 | 0.000 |\n\nPlease note that this is not a fair comparison. While we train the diffusion model and SRM for **inpainting of continuous** **visual Sudokus**, the LLM evaluation is a case of **few-shot discrete text completion**. Despite the simpler discrete representation instead of grids of MNIST numbers, LLMs perform poorly in Sudoku completion. Since image-based math-VQA with VLMs can be considered to be an even more difficult task, as it adds correct (implicit) discretization as a step before reasoning, we argue that (visual) Sudoku is not a “naive task at which VLMs excel”. We further support this by showing qualitatively that current systems like ChatGPT (Pro) are not able to perform the full image-to-image task of visual Sudoku (even with few-shot examples): [https://figshare.com/s/dd4d6ad9a281cad7bacf](https://figshare.com/s/dd4d6ad9a281cad7bacf) \n\n   \n\n3. The same goes for the Counting Polygons FFHQ dataset. Please note that the task is not to count polygons and their vertices in input images, but to **generate images** that correctly follow the rules of the data distribution.\n\n> Latest study (Sudoku-Bench) on visual reasoning with VLMs.\n\nThank you for the pointer to this interesting benchmark. We will add a discussion to the final version. Please note the following:\n\n1. The benchmark comprises puzzles with different rules that can be leveraged for evaluation of LLMs and VLMs but not for training and testing continuous denoising generative models.  \n2. Just like for our additional experiments with LLMs above, the motivation and experimental setup is different in terms of the representation for reasoning (discrete vs continuous) and in-context learning with given rules (L/VLMs) versus fitting of the correct distribution given training samples only and notably no explicit puzzle rules (SRMs).  \n3. The blog post and GitHub repository were published on the 21st of March 2025, i.e., 4 days before the beginning of the rebuttal phase. Therefore, we hope this will not be considered to be a weakness of our paper.\n\n> “Diffusion Forcing scheme”\n\nOur approach does not simply adopt Diffusion Forcing (DF). We propose a general framework for reasoning over sets of continuous random variables with many possible applications such as spatial domains, but also temporal sequences, which DF is limited to. While the order of sequentialization is fixed to be the temporal order in DF, we propose task-specific (based on a given dependency graph) as well as task-agnostic orders leveraging predicted uncertainty. Moreover, we propose a technique for noise level sampling with larger numbers of variables, whereas the naive approach from DF fails completely in such a setting.", "claims_and_evidence": "The key claims of this work are in introducing a SRM (Spatial Reasoning Model) to perform spatial reasoning given human-designed images, following principles such as Sudokus, counting pixels, or counting polygons. It requires advanced visual reasoning to address these tasks reliably. The authors only study the limitations of applying conventional parallel diffusion models and apply a diffusion forcing-like scheme to overcome these challenges.\n\nOne fundamental limitation of this paper is the question of why not use the latest VLMs, such as GPT-4o, Gemini-2.0, or other open-source VLMs, to address this symbolic visual reasoning task?", "methods_and_evaluation_criteria": "The major contribution of this work is studying a non-trivial visual reasoning task using generative models such as diffusion or flow matching techniques. A major concern is why not use an LLM or VLM to address such visual reasoning tasks, as Sudokus can be transformed into either a text-only mathematical reasoning task or an image-based math-VQA task. It is confusing for me to understand why the diffusion forcing scheme was chosen to perform denoising on the masked regions.", "theoretical_claims": "The authors provide a thorough theoretical analysis of the Generation Process, Graph-Sequential Sampling, Uniform Sampling, and Recursive Allocation Sampling; however, I have not carefully checked their correctness.", "experimental_designs_or_analyses": "The experiment section of this paper is relatively weak and far from convincing, although the spatial reasoning task is a valuable topic. The authors are encouraged to conduct more ablation experiments to analyze why the diffusion forcing scheme and the proposed improvements are important, and how the denoising scheme can understand and address a visual reasoning task like Sudoku.\n\nThe current experiments are too superficial and cannot provide deep insights for the community.", "relation_to_broader_scientific_literature": "The authors are encouraged to extend the proposed approach to study more challenging visual reasoning tasks at which modern VLMs perform poorly, rather than focusing on naive tasks at which VLMs excel.", "essential_references_not_discussed": "The authors are required to provide a discussion of related work on the latest study [1] that attempts to use the latest VLMs to address the constructed visual reasoning tasks.\n\n[1] https://github.com/SakanaAI/Sudoku-Bench", "other_strengths_and_weaknesses": "The major concern of this paper is why the diffusion forcing techniques are used to address a visual reasoning task that existing VLMs might excel at. As mentioned earlier, why not use an LLM or VLM to address such visual reasoning tasks, as Sudokus can be transformed into either a text-only mathematical reasoning task or an image-based math-VQA task? It is confusing for me to understand why the diffusion forcing scheme was chosen to perform denoising on the masked regions.", "other_comments_or_suggestions": "No other comments.", "questions_for_authors": "Please refer to the aforementioned weaknesses, and I would like to see how the authors address the major concerns regarding the motivation for using diffusion forcing to tackle the visual reasoning task.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hNWJUsoajY", "reviewer": "Reviewer_G4qT", "summary": "This paper introduces Spatial Reasoning Models (SRMs), a framework for performing reasoning over sets of continuous variables using denoising generative models. The authors observe that standard diffusion/flow models often collapse to hallucination when handling complex distributions. The key innovations include: enabling different noise levels for different spatial variables, a novel \"Uniform $\\bar{t}$\" noise sampling strategy that better aligns training and inference distributions, uncertainty estimation to guide the generation process, and various sequentialization strategies.\n\nThe paper introduces three benchmark tasks to evaluate reasoning capabilities: (1) MNIST Sudoku, where models must complete partially filled Sudoku grids using MNIST digits; (2) Counting Pixels, requiring balanced color distribution; and (3) Counting Polygons, testing understanding of relationships between numbers and visual elements. Through extensive experimentation, the authors demonstrate that sequentialization in generation and uncertainty-based ordering significantly improves reasoning capabilities, increasing accuracy from <1% to >50% on hard reasoning tasks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your constructive comments and are happy to see that you value the *originality* of our *novel ideas* achieving *significant improvements [...] highlighting the effectiveness of the proposed approach for spatial reasoning tasks*.  \n> W1. Real-world applicability \n\nWe agree that there is a gap between our benchmark and real-world reasoning scenarios. We designed our datasets to measure hallucination of denoising generative models, which is extremely difficult for real-world applications and already non-trivial for our counting polygons dataset. We share Reviewer FTjC’s opinion about it being “suitable for an early conceptual investigation as this one” and would like to investigate real-world applications in future work. Our framework enables everybody to perform research on a wide range of domains.\n> W2 + Q3. Validation of noise schedules\n\nWe provide results for models trained on MNIST Sudoku with the cosine noise schedule commonly used in diffusion or flow-based models: [https://figshare.com/s/8b72019bc22a66bce0c8](https://figshare.com/s/8b72019bc22a66bce0c8)   \nAll our conclusions regarding the benefits of sequentialization with a meaningful order hold for the cosine schedule too. We add this ablation to the final version.\n> W3. Architecture limitations\n\nWe provide results for models using a diffusion transformer (DiT B) with patch size 7 and 130M parameters, which roughly corresponds to the size of our 2D UNet: [https://figshare.com/s/2d87cf6657c1a948347d](https://figshare.com/s/2d87cf6657c1a948347d)   \nAll our conclusions hold for the DiT architecture too. However, we see generally worse performance than with the UNet, which we attribute to DiTs being less suited for denoising in pixel space. We will add this ablation to the final version.\n> W4 + Q1. ResNet Accuracy\n\nThe classifier used for the Counting Polygons (and Stars, see next paragraph) evaluation has an accuracy of \\>99.9% on a validation split such that reported metrics are reliable.\n> W5 + Q2. Inconsistent improvements\n\n1. We attribute this to fundamental differences in the tasks. For MNST Sudoku, all numbers have the same sizes such that, during sampling with a parallel denoising strategy, the commitment to individual digits has to happen at similar points in time. Due to the spatial dependencies of Sudoku, this is suboptimal and a spatially autoregressive strategy together with a good order can commit to the digits in a cell-by-cell fashion.   \n   For the Counting Polygons dataset, the numbers are high-frequency details compared to larger polygons. As a result, a coarse-to-fine generation with the diffusion baseline can first commit to a number of polygons of a certain vertex count and then generate matching digits.  \n     \n2. To further validate our hypothesis, we conducted additional experiments on a modified benchmark version “Counting Stars”, for which we replace polygons with stars (cf. [https://figshare.com/s/6d8b59fd96f56f05b571](https://figshare.com/s/6d8b59fd96f56f05b571)). The motivation is that stars are composed of higher frequencies, moving the “point of commitment” to numbers and stars closer together in time.\n\n| Sampling | Counting Accuracy | Star Consistency |\n| :---- | :---- | :---- |\n| Diffusion Model | 0.070 | 0.544 |\n| Ours, Parallel | 0.034 | 0.576 |\n| Ours, Predicted Order w/o Overlap | 0.076 | 0.844 |\n| Ours, Predicted Order + Overlap | **0.150** | 0.888 |\n| Ours, Random Order w/o Overlap | 0.080 | 0.872 |\n| Ours, Random Order + Overlap | 0.104 | **0.938** |\n\nFor the diffusion baseline, the accuracy of matching numbers and stars decreases significantly compared to the version with polygons (7% vs 13.2%), while sequential sampling with overlap and predicted order maintains the same performance. More interestingly, for parallel generation, we noticed hallucinations of samples with stars having inconsistent numbers of points (cf. star consistency column). For sequential sampling, the model can replicate stars after the first one has been generated, whereas in parallel sampling, the decisions over the number of points for all stars are again closer in time. We hypothesize that this behavior was not visible for polygons because of differences in terms of frequencies, with the diffusion model having more “time to correct itself” for low-frequency polygons, as their generation starts earlier than for the high-frequency stars.\n> W6 + Q4. Computational cost \n\nIn all our experiments, we set the total number of denoising steps to 1000\\. This means that the computational cost is equal for all sampling methods resulting in a fair comparison. This also means that individual variables have fewer steps the “more sequential” the sampling is (lower overlap) but it still performs better. We will thoroughly describe this aspect in the final version.\n> W7. Data Statistics \n\nThank you for pointing out this missing detail that we will add to the paper. For all evaluations, we sample 500 data points to compute metrics.", "claims_and_evidence": "Most of the claims made are supported by various empirical evidence through different experimets on the introduced benchmarks. However, there are few claims that lack adequate evidence or contain some inconsistencies.\n\n1. While the paper shows improvement on all three benchmarks, the performance on the most realistic task (Counting Polygons FFHQ) is modest (18.6% vs 13.2% baseline). The paper doesn't sufficiently demonstrate that these improvements would translate to more complex real-world reasoning scenarios. \n\n2. The paper claims their formulation works with various noise schedules (DDPM, rectified flows, etc.) in Section 3.1.1, but doesn't provide comparative results across different schedules to validate this.\n\n3. In Section 4.2, the authors state \"While SRMs are agnostic to different architectures,\" but provide experiments only with 2D UNets. \n\n4. For the Counting Polygons FFHQ evaluation, the authors rely on a ResNet classifier to determine correctness, but don't report this classifier's accuracy, making it impossible to assess the reliability of the reported metrics. Since the authors don't report the accuracy of this ResNet classifier itself, we don't know how reliable the evaluation is. If the classifier makes errors in detecting numbers or counting polygons/vertices, the reported model performance metrics ( for example reported numbers like 18.6%, 13.2%) would be affected.\n\n5. The paper doesn't adequately explain why sequentialization yields dramatic improvements on Sudoku (>50%) but only modest gains on Counting Polygons (~5%).", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "I checked the theoretical claims and proofs throughout the paper and its appendices, including the mathematical formulation in Section 3, Appendix A and appendix C. I did not find any notable mathematical errors in the theoretical claims and proofs throughout the paper. However, I do not consider myself as an expert in the domain, so I might have overlooked some detail.", "experimental_designs_or_analyses": "I examined the experimental designs for the three benchmarks (MNIST Sudoku, Even Pixels, Counting Polygons) and ablation studies. The Counting Polygons experiment relies on a separately trained classifier for evaluation, but the paper doesn't report this classifier's accuracy, introducing potential measurement bias. This issue is discussed earlier.", "relation_to_broader_scientific_literature": "The paper extends chain-of-thought reasoning concepts from language models to continuous spatial domains, while building upon sequential generation techniques from AR-Diffusion and Diffusion Forcing.The paper adequately covers all the relevant work.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "**Strengths**\n\n**S1. Originality**: The paper introduces several novel ideas including different noise levels for different spatial variables, a \"Uniform $\\bar{t}$\" noise sampling strategy that better aligns training and inference distributions, uncertainty estimation to guide generation, and various sequentialization strategies.\n\n**S2. Benchmarking**: The proposed benchmarks provide a systematic way to evaluate and quantify reasoning capabilities and hallucination in generative models.\n\n**S3. Experimentation and Performance**: The paper demonstrates significant improvements over baseline diffusion models (from <1% to >50% accuracy on hard Sudoku), highlighting the effectiveness of the proposed approach for spatial reasoning tasks.\n\n\n\n**Weaknesses**\n\n**W1. Real-world applicability**: While the paper shows improvement on all three benchmarks, the performance on the most realistic task (Counting Polygons FFHQ) is modest (18.6% vs 13.2% baseline). The paper doesn't sufficiently demonstrate that these improvements would translate to more complex real-world reasoning scenarios.\n\n**W2. Limited validation of noise schedules**: The paper claims their formulation works with various noise schedules (DDPM, rectified flows, etc.) in Section 3.1.1, but doesn't provide comparative results across different schedules to validate this.\n\n**W3. Architecture limitations**: In Section 4.2, the authors state \"While SRMs are agnostic to different architectures,\" but provide experiments only with 2D UNets.\n\n**W4. ResNet Accuracy**: For the Counting Polygons FFHQ evaluation, the authors rely on a ResNet classifier to determine correctness without reporting this classifier's accuracy, making it impossible to assess the reliability of the reported metrics (e.g., 18.6%, 13.2%).\n\n**W5. Inconsistent improvements**: The paper doesn't properly explain why sequentialization yields dramatic improvements on Sudoku (>50%) but only modest gains on Counting Polygons (~5%).\n\n**W6. Computational overhead**: There is a significant computational cost associated with sequential generation compared to parallel approaches, which the authors should have addressed more thoroughly.\n\n**W7. Data Statistics**: Details about the number of data points used in evaluating each benchmark is reported. Authors should consider adding some detail about them.", "other_comments_or_suggestions": "- Line 603: VLB should be written in the full form before being used as short form\n- Figure 7: the x-axis label \"Overlap\" would be clearer if it specified \"Overlap Ratio\" or something similar.\n- The term \"sharpness\" is used in Section 3.2.1 without proper definition until much later in the appendix.", "questions_for_authors": "Q1: For the Counting Polygons FFHQ benchmark, what is the accuracy of the ResNet classifier used for evaluation?\n\nQ2: What explains the dramatic improvement on Sudoku (\\~ 50% gain) compared to the modest gains on Counting Polygons (~5%)? Is this due to fundamental differences in the tasks, model capacity limitations, or other factors? \n\nQ3: Section 3.1.1 claims compatibility with various noise schedules (DDPM, rectified flows, etc.). Have you empirically verified performance across different schedules?\n\nQ4: Could you quantify the computational cost difference between the parallel baseline and your best sequential approaches (e.g., inference time for Sudoku)?\n\nI am open to increasing my score after satisfactory answers to all my questions and previously raised concerns.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sDkfutweLN", "reviewer": "Reviewer_KdD5", "summary": "This paper introduces “Spatial Reasoning Models” (SRMs), a framework for performing high-level reasoning across sets of continuous variables in diffusion/flow-based generative models. By allowing each spatial variable (e.g., an image patch) to have its own noise level, SRMs can systematically add or remove noise in different regions of an image in a sequential or partially sequential manner. This approach is shown to reduce “hallucinations” common to standard diffusion models on tasks such as Sudoku with MNIST digits, balancing pixel colors, or matching the digit count of polygons\n\n## update after rebuttal", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your review. We address your concerns as follows:\n\n> Justification for uniform distribution of mean noise level during training\n\n1. We would like to refer you to Fig. 8 of our paper’s Appendix. It shows that for the two extreme cases of parallel and autoregressive generation, the observed mean noise levels during inference form a uniform distribution. For all cases in between these, like our sequential sampling with overlap, it therefore also makes sense to train the model exactly for this mean noise level distribution. However, since this is an intuitive rather than rigorous theoretical justification, we are willing to reformulate this precisely in the paper to avoid using the term “ideally”.   \n2. Moreover, previous works like StableDiffusion3 have investigated different noise level weightings such as a logit normal distribution to oversample intermediate noise levels, which can significantly improve performance in certain tasks like image generation. Our two-step noise level sampling strategy during training is directly compatible with such weightings by replacing the uniform distribution for the mean with the logit normal distribution, for example. We will make sure to clarify the flexibility w.r.t. the choice of distribution for the mean noise level in the final version and would like to further investigate this direction in future work.\n\n> Role of Appendix A\n\nThank you for pointing out this lack of clarity.\n\n1. We agree that many previous works such as \\[1, 2, 3\\] introduce the concrete formulations of diffusion, flow matching, or unifying frameworks. Our Appendix A does not reinvent the wheel, but shows that the DDIM \\[1\\] formulation can be extended to noise schedules introduced in the context of flow matching. By considering Gaussian reverse distributions as in DDPM / DDIM instead of ODEs and SDEs with flow-based models, we explicitly show that improvements of diffusion models such as a learned variance \\[4\\] can be combined with non-diffusion noise schedules like from rectified flows. To the best of our knowledge, there have been no flow-based formulations incorporating a learned variance so far. We will clearly specify the role of Appendix A in the final version.\n\n- \\[1\\] Denoising Diffusion Implicit Models. ICLR 2021  \n- \\[2\\] Flow Matching for Generative Modeling. ICLR 2023  \n- \\[3\\] Stochastic Interpolants: A Unifying Framework for Flows and Diffusions. arxiv 2023  \n- \\[4\\] Improved Denoising Diffusion Probabilistic Models. ICLR 2021\n\n2. Please note that we have chosen not to claim this as one of our main contributions (cf. contributions and key findings in introduction), as we are aware of the many prior works introducing the theoretical formulations of denoising generative models. Our SRM framework can benefit from any improvements of the diffusion process for single continuous random variables, as our formulation of spatial reasoning is orthogonal to that.\n\n> Confusion w.r.t. spatially autoregressive approaches\n\nYou are correct that sequentialization **without overlap**, utilizing random, predicted uncertainty or manual graph-based order, results in spatially autoregressive patch-by-patch generation. Our framework enables the training of one model that can be used in combination with multiple different sampling strategies including parallel generation (synchronous denoising of all variables) as with usual diffusion models, spatially autoregressive generation, and everything in between, e.g., sequentialization **with overlap**. In addition to the amount of sequentialization as one investigated degree of freedom (cf. Fig. 3), we propose and evaluate different orders, in which variables (patches) are chosen for denoising, with the uncertainty- and graph-based orders being novel w.r.t. prior works for denoising-based spatially autoregressive generation. Our paper shows that, for the same trained model, different sampling strategies significantly impact the level of hallucination and that the best strategy depends on the data distribution, e.g., non-overlapping sequentialization for MNIST Sudoku and a high overlap for Even Pixels.  \nWe will make sure to clarify this in the final version.", "claims_and_evidence": "- In line 190, the paper suggests that the mean noise level ( \\bar{t} ) should ideally follow a uniform distribution with a sufficient number of inference steps. Is there a rigorous theoritical justification for it?", "methods_and_evaluation_criteria": "- What's the role of Appendix A? If it is about the diffusion process for the single (possibly higher-dimensional) continuous random variable,   the concrete formulation and proofs have already been well-defined in previous works.\n- The difference between the proposal and spatially autoregressive approaches is really confusing. If I understand correctly, when the proposed algorithm utilize random, predicted uncertainty, or manual graph-based order of sequentialization, it still generates each patch one by one.", "theoretical_claims": "- I think the paper does not provide any meaningful theoritical claims.", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "Reasoning in diffuison generation in high-depent multi-variables.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "No", "other_comments_or_suggestions": "No", "questions_for_authors": "No", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MEImCMTwQZ", "reviewer": "Reviewer_FTjC", "summary": "The authors investigate the application of diffusion models as solvers of probabilistic inference over continuous variables, which accommodates various problems in spatial reasoning.  A key consideration is the various possible decompositions of the joint distribution of unobserved variables, and how some decompositions may be beneficial.  Another consideration is the noise schedule used traditionally in diffusion models, being a fixed increment, and the drawbacks of diffusion forcing leading to undertraining for early/late inference steps.  By tackling those two aspects in tandem, the authors propose to explore the design space spanning parallel generation to autoregressive generation, proposing a \"recursive allocation sampling\" algorithm with controllable sharpness, making it adaptable to different inference scenarios.  The benefits of the proposed scheme is demonstrated on a challenging MNIST-Sudoku benchmark, with 3 difficulty levels, as well as an even-coloring experiment.  Finally, the authors attempt to highlight more realistic settings by counting a small number of geometric primitives overlayed on backgrounds sampled from the FFHQ dataset, highlighting some of the remaining challenges.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you very much for your time and positive feedback. We are happy to see that you agree with us that reasoning over complex domains is certainly relevant in science.\n\n> Missing Proposition, Lemma, or Theorem in Appendix A\n\nThank you for pointing out this lack of clarity. On a high level, the Appendix A shows that the DDIM \\[1\\] formulation can be extended to noise schedules introduced in the context of flow matching \\[2\\]. By considering Gaussian reverse distributions as in DDPM / DDIM instead of ODEs and SDEs with flow-based models, we explicitly show that improvements of diffusion models such as a learned variance \\[3\\] can be combined with non-diffusion noise schedules like from rectified flows \\[2\\]. To the best of our knowledge, there have been no flow-based formulations incorporating a learned variance so far. To answer the question about the individual proofs:\n\n1. The first proof shows that the chosen mean of the Gaussian reverse distribution for the next (less noisy) state $x\\_{t\\_{i-1}}$ given the current noisy state $x\\_{t\\_i}$ and the clean data $x\\_0$ ensures that the marginal distribution only conditioned on the clean data $x\\_0$ is of the desired form with the noise schedule $a, b$ defining the interpolation weights of data and noise, respectively.  \n2. Since the standard deviation in the reverse distribution is a free variable, we follow DDIM by defining it as an interpolation between deterministic sampling (zero standard deviation) and stochastic sampling with a Markovian forward process, which makes DDIM equivalent to DDPM. Analogously, the second proof validates that the specific choice of the standard deviation in the reverse distribution results in a Markovian forward process, but for our formulation with more general noise schedules.\n\nAs you suggested, we will mark these statements as clear Propositions in the final version of the paper.\n\n- \\[1\\] Denoising Diffusion Implicit Models. ICLR 2021  \n- \\[2\\] Flow Matching for Generative Modeling. ICLR 2023  \n- \\[3\\] Improved Denoising Diffusion Probabilistic Models. ICLR 2021", "claims_and_evidence": "Yes.  Remaining questions were adequately highlighted following the conclusions.", "methods_and_evaluation_criteria": "Yes, albeit a bit primitive suitable for an early conceptual investigation as this one.", "theoretical_claims": "The derivations in Appendix A seemed correct, although it wasn't immediately clear which \"statement/result\" the \"Proof\" aims to establish.\n- Please clarify, marking a clear Proposition, Lemma, or Theorem.", "experimental_designs_or_analyses": "Yes, read through Section 4.  Some implementation details and ablations are deferred to the appendices, which is common.", "relation_to_broader_scientific_literature": "Reasoning over complex domains is certainly relevant in science, although this particular paper is a bit more primitive or too early to be of broader relevance.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "- The FFHQ backgrounds is mentioned as a confusing factor for the \"counting polygons\" experiment.\n- - It would have been nice to include an ablation with e.g. blank backgrounds to test this further\n- - This is potentially significant as this experiment doesn't strongly support the claimed benefits of sample orderings", "other_comments_or_suggestions": "N/A", "questions_for_authors": "No further questions at this time.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes Spatial Reasoning Models (SRMs), a framework leveraging denoising generative models to perform continuous spatial reasoning tasks. It demonstrates that model accuracy improves significantly (e.g., from <1% to >50% on Sudoku) when sequentialization and uncertainty-guided ordering are used. Reviewers appreciated the novelty and careful exploration of denoising strategies but raised concerns about the limited real-world applicability, clarity of theoretical contributions, and modest improvements on more realistic benchmarks. The authors addressed these concerns effectively in the rebuttal, adding new results, clarifying methodology, and reporting comparative results with LLMs/VLMs. AC feels the paper offers meaningful insights into reasoning with generative models, supported by strong empirical improvements. Therefore, an accept decision is recommended.
Decision: Accept (poster)