Table of Contents
Fetching ...

Unsupervised Federated Domain Adaptation for Segmentation of MRI Images

Navapat Nananukul, Hamid Soltanian-zadeh, Mohammad Rostami

TL;DR

This work tackles MRI semantic segmentation under domain shift without persistent annotations by proposing Federated Multi-Source UDA (FMUDA). FMUDA trains separate single-source UDA models for each annotated source and aligns their target-domain embeddings in a shared latent space using a distribution-matching objective, then ensembles the per-source predictions using confidence-based weights to form a final segmentation while preserving data privacy. Theoretical analysis adapts existing generalization bounds to the multi-source setting, showing the target error is upper-bounded by a weighted combination of source errors, embedding-space distribution distances, sample-size terms, and an optimal joint-domain error. Empirically, FMUDA achieves state-of-the-art Dice scores on the MICCAI 2016 MS lesion dataset across multiple two-source scenarios, substantially outperforming baselines and demonstrating robustness to negative transfer; the approach supports scalable, privacy-preserving incorporation of new sites with minimal retraining.

Abstract

Automatic semantic segmentation of magnetic resonance imaging (MRI) images using deep neural networks greatly assists in evaluating and planning treatments for various clinical applications. However, training these models is conditioned on the availability of abundant annotated data to implement the end-to-end supervised learning procedure. Even if we annotate enough data, MRI images display considerable variability due to factors such as differences in patients, MRI scanners, and imaging protocols. This variability necessitates retraining neural networks for each specific application domain, which, in turn, requires manual annotation by expert radiologists for all new domains. To relax the need for persistent data annotation, we develop a method for unsupervised federated domain adaptation using multiple annotated source domains. Our approach enables the transfer of knowledge from several annotated source domains to adapt a model for effective use in an unannotated target domain. Initially, we ensure that the target domain data shares similar representations with each source domain in a latent embedding space, modeled as the output of a deep encoder, by minimizing the pair-wise distances of the distributions for the target domain and the source domains. We then employ an ensemble approach to leverage the knowledge obtained from all domains. We provide theoretical analysis and perform experiments on the MICCAI 2016 multi-site dataset to demonstrate our method is effective.

Unsupervised Federated Domain Adaptation for Segmentation of MRI Images

TL;DR

This work tackles MRI semantic segmentation under domain shift without persistent annotations by proposing Federated Multi-Source UDA (FMUDA). FMUDA trains separate single-source UDA models for each annotated source and aligns their target-domain embeddings in a shared latent space using a distribution-matching objective, then ensembles the per-source predictions using confidence-based weights to form a final segmentation while preserving data privacy. Theoretical analysis adapts existing generalization bounds to the multi-source setting, showing the target error is upper-bounded by a weighted combination of source errors, embedding-space distribution distances, sample-size terms, and an optimal joint-domain error. Empirically, FMUDA achieves state-of-the-art Dice scores on the MICCAI 2016 MS lesion dataset across multiple two-source scenarios, substantially outperforming baselines and demonstrating robustness to negative transfer; the approach supports scalable, privacy-preserving incorporation of new sites with minimal retraining.

Abstract

Automatic semantic segmentation of magnetic resonance imaging (MRI) images using deep neural networks greatly assists in evaluating and planning treatments for various clinical applications. However, training these models is conditioned on the availability of abundant annotated data to implement the end-to-end supervised learning procedure. Even if we annotate enough data, MRI images display considerable variability due to factors such as differences in patients, MRI scanners, and imaging protocols. This variability necessitates retraining neural networks for each specific application domain, which, in turn, requires manual annotation by expert radiologists for all new domains. To relax the need for persistent data annotation, we develop a method for unsupervised federated domain adaptation using multiple annotated source domains. Our approach enables the transfer of knowledge from several annotated source domains to adapt a model for effective use in an unannotated target domain. Initially, we ensure that the target domain data shares similar representations with each source domain in a latent embedding space, modeled as the output of a deep encoder, by minimizing the pair-wise distances of the distributions for the target domain and the source domains. We then employ an ensemble approach to leverage the knowledge obtained from all domains. We provide theoretical analysis and perform experiments on the MICCAI 2016 multi-site dataset to demonstrate our method is effective.
Paper Structure (19 sections, 3 theorems, 10 equations, 8 figures, 2 tables, 1 algorithm)

This paper contains 19 sections, 3 theorems, 10 equations, 8 figures, 2 tables, 1 algorithm.

Key Result

Theorem 1

Consider Algorithm algorithm for FMUDA under the explained conditions, then the following holds where $\mathcal{C}_k$ is the combined error loss with respect to domain $k$, and $h^*_k$ is the optimal model with respect to this loss when a shared model is trained jointly on annotated datasets from all domains simultaneously.

Figures (8)

  • Figure 1: Block-diagram of the proposed multi-UDA approach: (a) we train source-specific models for each source domain based on ERM. (b) we perform single-source UDA for adapting each source-trained model via distributional alignment in the shared embedding space (c) we aggregate the individual source-trained model predictions to make the final prediction on the target domain predictions according to their reliability.
  • Figure 2: Segmentation masks generated for a sample MRI image when Source "01" is used as the source domain in UDA. In each figure, the colored area shows the mask generated by each UDA model.
  • Figure 3: Distribution matching in the embedding space: we use UMAP for visualization of data representations when Source "07" in the dataset is used as the UDA source domain and Source "01" of the dataset is used as the UDA target domain: (Left) source domain; (Center) target domain prior to single-source model adaptation; and (Right) target domain after single-source model adaptation
  • Figure 4: Effect of the pretraining and adaptation process on the target domain performance (yellow curve) and the training loss (blue curve).
  • Figure 5: Mode Performance versus the value for the hyperparameter $\lambda$.
  • ...and 3 more figures

Theorems & Definitions (5)

  • Theorem 1
  • Theorem 2
  • Lemma 1
  • proof
  • proof