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

2025-05-26

AI Summary

Medical image segmentation models exhibit poor generalization across devices and clinical centers. Existing domain generalization (DG) methods either lack theoretical guarantees or yield limited improvements via data augmentation. This paper proposes a manifold-alignment data augmentation method grounded in Langevin dynamics: it is the first to integrate energy-based models (EBMs) with Langevin sampling to generate intermediate samples with consistent distributions across multiple source domains. We theoretically prove that the method bounds the Rademacher complexity of generalized linear models, with regularization strength governed by the intrinsic dimensionality of the underlying data manifold. Evaluated on benchmark tasks—including fundus image segmentation and 2D MRI prostate segmentation—the proposed approach significantly outperforms state-of-the-art DG methods. Moreover, it synergizes effectively with domain randomization to further enhance performance.

Technology Category

• Machine Learning: Learning with Manifolds
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Search and Retrieval-Augmented AI: Vertical and domain-specific search
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

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 $ extbf{Lang}$evin $ extbf{D}$ata $ extbf{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.

Problem

• Improves multi-source domain generalization in medical image segmentation
• Generates synthetic samples using Langevin dynamics and EBMs
• Addresses limitations of representation learning and data augmentation

Innovation

• Leverages Energy-Based Models for data augmentation
• Uses Langevin dynamics to generate intermediate samples
• Complements domain-randomization approaches effectively