Federated Conditional Conformal Prediction via Generative Models

TL;DR

This paper tackles the challenge of uncertainty quantification under federated data shifts by introducing Federated Conditional Conformal Prediction (Fed-CCP). It leverages invertible generative mappings (normalizing flows or diffusion models) to transform client distributions to a simple Gaussian calibration space $P_{XY}$, enabling conditional, instance-specific conformal prediction sets that are then mapped back to each client's space. The method preserves privacy since only gradient information is exchanged, and it uses a client-conditioned generator to capture cross-client heterogeneity. Empirical results across healthcare, IoT, insurance, traffic, and epidemic forecasting show Fed-CCP achieves near nominal coverage with tighter, more adaptive prediction sets than baselines and ablations, demonstrating practical impact for high-stakes, privacy-conscious FL settings.

Abstract

Conformal Prediction (CP) provides distribution-free uncertainty quantification by constructing prediction sets that guarantee coverage of the true labels. This reliability makes CP valuable for high-stakes federated learning scenarios such as multi-center healthcare. However, standard CP assumes i.i.d. data, which is violated in federated settings where client distributions differ substantially. Existing federated CP methods address this by maintaining marginal coverage on each client, but such guarantees often fail to reflect input-conditional uncertainty. In this work, we propose Federated Conditional Conformal Prediction (Fed-CCP) via generative models, which aims for conditional coverage that adapts to local data heterogeneity. Fed-CCP leverages generative models, such as normalizing flows or diffusion models, to approximate conditional data distributions without requiring the sharing of raw data. This enables each client to locally calibrate conformal scores that reflect its unique uncertainty, while preserving global consistency through federated aggregation. Experiments on real datasets demonstrate that Fed-CCP achieves more adaptive prediction sets.

Figures (1)

• Figure 1: Comparison of CQR, Fed-CCP (w/o conditioner), and Fed-CCP. The top row shows marginal coverage (red dashed line: target 0.9), and the bottom row shows average prediction set size. Fed-CCP attains near-nominal coverage with smaller sets, indicating better adaptiveness and efficiency.