What it actually does
One Gaussian describes a single blob; real data is usually several, so describe it as a weighted sum of Gaussians and let the fit decide how much of each. The obstacle is that you never observe which component produced a point — that assignment is the latent variable — so you alternate. Given the current components, compute how responsible each is for each point. Given those responsibilities, refit each component to the points it owns, weighted by ownership. The defining feature is that ownership is fractional: a point between two clusters belongs 60/40 rather than being forced into one. k-means is this same loop with the responsibilities rounded to 0 and 1 and every covariance forced spherical and equal, which is exactly why k-means cannot represent an elongated or overlapping cluster and this can.
- The data is generated by a finite mixture of Gaussians — an approximation, and a good one for smooth multi-modal densities, a poor one for heavy tails or bounded supports
- The number of components is known or selectable; the likelihood cannot choose it, since more components always fit better
- Enough points per component to estimate its covariance — a full covariance needs d(d+1)/2 parameters, which is the constraint that fails first as dimension grows
- Points are exchangeable: the model has no notion of order, sequence, or spatial adjacency
- Components are non-degenerate, which is enforced by a covariance floor rather than guaranteed by the data
Likelihood
How probable is the observed data under this mixture, summed over every possibility for the component that generated each point. The sum sitting inside the logarithm is what makes this hard: with the sum outside, each component would separate and the fit would be a closed-form Gaussian estimate per group. Inside, nothing factorizes. The whole purpose of the E-step is to get the sum back out — it replaces the log of a sum with an expected complete-data log-likelihood, which does factorize, at the cost of being only a lower bound on the thing you actually wanted.
- mixing weight — the prior probability a point came from component k, summing to one
- mean and covariance of component k; the covariance is what lets a cluster be elongated or tilted
- number of components, which is a choice rather than a fitted quantity
- the marginalization over the unobserved component assignment — and the reason there is no closed form
How it is minimized
Expectation-Maximization — alternating closed-form updates that ascend a lower bound on the log-likelihood
- responsibility: the posterior probability that point i came from component k — the E-step, and the whole latent variable
- the effective number of points owned by component k, which replaces a hard count everywhere a count would appear
- the M-step: ordinary weighted maximum-likelihood estimates, with responsibilities as the weights
- updated the same way — a responsibility-weighted scatter matrix about the new mean
Each EM iteration maximizes a lower bound on the log-likelihood that touches it exactly at the current parameters, which is what makes the guarantee unusually clean: the likelihood cannot decrease, there is no step size, and there is nothing to diverge. Both steps are closed form — the M-step is just weighted Gaussian maximum likelihood, which is why the algorithm is short. What is not guaranteed is the global optimum. The surface is non-convex, it has K-factorial equivalent modes from relabelling components, and — the part that matters in practice — the likelihood is unbounded above: shrink one component onto a single point and its density diverges. So a higher likelihood can be a degenerate solution rather than a better fit, which is why every serious implementation carries a covariance floor.
K (n_components)
The number of components. Not learnable by likelihood — more components always fit better — so it is chosen by BIC, or by what the components are supposed to mean (2 to 20 by BIC; beyond that a mixture is rarely the right model)
covariance_type
full, tied, diagonal, or spherical. This is the real capacity dial: full costs K·d(d+1)/2 parameters and diagonal costs K·d (diagonal above ~50 dimensions, full below)
reg_covar
Floor added to the covariance diagonal. Not a nicety — it is what stops a component collapsing onto one point and reporting infinite likelihood (1e-6 relative to the feature scale)
n_init / init_params
Restarts and initialization. k-means initialization plus several restarts is standard, because a single random start regularly converges somewhere bad
tol / max_iter
Stop on the relative change in log-likelihood; overlapping components converge slowly and can need hundreds of iterations (tol 1e-3 on average log-likelihood, 100 to 500 iterations)
Monotone by construction — the log-likelihood increases every iteration and converges to a stationary point. That stationary point is where the trouble is. Random initialization regularly lands in a poor local optimum, so restarts are standard practice rather than diligence. Component collapse is the named failure: with a full covariance and no floor, a component that captures a single point drives its determinant to zero, its density to infinity, and the reported likelihood to positive infinity — a "perfect" fit that is entirely an artefact. Label switching means component 2 in one run is component 5 in the next, so anything downstream keyed to component index breaks silently across refits. And heavily overlapping components converge very slowly, because the responsibilities barely move between iterations.
Per iteration: O(n·K·d^2) for full covariances, or O(n·K·d) for diagonal, plus O(K·d^3) for the Cholesky factorizations. Memory is O(n·K) for the responsibility matrix, which is the array that actually constrains batch size, plus O(K·d^2) for the covariances. Typically tens to low hundreds of iterations, so a GMM is comfortably cheaper than the kernel methods next door — its limit is dimension, not sample count.
How, where, and why it is used
Time-Series Forecasting
AdaptedNot as a forecaster on its own — it has no representation of order at all — but as the distributional component inside a temporal model. The two standard placements are as the emission density of a hidden Markov or state-space model, where the mixture describes what is observed in each regime, and as an unsupervised regime labeller over engineered features, whose component assignment becomes a covariate for a separate forecaster.
- Emission distributions in HMMs for regime-switching series — the standard formulation in speech and in financial regime models
- Volatility- or condition-regime identification, feeding a regime indicator into a downstream forecaster
- Mixture-density outputs where a forecast needs to be multi-modal, since a single Gaussian cannot represent "either a spike or nothing"
It earns its place wherever the predictive distribution is genuinely multi-modal, which a squared-error forecaster cannot express: the mean of two plausible futures is often a value that will never occur. What it cannot supply is the temporal structure — fit a GMM to a series directly and it treats the observations as exchangeable, discarding the ordering that made it a time series. It is a component of a temporal model, never the model.
- Fit on engineered features (returns, rolling volatility, spectral summaries) rather than raw levels, which drift and make the components non-stationary
- Standardize before fitting, since the covariance floor and the initialization are both scale-dependent
- Pin the component ordering after fitting — by mean, say — or label switching makes the regime indicator meaningless across refits
Out-of-sample log-likelihood on a temporal split, not on shuffled folds. Where the mixture feeds a forecaster, evaluate the forecaster: a regime split that improves BIC and does nothing for downstream error was a description, not a signal.
- Treating exchangeable-data component assignments as regimes, when nothing in the model prevents the regime from flipping every observation
- Non-stationarity moving the components underneath a model that assumes they are fixed
- Selecting K by likelihood on the training window, which always prefers more components
Anomaly Detection
Primary fitFit the mixture to normal data and score new points by their log-likelihood under it: an anomaly is a point the fitted density considers improbable. Because the model is a density rather than a boundary, the score is continuous and comparable across the space, and the per-component responsibilities also say which mode of normality the point failed to belong to.
- Multivariate process and equipment monitoring with several distinct normal operating regimes
- Fraud and account-behaviour models where legitimate behaviour is genuinely multi-modal
- Screening inputs to a deployed model for falling outside the training distribution
- Background modelling in vision, where a per-pixel mixture separates background from foreground
It is the natural choice when normal behaviour is several distinct things rather than one: a single-Gaussian or distance-to-centroid detector will place the boundary in the empty space between two legitimate modes and alarm on everything there. Being generative also makes the threshold interpretable, since it is a density level rather than an arbitrary score. The limits are dimensional and structural — density estimation degrades quickly above a few tens of features, and a full-covariance mixture needs enough points per component to estimate d-squared parameters, so on wide data an isolation forest or a reconstruction-based detector is the better answer.
- Fit on a confirmed-clean window; contamination pulls a component out to cover the anomalies and they become normal
- Reduce dimension first — PCA, or a learned embedding — because covariance estimation is what fails as d grows, well before the feature count looks alarming
- Set the covariance floor relative to the feature scale after standardization, not as an absolute constant
- Choose K by BIC on clean data, then check that no component owns a negligible share, which usually means it has latched onto a handful of points
Precision@k and PR-AUC against confirmed incidents, with the threshold read from the log-likelihood quantiles of a clean period. Hold out entire anomaly episodes. Report the effective size of each component too — a fit where one component owns 1% of the mass is usually one collapse away from a useless density.
- Component collapse producing an enormous likelihood at one point and a badly-shaped density everywhere else
- Choosing K to maximize likelihood, which drives K toward n and turns the model into a memoriser
- Assuming components correspond to meaningful regimes; a mixture approximates a density, and a curved single mode is routinely fitted by three Gaussians that mean nothing individually
Optimization
AdaptedEM is the canonical latent-variable optimization scheme, and the mixture is the example it is always taught on. The transferable content is the construction: build a lower bound on an intractable objective that touches it at the current parameters, maximize the bound instead, repeat. That is the majorize-minimize pattern, it is where the variational objective in a VAE comes from, and it explains both the monotonicity guarantee and its limits.
- The reference example for EM, coordinate ascent on a bound, and why monotone improvement does not imply a global optimum
- Soft clustering feeding a downstream allocation problem, where a fractional assignment is more useful to an optimizer than a hard one
- Model selection by BIC as a worked example of penalizing likelihood by parameter count
It is the smallest complete example of optimizing something you cannot write down. Ordinary gradient methods are available here and are rarely used, because the closed-form alternating updates are both faster and free of step-size tuning — a case where exploiting the structure of the problem beats a general solver outright. The counterweight is worth stating: the same structure that gives monotone ascent gives no protection against local optima, and the objective itself is unbounded, so "the optimizer improved the objective" is not evidence the fit got better.
- Standardize so the covariance floor and convergence tolerance mean the same thing in every dimension
- Initialize from k-means, which is the cheap approximation of the same objective and lands EM in a far better basin than random starts
Verify monotonicity of the log-likelihood across iterations — a decrease is a bug, and a specific one, almost always a missing responsibility normalization or a covariance update taken about the wrong mean. Compare the best of several restarts rather than trusting one run.
- Reading monotone improvement as convergence to the global optimum
- Chasing an unbounded objective into a degenerate solution because the covariance floor was omitted
- Stopping on a fixed iteration count when overlapping components need hundreds of iterations to separate
Computer Vision
AdaptedThe famous deployed use is background subtraction: fit an independent small mixture to the colour history of each pixel, and classify a new observation as background if it falls under a well-supported component and foreground otherwise. The mixture is what lets a pixel have several legitimate appearances — a swaying branch, a flickering display — without those being flagged as motion.
- Per-pixel background modelling in fixed-camera surveillance and traffic monitoring, still deployed in embedded systems
- Fisher-vector encodings, where a GMM over local descriptors provides the soft vocabulary that a classifier consumes
- Skin and colour segmentation with a small mixture in a chosen colour space
It works here because the per-pixel problem is genuinely low-dimensional — three colour channels — which is precisely the regime where density estimation is reliable, and because the multi-modality is real rather than assumed. It also updates online cheaply, which matters for a model running per pixel per frame. On whole images it is superseded outright: learned features handle the appearance variation that a mixture over raw pixels cannot.
- Work in a colour space where the channels are closer to independent, so a diagonal covariance is defensible and cheap
- Keep K small per pixel — three to five — since each pixel has few observations and more components simply memorize them
- Use an online update with a learning rate rather than a batch refit, which is what makes the per-pixel cost affordable at frame rate
Foreground F-measure against annotated frames, evaluated separately on the hard cases the model exists for — gradual illumination change, repetitive motion, and objects that stop and become background.
- A foreground object that stops moving is gradually absorbed into the background model, which is inherent to the online update rather than a bug
- Sudden global illumination change invalidating every pixel model at once
- Per-pixel independence ignoring spatial structure entirely, so the output needs morphological cleanup to be usable
Risk & Fraud
ViableUnsupervised segmentation of accounts or transactions into behavioural modes, then either scoring the density directly for novelty or handing the responsibilities to a supervised model as features. The soft assignment matters here: an account that is 50/50 between two behaviour profiles is itself informative, and a hard cluster label discards that.
- Customer and merchant behavioural segmentation ahead of a supervised fraud model
- Novelty detection on account activity where labelled fraud is scarce and delayed
- Population-shift monitoring — refit periodically and watch whether the mixing weights move
Fraud labels arrive late and incomplete, so an unsupervised model of normal behaviour has a genuine role, and legitimate behaviour is multi-modal in a way a single profile cannot capture. The reasons for caution are the usual ones plus a domain-specific one: fraudulent behaviour is often not low-density but rather a legitimate-looking mode of its own, and a mixture fitted on contaminated data will happily give it a component and call it normal.
- Aggregate to account-level behavioural features over a window rather than modelling raw transactions, which are too heavy-tailed for a Gaussian
- Transform skewed monetary features toward symmetry before fitting; a Gaussian mixture over untransformed amounts spends every component on the tail
- Standardize and reduce dimension, since covariance estimation is what fails first
Where labels exist, PR-AUC of the density score at a fixed alert budget on an out-of-time split. Where they do not, evaluate stability: refit on adjacent windows and check that components and mixing weights are recognizably the same, because an unstable segmentation cannot support an operational process.
- Contamination giving fraud its own component, after which it is scored as perfectly normal
- Label switching breaking any downstream rule that references a component by index
- Heavy tails forcing extra components that model outliers rather than structure — a t-distribution mixture is the honest fix
What it costs to run
Cheap and bounded: tens to low hundreds of iterations, each O(n·K·d^2). A mixture over a million points in twenty dimensions fits in seconds to minutes on one core. What costs is the restarts — several initializations are standard, and they multiply the wall time directly.
K Mahalanobis distances and a log-sum-exp, so O(K·d^2) per point with full covariances, or O(K·d) with diagonal. Microseconds, and the model is small — a handful of means and covariances — which makes it easy to embed and easy to update online.
Periodic, driven by how fast the population shifts rather than by cost. There is a genuine online variant that updates the sufficient statistics incrementally, which is what makes the per-pixel and streaming uses possible; the batch fit is the one that needs scheduling.
- Track average out-of-sample log-likelihood over time — a steady fall means the density no longer describes the population, and it moves before any downstream metric does
- Watch the mixing weights across refits; a component quietly losing all its mass is a structural change worth investigating rather than absorbing
- Monitor the smallest covariance eigenvalue per component, which is the early warning of a collapse the floor is about to clamp
- Alert if the component count selected by BIC changes on refit, since anything keyed to component identity is about to be wrong
- Component indices are not stable across refits — label switching is inherent, so downstream systems must key on a canonicalized ordering rather than the raw index
- The covariance floor must be persisted with the model and applied identically at scoring time; a model fitted with a floor and scored without one can return infinite densities
- Score in log space end to end: densities in high dimensions underflow to exactly zero, and a probability of zero for every component makes the responsibility a division of zero by zero
- The feature scaler is part of the model, because both the floor and the fitted covariances are expressed in scaled units
- A likelihood threshold calibrated at one K does not transfer to another; changing K silently changes the scale of the score
Pros and cons, in context
Soft assignment carries information that hard clustering discards
A point that is 60/40 between two components is described as such, which matters when the assignment feeds a decision. Irrelevant when all anyone consumes is a single cluster label, where k-means is cheaper and simpler.
Clusters may be elongated, tilted, and of different sizes
The covariance is a fitted parameter, which is exactly what k-means lacks — its spherical equal-variance assumption is why it splits a long cluster in half. The cost is d-squared parameters per component, so the advantage inverts in high dimensions.
The output is a full density, not just an assignment
One fitted model answers clustering, density estimation, novelty scoring, and sampling. That breadth is why it appears under three different domains here; it is worth nothing if only the labels are used.
EM gives monotone improvement with no step size to tune
No learning rate, no divergence, and a convergence check that is meaningful. A genuine practical advantage over gradient fitting of the same objective — as long as monotone improvement is not mistaken for a global optimum.
The likelihood is unbounded, so the objective can be driven to infinity
A component collapsing onto a single point reports infinite density and a perfect fit. Unlike most failure modes here this one rewards itself, which is why the covariance floor is mandatory rather than defensive.
Non-convex, with results that depend on initialization
Different seeds give materially different fits, so restarts are part of the method rather than a precaution. It also means reproducibility requires pinning the seed and the initialization scheme, not just the data.
Covariance estimation degrades quickly with dimension
Full covariances need d(d+1)/2 parameters per component, so beyond a few tens of features the choice is between diagonal covariances — losing the main advantage over k-means — and not using the model.
K must be chosen, and likelihood cannot choose it
More components always fit better, so selection needs BIC or an external criterion. On real data BIC frequently keeps rising with K, at which point the honest conclusion is that the density is not a small mixture and the model is a description rather than a discovery.
Intuitive, idiomatic, optimized
The same algorithm in three languages, each progressing through Kent Beck's three stages. Every optimization names the lever it pulls and the trade-off it accepts — cost profiles are illustrative, not measured benchmarks.
The math, transcribed literally. No dependencies.
"""Gaussian mixture by EM - the E and M steps, transcribed.
Diagonal covariances, so each component is an axis-aligned ellipse and the
density is a product of one-dimensional Gaussians. That keeps the arithmetic
visible: no matrix inverse, no determinant, just d scalar terms per component.
"""
import math
def gaussian_density(x, mean, variance):
"""N(x | mean, diag(variance)) as an explicit product over dimensions."""
density = 1.0
for j in range(len(x)):
difference = x[j] - mean[j]
density *= math.exp(-0.5 * difference * difference / variance[j]) / math.sqrt(
2.0 * math.pi * variance[j]
)
return density
def fit(X, k, iterations=100, floor=1e-6):
n = len(X)
d = len(X[0])
# Initialize: means on the first k points, unit variances, uniform weights.
weights = [1.0 / k for _ in range(k)]
means = [list(X[i]) for i in range(k)]
variances = [[1.0] * d for _ in range(k)]
for _ in range(iterations):
# ---- E-step: responsibilities gamma[i][c] -------------------------
gamma = [[0.0] * k for _ in range(n)]
for i in range(n):
total = 0.0
for c in range(k):
gamma[i][c] = weights[c] * gaussian_density(X[i], means[c], variances[c])
total += gamma[i][c]
# If every component gives this point density zero the posterior is
# 0/0. In double precision that happens routinely above ~40
# dimensions, which is what the log-space rewrite later fixes.
if total == 0.0:
for c in range(k):
gamma[i][c] = 1.0 / k
else:
for c in range(k):
gamma[i][c] /= total
# ---- M-step: weighted maximum likelihood --------------------------
for c in range(k):
effective_n = 0.0
for i in range(n):
effective_n += gamma[i][c]
weights[c] = effective_n / n
for j in range(d):
weighted_sum = 0.0
for i in range(n):
weighted_sum += gamma[i][c] * X[i][j]
means[c][j] = weighted_sum / effective_n
for j in range(d):
weighted_scatter = 0.0
for i in range(n):
difference = X[i][j] - means[c][j]
weighted_scatter += gamma[i][c] * difference * difference
# The floor stops a component collapsing onto one point, where
# the variance goes to zero and the likelihood to infinity.
variances[c][j] = max(weighted_scatter / effective_n, floor)
return weights, means, variancesO(iterations * n * k * d) in interpreter loops, and numerically fragile: densities are multiplied in linear space.
How this language actually wants it written.
"""Gaussian mixture by EM - full covariances, Cholesky, log space."""
from dataclasses import dataclass
import numpy as np
from numpy.typing import NDArray
from scipy.linalg import solve_triangular
Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]
Tensor = NDArray[np.float64]
@dataclass(frozen=True)
class GaussianMixture:
"""A fitted mixture. Full covariances, so components may be tilted."""
weights: Vector # (K,)
means: Matrix # (K, d)
covariances: Tensor # (K, d, d)
log_likelihood: float
converged: bool
def score_samples(self, X: Matrix) -> Vector:
"""log p(x) per point - the anomaly score, and it stays in log space."""
return _log_sum_exp(_log_component_densities(X, self), axis=1)
def predict_proba(self, X: Matrix) -> Matrix:
log_joint = _log_component_densities(X, self)
return np.exp(log_joint - _log_sum_exp(log_joint, axis=1)[:, None])
def _log_sum_exp(values: Matrix, axis: int) -> Vector:
"""Factor the max out before exponentiating; the naive version underflows."""
shift = values.max(axis=axis, keepdims=True)
return np.squeeze(shift, axis=axis) + np.log(
np.exp(values - shift).sum(axis=axis)
)
def _log_component_densities(X: Matrix, model: "GaussianMixture") -> Matrix:
"""(n, K) matrix of log(pi_k) + log N(x | mu_k, Sigma_k)."""
n_components, dimension = model.means.shape
out = np.empty((X.shape[0], n_components), dtype=np.float64)
for component in range(n_components):
# Cholesky, never an explicit inverse: solving with the triangular
# factor is both cheaper and far better conditioned, and the
# log-determinant falls out of its diagonal for free.
factor = np.linalg.cholesky(model.covariances[component])
centred = X - model.means[component]
solved = solve_triangular(factor, centred.T, lower=True)
mahalanobis = np.einsum("ji,ji->i", solved, solved)
log_det = 2.0 * np.log(np.diag(factor)).sum()
out[:, component] = (
np.log(model.weights[component])
- 0.5 * (mahalanobis + log_det + dimension * np.log(2.0 * np.pi))
)
return out
def fit(
X: Matrix,
n_components: int,
max_iter: int = 200,
tol: float = 1e-4,
reg_covar: float = 1e-6,
seed: int = 0,
) -> GaussianMixture:
"""Fit by EM. Raises ValueError on malformed input."""
if X.ndim != 2:
raise ValueError(f"X must be 2-D, got shape {X.shape}")
if not 1 <= n_components <= X.shape[0]:
raise ValueError(f"n_components must lie in [1, {X.shape[0]}]")
if reg_covar <= 0.0:
raise ValueError("reg_covar must be positive; without it a component can collapse")
n, dimension = X.shape
rng = np.random.default_rng(seed)
weights = np.full(n_components, 1.0 / n_components)
means = X[rng.choice(n, size=n_components, replace=False)].copy()
covariances = np.tile(np.cov(X.T) + reg_covar * np.eye(dimension), (n_components, 1, 1))
previous = -np.inf
converged = False
model = GaussianMixture(weights, means, covariances, previous, False)
for _ in range(max_iter):
# ---- E-step ------------------------------------------------------
log_joint = _log_component_densities(X, model)
log_evidence = _log_sum_exp(log_joint, axis=1)
responsibilities = np.exp(log_joint - log_evidence[:, None])
# ---- M-step ------------------------------------------------------
effective_n = responsibilities.sum(axis=0) + np.finfo(np.float64).eps
weights = effective_n / n
means = (responsibilities.T @ X) / effective_n[:, None]
covariances = np.empty((n_components, dimension, dimension))
for component in range(n_components):
centred = X - means[component]
scatter = (responsibilities[:, component, None] * centred).T @ centred
covariances[component] = scatter / effective_n[component] + reg_covar * np.eye(
dimension
)
current = float(log_evidence.mean())
model = GaussianMixture(weights, means, covariances, current, False)
if abs(current - previous) < tol:
converged = True
break
previous = current
return GaussianMixture(weights, means, covariances, previous, converged)What changed: Two changes, and only one of them is about style. The model changes: diagonal variances become full covariance matrices, so a component can be tilted rather than axis-aligned — which is the capability that distinguishes a GMM from k-means, bought at d(d+1)/2 parameters per component. The implementation changes to match: densities are computed in log space with a log-sum-exp that factors out the maximum, because multiplying densities in linear space underflows to exactly zero in double precision above roughly forty dimensions and turns the responsibility into 0/0. And the covariance is used through a Cholesky factor rather than an explicit inverse — the Mahalanobis distance becomes a triangular solve and the log-determinant falls out of the factor’s diagonal, which is cheaper and far better conditioned.
O(n*K*d^2) per iteration plus O(K*d^3) for the factorizations, executed in LAPACK rather than the interpreter.
Faster, and the justification for each change.
"""Gaussian mixture by EM - cached precisions, batched E-step, in place."""
import numpy as np
from numpy.typing import NDArray
from scipy.linalg import solve_triangular
Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]
Tensor = NDArray[np.float64]
def _cholesky_precisions(covariances: Tensor, reg: float) -> tuple[Tensor, Vector]:
"""Factor every covariance ONCE per iteration, not once per use.
Stores the inverse Cholesky factor, so the Mahalanobis distance becomes a
matrix product rather than a triangular solve per batch - the solve is
done K times per iteration instead of K times per batch.
"""
n_components, dimension, _ = covariances.shape
inverse_factors = np.empty_like(covariances)
log_dets = np.empty(n_components, dtype=np.float64)
identity = np.eye(dimension)
for component in range(n_components):
factor = np.linalg.cholesky(covariances[component] + reg * identity)
inverse_factors[component] = solve_triangular(factor, identity, lower=True)
log_dets[component] = -2.0 * np.log(np.diag(factor)).sum()
return inverse_factors, log_dets
def fit_fast(
X: Matrix,
n_components: int,
max_iter: int = 200,
tol: float = 1e-4,
reg_covar: float = 1e-6,
batch: int = 65_536,
) -> tuple[Vector, Matrix, Tensor]:
"""EM with the two accelerations that matter at scale.
The responsibility matrix is (n, K) and is the array that actually bounds
batch size, so the E-step is chunked and the M-step sufficient statistics
are accumulated across chunks rather than materialized whole.
"""
design = np.ascontiguousarray(X, dtype=np.float64)
n, dimension = design.shape
weights = np.full(n_components, 1.0 / n_components)
means = np.ascontiguousarray(design[:n_components].copy())
covariances = np.tile(np.cov(design.T), (n_components, 1, 1))
log_2pi = dimension * np.log(2.0 * np.pi)
previous = -np.inf
# Sufficient statistics, allocated once and reused every iteration.
stat_n = np.empty(n_components, dtype=np.float64)
stat_sum = np.empty((n_components, dimension), dtype=np.float64)
stat_scatter = np.empty((n_components, dimension, dimension), dtype=np.float64)
for _ in range(max_iter):
inverse_factors, log_dets = _cholesky_precisions(covariances, reg_covar)
log_weights = np.log(weights)
stat_n.fill(0.0)
stat_sum.fill(0.0)
stat_scatter.fill(0.0)
total_log_evidence = 0.0
for start in range(0, n, batch):
block = design[start : min(start + batch, n)]
log_joint = np.empty((block.shape[0], n_components), dtype=np.float64)
for component in range(n_components):
# (x - mu) L^-T, one GEMM; the squared row norms are the
# Mahalanobis distances, fused by einsum without an intermediate.
whitened = (block - means[component]) @ inverse_factors[component].T
log_joint[:, component] = log_weights[component] + 0.5 * (
log_dets[component] - log_2pi
) - 0.5 * np.einsum("ij,ij->i", whitened, whitened)
shift = log_joint.max(axis=1, keepdims=True)
np.subtract(log_joint, shift, out=log_joint)
np.exp(log_joint, out=log_joint) # now unnormalized
normalizer = log_joint.sum(axis=1, keepdims=True)
total_log_evidence += float((np.log(normalizer) + shift).sum())
np.divide(log_joint, normalizer, out=log_joint) # responsibilities
# Accumulate sufficient statistics; the (n, K) matrix never exists.
stat_n += log_joint.sum(axis=0)
stat_sum += log_joint.T @ block
for component in range(n_components):
weighted = log_joint[:, component, None] * block
stat_scatter[component] += weighted.T @ block
weights = stat_n / n
means = stat_sum / stat_n[:, None]
for component in range(n_components):
outer = np.outer(means[component], means[component])
covariances[component] = stat_scatter[component] / stat_n[component] - outer
current = total_log_evidence / n
if abs(current - previous) < tol:
break
previous = current
return weights, means, covariancesWhat changed: Three structural changes. The Cholesky factorization moves out of the per-use path and is computed once per component per iteration, with the inverse factor cached so the Mahalanobis distance becomes a GEMM rather than a triangular solve repeated for every batch. The E-step is chunked and the M-step is rewritten to accumulate sufficient statistics across chunks, so the (n, K) responsibility matrix — the array that actually bounds how much data a GMM can see at once — is never materialized whole. And the log-sum-exp normalization runs in place through the same buffer, which removes three full-size temporaries per chunk from the innermost loop.
Vectorize to NumPy/BLAS (@, np.dot, einsum)
Whitening by the cached inverse Cholesky factor is a GEMM, and the squared row norms that follow are an einsum reduction that never materializes an intermediate.
Trade-off: Caching the inverse factor is a step back from the numerically cleanest form — a triangular solve is better conditioned than multiplying by a precomputed inverse — so this trades a little accuracy for a lot of throughput, and is wrong on a near-singular covariance.
Batch work to amortize interpreter overhead
Chunking the E-step bounds the responsibility block at (batch, K) instead of (n, K), which is what makes the fit possible on data far larger than that matrix.
Trade-off: The M-step must be restated in terms of accumulated sufficient statistics, which is a less obvious form and introduces a real numerical hazard — the scatter-minus-outer-product identity loses precision when the mean is large relative to the spread.
Pre-allocate output arrays and use in-place operations
Sufficient statistics are allocated once for the whole fit and the log-sum-exp normalization writes back through its own buffer, removing several full-size temporaries per chunk per iteration.
Trade-off: The same buffer holds log-joint, then unnormalized weights, then responsibilities — three meanings in three lines, which is fast and genuinely harder to read, and reordering those lines corrupts the E-step silently.
O(n*K*d^2) per iteration in BLAS, memory bounded by the chunk. Illustrative, not a measured benchmark.