Gaussian Mixture Models & EM Algorithm | AI & Machine Learning Engineering

While k-Means forces hard binary assignments ( $r_{i,k} \in \{0, 1\}$ ) and assumes isotropic spherical clusters, Gaussian Mixture Models (GMMs) provide a soft, probabilistic framework.

Each cluster is represented by a multivariate Gaussian density characterized by its own mean $\mu_k$ and covariance matrix $\Sigma_k$ , allowing clusters to capture arbitrary elliptical orientations and variances.

Mixing ProportionsProbability Theory

N_k / N

Optimal Gaussian mixture weights equal the normalized sum of sample responsibilities.

1. Probability Density Formulation#

A Gaussian Mixture Model represents the probability distribution of continuous feature vector $x \in \mathbb{R}^d$ as a convex combination of $K$ Gaussian components:

p(x | \theta) = \sum_{k=1}^K \pi_k \mathcal{N}(x | \mu_k, \Sigma_k)

Architecture Diagram

Rendering diagram...

Where:

$\pi_k \ge 0$ is the mixture weight (prior probability) of component $k$ , subject to $\sum_{k=1}^K \pi_k = 1$ .
$\mathcal{N}(x | \mu_k, \Sigma_k)$ is the multivariate Gaussian probability density function:

\mathcal{N}(x | \mu_k, \Sigma_k) = \frac{1}{(2\pi)^{d/2} |\Sigma_k|^{1/2}} \exp\left( -\frac{1}{2}(x - \mu_k)^T \Sigma_k^{-1} (x - \mu_k) \right)

The Incomplete-Data Log-Likelihood Dilemma#

For $N$ independent observations $X = \{x_1, \dots, x_N\}$ , the log-likelihood function is:

\ln p(X | \theta) = \sum_{i=1}^N \ln \left( \sum_{k=1}^K \pi_k \mathcal{N}(x_i | \mu_k, \Sigma_k) \right)

Notice that the summation over $K$ components appears inside the natural logarithm. Because $\ln(a + b) \ne \ln(a) + \ln(b)$ , setting the gradient of $\ln p(X|\theta)$ to zero yields no closed-form analytical solution.

2. Derivation of the Expectation-Maximization (EM) Algorithm#

To circumvent the log-sum intractability, Dempster, Laird, and Rubin introduced the Expectation-Maximization (EM) algorithm by positing unobserved latent categorical variables $z_i \in \{0, 1\}^K$ , where $z_{i,k} = 1$ indicates that sample $i$ was generated by component $k$ .

The complete-data log-likelihood (if latent assignments $Z$ were known) is:

\ln p(X, Z | \theta) = \sum_{i=1}^N \sum_{k=1}^K z_{i,k} \Big[ \ln(\pi_k) + \ln \mathcal{N}(x_i | \mu_k, \Sigma_k) \Big]

Because $Z$ is unobserved, we optimize the expected complete log-likelihood with respect to the posterior distribution of $Z$ .

Step 1: The Expectation Step (E-Step)

Compute the posterior probability (responsibility) that Gaussian component $k$ generated observation $x_i$ , given current parameter estimates $\theta^{(t)}$ :

\gamma_{i,k} \equiv \mathbb{E}[z_{i,k} | x_i, \theta^{(t)}] = P(z_{i,k} = 1 | x_i, \theta^{(t)})

Applying Bayes' Theorem:

\gamma_{i,k} = \frac{P(z_{i,k} = 1) \cdot p(x_i | z_{i,k} = 1)}{\sum_{j=1}^K P(z_{i,j} = 1) \cdot p(x_i | z_{i,j} = 1)} = \frac{\pi_k \mathcal{N}(x_i | \mu_k, \Sigma_k)}{\sum_{j=1}^K \pi_j \mathcal{N}(x_i | \mu_j, \Sigma_j)}

Step 2: Formulate the Auxiliary Q-Function

The expected value of the complete-data log-likelihood is:

Q(\theta, \theta^{(t)}) = \sum_{i=1}^N \sum_{k=1}^K \gamma_{i,k} \Big[ \ln(\pi_k) + \ln \mathcal{N}(x_i | \mu_k, \Sigma_k) \Big]

Step 3: The Maximization Step (M-Step) for Means $\mu_k$

Differentiate $Q$ with respect to vector $\mu_k$ :

\frac{\partial Q}{\partial \mu_k} = \sum_{i=1}^N \gamma_{i,k} \cdot \Sigma_k^{-1} (x_i - \mu_k) = 0

Multiply by $\Sigma_k$ and isolate $\mu_k$ :

\sum_{i=1}^N \gamma_{i,k} x_i - \left( \sum_{i=1}^N \gamma_{i,k} \right) \mu_k = 0

Let $N_k = \sum_{i=1}^N \gamma_{i,k}$ denote the effective number of data points assigned to cluster $k$ :

\mu_k^* = \frac{1}{N_k} \sum_{i=1}^N \gamma_{i,k} x_i

Step 4: The Maximization Step for Covariance $\Sigma_k$

Similarly, setting the matrix derivative with respect to $\Sigma_k^{-1}$ to zero:

\Sigma_k^* = \frac{1}{N_k} \sum_{i=1}^N \gamma_{i,k} (x_i - \mu_k^*)(x_i - \mu_k^*)^T

Step 5: The Maximization Step for Mixture Weights $\pi_k$

To optimize $\pi_k$ subject to $\sum_{k=1}^K \pi_k = 1$ , we introduce Lagrange multiplier $\lambda$ :

\mathcal{L}(\pi, \lambda) = \sum_{i=1}^N \sum_{k=1}^K \gamma_{i,k} \ln(\pi_k) + \lambda \left( 1 - \sum_{k=1}^K \pi_k \right)

Differentiating with respect to $\pi_k$ :

\frac{\partial \mathcal{L}}{\partial \pi_k} = \sum_{i=1}^N \frac{\gamma_{i,k}}{\pi_k} - \lambda = 0 \implies \pi_k = \frac{1}{\lambda} \sum_{i=1}^N \gamma_{i,k} = \frac{N_k}{\lambda}

Summing over all $K$ components:

\sum_{k=1}^K \pi_k = \frac{1}{\lambda} \sum_{k=1}^K N_k \implies 1 = \frac{N}{\lambda} \implies \lambda = N

Therefore:

\pi_k^* = \frac{N_k}{N}

3. From-Scratch Implementation: Zero Dependencies#

The following implementations construct an Expectation-Maximization Gaussian Mixture Model (with diagonal covariance matrices) from scratch in pure Python and pure Rust.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
import math
import random
from typing import List, Tuple

class GaussianMixtureFromScratch:
    def __init__(self, k: int = 2, max_iter: int = 100, tol: float = 1e-4):
        self.k = k
        self.max_iter = max_iter
        self.tol = tol
        self.weights: List[float] = []
        self.means: List[List[float]] = []
        self.variances: List[List[float]] = []

    @staticmethod
    def _gaussian_pdf(x: List[float], mean: List[float], var: List[float]) -> float:
        # Diagonal covariance multivariate Gaussian
        prob = 1.0
        for x_j, mu_j, v_j in zip(x, mean, var):
            v_safe = max(v_j, 1e-6)
            coeff = 1.0 / math.sqrt(2.0 * math.pi * v_safe)
            exp_term = math.exp(-((x_j - mu_j) ** 2) / (2.0 * v_safe))
            prob *= (coeff * exp_term)
        return max(prob, 1e-12)

    def fit(self, X: List[List[float]]) -> None:
        num_samples = len(X)
        num_features = len(X[0])
        random.seed(42)

        # 1. Initialize weights, means, and variances
        self.weights = [1.0 / self.k] * self.k
        self.means = [list(sample) for sample in random.sample(X, self.k)]
        
        # Initialize variances to global feature variance
        global_var = [
            sum((row[j] - sum(r[j] for r in X) / num_samples) ** 2 for row in X) / num_samples
            for j in range(num_features)
        ]
        self.variances = [[max(v, 0.1) for v in global_var] for _ in range(self.k)]

        prev_ll = -float("inf")

        for _ in range(self.max_iter):
            # E-Step: calculate responsibilities gamma[i][k]
            gamma: List[List[float]] = []
            log_likelihood = 0.0

            for i in range(num_samples):
                responsibilities_unnorm = [
                    self.weights[k] * self._gaussian_pdf(X[i], self.means[k], self.variances[k])
                    for k in range(self.k)
                ]
                total = sum(responsibilities_unnorm)
                log_likelihood += math.log(max(total, 1e-12))
                gamma.append([val / total for val in responsibilities_unnorm])

            # Check log-likelihood convergence
            if abs(log_likelihood - prev_ll) < self.tol:
                break
            prev_ll = log_likelihood

            # M-Step: update parameters
            for k in range(self.k):
                N_k = sum(gamma[i][k] for i in range(num_samples))
                self.weights[k] = N_k / num_samples

                # Update mean
                for j in range(num_features):
                    self.means[k][j] = sum(gamma[i][k] * X[i][j] for i in range(num_samples)) / N_k

                # Update diagonal variance
                for j in range(num_features):
                    self.variances[k][j] = sum(
                        gamma[i][k] * ((X[i][j] - self.means[k][j]) ** 2)
                        for i in range(num_samples)
                    ) / N_k

    def predict_proba(self, X: List[List[float]]) -> List[List[float]]:
        probs = []
        for row in X:
            unnorm = [
                self.weights[k] * self._gaussian_pdf(row, self.means[k], self.variances[k])
                for k in range(self.k)
            ]
            total = sum(unnorm)
            probs.append([u / total for u in unnorm])
        return probs

if __name__ == "__main__":
    X = [
        [0.8, 1.0], [1.2, 0.9], [1.0, 1.1],
        [5.0, 5.2], [5.2, 4.9], [4.8, 5.0]
    ]

    gmm = GaussianMixtureFromScratch(k=2)
    gmm.fit(X)
    print("Learned Means:", gmm.means)
    print("Cluster Probabilities for first sample:", gmm.predict_proba([X[0]]))

4. Active Knowledge Checks#

Concept check

Why can the incomplete-data log-likelihood of a Gaussian Mixture Model not be optimized directly using analytical derivatives?

Concept check

In the Expectation-Maximization algorithm, how does the E-step calculate the responsibility gamma_ik?