Implement Gaussian Mixture Model (GMM) M-step

Problem

Implement the M-step (Maximization step) of the Gaussian Mixture Model (GMM) EM algorithm. Given data and the responsibilities from the E-step, update the mixing coefficients, means, and covariance matrices.

Solution

import numpy as np

def gmm_m_step(
    X: np.ndarray,
    resp: np.ndarray
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    n, d = X.shape
    k = resp.shape[1]

    # Effective number of points per component
    Nk = resp.sum(axis=0)

    # Update mixing coefficients
    pi = Nk / n

    # Update means
    mu = np.zeros((k, d))
    for j in range(k):
        mu[j] = (resp[:, j:j+1].T @ X) / (Nk[j] + 1e-10)

    # Update covariances
    sigma = np.zeros((k, d, d))
    for j in range(k):
        diff = X - mu[j]
        sigma[j] = (diff.T @ (diff * resp[:, j:j+1])) / (Nk[j] + 1e-10)
        sigma[j] += np.eye(d) * 1e-6  # regularization

    return pi, mu, sigma

Explanation

1. Compute the effective count N_k for each component by summing responsibilities.
2. Update mixing coefficients: pi_k = N_k / n.
3. Update means: mu_k = (1/N_k) * sum(r_{ik} * x_i) — the responsibility-weighted average of all points.
4. Update covariances: sigma_k = (1/N_k) * sum(r_{ik} * (x_i - mu_k)(x_i - mu_k)^T) with a small regularization term for numerical stability.

Complexity

• Time: O(n k d^2) for covariance updates
• Space: O(k * d^2) for covariance matrices