neural-em-shapes

Greff, K., van Steenkiste, S., & Schmidhuber, J. (2017). Neural Expectation Maximization. NIPS 2017 (arXiv:1708.03498).

Problem

Unsupervised perceptual grouping. Given a binary image containing several non-overlapping objects, partition the foreground pixels into K slots so each slot binds to a single object — without ever showing the model a segmentation label.

The mechanism is a differentiable Expectation–Maximization loop. Each of the K slots carries a hidden state θ_k ∈ R^H that is decoded into a per-pixel Bernoulli mean μ_k = σ(W_dec θ_k + b_dec). One EM step is

E-step      γ_{k,i}  = softmax_k log p(x_i | μ_{k,i})        (uniform prior)
            r_{k,i}  = γ_{k,i} · (x_i − μ_{k,i})
M-step      θ_k_new  = tanh(W_x r_k + W_h θ_k + b_h)

The mixture negative log-likelihood is summed across T unrolled iterations and minimised end-to-end with Adam. Slot-binding emerges when the M-step amplifies tiny per-slot differences in μ_k so that each slot’s responsibility (γ) sharpens onto a single object.

This stub trains and evaluates on the static-shapes condition (Greff 2017, §4.1) re-implemented from scratch in numpy.

Dataset

24 × 24 binary canvas, 3 random shapes per image drawn from {square, disc, triangle} with half-size 2–4 px. Light overlap is permitted; pixel-level ground-truth labels record which shape generated each foreground pixel for evaluation only (the model never sees them). Foreground fraction ≈ 0.21.

Architecture

θ_0[b, k] = θ_init[k] + Gaussian(0, init_noise_std) per image. A bottleneck of H = 24 (vs. D = 576 pixels) is what stops the slots collapsing onto a single shared “predict-the-union” mode: each slot can only encode 24 dims of variation, so the K slots must cooperate to cover the 3 objects.

Files

Running

Headline (≈ 17 s on M-series CPU):

python3 neural_em_shapes.py --seed 0

This runs a numerical-gradient check (3 ms, ≤ 1e-5 relative error) and then 30 epochs over a 1024-image train set with batch 32.

Quick smoke (≈ 1 s, 3 epochs, 256 train images):

python3 neural_em_shapes.py --seed 0 --quick

Then regenerate viz:

python3 visualize_neural_em_shapes.py
python3 make_neural_em_shapes_gif.py

Results

Headline run, --seed 0 defaults (canvas=24, K=3, T=4, H=24, n_train=1024, batch=32, lr=3e-3, epochs=30, noise_p=0.10):

NMI rises sharply over the first ~7 epochs then partially collapses (see viz/nmi_curve.png). The N-EM loss continues to decrease even as NMI declines: the model trades slot specialisation for tighter overall reconstruction, so the best-NMI checkpoint (epoch 7) is what the headline visualisation uses.

Hyperparameters

Visualizations

Deviations from the original

Open questions / next experiments

• Full AMI rather than NMI. Greff 2017 reports AMI = 0.96 on static shapes. Re-deriving AMI in numpy and running the same comparison on this dataset would tell us how much of our 0.43 NMI is metric choice vs. capacity gap.
• Background slot. The paper’s K+1 setup with one dedicated “background” slot is the simplest fix for the slot-collapse drift. Adding it should let the foreground slots specialise harder, and we expect peak NMI to climb past 0.6.
• Larger M-step. A 2-layer or GRU-style recurrence (closer to the paper) is the natural next step. The minimal tanh we use here is the floor of expressiveness; what does the slot-collapse curve look like with more capacity?
• Bottleneck schedule. H is the single biggest knob — at H = 16 NMI is similar but loss is higher; at H = 64 there is no binding at all. A small scan over H × T would map the regime where binding is stable.
• Per-iteration loss weighting. Equal weighting across T encourages early iterations to converge to a usable θ. Up-weighting the final iteration (or final-only loss) marginally tightens reconstructions but accelerates collapse — there is probably a sweet spot.
• Recurrent N-EM (RNEM) on flying shapes. Once the static case is solid, the natural extension is the temporal version where slots track objects across frames. That is relational-nem-bouncing-balls in this catalog.
• ByteDMD instrumentation (v2). Each EM iteration re-reads the full image once per slot. The data-movement cost should scale roughly linearly with K × T at fixed image size; whether learned slot states reduce data movement vs. naive K-means is exactly the v2 question.