Gaussian Mixture with torch.distributions

Fit a Gaussian mixture by gradient ascent on the log likelihood, with torch.distributions doing the density maths.

EM is the classic way to fit a mixture. Gradient descent also works, and it is the method that carries over to mixtures inside a neural network, such as a mixture density network. torch.distributions gives you the model in two lines:

mixture = D.Categorical(logits=logits)                        # weights
components = D.MultivariateNormal(loc, scale_tril=scale_tril) # k Gaussians
gmm = D.MixtureSameFamily(mixture, components)
loss = -gmm.log_prob(X).mean()

The parameters must stay valid while the optimiser moves them freely. So train unconstrained tensors and map them: logits give weights through a softmax, and scale_tril is a lower triangular matrix whose diagonal goes through softplus so that it stays positive.

Fit k components to X with Adam for steps steps and return the fitted MixtureSameFamily.

Function signature

def fit_gmm(X, k, steps=150):

Example

Input

X = 600 points from three 2D Gaussian clusters with weights 0.5, 0.3, 0.2
gmm = fit_gmm(X, k=3)

Output

gmm.mixture_distribution.probs  ->  about tensor([0.5, 0.3, 0.2])  (any order)
gmm.log_prob(X_test).mean()     ->  about -3.06, the same as sklearn

Reasoning

The maximum likelihood mixture puts one component on each cluster, with weights equal to the share of points in it. sklearn's EM finds the same optimum, so the two average log likelihoods agree.

Constraints and notes

• X is a float32 tensor of shape (n, d). Use full covariance matrices.
• Start the means somewhere sensible. Random points can put two components in one cluster. The centres from a quick sklearn.cluster.KMeans are a good start.
• Return the distribution built from detached tensors, so that gmm.log_prob(x) does not track gradients.
• Graded on the result: the mean log likelihood on held-out data must be within 0.05 of sklearn's GaussianMixture.