Source-linked AI summary
Multi-Object Representation Learning with Iterative Variational Inference
Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Chris Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, Alexander Lerchner
TL;DR
Existing representation-learning methods often overlook scenes containing multiple objects or treat segmentation separately. IODINE jointly segments and represents objects through iterative variational inference, achieving unsupervised disentangled representations, occlusion inpainting, systematic generalization, and multiple decompositions for ambiguous inputs. Its scope is constrained by failures on some complex real-world datasets and by an independence assumption that limits modeling object dynamics and interactions.
Problem
Representation learning has focused mainly on single-object features, leaving the joint unsupervised segmentation and representation of multi-object scenes as an open problem.
Method
IODINE uses shared, permutation-invariant object slots in a spatial mixture generative model and refines their variational posterior iteratively with gradient-informed inputs.
Results
IODINE unsupervisedly segments scenes, learns disentangled object features, inpaints occluded regions, generalizes to more objects and unseen feature combinations, and represents multiple valid decompositions for ambiguous inputs.
Takeaways & Limitations
Joint multi-object representation learning can produce interpretable object decompositions while supporting combinatorial generalization and ambiguity handling within the demonstrated settings.
Takeaways & Limitations
On ImageNet IODINE segmented mostly by color rather than objects, while on Textured MNIST it mostly failed to capture foreground digits; its independent-object assumption also limits modeling interactions.
Abstract
from arXiv · showhide
Human perception is structured around objects which form the basis for our higher-level cognition and impressive systematic generalization abilities. Yet most work on representation learning focuses on feature learning without even considering multiple objects, or treats segmentation as an (often supervised) preprocessing step. Instead, we argue for the importance of learning to segment and represent objects jointly. We demonstrate that, starting from the simple assumption that a scene is composed of multiple entities, it is possible to learn to segment images into interpretable objects with disentangled representations. Our method learns -- without supervision -- to inpaint occluded parts, and extrapolates to scenes with more objects and to unseen objects with novel feature combinations. We also show that, due to the use of iterative variational inference, our system is able to learn multi-modal posteriors for ambiguous inputs and extends naturally to sequences.
1. Introduction
IODINE treats object discovery as part of unsupervised representation learning for scenes containing multiple, variably arranged and partially occluded objects. It combines multi-object latent structure with iterative variational inference to learn interpretable, disentangled object representations.
- Most unsupervised representation-learning work emphasizes single objects, whereas real scenes contain variable numbers of objects, spatial configurations, and partial occlusions.
- IODINE jointly discovers objects and learns their representations in a fully data-driven, unsupervised variational framework.The model uses a spatial mixture perspective and amortized iterative refinement of latent object representations.
- IODINE segments complex scenes and learns disentangled object features without supervision on datasets including CLEVR, Objects Room, and Tetris.
- The model generalizes to more objects than seen during training and to objects formed from unseen feature combinations.
- IODINE’s iterative refinement supports multiple valid decomposition solutions for ambiguous inputs and extends naturally to sequences.
- Figure 2 depicts a decoder that refines K object latents over N iterations and decodes them into pixel-wise means and masks.
2. Method
IODINE represents each scene with shared, permutation-invariant object slots and infers them through iterative refinement in a spatial mixture generative model. Gradient-informed updates progressively improve posterior estimates while training is performed through the unrolled iterations.
- 2.1. Multi-Object Representations: Multi-slot representations encode interchangeable objects using a shared format, supporting comparison, reuse, recombination, and combinatorial generalization.
- 2.1. Multi-Object Representations: Each scene is generated from K independent object latents whose shared mechanism makes the representation permutation-invariant.
- 2.1. Multi-Object Representations: Each slot decodes into a pixel-wise appearance mean and assignment mask within a spatial Gaussian mixture model.The masks assign pixels to objects, while the means describe object appearance.
- 2.2. Inference: Iterative inference refines posterior parameters from the input and samples using a trained refinement network rather than a single feed-forward encoder.
- 2.2. Inference: Slots are updated independently and in parallel, interacting only through their inputs; the refinement network combines convolutional processing with an LSTM.
- 2.2. Inference: Gradient information about the ELBO helps each slot determine what remains unexplained by the other slots.
- 2.3. Training: The decoder, refinement network, and initial posterior are trained by gradient descent through the unrolled inference iterations.
- 2.3. Training: Stopping gradients through gradient-based refinement inputs mitigates numerical instability caused by double derivatives.
3. Related Work
Prior work separately advances disentangled feature representations, image segmentation, and learned scene generation, but rarely addresses joint multi-object representation. IODINE is positioned as an unsupervised approach that combines segmentation, representation learning, inpainting, and ambiguity handling.
- Representation-learning advances largely target feature-level structure, including disentanglement and improved generalization, rather than multi-object scene structure.
- Image-segmentation methods often rely on supervised signals or hand-engineered features instead of jointly learning segmentation and representation.
- Full scene-representation approaches have produced appealing results but may require carefully engineered generative models that are not fully learned from data.
- Neural Expectation Maximization amortizes expectation maximization for spatial mixture models, but its variants fail to cope with colored scenes in the authors’ comparison.
- IODINE’s figures illustrate separate masked reconstructions and segmentation masks across CLEVR6, Multi-dSprites, and Tetris.
- Factor-regression evaluation on CLEVR6 compares prediction accuracy or R2 across position, color, shape, and size, including a deconvolutional ablation.
4. Results
IODINE jointly learns unsupervised object segmentation and representations, achieving strong segmentation, interpretable object features, disentanglement, and systematic generalization. Its iterative variational inference also supports multi-modal interpretations, while performance depends on sufficient slots and decoder design.
- Representation Quality: IODINE disentangles individual factors of single objects, whereas a standard VAE trained on multi-object scenes entangles factors across objects.The comparison is illustrated through single-latent traversals, including factors of an orange cylinder.
- Segmentation: IODINE achieves almost perfect foreground ARI of around 0.99 on CLEVR6 and Tetris, and 0.77 on Multi-dSprites.ARI measures similarity between predicted and ground-truth instance segmentations.
- Representation Quality: A linear mapping extracts object attributes such as color, position, shape, and size to high accuracy from IODINE’s latent representations.The predictor is trained after IODINE, without supervised fine-tuning of the representation itself.
- Generalization: A model trained with K = 7 slots on up to 6 objects generalizes to K = 11 slots and scenes containing 9 objects, including unseen green spheres.Segmentation performance decreases little when generalizing to more objects, while unseen feature combinations remain representable.
- Slot Capacity: Training with too few slots forces multiple objects into each slot, producing inefficient, entangled representations that increasing slots at test time cannot repair.When trained with sufficiently many slots, excess slots are typically left empty and test-time generalization is strong.
- Multi-Modality and Multi-Stability: Iterative refinement can yield different segmentation modes for ambiguous inputs because sampling across iterations implicitly supports multi-modal posteriors.Different evaluations can produce slot permutations or qualitatively different scene interpretations.
5. Discussion and Future Work
IODINE decomposes multi-object scenes into individual representations, handling occlusions, novel object counts and feature combinations, and ambiguous decompositions. The discussion identifies sequential data, object relations, and task-relevant selection as directions for extending the approach to richer scenes.
- IODINE decomposes partially occluded scenes into objects represented by individual properties such as color, size, and material.It can inpaint covered sections while performing unsupervised scene representation learning.
- IODINE generalizes beyond training distributions in both object numerosity and object-property combinations.The paper reports extrapolation to scenes with more objects and to objects formed from unseen feature combinations.
- Ambiguous object decompositions yield multiple valid solutions for the same input image.This supports representing uncertainty over alternative scene interpretations.
- Sequential inputs are a promising extension because movement and smooth factor changes provide rich statistics about objectness.Feeding a new frame at every iteration can also maintain object-to-slot associations.
- Current IODINE has limited abilities for modeling dynamics and assumes objects are placed independently of one another.The authors propose graph networks that support relations while preserving slot symmetry.
- Task-related signals may be needed to scale representation learning to scenes containing extremely many potential objects.Such signals could bias which objects are selected and how scenes are decomposed.
A.1. Sequences
IODINE extends its iterative inference procedure to sequences by feeding a new frame at each iteration. In the demonstrated setup, one iteration corresponds to one timestep in a 16-timestep sequence.
- 16 timesteps are shown with one inference iteration per timestep.Figure 12 sets the number of refinement iterations N equal to the number of timesteps in the data.
- IODINE can process sequences by feeding a new frame at every iteration instead of repeating one image.The sequence setup uses next-step prediction rather than reconstruction as part of the training objective.
A.2. Memory Limitations
IODINE’s memory use grows with both the number of slots and refinement steps, constraining training scalability, especially for long sequences. At inference, prior-step activations need not be retained, removing the dependence on refinement steps.
- Memory scaling: Memory consumption scales linearly with both the number of slots K and refinement steps T because training stores decoder and refinement activations.Each slot and refinement step requires activations for an entire decoder and refinement network.
- Empirical limit: 16 timesteps with batch-size 4 was the upper training limit on V100 GPUs with 16GB of RAM in the reported experiments.The limit also depends on input and network size.
- Sequential data: The memory constraint is particularly restrictive for sequential data, where the number of refinement steps can become very large.
- Inference: At inference, previous-timestep activations do not need to be stored, so dependence on T can be eliminated.
B.1. CLEVR
The CLEVR evaluation uses regenerated scenes with ground-truth segmentation masks, training on scenes containing three to six objects and evaluating on the full three-to-ten-object distribution. The dataset samples pair each scene with a background mask and object masks.
- Dataset construction: 70 000 CLEVR images were regenerated with ground-truth segmentation masks, center-cropped from 240 × 320 to 128 × 128 pixels.
- Scene composition: Scenes contain between three and ten objects characterized by shape, size, material, and color.
- Training and evaluation: Training uses CLEVR6, the subset containing three to six objects, while evaluation uses the full CLEVR distribution containing three to ten objects.
- Visual examples: Figure 13 presents each CLEVR6 sample as a scene, background mask, and subsequent ground-truth object masks.
B.2. Multi-dSprites
Multi-dSprites contains simple multi-object sprite scenes with varying shape, color, scale, position, rotation, and background brightness, alongside a binarized variant. The supplied passages also describe related generated shape and Tetromino datasets.
- Dataset: The Multi-dSprites dataset contains 60 000 images at 64×64 resolution.
- Dataset: Each Multi-dSprites image contains two to five random sprites varying in shape, color, scale, position, and rotation, with grayscale background brightness variation.
- Dataset variants: The binarized Multi-dSprites variant uses white sprites, black backgrounds, and two to three sprites per image.
- Related datasets: The generated Tetromino dataset contains 60 000 non-overlapping 35×35 images, each with three randomly colored Tetrominoes.Each Tetromino comprises four 5 × 5 blocks, with 17 rotation-counted shapes and six possible colors.
- Related datasets: The shapes dataset contains 60 000 binary 28×28 images with three random shapes drawn from triangle, inverted triangle, and square.
B.5. Objects Room
The sequential Objects Room experiments use RGB views of a cubic room containing randomly positioned objects, with object motion governed by fixed sequence velocities. The supplied implementation passages specify preprocessing, optimization, posterior initialization, refinement inputs, and model settings.
- Dataset: Objects Room provides 64x64 RGB views of a cubic room with randomly colored walls and floors and three randomly shaped objects, of which 1-3 are visible per frame.
- Dataset: Object position and angle change according to a fixed velocity throughout each sequence, with a random velocity sampled per sequence.
- Training: Training uses Adam with learning rate 0.0003, gradient clipping above norm 5.0, and output-distribution scale σ = 0.1.
- Posterior initialization: IODINE initializes the posterior independently of input data and makes the initial distribution parameters trainable rather than matching the prior.Matching the prior initially was found to create competing requirements for the latent space.
- Refinement inputs: The refinement network receives image-sized likelihood, mask, gradient, and coordinate inputs, then combines convolutional outputs with flat posterior parameters before an LSTM.
- Model settings: For CLEVR experiments, models used K = 7 slots and T = 5 iterations during training, increasing to K = 11 slots for full-CLEVR evaluation.
D. Additional Plots
Additional plots visualize decompositions, latent-space organization, latent traversals, and input-ablation effects across the evaluated datasets. They also expose reconstruction, segmentation, and occlusion-handling behavior in more detail.
- Decompositions: Additional samples show segmentation and object-reconstruction decompositions for Multi-dSprites, Tetris, and CLEVR6.The figures include masked reconstruction slots, object masks, ground-truth masks, and model-produced mask logits.
- Projections of Object Latents: Object latents are visualized with principal-component and t-SNE projections for CLEVR6, Multi-dSprites, and Tetris.Each plotted point represents an object latent and is colored by its corresponding ground-truth factor.
- Traversals: Additional randomly chosen CLEVR latent traversals illustrate changes to individual objects while revealing flaws in occlusion handling.The captions note that adjusting an unshown depth-ordering latent could address some occlusion errors, and one traversal exposes a bright background spot behind a moved object.
- Input Ablations: Input ablations evaluate how removing refinement-network inputs affects total loss, reconstruction error, KL divergence, and ARI segmentation performance on CLEVR6 and Tetris.Each curve corresponds to training without one particular input; ARI excludes background pixels, and reconstruction error uses mean squared error as a proxy for the full reconstruction loss.