Curso de ML EN

Capítulo 32 de 37 · avanzado

Modelos de mezclas gaussianas

Qué cubre este capítulo

k-means le daba a cada punto una etiqueta dura y dibujaba cada cluster como una bola redonda. Eso funcionó con cuatro blobs bien portados, y se cae en cuanto los grupos están estirados, inclinados o traslapados — o sea, casi siempre. Este capítulo se queda con la buena idea de k-means, que un cluster tiene un centro, y arregla las dos cosas que lo perjudican. Un cluster ahora puede ser una elipse inclinada y estirada en lugar de un círculo, y un punto puede pertenecer en parte a varios en vez de estar obligado a escoger uno.

Ese modelo es la mezcla gaussiana, y lo ajustamos con expectation-maximization — el mismo ciclo de asignar suave y luego reajustar del capítulo de EM, ahora en dos dimensiones y con una matriz de covarianza completa por componente para que cada cluster cargue su propia forma y su propio ángulo. Construimos todo a mano con NumPy: la gaussiana multivariada, el paso E que reparte responsabilidades suaves, el paso M que reajusta cada media, covarianza y peso, y la log-verosimilitud que el ciclo va escalando. Después dejamos que GaussianMixture de scikit-learn ajuste el mismo modelo y verificamos que nuestro número caiga sobre el suyo.

La estrella es una animación de las elipses moldeándose a los datos — tres blobs redondos que se despegan y se inclinan hasta acomodarse, puntos que se van degradando de un color a otro donde los clusters se traslapan, y un panel de log-verosimilitud subiendo abajo. Y el premio es un comparativo lado a lado que ya no vas a poder ignorar: los mismos datos agrupados por k-means duro y esférico y por GMM suave y elíptico, uno de ellos cortando de tajo por en medio de grupos que el otro envuelve limpiamente.

Un poco de historia

Las piezas son viejas. En 1894 Karl Pearson, trabajando con un conjunto de medidas de cangrejos que le salía ladeado en vez de acampanado, decidió que la muestra en realidad eran dos especies mezcladas y se puso a recuperar ambas campanas a partir de la mezcla. Lo hizo con el método de los momentos — igualando cinco momentos muestrales a los parámetros de una mezcla gaussiana de dos componentes, lo que lo dejó resolviendo a mano un polinomio de noveno grado. Es de las primeras veces que alguien ajusta un modelo de mezcla a datos reales, y lo hizo sin computadora, sin máxima verosimilitud y sin nada de la maquinaria que después volvería esto rutinario.

La maquinaria llegó en 1977, cuando Arthur Dempster, Nan Laird y Donald Rubin le pusieron nombre al algoritmo de expectation-maximization y mostraron que toda una familia enorme de problemas del tipo "los datos que tengo son una versión incompleta de datos que podría ajustar fácil" se reducía al mismo ciclo. Una mezcla gaussiana es el caso de libro de texto: si supieras qué componente generó cada punto, ajustar las gaussianas sería un promedio ponderado; si supieras las gaussianas, etiquetar los puntos sería una aplicación de la regla de Bayes. No sabes ninguna de las dos, así que EM alterna entre ellas. Los cangrejos de Pearson por fin tuvieron un método general, y la mezcla de gaussianas se volvió el primer ejemplo estándar de EM en todos los cursos desde entonces — incluido el anterior.

La intuición

Empecemos por lo que k-means no puede hacer. Mide todo con distancia simple a un centro, así que su idea de cluster es una bola: igual de ancha en todas direcciones, del mismo tamaño en todos lados, con fronteras que son líneas rectas a la mitad entre centros. Dale un cluster con forma de franja diagonal larga y no tiene manera de expresarlo. Planta un centro en medio y declara outliers a los extremos, o peor, se los regala a un vecino cuyo centro casualmente queda más cerca.

Una mezcla gaussiana reemplaza la bola por una gaussiana completa. Una gaussiana en dos dimensiones tiene una media, que es el centro, y una matriz de covarianza, que es la forma: qué tan ancha, qué tan alta y, sobre todo, qué tan inclinada. Esa covarianza es toda la mejora. Una sola varianza te devolvería un círculo; una covarianza diagonal te da una elipse alineada a los ejes; una covarianza completa, con el término fuera de la diagonal libre, te da una elipse rotada a cualquier ángulo. Cada cluster se envuelve en la elipse que de verdad le queda a sus puntos.

La segunda mejora es la suavidad. En lugar de asignar un punto a un solo cluster, el modelo pregunta: dadas las elipses actuales, ¿cuál es la probabilidad de que este punto haya salido de cada una? Un punto bien adentro de una elipse recibe una respuesta casi segura. Un punto que cae en el traslape entre dos recibe un reparto — 60% de esta, 40% de aquella — y el modelo conserva esa ambigüedad en vez de fingir que no existe. Esas probabilidades se llaman responsabilidades, y son a la vez la forma en que el modelo se ajusta y lo que te regresa.

Estos son los datos con los que vamos a trabajar. Tres blobs, pero nada redondos: pasaron por un shear que los estira e inclina hasta volverlos elipses diagonales largas, recargadas en la misma dirección y traslapadas en las costuras. Sin etiquetas, sin colores — nada más 300 puntos en un plano.

Tu ojo sigue encontrando tres grupos, pero fíjate que son diagonales y se tocan. Esos son los datos con los que k-means fue hecho para fallar y con los que la mezcla fue hecha para responder.

Las matemáticas

Una mezcla dice que los datos vienen de KK gaussianas combinadas. Cada componente kk tiene un peso de mezcla πk\pi_k (su porción de los datos, con los pesos sumando uno), un vector de medias μk\mu_k y una matriz de covarianza Σk\Sigma_k. La probabilidad que el modelo le asigna a un punto xx es la suma ponderada de lo que cada componente opina de él:

p(x)=∑k=1Kπk N(x∣μk,Σk)p(x) = \sum_{k=1}^{K} \pi_k \, \mathcal{N}(x \mid \mu_k, \Sigma_k)

El término de forma es la densidad gaussiana multivariada. Para un punto de DD dimensiones es:

N(x∣μ,Σ)=1(2π)D/2 ∣Σ∣1/2exp⁡ ⁣(−12(x−μ)⊤Σ−1(x−μ))\mathcal{N}(x \mid \mu, \Sigma) = \frac{1}{(2\pi)^{D/2}\,|\Sigma|^{1/2}} \exp\!\left(-\tfrac{1}{2}(x-\mu)^{\top}\Sigma^{-1}(x-\mu)\right)

La cantidad (x−μ)⊤Σ−1(x−μ)(x-\mu)^{\top}\Sigma^{-1}(x-\mu) es la distancia de Mahalanobis al cuadrado — la distancia cuadrada de siempre, deformada por la inversa de la covarianza, de modo que "lejos" se mide en unidades de la propia dispersión y dirección de la elipse. Cuando Σ\Sigma es la identidad esto colapsa a la distancia euclidiana y las curvas de nivel son círculos, que es exactamente la visión de k-means. Una Σ\Sigma completa es lo que inclina y estira esos círculos hasta volverlos elipses.

Ajustar significa elegir todos los πk,μk,Σk\pi_k, \mu_k, \Sigma_k para hacer los datos lo más probables posible. EM lo hace en dos pasos. El paso E calcula la responsabilidad del componente kk sobre el punto ii — la probabilidad posterior de que haya venido de ese componente, que no es más que su densidad ponderada normalizada entre componentes:

γik=πk N(xi∣μk,Σk)∑j=1Kπj N(xi∣μj,Σj)\gamma_{ik} = \frac{\pi_k \, \mathcal{N}(x_i \mid \mu_k, \Sigma_k)}{\sum_{j=1}^{K} \pi_j \, \mathcal{N}(x_i \mid \mu_j, \Sigma_j)}

Cada renglón γi⋅\gamma_{i\cdot} suma uno; esa es la asignación suave. El paso M después reajusta cada componente como una versión ponderada por responsabilidades de los estimadores muestrales de siempre. Sea Nk=∑iγikN_k = \sum_i \gamma_{ik} el conteo suave de puntos en el componente kk. El nuevo peso es su porción:

πk=NkN\pi_k = \frac{N_k}{N}

La nueva media es el promedio ponderado de los puntos:

μk=1Nk∑i=1Nγik xi\mu_k = \frac{1}{N_k} \sum_{i=1}^{N} \gamma_{ik}\, x_i

Y la nueva covarianza es la dispersión ponderada alrededor de esa media — la ecuación que le permite a cada elipse tomar su propia forma e inclinación:

Σk=1Nk∑i=1Nγik (xi−μk)(xi−μk)⊤\Sigma_k = \frac{1}{N_k} \sum_{i=1}^{N} \gamma_{ik}\, (x_i - \mu_k)(x_i - \mu_k)^{\top}

Cada punto contribuye a cada componente en proporción a qué tanto le pertenece. El número que el ciclo va escalando es la log-verosimilitud total — el logaritmo de la probabilidad de la mezcla, sumado sobre los puntos:

ℓ=∑i=1Nlog⁡∑k=1Kπk N(xi∣μk,Σk)\ell = \sum_{i=1}^{N} \log \sum_{k=1}^{K} \pi_k \, \mathcal{N}(x_i \mid \mu_k, \Sigma_k)

La garantía de EM es que ℓ\ell nunca baja de una iteración a la siguiente. Cuando deja de subir, paramos.

En qué es bueno y en qué no

La mezcla gana en todos los casos donde la suposición de bola redonda de k-means está mal. Las covarianzas completas le permiten ajustar clusters estirados e inclinados, de distintos tamaños y distintas orientaciones, y las responsabilidades suaves le permiten ser honesta con los puntos que caen en una frontera en vez de adivinar. Y como es un modelo de probabilidad de verdad, hace cosas que k-means no puede: te da una densidad contra la cual puedes evaluar puntos nuevos, te da una confianza por punto y te da una manera fundamentada de elegir el número de clusters. Suma la log-verosimilitud, penalízala por la cantidad de parámetros que gastaste y obtienes el criterio de información bayesiano — un score con un mínimo real, así que "cuántos clusters" deja de ser un juicio subjetivo sobre un codo y se convierte en un número que puedes optimizar.

El costo de esa expresividad es todo lo que esperarías de un modelo más grande. Una covarianza completa en DD dimensiones tiene D(D+1)/2D(D+1)/2 parámetros libres por componente, así que en dimensiones altas la mezcla tiene mucho que estimar y puede sufrir overfitting, sobre todo el componente que se queda con un puñado de puntos — su covarianza puede colapsar hacia un pico, mandando la verosimilitud al infinito en una solución degenerada. La protección típica es un ridge diminuto sumado a la diagonal de cada covarianza, que es justo lo que hacemos. EM también hereda el problema de los óptimos locales de k-means y con creces: la superficie de verosimilitud es accidentada, la respuesta depende de la inicialización y un mal arranque cae en una cuenca peor. El remedio es el mismo — siémbralo bien y reinícialo varias veces.

Los datos

Tres blobs gaussianos con make_blobs de scikit-learn, de tamaños desiguales (110, 100, 90 puntos), pasados después por un mismo mapeo lineal compartido — un shear que estira e inclina cada blob hasta volverlo una elipse diagonal larga. Esa transformación es deliberada: es el caso canónico donde k-means falla y una mezcla con covarianza completa acierta, porque los clusters quedan alargados en una dirección que deja el extremo lejano de cada uno más cerca del centro de un vecino que del propio. Guardamos los 300 puntos como data/blobs.csv junto con el id verdadero de blob de cada uno, que tratamos como un sobre sellado: la mezcla solo ve las coordenadas, y los ids salen al final nada más para calificar el agrupamiento. El generador es code/make_data.py, con semilla fija para que el snapshot se reproduzca idéntico.

Constrúyelo, una función a la vez

Ocho funciones. La primera es el término de forma que todas las demás llaman; luego los dos pasos de EM; luego el objetivo; luego la siembra y el ciclo; y al final las dos maneras de leer una respuesta de un modelo ya ajustado. Todo es NumPy puro.

Arrancamos con la densidad gaussiana multivariada. Dados los puntos, una media y una covarianza, queremos la densidad en todos los puntos de un jalón — sin loop de Python sobre los 300 renglones:

def gaussian(X, mean, cov):
    """Multivariate Gaussian density N(x | mean, cov) for every row of X.

    X is (N, D), mean is (D,), cov is (D, D); the result is (N,), the density
    at each point. This is the shape term of the whole model — a full DxD
    covariance is what lets the level sets be tilted ellipses instead of the
    circles a single variance would give you.
    """
    D = X.shape[1]
    diff = X - mean                                   # (N, D)
    inv = np.linalg.inv(cov)                           # (D, D)
    # Mahalanobis distance squared for every row, without a Python loop.
    maha = np.einsum("ni,ij,nj->n", diff, inv, diff)   # (N,)
    norm = 1.0 / (np.power(2 * np.pi, D / 2) * np.sqrt(np.linalg.det(cov)))
    return norm * np.exp(-0.5 * maha)

El einsum calcula la distancia de Mahalanobis de cada renglón en una sola llamada, y el inv y el det de la covarianza completa son lo que vuelve las curvas de nivel elipses inclinadas en vez de círculos. Esta única función es toda la diferencia con k-means; lo demás alrededor es papeleo.

El paso E apila esa densidad a lo largo de los componentes, pondera cada una por su peso de mezcla y normaliza cada renglón para que sume uno:

def e_step(X, weights, means, covs):
    """E-step: the responsibility of every component for every point.

    Returns gamma, an (N, K) matrix where gamma[i, k] is the posterior
    probability that point i was drawn from component k, given the current
    parameters. Each row is a soft assignment that sums to 1 — this is the
    'soft' in soft clustering, and the whole reason a GMM differs from k-means.
    """
    K = len(weights)
    weighted = np.stack([weights[k] * gaussian(X, means[k], covs[k])
                         for k in range(K)], axis=1)   # (N, K): pi_k * N(x|k)
    total = weighted.sum(axis=1, keepdims=True)         # (N, 1): the mixture density
    return weighted / np.maximum(total, 1e-300)

El resultado gamma es la matriz (N, K) de responsabilidades — la asignación suave. El piso de 1e-300 mantiene segura la división cuando un punto está tan lejos de todos los componentes que todas las densidades hacen underflow a cero. Con esas responsabilidades, el paso M reajusta cada parámetro como un estimador ponderado:

def m_step(X, gamma, reg=1e-6):
    """M-step: re-estimate every parameter as a responsibility-weighted fit.

    Given the responsibilities, each component's new mean, covariance, and
    mixing weight is just a weighted version of the ordinary sample estimate,
    where every point contributes in proportion to how much it belongs to that
    component. `reg` adds a tiny ridge to each covariance diagonal so a
    component that grabs only a few points can't collapse to a singular matrix.
    """
    N, D = X.shape
    K = gamma.shape[1]
    Nk = gamma.sum(axis=0)                              # (K,): soft count per component
    weights = Nk / N                                    # mixing weights
    means = (gamma.T @ X) / Nk[:, None]                 # (K, D): weighted means
    covs = np.empty((K, D, D))
    for k in range(K):
        diff = X - means[k]                             # (N, D)
        covs[k] = (gamma[:, k, None] * diff).T @ diff / Nk[k]
        covs[k] += reg * np.eye(D)                       # keep it non-singular
    return weights, means, covs

La media ponderada es una sola multiplicación de matrices; la covarianza ponderada es un loop corto sobre los KK componentes, cada uno una matriz de dispersión ponderada por responsabilidad, salida directo de las matemáticas. El término reg es el ridge que evita que un componente chico colapse a una covarianza singular. Ahora, el número que el ciclo escala:

def log_likelihood(X, weights, means, covs):
    """Total log-likelihood of the data under the current mixture.

    Sum over points of the log of the mixture density. EM can only push this
    up or leave it flat, never down — it is the number the whole loop climbs,
    and the flat spot is where we stop.
    """
    K = len(weights)
    weighted = np.stack([weights[k] * gaussian(X, means[k], covs[k])
                         for k in range(K)], axis=1)   # (N, K)
    return float(np.log(np.maximum(weighted.sum(axis=1), 1e-300)).sum())

La misma densidad de mezcla del paso E, pero en vez de normalizar tomamos el log de las sumas por renglón y los sumamos. Esto solo puede subir entre iteraciones; si alguna vez baja, hay un bug. Sigue la siembra, que aquí importa por la misma razón por la que importaba en k-means:

def init_params(X, K, rng):
    """Seed the mixture: k-means++ means, one shared covariance, equal weights.

    Spreading the initial means with k-means++ (the same trick from the k-means
    chapter) keeps EM out of the bad local optima a careless random start falls
    into. Every component starts as one big round copy of the data's overall
    covariance, so the animation opens with three identical blobs that then peel
    apart and mold to their own cluster's shape.
    """
    # k-means++ seeding of the means
    first = int(rng.integers(len(X)))
    means = [X[first]]
    for _ in range(1, K):
        d2 = np.min([((X - m) ** 2).sum(axis=1) for m in means], axis=0)
        means.append(X[int(rng.choice(len(X), p=d2 / d2.sum()))])
    means = np.array(means, dtype=float)

    shared = np.cov(X.T)                                 # one covariance for all
    covs = np.stack([shared.copy() for _ in range(K)])
    weights = np.full(K, 1.0 / K)                        # equal mixing weights
    return weights, means, covs

Tomamos prestado k-means++ para separar las medias iniciales, le damos a cada componente una copia de la covarianza global de los datos y arrancamos con pesos iguales. Por eso la animación abre con tres blobs redondos idénticos — todavía no encuentran su propia forma. El ciclo amarra todo: alternar E y M, guardando un snapshot en cada pasada, hasta que la log-verosimilitud se aplana:

def fit(X, K, rng, max_iter=100, tol=1e-4):
    """Run EM to convergence, recording a snapshot every iteration.

    Alternate the E-step and the M-step until the log-likelihood stops climbing
    by more than `tol`. Returns the final weights, means, covariances, and the
    full per-iteration history — each snapshot holds the parameters, the
    responsibilities, and the log-likelihood, which is exactly what the chapter
    replays frame by frame so you watch the ellipses form.
    """
    weights, means, covs = init_params(X, K, rng)
    gamma = e_step(X, weights, means, covs)
    ll = log_likelihood(X, weights, means, covs)
    history = [_snapshot(weights, means, covs, gamma, ll)]
    for _ in range(max_iter):
        weights, means, covs = m_step(X, gamma)         # re-estimate parameters
        gamma = e_step(X, weights, means, covs)         # recompute responsibilities
        new_ll = log_likelihood(X, weights, means, covs)
        history.append(_snapshot(weights, means, covs, gamma, new_ll))
        if new_ll - ll < tol:                           # likelihood stopped climbing
            ll = new_ll
            break
        ll = new_ll
    return weights, means, covs, ll, history

Cada snapshot guarda todos los parámetros, las responsabilidades y la log-verosimilitud — exactamente lo que la animación reproduce cuadro por cuadro. Por último, las dos maneras de leer el modelo ajustado. La etiqueta dura es el componente más responsable:

def predict(X, weights, means, covs):
    """Hard label: the single most responsible component for each point.

    argmax over the responsibilities collapses the soft assignment back to one
    label per point, so a GMM can stand in for k-means when you need a firm
    clustering — but you threw away the confidence to get it.
    """
    return np.argmax(e_step(X, weights, means, covs), axis=1)

Y la etiqueta suave es el vector completo de responsabilidades — eso que k-means nunca te va a poder dar:

def predict_proba(X, weights, means, covs):
    """Soft label: the full responsibility vector for each point.

    This is what a GMM has that k-means never will — a distribution over
    components per point, so you can see the points on a boundary that belong
    honestly to two clusters at once.
    """
    return e_step(X, weights, means, covs)

Míralo trabajar

Esta es toda la razón para quedarnos en dos dimensiones. Abajo hay una corrida real de la función fit de arriba, un cuadro por iteración de EM. Cada componente se dibuja como dos elipses inclinadas — los contornos de 1σ y 2σ, calculados a partir de una descomposición en eigenvalores real de la covarianza de ese componente, así que los ejes apuntan a lo largo de los eigenvectores verdaderos y sus longitudes son las raíces cuadradas de los eigenvalores. Aquí nada es un círculo falsificado. Las cruces son las medias. Los puntos están sombreados por su responsabilidad suave: un punto que pertenece a un componente toma su color, un punto repartido entre dos sale como una mezcla, así que los traslapes se leen como mezclas genuinas. El panel de abajo sigue a la log-verosimilitud escalando.

Dale play. Abre en la iteración 0 con tres contornos redondos idénticos sentados sobre las semillas de k-means++ — una covarianza compartida, nadie comprometido todavía. Después cada cuadro hace un paso E y un paso M, y vas viendo cómo las elipses se despegan, se inclinan hacia la diagonal de los datos y se estiran para cubrir su propio cluster mientras los puntos de las costuras van cambiando su mezcla. La log-verosimilitud se dispara en los primeros cuadros y luego avanza a gatas.

Fíjate en los primeros tres o cuatro cuadros — ahí es donde las elipses encuentran su ángulo, girando de redondas a diagonales en un par de pasos. De ahí en adelante es refinamiento: los contornos se acomodan tantito, los puntos de las costuras asientan sus mezclas, y la log-verosimilitud sube de alrededor de −1,081 en el arranque sembrado a −515 en la convergencia, veintitrés iteraciones después, donde la curva se aplana y la corrida se detiene. Reinicia y ponlo otra vez — la siembra y todo lo demás tiene semilla fija, así que moldea las mismas tres elipses cada vez.

Esa repetibilidad es una propiedad de esta semilla, no una promesa. La superficie de verosimilitud de EM es accidentada, y una inicialización distinta puede dejar varadas las elipses en un acomodo peor — me pasó exactamente eso con un par de semillas mientras armaba esto, donde la corrida enganchó dos componentes en un mismo cluster real y partió un tercero. Ese es el asunto de los óptimos locales hecho concreto, y la razón por la que la librería reinicia desde varias semillas por default.

La implementación completa

El archivo entero, sin librería, de arriba a abajo. Esto es exactamente lo que corrió la animación:

"""Gaussian mixture model with full covariances, fit by EM, from scratch.

k-means gave every point one hard label and drew every cluster as a round ball.
A Gaussian mixture keeps the idea of k centers but lets each cluster be a tilted,
stretched ellipse with soft edges: a point can belong 70% to one component and
30% to another. We fit it with expectation-maximization — the general method
lives in its own chapter; here it is specialized to the 2-D, full-covariance
mixture so we can watch the ellipses mold themselves to the data.

Pure NumPy. Every function below appears in the chapter one step at a time (the
`# region:` markers are what the book's include directives pull in).
"""

import numpy as np
import pandas as pd


# region: gaussian
def gaussian(X, mean, cov):
    """Multivariate Gaussian density N(x | mean, cov) for every row of X.

    X is (N, D), mean is (D,), cov is (D, D); the result is (N,), the density
    at each point. This is the shape term of the whole model — a full DxD
    covariance is what lets the level sets be tilted ellipses instead of the
    circles a single variance would give you.
    """
    D = X.shape[1]
    diff = X - mean                                   # (N, D)
    inv = np.linalg.inv(cov)                           # (D, D)
    # Mahalanobis distance squared for every row, without a Python loop.
    maha = np.einsum("ni,ij,nj->n", diff, inv, diff)   # (N,)
    norm = 1.0 / (np.power(2 * np.pi, D / 2) * np.sqrt(np.linalg.det(cov)))
    return norm * np.exp(-0.5 * maha)
# endregion


# region: e_step
def e_step(X, weights, means, covs):
    """E-step: the responsibility of every component for every point.

    Returns gamma, an (N, K) matrix where gamma[i, k] is the posterior
    probability that point i was drawn from component k, given the current
    parameters. Each row is a soft assignment that sums to 1 — this is the
    'soft' in soft clustering, and the whole reason a GMM differs from k-means.
    """
    K = len(weights)
    weighted = np.stack([weights[k] * gaussian(X, means[k], covs[k])
                         for k in range(K)], axis=1)   # (N, K): pi_k * N(x|k)
    total = weighted.sum(axis=1, keepdims=True)         # (N, 1): the mixture density
    return weighted / np.maximum(total, 1e-300)
# endregion


# region: m_step
def m_step(X, gamma, reg=1e-6):
    """M-step: re-estimate every parameter as a responsibility-weighted fit.

    Given the responsibilities, each component's new mean, covariance, and
    mixing weight is just a weighted version of the ordinary sample estimate,
    where every point contributes in proportion to how much it belongs to that
    component. `reg` adds a tiny ridge to each covariance diagonal so a
    component that grabs only a few points can't collapse to a singular matrix.
    """
    N, D = X.shape
    K = gamma.shape[1]
    Nk = gamma.sum(axis=0)                              # (K,): soft count per component
    weights = Nk / N                                    # mixing weights
    means = (gamma.T @ X) / Nk[:, None]                 # (K, D): weighted means
    covs = np.empty((K, D, D))
    for k in range(K):
        diff = X - means[k]                             # (N, D)
        covs[k] = (gamma[:, k, None] * diff).T @ diff / Nk[k]
        covs[k] += reg * np.eye(D)                       # keep it non-singular
    return weights, means, covs
# endregion


# region: log_likelihood
def log_likelihood(X, weights, means, covs):
    """Total log-likelihood of the data under the current mixture.

    Sum over points of the log of the mixture density. EM can only push this
    up or leave it flat, never down — it is the number the whole loop climbs,
    and the flat spot is where we stop.
    """
    K = len(weights)
    weighted = np.stack([weights[k] * gaussian(X, means[k], covs[k])
                         for k in range(K)], axis=1)   # (N, K)
    return float(np.log(np.maximum(weighted.sum(axis=1), 1e-300)).sum())
# endregion


# region: init_params
def init_params(X, K, rng):
    """Seed the mixture: k-means++ means, one shared covariance, equal weights.

    Spreading the initial means with k-means++ (the same trick from the k-means
    chapter) keeps EM out of the bad local optima a careless random start falls
    into. Every component starts as one big round copy of the data's overall
    covariance, so the animation opens with three identical blobs that then peel
    apart and mold to their own cluster's shape.
    """
    # k-means++ seeding of the means
    first = int(rng.integers(len(X)))
    means = [X[first]]
    for _ in range(1, K):
        d2 = np.min([((X - m) ** 2).sum(axis=1) for m in means], axis=0)
        means.append(X[int(rng.choice(len(X), p=d2 / d2.sum()))])
    means = np.array(means, dtype=float)

    shared = np.cov(X.T)                                 # one covariance for all
    covs = np.stack([shared.copy() for _ in range(K)])
    weights = np.full(K, 1.0 / K)                        # equal mixing weights
    return weights, means, covs
# endregion


# region: fit
def fit(X, K, rng, max_iter=100, tol=1e-4):
    """Run EM to convergence, recording a snapshot every iteration.

    Alternate the E-step and the M-step until the log-likelihood stops climbing
    by more than `tol`. Returns the final weights, means, covariances, and the
    full per-iteration history — each snapshot holds the parameters, the
    responsibilities, and the log-likelihood, which is exactly what the chapter
    replays frame by frame so you watch the ellipses form.
    """
    weights, means, covs = init_params(X, K, rng)
    gamma = e_step(X, weights, means, covs)
    ll = log_likelihood(X, weights, means, covs)
    history = [_snapshot(weights, means, covs, gamma, ll)]
    for _ in range(max_iter):
        weights, means, covs = m_step(X, gamma)         # re-estimate parameters
        gamma = e_step(X, weights, means, covs)         # recompute responsibilities
        new_ll = log_likelihood(X, weights, means, covs)
        history.append(_snapshot(weights, means, covs, gamma, new_ll))
        if new_ll - ll < tol:                           # likelihood stopped climbing
            ll = new_ll
            break
        ll = new_ll
    return weights, means, covs, ll, history
# endregion


# region: predict
def predict(X, weights, means, covs):
    """Hard label: the single most responsible component for each point.

    argmax over the responsibilities collapses the soft assignment back to one
    label per point, so a GMM can stand in for k-means when you need a firm
    clustering — but you threw away the confidence to get it.
    """
    return np.argmax(e_step(X, weights, means, covs), axis=1)
# endregion


# region: predict_proba
def predict_proba(X, weights, means, covs):
    """Soft label: the full responsibility vector for each point.

    This is what a GMM has that k-means never will — a distribution over
    components per point, so you can see the points on a boundary that belong
    honestly to two clusters at once.
    """
    return e_step(X, weights, means, covs)
# endregion


def _snapshot(weights, means, covs, gamma, ll):
    """One frame of the run: parameters, responsibilities, and log-likelihood."""
    return {
        "weights": weights.copy(),
        "means": means.copy(),
        "covs": covs.copy(),
        "gamma": gamma.copy(),
        "log_likelihood": ll,
    }


def load_data(path="../data/blobs.csv"):
    """Anisotropic blobs snapshot: 300 2-D points (x, y) plus the true blob id.

    The blob column is ground truth we NEVER fit on — the GMM sees only the
    coordinates. It exists so we can score the clustering afterwards.
    """
    return pd.read_csv(path)

La versión con librería

Nadie escribe a mano una mezcla gaussiana en producción, y una vez que armaste una entiendes exactamente qué está haciendo la librería. El GaussianMixture de scikit-learn con covariance_type="full" es el mismo ciclo de EM con covarianzas completas, inicialización con k-means, varios reinicios para que una semilla con mala suerte no decida la respuesta, y la misma protección de ridge contra componentes singulares. Optimiza la log-verosimilitud idéntica, así que con estos datos debería caer donde caímos nosotros:

def sklearn_gmm(X, k, seed=0):
    """Fit a full-covariance GMM with scikit-learn.

    Returns the hard labels, the component means, covariances, mixing weights,
    and the average per-sample log-likelihood (GaussianMixture.score) — the
    same things our fit() produces, so the two are directly comparable.
    """
    gm = GaussianMixture(n_components=k, covariance_type="full",
                         n_init=5, random_state=seed)
    labels = gm.fit_predict(X)
    return (labels, gm.means_, gm.covariances_, gm.weights_,
            float(gm.score(X)), gm)

La librería también abarata la selección de modelo. Como la mezcla es un modelo de probabilidad real, podemos calificar cada elección de kk con el criterio de información bayesiano — la log-verosimilitud penalizada por el número de parámetros — y quedarnos con la kk que lo minimiza. A diferencia de la log-verosimilitud, que siempre mejora conforme agregas componentes, BIC sí tiene un fondo de verdad:

def bic_sweep(X, ks, seed=0):
    """Fit a GMM for a range of k and record the Bayesian information criterion.

    Unlike log-likelihood, which always improves as you add components, BIC
    penalizes the parameter count — so it has a genuine minimum, and the k that
    minimizes it is the model-selection answer for 'how many clusters'.
    """
    out = []
    for k in ks:
        gm = GaussianMixture(n_components=k, covariance_type="full",
                             n_init=5, random_state=seed).fit(X)
        out.append((int(k), float(gm.bic(X))))
    return out

Ajusta la mezcla para kk de 1 a 8 y grafica el BIC de cada una:

El BIC baja en picada de 1654.1 en k=1 a 1404.8 en k=2 hasta un mínimo de 1126.7 en k=3, y de ahí voltea y sube — 1152.7 en k=4 y para arriba, cada componente extra costando más en parámetros de lo que gana en ajuste. El punto naranja en k=3 es la respuesta, y es la correcta porque construimos los datos con tres blobs. Esta es la pregunta del codo de k-means pero ahora con un objetivo real debajo: ya no hay que adivinar un doblez a ojo, nada más el mínimo de una curva.

Desde cero contra librería

Mismos datos, mismos tres componentes, ambos ajustados hasta converger — nuestro EM hecho desde cero contra el GaussianMixture de sklearn. Las barras son la log-verosimilitud promedio por punto, el número que ambos están maximizando:

Coinciden hasta el cuarto decimal — los dos en −1.7162. Eso no es suerte del redondeo; ambos ciclos optimizan la misma verosimilitud sobre los mismos datos y escalan al mismo óptimo, y desde un arranque con k-means++ la cuenca es la buena. Los dos etiquetados también coinciden perfecto: índice de Rand ajustado de 1.0 entre nuestros clusters y los de sklearn, salvo el intercambio arbitrario de a cuál componente le toca llamarse 0, 1 o 2. Y contra el sobre sellado de ids verdaderos de blob, ambos sacan un limpio 1.0 — cada punto regresó al grupo del que realmente salió.

Ahora sí, la comparación alrededor de la cual está construido todo este capítulo. Los mismos 300 puntos, agrupados de dos maneras. A la izquierda, k-means: etiquetas duras y un círculo en cada centro estirándose hasta la dispersión del cluster. A la derecha, nuestro GMM: etiquetas duras del mismo ajuste, cada cluster envuelto en su elipse inclinada de 2σ.

Mira las costuras. k-means dibuja sus círculos y, como juzga cada punto por distancia simple a un centro, corta una línea recta limpia entre los clusters — que atraviesa en diagonal los grupos reales, etiquetando mal los extremos alargados que se asoman más allá del centro de un vecino. Saca un índice de Rand ajustado de 0.60 contra la verdad: un tercio de los puntos en el grupo equivocado, y con total seguridad. Las elipses del GMM se acuestan a lo largo de la diagonal de cada cluster, la frontera de Mahalanobis se curva alrededor de la forma, y los extremos alargados se quedan con su propio componente. Saca 1.0. Mismos datos, misma k, misma garantía de convergencia — la única diferencia es que a un modelo se le permitió dibujar una elipse y al otro no.

Conclusiones

Una mezcla gaussiana es a lo que le tiras cuando k-means es el instinto correcto pero la forma equivocada. Sigues creyendo que los datos son un puñado de blobs alrededor de centros; nada más ya no crees que los blobs sean redondos, iguales ni estén bien separados. Las covarianzas completas te compran clusters estirados e inclinados como los datos lo pidan, las responsabilidades suaves te compran honestidad en las fronteras, y como es un modelo de probabilidad real obtienes cosas que el clustering solo nunca te da: una densidad contra la cual evaluar puntos nuevos, una confianza por punto y el BIC para elegir el número de componentes sin andar entrecerrando los ojos frente a un codo. Con los datos de aquí eso fue todo el partido — k-means en 0.60, la mezcla en un 1.0 perfecto, decidido enteramente por si un cluster tenía permiso de ser una elipse.

La cuenta llega por dos lados, y ambos vienen de la misma fuente: una covarianza completa son muchos parámetros. En dimensiones altas eso es mucho que estimar y una forma fácil de caer en overfitting, y un componente que se queda con muy pocos puntos puede colapsar su covarianza a un pico y reventar la verosimilitud — por eso el ridge en la diagonal no es opcional. Y EM cae en óptimos locales, más de los que cae k-means, así que la respuesta depende de la semilla y reinicias por si las dudas. Cuando tus clusters de verdad son redondos y están bien separados, k-means es más rápido y tiene menos cosas que puedan salir mal; guarda la mezcla para cuando la forma sea todo el problema. Es el mismo trueque que este curso hace una y otra vez — un modelo más flexible ajusta más, y te pide tener más cuidado para distinguir si ajustó o nada más memorizó. El agrupamiento suave y con forma de aquí también es una puerta: el ciclo de paso E y luego paso M, las responsabilidades, la log-verosimilitud que solo sube, son los mismos movimientos detrás de los modelos ocultos de Markov, los modelos de tópicos y la mitad del aprendizaje no supervisado. Construye la mezcla a mano una vez y ya construiste el patrón.