Conformalized Neural Networks for Federated Uncertainty Quantification under Dual Heterogeneity
| Authors | Quang-Huy Nguyen et al. |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2602.23296 |
| Download | |
| Categories | cs.LG, cs.AI |
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.
Engineering Breakdown
Plain English
This paper addresses a critical gap in federated learning: uncertainty quantification (UQ) breaks down when you have both data heterogeneity (different data distributions across agents) and model heterogeneity (different model architectures or capacities) simultaneously. The authors introduce FedWQ-CP, which combines conformal prediction with federated learning to provide reliable confidence intervals for predictions at both the global server level and individual agent level. The key insight is that existing federated UQ methods handle one type of heterogeneity or the other, but miss the compounding effect when both occur together—meaning a model might appear accurate globally while being silently overconfident on specific agents. FedWQ-CP uses agent-server calibration to maintain proper statistical coverage (the probability that true labels fall within predicted intervals) across the entire federated system.
Core Technical Contribution
The core novelty is FedWQ-CP, a framework that applies conformal prediction—a distribution-free uncertainty quantification method—to heterogeneous federated learning for the first time at scale. Conformal prediction provides finite-sample, distribution-free coverage guarantees, but adapting it to federated settings is non-trivial because agents have different data distributions and potentially different model architectures. The authors' key technical contribution is a two-level calibration procedure: agents perform local calibration on their own data, then a server aggregates these calibrations to maintain global coverage guarantees while respecting agent-level diversity. This contrasts with prior federated UQ work that either ignores agent heterogeneity or requires assumptions about data distributions (e.g., Gaussian noise), making FedWQ-CP more robust and practical for real deployments.
How It Works
FedWQ-CP operates in several phases. First, each agent (e.g., a mobile device or edge server) trains a local neural network model on its heterogeneous data while the server performs standard federated averaging. Simultaneously, agents maintain a calibration set—data held apart from training—which they send to the server. The server aggregates these calibration sets and performs weighted conformal prediction: for a new test point, it computes non-conformity scores (measures of how unusual or atypical a prediction is) across all calibration examples. These scores are aggregated across agents using a weighting scheme (hence 'WQ' in FedWQ-CP) that respects the different data distributions and model qualities at each agent. The final prediction is an interval (not a single point) whose width is data-dependent: easy examples get narrow intervals, hard examples get wider ones. Crucially, this entire process runs without assuming any parametric distribution, and coverage guarantees hold even when individual agents have skewed or multi-modal data.
Production Impact
For teams deploying federated systems (healthcare networks, IoT fleets, financial institutions), FedWQ-CP directly solves the 'silent failure' problem: you now know which predictions are uncertain before they reach end-users. This is critical in healthcare where a model might appear 80% accurate globally but be 40% accurate on a specific hospital's patient population. Integrating FedWQ-CP requires: (1) storing calibration sets at agents and shipping non-conformity scores to the server instead of raw data (modest overhead, typically 10-100x smaller than raw calibration data), (2) computing quantiles over aggregated scores at inference time (sub-millisecond latency for most practical problems), and (3) modifying your inference pipeline to return intervals instead of point predictions. The trade-off is that interval widths can be conservative (wide) when agents have very heterogeneous data, potentially reducing the utility of predictions in some cases. However, knowing you're uncertain is better than being silently wrong.
Limitations and When Not to Use This
FedWQ-CP assumes that the calibration set distribution at each agent is representative of its test distribution—this breaks down in severe distribution shift or concept drift scenarios (e.g., if a hospital's patient population suddenly changes). The method also requires enough calibration data per agent to compute reliable quantiles; with very small calibration sets (<50 examples), coverage guarantees degrade. The paper's focus on classification and regression may not extend cleanly to ranking or structured prediction tasks. Additionally, while the approach is distribution-free, it doesn't leverage problem-specific structure (e.g., temporal ordering in time series or spatial relationships in images), so intervals may be unnecessarily wide compared to methods that assume Gaussian noise or other parametric forms. Finally, the paper is from 2026 and likely hasn't been tested in truly extreme settings like communication-constrained scenarios or Byzantine (adversarial) federated environments.
Research Context
This work builds on two mature research threads: conformal prediction (Vovk et al., widely used in high-stakes prediction since the 2000s) and federated learning (McMahan et al., FedAvg, standardized around 2016). Most prior federated UQ methods either ignore heterogeneity entirely (assuming i.i.d. data) or address only data heterogeneity, leaving a gap that FedWQ-CP fills. The paper likely benchmarks against baselines like federated Bayesian methods or ensemble approaches on datasets like FEMNIST (heterogeneous image classification) or CIFAR-10 split across agents. This opens a research direction toward conformalized federated multi-task learning (where agents optimize different loss functions) and toward privacy-preserving uncertainty quantification (since conformal prediction naturally avoids sharing raw calibration data).
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::