Scalable Posterior Uncertainty for Flexible Density-Based Clustering

TL;DR

This paper presents a scalable, GPU-compatible framework for quantifying uncertainty in density-based clustering using martingale posterior distributions and differentiable density estimators.

Contribution

It introduces a nonparametric approach combining martingale posteriors with density estimators for uncertainty quantification in clustering.

Findings

Efficient density resampling with normalizing flows on GPU.

Principled inference on clustering uncertainty.

Application to image and single-cell RNA data.

Abstract

We introduce a novel framework for uncertainty quantification in clustering that combines martingale posterior distributions with density-based clustering. Unlike classical model-based approaches, which define clusters at the latent level of a mixture model, we treat clusters as explicit functionals of the data-generating density, without assuming any specific parametric form. To characterize density uncertainty, we obtain martingale posterior samples via a predictive resampling scheme driven by model score evaluations. This allows us to leverage state-of-the-art differentiable density estimators, such as normalizing flows, making density resampling efficient in large-scale settings and fully parallelizable on modern GPU hardware. Martingale posterior samples of the clustering structure are then obtained by applying density-based clustering to the density draws, enabling principled inference on any clustering-related quantity. Casting the inference target as a density functional further enables a rigorous theoretical analysis of the procedure's convergence properties. We apply our methodology to image and single-cell RNA sequencing data, demonstrating the computational efficiency afforded by its GPU compatibility as well as its ability to recover meaningful clustering structures, with associated uncertainty, across diverse domains.