---
kind: concept
slug: expectation-maximization
title: Expectation maximization
summary: >-
  Fitting a model with hidden assignments by alternating responsibilities (E-step) and parameter updates (M-step). Covers Gaussian-mixture EM over time-extended features for robot models, a gradient-step M-step for non-linear decoders, unrolled EM trained end to end, a recurrent variant, and latent-task-bucket EM for ranking robot policies. Three papers so far.
tags: [unsupervised-learning, clustering, mixture-models]
papers: ["1708.03498", "1810.04285", "2506.18123"]
related: [clustering-evaluation, gaussian-mixture-model, k-means-clustering, deep-unfolding, recurrent-neural-network, bradley-terry-model]
status: draft
updated_at: 2026-11-01
---

## 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
[arxiv: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 [arxiv: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
[arxiv: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 [arxiv: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
[arxiv: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
[arxiv: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 [arxiv: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 [arxiv:1810.04285 §3.8]. The paper reports lower
prediction error than Mean, Hist and FreMEn on four robot datasets, but only in
figures [arxiv:1810.04285] [fig. 2]–[fig. 5]. Card: [Warped Hypertime Representations
for Long-term Autonomy of Mobile Robots](/papers/1810.04285).

### 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 [arxiv:1708.03498 §2.2]
[eq. 4]. Card: [Neural Expectation Maximization](/papers/1708.03498).

### 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 [arxiv:1708.03498 §2.3]. The static-shape experiments use
15 EM steps [arxiv: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](/papers/1708.03498).

### 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 [arxiv:1708.03498 §2.3] [fig. 2]. Card: [Neural Expectation
Maximization](/papers/1708.03498).

### 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 [arxiv: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 [arxiv:2506.18123 alg. 2]. The paper's default is 60 EM iterations
[arxiv:2506.18123 tab. 1]. Card: [RoboArena: Distributed Real-World Evaluation
of Generalist Robot Policies](/papers/2506.18123). The ranking model and its
outcome probabilities are worked out on the [Bradley-Terry page](/concepts/bradley-terry-model).

## When to use which

| Situation | Use | Why |
|---|---|---|
| Samples with routine or periodic structure in time, and you need a full distribution over the value given place and time | Gaussian-mixture EM on hypertime features | The paper's periods are found from the model error, so no period has to be known in advance [arxiv:1810.04285 §3.7]. Its candidate period set is not stated |
| Fixed set of images, and you need the unrolled procedure to train a network | Unrolled EM | The update is derived from EM, and the loss is applied through the unrolled steps [arxiv:1708.03498 §2.3] |
| You need the data likelihood to keep converging across steps | Unrolled EM, not RNN-EM | The paper says RNN-EM's unrestricted weights remove the convergence guarantee [arxiv:1708.03498 §2.3] |
| Input is a sequence and you need next-step prediction | RNN-EM | The paper's sequence experiments use RNN-EM, with a next-step loss at each step [arxiv:1708.03498 §4.2] |
| Number of components is unknown | Set K above the expected count | The authors report that unused components stay empty and that a large K is not harmful [arxiv:1708.03498 §4.2] |
| Pairwise outcomes where the task or scene changes from trial to trial, and you need one ranking | Latent-bucket EM | The paper argues that a plain fit treats task effects as noise, and reports the closest agreement with an exhaustive oracle among the methods it compares [arxiv:2506.18123 §3.2] [§5.2] |
| Compute or memory budget is tight | Check before choosing the unroll length | The paper gives no runtime, memory or GPU figures for back-propagating through 15 to 20 steps |

## Implementation sketch

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

```pseudocode
# 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 [arxiv: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 [arxiv:1810.04285 §3.3]. Its period search is a
spectral scan over the error, not an EM update [arxiv: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
[arxiv: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 [arxiv:1708.03498 §4.1]. The
  appendix says K=4 [arxiv: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
  [arxiv:1708.03498 eq. 5]. The text describes the term as a penalty on
  predicting pixels a component does not own [arxiv:1708.03498 §2.4]. Read
  literally, the equation would reward the divergence.
- **Gradients through γ.** The paper stops gradients through γ in general
  [arxiv: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 [arxiv: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/2506.18123).

## Papers

- [Warped Hypertime Representations for Long-term Autonomy of Mobile Robots](/papers/1810.04285):
  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](/papers/1708.03498): 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](/papers/2506.18123):
  fits a task-aware Bradley-Terry model with latent task buckets by EM, to
  rank robot policies from distributed A/B trials.