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

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.

Computational and Storage Requirements

Thanks for this comment. We acknowledged this in our limitations, our method increases storage and computational costs.

• 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).
• 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.
• 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.
• 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).
• 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.

Stability of EBMs

We 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):

We 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.

Restricted Training distribution

Thanks for this comment. In general, all DG methods rely on diversity in source domains, as noted in Sec. 1, 2. This aligns with a fundamental assumption in DG - diverse source domains allow the model to learn representations for better generalization. When 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.

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.

Computational Cost Analyses

We 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):

While 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.

In 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.

Increased 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.

Effect of each component

Thanks for this suggestion. We provide ablations for several components of our method in the following tables:

• Under reasonable EBM hparams (smaller $\beta$, higher $k$), the performance of LangDAug is stable.
• The performance remains stable with varying EBM complexity.
• Smaller number of langevin samples is not effective as it doesn't capture the vicinal distributions.
• 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.

Application to 3D Medical Images

The 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. As 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. We 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.

[1] Xie, Jianwen et al. "Learning descriptor networks for 3D shape synthesis and analysis." CVPR 2018.

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.

Comparison with more Baselines

Thanks 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:

We observe that our method outperforms these methods which is consistent to our observation with methods like Fish, Fishr and Hutchinson.

Confusion for DA as Domain Adaptation

Thanks for this suggestion. We will change this abbrevation from DA to DAug to avoid any such confusion.

Benefit Specific to Medical Image Segmentation

Thanks 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:

• The effectiveness of LangDAug depends on the ability of EBMs to effectively traverse and interpolate between source domain distributions.
• 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].
• 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.

For these reasons, we use LangDAug in context of medical image segmentation.

[1] Zhou et al. "Generalizable medical image segmentation via random amplitude mixup and domain-specific image restoration." ECCV, 2022.

[2] Liu et al. "FedDG: Federated domain generalization on medical image segmentation via episodic learning in continuous frequency space." CVPR, 2021.

[3] Du et al. "Compositional visual generation with energy based models." NeurIPS, 2020.

[4] Tancik et al. "Fourier features let networks learn high frequency functions in low dimensional domains." NeurIPS, 2020.