Student-t and von Mises-Fisher Mixtures¶

This notebook demonstrates the Phase 4 (v0.4 roadmap) additions: StudentTMixture (Peel & McLachlan, 2000), robust to outliers, and VonMisesFisherMixture (Banerjee, Dhillon, Ghosh & Sra, 2005), for L2-normalized (directional) data -- both built on the same shared BaseMixture GaussianMixture is now rebased onto.

See the Other Mixture Families user guide page for background on both classes.

In [1]:

import numpy as np
import torch
import matplotlib.pyplot as plt
import os

os.chdir('../')

from tgmm import GaussianMixture, StudentTMixture, VonMisesFisherMixture
from tgmm.metrics import ClusteringMetrics
from tgmm.plotting import plot_gmm

random_state = 42

import numpy as np import torch import matplotlib.pyplot as plt import os os.chdir('../') from tgmm import GaussianMixture, StudentTMixture, VonMisesFisherMixture from tgmm.metrics import ClusteringMetrics from tgmm.plotting import plot_gmm random_state = 42

Part 1: StudentTMixture vs GaussianMixture on contaminated data¶

Two well-separated clusters, plus 5% uniform-noise outliers scattered across the plot. A Gaussian component's covariance gets dragged toward any outliers it's forced to "explain"; a Student-t component's heavier tails let it downweight them instead (via the per-sample u_ik weight -- see sample_weights below).

In [2]:

gen = torch.Generator().manual_seed(random_state)
n_per = 250
true_centers = torch.tensor([[5.0, 5.0], [-5.0, -5.0]])

X_clean = torch.cat([
    torch.randn(n_per, 2, generator=gen) * 0.6 + true_centers[0],
    torch.randn(n_per, 2, generator=gen) * 0.6 + true_centers[1],
])
n_out = int(0.05 * 2 * n_per)
outliers = (torch.rand(n_out, 2, generator=gen) - 0.5) * 40
X = torch.cat([X_clean, outliers])

plt.figure(figsize=(5, 5))
plt.scatter(X[:-n_out, 0], X[:-n_out, 1], s=10, label='clean')
plt.scatter(X[-n_out:, 0], X[-n_out:, 1], s=10, c='red', marker='x', label='outliers')
plt.legend()
plt.title(f'{2*n_per} clean points + {n_out} outliers (5%)')
plt.show()

gen = torch.Generator().manual_seed(random_state) n_per = 250 true_centers = torch.tensor([[5.0, 5.0], [-5.0, -5.0]]) X_clean = torch.cat([ torch.randn(n_per, 2, generator=gen) * 0.6 + true_centers[0], torch.randn(n_per, 2, generator=gen) * 0.6 + true_centers[1], ]) n_out = int(0.05 * 2 * n_per) outliers = (torch.rand(n_out, 2, generator=gen) - 0.5) * 40 X = torch.cat([X_clean, outliers]) plt.figure(figsize=(5, 5)) plt.scatter(X[:-n_out, 0], X[:-n_out, 1], s=10, label='clean') plt.scatter(X[-n_out:, 0], X[-n_out:, 1], s=10, c='red', marker='x', label='outliers') plt.legend() plt.title(f'{2*n_per} clean points + {n_out} outliers (5%)') plt.show()

In [3]:

gmm = GaussianMixture(n_components=2, random_state=random_state, n_init=5)
gmm.fit(X)

stm = StudentTMixture(n_components=2, covariance_type='full', dof='per_component',
                       random_state=random_state, n_init=5)
stm.fit(X)

def match_err(means, true_centers=true_centers):
    d = torch.cdist(means, true_centers)
    return min(d[0, 0] + d[1, 1], d[0, 1] + d[1, 0]).item() / 2

print(f'GaussianMixture mean error from truth:  {match_err(gmm.means_):.3f}')
print(f'StudentTMixture mean error from truth:  {match_err(stm.means_):.3f}')
print(f'StudentTMixture fitted dof_: {stm.dof_.tolist()}')

gmm = GaussianMixture(n_components=2, random_state=random_state, n_init=5) gmm.fit(X) stm = StudentTMixture(n_components=2, covariance_type='full', dof='per_component', random_state=random_state, n_init=5) stm.fit(X) def match_err(means, true_centers=true_centers): d = torch.cdist(means, true_centers) return min(d[0, 0] + d[1, 1], d[0, 1] + d[1, 0]).item() / 2 print(f'GaussianMixture mean error from truth: {match_err(gmm.means_):.3f}') print(f'StudentTMixture mean error from truth: {match_err(stm.means_):.3f}') print(f'StudentTMixture fitted dof_: {stm.dof_.tolist()}')

/home/asp/Downloads/HeaDS/tgmm/tgmm/mixtures/base.py:665: UserWarning: With n_init=5 and random_state=42, initializations will use random states [42, 43, ..., 46]. The best initialization's random state will be stored in best_random_state_.
  warnings.warn(
/home/asp/Downloads/HeaDS/tgmm/tgmm/mixtures/base.py:714: UserWarning: Best result from random_state=42. To reproduce this specific result: use random_state=42 with n_init=1.
  warnings.warn(

GaussianMixture mean error from truth:  0.258
StudentTMixture mean error from truth:  0.066
StudentTMixture fitted dof_: [1.4095878601074219, 18.697467803955078]

In [4]:

fig, axes = plt.subplots(1, 2, figsize=(11, 5))
plot_gmm(X, gmm=gmm, ax=axes[0])
axes[0].set_title(f'GaussianMixture (mean error {match_err(gmm.means_):.2f})')
plot_gmm(X, gmm=stm, ax=axes[1])
axes[1].set_title(f'StudentTMixture (mean error {match_err(stm.means_):.2f})')
plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(11, 5)) plot_gmm(X, gmm=gmm, ax=axes[0]) axes[0].set_title(f'GaussianMixture (mean error {match_err(gmm.means_):.2f})') plot_gmm(X, gmm=stm, ax=axes[1]) axes[1].set_title(f'StudentTMixture (mean error {match_err(stm.means_):.2f})') plt.tight_layout() plt.show()

Flagging outliers with sample_weights¶

sample_weights(X) returns each point's robustness weight u_ik at its predicted component -- small values are points the t-mixture effectively treated as outliers.

In [5]:

weights = stm.sample_weights(X)

plt.figure(figsize=(5, 5))
sc = plt.scatter(X[:, 0], X[:, 1], c=weights, cmap='viridis_r', s=12, vmin=0, vmax=1)
plt.colorbar(sc, label='sample weight u_ik (low = likely outlier)')
plt.title('StudentTMixture outlier weights')
plt.show()

print(f'Mean weight, clean points:    {weights[:-n_out].mean():.3f}')
print(f'Mean weight, injected outliers: {weights[-n_out:].mean():.3f}')

weights = stm.sample_weights(X) plt.figure(figsize=(5, 5)) sc = plt.scatter(X[:, 0], X[:, 1], c=weights, cmap='viridis_r', s=12, vmin=0, vmax=1) plt.colorbar(sc, label='sample weight u_ik (low = likely outlier)') plt.title('StudentTMixture outlier weights') plt.show() print(f'Mean weight, clean points: {weights[:-n_out].mean():.3f}') print(f'Mean weight, injected outliers: {weights[-n_out:].mean():.3f}')

Mean weight, clean points:    1.049
Mean weight, injected outliers: 0.027

Part 2: VonMisesFisherMixture on high-dimensional unit-norm embeddings¶

Simulating the kind of data a real text-embedding pipeline produces: 768-dimensional vectors, L2-normalized before clustering by cosine similarity. We compare VonMisesFisherMixture against GaussianMixture(covariance_type='diag') fit on the same normalized data -- a reasonable baseline, just one that doesn't know the data lives on a sphere -- using icl (within each family) and cosine silhouette (comparable across both).

In [6]:

d = 768
n_per = 150
K_true = 4

gen = torch.Generator().manual_seed(random_state)
true_mus = torch.randn(K_true, d, generator=gen)
true_mus = true_mus / true_mus.norm(dim=1, keepdim=True)
true_kappa = 300.0

# Generate actual vMF-distributed data by constructing a VonMisesFisherMixture
# with known parameters and using its own (Wood, 1994) rejection sampler --
# more faithful than an ad-hoc approximation, and exercises the same code path
# being demonstrated below.
true_model = VonMisesFisherMixture(n_components=K_true, n_features=d)
true_model.weights_ = torch.full((K_true,), 1.0 / K_true)
true_model.means_ = true_mus
true_model.kappa_ = torch.full((K_true,), true_kappa)

torch.manual_seed(random_state)
X_list, labels_list = [], []
for k in range(K_true):
    X_list.append(true_model._sample_component(k, n_per))
    labels_list.append(torch.full((n_per,), k))
X_emb = torch.cat(X_list)
labels_true = torch.cat(labels_list)

print(f'X_emb shape: {tuple(X_emb.shape)}, row norms (should be ~1): '
      f'{X_emb.norm(dim=1).min():.4f} - {X_emb.norm(dim=1).max():.4f}')

d = 768 n_per = 150 K_true = 4 gen = torch.Generator().manual_seed(random_state) true_mus = torch.randn(K_true, d, generator=gen) true_mus = true_mus / true_mus.norm(dim=1, keepdim=True) true_kappa = 300.0 # Generate actual vMF-distributed data by constructing a VonMisesFisherMixture # with known parameters and using its own (Wood, 1994) rejection sampler -- # more faithful than an ad-hoc approximation, and exercises the same code path # being demonstrated below. true_model = VonMisesFisherMixture(n_components=K_true, n_features=d) true_model.weights_ = torch.full((K_true,), 1.0 / K_true) true_model.means_ = true_mus true_model.kappa_ = torch.full((K_true,), true_kappa) torch.manual_seed(random_state) X_list, labels_list = [], [] for k in range(K_true): X_list.append(true_model._sample_component(k, n_per)) labels_list.append(torch.full((n_per,), k)) X_emb = torch.cat(X_list) labels_true = torch.cat(labels_list) print(f'X_emb shape: {tuple(X_emb.shape)}, row norms (should be ~1): ' f'{X_emb.norm(dim=1).min():.4f} - {X_emb.norm(dim=1).max():.4f}')

X_emb shape: (600, 768), row norms (should be ~1): 1.0000 - 1.0000

In [7]:

vmf = VonMisesFisherMixture(n_components=K_true, concentration='per_component',
                             random_state=random_state, n_init=3, max_iter=100)
vmf.fit(X_emb)

gmm_diag = GaussianMixture(n_components=K_true, covariance_type='diag',
                            random_state=random_state, n_init=3)
gmm_diag.fit(X_emb)

pred_vmf = vmf.predict(X_emb)
pred_gmm = gmm_diag.predict(X_emb)

ari_vmf = ClusteringMetrics.adjusted_rand_score(labels_true, pred_vmf)
ari_gmm = ClusteringMetrics.adjusted_rand_score(labels_true, pred_gmm)
sil_vmf = ClusteringMetrics.silhouette_score(X_emb, pred_vmf, metric='cosine')
sil_gmm = ClusteringMetrics.silhouette_score(X_emb, pred_gmm, metric='cosine')

print(f'VonMisesFisherMixture:  ARI={ari_vmf:.3f}  cosine silhouette={sil_vmf:.3f}  '
      f'ICL={vmf.icl(X_emb):.1f}  fitted kappa={vmf.kappa_.mean().item():.1f} (true={true_kappa})')
print(f'GaussianMixture(diag):  ARI={ari_gmm:.3f}  cosine silhouette={sil_gmm:.3f}  '
      f'ICL={gmm_diag.icl(X_emb):.1f}')

vmf = VonMisesFisherMixture(n_components=K_true, concentration='per_component', random_state=random_state, n_init=3, max_iter=100) vmf.fit(X_emb) gmm_diag = GaussianMixture(n_components=K_true, covariance_type='diag', random_state=random_state, n_init=3) gmm_diag.fit(X_emb) pred_vmf = vmf.predict(X_emb) pred_gmm = gmm_diag.predict(X_emb) ari_vmf = ClusteringMetrics.adjusted_rand_score(labels_true, pred_vmf) ari_gmm = ClusteringMetrics.adjusted_rand_score(labels_true, pred_gmm) sil_vmf = ClusteringMetrics.silhouette_score(X_emb, pred_vmf, metric='cosine') sil_gmm = ClusteringMetrics.silhouette_score(X_emb, pred_gmm, metric='cosine') print(f'VonMisesFisherMixture: ARI={ari_vmf:.3f} cosine silhouette={sil_vmf:.3f} ' f'ICL={vmf.icl(X_emb):.1f} fitted kappa={vmf.kappa_.mean().item():.1f} (true={true_kappa})') print(f'GaussianMixture(diag): ARI={ari_gmm:.3f} cosine silhouette={sil_gmm:.3f} ' f'ICL={gmm_diag.icl(X_emb):.1f}')

/home/asp/Downloads/HeaDS/tgmm/tgmm/mixtures/base.py:665: UserWarning: With n_init=3 and random_state=42, initializations will use random states [42, 43, ..., 44]. The best initialization's random state will be stored in best_random_state_.
  warnings.warn(

VonMisesFisherMixture:  ARI=1.000  cosine silhouette=0.111  ICL=-1791013.0  fitted kappa=311.3 (true=300.0)
GaussianMixture(diag):  ARI=1.000  cosine silhouette=0.111  ICL=-1777652.4

Both recover the 4 clusters well on this synthetic example (it's a fairly easy, well-separated case), but VonMisesFisherMixture's cosine silhouette and recovered kappa_ directly reflect the generative process, and its icl/bic/aic are only meaningful compared to other vMF fits -- not to GaussianMixture's, since the two have structurally different densities. On real embeddings with more overlap between classes, the gap between the two typically widens in VonMisesFisherMixture's favor, precisely because it isn't spending parameters modeling variation in vector length, which L2-normalization has already thrown away.

Summary¶

• StudentTMixture is a drop-in, more outlier-robust alternative to GaussianMixture -- reach for it when you suspect contamination, and use sample_weights to find out which points it downweighted.
• VonMisesFisherMixture is the right model for L2-normalized / cosine-similarity data -- reach for it instead of GaussianMixture whenever your features are directions, not points.
• Both share GaussianMixture's full EM toolkit (algorithm, temperature, step_size, n_init, bic/aic/icl) via the common BaseMixture base class.