Mixture of calibrated networks for domain generalization in brain tumor segmentation

Abstract

Recent advances in deep learning for brain tumor segmentation demonstrate good performance when the training data and test data share the same distribution. However, medical images are often faced with distribution shifts due to different grades of cancers, various imaging qualities, and data from different medical institutions. Model calibration is highly related to domain generalization ability, and well-calibrated deep models show good domain generalization ability. Besides, ensemble learning is an implicit method for model calibration where domain shift problems can be alleviated by combining multiple models. In this paper, we aim to improve the generalization ability by explicitly calibrating the model in an ensemble. We proposed Mixture of Calibrated Networks (MCN) where multiple networks are jointly learned on the source and the augmented domain. We introduce a temperature scale to each network in an ensemble, which can be seen as the uncertainty of each domain to calibrate the predicted probabilities. These temperature scales form an uncertainty-aware loss to adaptively weigh the losses of multiple networks. The expectation–maximization (EM) algorithm is used to learn the parameters and explicitly model the relationship between key parameters, achieving better interpretability and robust parameter estimation. Moreover, we introduce an orthogonal constraint between convolutional kernels from the corresponding layer of multiple sub-networks, which keeps the diversity of sub-networks. We conduct extensive experiments on the brain tumor segmentation dataset from BRATS 2018, BRATS 2019, BRATS 2020, and FeTS 2021. Our proposed approach consistently improves the generalization performance and shows better calibration over inter-domain methods and intra-domain methods.