Source-linked AI summary
Relational Abstractions for Spatial Reasoning with Diffusion Models
Ana Ezquerro, Ozan Özdenizci
TL;DR
Diffusion models can synthesize images but remain limited in satisfying structured spatial reasoning constraints, and pixel-level modeling does not capture the logical structures required. RDM combines unsupervised object discovery and relational abstractions to guide diffusion generation; across benchmark tasks it improves reasoning-driven image generation, including by up to 3× on challenging multi-context problems.
Problem
Diffusion models remain limited in reliably satisfying spatial reasoning constraints, while pixel-level reconstruction alone does not capture the logical structures required for reasoning.
Source: §1 · pp. 1–2Method
Results
Takeaways & Limitations
Object-centric relational abstractions provide a task-agnostic approach to guiding generative spatial reasoning, without positioning one model variant as the preferred recipe.
Source: §7 · p. 9The approach relies on object discovery with finite scalability and may not scale to unconstrained in-the-wild scenes containing thousands of entities.
Source: §7 · p. 10
Abstract
from arXiv · showhide
Diffusion models excel at image synthesis, but they remain limited in their ability to reliably satisfy structured spatial reasoning constraints. In conditional data distribution modeling tasks with implicit logical structure, such as puzzles defined by visible clues paired with consistent solutions, state-of-the-art generative models tend to approximate pixel-space distributions without learning the underlying logical rules required for inference. To address this limitation, we present a novel framework for spatial reasoning with diffusion models that leverages unsupervised object discovery and abstractions of object relations. We show that the relational knowledge derived from object-centric representations enriches diffusion models with structural primitives, allowing them to effectively guide the generative representation space during both training and inference, and enabling conditional image generation that satisfies reasoning constraints. Additionally, we introduce a large-scale generative spatial reasoning benchmark with four datasets inspired by human-solvable puzzles. Our results show that relational abstractions significantly improve reasoning capabilities of diffusion models on a variety of complex reasoning tasks, while enabling robust generalization in out-of-distribution settings.
1 INTRODUCTION
The introduction asks whether diffusion models can perform spatial reasoning and argues that doing so requires more than pixel-level reconstruction. It presents RDM and a four-dataset benchmark, reporting improved reasoning performance and out-of-distribution generalization.
1 INTRODUCTION
- Diffusion models capture structural and semantic regularities, but whether they can perform spatial reasoning remains an open question.Source: §1 · p. 1
- The introduction argues that complex visual reasoning requires higher-level context abstractions beyond refined pixel-level reconstruction.Source: §1 · p. 2
- RDM combines unsupervised object discovery with abstractions of object relations to guide diffusion models in generating images that satisfy reasoning constraints.The paper also introduces four puzzle-inspired datasets, each with up to one million paired examples.Source: §1 · p. 2
- RDM improves performance on complex reasoning tasks by up to 3× on challenging multi-context problems and demonstrates strong out-of-distribution generalization.Source: §1 · p. 2
2 RELATED WORK
Related work spans diffusion formulations, generative approaches to visual reasoning, and object-centric representations. The cited methods include grid-tailored diffusion sampling strategies and relational modules built on object-discovery representations.
2 RELATED WORK
- Diffusion models learn data distributions through denoising, score-based, or flow-based formulations, with continuous-time ODE formulations leading to continuous normalizing flows.Source: §2 · p. 2
- Generative visual reasoning includes RPM methods that synthesize missing panels and diffusion-based sampling strategies, including SRM’s patch-level uncertainty approach tailored to structured grids.The cited diffusion approaches include Wewer et al. (2025) and Lee et al. (2025); SRM is described as exploiting Sudoku-style grid structure.Source: §2 · p. 2
- Slot Attention represents visual components with feature and positional embeddings called slots, while relational modules use slots to model object interactions.Object-centric diffusion models extend these designs to generative settings, including downstream vision tasks and compositional image synthesis.Source: §2 · pp. 2–3
3 RELATIONAL ABSTRACTIONS FOR DIFFUSION MODELS
RDM models conditional spatial reasoning by pairing an unsupervised object-discovery model with a diffusion reasoner guided by relational abstractions. Its discoverer forms slot-based representations, while the reasoner uses clue and, during training, solution representations to guide generation.
3 RELATIONAL ABSTRACTIONS FOR DIFFUSION MODELS
- RDM models the conditional distribution of valid solutions given problem clues, using a discoverer to extract object-centric latents and a diffusion reasoner conditioned on those latents.Source: §3 · p. 3
3.1 THE Discoverer: UNSUPERVISED SLOT-BASED OBJECT DISCOVERY
- Slot Attention iteratively binds object-centric information to slots, producing feature and positional embeddings that represent the discovered objects.A slot-based autoencoder learns these representations from visual data through reconstruction.Source: §3.1 · pp. 3–4
- The original slot decoder’s latent bottleneck incurs computational overhead as the number of slots increases; RDM’s scalable decoder uses a positional grid to reduce backward-pass memory demand.Source: §3.1 · p. 4
3.2 THE Reasoner: RELATIONAL OBJECT ABSTRACTIONS AS DIFFUSION GUIDANCE
- The relational abstraction module processes object features and positions into embeddings that encode relations, guiding generation with contextual interactions as well as individual object features.It uses interleaved attention and feed-forward blocks to refine positional signals with feature projections.Source: §3.2 · p. 4
- The reasoner trains one parameterization on solution-guided and unguided distributions, dropping solution representations during training to support clue-only inference.The reported experiments use flow matching, though the framework also allows noise prediction or score matching.Source: §3.2 · p. 4
4 SPATIAL REASONING BENCHMARK
The benchmark comprises four puzzle-inspired datasets, each requiring models to infer implicit rules from conditional data. It covers grid-based completion under spatial constraints and multi-context generation from clues.
4 SPATIAL REASONING BENCHMARK
- The benchmark introduces four datasets, with up to one million images each, spanning grid-based completion and multi-context generation tasks.Their logical rules are implicit in the training data and must be inferred.Source: §4 · pp. 4–5
4.1 GRID-BASED IMAGE COMPLETION TASKS
- Akari and Coldoku are grid-based image-completion tasks, with reasoning framed as inpainting under global and local spatial constraints; Coldoku is a Sudoku puzzle played on a 9 × 9 grid.In Akari, bulbs illuminate every cell, no two share a line of sight, and numbered walls specify exact counts of adjacent bulbs. Coldoku's clue is a partially completed grid.Source: §4 · p. 5, §4.1 · p. 5
4.2 MULTI-CONTEXT IMAGE GENERATION TASKS
- Tangram and LogicFace require generating a solution from context images by inferring relations that connect the clues.Tangram uses colored pieces and a silhouette; LogicFace uses two face images and hidden logical operators over facial features.Source: §4 · p. 5, §4.2 · p. 6
5 EXPERIMENTS
Experiments use task-specific object-discovery configurations and relational conditioning, then evaluate generated solutions against puzzle criteria. The setup compares RDM with capacity-matched diffusion baselines and varies conditioning, model size, and sampling.
Model configurations
- The discoverer uses task-specific slot counts, while the reasoner injects relational embeddings from six Abstractor layers into a U-Net bottleneck via cross-attention.Tangram and LogicFace use 7 and 6 slots; Akari and Coldoku are tested with cell-count slots or 16 slots.Source: §5 · p. 6
Conditioning on relational abstractions
- Three conditioning settings use clue relations, solution relations, or both during training; at inference, clue conditioning is retained when used.Solution-derived conditioning is dropped during training because the solution is unavailable at inference.Source: §5 · p. 6
Training details
- The reasoner uses flow matching and deterministic reverse-time ODE sampling, while the capacity-matched Base omits relational abstractions.A diffusion-forcing variant is also compared with SRM using parallel and sequential sampling strategies.Source: §5 · pp. 6–7
- The chosen pdrop = 0.5 gives the best balance between learning solution-guided and unguided distributions, while smaller values degrade unconditional-distribution estimation at inference.The authors hypothesize that this setting allows both distributions to be learned evenly.Source: §5 · p. 6
Model evaluations
- Evaluation uses 5,000 held-out samples per task and overall accuracy, with additional task-specific metrics and difficulty or wall-density regimes.Table 1 reports Akari and Coldoku results; Coldoku includes digit and color accuracy across three difficulty levels.Source: §5 · p. 7, Table 1 · p. 7
6 EXPERIMENTAL RESULTS
RDM outperforms baselines on Akari, Tangram, and LogicFace, while SRM leads on the hardest Coldoku level. Results also show that performance depends on puzzle structure and configuration, and that relational abstractions contribute beyond slot-based guidance alone.
Overall results
- RDM outperforms baselines on Akari, Tangram, and LogicFace, while only SRM performs better on Coldoku’s hardest level.Source: §6 · p. 7
Akari
- With 16 slots, RDM clue accuracy remains 72.48 in Akari’s low regime, versus Base’s decline from 80.04 in-distribution to 39.48.The authors report that 16-slot disentanglement aids out-of-distribution generalization, while 121 slots slightly overfit; SRM also reaches 18.80 accuracy on Akari.Source: §6 · pp. 7–8
Coldoku
- RDM generally leads Coldoku’s easy and medium levels, but SRM performs better on hard; RDM’s best results use 16 slots.The 16-slot configuration reaches ACC 97.48 easy, 71.64 medium, and 12.58 hard; the 81-slot configuration reaches 90.36 on easy.Source: §6 · p. 8
Tildoku
- Tildoku random tilts disrupt grid-cell alignment, and SRM’s low-tilt advantage over baseline collapses when digits no longer align with fixed patches.Source: §6 · pp. 8–9
Tangram
- On Tangram, RDM (full) reaches ACC 46.76 versus baseline’s 15.46, with solution-view information improving shape and overall accuracy.The solution view supplies an explicit silhouette partition; shape accuracy reaches 65.12.Source: Table 2 · p. 8, §6 · p. 9
LogicFace
- On LogicFace, clue-conditioned RDM reaches accuracy 74.74, while standard diffusion reaches ACC 11.24.The task combines real-image modeling, multiple input attributes, and inference of logical operators.Source: Table 2 · p. 8, §6 · p. 9
Relational abstraction ablation
- Ablations indicate that explicit relational abstractions are critical for complex reasoning beyond slot-based diffusion guidance alone.Source: §6 · p. 9
Evidence of rule learning
- Relational-abstraction embeddings show evidence of near-optimal solution-conditioned inference, clue-free reasoning gains, and linear separation of valid from rule-violating puzzles.Source: §6 · p. 9
7 DISCUSSION
RDM uses object-centric relational abstractions to address diffusion models’ difficulty capturing logical structure, and the paper reports improved reasoning-driven generation across a diverse benchmark. The discussion also identifies finite object-discovery scalability and evaluator estimation errors as limitations.
7 DISCUSSION
- RDM uses unsupervised object discovery and relational abstractions to guide diffusion models toward logical structure without explicit symbolic computation.The authors describe the learned abstractions as moving diffusion models toward more symbolic forms of reasoning.Source: §7 · p. 9
- RDM significantly improves reasoning-driven image generation on a large-scale benchmark of diverse challenging tasks.The benchmark spans grid-based completion and multi-context generation tasks.Source: §7 · p. 9, A DETAILS ON THE SPATIAL REASONING… · p. 14
- Object-discovery architectures have finite scalability, and the authors caution that they may not scale to unconstrained scenes with thousands of entities.Neural benchmark evaluators may also make minor estimation errors, although consistent evaluation is used across models.Source: §7 · p. 10
REPRODUCIBILITY STATEMENT
- The reproducibility statement says that model and training details are provided and that the code and datasets are available online.The passage supplies the project repository URL.Source: REPRODUCIBILITY STATEMENT · p. 10
A DETAILS ON THE SPATIAL REASONING BENCHMARK DATASETS
- The benchmark is designed to test discrete geometric and logical constraints rather than simple pattern matching.Its stated goal is to assess generative models on challenging visual tasks.Source: A DETAILS ON THE SPATIAL REASONING… · p. 14
A.1 AKARI DATASET
- Akari puzzles place bulbs to illuminate every empty cell without allowing bulbs to share an unobstructed horizontal or vertical line of sight.The dataset contains one million 11 × 11 puzzles; 50% of walls are clued with adjacent bulb counts.Source: §A.1 · p. 14
- Akari is NP-complete with numbered-wall constraints, and the evaluation checks puzzle constraints using accuracy, bulb accuracy, unlit count, distance, and overlap.Without clued cells, the puzzle reduces to polynomial time.Source: §A.1 · pp. 14–15
A.2 COLDOKU: THE COLORED SUDOKU DATASET
Coldoku combines digit and color Sudoku constraints, while Tildoku adds bounded geometric distortion. The dataset descriptions also specify a two-stage evaluator for Coldoku samples.
A.2 COLDOKU: THE COLORED SUDOKU DATASET
- Coldoku pairs a Sudoku grid with an independent Sudoku that determines the color assigned to each digit, creating joint digit-and-color constraints.The dataset cross-combines the digit and color arrangements into a large configuration space.Source: §A.2 · p. 15
- Coldoku requires simultaneous resolution of two orthogonal Sudoku instances, preserving the uniqueness and NP-completeness properties of both layers.The model must maintain two independent sets of global constraints in a single inference pass.Source: §A.2 · p. 15
- Tildoku preserves Coldoku’s rules but randomly tilts samples within a bounded range, requiring detection of varying puzzle geometry.Tilting is limited to 15% corner displacements, and evaluation reverses the transformation before checking predictions.Source: §A.2 · p. 15
- Coldoku evaluation classifies each generated cell into a digit-color pair, then verifies Sudoku constraints for both arrangements.It reports digit accuracy, color accuracy, and joint accuracy.Source: §A.2 · pp. 15–16
- The Coldoku sample visualization places solutions above clues.The accompanying Tildoku figure shows solution samples with different tilting ratios.Source: Figure 8 · p. 16, Figure 9 · p. 16
A.3 TANGRAM DATASET
The Tangram dataset tests whether models can infer how seven discrete pieces map to an ambiguous silhouette under piece-use, overlap, and proportion constraints. LogicFace applies logical operators to facial attributes, and the architecture appendix outlines the discoverer and reasoner modules.
A.3 TANGRAM DATASET
- Tangram requires arranging all seven tans into a silhouette without overlap while preserving their area proportions, often despite ambiguous internal boundaries.The same layout may be structurally ambiguous, and multiple valid arrangements may exist.Source: §A.3 · p. 16, Figure 10 · p. 17
- Color-coded Tangram pieces make the task jointly spatial and logical by requiring each piece’s geometry to match its color and position in the solution.The shuffled pieces provide an initialization resembling the described human solving process; perception is robust to shuffles and silhouette rotations.Source: §A.3 · p. 17
- The Tangram evaluator checks silhouette mismatch before assessing piece shape and proportions, then whether each piece’s color matches its identity.A mismatch above 5% of the target layout yields zero accuracy.Source: §A.3 · pp. 17–18
A.4 LOGICFACE: LOGIC OPERATORS ON FACIAL FEATURES
- LogicFace presents two facial examples and requires a generated face’s binary attributes to satisfy fixed logical operators.The four attributes are smiling, male, glasses, and young; their operators are OR, AND, XOR, and IMPLIES, respectively.Source: §A.4 · p. 18, Figure 12 · p. 18
- LogicFace’s glasses and young attributes are imbalanced, and the authors acknowledge possible training-set replicas, especially for these attributes.Glasses appear in 6.5% of CelebA-HQ images, while young appears in 77.36%.Source: §A.4 · pp. 18–19
B DETAILS ON THE MODEL ARCHITECTURES
- The framework separates an independently trained object discoverer from a diffusion reasoner conditioned through learned relational embeddings.The reasoner models a conditional solution distribution from the clue and relational guidance.Source: B DETAILS ON THE MODEL ARCHITECTURES · p. 19
B.1 ARCHITECTURE OF THE Discoverer
The Discoverer uses slot attention to extract object-specific features and positions, which a relational abstraction module contextualizes for the diffusion reasoner. The appendix also reports that RDM works across model sizes, while larger models do not necessarily improve reasoning.
B.1 ARCHITECTURE OF THE Discoverer
- Slot attention iteratively refines slots and separates their information into object-specific features and positions.The reported configuration uses three refinement iterations.Source: §B.1 · p. 19
- The scalable decoder estimates attention masks before broadcasting and aggregates slots into a weighted feature map, decoupling reconstruction from slot count.This design avoids generating full candidate reconstructions for every slot.Source: §B.1 · p. 20
B.2 ARCHITECTURE OF THE Relational Abstraction MODULE
- The relational abstraction module uses interleaved attention and MLP layers to pool feature relations and produce contextualized, position-aware representations.It enriches slot representations with relational and positional information.Source: §B.2 · p. 20
B.3 CONFIGURATION OF THE Reasoner
- The reasoner models clue-conditioned solution generation, using relational embeddings from the discoverer to guide its diffusion model.Training includes solution-guided and unguided distributions, while inference samples the unguided distribution because the solution is unavailable.Source: §B.3 · p. 20
- The reasoner’s continuous normalizing flow starts from data and converges to a Gaussian prior through Gaussian probability paths.Sampling evolves Gaussian noise backward in time by solving an ODE.Source: §B.3 · pp. 20–21
- The authors found DiT difficult to optimize even under baseline configurations for these spatial tasks.They leave hybrid architectures combining convolutional inductive biases and Transformer scalability for future work.Source: §B.3 · p. 21
B.4 MODEL COMPLEXITY AND PARAMETERS
- RDM scales across 30M–120M-parameter backbones, but larger architectures do not necessarily benefit reasoning and smaller ones converge faster.The authors frame compatibility with multiple sizes as a reproducibility and feasibility benefit.Source: §B.4 · p. 21
- The small and medium configurations retain a Discoverer with fewer than 8M parameters, while their reasoners use 25M and 40M parameters, respectively.The large configuration expands the Discoverer dimension and matches a larger diffusion backbone.Source: §B.4 · p. 22
C.1 RESULTS OF THE Discoverer
The discoverer reconstructs puzzle inputs with varying accuracy, while its slots do not always align one-to-one with objects or attributes. Discrete features are captured more consistently than diffuse visual structure.
C.1 RESULTS OF THE Discoverer
- Akari reconstruction reached 97% and 98% with 16 and 121 slots; Coldoku digit accuracy was 82.64% and 99.79% with 16 and 81 slots, while color accuracy was 40.41% and 72.05%.The authors link lower Coldoku color accuracy to blurred reconstructions and errors when colors fall between palette values.Source: §C.1 · p. 22
- Tangram reconstruction reached 57% accuracy, with blurred borders contributing to errors when the evaluator estimates polygons.Source: §C.1 · p. 22
- LogicFace reconstruction matched 95% of selected features: smiling, male, glasses, and young.Source: §C.1 · p. 23
- On Tangram, some slots remained inactive while others grouped co-occurring tans, rather than representing each piece separately.The authors hypothesize that color and slight scale variation may lead slots to fit chromatic patterns instead of individual geometric entities.Source: §C.1 · p. 23
- On LogicFace, slots consistently isolated salient features such as glasses and smiling, but struggled with diffuse attributes and high-variance backgrounds.The authors nevertheless report that key features in slots remained useful for reasoning, with other attributes associable to them.Source: §C.1 · p. 23
C.2 RESULTS OF THE Reasoner
These results examine reasoner performance on Tangram and report detailed results across several puzzle datasets and sampling settings. Tangram experiments isolate geometric partitioning and color matching before considering the full task’s additional demands.
C.2 RESULTS OF THE Reasoner
- Detailed results cover Akari, Tildoku, Tangram, and LogicFace at sampling steps τ = 10 and 100, with τ = 100 omitted for Tildoku and some SRM results.The authors also report Akari-specific unlit-cell, bulb-distance, and overlap measures; SRM sequential results at τ = 100 are omitted in low and high regimes due to low ID performance and inference overhead.Source: §C.2 · pp. 23–24
- The full Tangram task combines geometric induction, color consistency across dynamic pieces, and latent spatial transformations to align tans with a target configuration.The authors contrast this with prior Sudoku-style diffusion tasks, which use localized relational rules in a fixed coordinate system.Source: §C.2 · pp. 24–25
- On silhouette-only geometric partitioning, the diffusion model reached 82% accuracy, compared with 1.93% for the convolutional autoencoder.The setup used a unique color for each tan; the authors report this as evidence that diffusion models learn complex partitions in this intermediate task.Source: §C.2 · p. 25
- The intermediate color-matching task fixed the piece view and varied colors, while the full Tangram task varies the tan view across samples.Source: §C.2 · p. 25
- With a fixed tan view and varying colors, Tangram partitioning accuracy was 54% and total accuracy was 52%.The authors infer that arbitrary color permutations tied to the same silhouette partition challenge correspondence between the conditional input and generated arrangement.Source: §C.2 · p. 25
C.2.1 ADDITIONAL EXPERIMENTS
Additional experiments compare relational abstractions with slot-only and non-object-centric guidance, test whether relational embeddings encode puzzle rules, and evaluate prior benchmarks. The reported results favor relational representations in these comparisons.
Impact of relational abstractions
- On Tangram, RDM outperformed slot-conditioned diffusion on clue accuracy, 34.37 versus 27.09; slot conditioning alone nevertheless improved on the standard baseline.The authors interpret the remaining gap as evidence that explicit relational abstraction between objects and spatial configurations matters for complex reasoning.Source: §C.2.1 · p. 26, Table 11 · p. 26
Empirical evidence of rule learning
- With relational solution embeddings available at inference, accuracy reached 93% on Tangram and 95% on hard Coldoku.This experiment provides evidence that RDM learns abstractions capable of solving both puzzles under that setting.Source: §C.2.1 · p. 26
- Without raw visual clues during Tangram training, relational abstractions achieved ACC 37.55, above the Base model’s 15.46.The authors report this as evidence that relational abstractions alone can replace the clue and improve reasoning.Source: §C.2.1 · p. 26
- Coldoku linear probes distinguished valid from invalid puzzles with 89.98% accuracy for repeated-column violations and 82.59% for shuffled invalid grids.The authors take this as evidence that learned relational representations encode information about logical rules.Source: §C.2.1 · p. 27
Evaluations on previous benchmarks
- On Counting Polygons, RDM scored 30.7 versus 10.1 for the comparison cited, and it outperformed SRM in MNIST-Sudoku medium and hard regimes with fewer than half its parameters.For MNIST-Sudoku medium, RDM slightly surpassed Base, 55.80 versus 54.18; the accuracy gap disappeared when sampling used 100 steps.Source: §C.2.1 · p. 27, Table 12 · p. 27
- DINOv2 and ViT-MAE non-object-centric features did not substantially improve MNIST-Sudoku medium or hard results, which remained below Base and RDM.Source: §C.2.1 · p. 27
- The original Counting Polygons evaluator biased predictions toward positive counts and exactly two digits, so the authors also evaluated with an unbiased evaluator.The revised evaluator handles cases with no digit or shape in the image.Source: §C.2.1 · p. 27