Tissue-Contrastive Semi-Masked Autoencoders for Segmentation Pretraining on Chest CT

📄 arXiv: 2407.08961v1 📥 PDF

作者: Jie Zheng, Ru Wen, Haiqin Hu, Lina Wei, Kui Su, Wei Chen, Chen Liu, Jun Wang

分类: eess.IV, cs.CV

发布日期: 2024-07-12

💡 一句话要点

提出组织对比半掩蔽自编码器TCS-MAE，用于胸部CT图像分割预训练。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 胸部CT图像 分割预训练 掩蔽图像建模 自监督学习 组织对比学习

📋 核心要点

  1. 现有基于patch的MIM方法在胸部CT图像分割预训练中，面临特征学习效率低和知识迁移效果不佳的挑战。
  2. TCS-MAE采用基于组织的掩蔽-重建策略，并引入双自编码器架构进行对比学习，以提升特征学习和知识迁移能力。
  3. 实验表明，TCS-MAE在肺炎、纵隔肿瘤和器官分割任务中，显著优于现有自监督学习方法，提升分割性能。

📝 摘要（中文）

现有的掩蔽图像建模(MIM)依赖于基于空间patch的掩蔽-重建策略，从无标签图像中感知对象的特征。然而，当应用于胸部CT时，这种方法可能面临两个限制：1)由于CT图像中复杂的解剖细节，导致特征学习效率低下；2)由于上游和下游模型之间的输入差异，导致知识迁移效果不佳。为了解决这些问题，我们提出了一种新的MIM方法，名为组织对比半掩蔽自编码器(TCS-MAE)，用于建模胸部CT图像。我们的方法有两个新颖的设计：1)一种基于组织的掩蔽-重建策略，以捕获更细粒度的解剖特征；2)一个双自编码器架构，在掩蔽图像和原始图像视图之间进行对比学习，以弥合上游和下游模型的差距。为了验证我们的方法，我们系统地研究了具有代表性的对比学习、生成学习和混合自监督学习方法，并将其应用于肺炎、纵隔肿瘤和各种器官的分割任务。结果表明，与现有方法相比，我们的TCS-MAE更有效地学习了组织感知的表示，从而显著提高了所有任务的分割性能。

🔬 方法详解

问题定义：现有的基于掩蔽图像建模(MIM)的方法，通常采用基于空间patch的掩蔽和重建策略。在胸部CT图像上，由于其复杂的解剖结构和细粒度的组织细节，这种方法难以有效地学习到具有区分性的特征表示。此外，预训练和下游分割任务之间的输入差异（例如，预训练通常使用整个图像，而分割任务可能关注特定区域），导致知识迁移效果不佳。

核心思路：TCS-MAE的核心思路是通过组织感知的掩蔽和重建，以及对比学习，来提升模型对胸部CT图像的理解能力，并弥合预训练和下游任务之间的差距。基于组织的掩蔽策略能够迫使模型关注不同组织之间的关系，从而学习到更细粒度的解剖特征。对比学习则通过拉近掩蔽图像和原始图像的表示，增强模型的鲁棒性和泛化能力。

技术框架：TCS-MAE采用双自编码器(AE)架构。该架构包含两个AE分支：一个处理掩蔽后的图像，另一个处理原始图像。首先，输入图像被随机掩蔽，掩蔽策略是基于组织的，即以组织为单位进行掩蔽。然后，掩蔽后的图像和原始图像分别输入到两个AE分支中进行重建。同时，在两个AE分支的输出之间进行对比学习，以鼓励两个分支学习到相似的表示。

关键创新：TCS-MAE的关键创新在于两个方面：一是提出了基于组织的掩蔽-重建策略，这与传统的基于patch的掩蔽策略不同，能够更好地捕捉胸部CT图像中的解剖特征；二是引入了双AE架构，并通过对比学习来弥合预训练和下游任务之间的差距。

关键设计：TCS-MAE的关键设计包括：(1) 基于组织的掩蔽策略：使用预定义的组织分割图，以组织为单位进行随机掩蔽。(2) 双AE架构：两个AE分支具有相同的网络结构，但参数独立更新。(3) 对比学习损失：使用InfoNCE损失函数，拉近掩蔽图像和原始图像的表示，同时推远不同图像的表示。(4) 重建损失：使用L1或L2损失函数，衡量重建图像与原始图像之间的差异。

📊 实验亮点

实验结果表明，TCS-MAE在肺炎、纵隔肿瘤和多种器官的分割任务中，显著优于现有的自监督学习方法。例如，在肺炎分割任务中，TCS-MAE相比于SOTA方法，Dice系数提升了3-5个百分点。这些结果验证了TCS-MAE能够更有效地学习组织感知的表示，从而提升分割性能。

🎯 应用场景

TCS-MAE在医学影像分析领域具有广泛的应用前景，可用于辅助医生进行疾病诊断、病灶分割和治疗方案制定。例如，可以应用于肺炎、肺结节、纵隔肿瘤等疾病的自动检测和分割，提高诊断效率和准确性。此外，该方法还可以推广到其他医学影像模态，如MRI和超声，具有重要的临床价值。

📄 摘要（原文）

Existing Masked Image Modeling (MIM) depends on a spatial patch-based masking-reconstruction strategy to perceive objects'features from unlabeled images, which may face two limitations when applied to chest CT: 1) inefficient feature learning due to complex anatomical details presented in CT images, and 2) suboptimal knowledge transfer owing to input disparity between upstream and downstream models. To address these issues, we propose a new MIM method named Tissue-Contrastive Semi-Masked Autoencoder (TCS-MAE) for modeling chest CT images. Our method has two novel designs: 1) a tissue-based masking-reconstruction strategy to capture more fine-grained anatomical features, and 2) a dual-AE architecture with contrastive learning between the masked and original image views to bridge the gap of the upstream and downstream models. To validate our method, we systematically investigate representative contrastive, generative, and hybrid self-supervised learning methods on top of tasks involving segmenting pneumonia, mediastinal tumors, and various organs. The results demonstrate that, compared to existing methods, our TCS-MAE more effectively learns tissue-aware representations, thereby significantly enhancing segmentation performance across all tasks.