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

TL;DR

This work addresses the challenge of generalizing medical image segmentation models across unseen domains by proposing LangDAug, a Langevin data augmentation scheme that uses Energy-Based Models to traverse and interpolate between multiple source domains. By training EBMs with contrastive-divergence and sampling intermediate Langevin iterations, LangDAug generates domain-bridging samples that augment the standard ERM objective, yielding regularization effects and bounds tied to the data manifold's intrinsic dimensionality. Theoretical results show the augmented risk induces Hessian-smoothing and reduces Rademacher complexity, while empirical evaluations on fundus and 2D MRI prostate segmentation demonstrate state-of-the-art domain generalization and strengthened compatibility with domain randomization. LangDAug offers a principled, augmentation-based route to robust cross-domain medical image segmentation, with practical implications for clinical deployment and potential extensions to 3D data and shared-domain conditioning to improve scalability.

Abstract

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 (DAug). While representation learning methods seek domain-invariant features, they often rely on ad-hoc techniques and lack formal guarantees. DAug methods, which enrich model representations through synthetic samples, have shown comparable or superior performance to representation learning approaches. We propose LangDAug, a novel $\textbf{Lang}$evin $\textbf{D}$ata $\textbf{Aug}$mentation 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.

Figures (21)

• Figure 1: Overview of the proposed method: (a) We run Langevin dynamics (LD) using trained EBMs to transverse between different domains. The intermediate iterates of LD (called Langevin samples) are stored and used for augmentation, (b) The segmentation model is trained with original domain samples as well as Langevin samples, (c) The trained model is then deployed on unseen target domains.
• Figure 2: t-SNE visualization comparing domain distributions: (a) the original four domains, and (b) LangDAug generated augmented domains. Origin $x$ denotes samples generated by LangDAug starting from Domain $x$.
• Figure 3: Visualization of retinal fundus segmentation performance of different domain generalization methods.
• Figure 4: Visualization of prostate segmentation performance of different domain generalization methods.
• Figure 5: Ablation analysis across four domains. The first row shows the effect of varying the number of Langevin steps ($k$); the second row shows the effect of varying the step size ($\beta$); the third row depicts the impact of changing the number of convolutional blocks; and the fourth row presents the effect of varying the number of Langevin samples per chain. Metrics reported are mIoU and mDSC.

Theorems & Definitions (13)

• Theorem 4.1
• Corollary 4.2
• Theorem 4.3
• Corollary 4.4
• Theorem 1.1
• proof
• Corollary 1.2
• proof
• Theorem 1.3
• proof