Source-linked AI summary
Attend, Infer, Repeat: Fast Scene Understanding with Generative Models
S. M. Ali Eslami, Nicolas Heess, Theophane Weber, Yuval Tassa, David Szepesvari, Koray Kavukcuoglu, Geoffrey E. Hinton
TL;DR
Structured image models need efficient inference that can reason about variable numbers of objects without relying on labels. AIR uses attentive, iterative amortized inference with generative models, and learns interpretable scene representations that count, locate, classify, and reconstruct objects across 2D and 3D settings. The authors report fast inference and improved usefulness and generalization, while noting unresolved scaling, likelihood, and model-capacity challenges.
Problem
Inference in structured visual models is challenging because latent dimensionality can vary, probabilistic renderers can be expensive or get stuck in local minima, and 3D representations may include discrete variables.
Method
AIR performs amortized variational inference with a recurrent network that attends to one object at a time and learns the required number of inference steps for structured 2D and 3D generative models.
Results
The models learn without supervision to count, locate, classify, and reconstruct scene elements, producing useful representations and fast inference across partially specified 2D models and fully specified 3D renderers.
Takeaways & Limitations
Model structure provides an inductive bias for interpretable representations, while attentive, iterative, variable-length inference supports fast reasoning about scenes with varying numbers of objects.
Takeaways & Limitations
AIR remains sensitive to gradient variance and likelihood choice, and the authors do not claim it is an ideal model for all images because richer natural variability remains difficult to capture.
Abstract
from arXiv · showhide
We present a framework for efficient inference in structured image models that explicitly reason about objects. We achieve this by performing probabilistic inference using a recurrent neural network that attends to scene elements and processes them one at a time. Crucially, the model itself learns to choose the appropriate number of inference steps. We use this scheme to learn to perform inference in partially specified 2D models (variable-sized variational auto-encoders) and fully specified 3D models (probabilistic renderers). We show that such models learn to identify multiple objects - counting, locating and classifying the elements of a scene - without any supervision, e.g., decomposing 3D images with various numbers of objects in a single forward pass of a neural network. We further show that the networks produce accurate inferences when compared to supervised counterparts, and that their structure leads to improved generalization.
1 Introduction
The paper combines structured generative models with learned amortized inference to interpret scenes as collections of objects without label supervision. Its iterative framework reasons about variable-sized scenes and yields interpretable representations that support generalization and downstream usefulness.
- Framework: The framework imposes representational structure through partly or fully specified generative models rather than supervised labels.The authors emphasize that the objective is useful scene representations, not primarily good image reconstructions.
- Results: Structured representations from the framework are reported to be more useful for downstream tasks than those from VAEs and DRAW.This comparison concerns the usefulness of learned representations rather than reconstruction quality alone.
- Unsupervised structure: AIR can discover variable numbers of scene entities, including their appearance and locations, through an Occam’s-razor effect in unsupervised learning.The representation combines distributed high-dimensional features with interpretable latent variables such as affine pose.
- Framework: AIR performs inference iteratively, attending to one object at a time while learning the appropriate number of inference steps for each image.The recurrent inference network supports variable-dimensional latent spaces and is trained end-to-end with the generative model.
- Applications: The paper applies AIR to multi-object generative models, including MNIST digits, overlapping sprites, Omniglot glyphs, and 3D scenes.The authors report unsupervised decomposition and improved generalization from model structure as an inductive bias.
2 Approach
The approach represents scenes with variable numbers of object-specific latent variables and approximates otherwise intractable posterior inference using an attentive recurrent network. Sequential conditioning, learned stopping, and specialized gradient estimators enable end-to-end optimization for continuous and discrete scene variables.
- Scene model: The generative model combines a prior over scene descriptions with a likelihood describing how those descriptions render into images.The prior captures assumptions about scenes, while the likelihood models image formation; both vary across applications.
- Scene model: Scene descriptions contain a variable-length sequence of object groups, with each group encoding attributes such as type, appearance, and pose.The number of objects is sampled from a prior, after which the object variables form the scene descriptor.
- Model instances: The framework supports 2D objects with learned continuous shape representations and three-dimensional poses, and 3D objects with categorical identities, positions, and rotations.Thus the latent-variable types differ between the 2D and 3D settings.
- Inference: AIR addresses trans-dimensional posterior inference by using a recurrent network that infers one object at a time while conditioning on the image and previously inferred objects.This sequential conditioning also helps avoid explaining the same object twice.
- Inference: A unary presence code lets the network represent the object count and terminate inference when it predicts that no more objects should be described.The network runs for at most N steps, with the presence variable acting as an interruption signal.
- Learning: AIR jointly optimizes the generative model and inference network using a variational lower bound on image marginal likelihood.Continuous variables use pathwise reparameterization, while discrete variables use likelihood-ratio gradients with variance-reducing neural baselines.
3 Models and Experiments
AIR performs learned, amortized inference over variable-sized object representations, applying attention recurrently to 2D and renderer-defined 3D scenes. Experiments show unsupervised object discovery, adaptive computation, generalization, useful representations, and multi-object 3D inference.
- 3.1 Multi-MNIST: AIR learns to detect and generate multiple MNIST digits from unlabeled images, with inference able to stop after the first absent entity to save computation.The model was trained from scratch on 60,000 images without supervision, and its count detections became accurate.
- Models: In 2D AIR, each entity is represented by presence, appearance or class, and affine position-and-scale variables, while spatial attention extracts image patches.The generative model combines transformed decoded digits to form the image, and inferred pose guides attention and decoding.
- 3.1.1 Strong Generalization: DAIR generalizes well to both extrapolation and interpolation of digit counts, whereas AIR trained on up to two digits is effectively unable to infer three-digit counts.The reported explanation attributes AIR’s extrapolation failure to its LSTM learning never to expect more than two digits.
- 3.1 Multi-MNIST: AIR representations support downstream digit summation and ascending-order detection, achieving high accuracy with only a fraction of labeled data and outperforming less interpretable alternatives.Figure 5 compares AIR with CAE and DRAW despite comparable reconstructions.
- 3.2 3D Scenes: With a specified 3D renderer, AIR infers object counts, identities, poses, and camera position from unlabeled images, including scenes with repeated objects.The 3D setting addresses expensive or unstable posterior inference and renderer inputs that may be discrete or lack gradients.
4 Related Work
Related work spans structured inverse-graphics models, deep networks for image prediction, and amortized methods that accelerate inference. AIR combines and extends these ideas with counting, iterative attention, and renderer-based inference.
- 4 Related Work: Inverse-graphics approaches provide structured generative models, while deep networks offer sophisticated image predictions but typically rely on large labeled datasets.Structured models can face difficult test-time inference, whereas neural predictions often provide less interpretable structure.
- 4 Related Work: PICTURE uses a deep network to distill slow MCMC results, and variational auto-encoders interleave learning of amortized inference and generative models.These methods connect data-driven prediction with structured generative modeling.
- 4 Related Work: AIR draws upon, extends, and links prior ideas related to counting, pondering, and gradient estimation through renderers.The paper identifies their combination as enabling the full capabilities of the proposed approach.
5 Discussion
The framework combines structured representations with attentive, iterative, variable-length amortized inference to support unsupervised scene understanding and fast inference. It yields interpretable representations, while learning remains sensitive to gradient variance, likelihood choice, and model expressiveness.
- Discussion: AIR performs inference iteratively with an attentive recurrent network that processes one object at a time and learns how many steps each image requires.The framework uses end-to-end amortized learning for variable-dimensional latent spaces.
- Discussion: The models learn to count, locate, classify, and reconstruct scene elements without supervision, in addition to reconstructing images.The stated aim is scene understanding rather than reconstruction alone.
- Discussion: Model structure provides an inductive bias for interpretable representations, and fast inference remains possible for sophisticated models or renderers.The authors connect these outcomes to appropriate structure and amortized inference.
- Discussion: AIR learning is most successful when gradient variance is low and the likelihood is well suited to the data.The authors identify scaling with object count and alternative likelihoods as future questions.
- Discussion: The unary prior encodes object count with a variable-length code whose maximum index follows the desired distribution p(n).The construction makes the sampled vector a correct unary code and defines a joint distribution consistent with p(n).
C Details of 2D Experiments
The 2D experiments use fixed implementation settings and a DAIR variant that updates a partial reconstruction to guide inference of each object.
- Implementation: Experiments used batch size 64, learning rates 10^-4 for inference and decoders and 10^-3 for baselines, 256-unit LSTMs, and 50-unit object appearances.Images were normalized to [0,1], with a Gaussian likelihood of fixed standard deviation 0.3 and a sparse-favoring geometric object-count prior.
- DAIR Network: DAIR uses a canvas-reconstruction recurrent architecture that forms a partial reconstruction at each time step before inferring the next latent variables.It is described as AIR with additional structure, with the recurrent component fixed to a canvas-reconstruction network.
- DAIR Network: The renderer’s sufficient statistic is iteratively updated, and DAIR feeds the error between the partial reconstruction and the image into a network predicting the next latent variables.The link function maps the sufficient statistic to the likelihood mean for Gaussian or Bernoulli renderers.
E Details of AIR vs. CNN vs. CAE vs. DRAW Experiments
The experiments compare AIR and DAIR with DRAW and describe the architectures used for image reconstruction and downstream prediction. DRAW performs better on simpler digit counts, whereas DAIR generalizes better to three digits.
- Results: DRAW achieves lower scores than AIR and DAIR on up to 2 digits but is outperformed by DAIR when generalizing to 3 digits.This comparison is reported in Table 1.
- AIR: AIR reconstructions on Omniglot use variable numbers of strokes for digits of varying complexity.The figure visualizes reconstructions at every AIR time step.
- Architectures: The convolutional autoencoder uses three convolutional encoding layers with 2×2 max pooling and three convolutional deconvolution layers with nearest-neighbor upsampling.The filters are 64 × (6 × 6).
- Downstream prediction: AIR, DRAW, and CAE embeddings feed a four-layer 512-unit network for 19-way sum or 2-way order prediction.The downstream network receives embeddings from each model.
- Evaluation: The comparison reports free energy on test sets containing 0, 1, or 2 digits and images with precisely 3 digits.All likelihood models use a Gaussian with fixed standard deviation 0.3.
G Omniglot Experiments
On Omniglot, AIR is evaluated under a generative-model mismatch because the data arise from pen strokes not directly represented in the model. It nevertheless discovers spatially coherent stroke-like elements using variable inference lengths.
- Dataset: Omniglot contains 1,623 handwritten characters from 50 alphabets, with each character drawn online by 20 people.The dataset was generated through human-produced drawing sequences.
- Quantitative results: Figure 11 reports the variational lower bound over training and sprite count accuracy.The figure separates training progress from count evaluation.
- Method: AIR allows up to 4 inference time steps per image and uses different step counts for characters of varying complexity.This behavior is shown in the reconstructions at every time step.
- Results: Despite not using stroke labels or stroke physics, AIR discovers spatially coherent elements resembling strokes.The experiment tests learning under a generative-model misspecification because pen strokes are not directly reflected in the model structure.
H Sprites Experiments
On a 50×50 sprite dataset with zero, one, or two non-overlapping sprites, AIR learns to identify sprite causes, counts, locations, and reconstructions without supervision. Its learned scanning policies are spatial and attend to each object once.
- The dataset contains red circles, green squares, and blue diamonds, with zero, one, or two additively composed sprites per image.
- After unsupervised training, AIR learns sprite causes, counts, locations, and convincing reconstructions.
- AIR correctly detects two overlapping sprites of the same type and color and reconstructs each sprite despite the other being present.
- Successful counting and reconstruction require a scanning policy that attends to every object only once.
- Across three differently seeded runs, AIR learns a unique scanning policy that appears spatial rather than identity- or size-based.
- Reconstructions are shown at different training points, while counting accuracy and scanning policies are visualized separately.
I Details of 3D Scene Experiments
The 3D experiments use MuJoCo to render scenes and finite-difference its graphics engine with respect to scene configuration. This provides a generic gradient-estimation procedure while exploiting MuJoCo’s speed and parameterization.
- The 3D experiments use MuJoCo’s physics simulator for rendering.
- MuJoCo’s graphics engine is differentiated by forward finite-differencing with a constant ϵ = 10−4 relative to scene configuration.
- The finite-difference procedure is generic to graphics engines; MuJoCo was chosen for speed and convenient scene parameterization.
- Despite coarse 8-bit OpenGL output, quantization errors appeared to average out reasonably well over pixels.
I.2 Scene generation
The 3D scene experiments generate randomized tabletop scenes with up to three objects, using MuJoCo geometry and optionally a variable camera. The inference model represents object presence, type, pose, and, for variable-camera scenes, camera angle.
- Scene generation: Single-object scenes use a MuJoCo table, three possible object geometries, a fixed camera, and three object degrees of freedom.
- Scene generation: Tabletop scenes use a table and replicated cup, pan, and plate meshes, allowing arbitrary scenes with up to N = 3 objects.
- Scene generation: Object presence, position, rotation angle, and type are randomly sampled for tabletop scenes; absent objects are moved outside the camera’s field of view.
- Scene generation: The variable-camera condition samples camera position within ± 40 degrees of center, adding one camera degree of freedom.
- Scene generation: Scenes are rendered at 128 × 128 pixels and down-sampled to 32 × 32 pixels for network input.
- Model: Each object’s latent variables encode visibility, a three-valued object type, and continuous position and rotation; variable-camera scenes add zcam ∈R.
- Model: The likelihood is evaluated by rendering a high-resolution image, blurring both rendered and input images, and computing a Gaussian likelihood.
- Model: AIR uses a recurrent network that processes one object per step, with sampling distributions for presence, type, and pose variables.
J Inference Speed
AIR inference is fast in both MNIST and 3D scenes, reaching around 59 frames per second for three-digit MNIST images. It is also faster than gradient-based optimization, which requires many steps per object.
- MNIST inference takes 5.6 milliseconds per object step and up to around 17 milliseconds per three-digit image.
- At around 59 frames per second, MNIST inference runs significantly faster than real-time.
- 3D inference takes around 2.3 milliseconds per step and 8 milliseconds per image on a CPU.
- Gradient-based optimization takes 5 milliseconds per gradient step per object and tens or hundreds of steps per image.
- Training converges in around 2 days for MNIST and around 3 days for the 3D scenes.