Bayesian Sparse GMM

Bayesian Sparse Gaussian Mixture Model (GMM) implementation in Python.

Installation

To install the latest release:

pip install bayesian-sparse-gmm

Or for development (editable mode):

git clone https://github.com/Coalyx/bayesian-sparse-gmm.git
cd bayesian-sparse-gmm
pip install -e .

Quick Start

import numpy as np
from sklearn.datasets import make_blobs
from sklearn.preprocessing import StandardScaler
from bayesian_sparse_gmm import BayesianSparseGMM

# Append noise dimensions to true clusters to verify that the model successfully performs feature selection.
rng = np.random.default_rng(42)
X_clean, _ = make_blobs(n_samples=200, centers=3, n_features=2, cluster_std=0.5, random_state=42)
X_noise = rng.normal(loc=0.0, scale=1.0, size=(200, 8))
X = np.hstack([X_clean, X_noise])

# Standardize features to satisfy the zero-mean assumptions in the prior structure.
X = StandardScaler().fit_transform(X)

model = BayesianSparseGMM(
    K_max=5,
    n_iter=300,
    burn_in=100,
    lambda_0=10.0,
    lambda_1=0.05,
    random_state=42,
    verbose=0
)
model.fit(X)

print(f"Number of active clusters: {model.n_clusters_}")
print(f"Selected informative features: {model.selected_features_}")
print(f"Feature inclusion probabilities: {model.feature_probabilities_.round(3)}")

labels = model.predict(X)

Optimization Methods

The model supports two optimization architectures via the optimizer parameter:

1. Gibbs Sampling (MCMC) (optimizer="default"): The original, mathematically exact implementation. It explores the full posterior distribution but requires processing the entire dataset per iteration. Best for smaller datasets or when exact posterior distributions are required.
2. Stochastic Variational Inference (SVI) (optimizer="svi"): A highly-scalable approach using Coordinate Ascent Variational Inference (CAVI) with Natural Gradients. It processes data in mini-batches, making it orders of magnitude faster and capable of scaling to extremely large datasets ($N \gg 10,000$).

GPU / CUDA Acceleration

The model supports three backends, selected via the backend parameter:

Install GPU support

# CuPy (recommended — more flexible with driver versions)
pip install "bayesian-sparse-gmm[cuda]"
# or install directly for your CUDA version, e.g.:
pip install cupy-cuda12x   # CUDA 12.x
pip install cupy-cuda11x   # CUDA 11.x

Run with GPU (Python API)

model = BayesianSparseGMM(
    K_max=10,
    n_iter=500,
    backend="cuda",   # GPU via CuPy → Numba CUDA → NumPy fallback
    random_state=42,
)
model.fit(X)

Development and Testing

Install development dependencies:

pip install -e ".[dev]"

Run tests using pytest:

pytest

Algorithm Overview

Bayesian Sparse Gaussian Mixture Model (BSGMM) is a robust clustering algorithm designed specifically for high-dimensional data where the number of features significantly exceeds the number of samples ($p \gg n$). It integrates a Spike-and-Slab LASSO prior to perform simultaneous clustering and feature selection.

Suitable Use Cases

1. High-Dimensional Clustering ($p \gg n$): When dealing with datasets where traditional clustering algorithms (like K-Means or standard GMM) fail due to the "curse of dimensionality". Examples include bioinformatics (e.g., single-cell RNA-seq, genomics), text mining (high-dimensional TF-IDF matrices), and high-resolution images.
2. Automatic Feature Selection (Interpretability): When the goal is not only to cluster the samples but also to identify which specific features (biomarkers, keywords, pixels) drive the cluster assignments. The model automatically shrinks noisy features to exactly zero.
3. Unknown Number of Clusters: When the true number of clusters is unknown. BSGMM can dynamically infer the optimal number of clusters from the data (bounded by K_max).
4. Weak Signals in Noisy Backgrounds: When the discriminative signal is weak and dispersed among thousands of irrelevant features, the model's sparsity mechanism is highly effective at pooling signals.

Hyperparameters

Understanding the key hyperparameters is crucial for fine-tuning the model's sparsity and clustering behavior:

MCMC Parameters (optimizer="default")

SVI Parameters (optimizer="svi")

Tip: For extremely high-dimensional datasets with heavy noise, tuning lambda_0 to be larger and theta to be smaller will encourage more aggressive feature selection.

Reference

@article{JMLR:v26:23-0142,
  author  = {Dapeng Yao and Fangzheng Xie and Yanxun Xu},
  title   = {Bayesian Sparse Gaussian Mixture Model for Clustering in High Dimensions},
  journal = {Journal of Machine Learning Research},
  year    = {2025},
  volume  = {26},
  number  = {21},
  pages   = {1--50},
  url     = {http://jmlr.org/papers/v26/23-0142.html}
}

Contributors

• Nam Nam (@Neeze) - Developer of the SVI (Stochastic Variational Inference) optimizer, GPU/CUDA acceleration, and benchmarking suite.