[Paper Review] Neural Expectation Maximization

Klaus Greff, Sjoerd van Steenkiste+1 more|arXiv (Cornell University)|Aug 11, 2017

Gaussian Processes and Bayesian Inference31 references50 citations

TL;DR

Neural Expectation Maximization (N-EM) learns to cluster and represent multiple objects in images in an unsupervised, differentiable EM framework, extended to sequential data via RNN-EM.

ABSTRACT

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.

Motivation & Objective

• Motivate learning separate, distributed representations for multiple objects to address the binding problem.
• Formalize object representations as neural-network-parametrized components in a spatial mixture model.
• Derive a differentiable EM procedure to cluster pixels into objects and learn object-specific representations.
• Extend the framework to sequential data for next-step prediction and improved grouping.
• Provide unsupervised training without labeled segmentations and evaluate on synthetic datasets.

Proposed method

• Model images as a K-component spatial mixture with per-component parameters theta_k mapped to pixel likelihoods psi_i,k via a differentiable f_phi.
• Compute E-step to obtain soft pixel assignments gamma_i,k based on current psi and x.
• Perform M-step by gradient ascent on Q using the differentiable f_phi to update theta_k (Eq. 4).
• Unroll EM iterations to create an end-to-end differentiable clustering procedure (N-EM) trained by backpropagation through time.
• Introduce RNN-EM by replacing the M-step with a learned recurrent network to handle sequential data and improve grouping.
• Train with a two-term loss: intra-cluster reconstruction weighted by gamma and inter-cluster KL penalty to constrain non-assigned pixels.

Experimental results

Research questions

• RQ1 Can unsupervised neural networks discover and represent multiple objects as separate, disentangled components within a spatial mixture model?
• RQ2 Does differentiable EM allow end-to-end training that yields useful object-centric representations for next-step prediction?
• RQ3 How well does the approach extend to sequential data and provide robust perceptual grouping over time?

Key findings

• N-EM and RNN-EM can recover individual shapes in static grouping tasks when objects are separated.
• RNN-EM generally yields stronger grouping performance (AMI) than N-EM, especially under occlusion.
• On flying shapes, AMI scores remain high even with more objects and show good generalization to unseen sequence lengths and numbers of objects.
• Next-step prediction benefits from multi-object representations, with RNN-EM achieving lower prediction error than a single-component recurrent autoencoder when objects are present.
• On Flying MNIST, RNN-EM achieves AMI up to 0.917±0.005 on test with two digits and generalizes to three digits without retraining.