On-the-Fly Test-time Adaptation for Medical Image Segmentation

Adapting the source model to target data distribution at test-time is an efficient solution for the data-shift problem. Previous methods solve this by adapting the model to target distribution by using techniques like entropy minimization or regularization. In these methods, the models are still updated by back-propagation using an unsupervised loss on complete test data distribution. In real-world clinical settings, it makes more sense to adapt a model to a new test image on-the-fly and avoid model update during inference due to privacy concerns and lack of computing resource at deployment. To this end, we propose a new setting - On-the-Fly Adaptation which is zero-shot and episodic (i.e., the model is adapted to a single image at a time and also does not perform any back-propagation during test-time). To achieve this, we propose a new framework called Adaptive UNet where each convolutional block is equipped with an adaptive batch normalization layer to adapt the features with respect to a domain code. The domain code is generated using a domain prior generator specially trained on medical images. During test-time, the model takes in just the new test image and generates a domain code to adapt the features of source model according to the test data instance. We validate the performance on both 2D and 3D data distribution shifts where we get a better performance compared to previous test-time adaptation methods while not performing back-propagation during test-time.