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

TL;DR

This paper introduces a probabilistic and contrastive learning framework that enhances domain generalization in medical image segmentation by reducing reliance on domain alignment and improving robustness to domain shifts.

Contribution

It proposes a novel combination of uncertainty modeling, contrastive learning, and frequency-domain structural enhancement to improve domain generalization in medical image segmentation.

Findings

Significant improvement in segmentation accuracy across domains.

Enhanced robustness to style variations and domain shifts.

Effective preservation of structural details in images.

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.