CP4Gen: Conformal Prediction for Generative Models

Code for Conformal Prediction for Generative Models via Adaptive Cluster-Based Density Estimation, Transactions on Machine Learning Research (2026). OpenReview

Conditional generative models produce realistic samples but no calibrated uncertainty. CP4Gen turns an ensemble of generated samples into a prediction set with a finite-sample coverage guarantee, and does so in a way that keeps the set small and simple enough to interpret and optimise over.

The idea

Given an ensemble Y_hat = {Y_hat^1, ..., Y_hat^M} drawn from q(Y | X), CP4Gen

1. clusters the ensemble into K groups with K-means,

2. treats each cluster as one mode of a Gaussian mixture, taking its sample mean mu_k, covariance Sigma_k and weight w_k,

3. scores an observation by the negative log-density of the dominant component

   s(X, Y) = -log max_k w_k N(Y; mu_k, Sigma_k + beta^2 I),

4. and inverts that score at the calibration quantile Q, giving a prediction set that is an explicit union of K ellipsoids:

   C(X) = union_k { y : (y - mu_k)^T (Sigma_k + beta^2 I)^{-1} (y - mu_k) <= r_k }, with r_k = 2 (Q + log w_k) - d log(2 pi) - log|Sigma_k + beta^2 I|.

Using the dominant component rather than the full mixture is what makes step 4 closed-form. K is a design choice, not a fitted quantity: K = M recovers PCP (one ball per sample), and smaller K trades a little volume for a far simpler set. Any K fixed before calibration keeps the coverage guarantee.

Because PCP is the K = M corner of the same construction, this repository runs both through one code path, so coverage, volume and structural complexity are measured identically for the two methods.

Install

python -m venv venv && source venv/bin/activate
pip install -r requirements.txt

Python 3.10+. A GPU is optional — it is used only to train the generative model, and the climate datasets need no training at all.

Quick start

The climate datasets ship with pre-generated ensembles, so they run immediately from a clean checkout — no generative model to train:

python main.py --dataset precip_2 --CP_type PCP     --output_saving_path output/precip_2
python main.py --dataset precip_2 --CP_type CP4Gen  --output_saving_path output/precip_2
python summarise_results.py --results_path output/precip_2

Any other dataset trains a conditional flow-matching generator first, then caches its ensembles under --model_path so later runs reuse them:

python main.py --dataset s_curve --CP_type CP4Gen --n_epochs 20000

Seeing the method itself

notebooks/cp4gen_from_scratch.ipynb implements CP4Gen in about sixty lines of numpy, without importing anything from this repository, and runs it on the shipped climate ensembles. Start there if you want to see what the method actually does before reading the packaged version.

Repository layout

File	Contents
cp4gen.py	the method: mixture fitting, the Eq. 6 score, calibration and inversion
metrics.py	coverage, volume and structural complexity; the prediction set geometry
dataset.py	every dataset, synthetic through climate
flow_matching.py	the conditional flow-matching generator (Appendix B.2)
generative_models.py	training and sampling, wrapping the generator
main.py	run one dataset with one CP method
summarise_results.py	collect a K sweep and select K
ablations.py	the Appendix E studies
job_launcher.py	submit the experiment grid to Slurm
selected_configs.py	per-dataset settings chosen during the experiments
figures/	regenerate Figure 3
notebooks/	a from-scratch walkthrough of the method, plus climate data preparation
data/	all datasets, documented in data/README.md

Methods

--CP_type selects the method:

• CP4Gen sweeps K over a coarse grid and saves every K; use summarise_results.py to pick one.
• PCP is the K = M corner — one ball per ensemble member (Wang et al., 2022).
• HD-PCP is PCP after discarding the least confident ensemble members. Confidence is the log-density of the Gaussian noise seed each sample was integrated from; because the flow is a deterministic map out of that noise, samples seeded nearer the noise mode are the ones the model places in its high-density region. This needs scores the generator records, so it is unavailable for the pre-generated climate ensembles.

Metrics

• Coverage — fraction of test observations inside the prediction set. Expected near 1 - alpha = 0.9.
• Volume — computed exactly in 1-d by merging the component intervals, on a grid in 2-d, and by Monte Carlo in 3-d and above, where --mc_points sets the budget. --volume_method exact switches to an unbiased lower-variance estimator that samples inside each ellipsoid rather than a bounding box; it is more accurate but shifts numbers slightly against the published tables, so paper is the default.
• Structural complexity — the number of convex sets whose union forms the prediction set: M for PCP, and K for CP4Gen minus any component whose ellipsoid came out empty. It measures how interpretable the set is, and bounds the cost of optimising over it, which grows linearly in the number of convex pieces. For 1-d responses the number of disjoint intervals after merging overlaps is also reported; that is the count histogrammed in Figure 4, and it is smaller than the complexity whenever components overlap.

Reproducing the paper

This repository reproduces Table 3 and Figure 3, the climate emulation results. They were chosen because the climate ensembles come from a separate emulator and ship with the repository, so both run end to end on a laptop without training anything.

The rest of the paper's tables and figures need a trained generative model per dataset and a full sweep across all 21 datasets — a cluster-scale job. That code is here and functional (main.py, ablations.py, job_launcher.py), but running it is left to the reader rather than reproduced out of the box.

Figure 3

python figures/figure3_climate_case_study.py --combined --formats pdf png

Top row is Precip-2 (hemispheric means, weakly correlated), bottom row Precip-2-N (Key West and Miami, strongly correlated); PCP on the left, CP4Gen on the right.

The contrast is the paper's point. PCP places one fixed-radius ball on every ensemble member, so its set is a scalloped union that spends volume on isolated outliers — most visible bottom-left, where two stray samples each get their own ball. CP4Gen fits one Gaussian to the ensemble, so its set is a single ellipse that tilts with the correlation between the two cities and covers the same mass in less volume.

Volumes reproduce the published panels:

Panel	This repository	Paper
Precip-2, PCP	0.0086	0.0086
Precip-2, CP4Gen	0.0081	0.0081
Precip-2-N, PCP	6.5515	6.5523
Precip-2-N, CP4Gen	5.6292	5.6304

Table 3

for d in precip_2 precip_2n precip_3 precip_4; do
  for m in PCP CP4Gen; do
    python main.py --dataset $d --CP_type $m --output_saving_path output/$d
  done
done

Full test sets, M = 30, alpha = 0.1. Reproduced values against the published table:

Dataset	Method	Coverage	Paper	Volume	Paper	Complexity	Paper
Precip-2	PCP	0.851	0.85	0.0082	0.0082	30	30.0
Precip-2	CP4Gen	0.842	0.84	0.0072	0.0072	1	1.0
Precip-2-N	PCP	0.911	0.91	5.8275	5.8275	30	30.0
Precip-2-N	CP4Gen	0.912	0.91	5.4204	5.4204	1	1.0
Precip-3	PCP	0.923	0.92	0.0006	0.0006	30	30.0
Precip-3	CP4Gen	0.924	0.92	0.0005	0.0005	1	1.0
Precip-4	PCP	0.887	0.89	0.0003	0.0003	30	30.0
Precip-4	CP4Gen	0.879	0.88	0.0002	0.0002	1	1.0

Volumes are shown at the paper's precision; the Precip-3 and Precip-4 rows were run with --mc_points 100000 and their unrounded values are 0.00064895 / 0.00048678 and 0.00025509 / 0.00023535.

Monte Carlo budget in 3-d and 4-d

Precip-3 and Precip-4 volumes come from Monte Carlo over a bounding box rather than the 2-d grid, and the published budget is 10^6 points per test point in 3-d — enough to make a full run take hours on one core.

That budget is far more than the reported number needs. The table reports a mean over thousands of test points, so independent per-point Monte Carlo error largely cancels. Dropping to 10^5 points changes the mean volume by about 0.001% and runs roughly ten times faster:

python main.py --dataset precip_3 --CP_type PCP --mc_points 100000

Omit --mc_points to use the published budget exactly.

Both methods hold coverage near the nominal 0.9 while CP4Gen cuts volume and reduces structural complexity from 30 convex pieces to 1. Precip-2's coverage sits at 0.85 for both methods; results come from one fixed calibration/test split, and the guarantee is marginal over calibration and test draws rather than pointwise.

Notes on reproduction

• Choosing K. Section 3.4.2 selects K on the preliminary split. main.py records every metric on both the calibration and test splits, and summarise_results.py selects K on the calibration column by default, so the choice never sees the test set; --select_on_split test judges it on the test column instead. Coverage is guaranteed either way — Theorem F.1 only requires K to be fixed before calibration — so this affects reported volume and complexity, not validity.
• Single split. Results come from one fixed train/calibration/test split, so empirical coverage fluctuates by a point or so around 0.9.
• Sampling integration. The Euler sampler takes one step at each of timesteps grid points including t = 1, so it advances slightly past t = 1. Kept as-is because it is what produced the published results.

Running the full experiment grid

job_launcher.py builds the grid for a Slurm cluster. Edit ClusterConfig at the top for your site, then:

python job_launcher.py --grid climate --dry_run   # inspect the commands first
python job_launcher.py --grid synthetic
python job_launcher.py --grid real_world
python job_launcher.py --grid ensemble_size

Per-dataset training epochs are in selected_configs.py.

Ablations

python ablations.py --study k        --dataset 25-Gaussians   # E.1, sensitivity to K
python ablations.py --study ensemble --dataset 25-Gaussians   # E.3, ensemble size
python ablations.py --study em       --dataset 25-Gaussians   # E.2, K-means vs EM

These reuse the ensembles cached by a prior main.py run, so only the quantity under study varies. The EM comparison uses the identical dominant-component score as K-means, so it isolates the mixture-fitting step.

Citation

@article{yang2026cp4gen,
  title   = {Conformal Prediction for Generative Models via Adaptive Cluster-Based
             Density Estimation},
  author  = {Yang, Qidong and Zhu, Qianyu and Giezendanner, Jonathan and
             Marzouk, Youssef and Bates, Stephen and Wang, Sherrie},
  journal = {Transactions on Machine Learning Research},
  year    = {2026},
  url     = {https://openreview.net/forum?id=goxeVsh9Po}
}

License

MIT — see LICENSE.