LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
Authors: Piyush Tiwary
Kinjawl Bhattacharyya
Prathosh AP
Venue: ICML 2025 poster (2025)
Keywords: Medical Image Segmentation
Domain Generalization
Energy-based Models
Primary area: applications->health_medicine

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 (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 **Lang**evin **D**ata **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.

LangDAug:
Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
Piyush Tiwary' Kinjawl Bhattacharyya'
Prathosh A.P.
Abstract
supervised learning approaches leveraging Convolutional Neural Networks (CNNs) (Milletari et al., 2016; Falk et al., 2019; Zhou et al., 2019; Jin et al., 2019; Isensee et al., 2020; Shen et al., 2017; Olabarriaga & Smeulders, 2001; Asgari Taghanaki et al., 2021; Tiwary et el., 2023). However, persistent challenge remains: these models often struggle to generalize from the training set (source domain) to unseen test sets (target domain) (Ganin et al., 2016; Kamnitsas et al., 2017: Li et al., 2018b).
Medical image segmentation models often strug- gle to generalize across different domains due to various reasons. Domain Generalization (DG) methods overcome this either through representa- tion learning or data augmentation (DAug). While representation learning methods seek domain- invariant features, they often rely on ad-hoc tech- niques and lack formal guarantees. DAug meth- ods, which enrich model representations through synthetic samples, have shown comparable or superior performance to representation learning approaches. We propose LangDAug, novel Langevin Data Augmentation for multi-source domain generalization in 2D medical image seg- mentation. LangDAug leverages Energy-Based Models (EBMs) trained via contrastive divergence to traverse between source domains, generating intermediate samples through Langevin dynam- ies. 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 seg- mentation and 2D MRI prostate segmentation benchmarks, we show that LangDAug outper- forms sfate-of-the-art domain generalization meth- ods and effectively complements existing domain- randomization approaches. The codebase for our method is available at https: //gi thubb. com/ba comm//backoroopggattr//aagggao.
The generalization problem stems from the limitations of the Empirical Risk Minimization (ERM) framework, which often performs poorly on outto--distribution (OOD) test or target domains (Shalev-Shward & Ben-David, 2014; Mur- phy, 2022; Hart al., 2000; Hajek & Raginsky, 2021). This is primarily due to the i.i.d. assumption underlying ERM, where models minimize the loss solely with respect to the training data distribution. As result, ERM-trained models exhibit degraded performance on test data from different distributions, violating the i.i.d. assumption and leading to sub-optimal generalization unseen domains (Vapnik, 1999).
The challenge is particularly pronounced in medical imag- ing, where domain shifts significantly impact model per- formance. These shifts arise from factors such as differ- ences in imaging protocols across institutions, variations in equipment specificatiion, diverse patient attributes (e.g., age, gender, demographics), and disparities between modal- ities (e.g., CT, MRI, Ultrasound). As noted by Guan & Liu (2021), such shifts hinder the deployment of machine learning models in new clinical settings. The discrepancies between training and target domains can result in incon- sistent performance, potentially compromising the model's clinical utility and patient care. Addressing these issues critical for developing robust and generalizable models.
1. Introduction
Recent years have seen remarkable advancements in the field of Medical image segmentation, particularly through
Various approaches have been proposed address domain generalization. A key strategy is representation learning, which seeks to learn either domain-invarianl invariant (DI) features or sufficiently diverse features to generalize to unseen do- mains. These methods often disentangle and domain-specific (DS) features, with some approaches discarding Nomain-specine features entirely and others com- bining DI-DS pairs to create richer features. Another ap-
'Deaartment of Electrical Communication Engineering, Indian Institute of Science, Bangalore, India, Correspondence to: Piyush Tiwary <piyushtiwary@ @ iisc ac.......
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

LungDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
(a) LangDAug
(b) Training
Domain
Domain
B² Xt+1 Xt vx (x,) Be
Source LangDAur
(c) Inference
Targe
Domain
Figure I: 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.
proach data augmentation, which enhances model rep- resentations by supplementing original samples with aug- mented ones. As shown by (Gulrajani & Lopez-Pazz 2021), data augmentation can achieve comparable performance to
We provide theoretical insights into LangDAug's ef- fects, demonstrating its regularization impact ERM training. This leads to minima, aiding generalization to unseen domains.
and flatter
representation learning while being simpler to implement.
3. We conduct rigorous evaluations on two medical im- age segmentation benchmarks: fundus segmentation and 2D MRI image segmentation. Our results show that LangDAug outperforms state-of-the-art domain generalization methods and can enhance the perfor- mance of approaches.
Our work adopts the data augmentation approach to im- prove multi-source domain generalization in 2D medical image segmentation. We propose novel Langevin Data Augmentation (LangDAug) scheme that leverages Langevin dynamics (LD) iterates augment the original dataset. Specifically, we train Energy-Based Models (EBMs) to tra- verse between different source domains using a contrastive- divergence (CD) objective. This traversal achieved via iterative gradient-based MCMC methods like LD, with the iterates representing samples from the manifold bridging the source domains. These samples then incorporated into the ERM framework for downstream task training, poten- tially enhancing the model's ability to generalize to unseen domains. Our key contributions are:
based generalization
2. Related Works
2.1. Domain Generalization
Domain Generalization (DG) frameworks (Motiian et al., 2017; Wang et al., 2022; Zhou al., 2022a; Dou et al., 2019; Gu al., 2021; Cong al., 2022) aim to train single model using one multiple source domains achieve robust performance unseen target domains. These methods typically focus on either learning domain-invariant features or developing diverse features capable of generalizing to target domains. Meta-learning DG methods (Li et al., 201 8a; Liu et al., 2021b;a; 2020; Finn et al., 2017; Balaji et al., 2018: Du et al., 2020b; Tiwary al., 2024; 2025) simulate
1.
We introduce LangDAug. augmentation technique for multi-source domain generalization in 2D medi- cal image segmentation. By training EBMs for each pair of source domains via CD, we generate intermedi- ate domain samples that augment the source domains within the ERM framework.
2

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
domain shifts by splitting source data into meta-train/test sets.
(Xu et al., 2021; Liu al., 2022; Peng et al., 2021; Zhang et al., 2020b; Zhou al., 2020c; Rohey al., 2021; Liu et al., 2018; Fick et al., 2021; Li et al., 2022a; Zhou et al., 2022b; Liu et al., 2021a) are often sought in medical settings, to enhance training set diversity. These methods are catego- rized (Wang et al., 2022) into (i) data augmentation (DAug) and (ii) data generation. DAug techniques (Honarvar Nazari & Kovashka, 2020: Prakash et al., 2019; Zhou al.. 2020a; 2023) apply some combination of operations such as flip- ping, rotation, scaling, cropping, or model-based transfor- mations to reduce overfitting, while data generation meth- ods (Rahman et al., 2019; Somavarapu et al., 2020; Zhang, 2017: Zhou et al., 2020b; Li et al., 2021a) utilize generative models to create diverse data that mirrors the training dis- tribution. MTS (Li et al., 2022a) enhances model-agnostic meta-learning for DG via task augmentation through mixed task sampling and meta-update objective, effectively mit- igating task-level overfitting. RandConv(Xi et al., 2021) uses multi-scale random convolutions as a DAug technique to improve model robustness. FedDG (Liu et al., 2021a) adopts a federated DG strategy that preserves privacy by exchanging amplitude spectra between clients and employs continuous frequency space interpolation and boundary- oriented episodic learning to bolster model generalization to unseen domains.
Current state-of of-the-art DG approaches can be broadly cate- gorized into two main types (Wang et al., 2022): representa- tion learning-based methods and data manipulation- methods. Representation learning methods (Li et al., 2018c; Muandet et al.. 2013; Upchurch et al., 2017: Gardner et al., 2015; Zhang et al., 2022: Li et al., 2022b; b021b; Zhou et al., 2021; Wang et al., 2020; Chen et al., 2023; Pan et al., 2018; Hu et al., 2022) are further subdivided into domain-invariant (DI) representation learning and feature disentanglement approaches. Domain-invariant methods (Blanchard et al., 2021; Hu et al., 2020; Li et al., 2018b; Garg et al., 2021; Pan et al., 2018; Arjovsky et al., 2019) minimize cross-domain feature discrepancies--e.gg Fish aligns inter-domain gradi- ents, Fishr regularizes gradient variances, and Hutchinson matches Hessian matrices. While effective, these may sac- rifice domain-specift (DS) features to balance single- and multi-domain performance. Another line of work uses test time adaptation (TTA) for DAug using EBMs (Xiao et al., 2023; Yuan et al., 2024; Tiwary et al., 2023). F.g., Xiao al. (2023) first trains EBM on source domain and during testing, they use LD convert test samples using these trained EBMs. However, this method has a fundamental flaw akin to that of data coverage problem faced in score- based models (Song& & Ermon, 2019). Specifically, since the trained EBMs have never seen the test data, the assigned energy to test data manifold could be arbitrary. This renders TTA unreliable for domain adaptation.
3. Proposed Methodology
3.1. Problem Setup and Method Overview
We consider the standard domain generalization setup in which, we have n domains (also called source domains) (De)== where each domain admits and underlying prob- ability density Dj PD. (z, on common support X' x y. Further, we are given n, samples from do- main D; that are independent and identically distributed,
Feature disentanglement approaches (Zunino et al., 2021; Peng et al., 2020; Ilse et al., 2020; Qiao et al., 2020; Zhang et al., 2021; Mahajan et al., 2021) aim to separate repre- sentations into domain-invariant and domain-specific com- ponents, allowing combinations that yield diverse repre- sentations for downstream tasks. Mixstyle (Zhou et al., 2021) facilitates generalization by blending instance-level feature statistics across domains, generating novel styles in feature space. RAM (Zhou et al., 2022b) promotes seg- mentation generalization by mixing amplitude spectra from different domains and incorporating a domain-specific im- age restoration module. TriD (Chen et al., 2023) supports generalization in 2D medical image segmentation through uniform sampling and channel-wise style mixing of feature statistics, broadening feature space coverage and enhancing robustness. Our work focuses on this promising stream of data manipulation methods and propose novel approach called Langevin Data Augmentation (LangDAug) for do- main generalization.
1e } is to learn parametric mapping fo(-): : X + y such that it minimizes the empirical risk. Particularly, the parame- ters 0 are obtained via solving the following optimization problem:
{x,,
Lid p (x, The goal of an ERM-framework
N Σ
arg min ([(())))) vo arg min
N (fa(()))
#
where N N= E. Formally, the True risk is approximated with an unbiased estimate using law of large numbers to get the Empirical risk. This is done in order to obtain a tractable objective function which can be minimized using gradient-based optimization methods. However, since the approximation only takes the provided n domains into ac- count, it is not able to generalize on unseen domain (also called target domain) D, & {D.)i- The goal of do- main generalization is to aleviate this issue by obtaining
2.2. Data Manipulation based DG
Owing to the difficulty in obtaining large datasets, data manipulation-bast based domain generalization (DG) methods
3

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
parameter 0 that is able to generalize well on unseen do- mains.
D, and then runs for K steps using Egg as the kernel to approximate sampling from Pa.. Mathematically,
To address this issue, we adopt a aatta-aggmenttiinnnbo based approach by introducing new samples in the ERM frame- work to cover broader support in the X' x y domain, improv- ing the True risk estimate. Specifically, we train Energy- Based Models (EBMs) to traverse between source domain pairs using gradient-based Markov Chain Monte Carlo (MCMC) methods, such as Langevin dynamics (LD). The intermediate iterates generated during this process act as "bridge' samples between domains. We leverage these in- termediate samples as augmentation data, referred to as Langevin-dos data Augmentation, to enhance the model's gener- alization to unseen target domains.
a² Vx, Eo.g (Xx) + aee X = X 2
where, Xo ~ PD. (z)). Thus, we learn B,y using Eq. 1, where the second term is approximated using samples from LD initialized with a sample from D, The LD iteration can be interpreted as starting from sample under distribution Pd; and gradually transforming into a sample from PD. In other words, the LD updates traverse from domain D; to domain Dj. The intermediate samples generated during the I.D process can be viewed as forming bridge between the two domains, We train such an EBM for all 2(()) possible pairs within n source domains.
While LangDAug can be potentially used for normal DG scenarios, we note that there are certain factors that guided our application design. The effectiveness of LangDAug de- pends on the ability of EBMs to effectively traverse and in- terpolate between source domain distributions. Such traver- sal becomes especially natural when source domains share structured similarities or consistent underlying factors of variation. Medical imaging data typically demonstrates structured variations. predominantly reflected through dif- ferences in amplitude spectra across domains (Zhou et al., 2022b; Liu et al., 2021a). EBMs have been shown to excel at capturing and modeling variations in amplitude spectra (Tan- cik et al., 2020; Du et al., 2020a), making them suited for the domain variations characteristic of medical imaging data. Thus, LangDAug's 's capability aligns well with the domain shift challenges encountered specifically in medical image segmentation.
3.3. Langevin Data Augmentation
As explained previously, we learn an EBM for transversal between each pair of domains. Here, we explain how to generate the augmentation samples. Consider the EBM Fe.,, which can transverse from domain D, to domain Dy. To generate samples for augmentation, we take sample {z E DD and run LD for K steps, starting from Xy:
B
x x)
x For (xx) + Be, where, x² = Xj
2
(4)
The above process gives rise to K-intermediate samples (x) We use the samples {x),,, V,) P..K " (named 'Langevin' data) for augmentation during FRM-based learn- ing On a closer look, we can see that for a given k. D A {xx). yx 1;= follow the same distribution they are obtained via k-step LD updates. This can be viewed as samples from a new domain altogether. Hence, we can see that we are able to obtain data from new novel domains by transversing between the original source domains. While the original dataset consisted of Dare A {
3.2. EBMs for Inter-Domain Transversal
We train Energy-Based Models (EBMs) (Du & Mordatch, 2019; Zhao & Chen, 2021) to traverse between pairs of source domains. Consider two distinct domains D; and DDj where i * j. To train an EBM, Eo., that can transverse from Dj to Dy. we employ the contrastive divergence loss objec- tive (Hinton, 2002), LCD. The gradient of this objective is:
U.., the augmented dataset consists of Daug A U D This shows that we
inf
can generate novel data-points by simply transversing be- tween source domains in meaningful way. We refer the reader to Supplementary for all the implementation details.
Ve, = E [Vo,, (x)] - JE Pass
Eo,, (x))
Po;
exp (- Eo,, (x))
4. Theoretical Analysis
where,Pat,
Zoo,, = Ix exp (- Eo., (xx) dx)
Analysis Overview: Our theoretical analysis reveals how LangDAug naturally induces regularization and improves generalization. We follow steps similar to Arora et al. (2021); Zhang et al. (2020a) by first demonstrating that
We approximate the expectations in Eq. I using finite sample averages from their respective distributions. Since sampling directly from Po., is intractable due to the partition function Ze., we employ MCMC methods, specifically Langevin dynamics (LD).
Note that we retain original labels while augmentation be- cause, empirically observe that the Langevin samples preserve the content of x, (see Supplementary for visualizations).
The LD iteration is initialized with sample from domain

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
LangDAug regularizes the directional derivatives of fo(-) for any parametric model (Thm. 4. This insight leads us to examine generalized linear models (GLMs), where we characterize these regularization effects more precisely (Cor. 4.2). We then establish bounds on the Rademacher complexity the resulting function class (Thm. 4.3) and use these to bound the generalization gap between true and empirical risk (Cor. 4.4).
The regularization terms reveal how LangDAug controls model behavior through derivatives of the prediction func- tion. To understand these effects more concretely, we ex- amine GLMs, which model predictive densities using ex- ponential families with linear sufficient statistics (Murphy, 2022):
p(y1x,0) r(x) exp (y0*x x A(6 x))
(8)
Notations: We denote general parametric loss function (0,2% z), where 0 E θ C Rd and z (x, y) denotes the input output pair. Further, we consider the training dataset D = yx E X' x y such that they follow an un-
where r(x) is the base measure and A(-) is the log parti- tion function. The negative log- oggliikelhhod takes the form A(f y0 allowing us to apply Thm. 4.1 to obtain: Corollary 4.2. For a GLM, if A() is twice differentiable. then the regularization terms obtained via second-order approximation is given by:
derlying distribution, zi yx) lid PP(()). Further, we denote (3) x,
gt x log p(x,) + e 10, which is the sample obtained by running one-step LD start- ing from Xy ~ (z( (we drop the dependence on / for brevity). We denote 2; = yz) as the data pair used for augmentation. Further, we denote the model parameterized by as fo X - V/ the gradient w.r.t x and 0 is denoted by fzx and (o)) respectively.
2 V
RGLM & 82 k (A" (1000x)))' x,) A'(F s(x:))
Σ
2n 1=1
where, s(x,) V log p(x:) is the Stein's 's score function.
Having obtained the above regularization term, we use the techniques in Arora et al. (2021); Zhang et al. (2020a) to bound the Rademacher complexity of the following class of functions which are used in optimization of dual of above regularization term:
Lastly, we denote the true risk as L()) [[[((, z)] where the expectation is w.r.t true data distribution. The standard empirical and the augmented empirical risks are denoted by:
L (0,
Σ
2(0,z:)
(5)
W/ 4
x x 0 (0 $n}
1:
L (0,
1
Σ
E (0.1)
[e(0, 2.)]
(6)
The following provides an upper-bound on Rademacher complexity of such class of functions:
k
Theorem 4.3. Assume that the distribution of x, is p retentive, and let Σ, = E[xx"] have bounded singular values. Further assume that the norm of the parameters and E log p(x)][?] are bounded. Then the empirical Rademacher complexity of W,, satisfies:
We start our alalysis by characterizing the regularization effect of LangDAug:
Theorem 4.1. Consider a real-valued loss function of the form E(0, y)) = h((f(xx)) (x)) where h(-) and fo(-) are twice differentiable for all 0 E Q Then given dataset D = (i). the augmented empirical risk as denoted in Eq. 6 can be written as:
/ (rank(Enn)
Rad(W,, D) S D) SCV
(9)
k
where, C= 7
1/2
((2
7
1/1
3 cas (0. = Ltd (dd (0, + Σ D)
(10)
(7)
po
i=1
here, a denotes the lowest singular values of Ex.
where,
Now, using this result, we can directly use the results of Bartlett & Mendelson (2002) to bound the generalization 3 - log pap as follows:
&
1
(0, =
T
i=1
Corollary 4.4. If A(-) is L Lipcchhit continuous, X. 3, θ are all bounded, then there exists constant L, B > 0 such that for all 9 satisfying the constraint in Way we have:
R2(0, (
rank(2 (cvronnk2)))
log (1/8)
R2 = -0
L < Ltd + 2LLA
B
k
2k
with probability atleast 8.
5

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
We note that the appearance of rank(EE)) rather than the full dimensionality of x suggests that LangDAug's gener- alization bound depends on how many "true" degrees of freedom the data has (intrinsic dimension) rather than how many numbers we use to store each data point (ambient dimension). The dependence on intrinsic dimensionality rather than ambient dimension particularly valuable for high-dimensional problems where the data lies on lower- dimensional manifold (Gorban & Tyukin, 2018 8; Cayton et al., 2008).
1.00
1.00
t-SNE Dimensionss
6-ONE Dimension'
0.50
0.50
origl
Drgae
Organ
Doran
0.00
0.00
Darano
0.00
1-SNE Dimension
1:5NE Dimension 1
(4) Original Domains
(9) Original Domains + LangDAug
Figure t-SNE visualization comparing domain distribu- tions: (a) the original four domains, and (b) LangDAug generated augmented domains. Origin x' denotes samples
generated by LangDAug starting from Domain
5. Experiments
5.1. Datasets and Baselines
Domain D Domain C Domain B Domain
Datasets and Metrics: We evaluate LangDAug's effective- ness under domain shift using two widely-used medical imaging datasets: Retinal Fundus Segmentation (RFS) (Al- mazroa al., 2018; Orlando al., 2020; Zhang et al., 2010; Sivaswamy et al., 2014) and Prostate MRI Segmen- tation (Liu et al.. 2020). For both datasets, we employ a leave-one-out protocol (Gulrajani & Lopez-Paz, 2021) to assess domain generalization performance.
The RFS dataset challenges models to segment the optical cup (OC) and optical disc (OD) in retinal images. It contains samples from four clinical sites, each representing distinet domain with its own train/test split. We use the training splits for training and resize all images to 256 x 256 pixels during preprocessing. Segmentation accuracy is evaluated using Intersection-over-nneee Union (IoU) and Dice Similarity Score (DSC) metrics.
GT
RandCony
MixStyle
FedDG
RAM
TriD
Ours
Figure Visualization of retinal fundus segmentation per- formance of different domain generalization methods.
employs federated learning for domain generalization; and TriD (Chen et al., 2023), which exploits distribution-level information. Since, we follow similar steps Zhang et al. (2020a) for our analysis² we also compare method with RAM (Zhou et al., 2022b) which is a variant of MixUp designed for medical image segmentation. Additionally, we provide comparisons with CORAL (Sun & Saenko, 2016), RSC (Huang et al., 2020). For fair comparison we use ResNet34 model (He al., 2016) the segmentation back- bone for all the methods. All the implementation details including hyperparameters and Ablation studies are pro- vided in Supplementary.
The Prostate MRI dataset includes 116 T2-weighted MRI scans from six clinical sites (domains). Preprocessing in- volves cropping each scan to the 3D prostate region, extract- ing 2D axial slices, and resizing them to 384 x 384 pixels. We train on these 2D slices and evaluate using: (1) DSC between predicted and ground truth slices, and (2) Average Surface Distance (ASD) for reconstructed 3D volumes, ob- tained by concatenating model-predicted slices. Results are reported the mean of three independent runs.
Baselines: We evaluate our method against two categories of domain generalization approaches. First, we consider three recent methods from the DomainBed benchmark (Gul- rajani Lopez-Paz, 2021): Hutchinson (Hemati et al., 2023), which matches domain statisties using Hessian in- formation; Fish (Shi et al., 2022), which employs gradient- based domain alignment; and Fishr (Rame al., 2022), which utilizes gradient variance matching. Further, we compare against recent approaches specifically tailored for medical image segmentation: RandConv (Xu et al., 2021), which leverages random convolution for style in- variance; MixStyle (Zhou al., 2021), which performs feature-level style mixing: FedDG (Liu al., 2021a), which
5.2. Results
t-NE Visualization: Before presenting the main results, we visualize the RFS dataset samples in the feature space using t-SNE plots (Figure 2). The original domain sam- ples show a concentrated distribution, while Langevin sam- ples significantly expand feature space coverage. Notably, Langevin samples bridge gaps between domains, consistent with the framework in Section 3.2. This inter-domain traver- sal property allows the model to learn a more comprehensive
whill the proof steps share similarities, niques and arrive an independent verification the result
different tech-
6

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
representation of the domain space. These observations sup- port our hypothesis on the effectiveness of LangDAuu in im- proving domain generalization. Visualizations of Langevin samples are provided in Supplementary.
Domain F Domain Domain D Domain C Domain B Domaln A
5.2.1. FUNDUS SEGMENTATION
We provide the results of optical cup (OC) and optical disc (OD) segmentation on retinal fundus images in Table We present the metries in the form of (OC, OD) for con- venience, further, we also provide the mean metric across OC and OD. We also provide visual qualitative compari- son of segmentation results in Figure 3. We observe that LangDAug performs most consistently across all domains.
Among the DomainBed approaches (Hutchinson, Fish, and Fishr), performance is generally underwhelming. Hutchin- son performed the poorest, with average IoU and DSC of 67.39 and 78.14 respectively. While Fishr showed some im- provement over the others, reaching average IoU and DSC of 72.28 and 82.16 respectively. Overall these methods lagged behind other approaches. This pattern was consistent across different domains and metrics.
GT
RandCony
MisStyle
FedDG
RAM
TriD
Ours
Figure 4: Visualization of prostate segmentation perfor- mance of different domain generalization methods.
mDSC values dropping to approximately 78. In contrast, LangDAug maintains remarkable consistency across both domains, achieving competitive mDSC of 86.99 on Domain A while maintaining its performance with mDSC of 85.87 on Domain B.
Domain B emerges as the most challenging for generaliza- tion, with most methods failing to achieve mDSC above 80. The reason for this is significantly different conditions under which the fundus images are acquired. In other words, the distribution of Domain B is very different from that of other domains. In this domain, LangDAug demonstrates superior performance with an mIoU of 75.05, surpassing the next best method (RAM) by 1. 1.26%. For optical disc seg- mentation, LangDAug achieves the highest DSC of 91.64, followed by RAM at 90.52. While RAM leads in optical cup segmentation with a DSC of 82.82, LangDAug maintains competitive performance at 80. 11.11, demonstrating consistent effectiveness across all metrics in this challenging domain. This also shows robustness of LangDAug in limited source distribution setting where source domains are not able to cover the target domain. In such scenarios, previous method fail to generalize effectively whereas LangDAug maintains edge over these methods.
LangDAug demonstrates superior overall performance, achieving the highest average loU (78.84%) and DSC (87.61%) across all domains. Notably, exhibits better stability in cross-domain performance, as evidenced by sig- nificantly lower standard deviations (mIoU: =2.43, mDSC: 1.89) compared to its counterparts TriD (mloU: #3.76. 76, mDSC: :.511 12) and RAM (mloU: +22.9, mDSC: +3.42). This reduced variance underscores LangDAug's consistent generalization capabilities.
5.2.2. PROSTATE SEGMENTATION
Table presents our results for 2D MRI prostate segmenta- tion. We provide visual results in Figure 4. Consistent with the findings from fundus segmentation, the DomainBed methods (Hutchinson, Fish, and Fishr) generally demon- strate lower performance, with Fishr showing the best per- formance and Hutchinson showing the worst performance among the three.
In Domain C, LangDAug demonstrates good performance with the highest mloU of 81.01, while maintaining com- petitive mDSC metrics second only to TriD. The method shows balanced effectiveness across both optical cup and disc segmentation, surpassing established approaches like FedDG and RAM. This performance extends to Domain D, where LangDAug leads in both mloU and mDSC metrics, outperforming comparable methods like TriD and RAM.
LangDAug demonstrates best performance across Domains A through D. consistently outperforming the second-best methods in both ASD and DSC. Specifically, in Domains A and B, LangDAug achieves approximately 0. 0.1mm mm lower ASD and nearly 0.6% higher DSC compared to the next best method, TriD. In Domain D, LangDAug surpasses other methods with a substantial reduction of 0.14mm in ASD and an increase in DSC by over 4%. Even in the
While established methods MixStyle, RandConv, and TriD demonstrate marginally better performance in Domain A (mDSC: 88.62, 87.90, and 87.80 respectively), they exhibit substantial performance deterioration in Domain B, with
7

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
Table 1: Cross-domain generalization performance on Retinal Fundus Segmentation. IoU (%) (1) and DSC (%) (f) are presented in pairs (zz y) where, x denotes OC and y denotes OD segmentation metric. Bold (x) and Underlined (2) values indicate best and second-bes besttd-bstttse performance respectively for each metric in their respective columns. Average denotes the average metric across all domains.
Domain
Average
mloU
DSC
BDDSC
mlol:
mDDSC
milol
MDSCC
mloU
D88C
waDISC
IoU
DSC
CORAL
(06 30, 901
75:37
(78-79,
(41.11)
86.55
(54.3) T2,
65:99
(95.155
70.20
(7) (1,.34,0 86.54)
T8.
[62 82,
5M1
67.17
(6T
885
75.85
(50. 01:
15.05
74.35 11:17
83.51
RSC
76.82
94.24)
(53.
71.74:
(64.94
(72.75
12)
(68.19.
74.70
(50.70)
73.45
92.97
SagNet
(6c.45,
P5.02
12.61)
54:51
52
00,75,05
90.90...
[55.1 79.
07.31
(70.04,
SWAD
(8.032.87.00)
76.19
(77.00,
22.50)
$4.90
(56.60, 76.011
(63.21,
76.07
(07.91. 85.88
76.01
(80
66.90
(61.57.
72.00
(75.57,
82.82
72.63
82.54
Hutchinsnn
(55 78:27)
08.73
(88-10)
s0.51)
77:46
(55
78-
06
(63.
85.81)
77. 16
[56: 82.54
08:30
170
17)
99:17
(50.45, 72.300
08.72
(64.47) 83 63-007
77.16
07:33
T6.14
Fish
(S1.16.85.77) 55.73)
74.44
176.76,
91.36)
83.96
(48.91. 74.87)
04.44
[61.35 87.01)
70.70
(63.)). 82.92)
73.25
(79.99,
88.87
84.43
(36.27.
80 (00
72.50
(70.54, 80, (4)
53.24
11.16
$1.81
Fishe
(65.35, 50.35)
75.85
(77.55,
(1.16)
$4.30
(54.28, 73 22)
64.08
(65.17.82.592)
73.75
(68.3 62.95
76.63
[60 00, 65.45)
61.00
(6:.00
53.11)
73.60
(70.90,
01 (02)
85.50
72.28
82.16
RandCCov
(71.46. (1)
79:06
(0.51)
(50 00, 80 285
6T 022
(95:35,
113
78:23
(07. 05 055
70.73
[801
14
95.07
(6T
55 61)
T4.73
(7A.
15, 11 (20)
84.84
14.76
84.20
MisSiste
171.76,59.77
80.76
(82.08, 34.56)
68.60 62
(88.16, 17.11.77.2 23)
67.69
86.71
79.96
(71.65, 87.93)
70.79
(68.38
08.39 10)
58.98
167.56. 56.62)
77.09
(50.00,
02, 72)
80.37
75.03
85.08
FedDG
(75.35 50.01)
76.05
(78.52, $1.12)
81.47
(68
75
72.11 8:
(17.8),
85.55)
53.08
(00 02, 62.25)
10
(65 135.11.33 87 32)
85.21
(60.01 52.91)
75.96
(8)
88
84.90
75.21
83.07
RAM
(67.28, 57.
77:42
(72.2) M6.11
85.67
78.74
(77.
77.30
85,20
TriD
171.58.
60.92
76.50
72
(65.01.
90.421
(72.66.
08
(70.52
(53.08,
87.76
17.92
65.05
Langhhing
(08.31, (9))00)
79.79
(79.75, 941.21)
(70.61, T9. 50)
55.83
(73..7
94]
15.00
70,
80.
93)
Table 2: Cross-domain generalization performance on T2-weighted Prostate MRI Segmentation. ASD (mm) (L) and DSC
(%) (1) are presented in pairs (ar, v) where, x denotes OC and y denotes OD segmentation metric. Bold (æ) and Underlined (z) values indicate best and second-best performance respectively for each metric in their respective columns. Average denotes the average metric across all domains.
B
C
D>
F
F
Average
Unseen Domain
ASD
DSC
ASD
DSC
ASD
DSC
ASD
DSC
ASD
DSC
ASD
DSC
ASD
DSC
Hutchinson
3.28
82.26
1.48
85.85
2.07
72.65
3.98
77-41
2.78
71.75
1.04
81.79
2.54
78.62
Fish
1.58
84.40
1.46
85.00
3.74
7U.39
3.37
78.45
3.72
74.26
0.76
83.43
2.44
79.37
Fishr
1.10
88.03
0.94
88.41
2.07
83:21
1.08
83.94
1.78
81.41
0.61
84.41
1.26
81.90
RandCony
1.27
67.90
1.18
88.24
1.90
$1.96
0.84
84.44
1.54
83.04
0.44
86.37
1.20
85:33
MixStyle
0.72
88.59
0.88
88. 12
1.62
87.10
0.65
85.62
1.59
82,18
0.51
85.72
1.00
86.27
FedDG
86.10
87.03
98.23
0.88
87.32
1.73
81.35
0.50
85,27
1.07
85.95
RAM
0.93
68.72
0.98
87.20
1,26
88.78
0.74
80.18
1.78
81.01
0.32
87
1.00
87.02
TriD
91:52
0.72
87.71
0.71
89.45
1.48
82.41
85.41
0.90
87.68
LangDAng
0.58
92.15
0.64
90.31
1.21
89.53
0.57
93.41
1.49
83.17
0.38
86.41
0.81
89.16
Table 3: Cross-domain generalization performance on Retinal Fundus Segmentation of Domain Randomization method, with and without LangDAug. Values in parentheses for Average columns show absolute improvement over base method.
Uaseen Damain
Average
mlee
mBSC
adMi
RDBCC
DDSC
mnleU
DSC
nDBSC
aBBC
FeIDG
(73.3) 80.013
75:61
(75.5% 5L 12)
51.47
(16-7):
75 16)
7B. 14
(77
88 05)
80.00
(00.00
e0 76)
76 10
(83 11, 57.33)
85-91
(00.0) 83:913
75.96
(S1 43, 56 IT)
$1.50
75.91
50 07
FelDG + Ours
32.37)
78.82
79.066 85.29)
82.47
(70.03,
76.10)
75.06
(70
12.
90
03)
84
(71 10.40.44 84 411
78.10 (69.68, 58.88
56.18
(70.37,
00.87.58.22 58.22)
76,50
(82. (1.87)
57.05
70.89 (+1.34) 58.14
(1147
(07 83.336
77.42
(7300
(58.98, 75.0)
T8
(72.9)
III
(7) 85
T6.74
(TTUT) 91.54)
90.31)
(+1.0)
TrID
(71.81.9001)
90.9
(78.4) 86.00
57.50
(58.3).
72.45
[60.0L
00 42)
78.21
(72.6) 08)
70.34 (96.22 21.83)
90/12
(00.2287
75.96
(55.05, (8.40)
S7.76
77.50
85.06
Trill Ouns
(72.8)
82.81
(80.54, N7 N2) Wo.le
(41.67,
52) H0.09
(73 95)
81.98
(75 W7.33
80.30 (67 06.35;
91.63
(71 69.04)
80.42
(85. 60, (4.00)
NO.74
SU.T (+11)) 8M 36 (+2.33)
Table 4: Cross-domain generalization performance on 2D MRI Prostate Segmentation of Domain Randomization method, with and without LangDAug. Values in paren- theses for Average columns show absolute improvement over base method.
challenging Domain C, LangDAug maintains significant performance advantage, outperforming other methods by O. 18mm in ASD and 1.82% in DSC. Furthermore, in the most difficult Domain E (identified by the highest average ASD and lowest average DSC across methods), LangDAug secures the best DSC, demonstrating its performance con- sistency across domains.
ASD
DSC
ASD
ASD
PHEEU
01:0
1.5
0.5
1/3
16:21
THEDG
0.50
16.24
8:55
60:33
135
0623
175
86.1
135
13.00
3:45
86.3
While RAM achieves the best performance in Domain F (ASD: 0.32 mm, DSC: 87 87.25%), LangDAug remains com- petitive in this domain (ASD: 0.38 mm, DSC: 86.41% 41%) and demonstrates the most consistent performance across all domains. This consistency is highlighted by LangDAug's
BAM
4.33
M1.2
11.3
9.7
321
16.30
1.3
1/2
0.
30115
M111
1:15
3810
143
1.4
10.09
131
(-44))
8

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
best overall average performance across domains, which is lower by 0.09 mm in ASD and higher by 1.48% in DSC compared to the second-best TriD, similar to its trend in Fundus segmentation results.
to 0.66mm, DSC increase: 87.03% to 89.31%), with overall metrics improving by mm.13mm ASD and +1.40% DSC. Notably, while Domain E looks most challenging across all methods, LangDAug consistently enhances performance, underscoring its efficacy in challenging cross-domain sce- narios.
5.3. LangDAug with Domain Randomization
LangDAug, described in Section 3. generates samples from intermediate domains using Langevin dynamics. This can enhance domain randomization methods, which combine domain-invariant (DI) and domain-specific (DS) features from source domains. The key assumption is that test domain features lie within the feature space formed by mixing DI-DS features from source domains. Thus, more source domains produce more diverse features, improving target domain performance. By creating new domains via Langevin dynamics. LangDAug enhances domain random- ization methods. We validate this on retinal fundus and prostate segmentation datasets, testing with three leading methods: FedDG (Liu et al., 2021a), RAM (Zhou et al., 2022b), and TriD (Chen al., 2023).
These observations conclusively show that LangDAug sig- nificantly boosts the performance of all domain randomiza- tion baselines.
6. Conclusion
In this work, we propose LangDAug, a novel Langevin data augmentation method for multi-source domain generaliza- tion in 2D medical image segmentation. LangDAug lever- ages Energy-Based Models (EBMs) trained via contrastive divergence to generate intermediate samples between source domains using Langevin dynamics. These samples act as bridges across domain distributions. We establish a theoreti- cal foundation for LangDAug, showing its induced regular- ization effect on parametric models. For Generalized Linear Models (GLMs), we prove that LangDAug's 's regularization terms upper-bound the Rademacher complexity based on the data manifold's ' intrinsic dimension. Comprehensive experiments on retinal fundus and prostate MRI segmen- tation demonstrate LangDAug's superiority over existing methods. Moreover, Lang effectively complements domain randomization approaches, achieving ofate-of-the- art performance in domain generalization tasks.
Fundus Segmentation: Table 3 presents the performance comparison of domain randomization methods augmented with LangDAug on fundus segmentation. LangDAug demonstrates consistent performance improvements across all baseline methods and domains. Specifically, FedDG aug- mented with LangDAug shows significant enhancements of +1.38% in average IoU and +1.47% in average DSC across all domains. RAM exhibits even more substantial gains, with improvements of + +.86% and +2.04% in aver- age IoU and DSC respectively, notably achieving an mDSC increase from 84.87% to 87.71% on Domain A. Most re- markably. TriD combined with LangDAug demonstrates the highest improvements of +3.05% and +2.21% in aver- age loU and DSC respectively, with particularly significant enhancement in Domain B (mDSC increase from 78.21% to 81. 81.98%). These consistent improvements across methods and domains demonstrate LangDAug 's effectiveness in enhancing domain randomization methods.
Limitations: While LangDAug delivers strong perfor- mance, it has two primary limitations. First, the larger number of training samples results in longer training times, though selective sampling can mitigate this. Second, the number of required EBMs scales with the source domains, posing scalability challenges. Future work could address this by adopting shared architectures with domain condi- tioning. We provide detailes quantitative computational cost in Supplementary. Further, we process 3D volumes 2D slices. Future work could explore training EBMs directly on 3D data to bypass slicing and better model 3D spatial relationships.
Prostate Segmentation: Table 4 demonstrates LangDAug' effectiveness in enhancing cross-domain generalization for prostate segmentation. Across all six domains, LangDAAg consistently improves both ASD and DSC metrics, paral- leling the improvements observed in fundus segmentation. The most substantial improvements are observed with TriD, yielding average ASD reduction of 0.14mm and DSC im- provement of +2.90%. RAM augmented with Lang shows similar gains, with average ASD and DSC improve- ments of -0.1 7mm and +1.67% respectively, particularly notable in Domain E where ASD decreased from 1.78 to 1.41mm. FedDG, despite moderate baseline performance, demonstrates significant improvement when combined with LangDAug, especially in Domain B (ASD reduction: 0.93
Acknowledgements
This work was supported (in part for setting up the GPU compute) by the Indian Institute of Science through a start- up grant. Piyush is supported by Government of India via Prime Minister's Research Fellowship. Kinjawl is supported by Pre-doctoral fellowship provided by Kotak IISc AI-MI. Center (KIAC). Prathosh is supported by Infosys Foundation Young investigator award.
9

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
Impact Statement
learning. Journal of machine learning research, 22(2): 1-55, 2021.
LangDAug advances domain generalization in medical im- age segmentation by enabling robust adaptation to diverse imaging environments, a critical challenge in deploying AI models across clinical settings. By generating domain- bridging samples through Langevin dynamics and Energy- Based Models, the method enhances segmentation accuracy in scenarios affected by domain shifts (e.g., variations in scanners or protocols), directly supporting more reliable di- agnostics and treatment planning. Its theoretical grounding in bounding model complexity via intrinsic data dimension- ality offers a principled framework for future research.
Cayton, L. et al. Algorithms for manifold learning. eSchol- arship, University of California, 2008.
Chen, Z., Pan, Y., Ye, Y. Cui, H., and Xia, Y. Treasure in distribution: a domain randomization based multi-source domain generalization for 2d medical image segmenta- tion. In International Conference on Medical Image Com- puting and Computer-Assisted Intervention, pp. 89-99. Springer, 2023.
Cong. C., Liu, S., Di leva, A., Pagnucco, M., Berkovsky, S., and Song, Y. Colour adaptive generative networks for stain normalisation of histopathology images. Medical Image Analysis, 82:102580, 2022. ISSN 1361-8415, doi: URL http://db...gg11..000///mmmddddd 102580. acbeoee/artcce///////11111222222
While computational costs and scalability limitations cur- rently constrain its adoption, LangDAug's compatibility with domain randomization and its ability to preserve anatomical fidelity position it as a versatile tool for real- world applications. Future efforts to optimize training ef- ficiency and extend the framework 3D volumetric data could further unlock its potential, paving the way for Al systems that generalize seamlessly across evolving medical imaging technologies.
Dou, Q. Coelho de Castro, D., Kamnitsas, K., and Glocker, B. Domain generalization via model-agnostic learning of semantic features. Advances in neural information processing systems, 32, 2019.
References
Du, Y. and Mordatch, I. Implicit generation and modeling with energy based models. Advances in Neural Informa- tion Processing Systems, 32, 2019.
Almazroa, A., Alodhayb, S., Osman, E., Ramadan, E., Hum- madi, M., Dlaim, M., Alkatee, M., Raahemifar, K., and Lakshminarayanan, V. Retinal fundus images for glau- coma analysis: the riga dataset. In Medical Imaging 2018: Imaging Informatics for Healthcare, Research, and Ap- plications, volume 10579, pp. 55 SPIE, 2018.
Du, Y., Li, S., and Mordatch, I. Compositional visual gen- eration with energy based models. Advances in Neural Information Processing Systems, 33:6637-6647, 2020a.
Arjovsky, M., Bottou, L., Gulrajani, I., and Lopez- Paz, D. Invariant risk minimization. arXiv preprint arXiv: /907:029900089 2019.
Du, Y., Xu, J., Xiong, H., Qiu, Q., Zhen, X., Snoek, C. G., and Shao, L. Learning to learn with variational informa- tion bottleneck for domain generalization. In Computer Vision-ECCV 2020: 16th European Conference, Glas- gow, UK, August 23-28, 2020, Proceedings, Part X 16, pp. 200-216. Springer, 2020b.
Arora, R., Bartlett, P., Mianjy, P., and Srebro, N. Dropout: Explicit forms and capacity control. In International Conference on Machine Learning, pp. 351-361. PMLR, 2021.
Falk, T., Mai, D.,, Bensch, R., Çiçek, Ö., Abdulkadir, A., Marrakchi, Y., Böhm, A., Deubner, J., Jäckel, Z.,, Seiwald, K., et al. U-net: deep learning for cell counting, detection, and morphometry. Nature methods, 16(1):67-70, 2019.
Asgari Taghanaki, S., Abhishek, K., Cohen, J. P., Cohen- Adad, J., and Hamarneh, G. Deep semantic segmentation of natural and medical images: review. Artificial Intelli- gence Review, 54:137-178, 2021.
Fick, R. H., Moshayedi, A., Roy, G., Dedieu, J., Petit, S., and Hadj. S. B. Domain-specific cycle-gan augmentation improves domain generalizability for mitosis detection. In International Conference on Medical Image Computing and Computer-Assisted Intervention, pp. 40-47. Springer, 2021.
Balaji, Y., Sankaranarayanan, S., and Chellappa, R. Metareg: Towards domain generalization using meta-regularization. Advances in neural information processing systems, 31, 2018.
Bartlett, P. L. and Mendelson, S. Rademacher and gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research, 3(Nov):463-482, 2002.
Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta- learning for fast adaptation of deep networks. In Interna- tional conference on machine learning, pp. 1126-1135. PMLR, 2017.
Blanchard, G., Deshmukh, A. A., Dogan, U., Lee, G., and Scott, C. Domain generalization by marginal transfer
10

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
Ganin, Y., Ustinova, E., Ajakan, H., Germain, P., Larochelle, H., Laviolette, F., March, M., and Lempitsky, V. Domain- adversarial training of neural networks. Journal of ma- chine learning research, 17(59):1-35, 2016.
Honarvar Nazari, N. and Kovashka, A. Domain generaliza- tion using shape representation. In European Conference Computer Vision, pp. 666-670. Springer, 2020.
Hu, S., Zhang, K., Chen, Z., and Chan, Domain gener- alization via multidomain discriminant analysis. In Un- certainty in artificial intelligence, pp. 292-3922 PMLR, 2020.
Gardner, J. R., Upchurch, P., Kusner, M. J., Li, Y., Wein- berger, K. Q.. Bala, K., and Hopcroft, J. E. Deep manifold traversal: Changing labels with convolutional features. arXiv preprini orXiv:1511.000
2015.
Hu, S., Liao, Z., Zhang. J., and Xia, Y. Domain and content adaptive convolution based multi-source domain general- ization for medical image segmentation. IEEE Transac- tions on Medical Imaging, 42(1):233-2444 2022.
Garg, V., Kalai, A. T., Ligett, K., and Wu, S. Learn to expect the unexpected: Probably approximately correct domain generalization. In International Conference on Artificial Intelligence and Statistics, pp. 35842. PMLR, 2021.
Huang, Z., Wang, H., Xing, P., and Huang, D. Self- challenging improves cross-domain generalization. In Computer DesionEECVV 2020: 16th European conference, Glasgow, UK, August 23-28, 2020, proceedings, part Il 16, pp. 124-140. Springer, 2020.
Gorban, A. N. and Tyukin, I. Y. Blessing of dimensionality: mathematical foundations of the statistical physics of data. Philosophical Transactions of the Royal Society A: Mathemutical, Physical and Engineering Sciences, 376 (21 (2118):2170237 2018.
Ilse, M., Tomezak, J. M., Louizos, C., and Welling, M. Diva: Domain invariant variational autoencoders. Medical Imaging with Deep Learning, pp. 322-348. PMLR, 2020.
Gu, R., Zhang, J., Huang, R., Lei, W., Wang, G., and Zhang, S. Domain composition and attention for unseen-domain generalizable medical image segmentation. In Medical Image Computing and Computer Assisted Intervention- MICCAI 2021: 24th International Conference, Stras- bourg, France, September 27-October 1, 2021, Proceed- ings, Part III 24, pp. 241-250. Springer. 2021.
Isensee, F,, Jaeger, P. F., Kohl, S. A. A., Petersen, J., and Maier-Hein, K. nnu-net: a self-configuring method for deep learning-based biomedical image segmentation. Nature Methods, 18:203 211, 2020. URL https: apii www.nt org/ /ttS:/oooST:::::999909000 774477.
Guan, H. and Liu, M. Domain adaptation for medical image analysis: survey. IEEE Transactions on Biomedical Engineering, 69(3):113311185, 2021.
Jin, Q.. Meng, Z., Pham, D., Chen, Q., Wei, L., and Su, R. Dunet: A deformable network for retinal vessel segmentation. Knowledge-Base Systems, 188:149-162 2019.
Gulrajani, I. and Lopez-Paz, D. In search of lost domain generalization. In International Conference on Learning Representations, 2021. URL https: //ppeeeeeeeeeeeeeeeeeeeee net/ met/forue
Kamnitsas, K., Baumgartner, C., Ledig, C., Newcombe, V., Simpson, J., Kane, A., Menon, D., Nori, A., Criminisi, A., Rueckert, D., et al. Unsupervised domain adaptation in brain lesion segmentation with adversarial networks. In Information Processing in Medical Imaging: 25th International Conference, IPMI 2017, Boone, NC, USA, June 25-30, 2017, Proceedings 25, pp. 597-609. Springer, 2017.
Hajek, B. and Raginsky, M. Ece 543: Statistical learning theory. Class notes for ECE, 2021.
Hart, P. E., Stork, D. G., and Duda, R. O. Pattern classifica- fion. Wiley Hoboken, 2000.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual leam- ing for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Li, C., Lin, X., Mao, Y., Lin, W., Qi Q. Ding, X., Huang, Y., Liang, D., and Yu, Y. Domain generalization on medical imaging classification using episodic training with task augmentation. Computers in biology and medicine, 141: 105144, 2022a.
Hemati, S. Zhang, G., Estiri, A., and Chen, X. Under- standing hessian alignment for domain generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision pp. 19004419014, 2023.
Li, D., Yang, Y., Song, Y.-Z., and Hospedales, T. Learning to generalize: Meta-learning for domain generalization. Proceedings of the AAAI conference on artificial intel- ligence, volume 32, 201 8a.
Hinton, G. E. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8):1711- 1800, 2002.

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
Li, H., Pan, S. J., Wang, S., and Kot, A. C. Domain general- ization with adversarial feature learning. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 1400-5409, 2018b.
Mahajan, D., Tople, S., and Sharma, A. Domain generaliza- tion using causal matching. In International conference machine learning, pp. 7313-7324. PMLR, 2021.
Milletari, F., Navab, N., and Ahmadi, S.-A. V-net: Fully con- volutional neural networks for volumetric medical image segmentation. In 2016 fourth international conference on 3D vision (3DV), pp. 565-571. Icee, 2016.
Li, L., Gao, K., Cao, J., Huang, Z., Weng. Y., Mi, X., Yu, Z., Li, X., and Xia, B. Progressive domain expansion network for single domain generalization. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 224-233, 2021a.
Motian, S., Piccirilli, M., Adjeroh, D. A., and Doretto, G. Unified deep supervised domain adaptation and gen- eralization. In Proceedings of the IEEE international conference computer vision, pp. 5715-5725, 2017.
Li, P., Li, D., Li, W., Gong, S., Fu, Y., and Hospedales, T. M. A simple feature augmentation for domain general- ization. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 8886-889) 0866-889,22 2021b.
Muandet, K., Balduzzi, D., and Schölkopf, B. Domain generalization via invariant feature representation. In International conference on machine learning, pp. 10-18. PMLR, 2013.
X., Dai, Y., Ge, Y., Liu, J., Shan, Y., and Duan, L.-Y. Un- certainty modeling for out-of -f--orribittttiionnnn generalization. arXiv preprint arXiv: 033588, 2022h.
Murphy, K. P. Probabilistic Machine Learning: An intro- duction. MIT Press, 2022. URL probml .ai.
Li, Y., Gong, M., Tian, X., Liu, T., and Tao, D. Domain generalization via conditional invariant representations. In Proceedings of the AAAI conference on artificial intel- ligence, volume 32, 2018c.
Olabarriaga, S. D. and Smeulders, A. W. Interaction in the segmentation of medical images: A survey. Medical image analysis, 5(2):127- 142, 2001.
Liu, A. H., Liu, Y.C.., Yeh, Y.-Y., and Wang, Y.C. A unified feature disentangler for multi-domain image trans- lation and manipulation. Advances in neural information processing systems, 31, 2018.
Orlando, J. I., Fu, H., Breda, J. B., Van Keer, K., Bathula, D. R., Diaz- Diaz-Pinto, A., Fang, R., Heng, P.-A., Kim, J., Lee, J., et al. Refuge challenge: A unified framework for evaluating automated methods for glaucoma assessment from fundus photographs. Medical image analysis, 59: 101570, 2020.
Liu, Q. Dou, Q. and Heng. P.-A. Shape-aware meta- learning for generalizing prostate mri segmentation to unseen domains. In Medical Image Computing and Com- puter Assisted Intervention-MICCAl 2020: 23rd Inter- national Conference, Lima, Peru, October 4-8, 2020, Proceedings, Part 23, pp. 475-485, Springer, 2020.
Pan, X., Luo, P., Shi, J., and Tang, X. Two at once: En hancing learning and generalization capacities via ibn-net. Proceedings of the european conference on computer vision (ECCV), pp. 464 479, 2018
Peng, D., Lei, Y., Liu, L., Zhang, P., and Liu, J. Global and local texture randomization for synthetic-to-real semantic segmentation. JEEE Transactions on Image Processing, 30:6594 6608, 2021.
Liu, Q., Chen, C., Qin, J., Dou, Q., and Heng, P-A. Feddg: medical image seg-
Federated domain generalization mentation via episodic learning in continuous frequency space. In Proceedings of the JEEE/CVF conference on computer vision and pattern recognition, pp. 1013-1023, 2021a.
Peng. X., Li, Y., and Saenko, K. Domain2vec: Domain embedding for unsupervised domain adaptation. In Com- puter Vision-ECco 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part VI 16, pp. 756-774. Springer, 2020.
Liu, X., Thermos, S., O'Neil, A., and Tsaftaris, S. A. Semi-supervised meta-learning with disentanglement for Nomain-generalisse medical image segmentation. In Medical Image Computing and Computer Assisted Intervention-MICCA. 2021: 24th International Confer- ence, Strasbourg, France, September 27-October 1, 2021, Proceedings, Part 11 24, pp. 307-317. Springer, 2021b.
Prakash, A., Boochoon, S., Brophy, M., Acuna, D., Camer- acci, E., State, G., Shapira, O., and Birchfield, S. Struc- tured domain randomization: Bridging the reality gap by context-aware synthetic data. In 2019 International Con- ference on Robotics and Automation (ICRA), pp. 7249- 7255. IEEE, 2019.
Liu, X.-C., Yang, Y.-L., and Hall, P. Geometric and textural augmentation for domain gap reduction. In Proceedings of the JEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 14340-143500 2022.
Qiao, F., Zhao, L., and Peng, X. Learning to learn single domain generalization. In Proceedings of the IEFF/CVF
12

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
conference on compuler vision and pallern recognition, pp. 13556-12565, 2020.
Tiwary, P., Bhattacharyya, K., and AP, P. Cycle con- sistent twin energy-based models for image-to-image translation. Medical Image Analysis (MedIA), 91, 2023. doi: http:///o..ngg//0..lo///mmmttdt: URL. scterce//artideee//////111111111
Rahman, M. M., Fookes, C., Baktashmotlagh, M., and Srid- haran, S. Multi-component image translation for deep domain generalization. In 2019 IEEE Winter Confer- ence Applications of Computer Vision (WACV), pp- 579-588. IEEE, 2019.
Tiwary, P., Shubham, K., Kashyap, V., and AP, P. Bayesian pseudo-coresets via contrastive divergence. In Confer- ence on Uncerlainity in Artificial Intelligence (UAI), 2024, URL https: http//aaxvvv..//////..///.. org/ 20bb/2330 /tt//abbb/33.12222 11278.
Rame, A., Dancette, C., and Cord, M. Fishr: Invariant gradient variances for duitt---istiiuutiin generalization. In International Conference on Machine Learning, pp- 18347-18377. PMLR, 2022.
Tiwary, P. Guha, A., Panda, S., and AP, P. Adapt then unlearn: Exploring parameter space seman- ties for unlearning in generative adversarial net- works. Transactions on Machine Learning Research, 2025, URL https: / http://wwweeiii.eeeeeee nnt forum? id= (o--jAHBBvvbbose 739258832239..
Robey, A., Pappas, G. J., and Hassani, H. Model-based domain generalization. Advances in Neural Information Processing Systems, 34:20210-2023) 2021.
Shalev-Shwartz, S. and Ben-David, S. Understonding Ma- chine Learning: From Theory to Algorithms. Cambridge University Press, 2014.
Upchurch, P., Gardner, J.. Pleiss, G., Pless, R., Snavely, N., Bala, K., and Weinberger, K. Deep feature interpolation for image content changes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7064-7073, 2017.
Shen, D., Wu, G., and Suk, H.-I. Deep learning in medical image analysis. Annual review of biomedical engineering, 19(1):221-2444 2017.
Shi, Y., Seely, J.. Torr, P., N. S., Hannun, A., Usunier, N., and Synnaeve, G. Gradient matching for domain gen- eralization. In International Conference on Learning Representations, 2022. URL ht.ts:::/// http////eeeeee.eee//e/eee/ net/ BWW 9HmO
Vapnik, V. N. An overview of statistical learning theory. IEEE transactions on neural networks, 10(5):988-99) 1999.
Wang, J., Lan, C.,, Liu, C., Ouyang, Y., Qin, T., Lu, W., Chen, Y., Zeng, W., and Philip, S. Y. Generalizing to unseen domains: A survey on domain generalization. IEEE transactions on knowledge and data engineering, 35(8):8052-807) 2022.
Sivaswamy, J., Krishnadas, S., Joshi, G. D., Jain, M., and Tabish, A. U. S. Drishti-gs: Retinal image dataset for optic nerve head (onh) segmentation. In 2014 IEEF Ith international symposium on biomedical imaging (ISBI), pp. 53-56. IEEE, 2014.
Wang, S., Yu, L., Li, K., Yang, X., Fu, C.-W., and Heng, P.-A. Dofe: Nomain-oricnter feature embedding for gen- eralizable fundus image segmentation on unseen datasets. IEEE Transactions on Medical Imaging, 39(12):4237- 4248. 2020.
Somavarapu, N., Ma, C.-Y., and Kira, Z. Frustratingly simple domain generalization via image stylization. arXiv preprint Xii:20006.11007 2020.
Song. Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Xiao, Z., Zhen, X., Liao, S., and Snoek, C. G. M. Energy- based test sample adaptation for domain generalization. In The Eleventh International Conference on Learning Representations, 2023. URL https: pttttp:://ppezevviii /wwp//veeenee.eeeeeeeeeeee net / forum? 3dw.ek.kedddddddd....
Sun, B. and Saenko, K. Deep coral: Corrllatinn alignment for deep domain adaptation. In Computer vision-ECCV 2016 workshops: Amsterdam, the Netherlands, October 8-10 and 15-16, 2016, proceedings, part III 14, pp. 443- 450. Springer, 2016.
Xu, Z.. Liu, D., Yang. J., Raffel, C., and Niethammer, M. Robust and generalizable visual representation learning via random convolutions. In International Conference on Learning Representations, 2021. URL https:// openreview. net /eee/foornee forum? Bat.rmmmm.ddBVSS003.....
Tancik, M., Srinivasan, P., Mildenhall, B., Fwi.ooiiioooaaaa Keil, S., Raghavan, N., Singhal, U., Ramamoorthi, R., Barron, J., and Ng, R. Fourier features let networks learn high fre- quency functions in low dimensional domains. Advances in neural information processing systems, 33:7537-7547, 2020.
Yuan, Y., Xu, B., Hou, L.. Sun, F., Shen, H., and Cheng, X. Tea: Test-time energy adaptation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 23901-23911, 2024.
13

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
Zhang, H. mixup: Beyond empirical risk minimization. arXiv preprint arXiv: 1710.00991 2017.
Zhou, K., Liu, Z., Qiao, Y., Xiang, T., and Loy, C. C. Do- main generalization: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(4):4336- 4415, 2022a.
Zhang, L., Deng, Z., Kawaguchi, K., Ghorbani, A., and Zou, J. How does mixup help with robustness and generaliza- tion? arXiv preprint arXiv:2010.04819 2020a.
Zhou, K., Loy, C. C., and Liu, Z. Semi-supervised domain generalization with stochastic stylematch. International Journal of Computer Vision, 131(0):2377-2337, 2023.
Zhang, L., Wang, X., Yang, D., Sanford, T., Harmon, S., Turkbey, B., Wood, B. J., Roth, H., Myronenko, A., Xu, D., et al. Generalizing deep learning for medical image segmentation to unseen domains via deep stacked trans- formation. IEEE transactions on medical imaging, 39(7): 2531 2531, 2020b.
Zhou, Z Siddiquee, M. M. R., Tajbakhsh, N., and Liang. J.
Unet++: Redesigning skip connections to exploit multi- scale features in image segmentation. IEEE transactions medical imaging, 3966, 2019.
Zhou, Z., Qi, L.,, and Shi, Y. Generalizable medical image segmentation via random amplitude mixup and domain specific image restoration. In ECCV, 2022b.
Zhang, X., Cui, P., Xu, R., Zhou, I.., He, Y., and Shen, Z. Deep stable learning for out-of-distribution general- ization. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition. pp. 5372- 5382, 2021.
Zunino, A., Bargal, S. A., Volpi, R., Sameki, M., Zhang, J., Sclaroff, S., Murino, V., and Saenko, K. Explainable deep classification models for domain generalization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 2233-3242, 2021.
Zhang, Y., Li, M., Li, Jia, K., and Zhang, L. Exact fea- ture distribution matching for arbitrary style transfer and domain generalization. Proceedings of the IEEE/VV conference on computer vision and pattern recognition, pp. 8035-8045 2022.
Zhang, Z., Yin, F. S., Liu, J., Wong, W. K., Tan, N. M., Lee, B. H., Cheng, J., and Wong, Y. Origa-light: An online retinal fundus image database for glaucoma analysis and research. In 2010 Annual international conference of the IEEE engineering in medicine and biology, pp. 3065- 3068. IEEE, 2010.
Zhao, Y. and Chen, C. Unpaired image-to-image transla- tion via latent energy transport. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 16418-16427, 2021.
Zhou, K., Yang, Y., Hospedales, T., and Xiang, T. Deep dowain-adversarril image generation for domain general- isation. In Proceedings of the AAAI conference on artifi- cial intelligence, volume 34, pp. 13025-13032, 2020a.
Zhou, K., Yang, Y., Hospedales, T., and Xiang, T. Learning to generate novel domains for domain generalization. In Computer Vision ECcn-EECCV/ 2020: 16th European Confer- ence, Glasgow, UK, August 23-28, 2020, Proceedings, Part XVI 16, pp. 561-578. Springer, 2020b.
Zhou, K., Yang, Y., Hospedales, T., and Xiang, T. Deep domain-adversarial image generation for domain general- isation. Proceedings of the AAAI conference on artifi- cial intelligence, volume 34, pp. 13025-13032 2020c.
Zhou, K., Yang, Y., Qiao, Y., and Xiang, T. Domain gener- alization with mixstyle. ICLR, 2021.
14

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
A. Proof of Theoretical Results
Next, for the second derivative, we differentiate the above expression:
Theorem A.1. Consider a loss function of the form E(O, (x, = h(fo(x))) {()) where h(-) and fo(-) are twice differentiable for all E O. Then given a dataset D {x..y.... the augmented empirical risk as denoted in Eq. 6 can be written as:
1e (B) = 1n" (6666)))
+ (N'((O(8))) - vi (o
-[N(fo(3))) -y. OV log p()) )' (r()))))
3 CaMs (0,D) = = (0, (0, + Σ R,(O,D)
(II)
+ Vlogg((x; )Vogg(x))
(21)
1=1
= (((()()))))( (V. )7) re(
where,
(n(((())))))))))) - v1)' ( fo (2(
1 k
(1(0,D)
Σ
8 (K'(66(x)) -xV16(()))))) Volgp(x))
-[N((F((2)) yY. V.) V log)(x;)
T
(22)
R2(0,D) = k 1=1
Σ
Substituting B = in the above gives:
= 11" =0) -('([(x)))e""" (V.()))))))
D)) = 1
ox)
+ (f((())))) (V (o=))(x)))
- -[n'(Joxx)) - y.) (xx(xx) ''V V logg(xx) (23)
Proof. By definition, we have:
Taking expectation w.t.e:
L (0, - Σ E [e(0,2, 2)]
(12)
Elch (0) = n"" (fo(x.)) (VV
k (~.()
+ - - VV log p(x,)
fo(x,)
Consider one term of the above summation:
E(0, 2.) = h (x. - (f - (3) (13)
(24)
82 2 w/ (0) + (B) (14)
(0) + (0)
Now, combining results from Eq. 15, 20, 24 into Eq. 14 and Eq. 6:
where, we have used Taylor expansion to approximate she around 3 = 0 and lim s03 = O. We will evaluate each of
L (0, Σ ((((..2))
8-0
the above terms below.
Recall that (3) x,
8² ((x + Bs, hence,
8²
(h' (f(((x))) (f.
VV
(0) = x, Using this, we have:
82
(0) = h(fe(xx)) - vx(f(xx,)) = E(0,z,) z,)
(15)
h" ('(f6(x)) (T
2
Now,
82
(h' (o((o(x)) y;) Tr Tr TT fo(x;))) (25)
+
T ax;
2
- -1 (fo(z))))
0x;
03
= L (0, L.dd (0, D) + Σ R,(0, D)
(xx:))))
T
0x;
yi
0%:
03
(16)
(26)
[h' (fo(x))) - y. y (xx )²r
3V log p(x))]
(17)
[n' (fo(x.) y.] (x. T (B)
(18)
Corollary A.2. For a GLM, if A(-) is twice differentiable. then the regularization terms obtained via second-order approximation is given by:
where, re() _ = - sV log lxg = 08. Hence, we
have:
(0) = (h' (fo(x:)) y.) \z,
(19)
B2
(A" xx)00' 0 A' x:)0= s(xi))
(27)
= B0)) (1 (fo(x;)) - Vo)x [T =
2m
(20)
where, s(x,) = V p(ox)(x;) the Stein's 's score function.

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
Proof. Using the result of Theorem 4.1 with h(-)) A())
Combining this with constraint in Eq. 36, we get:
and fo(x) 0 xxx,,
x, we get:
(1,02E,0} S mmax { 2
|1013 - p (1,
8a
(0 p0 E
E(8, 2,) = 2(0,2;)
A (0 x x,) 0² V p(x.)
2
(39)
3² A" (00 x²)
(28)
1101|
7 V P
7
(40)
2
pa
where we use the fact that for GLMs, so(- = 0 and Tr(00')= = o This completes the proof.
where a is the lowest singular value of Ez. Now, the empir- ical Rademacher complexity is given by:
Theorem A.3. Assume that the distribution of x; is p- retentive, and let Σg = E[xx"] have bounded singular values. Further assume that the norm of the parameters and Ey log p(x)]l² are bounded. Then the empirical Rademacher complexity of Wy satisfies:
Rod(W,,, D) = E 5
sup Wi,
Σ
5:0"
Xi
(41)
To
EsOTxx X;
1F
sup
(42)
8
|10||
T:
(ank(EE))
Rad(W, D) S CV.
(29)
1/2
E
Ex,,I |1
where, C=
(2)
1/2
V
1/4
(30)
po
1=1
(43)
po
1 C
here, a denotes the lowest singular values of Σx.
S
E ||
Σ
ExxI1
(44)
Proof. The function class under consideration is:
S C
1
Σ
xi
(45)
W, 2 x {x- 0- TT E A" (0=x)) A' (0(T00))(0 x)x( s $
Taking expectation over dataset distribution completes the proof:
The above constraint can be written as:
y E[A" (0"x) x))0=0 - EJA' (0 x)0 sx))
(31)
C
Rad(W D) = Raddd D)}=
Σ E E x1
E[A" (0 x))0000 - V (0T x E(OT s(x)
(46)
(32)
C
2 E[A" (0x x))000 - BLAA (0T x
1012
[1
EE
|3]
E(Tr(xxx2") )
(47)
a
x
(33)
C
2 E[A" (0"xx)00 x)]00000-000 - 1110 101001
(34)
k
Tr(E.)
(48)
E[A" (0 x))000
1 11011
B( ((AA'(TTx))
(35)
122
C
(49)
kc
2 |01|2 (ETA"(0Tx2)
Kl
B(A ((A(()))))3
(36)
K2
where, R. [1]))]]]
S Ko and 11012 z K2
From p-
Lemma A.4 (Bartlett & Mendelson (2002)). For any B- uniformly bounded and FLLipchit,fuunction G. for all o E 4 with probability atleast -8:
retentiveness, we have:
pmin x)) < E A (0
A"
x)
K1 A' (0T x)
K2
log( (1/8)
(37)
B[((o(x)) S
Σ ((o(xx)) + 2aRd + BV
K
2k:
pmin(1,0')) 0'',0}} < S A" (0T x)
K1
E
[A'
(0)x)2]
(50)
K2
(38)
2

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
Corollary A.5. If A(-) is LA Lipchitz continuous, X, Y θ are all bounded, then there exists constant I,, B 0 such that for all 0 satisfying the constraint in Wy, we have:
samples. Specifically, we store 13 samples for fundus seg- mentation (uniformly at iterations i = 3, .,.6....300 ,3) and 5 samples for 2D MRI prostate segmentation (at iterations k = 5. ... 40). each EBM to save computational time in downstream seg- mentation network training. We provide a pseudo-code in Algorithm
These samples are precomputed for
L S Cald +2LLA (cv
(ann (
log log 1/))
+ B
2k
(51)
Segmentation Network: For downstream segmentation, employ ResNet34-based network (He al., 2016)². It comprises of a ResNet34 encoder and ASPP-based decoder. The network is trained using a combination of cross-entropy and DICE loss, with the AdamW optimizer. The learning rate (Ir) is selected from {1 x 10 10-4,1 x 10-9), and the batch size (bs) is chosen from (8,32, 64). The running average parameters are set to 31 0.9 and 32 0.99. Empirically, we observed that (Ir = x 10-4, bs 8) yields optimal performance for fundus segmentation, whereas (Ir = 1 x 10-6, bs = 32/64) works best for 2D MRI prostate segmentation.
with probability atleast 1-6. 8.
Proof. The results follow directly from Lemma A.4 and result of Theorem 4.3.
0
B. Implementation Details
This section outlines the implementation specifics of LangDAug. The code and relevant resources are avail- able https: htmpp:///qitee siighuuu.... wwt.ccom/ccc..ece cco/bbakkroopggtoox. LangDAug.
Baselines: For all the baselines, we use the same ResNet34- based segmentation network. The DomainBed methods (Fish, Fishr, and Hutchinson) are primarily used in classi- fication tasks. For segmentation, methods to ensure compatibility with the retinal fundus and prostate datasets, utilizing the ResNet-34 backbone. These adaptations were closely aligned with the original imple- mentations provided in DomainBed In addition to the default implementation, we re-scaled the losses for better stability and training of network.
Energy-based Models: LangDAug leverages trained Energy-Based Models (EBMs) to generate augmented sam- ples, as described in the main text. Following the approach in (Zhao& & Chen, 2021), we train EBMs in the latent space of a VQ-VAE model?. The energy model implemented using the discriminator architecture of StyleGAN2
re-implemented these
First, the VQ-VAE 2 model is trained on all source do- mains. Subsequently, EBMs are trained in the latent space of this VQ-VAE 2 model to enable domain traversal using Langevin dynamics. The VQ-VAB 2 configuration includes a codebook dimension of 32 and a codebook size of 256. Similar to (Zhao & Chen, 2021), we train the VQ-VAE us- ing only the reconstruction loss, omitting the sttge PixelCNN training. The EBMs are optimized using the Adam optimizer with a learning rate of 0.001. During do- main traversal, the Langevin step size (3 in Eq. 4) is set to 3 = 1, and the number of Langevin steps (K) is set to K = 40.
We use the publicly available codebase for RandConv 9 MixStyle" FedDG" RAM and ' For MixStyle, we add additional layers after the first two ResNet blocks of segmentation model to transfer the instance level fea- ture statistics between source domains. Further, for RAM and FedDG, instead of UNet encoder-decoder, we use the ResNet34 blocks and ASPP decoder for training.
FFor RFS dataset, we use an L-channel replacement strategy to maintain the position of optical disc and optical cup while running LD. Specifically, after each LD step, we replace the L channel of the updated sample with L channel of original sample. This ensures that the position of OC and OD are preserved and segmentation masks remain valid.
LangDAug: For n source domains, we train 2(x) EBMs to model domain pairs in both directions To generate Langevin samples, as outlined in Section 3.3, we execute Langevin Dynamics (LD) and store intermediate samples as per Eq. 4. Instead of retaining all K = 40 intermedi- ate samples, we store a subset of these samples for two reasons: (a) to manage storage and computational over- head, and (b) to reduce the correlation between MCMC
IImplementation taken from https::/github. com/ http:////eeepaapppooootttteeeeee
1 Libbs/mode /mwdeld
Domeinieed/rrre/mmiinn
'Implementation taken from https://github... com/ rosi na 'Implementation taken from http::://gihbb..../ com/ rosina nal 1 telllty/stte tttp:/ttlllgnnn-yttrrr
httrpe//ggthab.ccc///aaaaagggg/ mixstyle-releasee
"s mentioned in limitations. this can be problematic if the number of source domains is too large. In that case, shared archi- tectures can be explored with conditioning on desired domains
pnttps:/
3

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
The datasets can be found publicty. The Retinal Fundus Dataset can be found at: gowgle google. www/file//lllooocoooooooo
Datasets:
Algorithm Langevin Data Augmentation
Input: Source domains {D,) Langevin step size B,
https: /tt.eerdieeeee///eee//eeeeeeeeee
Number of Langevin steps K. 1...
Output: Augmentation datasets (
2D MRI Prostate Dataset can be found at: https:: http:///lluuandeeeee Sengguande.githut. github..hb.... Lo//SML///// have used the same convention as earlier works (Liu al., 2021a; Zhou al., 2022b; Chen al., 2023) for naming the datasets.
I: Initialize parameters {0:z [8syli;j=1> where i i
We
Energy-Based Model (EBM) Training
8 Loop over all domain pairs for i E n]: (I....,.]) : do
C. Additional Results
repeat
C.1. Ablation Studies on EBM Hyperparameters
Sample x Dj. Xo D,.
$ Perform Langevin Dynamics
We analyze the impact of various Energy-Basse Model (EBM) hyperparameters on cross-domain generalization for the Retinal Fundus Segmentation task. The results are summarized below:
6:
for k: = 1, K do
7:
Sample e ~ NO /).
= 8 Vx--1 Ear, (xk-1) + Bc.
end for
10:
x XK.
Langevin Steps (K): First row of Figure 5 shows the performance across different numbers of Langevin steps used during EBM sampling. LangDAug main- tains stable performance higher values of K.
& Compute gradient
11:
Va, -= C= Eo,, (x) - Vol, Eo,, (x).
12:
Oij t 0ij \Vo
13:
until convergence
14:
end for
Step Size (5): Second row of Figure 5 illustrates the effect of varying the step size in Langevin dynamics. Lower values of / lead to more stable and reliable performance.
15: Generation of Langevin Samples
$ Storage for augmented datasets
EBM Complexity (#conv blocks): We vary the number of convolutional blocks in the EBM to control model complexity. As shown in the third row of Fig- ure 5, performance remains consistent, suggesting that even lightweight EBMs are sufficient.
16: Initialize empty sets (D pnm. (D =1,
Loop over all domain pairs
17:
for i pal...,,j::. {l,... nl: do
17:
Loop through samples in source doma in
Number of Augmented #samples/chain) Fourth row of Figure 5 examines the number of intermediate samples stored per Langevin chain. Too few samples fail to capture vicinal distributions, while too many introduce high auto-correlation between langevin samples leading to biased learning. A moderate, well-spaced selection of samples yields the best results.
Samples
18:
for x; D2 do
19:
Xo + Xx.
$ Perform Langevin steps
20:
for k: = 1,..., K do
21:
Sample e ~ ~N(0
X = 8² Vx-_, Ed, (x+-1) + Bc.
22:
$ Store augmented sample appendd))))
23:
24: end for
25: end for
C.2. Computational Cost Analysis
26: end for
Return all augmented datasets
We acknowledge the increased computational cost of Lang- DAug our limitations. In Table 5, we provide a detailed comparison of training time and peak memory usage across methods on the Retinal Fundus Segmentation (RFS) task. All experiments were conducted on a single NVIDIA A6000 GPU with 48GB memory, and results are averaged across domains.
27: return (D) pna.K
1 -=1
more efficient than FedDG (4.60 hrs), while providing supe- rior performance.
Additionally, the average Energy- Model (EBM) train- ing time is approximately 0.357 hours per source-target domain pair. The inference cost of running a Langevin
Although LangDAug increases training time, remains comparable methods like RAM (2.75 hrs) and notably

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
Domain A
Domain B
Domain C
Domain D
90
90
90
90
Score (%)
80
80
80
80
70
70
70
20
40
60
80
20
40
60
80
20
40
60
80
20
40
60
80
K (Langevin steps)
90
90
90
90
Score
80
80
80
80
70
70
70
70
10
10P
10
10
100
10
10
100
10
10
10"
10
e (Step size)
90
90
90
90
Score (%)
80
70
70
70
70
7
#conv blocks
90
90
90
Score (%)
80
80
80
80
70
70
70
70
2
13
40
2
13
40
2
13
40
13
40
#samples/ chaiin
mloU mDSC
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 (3); 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 mloU and mDSC.
5

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
Table 5: Training time (in GPU hours) and peak memory usage (in GB) for all methods. "+Ours" refers to integration of LangDAug into the baseline.
Metric
ERM
Ours
FedDG
FudDG+Ours
RAM
RAM+Ours
TriD
TriD+Ours
GPU hrs
1.51
3.14
4.60
6.13
2.75
3.77
5.53
7.49
Memory (GB)
10.36
19.41
16.77
23.16
12.58
20.24
24.87
30.11
Dynamics (LD) chain is minimal, taking roughly 2 seconds.
Apart from this, the proposed method also requires addi- tional storage requirements to store the Langevin samples. Particularly, for data points, and a saving frequency of f. the additional data points will equal nk/f.
As with many domain augmentation (DA) methods that rely on synthetic sample generation, increased training cost is an inherent trade-off. Potential directions for optimization include selective sampling strategies (e.g., coresets) and architectural sharing for EBMs to enable conditioning across domains.
D. Inter-Domain Transversal Visualizations
We present visual examples of the inter-domain transver- sal process over K = 40 Langevin dynamics steps. The transversal sequences for the retinal fundus dataset are de- picted in Figures to while those for the prostate dataset are shown in Figures 10 15. In each figure, rows corre- spond distinct samples starting from the source domain, with each capturing the intermediate Langevin samples, while columns represent the specific steps in the Langevin Dynamics process at which the transversal are recorded.
6

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
Gap
(a) Ex. 1: Domain A Domain
(b) Ex. 2: Domain
Domain B
Benin
See
(c) Ex. 1: Domain A
+ Domain
(d) Ex. 2: Domain A
Domain C
Yeu
Wwzi
(e) Ex. 1: Domain A - Domain
(I) Ex. 2: Domain A + Domain D
Figure 6: Examples of translation from Domain A to Domains B, and D using the proposed method on the retinal fundus dataset.

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
Drg
Gep
(a) Ex. 1: Domain B Domain
(b) Ex. 2: Domain
Domain A
Tro
0ry
Gep
Smgll
(c) Ex. I: Domain B + Domain C
(d) Ex. 2: Domain Domain C
(c) Ex. 1: Domain B Domain
(f) Ex. 2: Domain B Domain D
Figure 7: Examples of translation from Domain B to Domains A, and D using the proposed method on the retinal fundus dataset.
8

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
(a) Ex. 1: Domain C
Domain
(b) Ex. 2: Domain
Domain A
Dry
(c) Ex. 1: Domain
C Domain
(d) Ex. 2: Domain
Domain B
Bezt
(c) Ex. 1: Domain C
Domain
( Ex. 2: Domain C Domain D
Figure 8: Examples of translation from Domain C to Domains A, and D using the proposed method on the retinal fundus dataset.

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
(a) Ex. 1: Domain D Domain
(b) Ex. 2: Domain
Domain A
(c) Ex. 1: Domain
D
+ Domain
(d) Ex. 2: Domain Domain B
Swzt
(c) Ex. 1: Domain D > Domain
(f) Ex. 2: Domain D - Domain C
Figure 9: Examples of translation from Domain to Domains A, B and C using the proposed method on the retinal fundus dataset.
10

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
(a) Ex. 1: Domain A
Domain
(b) Ex. 2: Domain
Domain B
Tog
(c) Ex. 1: Domain A + Domain
(d) Ex. 2: Domain A + Domain C
(e) Ex. 1: Domain A - Domain
(f) Ex. 2: Domain A + Domain D
Figure 10: Examples translation from Domain A to Domains B, and using the proposed method on the prostate dataset.

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
Ex. 1: Domain A
Domain
(h) Ex. Domain
Domain E
(i) Ex. 1: Domain A Domain F
(i) Ex. 2: Domain Domain F
Figure 10: Examples of translation from Domain A to Domains E and F using the proposed method on the prostate dataset.
12

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
(a) Ex. 1: Domain B
Domain
(b) Ex. 2: Domain + Domain A
(c) Ex. 1: Domain B + Domain C
(d) Ex. 2: Domain Domain C
(c) Ex. 1: Domain B -> Domain
(i) Ex. 2: Domain B Domain D
Figure 11: Examples translation from Domain B to Domains A, and using the proposed method on the prostate dataset.
13

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
(g) Ex. 1: Domain B
Domain
(h) Ex. Domain
Domain E
(i) Ex. 1: Domain B Domain F
(j) Ex. 2: Domain B Domain
Figure 11: Examples of translation from Domain B to Domains E and F using the proposed method on the prostate dataset.
14

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
(a) Ex. 1:
Domain C
Domain A
(b) Ex. 2: Domain C - Domain A
(c) Ex. 1: Domain C + Domain
(d) Ex. 2: Domain
- Domain B
(c) Ex. 1: Domain C -> Domain
(f) Ex. 2: Domain C Domain D
Figure 12: Examples of translation from Domain C to Domains A, B, and D using the proposed method on the prostate dataset.
15

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
(g) Ex. 1: Domain C
Domain
(h) Ex. Domain C
Domain E
(i) Ex. 1: Domain C Domain F
(j) Ex. 2: Domain C Domain
Figure 12: Examples of translation from Domain C to Domains E and F using the proposed method on the prostate dataset.
16

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
(a) Ex. 1: Domain D
Domain
(b) Ex. 2: Domain + Domain A
(c) Ex. 1: Domain D +
Domain
(d) Ex. 2: Domain
Domain B
(c) Ex. 1: Domain D > Domain
(i) Ex. 2: Domain D - Domain C
Figure 13: Examples translation from Domain to Domains A, B and using the proposed method on the prostate dataset.
17

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
Ex. 1: Domain D
Domain E
(h) Ex. 2: Domain
Domain E
(i) Ex. 1: Domain D Domain F
(i) Ex. 2: Domain D Domain F
Figure 13: Examples of translation from Domain D to Domains E and F using the proposed method on the prostate dataset.
18

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
(a) Ex. I: Domain E - Domain A
(b) Ex. 2: Domain
- Domain A
(c) Ex. 1: Domain E +
Domain
(d) Ex. 2: Domain
Domain B
Yum
(c) Ex. 1: Domain E - Domain C
(f) Ex. 2: Domain E -> Domain
Figure 14: Examples of translation from Domain to Domains A, and using the proposed method on the prostate dataset.
19

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
(g) Ex. 1: Domain E
Domain D
(h) Ex. 2: Domain E
Domain D
(i) Ex. 1: Domain E Domain F
(i) Ex. 2: Domain E
Domain
Figure 14: Examples of translation from Domain E to Domains D and F using the proposed method on the prostate dataset.
20

---

LangDAug: Langevin Data Augmentation for Multi-Source Domain Generalization in Medical Image Segmentation
(a) Ex. 1: Domain F +
Domain A
(b) Ex. 2: Domain
+ Domain A
(c) Ex. 1: Domain F
+ Domain B
(d) Ex. 2: Domain
+ Domain B
(c) Ex. 1: Domain F - Domain C
(f) Ex. 2: Domain F - Domain
Figure 15: Examples of translation from Domain to Domains A, and C using the proposed method on the prostate dataset.
21

---

LangDAug: Langevin Data Augmentation for Multi-Souree Domain Generalization in Medical Image Segmentation
(g) Ex. 1: Domain F
Domain D
(h) Ex. 2: Domain F
Domain D
(i) Ex. 1: Domain F
Domain E
(i) Ex. 2: Domain F Domain E
Figure 15; Examples of translation from Domain F to Domains D and E using the proposed method on the prostate dataset.
22


Reviews:
[{"review_id": "mf0UINVDAv", "reviewer": "Reviewer_G9sU", "summary": "This paper proposes a data augmentation method via Langevin dynamics with theoretical analyses. They also conduct experiments to show its usefulness.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments. We have tried our best to address your concerns by adding more baseline comparisons with your suggest methods. We will include these in the final version of the manuscript. We will be happy to answer any follow-up questions you might have.\n\n## Comparison with more Baselines\nThanks for this suggestion. The mentioned literature are slightly older DomainBed methods primarily designed for classification tasks. We have compared our method against methods like Fish, Fishr and Hutchinson which fall in same category. We provide the comparison of our method with the mentioned methods on RFS dataset (due to character limit) in the table below:\n\n| Baselines  | Domain 1 |       | Domain 2 |       | Domain 3 |       | Domain 4 |       |\n|------------|----------|-------|----------|-------|----------|-------|----------|-------|\n|            | mIoU     | mDSC  | mIoU     | mDSC  | mIoU     | mDSC  | mIoU     | mDSC  |\n| CORAL      |78.37     |86.55  |63.94     |74.20  |78.89     |87.71  |75.85     |85.65  |\n| RSC        |76.82     |86.11  |62.54     |72.42  |79.73     |87.92  |74.75     |85.04  |\n| SagNet     |75.42     |84.99  |54.51     |63.69  |75.52     |85.29  |67.31     |79.20  |\n| SWAD       |76.19     |84.99  | 65.32    |75.07  |76.91     |86.30  |72.09     |82.82  |\n| LangDAug       |**78.79**     |**87.00**  | **75.05**    |**85.87**  |**81.00**     |**88.91**  |**80.51**     |**88.68**  |\n\n\nWe observe that our method outperforms these methods which is consistent to our observation with methods like Fish, Fishr and Hutchinson.\n\n\n## Confusion for DA as Domain Adaptation\nThanks for this suggestion. We will change this abbrevation from DA to DAug to avoid any such confusion. \n\n## Benefit Specific to Medical Image Segmentation\nThanks for this comment. While LangDAug can be potentially used for normal DG scenarios, we note that there are certain factors that guided our application design:\n- The effectiveness of LangDAug depends on the ability of EBMs to effectively traverse and interpolate between source domain distributions.\n- Such traversal becomes especially natural when source domains share structured similarities or consistent underlying factors of variation. Medical imaging data typically demonstrates structured variations, predominantly reflected through differences in amplitude spectra across domains [1,2].\n- EBMs have been shown to excel at capturing and modeling variations in amplitude spectra [3,4], making them suited for the domain variations characteristic of medical imaging data. Thus, LangDAug’s capability aligns well with the domain shift challenges encountered specifically in medical image segmentation.\n\nFor these reasons, we use LangDAug in context of medical image segmentation.\n\n[1] Zhou et al. \"Generalizable medical image segmentation via random amplitude mixup and domain-specific image restoration.\" ECCV, 2022.\n\n[2] Liu et al. \"FedDG: Federated domain generalization on medical image segmentation via episodic learning in continuous frequency space.\" CVPR, 2021.\n\n[3] Du et al. \"Compositional visual generation with energy based models.\" NeurIPS, 2020.\n\n[4] Tancik et al. \"Fourier features let networks learn high frequency functions in low dimensional domains.\" NeurIPS, 2020.", "claims_and_evidence": "In Line 019-023, \"DA methods, which enrich model representations through synthetic samples, have shown comparable or superior performance to representation learning approaches.\" Is there empirical evidence for this?", "methods_and_evaluation_criteria": "N/A", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "There seem to be too few baselines considering the abundant literature of domain generalization. To name a few: CORAL [1], RSC [2], SWAD [3], SagNet [4], etc.\n\n\n\n[1] Sun B, Saenko K. Deep coral: Correlation alignment for deep domain adaptation[C]//Computer vision–ECCV 2016 workshops: Amsterdam, the Netherlands, October 8-10 and 15-16, 2016, proceedings, part III 14. Springer International Publishing, 2016: 443-450.\n\n[2] Huang Z, Wang H, Xing E P, et al. Self-challenging improves cross-domain generalization[C]//Computer vision–ECCV 2020: 16th European conference, Glasgow, UK, August 23–28, 2020, proceedings, part II 16. Springer International Publishing, 2020: 124-140.\n\n[3] Cha J, Chun S, Lee K, et al. Swad: Domain generalization by seeking flat minima[J]. Advances in Neural Information Processing Systems, 2021, 34: 22405-22418.\n\n[4] Nam H, Lee H J, Park J, et al. Reducing domain gap via style-agnostic networks[J]. arXiv preprint arXiv:1910.11645, 2019, 2(7): 8.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "In this paper, \"DA\" is short for data augmentation. However, in OOD literature, DA usually stands for domain adaptation. It is better to change another abbreviation for data augmentation to avoid confusion.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "What is the benefit of the proposed method in domain generalization specifically in medical image segmentation? It seems that it can be applied to normal domain generalization tasks as well.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TDOyWKgti7", "reviewer": "Reviewer_1jQ9", "summary": "The authors propose LangDAug, a Langevin Data Augmentation technique to improve multi-source domain generalization in medical image segmentation tasks. The core idea is to train Energy-Based Models (EBMs) via contrastive divergence to model the transitions between pairs of source domains and use Langevin dynamics (LD) to generate intermediate samples that serve as augmented data bridging domain gaps. Theoretically, the authors demonstrate that this augmentation method induces a regularization effect, leading to smoother optimization landscapes and better generalization by bounding model complexity based on intrinsic data manifold dimensionality. Experimental evaluations conducted on retinal fundus and prostate MRI segmentation datasets show that LangDAug outperforms previous methods across unseen domains.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your appreciation and in-depth review. We have tried our best to answer your concerns. We will include these in final version of the paper. Please let us know if you have any other questions.\n### Computational Cost Analyses\nWe acknowledged the increased computational cost of LangDAug in our limitations. Below, we provide details of the training time and memory usage across methods on the RFS data (averaged across domains on single 48GB A6000 card):\n| Metric | ERM | Ours | FedDG | FedDG+Ours | RAM | RAM+Ours | TriD | TriD+Ours |\n|--------|---------|------|-------|------------|-----|----------|------|-----------|\n| GPU hrs/train. time | 1.508 | 3.141 | 4.604 | 6.129 | 2.746 | 3.772 | 5.528 | 7.494 |\n| Peak Memory (GB) | 10.36 | 19.41 | 16.77 | 23.16 | 12.58 | 20.24 | 24.87 | 30.11 |\n\nWhile LangDAug increases training time, it's important to note that LangDAug's time (3.14 hrs) is comparable to existing methods like RAM (2.75 hrs) and less than FedDG (4.60 hrs). Moreover, LangDAug performs better than these methods.\n\nIn addition to above, we note an average EBM training time of 0.357 hr/pair. Further the inference cost (to run one LD chain) is very minimal ~2 sec. \n\nIncreased training cost is an inherent limitation of DA methods that use synthetic samples. Future optimizations could include selective sampling (e.g., coresets) and shared architectures to jointly train EBMs with conditioning on source/target domains.\n## Effect of each component\nThanks for this suggestion. We provide ablations for several components of our method in the following tables:\n|                         | Domain 1 |       | Domain 2 |       | Domain 3 |       | Domain 4 |       |\n|-------------------------|----------|-------|----------|-------|----------|-------|----------|-------|\n|                         | mIoU     | mDSC  | mIoU     | mDSC  | mIoU     | mDSC  | mIoU     | mDSC  |\n| $k$                     |          |       |          |       |          |       |          |       |\n| 20                      | 66.18    | 71.07 | 70.39    | 80.29 | 77.67    | 82.45 | 75.39    | 84.31 |\n| 40                      | 78.79    | 87.00 | 75.05    | 85.87 | 81.00    | 88.91 | 80.51    | 88.68 |\n| 60                      | 78.59    | 87.39 | 74.11    | 85.85 | 81.46    | 88.06 | 80.94    | 87.72 |\n| 80                      | 75.92    | 83.82 | 72.35    | 82.84 | 79.03    | 85.22 | 78.95    | 85.59 |\n| $\\beta$                 |          |       |          |       |          |       |          |       |\n| 0.1                     | 77.32    | 85.71 | 74.17    | 85.44 | 79.56    | 88.07 | 79.99    | 86.95 |\n| 1                       | 78.79    | 87.00 | 75.05    | 85.87 | 81.00    | 88.91 | 80.51    | 88.68 |\n| 10                      | 71.97    | 79.05 | 69.07    | 78.11 | 76.10    | 83.62 | 73.72    | 81.24 |\n| #Conv blocks/resolution |          |       |          |       |          |       |          |       |\n| 1                       | 78.79    | 87.00 | 75.05    | 85.87 | 81.00    | 88.91 | 80.51    | 88.68 |\n| 4                       | 79.81    | 88.08 | 74.80    | 85.28 | 81.56    | 88.67 | 80.59    | 87.19 |\n| 7                       | 78.725   | 86.80 | 73.48    | 84.70 | 79.29    | 86.69 | 79.30    | 86.59 |\n| #Augmented Samples      |          |       |          |       |          |       |          |       |\n| 2/chain                 | 74.78    | 80.09 | 69.23    | 79.33 | 75.70    | 84.63 | 73.90    | 84.01 |\n| 13/chain                | 78.79    | 87.00 | 75.05    | 85.87 | 81.00    | 88.91 | 80.51    | 88.68 |\n| 40/chain                | 71.32    | 79.12 | 66.09    | 77.28 | 74.09    | 81.20 | 72.05    | 82.28 |\nWe make the following observations:\n- Under reasonable EBM hparams (smaller $\\beta$, higher $k$), the performance of LangDAug is stable.\n- The performance remains stable with varying EBM complexity. \n- Smaller number of langevin samples is not effective as it doesn't capture the vicinal distributions.\n- Large number of langevin samples also harm performance due to high auto-correlation between langevin samples leading to biased learning. Hence, a well spaced selection of langevin samples is desired.\n## Application to 3D Medical Images\nThe primary challenge in extending LangDAug to 3D medical images is the current lack of mature methodologies for generating 3D volumes using EBMs. Current approaches [1] are limited to primitive 3D shapes and don't scale well to complex anatomical structures.\nAs a partial evaluation, we applied LangDAug to the 3D Prostate MRI dataset by processing and augmenting 2D axial slices. While this doesn't capture full volumetric continuity, it provides evidence of LangDAug's effectiveness in 3D imaging contexts.\nWe view full 3D EBM generation as an important direction for future work, which could be enabled by recent advances in score-based modeling and memory-efficient architectures.\n\n[1] Xie, Jianwen et al. \"Learning descriptor networks for 3D shape synthesis and analysis.\" CVPR 2018.", "claims_and_evidence": "Overall, the paper’s main claims regarding the effectiveness and robustness of the proposed LangDAug approach are convincingly supported by empirical evidence (quantitative benchmarks and visualizations) and theoretical analysis. Some limitations exist concerning scalability and detailed validation of anatomical fidelity, but these do not substantially undermine the validity or clarity of the core contributions.", "methods_and_evaluation_criteria": "The proposed methodology and evaluation criteria used in the submission are clearly appropriate and well-justified for the stated objectives of the paper; however, I would also want to know if the proposed LangDAug could be used for 3D medical image segmentation tasks and the efficiency of doing so.", "theoretical_claims": "Overall, the theoretical claims made by the authors are clearly articulated, logically consistent, and mathematically rigorous. The authors provide a coherent theoretical framework that convincingly justifies why LangDAug would improve domain generalization. There are no substantial mathematical errors or inconsistencies evident from the provided derivations.", "experimental_designs_or_analyses": "The experimental design and analyses used in this submission are sound, valid, and appropriately rigorous for assessing domain generalization performance.\n\nHowever, there are still some limitations:\n- Missing detailed computational cost analyses as the proposed LangDAug could increase training time.\n- It is better to analyze the contribution of individual components of LangDAug (e.g., varying Langevin steps, EBMs complexity, or the number of augmented samples).\n- How about applying LangDAug to 3D medical image segmentation tasks?", "relation_to_broader_scientific_literature": "- The key contributions (LangDAug) of this paper closely align with and extend several ideas from the broader scientific literature, specifically related to domain generalization, energy-based models, Langevin dynamics, and data augmentation strategies in machine learning, particularly in medical imaging. \n\n- LangDAug directly extends prior work on domain generalization, especially the notion that models trained under the Empirical Risk Minimization (ERM) paradigm often fail to generalize under domain shifts (Gulrajani & Lopez-Paz, 2021; Ganin et al., 2016). It is positioned within data manipulation approaches, specifically leveraging the concept of intermediate-domain traversal via Langevin dynamics, distinguishing itself from prior DG methods that primarily relied on style-mixing (MixStyle), random convolutions (RandConv), or frequency domain augmentation (FedDG, RAM, TriD).", "essential_references_not_discussed": "The manuscript thoroughly references relevant literature on domain generalization (DG), energy-based models (EBMs), Langevin dynamics, and medical image segmentation.", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "- Could the authors provide a detailed computational cost analysis (e.g., GPU hours, memory requirements, inference and training times) of the proposed LangDAug method compared to the baseline domain generalization approaches?\n\n- The current implementation and validation of LangDAug focus exclusively on 2D medical image segmentation. Could the authors explain any specific technical or methodological reasons why you have not extended LangDAug directly to 3D medical image segmentation tasks? Are there fundamental challenges, such as increased complexity, computational demands, or theoretical limitations, that prevented immediate extension?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "34pf49SRRP", "reviewer": "Reviewer_WZCC", "summary": "This paper presents LangDAug, a novel data augmentation method designed to address multi-source domain generalization challenges in medical image segmentation. Leveraging Langevin Dynamics (LD) and Energy-Based Models (EBMs), LangDAug generates intermediate samples bridging source domains to enhance model generalization across unseen target domains. The approach involves training EBMs to capture energy distributions between source domains, using LD to create bridging samples, and incorporating these samples into the segmentation model's training process. Theoretical analysis demonstrates that LangDAug induces regularization, reducing Rademacher complexity and improving generalization. Experiments on retinal fundus segmentation and 2D MRI prostate segmentation datasets show that LangDAug outperforms existing domain generalization methods and effectively complements domain randomization techniques like FedDG and RAM.\n\n## update after rebuttal\nWhile this work is application-oriented, after thoroughly reviewing all the responses, I remain concerned that the high computational cost and storage requirement could significantly limit the practical applicability of the proposed method. Therefore, I maintain my original score as 2.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your comments. We have tried our best to answer your concerns. We will be happy to engage in follow-up discussion if you have further questions.\n\n## Computational and Storage Requirements\nThanks for this comment. We acknowledged this in our limitations, our method increases storage and computational costs. \n-  Additional storage is needed for Langevin samples - specifically $(k/f)×n$ samples per domain pair (where $f$ is saving frequency, $n$ is total data points).\n- Even though we train EBMs for each pair, our EBMs are lightweight and require only 0.357 hr/pair for training. Further, this can be parallelized to reduce this time. \n- EBM training is performed offline, and models can be discarded after generating samples. This doesn't affect the main segmentation network's training or inference requirements.\n- As shown in response to reviewer 1jQ9, LangDAug's training time (3.14 hrs) are comparable to existing methods like RAM (2.75 hrs) and less than FedDG (4.60 hrs).\n- Increased training cost is an inherent limitation of DA methods that use synthetic samples. Solutions could include selective sampling (e.g., coresets) and shared architectures to jointly train EBMs with conditioning on source/target domains.We leave these techniques to be explored in future works.\n## Stability of EBMs\nWe appreciate the reviewer’s comment regarding the importance of EBM training stability in performance. While our submission mentioned training details, we didn't observe any instability. To show this, we conduct additional ablations to assess the stability of the EBMs used in our work. Specifically, we monitor the segmentation performance of our method against different EBM hparam settings. We observe the effect of number of langevin steps ($k$), step size ($\\beta$) and EBM complexity (by varying the number of conv blocks):\n|                         | Domain 1 |       | Domain 2 |       | Domain 3 |       | Domain 4 |       |\n|-------------------------|----------|-------|----------|-------|----------|-------|----------|-------|\n|                         | mIoU     | mDSC  | mIoU     | mDSC  | mIoU     | mDSC  | mIoU     | mDSC  |\n| $k$                     |          |       |          |       |          |       |          |       |\n| 20                      | 66.18    | 71.07 | 70.39    | 80.29 | 77.67    | 82.45 | 75.39    | 84.31 |\n| 40                      | 78.79    | 87.00 | 75.05    | 85.87 | 81.00    | 88.91 | 80.51    | 88.68 |\n| 60                      | 78.59    | 87.39 | 74.11    | 85.85 | 81.46    | 88.06 | 80.94    | 87.72 |\n| 80                      | 75.92    | 83.82 | 72.35    | 82.84 | 79.03    | 85.22 | 78.95    | 85.59 |\n| $\\beta$                 |          |       |          |       |          |       |          |       |\n| 0.1                     | 77.32    | 85.71 | 74.17    | 85.44 | 79.56    | 88.07 | 79.99    | 86.95 |\n| 1                       | 78.79    | 87.00 | 75.05    | 85.87 | 81.00    | 88.91 | 80.51    | 88.68 |\n| 10                      | 71.97    | 79.05 | 69.07    | 78.11 | 76.10    | 83.62 | 73.72    | 81.24 |\n| #Conv blocks/resolution |          |       |          |       |          |       |          |       |\n| 1                       | 78.79    | 87.00 | 75.05    | 85.87 | 81.00    | 88.91 | 80.51    | 88.68 |\n| 4                       | 79.81    | 88.08 | 74.80    | 85.28 | 81.56    | 88.67 | 80.59    | 87.19 |\n| 7                       | 78.725   | 86.80 | 73.48    | 84.70 | 79.29    | 86.69 | 79.30    | 86.59 |\n\nWe see that under reasonable hparam values, the performance of LangDAug is stable. Particularly, smaller $\\beta$, moderate $k$ lead to good performance. Whereas, the EBM complexity has very minimal effect on performance, indicating lightweight EBMs are sufficient. From this, we conclude stable performance of LangDAug under reasonable hparam settings.\n## Restricted Training distribution\nThanks for this comment. In general, all DG methods rely on diversity in source domains, as noted in Sec. 1, 2. \nThis aligns with a fundamental assumption in DG - diverse source domains allow the model to learn representations for better generalization.\nWhen the original distribution is restricted or sparse in certain regions, LangDAug helps fill these gaps, enriching the training distribution. While extremely limited diversity would challenge any DG method, our experiments show competitive performance even in difficult cases. We refer to Tab. 1 & 2 where Domain B (Tab. 1) & Domain C (Tab. 2) are the most difficult domains because of significantly different intstruments used for image acquisition. In this case, the source domains are ineffective in capturing the corresponding target domain distribution. Despite this challenge, LangDAug outperforms other methods in several metrics for these domains. This demonstrates that while limited diversity in source domains is a fundamental challenge for all DG methods, our approach remains effective even under these constrained conditions.", "claims_and_evidence": "N/A", "methods_and_evaluation_criteria": "N/A", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "N/A", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "**Strengths:**\n\n1. LangDAug introduces a new data augmentation method combining Langevin Dynamics and Energy-Based Models, offering a new perspective for solving multi-source domain generalization problems in medical image segmentation.\n2. The method is supported by theoretical analysis, demonstrating its regularization effect and its ability to reduce Rademacher complexity, enhancing the reliability and scientific foundation of the approach.\n3. This method can be used independently or integrated with existing domain randomization techniques (such as FedDG and RAM), showing its flexibility and practical applicability.\n\n**Weaknesses:**\n\n1. Training an EBM for every pair of source domains results in significant computational and storage requirements, especially when the number of source domains increases. Generating a large number of intermediate samples substantially increases training time, which may hinder the method's practicality for large-scale datasets.\n2. The effectiveness of LangDAug heavily relies on the quality of the trained EBMs, but the paper lacks an in-depth discussion of EBM training stability and its potential impact on results.\n3. Generating inter-domain samples to improve generalization is a reliable solution. However, this requires that the samples in the training set cover the sample space uniformly to obtain good interpolation results. What happens to the performance of the model if the distribution of samples in the training set is restricted?", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Pls refer to Weaknesses", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes a Langevin Data Augmentation for multi-source domain generalization in 2D medical image segmentation. While the work presents an interesting approach, several critical concerns remain unresolved. During the rebuttal period, the concern that high computational requirement limits the application of the proposed method has not been adequately addressed. Besides, after reading the paper, reviews, and rebuttal, the submitted version misses some important content, like comparing it with other domain generalization methods, the ablation study of the proposed method, and broader experimental validation beyond fundus segmentation and 2D MRI datasets. Therefore, in my opinion, the current version does not meet the acceptance criteria for publication.
Decision: Accept (poster)