[论文解读] Disruptive Autoencoders: Leveraging Low-level features for 3D Medical Image Pre-training

Jeya Maria Jose Valanarasu, Yucheng Tang+9 more|arXiv (Cornell University)|Jul 31, 2023

Radiomics and Machine Learning in Medical Imaging被引用 5

一句话总结

该论文提出了一种新型的自监督预训练框架——破坏性自编码器（Disruptive Autoencoders, DAE），用于3D医学影像，通过结合局部掩码（在通道嵌入上掩码而非标记）与低层次扰动（如噪声和下采样），增强低层次特征的学习。该方法在多个3D医学影像分割基准上实现了最先进性能，包括在BTCV多器官分割挑战赛的公开测试排行榜上排名第一。

ABSTRACT

Harnessing the power of pre-training on large-scale datasets like ImageNet forms a fundamental building block for the progress of representation learning-driven solutions in computer vision. Medical images are inherently different from natural images as they are acquired in the form of many modalities (CT, MR, PET, Ultrasound etc.) and contain granulated information like tissue, lesion, organs etc. These characteristics of medical images require special attention towards learning features representative of local context. In this work, we focus on designing an effective pre-training framework for 3D radiology images. First, we propose a new masking strategy called local masking where the masking is performed across channel embeddings instead of tokens to improve the learning of local feature representations. We combine this with classical low-level perturbations like adding noise and downsampling to further enable low-level representation learning. To this end, we introduce Disruptive Autoencoders, a pre-training framework that attempts to reconstruct the original image from disruptions created by a combination of local masking and low-level perturbations. Additionally, we also devise a cross-modal contrastive loss (CMCL) to accommodate the pre-training of multiple modalities in a single framework. We curate a large-scale dataset to enable pre-training of 3D medical radiology images (MRI and CT). The proposed pre-training framework is tested across multiple downstream tasks and achieves state-of-the-art performance. Notably, our proposed method tops the public test leaderboard of BTCV multi-organ segmentation challenge.

研究动机与目标

• 为解决在3D医学影像中学习鲁棒低层次表征的挑战，这对器官和病灶分割等细粒度任务至关重要。
• 克服原始掩码自编码器（MAE）在重建3D放射影像中精细解剖结构方面的局限性。
• 设计一种预训练框架，以统一且对比学习感知的方式有效利用多种医学影像模态（如MRI、CT）。
• 整理并利用大规模、多模态的3D医学影像数据集，以实现有效的预训练。
• 证明更强的低层次特征学习可提升下游医学影像分割任务的性能。

提出的方法

• 提出一种新颖的局部掩码策略，即在通道嵌入上进行掩码而非空间标记，以增强模型学习局部、细粒度特征的能力。
• 引入破坏性扰动——在标记化前施加加性噪声和下采样，以进一步激发低层次表征学习。
• 将这些破坏性操作与基于Transformer的编码器及卷积解码器结合，从受损的潜在表征中重建原始3D图像。
• 引入跨模态对比损失（Cross-Modal Contrastive Loss, CMCL），在预训练过程中对齐不同影像模态（如MRI、CT）的特征，实现模态感知的表征学习。
• 使用大规模、经整理的3D MRI和CT影像体积数据集对模型进行预训练，以覆盖多样的解剖结构和扫描参数。
• 采用自监督预训练结合下游分割任务的微调策略，并在BTCV和PROMISE12等标准基准上进行评估。

Figure 1: Disruptive Autoencoders: Here, we disrupt the 3D medical image with a combination of low-level perturbations - noise and downsampling, followed by tokenization and local masking. These disrupted tokens are then passed through a transformer encoder and convolutional decoder to learn to reco

实验结果

研究问题

• RQ1 一种通过有针对性破坏强调低层次特征学习的预训练框架，能否改善3D医学影像的表征学习？
• RQ2 局部掩码（在通道嵌入上掩码而非空间标记）是否相比标准标记掩码，能带来更好的重建效果和下游性能？
• RQ3 组合的低层次扰动（噪声、下采样）与局部掩码相比单一扰动，在预训练有效性上表现如何？
• RQ4 跨模态对比学习在多模态3D医学影像预训练中能在多大程度上提升性能？
• RQ5 在更强低层次特征学习的预训练下，是否能实现更好的泛化能力，尤其是在低数据设置下？

主要发现

• 破坏性自编码器在BTCV多器官分割挑战赛的公开测试排行榜上实现了最先进性能，优于先前方法（包括MAE和Swin-UNETR）。
• 在BTCV数据集上，DAE的Dice分数达到0.8472，甚至超过Swin-UNETR中三种掩码任务的组合方法。
• 在PROMISE12数据集上，DAE在所有数据设置下均优于随机初始化和MAE，尤其在20%和50%训练数据下表现显著提升，展现出强大的数据效率。
• 消融研究证实，局部掩码与低层次扰动的组合性能最佳，且仅使用局部掩码的性能也优于单独使用噪声或下采样。
• CKA分析显示，从预训练到微调模型，网络早期层保留了更多低层次特征，证实低层次表征更稳定，且对下游性能至关重要。
• 引入跨模态对比损失（CMCL）可提升微调性能，表明模态感知的预训练能增强跨不同影像模态的泛化能力。

Figure 2: Comparison of reconstruction quality. It can be observed that masked image modelling produces a coarse reconstruction for radiology without local context while the proposed disruptions in this work obtain sharper reconstructions recovering meaningful fine details.