CertainTTA: Estimating uncertainty for test-time adaptation on medical image segmentation

Abstract

Cross-site distribution shift in medical images is a major factor causing model performance degradation, significantly challenging the deployment of pre-trained semantic segmentation models for clinical adoption. In this paper, we propose a novel framework, CertainTTA, to maximally exploit a pretrained model for test time adaptation. Firstly, we leverage variational inference and innovatively construct a probabilistic source model by incorporating Gaussian priors on the network parameters of the pre-trained source model. A predictive posterior distribution is computed at test time for the target image, which is then used to estimate the uncertainty of the target prediction based on entropy measure. In the meantime, a novel adaptive score is also constructed to measure the source model uncertainty on its adaptability for a target image based on the mutual information between the target prediction and the target input. Both output uncertainty and model uncertainty are incorporated at test time, where the former is minimized against a low-frequency prompt which optimally reduces the domain shift at image level, and the latter is used to select the target prediction with the best model adaptability during the prompt optimization process. CertainTTA overcomes the weakness of existing entropy minimization methods where the latter becomes unreliable under biased target scenarios and tends to yield overconfident predictions. To the best of our knowledge, CertainTTA also serves as the first solution to trace model adaptability in a CTTA setting. We conduct TTA and CTTA experiments on three medical semantic segmentation benchmarks, achieving average $DSC$ improvements of 2.94%, 4.06%, and 3.49% under the TTA scenario over the state-of-the-art method on the OD/OC, polyp, and MRI Prostate segmentation datasets, respectively.