Neural Expectation Maximization

In plain words

Before a computer can reason about objects, it has to decide which pixels belong together. This paper learns that grouping without labels. The model assumes the picture is made of a few parts, each generated by its own small neural network, and it keeps adjusting which pixels each part explains.

The grouping step follows Expectation Maximization, a classic method that alternates between guessing which part explains each pixel and improving each part's description. The authors write those steps as a loop of ordinary differentiable operations, so the whole loop can be trained by gradient descent. A recurrent version, RNN-EM, replaces the hand-derived update with a learned recurrent network.

The method is tested on synthetic pictures of shapes and digits, including short videos where the shapes move. The learned groups match the true objects well, and the representations help predict the next video frame.

TL;DR

• Grouping pixels into objects is cast as inference in a spatial mixture model whose K components are neural networks §2.1 eq. 1.
• Expectation Maximization is unrolled into a differentiable clustering, trained through time with a loss that combines a responsibility-weighted reconstruction term and a KL term §2.3 §2.4 eq. 5.
• N-EM uses the gradient step derived from EM; RNN-EM replaces it with an unrestricted RNN, which gives up the convergence guarantee and allows next-step prediction on sequences §2.3.
• Author-reported: static shapes AMI 0.826 ± 0.005 (RNN-EM) vs 0.475 ± 0.043 (N-EM); flying shapes with 3 objects AMI 0.970 ± 0.005; flying MNIST with 2 digits AMI 0.819 ± 0.022 direct and 0.917 ± 0.005 stage-wise §4.1 §4.2 §4.3 tab. 1.
• The paper and its appendix disagree on K for static shapes, and the code differs from the paper on several loss details. Read Paper vs code and Pitfalls before reproducing numbers.

Problem

Learning useful representations of objects needs a way to split each input into its entities, a step called perceptual grouping. Superimposed distributed representations of several objects are ambiguous, which is the binding problem §1. The authors want one representation per entity, in a shared format, with the grouping learned without labels from the statistics of the data §1 §2.

The method assumes that pixels of one object depend on each other, while pixels of different objects are roughly independent given their representations §2. It needs the number of components K as a hyperparameter. The authors report that using too large a K is not harmful, because unused components stay empty §4.2.

Method

Model. Each image x is a spatial mixture of K components with representations θ_1 … θ_K. A shared neural network f_φ maps each θ_k to per-pixel parameters ψ_k. Each pixel has a latent assignment z_{i,k} that says which component generated it §2.1. Pixels are independent given the representation, but not identically distributed §2.1.

EM and its gradient step. Maximising the marginal likelihood directly is hard, so EM optimises the expected complete-data log-likelihood Q instead §2.2 eq. 2. The E-step computes responsibilities γ from the current predictions §2.2 eq. 3. The M-step cannot be solved in closed form because f_φ is non-linear, so the paper takes one gradient step on Q with step size η §2.2 eq. 4. The paper writes the gradient for each component as a responsibility-weighted residual times the Jacobian of ψ_k with respect to θ_k §2.2 eq. 4.

Unrolling. Repeating the E- and M-steps a fixed number of times gives a differentiable clustering. Its loss is back-propagated through time into φ §2.3. This trainable procedure is N-EM §2.3.

RNN-EM. The M-step's residual input and its Jacobian-based encoder are replaced by an unrestricted recurrent network with hidden states θ_k that receive γ_k ⊙ (ψ_k − x) as input each step. The paper does not restrict the weights, so convergence of the data likelihood is no longer guaranteed. The fully parametrised recurrent matrix makes it applicable to sequences §2.3 fig. 2.

Loss. The training loss has an intra-cluster term, the responsibility-weighted negative log-likelihood, and an inter-cluster term that penalises each component for predicting pixels it does not own. The inter-cluster term is a KL divergence between the pixel prior and each prediction, weighted by 1 − γ §2.4 eq. 5. The authors add that standard tricks, such as a bottleneck or noisy inputs, help avoid trivial solutions, and that sequences can use a next-step prediction loss §2.4. In practice they stop gradients through γ, so the loss does not reward changing responsibilities §2.4. The appendix applies this stop only to flying shapes app. A.2.

Sequences. At each time step the input to RNN-EM is γ_k(ψ_k^{(t−1)} − x̃^{(t)}), where x̃ is the current frame with bitflip noise. The next-step loss replaces x with x^{(t+1)} in eq. 5 and is evaluated at every step §4.2 app. A.2. For the test-time next-step loss, the paper uses the maximum over components of ψ instead of the responsibility-weighted mixture, so the reported loss is an upper bound §4.2.

Pseudocode. The following follows the paper's equations and the official code. Where they differ, the comment says which is which.

# One (RN)N-EM step on a single input x  [§2.2] [§2.3]
# Input:  x (D,), thetas (K, M), preds (K, D), gammas (K, D)
# Output: thetas (K, M), preds (K, D), gammas (K, D)

def em_step(x, thetas, preds, gammas):
    # Gamma-weighted residual per component. The paper writes
    # gamma * (psi - x); the code uses gamma * (x - psi)  [§2.3] [eq. 4]
    # [code: nem_model.py#L100] [code: nem_model.py#L113-L116]
    inputs = gammas * (x - preds)                    # (K, D)
    # Update each component representation  [§2.3]
    thetas, preds = update_components(inputs, thetas)
    # E-step: per-pixel responsibilities  [§2.2] [eq. 3]
    gammas = e_step(x, preds)                        # (K, D)
    return thetas, preds, gammas

# E-step: responsibilities for each pixel  [§2.2] [eq. 3]
# Mixing weights pi are not used; equal weights are implied
# [code: nem_model.py#L131-L166]

def e_step(x, preds):
    # Likelihood of each pixel under each component, with a small
    # epsilon added so no gamma is exactly zero
    # [code: nem_model.py#L153-L155]
    likelihoods = pixel_likelihood(x, preds) + 1e-6   # (K, D)
    # Normalise over components, so the K gammas of a pixel sum to 1
    # [code: nem_model.py#L164]
    return likelihoods / sum_over_components(likelihoods)

# Training loss for one frame and one prediction step  [§2.4] [eq. 5]
# The code adds the inter term with a positive sign  [code: nem_model.py#L253]

def nem_loss(x, preds, gammas, prior, inter_weight):
    # Intra-cluster term: gamma-weighted pixel negative log-likelihood
    # [§2.4] [code: nem_model.py#L234] [code: nem_model.py#L246-L247]
    intra = gammas * neg_log_likelihood(x, preds)    # (K, D)
    # Inter-cluster term: KL between pixel prior and prediction, weighted
    # by (1 - gamma)  [§2.4] [code: nem_model.py#L235] [code: nem_model.py#L248]
    inter = (1 - gammas) * kl_divergence(prior, preds)  # (K, D)
    # Sum over pixels and components; the code divides by batch size
    # [code: nem_model.py#L242-L243]
    return sum(intra) + inter_weight * sum(inter)

Equations (numbering from the LaTeX source; the paper has five numbered equations).

P(\mathbf{x}\mid\boldsymbol\theta)=\prod_{i=1}^{D}\sum_{k=1}^{K}\pi_k\,P(x_i\mid\psi_{i,k},z_{i,k}=1) eq. 1

\mathcal{Q}(\boldsymbol\theta,\boldsymbol\theta^{\text{old}})=\sum_{\mathbf z}P(\mathbf z\mid\mathbf x,\boldsymbol\psi^{\text{old}})\log P(\mathbf x,\mathbf z\mid\boldsymbol\psi) eq. 2

\gamma_{i,k}=P(z_{i,k}=1\mid x_i,\psi_{i}^{\text{old}}) eq. 3

\boldsymbol\theta^{\text{new}}=\boldsymbol\theta^{\text{old}}+\eta\frac{\partial\mathcal{Q}}{\partial\boldsymbol\theta},\qquad \frac{\partial\mathcal{Q}}{\partial\boldsymbol\theta_k}\propto\sum_{i=1}^{D}\gamma_{i,k}(\psi_{i,k}-x_i)\frac{\partial\psi_{i,k}}{\partial\boldsymbol\theta_k} eq. 4

L(\mathbf x)=-\sum_{i=1}^{D}\sum_{k=1}^{K}\gamma_{i,k}\log P(x_i,z_{i,k}\mid\psi_{i,k})-\sum_{i=1}^{D}\sum_{k=1}^{K}(1-\gamma_{i,k})\,D_{KL}[P(x_i)\,\|\,P(x_i\mid\psi_{i,k},z_{i,k})] eq. 5

The sign of the second term in eq. 5 is discussed under Pitfalls.

Network architectures (layer lists in the appendix; the code uses the same layers, see Paper vs code).

• Static shapes: RNN-EM has a 250-unit sigmoid RNN, which receives the γ-weighted residual of the 28×28 pixels, and an equivalent output layer with a sigmoid for each pixel §4.1 app. A.1. N-EM uses a single fully connected layer from θ to the per-pixel outputs with a sigmoid app. A.1.
• Flying shapes: input is a 28×28 binary frame, then 4×4 stride-2 convolution with 32 ELU units and layer norm, 4×4 stride-2 convolution with 64 ELU units and layer norm, a 512-unit fully connected ELU layer with layer norm, a 100-unit sigmoid RNN with layer norm on its output, then the decoder: 512-unit ReLU fully connected, 7×7×64 ReLU fully connected, and two nearest-neighbour upsamplings each followed by a 4×4 convolution (32 ReLU units with layer norm, then one sigmoid output channel) app. A.2.
• Flying MNIST: the masked differences fed to the network are normalised to zero mean and unit standard deviation app. A.3. Three 4×4 stride-2 convolutions with 32, 64 and 128 ELU units and layer norm, a 512-unit ELU fully connected layer, a 250-unit sigmoid RNN with layer norm on its output, then a 512-unit ReLU fully connected layer, a 3×3×128 ReLU fully connected layer, and three upsample-and-convolve decoder blocks (64 ReLU, 32 ReLU, 1 linear output channel) app. A.3.

Hyperparameters

Results

Author-reported. Setup: synthetic datasets built from shapes and MNIST digits, AMI against ground-truth groups, averages over five runs §4.

Static shapes §4.1:

The paper says Tagger uses about twenty times more parameters §4.1.

Flying shapes, Table 1 tab. 1. Train and test columns use the same configuration. The generalisation columns test with a different number of objects or K.

Flying MNIST §4.3:

On flying shapes, next-step prediction error is lower for RNN-EM with K=5 than for the K=1 recurrent autoencoder, and the gap grows with the number of objects fig. 5. The text gives no numbers for this comparison. The AMI curve over 50 time steps on flying MNIST is in fig. 6.

Paper vs code

Official repository: github.com/sjoerdvansteenkiste/Neural-EM, linked from the paper's footnote §4, MIT licence, copyright Sjoerd van Steenkiste. Commit 062437c027d1e2e0a8699eb9c798258a6f5b62ba. Read: nem.py, nem_model.py, network.py, README.md. datasets.py and utils.py were not read, so dataset generation and the AMI routine were not compared.

• K for static shapes. The main text says K=3 §4.1. The appendix and README use K=4 app. A.1 README.md README.md. The code default is K=3 nem_model.py.
• Gradient through γ. The paper stops γ-gradients as a general practice §2.4 and in the flying-shapes appendix app. A.2. The code default does not stop them nem_model.py. Only the sequential config does nem_model.py. The static-shape commands in the README do not set it, so gradients do flow through γ there README.md.
• Gaussian variance. The paper uses a fixed σ² = 0.25 for flying MNIST §4.3. The code uses σ = 1.0 (variance 1) in the outer loss nem_model.py and e_sigma = 0.25 as the standard deviation in the E-step, so the E-step variance is 0.0625 nem_model.py nem_model.py.
• Sign of the residual. The paper's residual is γ(ψ − x) §2.3 eq. 4. The code computes data - predictions nem_model.py. For RNN-EM the learned RNN can absorb the sign. For N-EM it is absorbed by the learned scalar learning rate network.py.
• N-EM weight tying. The paper says f_φ is a single fully connected layer with a sigmoid output app. A.1 and that an extra weight learns the learning rate app. A.1. The code also ties the decoder to the transpose of the encoder weight network.py and applies the sigmoid-derivative factor in the update network.py. The paper does not state either choice.
• Metric during training. The paper's metric is AMI §4. The code computes ARI in the training graph nem.py and computes AMI only in run_from_file nem.py.
• Step weighting. The static default outer loss is on the last step only nem_model.py nem_model.py. The sequential config sums the loss over all steps with weight 1 nem_model.py nem_model.py. The paper does not say how steps are combined.
• Flying shapes inter-cluster weight. The paper gives 0.2 only for flying MNIST app. A.3. The code default of 1.0 applies to flying shapes nem_model.py.
• Input noise. The code uses bitflip for binary data and masked uniform for grey-scale data, selected from the dataset type nem.py. This matches the paper's text.
• Early stopping. The code stops when the best validation loss so far is below the minimum of the last ten validation losses nem.py. The paper says "not improved for 10 epochs" §4. These are close but not identical.

Pitfalls

• K for static shapes is ambiguous. The main text says K=3, the appendix and README say K=4. The paper says a K above the number of objects is not harmful §4.2 tab. 1, but the static-shape numbers depend on the choice. Check the run config.
• Eq. 5 has the wrong sign on the KL term, or the paper's text does. Eq. 5 subtracts the (1 − γ)-weighted KL, which as a loss would reward a larger divergence. The text says the term penalises the divergence §2.4, and the code adds it with a positive sign nem_model.py. Implement the positive version.
• Residual sign. The paper's γ(ψ − x) and the code's γ(x − ψ) differ by a sign. This is harmless only if the learned parts absorb it. A hand-written N-EM update must follow the sign of the implementation you use.
• Gaussian σ. σ² = 0.25 in the paper §4.3 versus σ = 1 outer and σ = 0.25 (std) in the code's E-step nem_model.py nem_model.py. The E-step variance then differs from the paper's by a factor of four.
• Mixing weights are dropped. Eq. 1 and eq. 3 include π_k §2.1 §2.2, but the paper never sets π. The code normalises without π, which is the same as uniform weights nem_model.py.
• Initialisation is not in the paper. The code draws the initial γ as absolute standard normal values normalised over K, and starts predictions at zero nem_model.py nem_model.py. Components can therefore start from random responsibilities, so run-to-run variance may depend on the seed. The paper reports five runs but does not state seeds.
• Loss scale. The loss is summed over pixels and components and divided by batch size nem_model.py, not averaged over pixels. Changing the inter-cluster weight therefore changes the balance of the two terms, not only their scale.
• Gradient clipping and epoch cap. The paper does not give either. The code disables clipping and caps training at 500 epochs nem.py.
• Staged training for flying MNIST. The stages (20, 500, 50,000 digits) §4.3 use learning rate 0.0005 only from stage two on app. A.3. The README runs follow this order README.md.
• Figure and equation numbers were counted from the LaTeX source, since no compiled PDF was available. Figures 1 to 7 and equations 1 to 5 follow source order, and the floating figures may be numbered differently in the compiled paper. Check against the arXiv PDF before citing.
• Dependencies are old. The README pins TensorFlow 1.2.1 and sacred 0.7.0 README.md. The datasets are not in the repository; the README links an external download README.md.
• Costs not stated. The paper gives no runtime, memory or GPU figures for unrolling 15 to 20 steps with back-propagation through time.

Lineage

• Builds on: Tagger, which the paper cites for the idea of unsupervised perceptual grouping with K networks §3 §4.1. No arXiv id appears in the source, and Tagger has no card here.
• Related: RTagger, a recurrent Tagger extension developed in parallel (footnote, §3); learned unrolled sparse-coding inference (Gregor et al., 2010; bibliography key gregor2010learning) §3. N-EM extends this line by deriving the update from EM §3.
• Compared against: Tagger on static shapes §4.1.
• Superseded by: not known from content in this wiki.