PaperlayerView exact revision on arXiv

Source-linked AI summary

Distribution-Free Federated Learning with Conformal Predictions

Charles Lu, Jayasheree Kalpathy-Cramer

arXiv:2110.07661v2cs.LGeess.IV

TL;DR

Healthcare federated models need privacy-preserving collaboration without sacrificing calibrated, interpretable uncertainty. This paper combines federated learning with adaptive conformal prediction, finding tighter prediction sets than local conformal calibration across six MedMNIST datasets while relating set size to class entropy.

  • Problem

    Privacy-preserving federated learning still needs reliable uncertainty and calibration for high-stakes clinical decision-making.

    Source: Introduction · p. 1

  • Method

    The method combines federated learning with adaptive conformal calibration to produce distribution-free prediction sets without modifying model training.

    Source: Introduction · p. 1, Methods · p. 2, Conclusion · p. 5

  • Results

    Federated conformal prediction yields lower cardinality than local conformal prediction across MedMNIST datasets, including under 30% validation-label noise.

    Source: Discussion · p. 4

  • Takeaways & Limitations

    Conformal set size corresponds intuitively with prediction difficulty as measured by predictive class entropy.

    Source: Experiments · p. 3, Discussion · p. 5, Conclusion · p. 5

    Coverage guarantees require exchangeable data, and the entropy relationship is evaluated using an auxiliary classifier and dataset-specific assumptions.

    Source: Methods · p. 2, Experiments · p. 3

Abstract

from arXiv · show

Federated learning has attracted considerable interest for collaborative machine learning in healthcare to leverage separate institutional datasets while maintaining patient privacy. However, additional challenges such as poor calibration and lack of interpretability may also hamper widespread deployment of federated models into clinical practice, leading to user distrust or misuse of ML tools in high-stakes clinical decision-making. In this paper, we propose to address these challenges by incorporating an adaptive conformal framework into federated learning to ensure distribution-free prediction sets that provide coverage guarantees. Importantly, these uncertainty estimates can be obtained without requiring any additional modifications to the model. Empirical results on the MedMNIST medical imaging benchmark demonstrate our federated method provides tighter coverage over local conformal predictions on 6 different medical imaging datasets for 2D and 3D multi-class classification tasks. Furthermore, we correlate class entropy with prediction set size to assess task uncertainty.

Introduction

Federated learning supports collaborative medical imaging while preserving institutional data privacy, but calibration, uncertainty estimation, and interpretability remain barriers to clinical use. The paper combines federated learning with adaptive conformal prediction to provide distribution-free prediction sets and studies their coverage, efficiency, and relationship to task uncertainty.

  • Limited labeled cases and privacy risks motivate federated learning for collaborative medical imaging model development.Source: Introduction · p. 1
  • Poor calibration and limited interpretability in deep CNNs can undermine trust during high-stakes medical decision-making.Source: Introduction · p. 1
  • Conformal prediction provides distribution-free uncertainty through prediction sets that align with differential-diagnosis-style clinical reasoning.Source: Introduction · p. 1
  • The paper investigates adaptive versus non-adaptive sets, federated versus local quantile calibration, and the relationship between set size and class entropy.Source: Introduction · p. 1

Related Work

Prior work establishes conformal prediction as a distribution-free uncertainty framework, while alternative CNN confidence methods impose modeling assumptions, require model-specific designs, or modify training. Federated conformal prediction remains insufficiently studied for CNN-based medical imaging.

  • Conformal inference provides marginal coverage without specifying a prior distribution and has been adapted to several prediction tasks.Source: Related Work · p. 1
  • Bayesian, model-specific, and loss-modification approaches provide alternative CNN uncertainty estimates but involve stronger assumptions or training changes.Source: Related Work · p. 2
  • CNN Softmax outputs often require post-processing calibration techniques such as temperature scaling or isotonic calibration.Source: Related Work · p. 2
  • Federated learning research has addressed Bayesian uncertainty and non-IID distributions, but distributed conformal prediction has not targeted federated CNN medical imaging.Source: Related Work · p. 2

Methods

The method calibrates Softmax-based conformal prediction sets using held-out data and adapts split conformal prediction to federated averaging. Its coverage guarantee relies on exchangeability, while federated quantiles combine calibration information across institutions without changing model training.

Methods

  • Naive Softmax sets may omit the coverage guarantee, whereas held-out calibration produces formally valid conformal predictions.Source: Methods · p. 2
  • A conformal set includes classes whose monotonic scores exceed 1 − q̂, where q̂ is estimated from a calibration set.Source: Methods · p. 2
  • The federated algorithm aggregates quantile estimates from multiple validation sets and returns a prediction set using the averaged threshold.Source: Methods · p. 2
  • The marginal coverage guarantee assumes exchangeable data and uses a finite-sample correction that can become near-tight through symmetry.Source: Methods · p. 2
  • FedAvg is used as a simple proof-of-concept training procedure, although the conformal framework is not tied to a specific federated learning method.Source: Methods · p. 2

Experiments

Experiments compare naive, local-conformal, and federated-conformal prediction sets across six MedMNIST datasets using shared trained models and institution-partitioned calibration data. They evaluate coverage, set cardinality, robustness to calibration-label noise, and the association between entropy and set size.

Experiments

  • Naive prediction sets are compared with local and federated conformal sets using coverage and average set cardinality.Source: Experiments · pp. 2–3
  • Six MedMNIST datasets are evaluated with ResNet-18 models trained from five random initializations across 2D and 3D medical imaging tasks.Source: Experiments · p. 3, Figure 2 · p. 3
  • The same trained models support all methods; local calibration uses one institution, while federated calibration averages quantiles across institutions and includes a 30% label-shuffling stress test.Source: Experiments · p. 3
  • Figure 1 illustrates 90%-confidence prediction sets, where larger sets indicate greater uncertainty and smaller sets indicate greater confidence.Source: Figure 1 · p. 3, Experiments · p. 3
  • Class distributions vary across the six MedMNIST datasets, providing heterogeneous benchmark settings for comparing prediction-set behavior.Source: Figure 2 · p. 3, Experiments · p. 3
  • Entropy is used as a model-uncertainty indicator and is expected to increase with conformal prediction-set size.Source: Experiments · p. 3

Discussion

Across six MedMNIST tasks, conformal prediction produced tighter, coverage-guaranteed sets than the naive baseline, while federated calibration generally reduced set size relative to local calibration. Set cardinality also increased with class entropy, including under moderate calibration-label corruption.

Discussion

  • Federated conformal prediction sets had lower cardinality than local sets across MedMNIST at most confidence levels, including with 30% calibration-label noise.The comparison is shown in Figure 4 and summarized across datasets in the discussion.Source: Discussion · p. 4, Figure 4 · p. 4
  • At 90% coverage, federated and local conformal methods produced much tighter sets than the naive baseline while satisfying coverage properties.The naive baseline was exceptionally conservative, often outputting all classes and reaching 100% coverage without useful discrimination.Source: Discussion · pp. 3–4, Table 1 · p. 4
  • Federated conformal sets achieved average sizes of 2.3 ± 0.1 on BloodMNIST and 2 ± 0.1 on PathMNIST, with lower sizes than local conformal sets at lower α thresholds.These datasets had mostly balanced class distributions.Source: Discussion · p. 4
  • On RetinaMNIST and OrganMNIST3D, federated and local set sizes were similar, while the smaller datasets showed larger variance in coverage and set size.The paper attributes these differences to the datasets having 10–200 times fewer cases than the others.Source: Discussion · p. 4, Figure 5 · p. 4
  • Prediction-set cardinality had a positive relationship with epistemic model uncertainty measured by class entropy, and conformal uncertainty remained robust to moderate calibration-label corruption.The entropy analysis used three percentiles across four institutions, including institutions trained with 30% random labels.Source: Discussion · pp. 4–5, Figure 6 · p. 4

Conclusion

The paper combines federated learning with conformal prediction to provide distribution-free uncertainty estimates for collaborative healthcare modeling. Evaluation on MedMNIST demonstrates feasibility, links set size with predictive class entropy, and requires no modification to model training.

Conclusion

  • The proposed method combines federated learning with conformal prediction sets and is evaluated on MedMNIST medical imaging benchmarks.The evaluation demonstrates feasibility for conformal uncertainty in healthcare.Source: Conclusion · p. 5
  • The method links conformal uncertainty with prediction-task difficulty through predictive class entropy.It provides a simple, flexible uncertainty-quantification technique without directly modifying model training.Source: Conclusion · p. 5
Loading 2110.07661v2…