Inference for Clustering: Conformal Sets for Cluster Labels

Abstract: While clustering is ubiquitously used across science and industry, uncertainty in cluster assignments is rarely quantified with rigorous guarantees. We propose a novel conformal inference framework for clustering that returns confidence sets for cluster labels. The key challenge is that labels are unobserved and estimated from data, so naively using deterministic cluster labels can violate exchangeability and induce severe under-coverage. To address this, we propose split conformal clustering with stochastic labels, which samples labels from soft cluster labels, fits a soft classifier to predict these stochastic labels, and calibrates conformal scores to construct confidence sets for cluster labels at any query point. We establish a finite-sample lower bound on marginal coverage that reveals how under-coverage is controlled by two properties of the clustering algorithm: consistency of estimated soft labels and replace-one stability. Under mild conditions, we prove asymptotic coverage and verify these conditions for correctly specified parametric mixture models. Simulations for mixture models show that our method attains target coverage with informative set sizes, validating our theoretical results. Applications to clustering cell types in single-cell RNA-seq data demonstrate the practical utility and interpretability of our approach to quantifying cluster label uncertainty.

• The paper introduces a novel framework that constructs finite-sample valid confidence sets for cluster labels using stochastic labeling and conformal prediction.
• It employs a split conformal approach with stochastic labels to overcome the limitations of traditional methods, yielding robust uncertainty quantification.
• Empirical results demonstrate that the method achieves tight coverage and informative set sizes across various clustering models and noise levels.

Conformal Inference for Clustering: Rigorous Uncertainty Quantification via Confidence Sets

Motivation and Problem Statement

While clustering is a central tool in unsupervised learning with extensive applications in science and industry, existing techniques offer limited means for quantifying uncertainty in cluster allocations. Ambiguities in cluster assignment, particularly at cluster boundaries or for high-noise data, pose substantial challenges for interpretability and robustness. Traditional methods predominantly focus on uncertainty related to cluster existence or centroids but neglect uncertainty in the cluster labels assigned to individual data points. This paper addresses this gap by introducing a novel conformal inference framework designed to provide distribution-free, finite-sample valid confidence sets for cluster labels.

The target is a set-valued function $\hat{C}(x) \subset \{1, \ldots, K\}$ for any input $x$ such that, with high probability, the set contains the true (unobserved) cluster label for a new data point. This construction remains agnostic to the source clustering algorithm, seeking broad applicability and robust theoretical coverage guarantees.

Previous Methods and Shortcomings

Approaches for uncertainty quantification in clustering have traditionally fallen into three categories: (1) centroid or structural uncertainty estimation (e.g., via bootstrap), (2) testing for the existence of clusters, and (3) Bayesian credible sets over partitions. None of these provide finite-sample, distribution-free guarantees for uncertainty in the cluster label itself.

Two naive approaches are considered:

1. Cutoff for Generalizable Clustering: Constructs confidence sets from soft labels (posterior cluster probabilities) by thresholding to achieve the desired nominal coverage. This method can produce uninformative sets (over-coverage) or significant under-coverage, especially for finite samples or when label probabilities are miscalibrated.
2. Split Conformal Clustering with Deterministic Labels: Mimics split conformal classification by training a classifier on cluster labels, then constructing confidence sets via calibration. However, standard deterministic clustering labels are inherently unstable and violate the exchangeability required for conformal inference, leading to serious under-coverage.

Proposed Method: Split Conformal Clustering with Stochastic Labels

To overcome these challenges, the paper introduces a meta-algorithm—split conformal clustering with stochastic labels. The essential innovation is to inject stochasticity into cluster label assignments, producing labels sampled from soft cluster assignments (e.g., posteriors from mixture models or fuzzy memberships). This process, combined with a supervised soft classifier and conformal calibration, yields valid set-valued predictions for new points, with explicit, provable coverage guarantees.

The method's workflow is:

1. Split the dataset into disjoint training and calibration sets.
2. On each, perform stochastic clustering by sampling label assignments from the soft cluster probability vectors.
3. Train a soft classifier on the sampled training labels.
4. Calibrate conformity scores on the calibration set to construct confidence sets for new points.
5. Employ permutation alignment to address label symmetry before set construction.

This procedure is flexible and can accommodate generalizable soft clustering algorithms (e.g., GMM, FCM) and various classifiers, and can be adapted in projected spaces for visualization.

Figure 1: Data points colored by true cluster (left) and regions of high-confidence cluster assignment (right) produced by split conformal clustering with stochastic GMM labels.

Theoretical Guarantees

The paper’s theoretical contribution is twofold:

1. Finite-Sample Lower Bound on Coverage: Coverage is lower-bounded by $1-\alpha$ minus terms reflecting the consistency of estimated soft labels and the replace-one stability of the clustering algorithm, both of which control the deviation from exchangeability.
2. Asymptotic Coverage: Under mild regularity conditions, including consistent estimation of cluster probabilities and diminishing influence of single data points (asymptotic replace-one invariance), the algorithm achieves asymptotic coverage at the nominal level.

These conditions are proved to hold for well-specified parametric mixture models with identifiable, efficiently estimable parameters.

Numerical Results and Empirical Performance

Simulations on both low- and high-dimensional mixture models (GMM, GaMM) illustrate the method's empirical tightness of coverage and informativeness of set size, contrasting it with naive cutoff and deterministic-label split conformal methods.

Figure 2: Visualization of confidence sets for GMM and GaMM simulations ($K{=}3$ in $\mathbb{R}^2$); mixed colors indicate regions where multiple labels are plausible under the calibrated confidence set.

Figure 3: Empirical coverage and average set size vs. noise level (top) and sample size (bottom) for GMMs; the proposed stochastic label method matches target coverage while naive methods severely underperform.

Figure 4: Analogous results for GaMMs: informative confidence sets and tight coverage are preserved in heavy-tailed, asymmetric settings.

Results confirm:

• The stochastic label conformal clustering method reliably achieves the nominal coverage across a range of noise levels and dimensions, whereas deterministic and cutoff methods fail (naive approaches under- or over-cover).
• Confidence set sizes decrease with reduced noise and increasing sample size, reflecting improved certainty, while set sizes for naive cutoff remain unnecessarily large across all regimes.

Application to Single-Cell RNA-Seq Data

The framework is applied to a standard single-cell RNA-seq data set (PBMC-3K), with confidence sets indicating regions and cell types where assignments are robust versus ambiguous—mirroring underlying biological continuity and uncertainty.

Figure 5: Confidence set sizes for cell type annotations in single-cell RNA-seq data; marker shape and color blend represent the number and identity of plausible cluster labels.

Methodological and Practical Implications

This work rigorously demonstrates that accurate uncertainty quantification for clustering labels—under minimal distributional assumptions—is feasible at non-asymptotic sample sizes, provided appropriate stochasticity and calibration. The methodology is agnostic to the underlying clustering and classifier models so long as the required consistency and stability properties are met; it is therefore extensible to more complex or domain-specific clustering procedures, including nonparametric methods and deep generative model-based clustering.

Furthermore, practitioners are able to directly interpret regions of high or low certainty, enabling informed decisions and robust downstream analysis, especially in applications such as cell-type annotation, anomaly detection, and astronomical data clustering.

Prospects for Future Development

Potential extensions include:

• Robustifying against misspecification in the choice of the number of clusters $K$, ideally incorporating post-selection inference.
• Aggregation across multiple stochastic labelings (cross-conformal methods) to improve efficiency.
• Adapting the framework to cluster-conditional coverage and to hard-label clustering via ensemble approaches (bootstrapping/subsampling).
• Correcting for distribution shift induced by label estimation using weighted conformal inference.

Conclusion

The paper establishes a comprehensive, theoretically justified, and empirically validated approach for uncertainty quantification in clustering by constructing finite-sample valid confidence sets for cluster labels using conformal prediction with stochastic cluster assignments. This paradigm enhances the reliability, interpretability, and subsequent utility of clustering analyses in both scientific research and applied data science.