[论文解读] Correlation-Aware Mutual Learning for Semi-supervised Medical Image Segmentation

Shengbo Gao, Ziji Zhang+3 more|arXiv (Cornell University)|Jul 12, 2023

Advanced Neural Network Applications被引用 4

一句话总结

该论文提出了一种相关性感知的互学习方法（Correlation-Aware Mutual Learning, CAML），这是一种新颖的半监督医学图像分割框架，通过显式利用标注数据来指导未标注数据的特征学习，从而提升模型性能。CAML采用一种协同教学机制，包含两个模块——用于密集跨样本注意力的跨样本互注意力（Cross-sample Mutual Attention, CMA）和用于标注与未标注特征间一致性正则化的全相关性一致性（Omni-Correlation Consistency, OCC），在低标注数据设置下实现了最先进性能，当仅使用5%标注数据时Dice分数达到87.34%，在20%标注数据时接近全监督性能。

ABSTRACT

Semi-supervised learning has become increasingly popular in medical image segmentation due to its ability to leverage large amounts of unlabeled data to extract additional information. However, most existing semi-supervised segmentation methods only focus on extracting information from unlabeled data, disregarding the potential of labeled data to further improve the performance of the model. In this paper, we propose a novel Correlation Aware Mutual Learning (CAML) framework that leverages labeled data to guide the extraction of information from unlabeled data. Our approach is based on a mutual learning strategy that incorporates two modules: the Cross-sample Mutual Attention Module (CMA) and the Omni-Correlation Consistency Module (OCC). The CMA module establishes dense cross-sample correlations among a group of samples, enabling the transfer of label prior knowledge to unlabeled data. The OCC module constructs omni-correlations between the unlabeled and labeled datasets and regularizes dual models by constraining the omni-correlation matrix of each sub-model to be consistent. Experiments on the Atrial Segmentation Challenge dataset demonstrate that our proposed approach outperforms state-of-the-art methods, highlighting the effectiveness of our framework in medical image segmentation tasks. The codes, pre-trained weights, and data are publicly available.

研究动机与目标

• 为解决现有半监督医学图像分割方法中对标注数据的利用不足问题，这些方法通常仅将标注数据视为初始监督信号。
• 通过显式建模标注数据与未标注数据之间的跨样本相关性，以实现先验知识的迁移，从而提升模型泛化能力。
• 开发一种互学习框架，通过联合优化标注与未标注数据，增强特征表示能力。
• 即使在仅提供少量标注图像的情况下，也能实现稳健且一致的性能表现。
• 通过最大化从有限标注数据中提取的信息，减少临床应用中对大规模标注数据集的依赖。

提出的方法

• CAML采用协同教学框架，通过两个学生网络在相互伪监督下进行训练，以促进一致性与知识蒸馏。
• 跨样本互注意力（Cross-sample Mutual Attention, CMA）模块在一组样本之间建立密集的双向注意力机制，实现标注与未标注数据之间的协同显著特征学习。
• 全相关性一致性（Omni-Correlation Consistency, OCC）模块构建一个标注特征的记忆库，并计算未标注特征与该参考集之间的全相关性，以在特征层面建模相似性。
• OCC强制两个子模型的全相关性矩阵保持一致性，从而正则化学习过程，并促进未标注数据中统一的特征分布。
• 该框架在互学习循环中集成两个模块，同步更新两个网络，同时实现从标注样本到未标注样本的知识传播。
• 该方法可端到端训练，且不增加推理开销，适用于临床部署。

实验结果

研究问题

• RQ1 能否有效利用标注数据，以提升半监督分割中从未标注医学图像中学习表征的能力？
• RQ2 如何建模标注与未标注数据之间的跨样本相关性，以增强特征迁移与一致性？
• RQ3 与标准的一致性正则化相比，引入多样本间的互注意力是否能提升分割性能？
• RQ4 标注与未标注特征之间的全相关性一致性是否能带来更鲁棒、更具泛化能力的特征学习？
• RQ5 在极端标注稀缺情况下（如5%或10%标注数据），所提出的CAML框架表现如何？

主要发现

• 在5%标注数据下，CAML的Dice分数达到87.34%，标准差为0.05，显著优于基线方法，表现出高度稳定性。
• 在10%标注数据下，CAML的Dice分数达到89.62%，已接近全监督模型的性能上限（90.98%）。
• 在20%标注数据下，CAML的Dice分数达到90.78%，与全监督性能上限的差距极小。
• 消融实验证实，CMA与OCC模块均不可或缺：二者结合在5%标注数据下使Dice分数绝对提升6.42%。
• 该方法在各次运行中均保持较低标准差，表明其具有高度鲁棒性与泛化能力，尤其在低标注场景下表现突出。
• 定性结果表明，CAML生成的分割掩码比所有对比方法更精确、更完整，尤其在边界区域与小结构区域表现更优。