Relational Neural Expectation Maximization: Unsupervised Discovery of Objects and their Interactions

In plain words

Watching a video of balls bouncing off each other, a person can tell which parts of the picture belong to which ball and predict what happens next. This paper trains a neural network to do the same thing with no labels at all: it only sees the frames.

The network keeps one internal "slot" per object. It learns to sort the pixels of each frame into these slots, so that each slot describes one ball (or one piece of the scene). Then, at every step, each slot is updated by looking at the other slots. A learned attention score decides how much each other object should matter, which fits the idea that balls only affect each other when they touch. Because every object is handled by the same small network, the same rules can be applied to scenes with more or fewer objects than seen in training.

The approach is tested on bouncing balls, on balls behind an invisible curtain, and on frames from the Atari game Space Invaders. It predicts the next frame better than plain recurrent baselines and keeps track of hidden balls.

TL;DR

• R-NEM extends Neural Expectation Maximization (N-EM) with a relational interaction function that updates each object's state from pairwise effects §2.2 eq. 5.
• Training is unsupervised: the network predicts the next frame, and the loss is an EM-style per-pixel likelihood term with a responsibility weighting §2.1 eq. 3.
• Author-reported: on 4-ball bouncing balls the soft assignments reach an ARI of about 0.8 §4; with a hidden curtain R-NEM's test BCE is 46.22 against 94.64 for an RNN §4.
• Use it when you want unsupervised object slots with explicit interactions. Expect fragile details: the attention sum depends on K, and the reported numbers come from synthetic 2D data.

Problem

Plain autoencoders and RNNs learn one distributed code for the whole frame. That code mixes information about all objects, so it cannot factor the interaction between a pair of objects, and learning these efficiently lacks an essential inductive bias §2 §2.2. Supervised object-based models do learn interactions, but they need object states from a simulator or other external source that is usually unavailable §1. An MLP over all object states removes that dependency but loses the ability to extrapolate to other numbers of objects §2.2.

Method

Neural Expectation Maximization (N-EM). A frame is modelled as a spatial mixture of K components, each with a parameter vector theta_k. A network f_phi maps each theta_k to per-pixel Bernoulli means psi_k. Pixel responsibilities gamma are the posterior of the hidden assignments §2.1 eq. 1 eq. 2. N-EM unrolls generalised EM and trains through the unrolled steps with back-propagation through time §2.1 eq. 3. The recurrent variant RNN-EM replaces the M-step with a recurrent network whose input is the masked residual gamma_k (x - psi_k) §2.1.

Relational N-EM (R-NEM). The recurrent update is generalised so that the new state of object k depends on all other object states through an interaction function Upsilon §2.2 eq. 4:

\theta_k^{(t)} = \sigma\big(W\,\tilde{x}^{(t)} + R\,\Upsilon_k(\theta^{(t-1)})\big) \qquad [\text{eq. 4}]

Upsilon_k is defined as §2.2 eq. 5:

\Upsilon_k(\theta) = [\hat\theta_k;\, E_k], \quad \hat\theta_k = \mathrm{MLP}^{enc}(\theta_k), \quad E_k = \sum_{i \neq k} \alpha_{k,i}\, e_{k,i}

\alpha_{k,i} = \mathrm{MLP}^{att}(\xi_{k,i}), \quad e_{k,i} = \mathrm{MLP}^{eff}(\xi_{k,i}), \quad \xi_{k,i} = \mathrm{MLP}^{emb}([\hat\theta_k;\hat\theta_i]) \qquad [\text{eq. 5}]

The paper states that without attention (R-NEM no att) every alpha equals 1 app. A.1. Setting the attention to 1 removes the relevance weighting but keeps the pairwise sum.

Training: each input frame is corrupted with bit-flip noise (p = 0.2 for the balls, 0.02 for Atari in the paper) app. A.1 app. A.2. The network predicts the next frame, with the prior for each pixel set to Bernoulli with p = 0 and no gradient through gamma app. A.1. Sequences are 30 steps long for the balls and 25 for Space Invaders §4. The paper gives the architecture in app. A.1 and app. A.2. The balls encoder is three 4x4 stride-2 convolutions (16, 32, 64 ELU channels, layer norm) followed by a 512-unit ELU layer; the recurrent state has 250 sigmoid units with layer norm on the output; the decoder mirrors the encoder with nearest-neighbour upsampling and 4x4 convolutions app. A.1.

Pseudocode for one step of the unrolled recurrence, with the E-step, follows §2.1 §2.2:

# One R-NEM time step  [§2.2] [eq. 4]
# Inputs:  frame (B, H, W, 1) noisy; next_frame (B, H, W, 1)
#          theta (B, K, 250); psi (B, K, H, W, 1); gamma (B, K, H, W, 1)
# Outputs: theta, psi, gamma for the next time step

def r_nem_step(frame, next_frame, theta, psi, gamma):
    # 1. Residual of each component, masked by its responsibility  [§2.1]
    #    gamma is treated as a constant here  [code: nem_model.py#L155-L158]
    residuals = stop_gradient(gamma) * (frame - psi)
    # 2. Encode each component's residual map  [app. A.1]
    #    [code: network.py#L17-L25]
    encoded = encoder(residuals)                  # (B, K, 512)
    # 3. Interaction function Upsilon  [eq. 5]
    #    [code: network.py#L384-L434]
    interaction = upsilon(theta)                  # (B, K, 500)
    # 4. Recurrent update: one linear map, layer norm, sigmoid  [eq. 4]
    #    [code: network.py#L445-L447] [code: network.py#L496-L498]
    theta = sigmoid(layer_norm(linear(concat(encoded, interaction))))
    # 5. Decode each object's reconstruction  [app. A.1]
    #    [code: network.py#L42-L50]
    psi = decoder(theta)                          # (B, K, H, W, 1)
    # 6. E-step: responsibilities from the next-frame likelihood  [§2.1]
    #    [code: nem_model.py#L140-L147]
    gamma = e_step(psi, next_frame)               # (B, K, H, W, 1)
    return theta, psi, gamma

The interaction function, with the attention term written out eq. 5:

# Interaction function Upsilon  [eq. 5]  [code: network.py#L384-L434]
# Input: theta (B, K, 250); output: (B, K, 500)

def upsilon(theta):
    # 1. Encode every object state with MLP_enc  [eq. 5]
    theta_hat = mlp_enc(theta)                    # (B, K, 250)
    summed_effects = []
    for focus in range(K):
        total = zeros(250)
        for context in range(K):
            if context == focus:
                continue
            # 2. Pair embedding of focus and context  [eq. 5]
            xi = mlp_emb(concat(theta_hat[focus], theta_hat[context]))
            # 3. Effect of context on focus, and attention weight  [eq. 5]
            effect = mlp_eff(xi)
            alpha = mlp_att(xi)                   # scalar in (0, 1)
            # 4. Attention-weighted sum over the K - 1 context objects
            total = total + alpha * effect        # [eq. 5]
        summed_effects.append(total)
    # 5. Concatenate the encoded state and the summed effect  [eq. 5]
    return concat(theta_hat, stack(summed_effects))

The loss for one step, with the detached responsibilities eq. 3:

# Loss for one time step  [eq. 3]  [code: nem_model.py#L186-L222]
# Prior p = 0 for every pixel  [app. A.1]

def step_loss(psi, gamma, next_frame):
    # 1. Per-pixel binary cross-entropy of each component  [eq. 3]
    #    [code: nem_model.py#L186-L189]
    bce = binary_cross_entropy(psi, next_frame)   # (B, K, H, W, 1)
    # 2. Intra-cluster term, weighted by responsibility  [eq. 3]
    #    gamma is detached in the code  [code: nem_model.py#L216]
    intra = stop_gradient(gamma) * bce
    # 3. Inter-cluster term: (1 - gamma) times KL(prior || psi)  [eq. 3]
    #    With p = 0 this equals -log(1 - psi) per pixel
    #    [code: nem_model.py#L193-L196] [code: nem_model.py#L217]
    inter = (1 - stop_gradient(gamma)) * kl_bernoulli(0.0, psi)
    # 4. Sum over pixels and components, average over the batch
    #    [code: nem_model.py#L207-L219]
    return sum(intra + inter) / batch_size

Unrolled training runs the step above for 30 time steps, sums the step losses with equal weight and divides by the number of steps nem_model.py. The optimiser is Adam with default parameters, batch size 64, and early stopping after 10 epochs without improvement §4 app. A. The paper does not state a learning rate; the code sets 0.001 nem.py.

Hyperparameters

Results

Author-reported. The paper reports results as figures and short text quotes; it gives no full results table. Setting: bouncing balls, 4-ball test set, 64 x 64 frames §4 fig. 3.

Figure 3's loss values are shown relative to a baseline that always predicts the current frame; the text does not give those values.

Paper vs code

Code: official repository, sjoerdvansteenkiste/Relational-NEM, linked from the paper's experiments section §4 (footnote). Licence file is MIT (copyright Sjoerd van Steenkiste) LICENSE. The commit is the HEAD of the shallow clone used for this card. The repo pins TensorFlow 1.2.1 and sacred 0.7.2 README.md. I did not run it.

• Bouncing-ball architecture, interaction function, attention, loss, E-step and noise match the paper. Compared: app. A.1 against network.py, network.py, network.py, nem_model.py, nem_model.py, nem.py.
• Atari encoder and decoder differ. The paper lists four convolutions in the encoder (16, 32, 32, 32 channels) and a 512-unit fully connected layer, and a decoder with a 512-unit layer, an 8x8x64 reshape, three upsampling convolutions and a one-channel output convolution app. A.2. enc_dec_84_atari has two convolutions (16 with an 8x8 kernel and stride 4, then 32 with a 4x4 kernel and stride 2), a 250-unit fully connected layer, and a decoder with a 250-unit layer, a 10x10x32 reshape and two output convolutions network.py. The README's Atari command uses this config README.md.
• Atari input noise differs. The paper says 0.02 app. A.2; the README command sets noise.prob=0.002 README.md.
• Atari batch size differs. The paper uses 64 for all experiments app. A; the README command sets training.batch_size=32 README.md.
• Learning rate. The paper says "default parameters" app. A; the code sets 0.001 explicitly nem.py. This matches the TensorFlow Adam default, but the paper does not say so.
• Numerical floors not in the paper. The E-step adds 1e-6 to each likelihood before normalising nem_model.py. The BCE clips predictions to [1e-6, 1 - 1e-6] nem_model.py, and the KL clips inside its logs nem_model.py.
• Simulation threshold matches. The paper says the max over components is thresholded at 0.1 app. A.1; the rollout does that at nem.py.
• Upper-bound loss. The reported "upper bound" replaces the E-step with a max over components §4 footnote nem_model.py.
• No code for the data. The repository has a dataset loader and points to a Dropbox folder with data and pretrained models README.md. The simulator that produced the balls data is not in the repo, and the paper does not state how velocities are sampled.

Pitfalls

• The attention sum is not normalised. Attention weights are sigmoid outputs summed over K - 1 context objects network.py. The paper attributes worse extrapolation without attention to this sum changing range with K §4 Attention. Test-time K different from training K changes the scale of the input to the next state.
• Responsibilities are constants for the gradient. gamma is detached in the residual masking nem_model.py and in both loss terms nem_model.py. Letting gradients flow through gamma would change training; the paper states only that gradients are not back-propagated through gamma app. A.1.
• Initialisation is not stated in the paper. The code starts the recurrent state at zero, predictions at 0.0, and gamma as the normalised absolute value of Gaussian noise over K nem_model.py.
• The inter-cluster term with p = 0 is -log(1 - psi) per pixel. Derived from the KL formula in the code; the paper writes it as KL with the prior eq. 3. The probability clipping matters when psi is near 1 nem_model.py.
• Loss scale is not stated. Eq. 3 has no batch or time averaging; the code divides by batch size and by the number of steps nem_model.py nem_model.py. Adam is roughly invariant to a constant rescaling, but the reported loss values depend on this choice.
• K must be chosen. The paper uses K = 5 for training four balls, K = 8 at test time on 6-8 balls, and K = 4 on Atari app. A.1 app. A.2. It notes that extra empty components can improve grouping §4 Comparison, but does not give a rule for choosing K.
• Rollouts depend on thresholding. Simulation replaces the E-step with a threshold on the max prediction app. A.1 nem.py. The paper reports that failures happen when one component models two balls §4 Simulation.
• ARI excludes background and overlaps app. A. Quoted ARI values are not comparable to ARI on all pixels.
• Results come from synthetic 2D data. The balls are binary 64 x 64 renders, and the Atari frames are binarised DQN recordings §4. The paper does not test natural images.
• Old framework. The code uses tf.contrib and TensorFlow 1.2.1 README.md. Porting to current TensorFlow or PyTorch needs an explicit check of the recurrent cell and the layer-norm placement.
• Compute and runtime are not stated. The paper gives no training time or hardware use.

Lineage

• Builds on: Neural Expectation Maximization (Greff et al.), which R-NEM uses as its object-grouping backbone §2.1 §3. R-NEM adds the relational interaction function Upsilon §2.2.
• Inspired by: relational interaction models for physics and compositional object-based approaches §2.2, including the distance-based context selection that the attention term replaces §2.2 §4 Attention. No card in this wiki.
• Compared against: RNN, LSTM, RNN-EM and R-NEM without attention §4.
• Superseded by: not known from wiki content.