[Paper Review] Multi-Object Representation Learning with Iterative Variational Inference

Klaus Greff, Raphaël Lopez Kaufman+7 more|arXiv (Cornell University)|Mar 1, 2019

Domain Adaptation and Few-Shot Learning33 references176 citations

TL;DR

IODINE learns unsupervised multi-object scene decomposition by iterative variational refinement, enabling object-friendly representations, inpainting, and generalization to more objects and unseen feature combinations.

ABSTRACT

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.

Motivation & Objective

• Motivate learning representations that segment and represent multiple objects jointly rather than as a preprocessing step.
• Propose a multi-slot, spatial mixture model that encodes scenes as independent object slots with shared decoding.
• Develop an iterative amortized inference procedure to infer object-level latents and segmentation without supervision.
• Demonstrate unsupervised segmentation, inpainting of occluded parts, and generalization to unseen object counts and feature combinations.

Proposed method

• Represent each scene with K latent object vectors z_k that collectively generate the image via a spatial Gaussian mixture model with per-slot masks m_k.
• Decode each z_k into pixel-wise appearance μ_k and mask logits, normalize masks across slots with a softmax, forming the overall likelihood p(x|z).
• Use an amortized iterative refinement network f_φ to update posterior parameters λ_k through T iterations, with additive updates and auxiliary inputs including gradients of the ELBO.
• Train end-to-end by unrolling T refinement steps and minimizing a weighted sum of ELBO terms over iterations to stabilize learning.
• Incorporate a broadcast decoder to encourage disentanglement of position from other features and enforce slot symmetry for interchangeable objects.

Experimental results

Research questions

• RQ1 Can a multi-slot, unsupervised model segment a scene into interpretable objects while learning their disentangled features?
• RQ2 Does iterative variational inference enable handling of occlusions, multi-modality, and permutation-invariant object representations in a fully data-driven way?
• RQ3 How well does the model generalize to scenes with more objects or unseen object-feature combinations beyond the training distribution?
• RQ4 What is the impact of architectural choices (e.g., spatial broadcast decoder) and input signals on segmentation quality and disentanglement?
• RQ5 Can the learned object representations support downstream tasks like predicting object attributes via linear mappings?

Key findings

• IODINE achieves near-perfect instance segmentation on CLEVR6 (ARI ~0.99) and Tetris, and strong segmentation on Multi-dSprites (ARI ~0.77).
• MONet achieves ARI ~0.96 on CLEVR6 and ~0.90 on Multi-dSprites under certain settings, providing a competitive baseline.
• The object latent representations enable linear predictability of ground-truth factors (color, position, shape, size) from per-object latents.
• Disentanglement is improved when objects are represented via separate slots; standard VAEs entangle factors across multiple objects.
• Model generalizes to more slots than seen in training and to scenes with more objects, maintaining segmentation quality in many cases.
• Iterative refinement yields multi-modal posteriors emergently due to sampling and slot interactions, allowing multi-stable segmentations under ambiguity.