RemInD: Remembering Anatomical Variations for Interpretable Domain Adaptive Medical Image Segmentation

TL;DR

Experiments on two public datasets demonstrate the superiority of RemInD, which achieves state-of-the-art performance using a single alignment approach, outperforming existing methods that often rely on multiple complex alignment strategies.

Abstract

This work presents a novel Bayesian framework for unsupervised domain adaptation (UDA) in medical image segmentation. While prior works have explored this clinically significant task using various strategies of domain alignment, they often lack an explicit and explainable mechanism to ensure that target image features capture meaningful structural information. Besides, these methods are prone to the curse of dimensionality, inevitably leading to challenges in interpretability and computational efficiency. To address these limitations, we propose RemInD, a framework inspired by human adaptation. RemInD learns a domain-agnostic latent manifold, characterized by several anchors, to memorize anatomical variations. By mapping images onto this manifold as weighted anchor averages, our approach ensures realistic and reliable predictions. This design mirrors how humans develop representative components to understand images and then retrieve component combinations from memory to guide segmentation. Notably, model prediction is determined by two explainable factors: a low-dimensional anchor weight vector, and a spatial deformation. This design facilitates computationally efficient and geometry-adherent adaptation by aligning weight vectors between domains on a probability simplex. Experiments on two public datasets, encompassing cardiac and abdominal imaging, demonstrate the superiority of RemInD, which achieves state-of-the-art performance using a single alignment approach, outperforming existing methods that often rely on multiple complex alignment strategies.

Figures (4)

• Figure 1: The proposed framework RemInD. (a) Based on inferred anchor weights (yellow arrows), images are mapped onto a latent manifold (red arrows), and segmentations are warped by spatial transformations to produce final predictions (green arrows). (b) Inference (pink) and generative (green) models (\ref{['sec:infer']}), with observations shaded.
• Figure 2: Network architecture of RemInD, illustrated with $L=3$ and $M=4$ as an example, including image feature encoding (orange), atlas inference (red), spatial transformation inference (blue), segmentation & reconstruction decoding (green), and loss calculations (purple). The feature maps parameterizing posterior distributions contain one half of channels for the mean and the other half for the log variance. Values around arrows indicate channel numbers.
• Figure 3: Qualitative comparison between RemInD and the best baselines. The spatial deformations $\boldsymbol{\phi}$ from RemInD are also displayed. Yellow arrows indicate regions where one method produces inferior predictions compared to the other.
• Figure 4: 3D visualization of the positive orthant of the unit sphere, with the shape blending weights for images in the MS-CMR dataset. (a) Distribution of the transformed weights $\boldsymbol{w}^\dagger$. Each point is based on a 2D image slice, with colors indicating its relative location among the total number of slices for the corresponding patient. (b) Manipulating $\boldsymbol{w}^\dagger$ (thus $\boldsymbol{z}$) along geodesics induces gradual variations in predicted segmentation (before warped by $\boldsymbol{\phi}^{-1}$). Endpoints of the shown six geodesics: (1,0,0), (0,1,0), (0,0,1), (0.99,0.1,0.1), (0.1,0.99,0.1), (0.1,0.1,0.99). Note that the segmentations with broken shapes could still be valid, as a few similar ground-truth labels exist.