Source-linked AI summary

Unsupervised domain adaptation for medical imaging segmentation with self-ensembling

Christian S. Perone, Pedro Ballester, Rodrigo C. Barros, Julien Cohen-Adad

arXiv:1811.06042v2cs.CV

TL;DR

Medical imaging models trained on a single domain may generalize poorly across domains, while expert annotation is costly. This paper adapts self-ensembling unsupervised domain adaptation to semantic segmentation and evaluates it on a small realistic MRI dataset, finding improved generalization from unlabeled data.

  • Problem

    Medical imaging models face domain shifts across centers and costly expert annotation, limiting generalization from labeled source domains.

  • Method

    The paper extends self-ensembling unsupervised domain adaptation to semantic segmentation and evaluates consistency-loss designs with unlabeled target data.

  • Results

    Self-ensembling improved generalization on unseen domains, with ablation evidence attributing improvements to unlabeled data rather than only the exponential moving average.

  • Takeaways & Limitations

    Unsupervised domain adaptation can increase medical-imaging model performance across multiple centers without relying on annotations.

  • Takeaways & Limitations

    The study evaluated only gray matter segmentation and did not evaluate adversarial domain-adaptation methods.

Abstract

from arXiv · show

Recent advances in deep learning methods have come to define the state-of-the-art for many medical imaging applications, surpassing even human judgment in several tasks. Those models, however, when trained to reduce the empirical risk on a single domain, fail to generalize when applied to other domains, a very common scenario in medical imaging due to the variability of images and anatomical structures, even across the same imaging modality. In this work, we extend the method of unsupervised domain adaptation using self-ensembling for the semantic segmentation task and explore multiple facets of the method on a small and realistic publicly-available magnetic resonance (MRI) dataset. Through an extensive evaluation, we show that self-ensembling can indeed improve the generalization of the models even when using a small amount of unlabelled data.

1 Introduction

Deep learning has become central to medical imaging, but models trained on one domain can generalize poorly across imaging centers and protocols. This paper extends self-ensembling unsupervised domain adaptation to semantic segmentation and evaluates it on a realistic small MRI dataset.

  • Deep neural networks have become widely used in medical imaging, where segmentation is especially important for quantifying anatomical structures and pathologies.
  • High sample complexity and expert-dependent annotation make the data requirements of medical imaging a persistent challenge.
  • Imaging protocols, parameter choices, vendors, and centers can create substantial data distribution shifts despite imaging the same object.
  • Samples from UCL, Montreal, Zurich, and Vanderbilt illustrate the multi-center MRI setting used to motivate domain-shift evaluation.
  • Homogeneous train-test splits can yield over-optimistic evaluations, whereas new centers commonly lack labeled data before deployment.
  • The paper extends unsupervised self-ensembling domain adaptation to semantic segmentation and explores consistency losses, unlabeled-data effects, and model dynamics.

2 Related work

Related work spans medical image segmentation, general deep domain adaptation, and studies documenting cross-institutional performance degradation. The paper positions its contribution as the first extension of self-ensembling unsupervised domain adaptation to semantic segmentation in medical imaging.

  • Medical image segmentation has progressed from patch-based convolutional methods toward Fully Convolutional Network architectures.
  • Deep domain adaptation methods include domain-invariant representations, higher-order statistics, explicit discrepancies, adversarial training, GANs, disentanglement, and self-ensembling.
  • Adversarial and generative approaches have been applied directly to medical imaging domain adaptation, including methods for segmentation and training augmentation.
  • The authors describe their work as the first to extend unsupervised self-ensembling domain adaptation to semantic segmentation tasks in medical imaging.

3 Semi-supervised learning and unsupervised domain adaptation

Semi-supervised learning combines labeled and unlabeled data, while unsupervised domain adaptation uses labeled source data and unlabeled target data under domain shift. The shared challenge is that unlabeled data may not follow the labeled data distribution.

  • Semi-supervised learning learns from labeled and unlabeled datasets to improve supervised or unsupervised learning tasks.
  • Semi-supervised learning generally relies on unlabeled data coming from a distribution similar to the labeled data.
  • Domain shift is the difference between the input and label-related distributions of training examples and test examples.
  • Unsupervised domain adaptation uses labeled source data and unlabeled target data, with target labels unavailable during adaptation.

4 Method

The method adapts Mean Teacher self-ensembling to unsupervised medical-image segmentation by combining supervised source-domain learning with target-domain consistency training. It uses EMA teacher updates, segmentation-specific losses, and matched spatial transformations to align student and teacher predictions.

  • Self-ensembling and mean teacher: Mean Teacher updates a teacher model from student weights with an exponential moving average, producing a temporally ensembled target model.The momentum term α controls the contribution of current versus previous model weights, with later training favoring stronger historical averaging.
  • Self-ensembling and mean teacher: Each training step combines a supervised task loss on labeled data with a consistency loss comparing student and teacher predictions for labeled and unlabeled samples.The overall objective also includes regularization, while the consistency weight follows a sigmoid ramp-up.
  • Unsupervised domain adaptation: The adaptation framework uses source-domain labels for the task loss and unlabeled target-domain images for consistency training.Batches contain images from both domains, enabling supervised source learning alongside target-domain prediction agreement.
  • Segmentation losses: Dice loss replaces cross-entropy as the segmentation task loss because it is less sensitive to class imbalance.Dice is computed over the entire batch for both baseline and adaptation experiments.
  • Segmentation-specific adaptation: Spatial augmentations use matching parameters on student inputs and teacher predictions to prevent spatial misalignment during consistency training.The target input and teacher prediction receive the same transformation, while student and teacher branches may otherwise use distinct parameterizations.
  • Consistency losses: Weighted MSE with thresholded teacher predictions did not improve overall results, while Dice and Tversky are problematic consistency losses for soft teacher targets.The paper leaves modified loss formulations for non-binary labels to future work because suitable thresholds and class weights are difficult to determine.

5 Materials

The study uses a small, publicly available, multi-center MRI dataset and a domain-adaptation split that keeps test-center data separate from training domains.

  • The dataset contains 80 healthy subjects, with 20 subjects from each of four centers.
  • Three MRI systems with distinct acquisition parameters produced voxel resolutions from 0.25 × 0.25 × 2.5 mm to 0.5×0.5×5.0 mm.
  • Each colored square in Figure 6 represents one subject containing multiple axial slices.
  • Centers 1 and 2 formed the training set, center 3 the validation set, and center 4 the test set.
  • Unlabeled center 4 test data supplied the target-domain adaptation data, while labeled center 4 training data evaluated final performance.

6 Experiments

Experiments evaluate adaptation across centers, consistency-loss choices, and training stability using supervised source data from centers 1 and 2 and unlabeled target-domain data.

  • 6 Experiments: The experiments train on centers 1 and 2, then adapt separately to centers 3 and 4 to study prediction changes, adaptation, and unseen-center generalization.
  • 6.1 Adapting to different centers: For both adapted centers 3 and 4, all metrics except recall outperform the baseline on the source-domain evaluation.
  • 6.1 Adapting to different centers: Evaluation on the same center used for adaptation gives the highest values for almost all metrics, except recall.
  • 6.1 Adapting to different centers: Cross-center evaluation after adaptation again gains over baseline for most metrics, suggesting improved generalization to unseen centers.
  • 6.2 Varying the consistency loss: The consistency-loss comparison considers cross-entropy, mean squared error, and Dice without additional hyperparameters.
  • 6.4 Training stability: Cross-entropy consistently fails across tested weights, and its training diverges after several iterations in the illustrated center-3 experiment.
  • 6.4 Training stability: MSE is relatively invariant to consistency weight, making it robust when target-center annotations are unavailable.

7 Ablation studies

The ablation isolates the exponential moving average from unlabeled-data consistency training to test whether EMA alone explains the observed improvement.

  • 7 Ablation studies: The ablation keeps EMA active while setting the consistency weight to zero, removing the unlabeled-data consistency signal.
  • 7.1 Exponential moving average (EMA): The authors therefore attribute the improvement observed in the full Mean Teacher setup specifically to the unlabeled-data consistency training rather than EMA alone.
  • 7.1 Exponential moving average (EMA): With consistency disabled, the EMA teacher shows no gains over the non-averaged supervised baseline.

8 Domain shift visualization

t-SNE visualizations compare supervised and adapted prediction features across centers, showing that unlabeled adaptation changes the organization of center-specific representations.

  • 8 Domain shift visualization: The analysis uses t-SNE to assess changes in the network’s predictive representation of unsupervised data, excluding the visualized data from training.
  • 8 Domain shift visualization: The t-SNE configurations use learning rate 10, perplexity 30, and approximately 1,000 iterations as a trade-off between group structure and interpretability.
  • 8 Domain shift visualization: In supervised training, data from centers 1 and 2 separate clearly from unseen centers 3 and 4.
  • 8 Domain shift visualization: After unsupervised adaptation, predictions become more diffuse for training centers and form clusters with domains observed only through unlabeled data.
  • 8 Domain shift visualization: Centers 1 and 4, and centers 2 and 3, form nearby clusters, potentially reflecting similar intensity distributions within each pair.

9 Conclusion and limitations

The study finds that unsupervised domain adaptation with self-ensembling improves medical-imaging generalization across domains using unlabeled data, while highlighting consistency-loss stability issues and several scope limitations.

  • Unsupervised domain adaptation increased model performance across multiple medical-imaging centers without relying on annotations.
  • Self-ensembling improved generalization on unseen domains by leveraging unlabeled data from multiple domains.An ablation study attributed the improvements to introducing unlabeled data, not solely to the exponential moving average.
  • Cross-entropy consistency loss failed to maintain training stability as training progressed, creating potential problems in more challenging multi-center scenarios.The study also discussed issues with Dice loss as a consistency loss.
  • Limitations: The study did not evaluate adversarial domain-adaptation methods, leaving comparison in the realistic small-data regime for future work.
  • Limitations: Evaluation was limited to single-task gray-matter segmentation, so broader confirmation requires more centers and tasks.
  • Limitations: Future methods could measure adaptation risk for particular centers or domains to clarify domain-adaptation limitations.
  • The authors identify proper multi-domain evaluation as important because medical-imaging variability is rarely represented in test sets from different centers or machines.

Appendix Extended visualizations

The appendix extends the domain-adaptation visualization with t-SNE embeddings and raw-intensity distributions, showing how similar pixel distributions cluster in prediction space.

  • Figure A1 extends the t-SNE visualization from the Figure 8b domain-adaptation scenario.
  • The visualization places each center’s pixel distribution alongside its corresponding prediction-space clusters.
  • Similar pixel distributions tend to form clusters in prediction space.
Loading 1811.06042v2…