[论文解读] Rethinking Data Augmentation for Single-source Domain Generalization in Medical Image Segmentation

Zixian Su, Kai Yao+4 more|arXiv (Cornell University)|Nov 27, 2022

Domain Adaptation and Few-Shot Learning被引用 8

一句话总结

本文提出了一种新颖的数据增强框架——显著性平衡的定位-尺度增强（SLAug），用于单源域泛化下的医学图像分割。通过结合类别级别的定位-尺度增强与显著性引导的融合机制，SLAug提升了增强数据的多样性与信息量，在腹部CT-MRI任务上相比ERM实现4.3%的Dice分数提升，在心脏bSSFP-LGE任务上提升2.1%，并理论上界定了泛化风险。

ABSTRACT

Single-source domain generalization (SDG) in medical image segmentation is a challenging yet essential task as domain shifts are quite common among clinical image datasets. Previous attempts most conduct global-only/random augmentation. Their augmented samples are usually insufficient in diversity and informativeness, thus failing to cover the possible target domain distribution. In this paper, we rethink the data augmentation strategy for SDG in medical image segmentation. Motivated by the class-level representation invariance and style mutability of medical images, we hypothesize that unseen target data can be sampled from a linear combination of $C$ (the class number) random variables, where each variable follows a location-scale distribution at the class level. Accordingly, data augmented can be readily made by sampling the random variables through a general form. On the empirical front, we implement such strategy with constrained B$\acute{ m e}$zier transformation on both global and local (i.e. class-level) regions, which can largely increase the augmentation diversity. A Saliency-balancing Fusion mechanism is further proposed to enrich the informativeness by engaging the gradient information, guiding augmentation with proper orientation and magnitude. As an important contribution, we prove theoretically that our proposed augmentation can lead to an upper bound of the generalization risk on the unseen target domain, thus confirming our hypothesis. Combining the two strategies, our Saliency-balancing Location-scale Augmentation (SLAug) exceeds the state-of-the-art works by a large margin in two challenging SDG tasks. Code is available at https://github.com/Kaiseem/SLAug .

研究动机与目标

• 解决仅有一个源域可用时医学图像分割中的域偏移问题。
• 克服仅全局或随机增强策略缺乏多样性和信息量的局限性。
• 利用医学图像中类别级别的表征不变性与风格可变性，设计更鲁棒的增强策略。
• 在所提出的增强框架下，理论上界定了泛化风险，以验证其有效性。
• 开发一个即插即用的模块，在无需目标域数据的情况下提升模型鲁棒性。

提出的方法

• 提出一种线性组合模型，其中增强图像从每个类别的定位-尺度分布中采样：$ x_{U}^{c} = \alpha^{c}x_{S}^{c} + \beta^{c} $，其中 $ \alpha^{c} $ 和 $ \beta^{c} $ 为可学习的缩放与定位参数。
• 实施全局定位-尺度增强（GLA），以模拟整个图像的全局强度偏移。
• 引入局部定位-尺度增强（LLA），利用真实标注掩码对每个类别应用特定变换，以模拟类别级别的域偏移。
• 设计显著性平衡融合（SBF）机制，利用基于梯度的显著性图引导融合过程，确保融合结果具有合适的朝向与幅度。
• 在GLA与LLA中均应用受限的Bézier变换，以确保图像变换的平滑性与真实性。
• 理论分析证明，所提出的增强策略在SGD设置下可导出泛化风险的上界。

实验结果

研究问题

• RQ1 能否利用医学图像中类别级别的表征不变性与风格可变性，设计出更有效的域泛化数据增强策略？
• RQ2 与仅全局或随机增强相比，从定位-尺度分布中抽取的类别级别随机变量的线性组合，是否能更好地捕捉未见域中的分布偏移？
• RQ3 通过显著性引导融合增强样本，能否通过聚焦于高梯度、与决策边界相关区域，提升模型泛化能力？
• RQ4 所提出的增强策略是否能在单源域泛化设置下实现理论上可界定的泛化风险？
• RQ5 各个组件（GLA、LLA、SBF）对整体性能的贡献如何？它们之间是否具有协同增益？

主要发现

• 在腹部CT-MRI任务上，SLAug相比ERM实现4.3%的Dice分数提升，腹部Dice达88.63，心脏Dice达86.69。
• 在心脏bSSFP-LGE任务上，SLAug相比ERM提升2.1%性能，平均Dice分数达到87.66%。
• 消融实验确认GLA、LLA与SBF三者均不可或缺：移除任一组件均导致性能显著下降。
• 变体4（GLA + SBF）优于变体1（仅GLA），表明GLA增强图像的显著性图能有效引导分布扩展。
• 随机融合性能较差（比SBF低2.9%），证实基于显著性的融合更具信息量且更具可解释性。
• SBF在保留源域知识方面优于ERM或随机融合，分别将源域测试性能提升0.8和4.0个Dice点。