Stochastic EM, Deterministic Annealing, and SAEM¶
This notebook demonstrates the Phase 2 (v0.4 roadmap) algorithm family: algorithm='sem'
(Stochastic EM, Celeux & Diebolt 1985), temperature (Deterministic Annealing EM, Ueda &
Nakano 1998), and algorithm='saem' (Stochastic Approximation EM, Delyon, Lavielle & Moulines
1999) -- all compared against standard EM and CEM from the same deliberately bad
initialization.
import torch
import matplotlib.pyplot as plt
import os
os.chdir('../')
from tgmm import GaussianMixture
random_state = 42
A known saddle initialization¶
Three well-separated true clusters, but all three components are initialized jittered around a single one of them (cluster 0) -- a classic EM failure mode: two components end up fighting over the one cluster they started near, while the third true cluster is left with no component to claim it. Standard (soft) EM and CEM both get stuck here; SEM's stochastic sampling and DAEM's tempered E-step both escape it, as shown below.
torch.manual_seed(0)
centers = [(0.0, 0.0), (8.0, 0.0), (4.0, 7.0)]
X = torch.cat([torch.randn(150, 2) + torch.tensor(c) for c in centers])
torch.manual_seed(3)
init_means = torch.tensor([0.0, 0.0]).unsqueeze(0).repeat(3, 1) + torch.randn(3, 2) * 0.3
init_covariances = torch.stack([torch.eye(2)] * 3)
plt.figure(figsize=(5, 5))
plt.scatter(X[:, 0], X[:, 1], s=8, alpha=0.4, label='data')
plt.scatter(init_means[:, 0], init_means[:, 1], c='red', marker='x', s=120, label='initial means (all 3)')
plt.legend()
plt.title('Saddle initialization: all 3 components start near one cluster')
plt.gca().set_aspect('equal')
plt.show()
Fit EM, CEM, SEM, and DAEM from the identical initialization¶
common_kwargs = dict(
n_components=3, n_features=2, random_state=random_state, max_iter=100,
init_means=init_means.clone(), init_covariances=init_covariances.clone(),
)
gmm_em = GaussianMixture(algorithm='em', **common_kwargs)
gmm_em.fit(X)
gmm_cem = GaussianMixture(algorithm='cem', **common_kwargs)
gmm_cem.fit(X)
gmm_sem = GaussianMixture(algorithm='sem', **common_kwargs)
gmm_sem.fit(X)
gmm_daem = GaussianMixture(temperature=(0.9, 30), **common_kwargs)
gmm_daem.fit(X)
for name, gmm in [('EM', gmm_em), ('CEM', gmm_cem), ('SEM', gmm_sem), ('DAEM', gmm_daem)]:
print(f'{name:5s} final log-likelihood: {gmm.lower_bound_:.4f}')
EM final log-likelihood: -4.4467 CEM final log-likelihood: -4.4475 SEM final log-likelihood: -3.9811 DAEM final log-likelihood: -3.9810
Log-likelihood traces, plotted together¶
EM and CEM plateau early, stuck at the same (worse) local optimum. SEM's trace is visibly
noisier -- note the "stationary band" it settles into rather than a smooth monotone curve, since
every iteration re-samples a hard assignment instead of converging to a fixed point -- but it
spends most of its time at the better optimum. DAEM's trace shows the tempered,
lower-likelihood early iterations (while beta < 1) before it locks onto the better solution as
beta reaches 1.0.
plt.figure(figsize=(8, 5))
plt.plot(gmm_em.log_likelihood_history_, label='EM', marker='.')
plt.plot(gmm_cem.log_likelihood_history_, label='CEM', marker='.')
plt.plot(gmm_sem.log_likelihood_history_, label='SEM', alpha=0.8)
plt.plot(gmm_daem.log_likelihood_history_, label='DAEM', linewidth=2)
plt.xlabel('iteration / sweep')
plt.ylabel('mean log-likelihood')
plt.title('EM vs CEM vs SEM vs DAEM from the same saddle initialization')
plt.legend()
plt.show()
Zooming in: SEM's stationary band¶
Unlike EM/CEM/DAEM, SEM never converges pointwise -- fit() always runs the full max_iter
sweeps (see algorithm='sem''s docstring). The tail of its trace fluctuates in a narrow band
around the best achievable log-likelihood instead of settling exactly, since each sweep samples a
fresh hard assignment from the responsibilities rather than taking the argmax.
tail = gmm_sem.log_likelihood_history_[-50:]
plt.figure(figsize=(8, 3))
plt.plot(tail, marker='.')
plt.axhline(gmm_daem.lower_bound_, color='C3', linestyle='--', label='DAEM/best optimum')
plt.xlabel('sweep (last 50)')
plt.ylabel('mean log-likelihood')
plt.title(f'SEM stationary band (std = {tail.std().item():.2e})')
plt.legend()
plt.show()
n_empty_components_ counts how many times a component emptied out and had to be re-seeded
during the run -- a non-zero count is a hint that n_components may be larger than the data
supports.
print('n_empty_components_:', gmm_sem.n_empty_components_)
print('annealing_schedule_ (first 10 values):', gmm_daem.annealing_schedule_[:10])
n_empty_components_: 0 annealing_schedule_ (first 10 values): [0.9, 0.9031663723922961, 0.9063438846891775, 0.9095325760829622, 0.9127324859038545, 0.9159436536204298, 0.9191661188401216, 0.9223999213097097, 0.9256451009158111, 0.928901697685371]
SAEM: the "SEM that converges"¶
SAEM blends each sweep's sampled-assignment statistics into a running average with a decaying
step size (step_size, defaulting to 'auto' for this algorithm), so -- unlike plain SEM --
tol-based convergence applies and fit() can stop early. Demonstrated on a fresh,
well-separated dataset here rather than continuing with the adversarial initialization above:
SAEM's own stochastic trajectory from that specific saddle point happens to settle into a
different, also-mediocre local optimum (not shown) -- a reminder that none of these methods
guarantee escaping every bad initialization, they each just change the odds. What SAEM
reliably does is match plain EM when both are given the same, reasonable starting point:
torch.manual_seed(1)
X2 = torch.cat([torch.randn(2000, 2) + c for c in
[torch.tensor([0.0, 0.0]), torch.tensor([10.0, 0.0])]])
gmm_em2 = GaussianMixture(n_components=2, n_features=2, algorithm='em', random_state=0, max_iter=200)
gmm_em2.fit(X2)
gmm_saem2 = GaussianMixture(n_components=2, n_features=2, algorithm='saem', random_state=0, max_iter=200)
gmm_saem2.fit(X2)
print('EM means:\n', gmm_em2.means_)
print('SAEM means:\n', gmm_saem2.means_)
print(f'SAEM converged: {gmm_saem2.converged_}, iterations: {gmm_saem2.n_iter_}')
plt.figure(figsize=(8, 4))
plt.plot(gmm_saem2.log_likelihood_history_, marker='.')
plt.xlabel('iteration')
plt.ylabel('mean log-likelihood')
plt.title('SAEM log-likelihood trace (converges, unlike plain SEM)')
plt.show()
EM means:
tensor([[ 3.5266e-03, 1.3034e-02],
[ 1.0016e+01, -3.9754e-02]])
SAEM means:
tensor([[ 3.5266e-03, 1.3034e-02],
[ 1.0016e+01, -3.9754e-02]])
SAEM converged: True, iterations: 1
Takeaways¶
- EM (default) and CEM (hard assignments) both got stuck at the same local optimum from the saddle initialization above.
- SEM escaped it by sampling hard assignments instead of taking the
argmax, at the cost of never settling to a single point -- usesem_estimate='best'/'mean'over the post-burn-in sweeps to get a final answer out of a noisy trace. - DAEM escaped it too, by tempering the E-step early on so the likelihood surface is smoother while the fit is still deciding where to settle.
- SAEM gets SEM's exploration with a decaying step size, so it converges like ordinary EM -- and matches EM closely when both start from a reasonable initialization.
See the Fitting Algorithms section of the
user guide for batch_size/step_size-driven mini-batch and online EM, which aren't shown in
this notebook.