MDCP

This repository contains the implementation for "Multi-Distribution Robust Conformal Prediction".

Table of Contents

• Installation
• Simulation experiments
• Real-data applications
• Figures
• Repository layout

Installation

Create and activate a virtual environment.

Windows (PowerShell):

python -m venv mdcp
.\mdcp\Scripts\Activate.ps1

macOS/Linux:

python -m venv mdcp
source mdcp/bin/activate

Install the dependencies:

pip install --upgrade pip
pip install -r requirements.txt

Simulation experiments

Simulation scripts live under notebook/ and write artifacts to eval_out/; see model/const.py for the exact directories. The paper reports averages over 100 independent runs per configuration, using contiguous seed ranges starting from the base seed.

Run the main simulation suites:

# Linear suite
python notebook/eval_linear.py --num-trials 1 --base-seed 34567

# Nonlinear suite
python notebook/eval_nonlinear.py --num-trials 1 --base-seed 23456

# Temperature suite
python notebook/eval_temperature.py --num-trials 1 --base-seed 12345 --temperatures 0.5 1.5 2.5 3.5 4.5

# Covariate shift suite
python notebook/eval_cov_shift.py --num-trials 1 --base-seed 34567 --delta-xs 0 0.5 1.5 2.5 3.5 4.5

# Covariate and concept shift suite
python notebook/eval_cov_and_concept_shift.py --num-trials 1 --base-seed 45678 --delta-xs 0 0.5 1.5 2.5 3.5 4.5

The linear, nonlinear, and temperature suites evaluate the aggregate and source-specific baselines, MDCP, and an Oracle that uses the simulator's true conditional densities.

The appendix scaling studies vary the number of sources, source-size balance, and total sample size. Each command runs all three sweeps:

# Linear scaling
python notebook/eval_linear_scaling.py --num-trials 1 --base-seed 34567

# Nonlinear scaling: linear, interaction, sinusoid, and softplus settings
python notebook/eval_nonlinear_scaling.py --num-trials 1 --base-seed 23456

# Temperature scaling
python notebook/eval_temperature_scaling.py --num-trials 1 --base-seed 12345

Scaling studies use gamma = 0. Tuned variants are available by combining --include-tuned-methods with a gamma grid.

Real-data applications

Three real-world evaluations are supported: FMoW, PovertyMap from WILDS, and MEPS. The paper reports averages over 100 runs per dataset, using contiguous seed ranges starting from the base seed.

Prerequisites

1. Clone the upstream WILDS repository, which provides the examples/ modules used by PovertyMap and FMoW:

   python -m wilds.setup_wilds_repo

   The helper places the checkout at external/wilds_upstream. Use --wilds-repo if it is stored elsewhere.

2. Download the WILDS datasets:

   python -m wilds.download_datasets --dataset poverty --root data --unpack
   python -m wilds.download_datasets --dataset fmow --root data --unpack

   If TLS inspection blocks the download, run python wilds/download_wilds_insecure.py --dataset <name> --root data --unpack and move the extracted folders under data/.

3. Prepare the MEPS CSV files by following the download guide, then place the three cleaned meps_*_reg.csv files in meps/data/.

MEPS

python -m meps.eval --panels 19 20 21 --num-trials 1 --base-seed 45678

Outputs are written to eval_out/meps/ by default.

PovertyMap

# 1. Create the 2014-2016 split
python wilds/poverty/create_training_split.py --repo-root . --train-frac 0.375 --seed 0 --years 2014 2015 2016

# 2. Train the ResNet18 multi-head regressor
python wilds/poverty/train_resnet.py --repo-root .

# 3. Predict on the holdout set
python wilds/poverty/predict_density.py --repo-root . --checkpoint eval_out/poverty/training/run_multihead_gaussian/best_model.pth

# 4. Learn lambda, calibrate, and evaluate
python wilds/poverty/analysis/eval.py --prediction-dir eval_out/poverty/predictions/run_multihead_gaussian --num-trials 1 --base-seed 90000

# 5. Reproduce the paper figures
python wilds/poverty/analysis/plot.py

FMoW

# 1. Create the 2016-only splits
python wilds/fmow/pipeline/cli.py split --root data/fmow_v1.1 --wilds-repo external/wilds_upstream --output eval_out/fmow/splits/2016_region --train-frac 0.375 --seed 0

# 2. Train DenseNet121 on the training set
python wilds/fmow/pipeline/cli.py train --root data/fmow_v1.1 --wilds-repo external/wilds_upstream --train-idx eval_out/fmow/splits/2016_region/train_idx.npy --holdout-idx eval_out/fmow/splits/2016_region/holdout_idx.npy --output eval_out/fmow/training/run_densenet121 --epochs 30 --batch-size 32

# 3. Predict on the holdout set
python wilds/fmow/pipeline/cli.py predict --root data/fmow_v1.1 --wilds-repo external/wilds_upstream --holdout-idx eval_out/fmow/splits/2016_region/holdout_idx.npy --checkpoint eval_out/fmow/training/run_densenet121/checkpoint_best.pt --output eval_out/fmow/predictions/run_densenet121 --batch-size 128 --save-embeddings

# 4. Learn lambda, calibrate, and evaluate
python wilds/fmow/analysis/eval.py --prediction-dir eval_out/fmow/predictions/run_densenet121 --num-trials 1 --base-seed 70000

# 5. Reproduce the paper figures
python wilds/fmow/analysis/plot.py

Figures

Run these commands after the corresponding evaluation artifacts exist. Outputs are saved under eval_out/paper_figures/.

Main simulation figures:

python notebook/paper_figures/plot_linear.py --eval-dir eval_out/linear --output-dir eval_out/paper_figures
python notebook/paper_figures/plot_nonlinear.py --eval-dir eval_out/nonlinear --output-dir eval_out/paper_figures
python notebook/paper_figures/plot_temperature.py --eval-root eval_out/temperature --output-dir eval_out/paper_figures
python notebook/paper_figures/plot_cov_shift.py --eval-root eval_out/cov_shift --output-dir eval_out/paper_figures
python notebook/paper_figures/plot_cov_and_concept_shift.py --eval-root eval_out/cov_and_concept_shift --output-dir eval_out/paper_figures

Scaling figures and appendix tables:

python notebook/paper_figures/plot_linear_scaling.py
python notebook/paper_figures/plot_nonlinear_scaling.py
python notebook/paper_figures/plot_temperature_scaling.py
python notebook/paper_figures/export_scaling_tables.py

Real-data figures:

# PovertyMap
python notebook/paper_figures/plot_poverty.py --input-dir eval_out/poverty/mdcp --output-dir eval_out/paper_figures

# PovertyMap target distribution
python notebook/paper_figures/plot_poverty_target.py --repo-root . --output eval_out/paper_figures/poverty_target.pdf

# FMoW
python notebook/paper_figures/plot_fmow.py --input-dir eval_out/fmow/mdcp --output-dir eval_out/paper_figures

# MEPS
python notebook/paper_figures/plot_meps.py --input eval_out/meps --output eval_out/paper_figures --coverage-target 0.9

Repository layout

MDCP/
|-- model/
|   |-- MDCP.py
|   `-- const.py
|-- notebook/
|   |-- eval_linear.py
|   |-- eval_nonlinear.py
|   |-- eval_temperature.py
|   |-- eval_cov_shift.py
|   |-- eval_cov_and_concept_shift.py
|   |-- eval_linear_scaling.py
|   |-- eval_nonlinear_scaling.py
|   |-- eval_temperature_scaling.py
|   `-- paper_figures/
|-- meps/
|   `-- eval.py
|-- wilds/
|   |-- poverty/
|   `-- fmow/
|-- data/       # downloaded datasets
`-- eval_out/   # experiment outputs and figures

Citation

If you find our method or package useful, please cite:

@misc{yang2026multi,
      title={Multi-Distribution Robust Conformal Prediction},
      author={Yuqi Yang and Ying Jin},
      year={2026},
      eprint={2601.02998},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2601.02998},
}