[论文解读] Contrastive learning of global and local features for medical image segmentation with limited annotations

Krishna Chaitanya, Ertunç Erdil+2 more|arXiv (Cornell University)|Jun 18, 2020

Radiomics and Machine Learning in Medical Imaging被引用 241

一句话总结

论文为医学影像分割提出领域和问题特定的对比学习预训练，增加局部对比损失并利用体积结构在少量标注下提升性能。

ABSTRACT

A key requirement for the success of supervised deep learning is a large labeled dataset - a condition that is difficult to meet in medical image analysis. Self-supervised learning (SSL) can help in this regard by providing a strategy to pre-train a neural network with unlabeled data, followed by fine-tuning for a downstream task with limited annotations. Contrastive learning, a particular variant of SSL, is a powerful technique for learning image-level representations. In this work, we propose strategies for extending the contrastive learning framework for segmentation of volumetric medical images in the semi-supervised setting with limited annotations, by leveraging domain-specific and problem-specific cues. Specifically, we propose (1) novel contrasting strategies that leverage structural similarity across volumetric medical images (domain-specific cue) and (2) a local version of the contrastive loss to learn distinctive representations of local regions that are useful for per-pixel segmentation (problem-specific cue). We carry out an extensive evaluation on three Magnetic Resonance Imaging (MRI) datasets. In the limited annotation setting, the proposed method yields substantial improvements compared to other self-supervision and semi-supervised learning techniques. When combined with a simple data augmentation technique, the proposed method reaches within 8% of benchmark performance using only two labeled MRI volumes for training, corresponding to only 4% (for ACDC) of the training data used to train the benchmark. The code is made public at https://github.com/krishnabits001/domain_specific_cl.

研究动机与目标

• 通过利用自监督学习，推动减少医学影像分割所需的标注数据量。
• 引入面向体积医学影像的全局与局部对比学习扩展。
• 结合来自 MRI/CT 的领域特定提示与局部区域区分，以改善像素级预测。
• 在三个 MRI 数据集上评估该方法，并与预文本任务、数据增强和半监督方法进行比较。

提出的方法

• 扩展全局对比损失，利用跨体积切片对齐作为领域特定的相似性线索。
• 引入在解码器特征图上运行的局部对比损失，以促进分割的区域级区分。
• 定义两种领域信息化全局对比策略（G^D- 与 G^D）和两种局部策略（L^R 与 L^D），用于构造正负对。
• 先使用全局损失对编码器进行预训练，然后冻结编码器，使用局部损失对解码器模块进行预训练，最后进行全流程的分割微调。
• 在分阶段的流程中结合全局和局部预训练，无需联合损失权重。

实验结果

研究问题

• RQ1 来自对齐体积的领域特定全局线索能否提升用于分割的对比预训练？
• RQ2 在解码器特征上添加局部对比损失是否比仅全局更能获得更好的像素级表示用于分割？
• RQ3 在标注有限的情形下，领域信息化对比策略与随机增强或预文本任务相比如何？
• RQ4 所提出的预训练是否与数据增强和半监督方法互补？

主要发现

• 全球策略（G^D）利用跨体积切片相似性优于随机增强基线。
• 在解码器上的局部预训练，使用 L^R 或 L^D，相较仅编码器预训练带来额外提升。
• 所提出的组合（G^D+ L^R）在 ACDC、Prostate、MMWHS 数据集上，在仅使用1–2个训练体积时取得显著改进。
• 该方法仅用两个带标签的体积就接近基准性能，并且在与数据增强或半监督方法结合时仍有增益。