Gaussian Mixture Model

A Gaussian Mixture Model (GMM) is a GPU-accelerated probabilistic clustering and density-estimation algorithm. It models a dataset as a weighted sum of n_components Gaussian components, learning a weight, mean, and covariance for each one with the Expectation-Maximization (EM) algorithm.

Use a GMM when you want soft cluster assignments (a probability that each row belongs to each component), a generative density model you can score new points against, or clusters that are elliptical rather than spherical. Unlike K-Means, which assigns every row to exactly one centroid, a GMM returns a full responsibility distribution over components and captures per-component shape through its covariance. Its primary outputs are the component weights, means, covariances (and their Cholesky-factored precisions), per-row labels, and the converged log-likelihood lower bound.

Example API Usage

C++ API

Fitting a mixture

Fitting learns the component weights, means, and covariances from a dataset on the device. The covariance-shaped outputs (covariances, precisions_chol, precisions) have a layout that depends on covariance_type; see the Covariance types section below for the exact shapes.

1
#include <cuvs/cluster/gmm.hpp>
2
3
#include <raft/core/device_mdarray.hpp>
4
#include <raft/core/resources.hpp>
5
6
using namespace cuvs::cluster;
7
8
raft::device_resources res;
9
raft::device_matrix_view<const float, int64_t> dataset = load_dataset();
10
11
gmm::params params;
12
params.n_components = 1024;
13
params.cov_type = gmm::covariance_type::FULL;
14
params.max_iter = 100;
15
params.tol = 1e-3;
16
17
// Output buffers; covariance-shaped extents follow params.cov_type.
18
auto [weights, means, covariances, precisions_chol, precisions, labels] =
19
    allocate_outputs(res, params, dataset);
20
21
float lower_bound;
22
int n_iter;
23
bool converged;
24
25
gmm::fit(res,
26
         params,
27
         dataset,
28
         weights.view(),
29
         means.view(),
30
         covariances.view(),
31
         precisions_chol.view(),
32
         precisions.view(),
33
         labels.view(),
34
         raft::make_host_scalar_view(&lower_bound),
35
         raft::make_host_scalar_view(&n_iter),
36
         raft::make_host_scalar_view(&converged));

Assigning labels and scoring

After fitting, reuse the learned weights, means, and precisions_chol to assign hard labels (predict), produce per-component responsibilities (predict_proba), or evaluate the per-row log-likelihood of new data (score_samples).

1
int64_t n_samples = dataset.extent(0);
2
3
// Hard label per row (argmax responsibility).
4
auto labels = raft::make_device_vector<int, int64_t>(res, n_samples);
5
gmm::predict(
6
  res, params, dataset, weights.view(), means.view(), precisions_chol.view(), labels.view());
7
8
// (n_samples, n_components) responsibility matrix.
9
auto resp = raft::make_device_matrix<float, int64_t>(res, n_samples, params.n_components);
10
gmm::predict_proba(
11
  res, params, dataset, weights.view(), means.view(), precisions_chol.view(), resp.view());
12
13
// Per-row log-likelihood under the fitted mixture.
14
auto log_prob = raft::make_device_vector<float, int64_t>(res, n_samples);
15
gmm::score_samples(
16
  res, params, dataset, weights.view(), means.view(), precisions_chol.view(), log_prob.view());

How GMM works

EM alternates between two steps until the average log-likelihood stops improving:

1. E-step: given the current parameters, compute the responsibility of each component for each row — the posterior probability that the row was generated by that component.
2. M-step: given the responsibilities, update each component’s weight, mean, and covariance to the responsibility-weighted statistics of the data.

The algorithm repeats until it reaches max_iter or the per-sample average log-likelihood changes by less than tol. Because both steps reduce to dense linear algebra over many rows and components, the GPU is well suited to the work.

Covariance types

covariance_type controls how much shape each component can express, trading flexibility for parameters and cost. The covariance-shaped buffers (covariances, precisions_chol, precisions) are passed as flat device vectors because their logical shape depends on the covariance type. With K = n_components and d = n_features the expected lengths are (row-major):

For full/tied, precisions_chol holds the upper-triangular factor U of each precision matrix (precision = U @ Uᵀ); for diag/spherical it holds reciprocal standard deviations. These conventions match scikit-learn’s GaussianMixture.

When to use

Use a GMM when soft, probabilistic assignments matter, when components are elliptical or have different shapes, or when you need a density model to score or compare new points. Prefer full or diag covariances when component shape is informative, and spherical or tied when data is limited or speed matters more than per-component shape. If you only need hard, roughly spherical partitions, K-Means is simpler and faster.

Configuration parameters

Tuning

Start with n_components and covariance_type. More components and richer covariances capture more structure but cost more memory and time, and can overfit when data is limited; raise reg_covar if covariances become ill-conditioned. Use kmeans initialization for robust default seeding, and increase n_init when different seeds produce noticeably different log-likelihoods. Tune max_iter and tol together: if n_iter regularly reaches max_iter, increase max_iter or relax tol.

Memory footprint

Fitting streams the E and M steps over tiles of rows, so it never materializes the full (n_samples, n_components) responsibility matrix — peak device memory stays bounded by the input data, the model parameters, and one responsibility tile, independent of n_samples. predict and score_samples likewise avoid the full responsibility matrix; only predict_proba materializes the (n_samples, n_components) output because that matrix is its result.