Investigating Test-Time Adaptation of Convolutional Neural Networks for Medical Image Analysis under Distribution Shifts

Abstract

Deep learning models often fail to generalize when classifying medical images that originate from a different data distribution than the training images. Their error rate increases, posing a severe risk to patients. Test-Time Training (TTT) with Masked Autoencoders (MAE) addresses distribution shifts by adapting a trained model to unlabeled images at inference time, using a self-supervised objective. However, we find that its success in the natural image domain does not immediately transfer to the medical domain. A key limitation is that TTT can degrade model performance when the auxiliary loss used for inference-time adaptation is uncorrelated with the main task. Thus, we tackle distribution shifts from a different angle, eliminating the need for a suitable loss function or weight adaptation during inference time. We propose WhiCo, a novel Test-Time Adaptation method that mitigates distribution shifts by aligning test features with the training feature distribution through whitening and coloring. We evaluate WhiCo on diverse medical tasks, including COVID-19 diagnosis, brain hemorrhage detection, and organ identification. WhiCo improves accuracy by 3.5 to 8.4 percentage points on COVID-19 diagnosis, 3.4 to 10.3 points on brain hemorrhage detection, and by up to 21.9 points on organ identification. Moreover, WhiCo is model-agnostic and requires no expert knowledge to choose test-time hyperparameters.