Domain Game: Disentangle Anatomical Feature for Single Domain Generalized Segmentation

TL;DR

The paper tackles single-domain generalization in medical image segmentation by addressing the ill-posedness of feature disentanglement through Domain Game, a two-encoder framework that separates $X$ (diagnostic features) from $\Delta$ (domain-specific features) using geometry-based transformations $\mathcal{T}_{\hat{\pi}}$. It couples a Dice-optimized diagnostic branch with a PSNR-guided reconstruction branch under a minmax objective and employs a Lasso regularizer to constrain the feature space, implemented with EfficientNet-B2 encoders and a DeeplabV3 decoder. Empirical results on prostate and BRaTS brain tumor segmentation show substantial cross-site generalization gains, with approximately $11.8\%$ Dice improvement in prostate and $10.5\%$ in brain tumor segmentation over the second-best methods, and ablation confirms the importance of the domain encoder, space constraint, and geometry transformations. Overall, the work provides a practical framework to mitigate domain shift in SDG for medical imaging, with potential to improve generalization in diverse clinical settings.

Abstract

Single domain generalization aims to address the challenge of out-of-distribution generalization problem with only one source domain available. Feature distanglement is a classic solution to this purpose, where the extracted task-related feature is presumed to be resilient to domain shift. However, the absence of references from other domains in a single-domain scenario poses significant uncertainty in feature disentanglement (ill-posedness). In this paper, we propose a new framework, named \textit{Domain Game}, to perform better feature distangling for medical image segmentation, based on the observation that diagnostic relevant features are more sensitive to geometric transformations, whilist domain-specific features probably will remain invariant to such operations. In domain game, a set of randomly transformed images derived from a singular source image is strategically encoded into two separate feature sets to represent diagnostic features and domain-specific features, respectively, and we apply forces to pull or repel them in the feature space, accordingly. Results from cross-site test domain evaluation showcase approximately an ~11.8% performance boost in prostate segmentation and around ~10.5% in brain tumor segmentation compared to the second-best method.

Figures (3)

• Figure 1: Domain Game overview. Colours indicates the relation to distinct features, i.e. purple - overall, red - diagnostic and blue - domain-specific features. (a) Image $I$ is drawn from singular domain ($C$) to update the parameters $D$ at each step according to utilities $U$. (b) Transformed set is generated by applying $\mathcal{T}_{\hat{\pi}}$ to $I$, and subsequently encoded into feature sets by two encoders. Then features within $\{\hat{\Delta}_{\hat{\pi}} \}$ are pull together to learn domain-specific character and introduce repulsion between $\{\hat{X}_{\hat{\pi}} \}$ and $\{\hat{\Delta}_{\hat{\pi}}\}$ to enforce disentanglement. (c) The quality of $\hat{X}_{\hat{\pi}}$ is measeured through segmentation utility. (d) The disentanglement is achieved by fusing randomly paired $\hat{X}$ and $\hat{\Delta}$ to reconstruct the input.
• Figure 2: Qualitative visualization through (a) broad and (b) intricate comparisons.
• Figure 3: Qualitative visualization of brain tumor segmentation.