Source-linked AI summary
SPACE: Unsupervised Object-Oriented Scene Representation via Spatial Attention and Decomposition
Zhixuan Lin, Yi-Fu Wu, Skand Vishwanath Peri, Weihao Sun, Gautam Singh, Fei Deng, Jindong Jiang, Sungjin Ahn
TL;DR
Unsupervised object-oriented scene representation remains challenging for scenes with many occluded objects, partial observability, and complex backgrounds, while prior spatial-attention and scene-mixture approaches each have limitations. SPACE combines both through a unified probabilistic model with parallel foreground processing. Experiments on Atari and 3D-Rooms show factorized foreground objects, complex background decomposition, and improved scalability without significant performance degradation.
Problem
Unsupervised structured representation of scenes with many objects, occlusion, partial observability, and complex backgrounds remains an unsolved challenge, with prior approaches limited in combining object structure and scalability.
Method
SPACE is a unified probabilistic latent-variable model that uses parallel spatial attention for foreground objects and component mixtures for complex background areas.
Results
SPACE provides factorized foreground object representations, decomposes complex background segments, and achieves significant speed-ups with scalability to larger object counts without performance degradation.
Takeaways & Limitations
SPACE combines object-wise disentanglement, flexible background decomposition, and scalable processing in a single scene representation model.
Takeaways & Limitations
SPACE’s rendering has worst-case complexity O(hw × HW), and the authors identify sequential background processing and natural-image modeling as future directions.
Abstract
from arXiv · showhide
The ability to decompose complex multi-object scenes into meaningful abstractions like objects is fundamental to achieve higher-level cognition. Previous approaches for unsupervised object-oriented scene representation learning are either based on spatial-attention or scene-mixture approaches and limited in scalability which is a main obstacle towards modeling real-world scenes. In this paper, we propose a generative latent variable model, called SPACE, that provides a unified probabilistic modeling framework that combines the best of spatial-attention and scene-mixture approaches. SPACE can explicitly provide factorized object representations for foreground objects while also decomposing background segments of complex morphology. Previous models are good at either of these, but not both. SPACE also resolves the scalability problems of previous methods by incorporating parallel spatial-attention and thus is applicable to scenes with a large number of objects without performance degradations. We show through experiments on Atari and 3D-Rooms that SPACE achieves the above properties consistently in comparison to SPAIR, IODINE, and GENESIS. Results of our experiments can be found on our project website: https://sites.google.com/view/space-project-page
1 INTRODUCTION
Unsupervised scene representation must decompose many-object scenes into meaningful entities despite occlusion, partial observability, and complex backgrounds. SPACE unifies spatial attention and scene mixtures to combine object-level structure, flexible background decomposition, and improved scalability.
- Motivation: Structured scene representations decompose visual scenes into abstract entities such as objects and spaces, supporting downstream reasoning and learning.The paper identifies potential benefits including improved sample efficiency, visual variable binding, reasoning, and causal inference.
- Prior approaches: Scene-mixture models provide flexible segmentation maps for objects and background segments with complex morphology.Their components correspond to full-scale images rather than necessarily local areas.
- Prior approaches: Scene-mixture inference is inherently sequential because components refer to one another, limiting scalability as object counts increase.The scaling limitation arises from processing scene components sequentially to obtain a complete scene.
- Prior approaches: Spatial-attention models explicitly represent object geometry such as position and scale, but rectangular attention struggles with morphologically complex objects and background segments.These geometric features are grounded in physical semantics and may support interpretability, geometric reasoning, sample efficiency, and transferability.
- Proposed approach: SPACE processes bounding-box-compatible foreground objects with parallel spatial attention and decomposes the remaining complex area using component mixtures.This design provides disentangled foreground object representations with position and scale while retaining flexible background decomposition and addressing scalability.
- Contributions: The paper introduces a unified probabilistic latent-variable model, a spatially parallel multi-object module, and extensive comparisons of prior methods’ capabilities and limitations.The comparisons evaluate the proposed framework’s ability to mitigate scalability problems alongside its representational properties.
2 THE PROPOSED MODEL: SPACE
SPACE is a unified probabilistic generative model that decomposes scenes into parallel-attended foreground objects and mixture-modeled background components. Its mean-field cell inference supports parallel processing, while a boundary loss helps prevent large objects from being split across boxes.
- Generative process: SPACE combines spatial attention for foreground objects with scene-mixture modeling for morphologically complex background segments.The foreground and background distributions are combined with a pixel-wise mixture model to reconstruct the complete image.
- Background model: The background is decomposed into K segments whose latents model component mixing probabilities and RGB distributions.The components are combined through a mixture model and reconstructed with a VAE.
- Foreground model: The foreground latent structure assigns each H × W cell at most one nearby object with presence, location, depth, and appearance variables.zpres indicates object presence; zwhere encodes size and relative location; zdepth resolves occlusions; and zwhat models appearance and the mask.
- Inference and training: SPACE uses a variational approximation, an ELBO objective, the reparameterization trick, and stochastic gradient descent for training.The variational approximation is needed because the generative-model integrals over continuous foreground and background latents cannot be evaluated analytically.
- Parallel inference: Mean-field inference makes each cell’s foreground latents independent of other cells, allowing local object detection and spatial attention in parallel.The paper contrasts this with SPAIR’s autoregressive row-major inference, whose cost becomes prohibitive as object counts increase.
- Preventing box-splitting: The boundary loss penalizes object masks overlapping a border around each glimpse, discouraging large objects from being split across multiple bounding boxes.The method addresses the trade-off between bounding-box-size priors that may split large objects or miss small ones.
3 RELATED WORKS
Prior unsupervised object-oriented scene models use recurrent attention, convolutional spatial attention, or scene-mixture decomposition. These approaches differ in scalability, geometric object representations, and their ability to model complex background morphology.
- Spatial-attention models: AIR recurrently attends to objects one at a time using what, where, and presence latents.Its number of recurrent steps scales with the number of attended objects, limiting scalability to images with many objects.
- Spatial-attention models: SPAIR replaces AIR’s recurrent network with a convolutional network and specifies object locations relative to local grid cells.This supports spatially invariant computations, but the related passage does not establish that it fully resolves scaling limitations.
- Scene-mixture models: Scene-mixture models explain a visual scene as a mixture of a finite number of component images, supporting flexible segmentation maps.MONet uses recurrent attention to produce pixel-wise component masks, followed by a VAE for each component; IODINE uses iterative refinement of latent object representations.
4 EVALUATION
SPACE is evaluated on Atari and 3D-Room scenes against scene-mixture and spatial-attention baselines, including settings with many objects. It combines object detection, complex background segmentation, and parallel inference while retaining comparable bounding-box quality.
- Experimental setup: SPACE is evaluated on Atari and generated 3D-Room datasets, including small scenes with 4–8 objects and large scenes with 18–24 objects.The 3D-Room images vary camera angle and the colors of objects, walls, floor, and sky.
- Experimental setup: The evaluation compares SPACE with IODINE, GENESIS, and SPAIR, while also testing SPAIR-P under patch-based training.SPAIR receives an additional VAE for background processing because it lacks an explicit background component.
- Qualitative comparison: IODINE and GENESIS segment some 3D-Room components but struggle with object decomposition and Atari foregrounds, especially as object counts increase.In Atari, both may fail to capture foreground properly or encode all objects in one component.
- Qualitative comparison: SPAIR-P often misses proper object boxes, duplicates detections, and labels background as foreground because each patch sees only a 32 × 32 glimpse.This limitation is illustrated by detecting small aliens while treating long background ground as foreground in Space Invaders.
- Qualitative comparison: SPACE accurately detects most 3D-Room objects and cleanly segments walls, ground, and sky despite variation in object positions, colors, and shapes.In River Raid, it segments the changing blue river while detecting foreground objects; SPAIR often cannot separate foreground and background there.
- Qualitative comparison: Across 10 jointly trained Atari games, SPACE correctly detects foreground objects and cleanly segments complex backgrounds.The joint-training demonstration shows six games with complex backgrounds.
- Quantitative comparison: SPACE and SPAIR converge faster than IODINE and GENESIS, with SPACE becoming faster than SPAIR at the 16×16 setting.At 4×4 and 8×8, SPACE and SPAIR converge equally fast.
5 CONCLUSION
SPACE combines object-oriented spatial attention with component-mixture scene decomposition, while improving scalability to scenes with many objects. The paper demonstrates these properties on Atari and 3D-Rooms.
- SPACE provides factorized representations for foreground objects and decomposes complex background segments.
- SPACE achieves significant speed-up without performance degradation as the number of objects increases.
- SPACE produces more intuitive detected objects than other methods.
- Experiments on Atari and 3D-Rooms demonstrate SPACE’s stated properties.
- Future directions include parallelizing background processing, improving performance on natural images, and applying SPACE to object-oriented model-based reinforcement learning.
A ADDITIONAL RESULTS OF SPACE
Additional results illustrate SPACE’s object detection and background segmentation on Atari and 3D-Room datasets, including scenes with small and large numbers of objects.
- Figure 6 compares object-detection behavior between SPACE and SPAIR on Montezuma’s Revenge.
- Figure 7 demonstrates SPACE trained jointly on a selection of 10 Atari games.
- Figure 8 shows SPACE detecting objects and segmenting backgrounds in 3D-Room scenes with few objects.
- Figure 9 shows SPACE detecting objects and segmenting backgrounds in 3D-Room scenes with many objects.
B ELBO DERIVATIONS
The ELBO derivations detail KL-divergence terms for SPACE’s foreground and background latents and the relaxed treatment of foreground-presence variables.
- The section derives the evidence lower bound for the log-likelihood log p(x).
- The foreground latent KL-divergence term is evaluated under SPACE’s approximate inference.
- The background latent KL term is evaluated under GENESIS-like inference modeling.
- SPACE models the Bernoulli foreground-presence variable zpres_i with the Gumbel-Softmax distribution during implementation.
- The boundary-loss kernel uses a 32 × 32 glimpse, boundary gap b = 6, negative uniform boundary weights, and zero weights between boundary and glimpse.
D IMPLEMENTATION DETAILS
SPACE’s implementation uses parallel foreground inference over grid cells and sequentially organized background inference with component processing and rescaling operations.
- Rescaling: The rescale_i function converts a shift latent using the cell index and grid dimensions into a rescaled shift latent.
- Foreground inference: Foreground variables are processed in parallel across all HW grid cells.
- Foreground inference: Foreground inference takes an image as input and outputs a foreground mask, appearance, grid height, and grid width.
- Foreground inference: The foreground pipeline rescales local shifts, extracts glimpses with a Spatial Transformer, and computes glimpse masks and appearances.
- Background inference: Background inference encodes the image, iterates over K components, and uses a stick-breaking process for component probabilities.
D.2 TRAINING REGIME AND HYPERPARAMETERS
SPACE experiments use fixed image dimensions, optimizer settings, gradient clipping, and dataset-specific hyperparameters. Separate configurations are provided for large 3D-Room experiments and joint training on 10 Atari games.
- All experiments use 128 × 128 images, batch sizes of 12 to 16, and gradient clipping with maximum norm 1.0.
- The foreground module uses RMSProp with learning rate 1 × 10−5, except Figure 5, which uses 1×10−4 as SPAIR.
- The background module uses Adam with learning rate 1 × 10−3, while other experiment hyperparameters are individually finetuned by dataset.
- The 3D-Room Large and joint-Atari configurations specify different priors, latent-component counts, temperatures, and boundary-loss settings.The 3D-Room configuration uses K=5 and boundary loss, whereas joint Atari uses K=3 and no boundary loss.
D.3 MODEL ARCHITECTURE
The 16 × 16 SPACE architecture combines foreground and background modules with convolutional encoders, glimpse processing, latent-component decoding, and sequential background dependencies. The 8 × 8 variant changes the final image-encoder stride.
- The 8 × 8 grid-cell model matches the 16 × 16 architecture except for a stride-2 convolution in the image encoder’s final layer.
- Image encoders are fully convolutional and output H × W feature maps, while the glimpse encoder ends with a linear Gaussian-parameter layer.
- The background component decoder decodes zc_k into background components, complementing the foreground module’s glimpse-based processing.
- Softplus parameterizes Gaussian standard deviations, sigmoid produces reconstructions and masks, and modules use either GN/CELU or BN/ELU.
E DATASET DETAILS
The evaluation uses Atari datasets sampled from a pretrained agent and procedurally generated 3D-Room scenes. Atari experiments include both individual games and joint training across ten games, while 3D-Room varies object count.
- Individual Atari games use 60,000 sampled images split into 50,000 training, 5,000 validation, and 5,000 testing examples.Images are resized to 128 × 128 pixels with BGR color channels.
- Joint Atari training uses 8,000 training, 1,000 validation, and 1,000 testing images for each of 10 games.
- MuJoCo generates 3D-Room scenes containing randomly sized spheres, cubes, and cylinders in walled enclosures with randomized object and background colors.The small dataset contains 4-8 objects, while the large dataset contains 18-24 objects.