Expectation maximization

Overview

Expectation maximization (EM) fits a model that has hidden variables. Directly maximising the marginal likelihood is hard, so EM maximises the expected complete-data log-likelihood Q instead, computed under the current parameters 1708.03498 §2.2 eq. 2. Every variant on this page uses the same two steps. The E-step computes responsibilities γ, the probability that each hidden assignment is active, from the current predictions 1708.03498 §2.2 eq. 3. The M-step updates the parameters to increase Q.

Three cards cover this concept so far. The Gaussian-mixture variant comes from a robot-mapping paper, where the hidden variable is the cluster of each sample 1810.04285 §3.8. The three neural variants come from the neural-grouping paper. The latent-bucket variant comes from a robot-evaluation paper, where the hidden variable is the task of each comparison rather than a pixel assignment 2506.18123 §3.2.

Variants

Gaussian-mixture EM over time-extended features

The hypertime model fits a Gaussian mixture over vectors that hold a value, a location and a cosine and sine pair for each periodicity found so far 1810.04285 §3.5 eq. 1 eq. 2. The number of clusters k is fixed, and the fit uses the standard EM algorithm from an OpenCV implementation 1810.04285 §3.8. When the covariances become unstable, the paper checks their eigenvalues and restarts from new initial cluster positions. The threshold and the number of restarts are not stated 1810.04285 §3.8. The second clustering method in the paper, HyT-KM, starts from k-means centres and chooses k itself. Its EM step with cosine distance for the time dimensions is not fully specified 1810.04285 §3.8. The paper reports lower prediction error than Mean, Hist and FreMEn on four robot datasets, but only in figures 1810.04285 fig. 2–fig. 5. Card: Warped Hypertime Representations for Long-term Autonomy of Mobile Robots.

Gradient-step M-step

When the decoder f_φ is non-linear, the M-step has no closed form. The paper takes one gradient step on Q with step size η instead. For each component, the gradient is a responsibility-weighted residual times the Jacobian of the component's output with respect to its representation 1708.03498 §2.2 eq. 4. Card: Neural Expectation Maximization.

Unrolled EM (N-EM)

Repeat the E- and M-steps a fixed number of times and back-propagate the loss through all of them into the network weights φ. This makes the clustering trainable end to end 1708.03498 §2.3. The static-shape experiments use 15 EM steps 1708.03498 app. A.1. The authors report AMI 0.475 ± 0.043 for N-EM on three static shapes, against 0.826 ± 0.005 for RNN-EM [arxiv:1708.03498 §4.1]. Card: Neural Expectation Maximization.

Recurrent EM (RNN-EM)

Replace the hand-derived update with an unrestricted recurrent network. Its hidden states receive the γ-weighted residual at each step. The weights are not restricted, so convergence of the data likelihood is no longer guaranteed. The fully parametrised recurrent matrix is what makes the variant applicable to sequences 1708.03498 §2.3 fig. 2. Card: Neural Expectation Maximization.

Latent-bucket EM for pairwise ranking

Here the hidden variable is a task bucket for each pairwise trial. The paper argues that a plain Bradley-Terry fit treats task effects as noise, and it adds a latent bucket per trial to absorb them 2506.18123 §3.2. The E-step computes the posterior over buckets for each trial, from the bucket prior and the likelihood of the observed outcome under each bucket. The M-step takes a clipped Newton step on the policy strengths, the per-bucket offsets and the bucket difficulties, centring after each. It then updates the bucket prior and the tie rate 2506.18123 alg. 2. The paper's default is 60 EM iterations 2506.18123 tab. 1. Card: RoboArena: Distributed Real-World Evaluation of Generalist Robot Policies. The ranking model and its outcome probabilities are worked out on the Bradley-Terry page.

When to use which

Implementation sketch

The loop below follows the paper's equations. Where variants differ, the comment says so.

# Unrolled EM on one input x, trained end to end  [arxiv:1708.03498 §2.3]
# Input: image x, number of steps T, step size eta
# Output: training loss, back-propagated into the network weights phi

def unrolled_em_loss(x, T, eta):
    # Initial states are not stated in the paper; pick a scheme and report it
    thetas, gammas = initial_state()
    for t in range(T):
        # M-step: one gradient step on Q  [arxiv:1708.03498 eq. 4]
        thetas = thetas + eta * grad_Q(thetas, gammas, x)
        # E-step: responsibilities from the updated predictions  [eq. 3]
        gammas = e_step(x, predict(thetas))
    # Static images use the loss at the final step only  [app. A.1]
    # Sequences sum the loss over steps instead  [§4.2]
    return em_loss(x, thetas, gammas)    # [arxiv:1708.03498 eq. 5]

For RNN-EM, replace the body of grad_Q with one recurrent update whose input is the γ-weighted residual 1708.03498 §2.3.

The Gaussian-mixture variant does not unroll EM into a network. It runs EM as a standard mixture fit inside a loop over periodicities, and the paper gives its loop in its own pseudocode 1810.04285 §3.3. Its period search is a spectral scan over the error, not an EM update 1810.04285 §3.7.

For the latent-bucket variant, the E-step and M-step are the ones listed under its variant above, and the paper's update formulas are in its appendix 2506.18123 app. B.2. The full loop is on the Bradley-Terry page.

Contradictions

• K for static shapes. The main text says K=3 1708.03498 §4.1. The appendix says K=4 1708.03498 app. A.1. The numbers in the results depend on which one was used.
• Sign of the KL term. Eq. 5 subtracts the (1 − γ)-weighted KL divergence 1708.03498 eq. 5. The text describes the term as a penalty on predicting pixels a component does not own 1708.03498 §2.4. Read literally, the equation would reward the divergence.
• Gradients through γ. The paper stops gradients through γ in general 1708.03498 §2.4 and in the flying-shapes appendix [arxiv:1708.03498 app. A.2]. The authors' released code does not stop them by default, and the repository's static-shape commands do not stop them either.
• Latent-bucket EM in the paper and its code. The paper's default ranking is this EM fit 2506.18123 §3.2. The public leaderboard code has the EM routine disabled and fits a Bradley-Terry model with Davidson ties instead. The card records this, and that the code's strength gradients use the opponent's solve probability, while the appendix gives the own-policy form. The card notes that only the EM model is affected. See the paper card's Paper vs code section.

Papers

• Warped Hypertime Representations for Long-term Autonomy of Mobile Robots: fits a fixed-k Gaussian mixture by EM over time-projected features, to predict p(a | x, t) for robot states that follow routines.
• Neural Expectation Maximization: unrolls EM into a differentiable clustering of pixels into objects, and adds RNN-EM, which replaces the update with a recurrent network and handles sequences.
• RoboArena: Distributed Real-World Evaluation of Generalist Robot Policies: fits a task-aware Bradley-Terry model with latent task buckets by EM, to rank robot policies from distributed A/B trials.