PaperlayerView exact revision on arXiv

Source-linked AI summary

Test-Time Adaptable Neural Networks for Robust Medical Image Segmentation

Neerav Karani, Ertunc Erdil, Krishna Chaitanya, Ender Konukoglu

arXiv:2004.04668v4eess.IVcs.CVcs.LGstat.ML

TL;DR

Medical image segmentation CNNs can lose robustness when test scanners or protocols differ from those used for training. This paper adapts a per-image normalization module using a denoising-autoencoder prior, with substantial gains reported across brain and prostate datasets.

  • Problem

    CNN segmentation methods lack robustness to variations in imaging protocols and scanners between training and test images, impeding their large-scale adoption in practice.

    Source: §1 · p. 1

  • Method

    The approach places an adaptable image-normalization module before the segmentation CNN and updates it per test image using a denoising autoencoder to encourage plausible predictions.

    Source: §3.2.1 · p. 4, §3.2.2 · p. 5

  • Results

    Test-time adaptation provided substantial performance gains over competing methods across brain and prostate datasets; cardiac segmentation was already fairly good after source-domain training with augmentation.

    Source: §4.4.3 · p. 12

  • Takeaways & Limitations

    Experiments across multiple datasets and anatomies demonstrate the method's promise and generality relative to data augmentation, meta-learning, and unsupervised domain adaptation.

    Source: §6 · p. 16

    The method assumes scanner or protocol shifts preserve anatomy similar to the training data; very large modality or protocol changes may require atlas initialization before DAE-guided adaptation.

    Source: §3.2.2 · p. 4, §3.4 · p. 6

Abstract

from arXiv · show

Convolutional Neural Networks (CNNs) work very well for supervised learning problems when the training dataset is representative of the variations expected to be encountered at test time. In medical image segmentation, this premise is violated when there is a mismatch between training and test images in terms of their acquisition details, such as the scanner model or the protocol. Remarkable performance degradation of CNNs in this scenario is well documented in the literature. To address this problem, we design the segmentation CNN as a concatenation of two sub-networks: a relatively shallow image normalization CNN, followed by a deep CNN that segments the normalized image. We train both these sub-networks using a training dataset, consisting of annotated images from a particular scanner and protocol setting. Now, at test time, we adapt the image normalization sub-network for \emph{each test image}, guided by an implicit prior on the predicted segmentation labels. We employ an independently trained denoising autoencoder (DAE) in order to model such an implicit prior on plausible anatomical segmentation labels. We validate the proposed idea on multi-center Magnetic Resonance imaging datasets of three anatomies: brain, heart and prostate. The proposed test-time adaptation consistently provides performance improvement, demonstrating the promise and generality of the approach. Being agnostic to the architecture of the deep CNN, the second sub-network, the proposed design can be utilized with any segmentation network to increase robustness to variations in imaging scanners and protocols. Our code is available at: \url{https://github.com/neerakara/test-time-adaptable-neural-networks-for-domain-generalization}.

1 Introduction

The paper proposes a segmentation network with a test-time-adaptable image-normalization module, guided by a denoising autoencoder prior on plausible segmentations. It frames this approach as domain generalization for unseen scanners and protocols, without requiring training samples during inference.

1 Introduction

  • Scanner and protocol shifts alter image intensities and tissue contrast, degrading segmentation CNN performance and making separate labeled training sets for each domain impractical.The introduction frames domain generalization as practical because it needs source data only for initial training, unlike unsupervised domain adaptation, which requires source data during training for each target domain.Source: §1 · pp. 1–2
  • At inference, the method adapts a shallow image-to-normalized-image CNN for each test image using a denoising autoencoder prior, while keeping the segmentation CNN frozen.Both sub-networks are jointly supervised-trained on source-domain data; no source samples are required during test-time adaptation.Source: §1 · p. 2
  • The adaptable normalization module targets each test image independently, while its shallow design limits structural changes and the frozen segmentation network retains supervised-training benefits.The denoising autoencoder can use high-level cues such as context and shape to suggest corrections to predicted segmentations.Source: §1 · p. 2
  • Experiments on multi-center MRI datasets for brain, prostate, and heart show consistent, substantial improvements on unseen target domains over competing methods.The authors describe this as the first proposed test-time adaptation approach for cross-scanner robustness in CNN-based medical image segmentation.Source: §1 · p. 2

2 Related work

Prior work addresses domain shift through feature invariance, augmentation, shape constraints, unsupervised segmentation, test-time adaptation, and post-processing. The authors argue that test-time adaptability can help address unseen domains and complement a stronger base CNN.

Domain Invariant Features

  • Feature-invariance methods use multi-task autoencoders, MMD, adversarial losses, or meta-learning, but commonly require multiple source domains during training.The approaches aim to discourage reliance on domain-specific input signals.Source: §2 · p. 2

Data Augmentation

  • Data augmentation promotes domain invariance by simulating plausible test-time variations through heuristic transforms, data-generation knowledge, worst-case searches, or soft in-between domains [25].Source: §2 · pp. 2–3

Imposing Shape Constraints during Training

  • Shape and topology constraints can impose output-space invariance, but the discussed method categories do not adapt the CNN at test time.The paper notes that these approaches could be combined with multiple-source-domain training and that the methods are complementary to its approach.Source: §2 · p. 3

Unsupervised Segmentation

  • Probabilistic generative models infer unknown segmentations in a Bayesian framework, but have largely used priors over small pixel neighbourhoods and been applied mainly in neuroimaging.Recent works learn long-range spatial priors using high-order Markov random fields and variational auto-encoders.Source: §2 · p. 3

3 Method

The method trains a segmentation network on source-domain image-label pairs, then adapts its shallow image-normalization CNN separately for each test image while keeping the segmentation CNN fixed. A denoising autoencoder guides adaptation toward plausible labels, with an atlas-based initialization for large domain shifts.

3.1 Training a Segmentation CNN on the Source Domain

  • The method trains a source-domain segmentation CNN as a concatenation of an image-to-normalized-image CNN and a normalized-image-to-segmentation CNN, minimizing supervised label-prediction loss.Without adaptation, the trained concatenation directly predicts segmentations for new images.Source: §3.1 · p. 4

3.2 Test-Time Adaptation

  • At test time, the method adapts the image-normalization CNN for each image, using only that image because the pretrained mapping may be unreliable under distribution shifts.The pretrained network’s parameters provide the initialization for this adaptation.Source: §3.2 · p. 4

3.2.1 Which parameters of segCNN to adapt at test-time?

  • Only the shallow residual normalization CNN is adapted; the deeper segmentation CNN remains fixed, allowing image-specific contrast changes while limiting alterations to anatomical content.The normalization CNN uses stride-1 convolutions without spatial down- or up-sampling and outputs the input’s channel count.Source: §3.2.1 · p. 4

3.2.2 How to drive the test-time adaptation?

  • The DAE denoises predicted segmentations toward labels resembling the source-domain training set, and the normalization CNN is updated to bring its predictions closer to the DAE outputs.The update is iterative; the final normalization parameters are selected from those yielding the least dissimilarity between DAE input and output.Source: §3.2.2 · p. 5, Figure 1 · p. 5

3.3 DAE Training

  • The DAE is a 3D CNN trained to map corrupted segmentations to clean source-domain labels, with corruption generated by copying cubic patches between randomly chosen locations in each label image.The 3D design can learn relative anatomical locations and whole-shape information.Source: §3.3 · pp. 5–6

3.4 Atlas initialization for test-time adaptation for large domain shifts

  • For large shifts such as different imaging modalities or very different protocols, adaptation starts with an affinely registered atlas and switches to DAE guidance when threshold conditions indicate predictions have improved.The switching rule uses similarity between the prediction and atlas and between the DAE input and output.Source: §3.4 · p. 6

3.5 Integrating 2D Segmentation CNN with 3D DAE

  • For 2D segmentation CNNs, the method predicts successive slices of a 3D test image, denoises the resulting volume with the 3D DAE, then aggregates slice-batch gradients to update normalization parameters.When protocol changes are present, the atlas-or-DAE threshold rule can select the adaptation targets.Source: §3.5 · p. 7

4 Experiments and Results · 4.1 Datasets

The experiments test domain generalization across brain, prostate, and cardiac MRI, using labeled data from one source domain and evaluating transfer to target datasets. Across domains, anatomy is relatively consistent while image contrast varies.

4.1 Datasets

  • The study evaluates transfer from one labeled source domain, which may include images from multiple scanners, across brain, heart, and prostate MRI datasets.The authors note that aggregating data across imaging centers may be difficult in practice.Source: §4.1 · p. 7
  • Brain experiments use HCP-T1w as the source domain and ABIDE-Caltech-T1w and HCP-T2w as target domains, segmenting 15 labels with FreeSurfer-generated pseudo ground truth.FreeSurfer is scanner- and protocol-robust but can take up to 10 hours on a CPU to segment one 3D MR image.Source: §4.1 · p. 7
  • Prostate experiments use NCI as the source and PROMISE12 and USZ as targets, evaluating whole-gland segmentation and, where annotations permit, central gland and peripheral zone separately.NCI and USZ provide expert labels for background, central gland, and peripheral zone; PROMISE12 provides whole-gland labels.Source: §4.1 · p. 7
  • Cardiac experiments train on ACDC and adapt to RVSC, evaluating RV endocardium because it is the only structure annotated in both datasets.ACDC annotations include LV endocardium, LV epicardium, and RV endocardium; RVSC annotations include RV epicardium and RV endocardium.Source: §4.1 · pp. 7–8
  • Table 1 reports training, test, and validation splits by number of 3D images for each dataset.The supplied passages do not provide the split counts.Source: §4.1 · p. 8, Table 1 · p. 8
  • Examples across the datasets show relatively consistent underlying anatomy but varying image contrast across scanners and protocols.The examples include three brain datasets, three prostate datasets, and two cardiac datasets; prostate labels are separate sub-glands in two datasets and the whole gland in the third.Source: Figure 2 · p. 8

4.2 Pre-processing

Pre-processing standardizes image intensities and spatial dimensions, with separate 2D pipelines for segmentation and 3D pipelines for the denoising autoencoder (DAE). Brain images also undergo skull stripping, and predictions are evaluated at their original pixel size.

4.2 Pre-processing

  • Images undergo N4 bias-field removal and per-image percentile-based 0–1 intensity normalization with clipping; brain datasets additionally undergo skull stripping.Skull stripping sets all non-brain voxel intensities to 0.Source: §4.2 · p. 8
  • The 2D segCNN pipeline rescales images to anatomy-specific in-plane pixel sizes, then crops or zero-pads them to 256x256.Pixel sizes are 0.7mm2 for brain, 0.625mm2 for prostate, and 1.33mm2 for cardiac images; labels receive the same transformation, while predictions are rescaled to the original pixel size for evaluation.Source: §4.2 · pp. 8–9
  • The DAE is trained in 3D on segmentation labels rescaled and cropped or padded to fixed voxel sizes and image dimensions.Voxel sizes are 2.8x0.7x0.7mm3 for brain, 2.5x0.625x0.625mm3 for prostate, and 5.0x1.33x1.33mm3 for cardiac; image dimensions are 64x256x256 for brain and 32x256x256 for the other anatomies.Source: §4.2 · pp. 8–9

4.3 Implementation Details

The implementation combines a shallow normalization CNN with U-Net-style segmentation and denoising networks, and specifies their optimization, augmentation, DAE noise, atlas initialization, and evaluation procedures.

4.3.1 Network Architectures

  • The normalization CNN has three 3-channel-kernel convolutional layers with 16, 16, and 1 outputs, while segmentation and DAE use U-Net-style encoder-decoders with skip connections.The normalization CNN uses a trainable, channel-specific exponential activation; batch normalization, ReLU, and bilinear upsampling are used in the segmentation and DAE networks. The adaptation strategy, normalization module, and DAE are agnostic to the architecture of the segmentation CNN.Source: §4.3.1 · p. 9

4.3.2 Optimization Details

  • Test-time optimization uses 500 brain updates and 7,500 cardiac and prostate updates, refreshing denoised labels every 25 steps.All networks use Dice loss and Adam at learning rate 10^-3; the segmentation CNN and DAE are trained for 50,000 iterations. The fast variant initializes later images from the first image’s optimized normalization parameters and uses 100 brain or 1,500 other-anatomy updates; adaptation takes about 1 hour for the first image and 12 minutes thereafter on a TITAN X.Source: §4.3.2 · p. 9, §4.3.3 · p. 9

4.3.3 Data Augmentation

  • Segmentation training uses geometric and intensity augmentation, adds right-angle rotations and flips for cardiac images, and applies geometric augmentation to DAE labels with probability 0.25.Geometric transformations include translations, rotations, scaling, and elastic deformations; intensity transformations include gamma changes, brightness changes, and per-pixel additive Gaussian noise. Geometric transformations apply to images and segmentations, while intensity transformations apply only to images.Source: §4.3.3 · pp. 9–10

4.3.4 Noise Hyper-parameters for DAE Training

  • DAE training uses a maximum of 20 copied patches, and the best model is selected by denoising performance on corrupted SD validation images.Each validation image is corrupted 50 times using the described noising process; the validation set is part of the source domain, not the target domain.Source: §4.3.4 · p. 10

4.3.5 Hyper-parameters for atlas-based initial optimization in case of domain shifts involving a protocol change

  • For brain protocol shifts, atlas-based initialization averages one-hot source labels voxel-wise, then switches to DAE guidance using thresholding with α = 1.0 and β = 0.25.The source domain uses T1w images and the target domain uses T2w images from different HCP subjects; the images are already rigidly registered.Source: §4.3.5 · p. 10

4.3.6 Evaluation Criteria

  • Predicted segmentations are evaluated using Dice coefficient and 95th-percentile Hausdorff distance, averaged across foreground labels, test images, and three experiment runs.Source: §4.3.6 · p. 10

4.4 Results · 4.4.1 Baseline and Benchmark · 4.4.2 Relevant Methods

The results compare source-trained segmentation with target-specific training and evaluate domain-generalization methods, whose observed benefits vary with protocol match and method. The supplied passages describe comparison scope and qualitative results but do not provide quantitative values for the proposed method.

4.4 Results

  • Figure 3 presents qualitative test images, ground truths, normalized images, and predicted segmentations for SD, SD + DA, SD + DA + TTA, and TD training.Source: Figure 3 · p. 10

4.4.1 Baseline and Benchmark

  • Source-trained CNNs provide the baseline, while separately trained target-specific CNNs define the benchmark and expose the generalization gap domain-generalization methods aim to close.Table 2 reports mean Dice scores and 95th percentile Hausdorff distance over foreground labels, test subjects, and three experiment runs.Source: §4.4.1 · p. 11, Table 2 · p. 11

4.4.2 Relevant Methods

4.4.2.1. Meta Learning based Methods

  • Meta-learning substantially improved over baseline when source and target images used the same protocol, using gamma transformations to simulate meta-train and meta-test domains.The training also used the other augmentations described in Section 4.3.3.Source: §4.4.2.1 · p. 12

4.4.2.2. Data Augmentation

  • Extensive data augmentation substantially improved performance for same-protocol shifts but remained below the benchmark of training separately on target domains.The authors treat data augmentation as the strong baseline for their method.Source: §4.4.2.2 · p. 12

4.4.2.3. Post-Processing with DAEs

  • DAE post-processing substantially improved over SD + DA, especially for prostate, but degraded brain results and failed when source and target protocols or modalities differed.The proposed explanation is that DAE outputs may be plausible without matching the test image; corrupted predictions outside the DAE training input distribution could not be improved, as with T1w source and T2w brain TD2.Source: §4.4.2.3 · p. 12

4.4.3 Test-Time Adaptation · 4.4.4 Comparison with Unsupervised Domain Adaptation

Test-time adaptation improves segmentation across brain and prostate datasets while preserving cardiac performance, and performs comparably to the best UDA methods without requiring labelled source data during adaptation. Its advantages and limitations vary with the size of the domain shift.

4.4.3 Test-Time Adaptation

  • TTA substantially improves over competing methods for brain and prostate, preserves cardiac performance without further gains, and significantly beats direct DAE post-processing on 3 of 5 datasets.Significance was assessed with a paired permutation test using 100000 permutations.Source: §4.4.3 · p. 12
  • Qualitatively, TTA corrects contextually misplaced predictions, completes organ shapes, and removes outliers, with especially substantial improvements for brain and prostate.These examples are shown in Fig. 3 relative to SD + DA.Source: §4.4.3 · p. 12
  • Initializing later-image adaptation with optimal parameters from the first target-domain image reduces iterations and time, while source-training initialization gives similar results, showing adaptation is stable across images.The method therefore does not depend on a good first image from a new target domain.Source: §4.4.3 · p. 12

4.4.4 Comparison with Unsupervised Domain Adaptation

  • UDA typically adapts using unlabelled target-domain training images Xtr_TD and evaluates on separate test images Xts_TD; the study also evaluates Xtr_TD without using its labels, a setting proposed in.UDA is presented as assuming access to labelled source-domain data, a requirement the paper notes may be challenging in practice.Source: §4.4.4 · pp. 12–13
  • TTA achieves results comparable to the best UDA methods, especially for scanner shifts, without labelled source data during adaptation; cycleGAN does better for modality shifts but worse than baseline for scanner shifts.Table 3 presents the comparison. The paper describes UDA as assuming labelled source data is available, while TTA does not require it when adapting to each new target domain.Source: Table 3 · p. 13, §4.4.4 · p. 13
  • The two representative UDA methods struggle when applied across mismatched shift sizes: adversarial method for modality changes and cycleGAN for scanner changes both underperform SD + DA.The paper reports poorer accuracies than the SD + DA baseline in these respective settings.Source: §4.4.4 · p. 13
  • UDA performance gains on Xtr_TD are similar to those on Xts_TD, so using the same images for adaptation and testing brought no additional improvements.The target-domain training-set size matched the specification in Table 1.Source: §4.4.4 · p. 13

4.4.5 Analysis experiments · 4.4.6 Convergence of Test-Time Adaptation

The analysis experiments examine why test-time adaptation is restricted to the I2NI sub-network and driven by the DAE, while the convergence experiments assess its empirical reliability and the DAE-based selection criterion.

4.4.5 Analysis experiments

  • Adapting all segCNN parameters reduced Dice but improved Hausdorff distance in every target domain except prostate TD2.The authors attribute the Dice decline to less accurate organ edges and the Hausdorff improvement to better removal of outliers, favoring freezing most parameters.Source: §4.4.5 · p. 13
  • Oracle-guided I2NI adaptation improved accuracy over SD + DA, but brain target-domain results remained below separately trained target-domain models.This ablation was not run on PROMISE or RVSC because their annotations covered different organs or tissues from the corresponding source-domain datasets; the brain results suggest residual source-domain bias in NI2S.Source: §4.4.5 · pp. 13–14
  • Multiple DAE passes improved segmentation less than test-time adaptation and worsened accuracy on the brain datasets.Figure 4 shows initial predictions, repeated DAE outputs, and test-time adaptation results for ABIDE-Caltech and USZ.Source: §4.4.5 · p. 14, Figure 4 · p. 14
  • Freezing NI2S and keeping adaptable I2NI shallow tie predictions to the input image and limit the impact of DAE errors.The authors contrast this with repeated DAE post-processing, whose outputs may be plausible for source-domain labels without matching the particular input image.Source: §4.4.5 · p. 14

4.4.6 Convergence of Test-Time Adaptation

  • Although convergence is not theoretically guaranteed, adaptation converged across more than 100 test volumes, anatomies, target domains, and repeated runs.Figure 5 is described as showing reliable convergence across target domains.Source: §4.4.6 · p. 14
  • The Dice between DAE inputs and outputs correlated with predicted-versus-ground-truth Dice, supporting its use to select optimal I2NI parameters.The reported correlation links the adaptation criterion to segmentation accuracy.Source: §4.4.6 · p. 14

5 Discussion

The method showed promising robustness improvements on images from unseen scanners or protocols, while the discussion identifies assumptions, evaluation limits, and possible extensions that could help close the gap to scanner-specific training. Its per-image flexibility adds adaptation time.

Noising strategy used for DAE training

  • The DAE’s denoising reliability depends on incorrect unseen-domain predictions resembling its training inputs; the heuristic corruption strategy might improve with noisy labels from intensity-transformed source images.One proposed route is training the segmentation CNN without augmentation, then using its predictions on gamma-transformed source images as noisy segmentations.Source: §5 · pp. 14–15

Assumption that DAE outputs are samples from the posterior P(Z|Zn)

  • The assumption that DAE outputs represent samples from P(Z|Zn) is uncertain because training maps each noisy segmentation to one clean segmentation rather than multiple possible denoised outputs.The authors suggest exploring DAE training that produces such behavior.Source: §5 · p. 15

Prior on normalized images

  • Because the normalized image is an intermediate representation, test-time adaptation is not expected to make the normalization network a translator from target-domain to source-domain images.Future work could impose additional constraints on the intermediate representation to guide adaptation.Source: §5 · p. 15

Affine registration for large domain shifts

  • Atlas-based initialization for large domain shifts was evaluated only on brain images, and affine registration may be needed for other target domains or anatomies.The HCP source and target images were already rigidly aligned, so atlas creation could skip affine registration; evaluation beyond brain images remains future work.Source: §5 · p. 15

Test-time adaptation in PGM framework

  • A possible extension is test-time adaptation of a supervised CNN in a probabilistic graphical model, using explicit prior P(Z) and likelihood P(X|Z) to maximize posterior P(Z|X).Source: §5 · p. 15

Time required for test-time adaptation

  • About 12 minutes are required per 3D image after the first image from a scanner/protocol, adding per-image optimization despite the authors’ view that the time is reasonable for clinical use.This timing was measured with the experimental implementation; the authors note that efficiency could be improved.Source: §5 · p. 15

6 Conclusion

The work proposes domain generalization for medical image segmentation under scanner and protocol shifts by combining per-image normalization with test-time adaptation guided by denoising autoencoders. Experiments across multiple datasets and anatomies demonstrate the method’s promise and generality over data augmentation, meta-learning, and unsupervised domain adaptation.

6 Conclusion

  • Experiments across multiple datasets and anatomies demonstrate the method’s promise and generality over data augmentation, meta-learning, and unsupervised domain adaptation.Source: §6 · p. 16
  • The method introduces an adaptable per-image normalization module into a segmentation CNN for domain generalization under scanner and protocol changes.Source: §6 · p. 16
  • Test-time adaptation uses denoising autoencoders to incentivize plausible segmentation predictions.Source: §6 · p. 16

Related papers

Loading 2004.04668v4…