[论文解读] Boosting Semi-supervised Image Segmentation with Global and Local Mutual Information Regularization

Jizong Peng, Marco Pedersoli+1 more|arXiv (Cornell University)|Mar 8, 2021

Advanced Neural Network Applications参考文献 59被引用 22

一句话总结

该论文提出了一种半监督图像分割方法，通过在全局和局部特征表示上应用互信息（MI）正则化来提升性能。通过利用投影头最大化同一图像不同增强视图之间聚类特征嵌入的互信息，并在网络各阶段强制保持一致性，该方法在仅使用5%标注数据的情况下，在医学影像基准上实现了接近全监督的准确率。

ABSTRACT

The scarcity of labeled data often impedes the application of deep learning to the segmentation of medical images. Semi-supervised learning seeks to overcome this limitation by exploiting unlabeled examples in the learning process. In this paper, we present a novel semi-supervised segmentation method that leverages mutual information (MI) on categorical distributions to achieve both global representation invariance and local smoothness. In this method, we maximize the MI for intermediate feature embeddings that are taken from both the encoder and decoder of a segmentation network. We first propose a global MI loss constraining the encoder to learn an image representation that is invariant to geometric transformations. Instead of resorting to computationally-expensive techniques for estimating the MI on continuous feature embeddings, we use projection heads to map them to a discrete cluster assignment where MI can be computed efficiently. Our method also includes a local MI loss to promote spatial consistency in the feature maps of the decoder and provide a smoother segmentation. Since mutual information does not require a strict ordering of clusters in two different assignments, we incorporate a final consistency regularization loss on the output which helps align the cluster labels throughout the network. We evaluate the method on four challenging publicly-available datasets for medical image segmentation. Experimental results show our method to outperform recently-proposed approaches for semi-supervised segmentation and provide an accuracy near to full supervision while training with very few annotated images.

研究动机与目标

• 解决基于深度学习的语义分割在医学图像标注数据有限的挑战。
• 通过表示不变性和空间一致性来利用未标注数据，提升半监督学习性能。
• 克服现有基于一致性学习和对比学习方法的局限性，这些方法通常需要大批次训练或在高维连续特征上表现不佳。
• 开发一种方法，通过离散聚类和互信息来有效正则化全局和局部特征表示。

提出的方法

• 通过最大化同一图像不同增强视图的聚类分配之间的互信息（MI），在编码器中强制实现全局表示不变性。
• 使用投影头将连续特征嵌入映射到离散聚类分配，从而实现高效的MI估计，避免连续密度估计。
• 引入局部MI损失，通过最大化相邻空间特征之间的MI，促进解码器特征图中的空间一致性。
• 应用最终的一致性正则化损失，对齐不同网络阶段的聚类标签，确保语义分组的一致性。
• 在联合优化框架中结合监督损失与基于MI的全局和局部正则化，避免两阶段预训练-微调设置。
• 采用基于熵的对称MI估计方法，避免出现平凡聚类解。

实验结果

研究问题

• RQ1 在离散聚类分配上应用互信息正则化是否能提升半监督医学图像分割性能？
• RQ2 结合全局和局部MI正则化如何增强特征表示并提升分割准确率？
• RQ3 联合优化监督损失与基于MI的损失是否优于两阶段预训练-微调策略？
• RQ4 基于MI的方法是否能在极低标注数据量下实现接近全监督的性能？
• RQ5 与标准的一致性学习和对比学习方法相比，该方法在鲁棒性和效率方面表现如何？

主要发现

• 所提方法在四个公开医学图像分割数据集上优于最先进的半监督方法（如Mean Teacher和熵最小化）。
• 在ACDC数据集上，仅使用5%标注数据，该方法实现了86.7%的平均DSC，接近全监督性能。
• 与两阶段预训练-微调基线相比，联合优化策略将分割准确率提升了5.10%，证明了端到端MI正则化与标注数据结合的优势。
• 该方法在三次独立运行中表现出稳定且一致的性能，标准差较低，表明对随机初始化具有鲁棒性。
• 消融研究证实，全局和局部MI损失均对性能有显著贡献：全局损失提升表示不变性，局部损失增强空间平滑性。
• 通过投影头实现的离散聚类分配使MI估计更加高效，并避免了连续MI估计带来的计算负担。