Inference for Clustering: Conformal Sets for Cluster Labels

1. Problem Formulation

Clustering is a fundamental unsupervised learning task, yet it lacks rigorous methods for quantifying uncertainty in cluster assignments. While existing literature addresses uncertainty in cluster existence (e.g., SigClust) or centroid locations, there is no broadly applicable framework for uncertainty in cluster labels themselves.

The core challenge is that cluster labels are:

• Unobserved: They are estimated from data rather than given.
• Discrete and Permutation-Invariant: They cannot be treated as smooth parametric objects.
• Non-Exchangeable: Standard conformal inference relies on the exchangeability of data points. However, when labels are estimated from the same data used for calibration, the exchangeability assumption is violated, leading to severe under-coverage (the true label falls outside the confidence set more often than the nominal rate).

Goal: Construct a confidence set $\hat{C}(x) \subseteq \{1, \dots, K\}$ for any query point $x$ such that, with probability at least $1-\alpha$, the set contains the true (but unknown) cluster label $Y^*_{n+1}$, even when labels are estimated by an arbitrary clustering algorithm.

2. Methodology

The authors propose a novel framework called Split Conformal Clustering with Stochastic Labels.

2.1 Failure of Naive Approaches

The paper first analyzes two naive approaches that fail to provide valid guarantees:

1. Cutoff for Generalizable Soft Clustering: Using soft labels (posterior probabilities) directly and selecting labels until cumulative probability exceeds $1-\alpha$. This fails because soft labels are often poorly calibrated, leading to either under-coverage (small samples) or uninformative over-coverage (large sets).
2. Split Conformal Clustering (Deterministic): Splitting data into training/calibration, clustering both with a deterministic algorithm (e.g., K-means), and treating the resulting hard labels as ground truth. This fails because deterministic labels do not capture the inherent uncertainty of the clustering process, violating the exchangeability condition required for conformal inference.

2.2 Proposed Algorithm: Split Conformal Clustering with Stochastic Labels

The proposed method (Algorithm 1) introduces stochasticity into the labeling process to restore exchangeability and capture uncertainty.

1. Data Splitting: Randomly split the unlabeled data $X_1, \dots, X_n$ into a training set ($I_{tr}$) and a calibration set ($I_{ca}$).
2. Stochastic Labeling (Training): Apply a stochastic clustering algorithm $\gamma$ to the training data. Instead of assigning a single hard label, the algorithm samples labels $Y_i$ from a probability distribution (e.g., sampling from the posterior of a Gaussian Mixture Model or Fuzzy C-Means).
3. Soft Classifier Training: Train a soft classifier $\hat{\pi}$ (e.g., Random Forest, SVM with probability outputs) on the training data $(X_i, Y_i)$ to predict soft labels (probability vectors) for any input.
4. Stochastic Labeling (Calibration): Apply the same stochastic clustering algorithm to the calibration data to obtain sampled labels $Y_i$ for $i \in I_{ca}$.
5. Label Alignment: Since clustering is permutation-invariant, the labels in the calibration set may differ from the classifier's output by a permutation. Find the optimal permutation $\hat{\sigma}$ that minimizes disagreement between the classifier's predictions and the sampled calibration labels.
6. Conformal Score Calculation: Compute conformity scores for the calibration set using the aligned labels and the soft classifier.
7. Confidence Set Construction: For a new point $x$, construct the confidence set $\hat{C}(x)$ containing all labels whose conformity scores fall below the empirical quantile of the calibration scores.

3. Theoretical Guarantees

The paper establishes finite-sample and asymptotic coverage guarantees.

3.1 Finite-Sample Lower Bound

Theorem 1 provides a lower bound on the marginal coverage probability:
$$P(\hat{\sigma}^*_o(Y^*_{n+1}) \in \hat{C}(X_{n+1})) \ge 1 - \alpha - \frac{n}{n+2}E_{n/2}(\gamma) - \frac{n}{2(n+2)}S_{n/2}(\gamma)$$
where:

• Consistency ($E_n$): Measures the estimation error between the soft labels produced by the algorithm and the true conditional label distribution.
• Replace-One Stability ($S_n$): Measures how much the joint distribution of labels changes when a single data point is replaced.

The bound reveals that under-coverage is controlled by the consistency and stability of the underlying clustering algorithm. If these terms are small, the method achieves the target $1-\alpha$ coverage.

3.2 Asymptotic Coverage

Theorem 2 and Theorem 3 show that for parametric mixture models (e.g., Gaussian Mixture Models) where parameters are estimated consistently (e.g., via MLE/EM), the consistency and stability conditions hold asymptotically. Consequently, the coverage converges to $1-\alpha$ as $n \to \infty$.

4. Key Results

4.1 Simulations

The authors tested the method on Gaussian Mixture Models (GMM) and Gamma Mixture Models (GaMM) in both low and high dimensions.

• Coverage: The proposed method (Split CC Stochastic) consistently achieved the target 90% coverage across various variance levels and sample sizes.
• Comparison:
  • Naive Split Conformal (Hard Labels): Suffered from severe under-coverage (often < 70%) because hard labels lack stability.
  • Cutoff Method: Was overly conservative, producing uninformative confidence sets (large set sizes) to force coverage.
• Set Size: The proposed method produced informative confidence sets. In well-separated regions, sets were singletons; in boundary regions, sets contained multiple plausible labels, accurately reflecting uncertainty.

4.2 Real-World Application: Single-Cell RNA-seq

The method was applied to the PBMC-3K dataset to quantify uncertainty in cell type identification.

• Findings:
  • High Certainty: B cells and NK cells formed tight, well-separated clusters, resulting in singleton confidence sets (high certainty).
  • High Uncertainty: T cells (CD4/CD8) and monocytes/dendritic cells showed larger confidence sets (size $\ge 2$), correctly reflecting their biological ambiguity and overlapping transcriptional states.
• Visualization: The method successfully visualized regions of uncertainty in PCA space, highlighting boundaries where cell type assignment is ambiguous.

5. Key Contributions

1. First Model-Agnostic Framework: Provides the first distribution-free, finite-sample inference framework specifically for cluster labels in unsupervised learning.
2. Stochastic Clustering Innovation: Identifies that deterministic clustering breaks conformal exchangeability and proposes stochastic labeling (sampling from soft distributions) as the solution to restore validity.
3. Theoretical Characterization: Derives explicit bounds linking coverage failure to the consistency and stability of the clustering algorithm, offering a roadmap for future algorithm design.
4. Practical Utility: Demonstrates that the method produces interpretable confidence sets that distinguish between "reliable" assignments and "ambiguous" boundaries, validated on both synthetic and real biological data.

6. Significance

This work bridges the gap between unsupervised learning and rigorous statistical inference. By quantifying uncertainty in cluster labels, it addresses a critical limitation in fields like biomedicine (e.g., cell type discovery) and social network analysis, where clustering decisions often drive downstream scientific conclusions. The framework allows practitioners to identify which data points can be confidently assigned to a cluster and which remain inherently ambiguous, thereby improving the reliability and interpretability of unsupervised learning results.