FIVA: Federated Inverse Variance Averaging for Universal CT Segmentation with Uncertainty Estimation

2025-08-08

AI Summary

Privacy-preserving semantic segmentation of multi-center, heterogeneous abdominal CT data remains challenging due to data silos and distributional shifts across institutions. Method: This paper proposes a Bayesian-inspired federated learning framework that online models client model weight uncertainty using inherent noise in stochastic mini-batch gradient descent; estimates client contribution reliability from this uncertainty; employs inverse-variance-weighted aggregation on the server to enhance global model robustness; and propagates weight uncertainty into the prediction stage to yield interpretable, pixel-wise segmentation confidence scores. Results: Extensive experiments demonstrate that our method significantly outperforms existing federated segmentation baselines on multi-center abdominal CT datasets, achieving consistent improvements in segmentation accuracy, cross-domain generalization, and uncertainty calibration quality—thereby delivering more reliable and interpretable AI support for clinical decision-making.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Computer Vision: Segmentation
• Natural Language Processing: Safety and Robustness

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Different CT segmentation datasets are typically obtained from different scanners under different capture settings and often provide segmentation labels for a limited and often disjoint set of organs. Using these heterogeneous data effectively while preserving patient privacy can be challenging. This work presents a novel federated learning approach to achieve universal segmentation across diverse abdominal CT datasets by utilizing model uncertainty for aggregation and predictive uncertainty for inference. Our approach leverages the inherent noise in stochastic mini-batch gradient descent to estimate a distribution over the model weights to provide an on-the-go uncertainty over the model parameters at the client level. The parameters are then aggregated at the server using the additional uncertainty information using a Bayesian-inspired inverse-variance aggregation scheme. Furthermore, the proposed method quantifies prediction uncertainty by propagating the uncertainty from the model weights, providing confidence measures essential for clinical decision-making. In line with recent work shown, predictive uncertainty is utilized in the inference stage to improve predictive performance. Experimental evaluations demonstrate the effectiveness of this approach in improving both the quality of federated aggregation and uncertainty-weighted inference compared to previously established baselines. The code for this work is made available at: https://github.com/asimukaye/fiva

Problem

• Universal segmentation across diverse CT datasets
• Handling heterogeneous data while preserving privacy
• Improving federated aggregation with uncertainty estimation

Innovation

• Federated learning with uncertainty-based aggregation
• Bayesian inverse-variance server parameter aggregation
• Model weight uncertainty for prediction confidence