Conformal Prediction for Generative Models via Adaptive Cluster-Based Density Estimation

TL;DR

This work extends conformal prediction to conditional generative models by reframing density estimation on model-generated samples as the basis for nonconformity scores. CP4Gen uses a $K$-component Gaussian mixture fitted via $K$-means to capture multi-modal ensemble distributions, reducing prediction-set volume and structural complexity while preserving marginal coverage at $1-\alpha$. Compared with PCP, CP4Gen yields sharper, more interpretable prediction sets across synthetic, real-world, and climate-emulation tasks, with notable reductions in complexity (often >90%) and robust performance as response dimensionality grows. The approach is compatible with any conditional generator and offers a practical, scalable tool for calibrated uncertainty quantification in high-stakes applications, with clear paths for extensions such as online calibration and localization.

Abstract

Conditional generative models map input variables to complex, high-dimensional distributions, enabling realistic sample generation in a diverse set of domains. A critical challenge with these models is the absence of calibrated uncertainty, which undermines trust in individual outputs for high-stakes applications. To address this issue, we propose a systematic conformal prediction approach tailored to conditional generative models, leveraging density estimation on model-generated samples. We introduce a novel method called CP4Gen, which utilizes clustering-based density estimation to construct prediction sets that are less sensitive to outliers, more interpretable, and of lower structural complexity than existing methods. Extensive experiments on synthetic datasets and real-world applications, including climate emulation tasks, demonstrate that CP4Gen consistently achieves superior performance in terms of prediction set volume and structural simplicity. Our approach offers practitioners a powerful tool for uncertainty estimation associated with conditional generative models, particularly in scenarios demanding rigorous and interpretable prediction sets.

Figures (5)

• Figure 1: The two phases of CP4Gen: calibration and inference. (a) $K$-means is applied to prediction ensemble members. Each identified cluster is treated as one mode of a Gaussian Mixture Model, whose sample mean $\boldsymbol{\mu}_k$ and covariance matrix $\boldsymbol{\Sigma}_k$ are calculated. Nonconformity score $s_{i}$ is defined as the minimal $s_{i}^{k}$ between observation $Y_{i}$ and each cluster. Nonconformity score quantile $Q_{1-\alpha}$ is then computed on the calibration set. (b) Given a new set of covariates, ensemble members are sampled and $K$-means is applied on top. Prediction set is obtained by inverting nonconformity score function with $Q_{1-\alpha}$ resulting as a union of $K$ ellipsoids.
• Figure 2: A comprehensive performance summary of CP methods. CP methods ($\alpha=0.1$) are extensively evaluated on various datasets including synthetic, real-world, and precipitation emulation tasks. The performance of PCP and CP4Gen are compared in terms of prediction set coverage rate (a), structural complexity (b), and volume (c). The results of structural complexity and volume are normalized to $[0, 1]$ for visualization convenience. Results in original scale are reported in the appendix Tables \ref{['tab:results:Synthetic']}, \ref{['tab:results:Real']}, and \ref{['tab:results:Precip']}. Both CP methods achieve desired coverage rate close to $1-\alpha=0.9$ on all datasets. CP4Gen consistently outperforms PCP in terms of both prediction set volume and structural complexity. In particular, the CP4Gen produced prediction sets' structural complexity is significantly lower than PCP's, reducing complexity by more than $90\%$ on most datasets.
• Figure 3: Dimensionality Reduction and Prediction Set Demonstration on (a) Precip-2 and (b) Precip-2-N. (a$_{1}$) Precip-2 is obtained by taking average of precipitation values over north (blue shaded) and south (red shaded) hemispheres. (b$_{1}$) Precip-2-N is by taking precipitation values at Miami and Key West. (a$_{\text{2-3}}$ and b$_{\text{2-3}}$) On both datasets, CP4Gen gives a smaller and simpler prediction set. The volume advantage is more evident when response two dimensions are correlated.
• Figure 4: Prediction set ($\alpha=0.1$) visualization on synthetic datasets: S-Curve and 25-Gaussians. Results for (a) S-Curve and (b) 25-Gaussians. Prediction sets by PCP and CP4Gen are visualized for 100 test data samples at the first two columns, (a$_{1}$, b$_{1}$) for PCP and (a$_{2}$, b$_{2}$) for CP4Gen. In addition, the last two columns show the histogram of the number of disjoint intervals constructing each prediction set (a$_{3}$, b$_{3}$) and the histogram of the volume of each final prediction set (a$_{4}$, b$_{4}$). Clearly, CP4Gen produces prediction sets with lower structural complexity (i.e., fewer disjoint intervals) and smaller volume than PCP.
• Figure 5: A study on the effect of prediction ensemble size. Results for (a) 25-Gaussians and (b) Bio. From left to right, prediction set coverage rate (a$_{1}$, b$_{1}$), volume (a$_{2}$, b$_{2}$), and structural complexity (a$_{3}$, b$_{3}$) are plotted separately each as a function of prediction ensemble size. The prediction ensemble size is varied on a grid of $[10, 20, 30, \dots, 100]$. Both CP methods' coverage rate stays flat close to $90\%$ ($\alpha=0.1$) for all ensemble sizes . As ensemble size increases, both methods' prediction set volume decreases. Interestingly, CP4Gen's set complexity converges to a constant (5 for 25-Gaussians and 2 for Bio) when ensemble size is larger than 20, while PCP's complexity scales linearly with ensemble size.

Theorems & Definitions (8)

• Proposition 3.1
• Remark 1.1
• Theorem 6.1
• Lemma 6.2
• Lemma 6.3
• proof
• Proposition 6.4
• proof