LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation

Medical image segmentation models often struggle to generalize across different domains due to various reasons. Domain Generalization (DG) methods overcome this either through representation learning or data augmentation (DA). While representation learning methods seek domain-invariant features, they often rely on ad-hoc techniques and lack formal guarantees. DA methods, which enrich model representations through synthetic samples, have shown comparable or superior performance to representation learning approaches. We propose LangDAug, a novel Langevin Data Augmentation for multi-source domain generalization in 2D medical image segmentation. LangDAug leverages Energy-Based Models (EBMs) trained via contrastive divergence to traverse between source domains, generating intermediate samples through Langevin dynamics. Theoretical analysis shows that LangDAug induces a regularization effect, and for GLMs, it upper-bounds the Rademacher complexity by the intrinsic dimensionality of the data manifold. Through extensive experiments on Fundus segmentation and 2D MRI prostate segmentation benchmarks, we show that LangDAug outperforms state-of-the-art domain generalization methods and effectively complements existing domain-randomization approaches. The codebase for our method is available at https://github.com/backpropagator/LangDAug.

When hospitals adopt AI tools for medical image analysis, they often discover a frustrating problem: systems that work perfectly at one facility fail at another. This happens because medical scanners, like cameras with different settings, produce images with subtle but important variations.We tackled this challenge by developing LangDAug, a technique inspired by how physicists model particle movement. Instead of trying to eliminate differences between imaging styles, we embrace them. LangDAug creates a spectrum of synthetic training images that gradually morph from one hospital's imaging style to another's, teaching AI to handle the full range of natural variations.Our method proved highly effective in experiments with retinal scans for eye disease and MRI scans for prostate analysis. AI systems trained with LangDAug maintained their accuracy even when tested on images from completely new hospitals they had never encountered. This breakthrough could help deploy medical AI more broadly, ensuring that advanced diagnostic tools benefit patients everywhere, not just those at well-resourced medical centers.