Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Gaussian Mixture Models and Expectation-Maximization › Implementation in NumPy and SciPy

Michael BrenndoerferOctober 29, 202613 min read

Machine Learning from ScratchAll chapters

Part 9 of 15Unsupervised Learning

Chapter 66Gaussian Mixture Models and Expectation-Maximization › Implementation in NumPy and SciPy

Reading section 6 of 10 · Implementation in NumPy and SciPy · 13 minSection 6 of 10 · 13 min

Section 6 of 10 · 13 min read

Gaussian Mixture Models and Expectation-Maximization › Implementation in NumPy and SciPySection view · Sections in this chapter6 / 10 ↓

Sections in this chapter10

Reading tools Definitions, links & notes

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Notes & comments

Notes & commentsLeave notes and join the discussionSelection toolbarShow comment actions when selecting text. Turn off for the browser menu only.Comment indicatorsShow small counts beside annotated passagesPassage highlightsGently shade annotated text selections

We now build the algorithm from primitives, keeping every step visible. We use NumPy for a Cholesky factorization (to evaluate multivariate normal densities without explicitly forming an inverse) and logsumexp from SciPy (to compute the normalizing constant of the responsibilities in a numerically safe way). The linear algebra functions np.linalg.cholesky and np.linalg.solve are standard NumPy.

PythonIn [7]
CollapseExpanddef log_normal_pdf(X, mean, cov):
    """Log density of a multivariate normal, evaluated row-wise on X."""
    n, d = X.shape
    L = np.linalg.cholesky(cov)  # cov = L @ L.T
    diff = X - mean  # (n, d)
    sol = np.linalg.solve(L, diff.T)  # (d, n), solves L @ w = diff_i
    maha = np.sum(sol**2, axis=0)  # squared Mahalanobis distance
    log_det = 2.0 * np.sum(
        np.log(np.diag(L))
    )  # log|cov| from the Cholesky factor
    return -0.5 * (d * np.log(2.0 * np.pi) + log_det + maha)


def log_component_density(X, weights, means, covs, reg_covar):
    """log(pi_k) + log N(x_i | mu_k, Sigma_k) for every point and component."""
    n, d = X.shape
    K = means.shape[0]
    out = np.empty((n, K))
    for k in range(K):
        cov = covs[k] + reg_covar * np.eye(d)  # regularize before factorizing
        out[:, k] = np.log(weights[k]) + log_normal_pdf(X, means[k], cov)
    return out


def gmm_e_step(X, weights, means, covs, reg_covar):
    """Return responsibilities and the total log-likelihood."""
    log_prob = log_component_density(X, weights, means, covs, reg_covar)
    log_norm = logsumexp(
        log_prob, axis=1
    )  # log of the mixture density per point
    if not np.isfinite(log_norm).all():
        raise FloatingPointError(
            "Non-finite log density; inspect scale, weights and covariances."
        )
    resp = np.exp(log_prob - log_norm[:, None])
    return resp, float(log_norm.sum())


def gmm_m_step(X, resp):
    """Weighted mean, weighted scatter, and weight updates."""
    n, d = X.shape
    K = resp.shape[1]
    Nk = resp.sum(axis=0)  # effective counts
    if np.any(Nk <= np.finfo(float).tiny):
        raise FloatingPointError(
            "A component lost its effective weight; reinitialize or reduce K."
        )
    means = (resp.T @ X) / Nk[:, None]  # weighted averages
    covs = np.empty((K, d, d))
    for k in range(K):
        diff = X - means[k]
        covs[k] = (
            (resp[:, [k]] * diff).T @ diff / Nk[k]
        )  # weighted outer products
    weights = Nk / n
    return weights, means, covs

The log_normal_pdf function is worth reading line by line. Factorizing $\boldsymbol{\Sigma} = \mathbf{L}\mathbf{L}^{\top}$ avoids explicitly computing its inverse, and it gives both quantities we need at once: $\log\lvert\boldsymbol{\Sigma}\rvert = 2\sum_i \log L_{ii}$ , and the squared Mahalanobis distance as $\lVert \mathbf{L}^{-1}(\mathbf{x} - \boldsymbol{\mu}) \rVert^2$ , which we obtain by solving the triangular system $\mathbf{L}\mathbf{w} = \mathbf{x} - \boldsymbol{\mu}$ for every point at once. The transpose inside np.linalg.solve exists because solve expects the right-hand side as columns; each column of sol is one point's whitened residual.

PythonIn [8]
CollapseExpanddef gmm_init(X, K, rng, method="kmeans", reg_covar=1e-6):
    """Initial parameters from k-means, or from random observations."""
    n, d = X.shape
    if method == "kmeans":
        km = KMeans(
            n_clusters=K, n_init=10, random_state=int(rng.integers(1_000_000))
        )
        labels = km.fit_predict(X)
        means = km.cluster_centers_.copy()
    else:
        means = X[rng.choice(n, size=K, replace=False)].copy()
        dist = ((X[:, None, :] - means[None, :, :]) ** 2).sum(axis=2)
        labels = dist.argmin(axis=1)

    counts = np.bincount(labels, minlength=K).astype(float)
    weights = (
        np.maximum(counts, 1.0) / n
    )  # floor avoids log(0) for empty clusters
    weights = weights / weights.sum()

    global_cov = np.atleast_2d(np.cov(X.T, bias=True)) + reg_covar * np.eye(d)
    covs = np.empty((K, d, d))
    for k in range(K):
        mask = labels == k
        if mask.sum() > d:
            covs[k] = np.atleast_2d(np.cov(X[mask].T, bias=True))
        else:
            covs[k] = global_cov.copy()
    return weights, means, covs


def fit_gmm(X, K, init="kmeans", n_iter=300, tol=1e-6, reg_covar=1e-6, seed=0):
    """Alternate E and M steps until the log-likelihood stops improving."""
    X = np.asarray(X, dtype=float)
    if X.ndim != 2 or not X.shape[1] or not np.isfinite(X).all():
        raise ValueError("X must be a finite two-dimensional feature matrix.")
    if (
        not isinstance(K, (int, np.integer))
        or isinstance(K, bool)
        or not 1 <= K <= len(X)
    ):
        raise ValueError(
            "K must be an integer between one and the number of rows."
        )
    if (
        not isinstance(n_iter, (int, np.integer))
        or isinstance(n_iter, bool)
        or n_iter < 1
    ):
        raise ValueError("n_iter must be a positive integer.")
    if (
        init not in {"kmeans", "random"}
        or not np.isfinite([tol, reg_covar]).all()
        or tol <= 0
        or reg_covar <= 0
    ):
        raise ValueError(
            "Use a supported initialization and finite positive tolerance and covariance floor."
        )
    rng = np.random.default_rng(seed)
    weights, means, covs = gmm_init(X, K, rng, method=init, reg_covar=reg_covar)
    history, prev_ll, converged = [], -np.inf, False
    for _ in range(n_iter):
        resp, ll = gmm_e_step(X, weights, means, covs, reg_covar)
        history.append(ll)
        if np.isfinite(prev_ll):
            gain = ll - prev_ll
            if gain < -1e-8 * (1.0 + abs(prev_ll)):
                raise FloatingPointError(
                    "Log-likelihood decreased materially; inspect regularization and initialization."
                )
            if abs(gain) < tol:
                converged = True
                break
        prev_ll = ll
        weights, means, covs = gmm_m_step(X, resp)
    if not converged:  # the loop ran out of iterations
        resp, ll = gmm_e_step(X, weights, means, covs, reg_covar)
        if ll - prev_ll < -1e-8 * (1.0 + abs(prev_ll)):
            raise FloatingPointError(
                "Log-likelihood decreased materially; inspect regularization and initialization."
            )
        history.append(ll)
    return {
        "weights": weights,
        "means": means,
        "covs": covs,
        "resp": resp,
        "log_likelihood": ll,
        "history": np.array(history),
        "n_iter": len(history),
        "converged": converged,
        "reg_covar": reg_covar,
    }

The convergence test uses the absolute change in total log-likelihood, not a relative change and not distance to an optimum. A material decrease raises an error instead of being labelled convergence; changes within the stated roundoff allowance can still be small. The diagonal-addition update is a practical safeguard, not the exact unconstrained M-step proved earlier. A small change or an exhausted iteration budget does not certify a stationary point or a global optimum.

Now we generate two Gaussian components with different spreads and orientations. Their overlap makes component membership uncertain from the observed features alone, so we can compare hard assignments with posterior responsibilities.

PythonIn [9]
CollapseExpandrng = np.random.default_rng(7)


def rotated_cov(sd, angle_deg):
    """Covariance matrix with given axis standard deviations, rotated by angle."""
    t = np.deg2rad(angle_deg)
    R = np.array([[np.cos(t), -np.sin(t)], [np.sin(t), np.cos(t)]])
    return R @ np.diag(np.square(sd)) @ R.T


true_means = np.array([[1.6, 1.2], [4.1, 2.4]])
true_covs = np.stack(
    [rotated_cov([1.55, 0.55], 35.0), rotated_cov([1.25, 0.85], -25.0)]
)
true_weights = np.array([0.55, 0.45])

n_obs = 400
z_true = rng.choice(2, size=n_obs, p=true_weights)
X = np.empty((n_obs, 2))
for i in range(n_obs):
    X[i] = rng.multivariate_normal(true_means[z_true[i]], true_covs[z_true[i]])

x_lo, x_hi = X[:, 0].min() - 0.6, X[:, 0].max() + 0.6
y_lo, y_hi = X[:, 1].min() - 0.6, X[:, 1].max() + 0.6

PythonIn [10]
CollapseExpanddef fit_gmm_best(X, K, n_init=5, **kwargs):
    """Run EM from several starting points and keep the best log-likelihood."""
    if (
        not isinstance(n_init, (int, np.integer))
        or isinstance(n_init, bool)
        or n_init < 1
    ):
        raise ValueError("n_init must be a positive integer.")
    runs = [fit_gmm(X, K, seed=s, **kwargs) for s in range(n_init)]
    return max(runs, key=lambda m: m["log_likelihood"]), runs


model, runs = fit_gmm_best(X, 2, n_init=5, reg_covar=1e-6)

ConsoleOut [11]
CollapseExpandconverged: True, iterations: 46
log-likelihood: -1231.566
weights: [0.535 0.465]
means:
[[1.865 1.455]
 [4.224 2.276]]
covariances:
[[[ 2.174  1.161]
  [ 1.161  0.931]]

 [[ 1.227 -0.348]
  [-0.348  0.784]]]

The recovered means sit close to the generating values, and the covariance matrices show the differing orientations we built in. Sort the components by their first coordinate before printing if you want to compare runs directly, since the fitting order is arbitrary.

To summarize ambiguity under the fitted model, we can count the observations whose responsibility falls in the middle of the range. K-means still assigns these points by their nearest centroid; it does not return a posterior ambiguity measure.

PythonIn [12]
CollapseExpandamb = ((model["resp"][:, 0] > 0.3) & (model["resp"][:, 0] < 0.7)).sum()
km = KMeans(n_clusters=2, n_init=10, random_state=0).fit(X)
ours = model["resp"].argmax(axis=1)
agreement = max(np.mean(ours == km.labels_), np.mean(ours != km.labels_))

ConsoleOut [13]
CollapseExpandobservations with responsibility between 0.3 and 0.7: 91 of 400
hard-label agreement between GMM and k-means: 0.863

The printed agreement summarizes the two hard partitions after allowing their binary labels to swap. It does not measure accuracy against the generating labels, and agreement alone does not locate the disagreements; the plots let us inspect their geometry.

A reference implementation is the natural cross-check. Scikit-learn's GaussianMixture also uses posterior responsibilities and diagonal covariance regularization. We compare log densities on the same coordinates, but initialization, stopping conventions and covariance bookkeeping can differ; close results are a useful check rather than a guarantee.

PythonIn [14]
CollapseExpandsk_model = GaussianMixture(
    n_components=2,
    covariance_type="full",
    reg_covar=1e-6,
    n_init=5,
    max_iter=300,
    tol=1e-9,
    random_state=0,
).fit(X)
sk_ll = sk_model.score(X) * len(X)
sk_labels = sk_model.predict(X)
sk_agree = max(np.mean(ours == sk_labels), np.mean(ours != sk_labels))

ConsoleOut [15]
CollapseExpandour log-likelihood:      -1231.566
sklearn log-likelihood:  -1231.566
assignment agreement:    1.000
sklearn BIC:             2529.039

Small differences are expected: the two implementations initialize differently, stop at slightly different points, and handle the regularization term in slightly different places. A large gap could reflect different local fits, stopping behaviour, regularization, or an implementation error, which is a reminder to use several restarts whichever library is in play.

Stable implementation and initializationLink Copied

The formulas above are mathematically correct and numerically fragile if taken literally. Two things must be changed before they will survive real data.

The first is the responsibility calculation. Forming $\pi_k \mathcal{N}(\mathbf{x}_i \mid \boldsymbol{\mu}_k, \boldsymbol{\Sigma}_k)$ directly and then dividing produces underflow for points far from every component: an exponential of $-800$ is exactly zero in double precision, and dividing zero by zero yields nan. The fix is to compute everything in log space and use the log-sum-exp identity,

\begin{aligned} \log \sum_{j=1}^{K} \exp(a_j) &= \log\!\left( e^{m} \sum_{j=1}^{K} \exp(a_j - m) \right) && \text{(collect the common factor)} \\ &= m + \log \sum_{j=1}^{K} \exp(a_j - m) && \text{(use } \log(e^m \cdot s) = m + \log s \text{)} \end{aligned}

where $a_j$ are finite real log-joint values and $m = \max_j a_j$ .

Factoring out the largest term is a legal algebraic rearrangement: we divide each term by $e^m$ and multiply the total back by $e^m$ . It guarantees that the largest exponent is zero, so no term overflows and at least one term is exactly $1$ . The responsibility then becomes $\exp(a_k - \text{logsumexp}(\mathbf{a}))$ , which is always between zero and one even when the raw densities are unrepresentable. SciPy's logsumexp supplies this stable finite-input calculation. An all-infinite or nan input still needs explicit handling; the identity does not make invalid model parameters meaningful.

The second issue is covariance singularity. A weighted scatter is positive definite when its positive-weight observations span all $d$ feature directions after centring. For a hard-assigned component that requires at least $d+1$ affinely independent points. The responsibility sum $N_k$ is not a rank test: many tiny positive weights can still give a full-rank scatter. Adding $\epsilon\mathbf{I}$ with $\epsilon>0$ bounds the eigenvalues away from zero and makes factorization possible. It also bounds the density of these regularized fits, although adding it after an unconstrained M-step does not guarantee ordinary-likelihood monotonicity.

Initialization deserves as much care as the numerics. Exact bound-increasing updates guarantee likelihood nondecrease, not global optimality. Choose starting values and a restart budget together:

Random means chosen from the data are simple; some runs land in poor fits.
K-means partitioning often gives a useful structured start. Fit k-means, use its centroids as the means, its labels to compute initial covariances, and its cluster proportions as initial weights. It can still miss an elongated or overlapping component.
Multiple restarts with the best final likelihood retained are the pragmatic default; n_init controls the number of starts in scikit-learn.

A useful diagnostic is to track the log-likelihood of every restart: if converged runs have substantially different final values, initialization matters. If some runs have not converged, more iterations may explain part of the spread.

Two further safeguards are worth knowing. Component weights can collapse toward zero when a component becomes responsible for almost nothing, which can reflect a redundant component or a poor initialization but produces $\log \pi_k \to -\infty$ if it reaches exactly zero; flooring the weights at a small positive value keeps the arithmetic finite. And the convergence threshold should be applied to the log-likelihood, not to the parameters, since parameters can drift while the likelihood is flat along a ridge.

Key ParametersLink Copied

The parameters of a Gaussian mixture fit control capacity, shape, and computational effort, and each has a practical consequence.

n_components ( $K$ ): the number of mixture components. The best attainable fit in nested model families cannot worsen as components are added, but finite EM runs may return worse local fits. Training likelihood alone does not penalize complexity.
covariance_type: "full" gives each component its own unrestricted matrix ( $K d(d+1)/2$ covariance parameters), "tied" shares one matrix across components, "diag" restricts each component to axis-aligned spreads, and "spherical" gives each a single variance. Constraining covariance is the main defence against overfitting when the data are limited, and a shared spherical variance with equal weights has the nearest-centroid hard-assignment rule. Finite-variance soft EM remains a different procedure.
reg_covar: the $\epsilon$ added to the diagonal. Scikit-learn defaults to $10^{-6}$ , the starting value used in this example. Inspect that setting relative to feature scales and degeneracy warnings; a larger value may be needed for an unstable fit.
init_params and n_init: how the starting point is chosen and how many restarts are attempted. K-means initialization with several restarts is the sensible default.
max_iter and tol: the iteration cap and the stopping threshold on the log-likelihood. Scikit-learn applies tol to the average lower-bound gain, whereas our teaching implementation checks total log-likelihood change. Inspect the convergence flag and increase the budget if it is exhausted.
weights_init, means_init, precisions_init: optional user-supplied starting values, useful when there is prior knowledge about where the components should be, or when a specific initialization must be reproduced.

Key MethodsLink Copied

Once fitted, a GaussianMixture object exposes the different ways of using the model, and the distinction between them matters.

fit(X) estimates the parameters; fit_predict(X) additionally returns the hard labels.
predict(X) returns the component index with the highest responsibility, which is the hard clustering. It discards the uncertainty that the model went to the trouble of computing.
predict_proba(X) returns the full responsibility matrix, one row per observation and one column per component. This is usually the more useful output.
score_samples(X) returns the log mixture density at each observation, which is the natural anomaly score.
score(X) returns the average log density per observation, the quantity we compared against our own implementation after multiplying by $n$ .
bic(X) and aic(X) return penalized likelihood criteria that can be compared across different values of n_components, since a lower value indicates a better trade-off between fit and complexity.
sample(n_samples) draws new observations from the fitted density, which is the generative use of the model.

The distinction between predict and predict_proba is central to using the model in practice: the first produces hard GMM labels, which can differ from k-means labels, while the second gives the information that made the mixture model worth fitting.

Comments

Add to the conversation

A question, a correction, or an example from your own work.

No comments yet. Be the first to share your thoughts!

Stable implementation and initializationLink Copied

Key ParametersLink Copied

Key MethodsLink Copied

Comments

Add to the conversation

Stay up to date