Spatial Reasoning with Denoising Models

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.

Missing Proposition, Lemma, or Theorem in Appendix A

Thank 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:

1. 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.
2. 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.

As you suggested, we will mark these statements as clear Propositions in the final version of the paper.

$1$ Denoising Diffusion Implicit Models. ICLR 2021

$2$ Flow Matching for Generative Modeling. ICLR 2023

$3$ Improved Denoising Diffusion Probabilistic Models. ICLR 2021

Thank you very much for your review. We address your concerns as follows:

Justification for uniform distribution of mean noise level during training

1. 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”.
2. 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.

Role of Appendix A

Thank you for pointing out this lack of clarity.

1. 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.

$1$ Denoising Diffusion Implicit Models. ICLR 2021

$2$ Flow Matching for Generative Modeling. ICLR 2023

$3$ Stochastic Interpolants: A Unifying Framework for Flows and Diffusions. arxiv 2023

$4$ Improved Denoising Diffusion Probabilistic Models. ICLR 2021

2. 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.

Confusion w.r.t. spatially autoregressive approaches

You 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.
We will make sure to clarify this in the final version.

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.

W1. Real-world applicability

We 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.

W2 + Q3. Validation of noise schedules

We 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
All 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.

W3. Architecture limitations

We 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
All 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.

W4 + Q1. ResNet Accuracy

The 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.

W5 + Q2. Inconsistent improvements

1. 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.
   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.

2. 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). The motivation is that stars are composed of higher frequencies, moving the “point of commitment” to numbers and stars closer together in time.

For 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.

W6 + Q4. Computational cost

In 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.

W7. Data Statistics

Thank 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.

We sincerely thank you for your valuable feedback. To address your concerns, we present a point-to-point response in the following:

Why not use the latest VLMs [...] to address this symbolic visual reasoning task?

Thank you for this important question. Our response is threefold:

1. 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.

2. TLDR: We tested LLMs and VLMs. They do not naively solve Sudoku.
   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 . 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:

Please 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

3. 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.

Latest study (Sudoku-Bench) on visual reasoning with VLMs.

Thank you for the pointer to this interesting benchmark. We will add a discussion to the final version. Please note the following:

1. 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.
2. 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).
3. 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.

“Diffusion Forcing scheme”

Our 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.