基于Langevin动力学的多源域泛化医学图像分割方法研究

作者：Piyush Tiwary1, Kinjawl Bhattacharyya1, Prathosh A.P.1
1印度科学理工学院电气通信工程系
发表于：Proceedings of the 42nd International Conference on Machine Learning, PMLR 267, 2025

研究背景

医学图像分割模型在实际应用中面临一个关键挑战：由于成像设备、协议和患者群体的差异，训练模型在新域（target domain）上的表现往往显著下降。这一问题源于传统经验风险最小化（Empirical Risk Minimization, ERM）框架对独立同分布假设的依赖。域泛化（Domain Generalization, DG）方法旨在解决这一问题，主要包括表示学习和数据增强（Data Augmentation, DAug）两大方向。

表示学习方法通过寻找域不变特征（domain-invariant features）来提高泛化能力，但往往依赖特定技巧且缺乏理论保证。相比之下，数据增强方法通过合成样本来丰富模型表示，已展现出与表示学习方法相当甚至更优的性能。本研究提出了一种新颖的Langevin数据增强（Langevin Data Augmentation, LangDAug）方法，专门针对2D医学图像分割中的多源域泛化问题。

研究方法

1. 总体框架

LangDAug的核心思想是利用基于能量的模型（Energy-Based Models, EBMs）在不同源域之间进行迁移，通过Langevin动力学（Langevin Dynamics, LD）生成中间样本。这些中间样本作为”桥梁”连接不同源域，从而扩展模型的域覆盖范围。整个方法包含三个主要步骤：

1. 跨域EBM训练：为每对源域训练一个EBM，使用对比散度（Contrastive Divergence, CD）目标函数
2. Langevin样本生成：通过LD迭代生成中间域样本
3. 增强训练：将原始样本与Langevin样本结合进行模型训练

2. 技术细节

2.1 EBM跨域迁移

对于两个不同的源域Di和Dj，训练一个EBM Eθij来实现在两域间的迁移。使用对比散度损失函数：

∇θij LCD = 𝔼_pDj[∇θij Eθij(x)] - 𝔼_pθij[∇θij Eθij(x)]

其中pθij = exp(-Eθij(x))/Zθij。由于直接从pθij采样不可行，研究采用Langevin动力学进行近似采样：

x^{t+1} = x^t - α/2 ∇x Eθij(x^t) + αε

初始化x^0 ~ pDi(x)，通过k步LD迭代生成样本。

2.2 Langevin数据增强

对于每个样本(xj,yj)∈Di，运行k步LD生成k个中间样本{xj^t}t=1^k。这些”Langevin数据”{(xj^t,yj)}被用于增强ERM训练。理论分析表明，经过k步LD后，数据Dij^k ≡ {xj^k,yj}j=1^ni可视为来自一个新域。

3. 理论贡献

研究提供了LangDAug的理论分析，主要贡献包括：

1. 正则化效应：证明了LangDAug对参数模型fθ(·)的方向导数具有正则化作用
2. 广义线性模型分析：对于GLMs，展示了正则化项的具体形式
3. Rademacher复杂度上界：证明了对于所考虑的函数类，LangDAug将Rademacher复杂度上界限制在数据流形的本征维度上

理论结果表明，LangDAug的性能提升源于其能够根据数据的”真实”自由度（本征维度）而非存储维度（环境维度）来限制模型的泛化能力。

实验结果

研究在两个医学图像分割基准上进行了评估：

1. 视网膜眼底分割(RFS)

使用来自四个临床站点的数据，采用留一法协议评估。LangDAug在所有域上都表现出最稳定的性能：

• 在最具挑战性的域B上，LangDAug的mIoU达到75.05%，优于次优方法(RAM)1.26%
• 平均而言，LangDAug实现了78.84%的mIoU和87.61%的mDSC，优于所有基线方法
• 性能标准差(mIoU: ±2.43, mDSC: ±1.89)显著低于其他方法，表明其跨域稳定性

2. 前列腺MRI分割

使用来自六个临床站点的116个T2加权MRI扫描：

• LangDAug在所有域上都保持领先性能
• 在最具挑战性的域E上，DSC达到83.17%，优于其他方法
• 平均ASD为0.81mm，DSC为89.16%，均为最佳结果

3. 与域随机化方法的结合

LangDAug能够有效提升现有域随机化方法的性能：

• FedDG结合LangDAug：mDSC提升1.40%
• RAM结合LangDAug：mDSC提升2.04%
• TRID结合LangDAug：mDSC提升2.90%

研究结论与价值

本研究提出了一种创新的Langevin数据增强方法，用于解决医学图像分割中的多源域泛化问题。主要贡献包括：

1. 方法创新：首次将Langevin动力学与EBMs结合用于域泛化，通过生成中间域样本来桥接不同源域
2. 理论突破：证明了该方法的正则化效应及其对Rademacher复杂度的限制
3. 性能优势：在两个医学图像分割任务上实现了state-of-the-art的性能
4. 兼容性：可有效增强现有域随机化方法的性能

该研究的科学价值在于： - 为域泛化问题提供了新的基于能量的方法视角 - 建立了Langevin动力学与域泛化性能之间的理论联系 - 为医学图像分析中的分布偏移问题提供了实用解决方案

临床应用价值包括： - 提高医学影像AI模型在新医院/新设备上的泛化能力 - 减少对大规模目标域标注数据的依赖 - 促进医疗AI模型的实际部署和临床转化

研究亮点

1. 方法新颖性：首次将Langevin动力学生成的中间样本用于域泛化
2. 理论深度：不仅提出实用方法，还提供了严格的理论分析
3. 性能卓越：在两个医学图像分割基准上均取得最佳结果
4. 实用性强：代码已开源，便于社区使用和扩展
5. 扩展性好：方法可与其他域适应技术结合使用

局限性与未来方向

作者指出了两个主要限制： 1. 计算成本：EBM数量随源域数量增加而增加 2. 2D处理：当前方法处理3D体积数据时需要切片

未来工作可探索： - 共享架构与域条件来提升可扩展性 - 直接处理3D数据以更好地建模空间关系 - 开发更高效的采样策略降低计算成本

这项研究为医学图像分析中的域泛化问题提供了新的思路和解决方案，具有重要的理论和实践意义。