[论文解读] Local contrastive loss with pseudo-label based self-training for semi-supervised medical image segmentation

Krishna Chaitanya, Ertunç Erdil+2 more|arXiv (Cornell University)|Dec 17, 2021

Advanced Neural Network Applications被引用 7

一句话总结

本文提出了一种基于伪标签的自训练新型局部对比损失，用于半监督医学图像分割，通过基于语义伪标签而非随机增强来定义正样本对，从而利用有标签和无标签数据。该方法在仅使用一个或两个3D有标签体积的情况下，在心脏和前列腺MRI数据集上实现了最先进性能，显著优于现有的半监督和对比学习方法。

ABSTRACT

Supervised deep learning-based methods yield accurate results for medical image segmentation. However, they require large labeled datasets for this, and obtaining them is a laborious task that requires clinical expertise. Semi/self-supervised learning-based approaches address this limitation by exploiting unlabeled data along with limited annotated data. Recent self-supervised learning methods use contrastive loss to learn good global level representations from unlabeled images and achieve high performance in classification tasks on popular natural image datasets like ImageNet. In pixel-level prediction tasks such as segmentation, it is crucial to also learn good local level representations along with global representations to achieve better accuracy. However, the impact of the existing local contrastive loss-based methods remains limited for learning good local representations because similar and dissimilar local regions are defined based on random augmentations and spatial proximity; not based on the semantic label of local regions due to lack of large-scale expert annotations in the semi/self-supervised setting. In this paper, we propose a local contrastive loss to learn good pixel level features useful for segmentation by exploiting semantic label information obtained from pseudo-labels of unlabeled images alongside limited annotated images. In particular, we define the proposed loss to encourage similar representations for the pixels that have the same pseudo-label/ label while being dissimilar to the representation of pixels with different pseudo-label/label in the dataset. We perform pseudo-label based self-training and train the network by jointly optimizing the proposed contrastive loss on both labeled and unlabeled sets and segmentation loss on only the limited labeled set. We evaluated on three public cardiac and prostate datasets, and obtain high segmentation performance.

研究动机与目标

• 通过减少对大规模有标签数据集的依赖，解决医学图像分割中高标注成本的挑战。
• 通过引入来自伪标签的语义标签信息，改善半监督分割中的局部表征学习。
• 克服现有局部对比方法依赖无监督代理标签或空间邻近性的局限性。
• 开发一种联合训练框架，结合在有限有标签数据上的分割损失和在有标签与无标签数据上的对比损失。
• 通过迭代式伪标签优化和一致性正则化，提升模型的泛化能力和鲁棒性。

提出的方法

• 定义一种局部对比损失，将具有相同伪标签的像素视为正样本对，不同伪标签的像素视为负样本对。
• 使用学生网络生成无标签图像的伪标签，并在训练过程中迭代优化。
• 通过联合优化有标签和无标签数据集上的对比损失，以及仅在有标签数据集上优化标准分割损失，实现自训练。
• 采用图像内和图像间表征匹配策略，比较同一图像不同增强视图之间的特征以及不同图像之间的特征。
• 通过基于变换视图之间重叠阈值的伪标签过滤，应用一致性正则化，以提高预测质量。
• 使用随机网络初始化或来自自监督预训练的预训练权重初始化模型。

实验结果

研究问题

• RQ1 与随机增强或空间邻近性相比，语义伪标签是否能改善半监督医学图像分割中的局部表征学习？
• RQ2 在有标签和无标签数据上联合优化对比损失是否能带来优于标准自训练的分割性能？
• RQ3 在缺乏真实标签的情况下，伪标签的质量如何影响对比损失的性能？
• RQ4 不同的表征匹配策略（图像内 vs. 图像间）对模型收敛性和准确率有何影响？
• RQ5 在无标签数据上进行预训练是否能进一步提升所提出的对比自训练框架的性能？

主要发现

• 在仅使用一个有标签3D体积的情况下，该方法在ACDC数据集上实现了0.881的平均DSC，优于最先进半监督方法。
• 使用八个有标签体积时，该方法在图像内匹配下达到0.885的DSC，在图像间匹配下达到0.897的DSC，表现出良好的鲁棒性和可扩展性。
• 使用Dice重叠为0.7或0.8的阈值过滤伪标签并未带来比使用所有伪标签（阈值=0）更高的性能，表明在此设置下高置信度过滤并无益处。
• 当使用两个有标签体积时，图像间匹配优于图像内匹配，表明跨受试者对齐有助于提升泛化能力。
• 该方法在低样本设置（一个或两个有标签体积）下，相较于同期的对比学习和半监督学习基线模型，取得了显著性能提升。
• 在无标签数据上进行预训练与所提出的对比自训练框架相结合，进一步提升了性能，证实了预训练与自监督的互补性。