Pay Attention to the Atlas: Atlas-Guided Test-Time Adaptation Method for Robust 3D Medical Image Segmentation

TL;DR

This work addresses the problem of performance degradation in 3D medical image segmentation under domain shifts across sites and scanners, compounded by privacy and labeling constraints. It introduces AdaAtlas, a test-time adaptation framework that leverages a learned atlas as a high-level shape prior and updates only adaptation blocks to align test-time predictions with the atlas in atlas space via a loss $\mathcal{L}_{atlas}$, computed after a fixed atlas-registration network deforms predictions to the atlas. A key innovation is incorporating dual attention blocks for more flexible, test-time adaptation, enabling channel- and spatial-wise feature recalibration. Experiments on multi-site prostate MRI demonstrate that AdaAtlas-Attention significantly outperforms Baseline and other TTA methods (≈+0.16–0.21 in mean Dice), highlighting robustness to large distribution shifts and the practical value of atlas-guided, attention-enabled TTA for cross-site medical image segmentation. Pure atlas supervision also reduces reliance on ground-truth labels in the target domain, with potential for extension to other organs and imaging modalities.

Abstract

Convolutional neural networks (CNNs) often suffer from poor performance when tested on target data that differs from the training (source) data distribution, particularly in medical imaging applications where variations in imaging protocols across different clinical sites and scanners lead to different imaging appearances. However, re-accessing source training data for unsupervised domain adaptation or labeling additional test data for model fine-tuning can be difficult due to privacy issues and high labeling costs, respectively. To solve this problem, we propose a novel atlas-guided test-time adaptation (TTA) method for robust 3D medical image segmentation, called AdaAtlas. AdaAtlas only takes one single unlabeled test sample as input and adapts the segmentation network by minimizing an atlas-based loss. Specifically, the network is adapted so that its prediction after registration is aligned with the learned atlas in the atlas space, which helps to reduce anatomical segmentation errors at test time. In addition, different from most existing TTA methods which restrict the adaptation to batch normalization blocks in the segmentation network only, we further exploit the use of channel and spatial attention blocks for improved adaptability at test time. Extensive experiments on multiple datasets from different sites show that AdaAtlas with attention blocks adapted (AdaAtlas-Attention) achieves superior performance improvements, greatly outperforming other competitive TTA methods.

Figures (4)

• Figure 1: AdaAtlas overview: Our TTA method utilizes an atlas-based shape prior loss $\mathcal{L}_{atlas}$ (Eq.\ref{['eq:atlas_loss']}) to adapt $\mathcal{S}_\theta$ on each single subject $\boldsymbol{x}_i$, which encourages the prediction $\hat{\mathbf{y}}_i$ to be aligned with a given atlas $\mathbf{y}_a$ after registration (reg.) in the atlas space (Sec. \ref{['sec:atlas_TTA']}). For efficiency, during TTA, only parts of the segmentation network, i.e., adaptation blocks (orange blocks, details in Sec. \ref{['sec:ada_blocks']}) are updated.
• Figure 2: Dual attention blocks with test-time adaptable parameters in red.
• Figure 3: Visualization of prostate MRI datasets from different sites. Sites A-F are target domains for testing. Site G is the training source domain.
• Figure 4: Visualization of segmentation results using Baseline model (without TTA) and different TTA methods at test time using the same segmentation backbone (U-Net+Dual attention). GT: ground truth.