[Paper Review] On-the-Fly Test-time Adaptation for Medical Image Segmentation

Jeya Maria Jose Valanarasu, Guo, Pengfei+2 more|arXiv (Cornell University)|Mar 10, 2022

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

This paper proposes On-the-Fly Test-Time Adaptation, a zero-shot, episodic framework for medical image segmentation that adapts a pre-trained model to a single test image without back-propagation. Using Adaptive UNet with domain-code-driven adaptive batch normalization, it achieves state-of-the-art performance on 2D and 3D domain shifts, outperforming existing test-time adaptation methods despite tighter clinical constraints.

ABSTRACT

One major problem in deep learning-based solutions for medical imaging is the drop in performance when a model is tested on a data distribution different from the one that it is trained on. 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 pre-trained encoder trained on a large corpus of 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. 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. Code is available at https://github.com/jeya-maria-jose/On-The-Fly-Adaptation

Motivation & Objective

• To address the performance drop in medical image segmentation when models face domain shifts between training and test data.
• To develop a clinically deployable adaptation method that avoids back-propagation and does not require access to the full target distribution at test time.
• To enable zero-shot, on-the-fly adaptation per individual image, aligning with real-world clinical constraints like privacy and limited compute.
• To design a framework that dynamically adapts model features using a domain code derived from a pre-trained encoder.
• To validate the method on diverse 2D and 3D medical image segmentation benchmarks under significant domain shifts.

Proposed method

• Proposes a new adaptation setting: On-the-Fly Test-Time Adaptation, where adaptation is episodic and zero-shot, with no model updates during inference.
• Introduces Adaptive UNet, which enhances standard UNet with adaptive batch normalization layers that adapt features using a domain code.
• Generates a domain code for each test image using a pre-trained encoder on a large corpus of medical images, capturing modality and scanner-specific characteristics.
• Uses the domain code to modulate the running statistics of batch normalization layers in both encoder and decoder, enabling instance-specific feature adaptation.
• Employs a domain prior from a pre-trained model to guide adaptation without requiring source data or test-time training.
• Applies the adapted model to a single image at a time, resetting the adaptation state after each inference, ensuring privacy and low resource use.

Experimental results

Research questions

• RQ1 Can a test-time adaptation method achieve high performance without any back-propagation during inference?
• RQ2 How effective is zero-shot, on-the-fly adaptation per image when the full target distribution is unavailable?
• RQ3 Can adaptive batch normalization with a learned domain code effectively handle domain shifts in 2D and 3D medical imaging?
• RQ4 Does the proposed method outperform existing test-time adaptation baselines under stricter clinical deployment constraints?
• RQ5 How does the method generalize across different modality shifts, such as between MRI sequences or fundus imaging devices?

Key findings

• On-the-Fly Adaptation with Adaptive UNet achieves a Dice score of 63.14 on 2D fundus image shifts, outperforming TENT (57.10) and RN+CR (59.15).
• For 3D MRI segmentation, the method achieves 60.66/58.73/39.30 Dice on T1→T1ce shift, surpassing TENT (55.47/52.89/39.17) and RN+CR (57.30/54.68/40.01).
• The method shows consistent improvement over direct testing and all baselines across 9 domain shifts in both 2D and 3D settings.
• Qualitative results show that Adaptive UNet produces segmentation masks closer to ground truth and oracle predictions, with reduced over-segmentation and misclassification.
• The method performs better in 2D than 3D shifts, likely due to the greater complexity of cross-modality MRI sequence shifts, yet still achieves strong performance.
• The ablation confirms that domain code-based adaptation is critical, as removing it leads to performance degradation similar to non-adapted models.