Source-linked AI summary
Neural Expectation Maximization
Klaus Greff, Sjoerd van Steenkiste, Jürgen Schmidhuber
TL;DR
The paper asks how to identify multiple conceptual entities and learn separate distributed representations for them without supervision. It models this problem with neuralized Expectation Maximization, deriving a differentiable clustering procedure that learns grouping and entity representations. The method accurately recovers constituent objects and produces representations useful for next-step prediction, while unsupervised segmentation may not always match the intended outcome.
Problem
Real-world tasks require identifying and manipulating multiple entities and relationships, but combining distributed representations can create ambiguities known as the Binding Problem.
Method
N-EM parametrizes a spatial mixture model with a neural network and unrolls generalized Expectation Maximization into an end-to-end differentiable clustering procedure.
Results
The method recovers constituent objects and learns localized, separate representations that are useful for next-step prediction with multiple objects.
Takeaways & Limitations
Representing each entity separately in a shared space avoids the Binding Problem and makes the learned representations usable as symbols for arbitrary entities in the dataset.
Takeaways & Limitations
Without supervision, the resulting segmentation may not always match the intended outcome because the notion of an object is ill-defined and task dependent.
Abstract
from arXiv · showhide
Many real world tasks such as reasoning and physical interaction require identification and manipulation of conceptual entities. A first step towards solving these tasks is the automated discovery of distributed symbol-like representations. In this paper, we explicitly formalize this problem as inference in a spatial mixture model where each component is parametrized by a neural network. Based on the Expectation Maximization framework we then derive a differentiable clustering method that simultaneously learns how to group and represent individual entities. We evaluate our method on the (sequential) perceptual grouping task and find that it is able to accurately recover the constituent objects. We demonstrate that the learned representations are useful for next-step prediction.
1 Introduction
The paper addresses the challenge of learning separate, symbol-like representations for multiple entities, whose relationships matter in reasoning and physical interaction. It formalizes unsupervised grouping and representation learning as a neural extension of Expectation Maximization.
- Real-world reasoning and physical interaction require identifying and manipulating multiple entities and their relationships.
- Superimposing distributed and disentangled representations can create ambiguities known as the Binding Problem.
- Perceptual grouping dynamically splits each input into its constituent conceptual entities.
- N-EM formalizes grouping as inference in a spatial mixture model whose components are parametrized by neural networks.
- The derived differentiable clustering method simultaneously learns to group and efficiently represent individual entities in an unsupervised manner.
2 Neural Expectation Maximization
Neural Expectation Maximization (N-EM) models images as spatial mixtures whose neural-network components represent individual objects, then makes EM differentiable so grouping and representations can be learned jointly. The method alternates soft pixel assignments with representation updates and trains the underlying statistical model by backpropagating through unrolled inference.
- 2.1 Parametrized Spatial Mixture Model: Each image is modeled as a mixture of K objects, with each pixel assigned to exactly one component and each object represented by θ_k.A neural network maps θ_k to parameters of separate pixel-wise distributions, typically Bernoulli or Gaussian.
- 2 Neural Expectation Maximization: N-EM learns the object-level distribution P(x|θ_k) unsupervised by parameterizing it with a differentiable neural network f_φ.Differentiability permits an outer loss to backpropagate into the network weights φ.
- 2.2 Expectation Maximization: EM alternates an E-step that computes soft pixel-to-component assignments with an M-step that improves representations by maximizing expected log likelihood.Because f_φ is nonlinear, the M-step uses gradient ascent rather than an analytical maximization.
- 2.2 Expectation Maximization: For sufficiently small η, the generalized EM procedure is guaranteed to converge to a local optimum of the data log likelihood.The guarantee applies to the generalized EM updates before introducing the more flexible recurrent variant.
- 2.3 Unrolling: Unrolling generalized EM yields an end-to-end differentiable clustering procedure whose statistical model can be trained with stochastic gradient descent and backpropagation through time.The model’s weights encode the statistical regularities used for clustering pixels into objects.
- 2.3 Unrolling: RNN-EM replaces constrained EM-like updates with recurrent weights, extending the procedure to sequential data while sacrificing the data-log-likelihood convergence guarantee.Its increased flexibility can produce more powerful clustering, but the method no longer has the guarantee described for generalized EM.
- 2.4 Training Objective: A two-term training loss combines an intra-cluster objective with an inter-cluster KL penalty so each network models one object while suppressing out-of-cluster predictions.Pixelwise weighting enables specialization, while the inter-cluster term prevents unconstrained predictions on pixels assigned elsewhere.
- 2.4 Training Objective: Stopping gradients through γ addresses conflicting incentives in which a component can reduce loss either by predicting a pixel better or by taking less responsibility for it.The conflict arises from the interaction between soft assignments γ and predictions ψ.
3 Related work
N-EM connects unsupervised perceptual grouping to prior work on neural reconstruction, binding mechanisms, segmentation, multi-causal modeling, and differentiable inference. RNN-EM is presented as a recurrent, more flexible extension, related to both Tagger and trainable unrolled optimization methods.
- Neural grouping: Tagger also learns unsupervised perceptual grouping with K cooperating neural networks, but additionally predicts the grouping and supports complex texture segmentation.N-EM instead maintains a closer connection to EM and relies on the model’s statistical structure.
- Binding problem: Binding-problem research includes synchronization-based oscillatory networks, complex-valued activations, and RNN mechanisms, although suitable mechanisms can be difficult to learn in practice.The passage identifies competition as one possible added mechanism.
- Neural grouping: RTagger, a recurrent extension of Tagger supporting sequential data, was developed concurrently with this work.The related-work passage identifies it as concurrent rather than as an earlier method.
- Unsupervised segmentation: Related unsupervised segmentation work spans random vectors, texture segmentation, images, video segmentation with generalized EM, optical-flow supervision, and multi-causal image models.These approaches address segmentation across several data and modeling settings.
- Unrolled inference: RNN-EM parallels work that unrolls iterative inference procedures such as ISTA and trains the resulting computation graph to produce sparse codes.Both approaches turn inference updates into trainable procedures, though RNN-EM derives its recurrence from N-EM.
4 Experiments
The experiments test RNN-EM on static and sequential perceptual grouping, measuring object recovery, prediction, robustness to component-count choices, and generalization. RNN-EM accurately separates objects and learns representations useful for next-step prediction across increasingly varied settings.
- Experimental setup: The evaluation uses generated static images and video with ground-truth clustering, measuring grouping quality using Adjusted Mutual Information (AMI).AMI ranges from 0 for random clustering to 1 for a perfect match; experiments also assess next-step prediction and the hyperparameter K.
- Static shapes: 0.826 ± 0.005 AMI for RNN-EM versus 0.475 ± 0.043 for N-EM on static shapes.RNN-EM is more stable under occlusion and its tunable degrees of freedom speed optimization; Tagger reports 0.79 ± 0.034 AMI, or 0.97 ± 0.009 with layer normalization, using about twenty times more parameters.
- Flying shapes: 0.970 ± 0.005 AMI for three flying shapes and 0.878 ± 0.003 for five cluttered shapes.The model separates individual objects in almost all five-shape cases, including sequences with substantial clutter and overlap.
- Flying shapes: RNN-EM produces significantly lower binomial cross-entropy errors than a recurrent autoencoder, especially as the number of objects increases.The comparison covers denoising and next-step prediction across different object counts; the reported RNN-EM BCE is an upper bound to the true BCE under the evaluation procedure.
- Flying shapes: Training with too many components is not detrimental because the network leaves excess groups empty, while too few components causes only a slight score drop when K is corrected at test time.The results indicate robustness to different choices of the unknown component-count hyperparameter K.
5 Discussion
Neural Expectation Maximization learns separate, localized representations for constituent entities, helping avoid the binding problem and supporting multi-object prediction. Its unsupervised segmentation remains limited because object identity can be ill-defined and task dependent.
- Neural Expectation Maximization groups pixels according to constituent objects and learns localized representations containing information relevant to each entity.
- Separate representations place entities in a common space, avoiding the binding problem and making them usable as symbols for arbitrary dataset entities.
- Learned representations benefit next-step prediction with multiple objects, including cases where overlapping objects challenge standard approaches.
- Unsupervised segmentation may not match the intended outcome because the notion of an object is ill-defined and task dependent in real-world images.
6 Conclusion
The paper presents Neural Expectation Maximization as an unsupervised, differentiable clustering framework for separating and representing entities. Experiments on synthetic images and video show that it recovers underlying objects and learns useful representations.
- The framework combines neural networks with generalized EM to form a trainable clustering algorithm.
- The method can be trained fully unsupervised to segment inputs into entities and represent them individually.
- Experiments on synthetic images and video verify recovery of underlying objects and learning of useful representations.
- The work aims to provide a theoretical foundation for unsupervised symbol-like representation learning and a first step toward practical solutions.
A Experiment Details
The experiments use fixed training settings and evaluate grouping quality against ground truth with Adjusted Mutual Information, using validation-based early stopping.
- Networks are trained with ADAM using default parameters, batch size 64, and 50 000 training, 10 000 validation, and 10 000 test inputs.
- Grouping quality is measured with Adjusted Mutual Information against ground truth while ignoring background and overlap regions.
- Training uses early stopping after validation loss fails to improve for 10 epochs.
A.1 Experiments on Static Shapes
The static-shapes experiment uses binary images containing three randomly positioned shapes and implements N-EM and RNN-EM with sigmoid-based neural networks trained through multiple EM steps with noise.
- Each input is a 28 × 28 binary image containing three regular shapes at random positions.
- N-EM uses a single-layer fully connected sigmoid network that maps a 250-dimensional vector to per-pixel Bernoulli parameters.
- RNN-EM uses a recurrent network with 250 sigmoid hidden units and a sigmoid output layer producing per-pixel Bernoulli parameters.
- Both networks use K = 4 for 15 EM steps, inject bitflip noise with probability 0.1, and apply the outer loss only at the final EM step.
A.2 Experiments on Flying Shapes
Flying Shapes uses sequential binary image inputs and a convolutional encoder-decoder with a recurrent bottleneck to model moving shapes. Training adds bitflip noise and uses next-step prediction with recurrent EM updates.
- Inputs are sequences of binary 28 × 28 images containing a fixed number of shapes moving along random trajectories for 20 steps.
- The model uses a convolutional encoder-decoder architecture with a recurrent neural network as its bottleneck.
- The surrounding fully connected layers use 512 units with ELU before the recurrent layer and ReLU after it, with layer normalization.
- The recurrent bottleneck has 100 Sigmoid units with layer normalization on its output.
- Nearest-neighbour upsampling followed by convolution replaces transposed convolutions to avoid frequency artifacts, without layer normalization on the recurrent connection.
- RNN-EM receives bitflip-noised inputs with p = 0.2, performs one step per timestep, and is trained using a next-step prediction objective.
A.3 Experiments on Flying MNIST
Flying MNIST uses sequential grayscale images of moving, down-sampled MNIST digits and a deeper encoder-decoder architecture. Its training procedure replaces bitflip noise with masked uniform noise and uses a learning rate of 0.0005.
- Inputs are grayscale 24 × 24 image sequences containing down-sampled MNIST digits moving from random positions along random trajectories for T timesteps.
- The Flying MNIST model uses a slightly deeper version of the architecture used for flying shapes.
- Its architecture stacks three convolutional layers, a 250-unit recurrent bottleneck, fully connected layers, and convolutional nearest-neighbour upsampling.
- Training replaces bitflip noise with masked uniform noise by interpolating masked pixels between the original image and Uniform samples over the data range.
- The learning rate is 0.0005 from the second stage onward when stage-wise training is used.