[论文解读] Shape-aware Meta-learning for Generalizing Prostate MRI Segmentation to Unseen Domains

Quande Liu, Qi Dou+1 more|arXiv (Cornell University)|Jul 4, 2020

Domain Adaptation and Few-Shot Learning参考文献 30被引用 12

一句话总结

该论文提出了一种新型元学习框架——形状感知元学习（SAML），通过在元优化过程中显式建模形状紧凑性和边界平滑性，提升了前列腺MRI分割中的域泛化能力。通过采用课程式训练模拟域偏移，并引入两种形状感知损失目标，SAML在六个未见临床中心的测试中均实现了最先进性能，Dice和ASD指标均有稳定提升。

ABSTRACT

Model generalization capacity at domain shift (e.g., various imaging protocols and scanners) is crucial for deep learning methods in real-world clinical deployment. This paper tackles the challenging problem of domain generalization, i.e., learning a model from multi-domain source data such that it can directly generalize to an unseen target domain. We present a novel shape-aware meta-learning scheme to improve the model generalization in prostate MRI segmentation. Our learning scheme roots in the gradient-based meta-learning, by explicitly simulating domain shift with virtual meta-train and meta-test during training. Importantly, considering the deficiencies encountered when applying a segmentation model to unseen domains (i.e., incomplete shape and ambiguous boundary of the prediction masks), we further introduce two complementary loss objectives to enhance the meta-optimization, by particularly encouraging the shape compactness and shape smoothness of the segmentations under simulated domain shift. We evaluate our method on prostate MRI data from six different institutions with distribution shifts acquired from public datasets. Experimental results show that our approach outperforms many state-of-the-art generalization methods consistently across all six settings of unseen domains.

研究动机与目标

• 为解决在一种扫描仪或协议下训练的模型在未见域上失效的域偏移挑战。
• 在训练阶段无需目标域数据的情况下提升模型泛化能力，实现新临床环境下的直接部署。
• 克服在未见域上分割结果的缺陷，如形状不完整和边界模糊等问题。
• 将形状特定的归纳偏置——紧凑性和平滑性——整合进元学习，以提升医学图像分割的泛化能力。
• 开发一种模型无关、基于梯度的元学习方法，以增强对多中心MRI数据分布偏移的鲁棒性。

提出的方法

• 该方法采用基于梯度的元学习框架，结合课程式训练，通过将源域划分为元训练和元测试集来模拟域偏移。
• 提出一种形状紧凑性损失（$\mathcal{L}_{compact}$），在域偏移下促使分割掩码保持完整、闭合的形状。
• 设计一种形状平滑性损失（$\mathcal{L}_{smooth}$），在潜在空间中促进域不变的轮廓嵌入，以提升边界清晰度。
• 元优化目标结合标准分割损失与$\mathcal{L}_{compact}$和$\mathcal{L}_{smooth}$，使模型能够学习在多种成像协议下均鲁棒的特征。
• 该方法为模型无关，可应用于任意分割主干网络，具备良好的灵活性和广泛部署潜力。
• 训练过程涉及在元训练和元测试任务之间交替优化，模型在元更新过程中适应模拟的域偏移。

实验结果

研究问题

• RQ1 通过显式引入形状感知正则化的元学习，是否能提升前列腺MRI分割在未见临床域中的泛化能力？
• RQ2 形状紧凑性和边界平滑性约束在域偏移下对分割性能有何影响？
• RQ3 将形状归纳偏置整合进元优化是否优于标准元学习和数据聚合基线方法？
• RQ4 源域数量的增加如何影响所提方法与现有方法相比的泛化性能？
• RQ5 形状感知损失是否能缓解未见域中常见的失败模式，如分割不完整和边界模糊？

主要发现

• SAML在所有六个未见域上均取得最高的平均Dice分数（87.67），显著优于基线方法DeepAll（85.52）及其他最先进方法。
• 在最具挑战性的站点（站点F）上，SAML的Dice分数达到88.34，领先于次优方法MASF（86.55）1.79分。
• 同时引入$\mathcal{L}_{compact}$和$\mathcal{L}_{smooth}$在所有六个站点上均带来一致性能提升，ASD较基线降低1.22分。
• 消融实验证实，两种形状感知损失独立且协同贡献，其中$\mathcal{L}_{smooth}$带来的性能增益最大。
• 随着源域数量的增加，SAML在与DeepAll的对比中持续保持优势，展现出稳定且可扩展的泛化能力。
• 在六个站点中有四个站点上，单域内训练的性能劣于SAML，表明多域数据比有限的本地数据更有利于模型的鲁棒学习。