Conformalized Neural Networks for Federated Uncertainty Quantification under Dual Heterogeneity

Quang-Huy Nguyen; Jiaqi Wang; Wei-Shinn Ku

cs.LGFeb 26, 2026

Conformalized Neural Networks for Federated Uncertainty Quantification under Dual Heterogeneity

Authors: Quang-Huy Nguyen, Jiaqi Wang, Wei-Shinn Ku

Organizations: Department of Computer Science and Software Engineering, Auburn University

Abstract

Federated learning (FL) faces challenges in uncertainty quantification (UQ). Without reliable UQ, FL systems risk deploying overconfident models at under-resourced agents, leading to silent local failures despite seemingly satisfactory global performance. Existing federated UQ approaches often address data heterogeneity or model heterogeneity in isolation, overlooking their joint effect on coverage reliability across agents. Conformal prediction is a widely used distribution-free UQ framework, yet its applications in heterogeneous FL settings remains underexplored. We provide FedWQ-CP, a simple yet effective approach that balances empirical coverage performance with efficiency at both global and agent levels under the dual heterogeneity. FedWQ-CP performs agent-server calibration in a single communication round. On each agent, conformity scores are computed on calibration data and a local quantile threshold is derived. Each agent then transmits only its quantile threshold and calibration sample size to the server. The server simply aggregates these thresholds through a weighted average to produce a global threshold. Experimental results on seven public datasets for both classification and regression demonstrate that FedWQ-CP empirically maintains agent-wise and global coverage while producing the smallest prediction sets or intervals.

Figures & tables

Figure 1: Framework of FedWQ-CP . Strong agents denote higher-capacity architectures trained more extensively, while weak agents use lower-capacity architectures with reduced predictive strength. See Appendix D for details.

Method	Classification	Regression	Federated	One-shot	No Iter. Opt.	No Het. Assump.
SplitCP	✓	✓	✗	✗	✓	✓
FCP	✓	✓	✓	✗	✓	✓
CPhet	✓	✓	✓	✗	✓	✗
DP-FedCP	✓	✗	✓	✗	✗	✗
FedCP-QQ	✓	✓	✓	✓	✓	✓
FedWQ-CP	✓	✓	✓	✓	✓	✓

Table 1: Comparison of UQ methods in FL settings. Classification and Regression indicate whether the method is directly applicable to classification and regression tasks, respectively. Federated indicates whether calibration can be performed in a decentralized manner without sharing raw calibration data. One-shot indicates whether calibration requires only a single round of communication between agents and the server. No Iter. Opt. (No Iterative Optimization) indicates that the method does not rely on solving a federated optimization problem (e.g., gradient-based quantile estimation) during calibration. No Het. Assump. (No Heterogeneity Assumption) indicates that the method does not require explicitly specifying or estimating a structural shift model (e.g., label shift or density ratios) across agents. A ✓ indicates that the method satisfies the stated property. FedCP-QQ computes the global calibration threshold via a quantile-of-quantiles scheme approximating the pooled mixture quantile, whereas FedWQ-CP aggregates local conformal quantiles using calibration-size-weighted averaging.

\widehat{q}=\sum_{k=1}^{M}\frac{n_{k}}{N}\widehat{q}_{k},\qquad N=\sum_{k=1}^{M}n_{k},

Algorithm 1 One-Shot Federated Conformal Calibration ( FedWQ-CP )

Dataset	Method	S0	S1	S2	W3	W4	W5	Avg	Runtime (s)
MNIST	DP-FedCP	$0.0610_{\scriptscriptstyle\pm 0.0131}$	$0.0668_{\scriptscriptstyle\pm 0.0109}$	$0.0675_{\scriptscriptstyle\pm 0.0146}$	$0.7251_{\scriptscriptstyle\pm 0.0319}$	$0.7214_{\scriptscriptstyle\pm 0.0317}$	$0.7191_{\scriptscriptstyle\pm 0.0258}$	$0.3957_{\scriptscriptstyle\pm 0.0033}$	$6.383_{\scriptscriptstyle\pm 0.109}$
	SplitCP	$1.0000_{\scriptscriptstyle\pm 0.0000}$	$1.0000_{\scriptscriptstyle\pm 0.0000}$	$1.0000_{\scriptscriptstyle\pm 0.0000}$	$0.9965_{\scriptscriptstyle\pm 0.0019}$	$0.9960_{\scriptscriptstyle\pm 0.0024}$	$0.9963_{\scriptscriptstyle\pm 0.0023}$	$0.9981_{\scriptscriptstyle\pm 0.0004}$	–
	FedCP-QQ	$1.0000_{\scriptscriptstyle\pm 0.0000}$	$1.0000_{\scriptscriptstyle\pm 0.0000}$	$1.0000_{\scriptscriptstyle\pm 0.0000}$	$1.0000_{\scriptscriptstyle\pm 0.0001}$	$1.0000_{\scriptscriptstyle\pm 0.0000}$	$1.0000_{\scriptscriptstyle\pm 0.0001}$	$1.0000_{\scriptscriptstyle\pm 0.0000}$	$2.887_{\scriptscriptstyle\pm 0.027}$
	FCP	$1.0000_{\scriptscriptstyle\pm 0.0000}$	$1.0000_{\scriptscriptstyle\pm 0.0000}$	$1.0000_{\scriptscriptstyle\pm 0.0000}$	$0.9999_{\scriptscriptstyle\pm 0.0001}$	$0.9999_{\scriptscriptstyle\pm 0.0001}$	$0.9999_{\scriptscriptstyle\pm 0.0001}$	$0.9999_{\scriptscriptstyle\pm 0.0001}$	$3.817_{\scriptscriptstyle\pm 0.124}$
	CPhet	$1.0000_{\scriptscriptstyle\pm 0.0000}$	$1.0000_{\scriptscriptstyle\pm 0.0000}$	$1.0000_{\scriptscriptstyle\pm 0.0000}$	$0.9967_{\scriptscriptstyle\pm 0.0017}$	$0.9961_{\scriptscriptstyle\pm 0.0025}$	$0.9966_{\scriptscriptstyle\pm 0.0024}$	$0.9982_{\scriptscriptstyle\pm 0.0004}$	$5.955_{\scriptscriptstyle\pm 0.061}$
	FedWQ-CP	$\mathbf{0.9978_{\scriptscriptstyle\pm 0.0012}}$	$\mathbf{0.9979_{\scriptscriptstyle\pm 0.0012}}$	$\mathbf{0.9976_{\scriptscriptstyle\pm 0.0011}}$	$\mathbf{0.9977_{\scriptscriptstyle\pm 0.0010}}$	$\mathbf{0.9973_{\scriptscriptstyle\pm 0.0010}}$	$\mathbf{0.9976_{\scriptscriptstyle\pm 0.0012}}$	$\mathbf{0.9977_{\scriptscriptstyle\pm 0.0009}}$	$\mathbf{2.831_{\scriptscriptstyle\pm 0.105}}$

Table 2: Empirical marginal coverage across seven benchmarks under dual heterogeneity. We report the agent-level coverage for strong ( S0–S2 ) and weak ( W3–W5 ) agents, alongside the global average ( Avg ). Results represent the median

\pm

95% CI over 10 independent runs (target error

\alpha=0.05

, partition

\text{Dir}(0.3)

). All runtimes are measured in seconds. FedWQ-CP empirically maintains coverage around the nominal level across all datasets with a one-shot runtime comparable to the most efficient baselines.

Dataset	Method	Avg Size (% $\downarrow$ )
MNIST	SplitCP	$6.88_{\scriptscriptstyle\pm 0.64}(56.4\%)$
	FedCP-QQ	$9.59_{\scriptscriptstyle\pm 0.21}(68.7\%)$
	FCP	$7.28_{\scriptscriptstyle\pm 1.56}(58.8\%)$
	CPhet	$6.90_{\scriptscriptstyle\pm 0.65}(56.5\%)$
	FedWQ-CP	$\mathbf{3.00_{\scriptscriptstyle\pm 0.24}}$
FashionMNIST	SplitCP	$4.68_{\scriptscriptstyle\pm 0.82}(25.4\%)$

Table 3: Empirical inefficiency measured by average prediction set size (classification) or interval length (regression). Smaller is better. Parentheses show percentage reduction relative to FedWQ-CP . Values are median

\pm

95% CI over 10 runs. FedWQ-CP consistently achieves lower inefficiency, producing substantially smaller prediction sets or intervals compared to existing federated UQ baselines.

Figure 2: Ablation on RetinaMNIST (denoted as FedAvgQ-CP ). Agent-level coverage for strong ( S0–S2 ) and weak ( W3–W5 ) agents. The red dashed line indicates the nominal 0.95 coverage level. Results represent the median

\pm

95% CI over 10 independent runs (target error

\alpha=0.05

, partition

\mathrm{Dir}(0.3)

). Unweighted quantile aggregation leads to systematic under-coverage on weak agents, underscoring the necessity of sample-size-aware aggregation in heterogeneous federated settings.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Dataset	Data Modality	Task (# Classes/Labels)	Image Size	Channels	# Samples	# Train / Val / Test
MNIST	Handwritten Digits	Multi-Class (10)	28 $\times$ 28	1	70,000	60,000 / – / 10,000
FashionMNIST	Apparel Images	Multi-Class (10)	28 $\times$ 28	1	70,000	60,000 / – / 10,000
CIFAR-10	Natural Images	Multi-Class (10)	32 $\times$ 32	3	60,000	50,000 / – / 10,000
DermaMNIST	Dermatoscope	Multi-Class (7)	28 $\times$ 28	3	10,015	7,007 / 1,003 / 2,005
BloodMNIST	Blood Cell Microscope	Multi-Class (8)	28 $\times$ 28	3	17,092	11,959 / 1,712 / 3,421
TissueMNIST	Kidney Cortex Microscope	Multi-Class (8)	28 $\times$ 28	1	236,386	165,466 / 23,640 / 47,280

Appendix

Table 4: Summary of datasets used in our experiments. We include standard vision benchmarks and diverse medical imaging datasets spanning classification and regression tasks.

Explore similar work

Multi-Agent Conformal Prediction with Personalized Statistical Validity
May 30, 2026Martin V. Vejling, Christophe A. N. Biscio, Adrien Mazoyer +2Weighted Conformal Prediction Conformal Prediction
Department of Electronic Systems Aalborg University, Aalborg, Denmark · Department of Mathematical Sciences Aalborg University, Aalborg, Denmark · Institut de Mathématiques de Toulouse Université de Toulouse, Toulouse, France
Quantization in Federated Learning: Methods, Challenges and Future Directions
Jun 25, 2026Farwa Ikram, Dipanwita Thakur, Antonella Guzzo +1Neural Network Quantization Federated Learning
DIMES - University of Calabria, Italy
Efficient Federated Conformal Prediction with Group-Conditional Guarantee
Mar 15, 2026Haifeng Wen, Osvaldo Simeone, Hong XingGroup-Conditional Conformal Prediction Conformal Prediction
1IoT Thrust, The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · Institute for Intelligent Networked Systems (INSI), Northeastern University London, London, UK · Department of ECE, The Hong Kong University of Science and Technology, HK SAR

May 30, 2026J/K move · Enter open · S save

Multi-Agent Conformal Prediction with Personalized Statistical Validity

Martin V. Vejling, Christophe A. N. Biscio, Adrien Mazoyer, Petar Popovski+1

Department of Electronic Systems Aalborg University, Aalborg, Denmark · Department of Mathematical Sciences Aalborg University, Aalborg, Denmark · Institut de Mathématiques de Toulouse Université de Toulouse, Toulouse, France

Weighted Conformal Prediction Conformal Prediction Privacy-Preserving ML

Open paper AskPDF arXiv

Uncertainty quantification is essential in high-stakes machine learning tasks. However, one of the principled solutions, conformal prediction, faces challenges under limited local calibration data, privacy constraints, and data heterogeneity. In multi-agent settings, existing works do not simultaneously and satisfactorily address these challenges with guarantees either limited to averages across agents or losing validity in heterogeneous settings. Hence, we propose personalized federated weighted conformal prediction (PFWCP), a framework that combines local density ratio weighting with weighted quantile aggregation to correct for heterogeneity while preserving privacy. The method yields asymptotically valid marginal and calibration-conditional coverage guarantees for each participating agent and supports protocols with one-shot communication. Theoretical analysis presents an adjustment to the coverage variance, governed by an effective sample size expression, which is necessary in the context of weighted conformal prediction, and experiments on synthetic and real datasets show improved calibration quality over state-of-the-art federated conformal baselines.

Abstract

Explore similar work

Multi-Agent Conformal Prediction with Personalized Statistical Validity

Quantization in Federated Learning: Methods, Challenges and Future Directions

Efficient Federated Conformal Prediction with Group-Conditional Guarantee