From Deterministic to Probabilistic: A Novel Perspective on Domain Generalization for Medical Image Segmentation

TL;DR

This work tackles domain generalization in medical image segmentation by shifting focus from domain alignment to data representation quality. It introduces probabilistic representation learning that combines deterministic features with mean and covariance to model uncertainty, paired with distribution-level contrastive learning to align feature distributions across domains. A wavelet-based structural preservation module preserves high-frequency structural details during style augmentation, mitigating distortions. Empirical results on Fundus and Prostate datasets demonstrate superior Dice and ASD metrics, validating improved cross-domain robustness and boundary delineation. The approach offers a practical path toward more reliable multi-site medical image analysis without heavy reliance on domain alignment.

Abstract

Traditional domain generalization methods often rely on domain alignment to reduce inter-domain distribution differences and learn domain-invariant representations. However, domain shifts are inherently difficult to eliminate, which limits model generalization. To address this, we propose an innovative framework that enhances data representation quality through probabilistic modeling and contrastive learning, reducing dependence on domain alignment and improving robustness under domain variations. Specifically, we combine deterministic features with uncertainty modeling to capture comprehensive feature distributions. Contrastive learning enforces distribution-level alignment by aligning the mean and covariance of feature distributions, enabling the model to dynamically adapt to domain variations and mitigate distribution shifts. Additionally, we design a frequency-domain-based structural enhancement strategy using discrete wavelet transforms to preserve critical structural details and reduce visual distortions caused by style variations. Experimental results demonstrate that the proposed framework significantly improves segmentation performance, providing a robust solution to domain generalization challenges in medical image segmentation.

Figures (3)

• Figure 1: Traditional domain alignment methods struggle to completely eliminate domain shifts, leading to sensitivity to feature distribution deviations in the model. In this figure, features of the same category exhibit significant distribution differences across different domains. Direct use of point estimation may result in inaccurate classifications. For example, the yellow sample in the middle is incorrectly classified into the green category because it is closer to the green prototype. However, by introducing probabilistic modeling and uncertainty estimation, the model can identify the differences in feature distributions. With the consideration of uncertainty, the yellow sample is more appropriately classified into the blue category, thereby improving the model’s adaptability and classification accuracy under domain shift conditions.
• Figure 2: The overall architecture of our method. (a)Probabilistic Representation Learning: We construct probabilistic representation learning by leveraging both the certainty and uncertainty in the latent space of images to explore the intrinsic relationships within the training data, mitigate domain shift, and fully realize the potential of the model. (b)Wavelet-Enhanced Structural Preservation: Using discrete wavelet transform (DWT) to separate style features from high-frequency structural information in the image, this method incorporates high-frequency components from source domain images into intermediate transformation results to enhance the expression of high-frequency information, thereby effectively preserving structural details.
• Figure 3: Visual comparison for Fundus and Prostate segmentation task. The red contours indicate the boundaries of ground truths while the green and blue contours are predictions.