[论文解读] M3AE: Multimodal Representation Learning for Brain Tumor Segmentation with Missing Modalities

Hong Liu, Wei Dong+4 more|arXiv (Cornell University)|Mar 9, 2023

Brain Tumor Detection and ClassificationNeuroscience被引用 3

一句话总结

本文提出M3AE，一种新颖的自监督多模态掩码自编码器框架，用于脑肿瘤分割，通过在预训练过程中同时掩码随机模态和图像块来处理缺失的MRI模态。它利用模型反演生成代表性全模态图像以替代缺失模态，并采用内存高效的自蒸馏方法统一异构缺失模态场景下的知识，仅使用单一紧凑模型即在BraTS 2018和2020上实现了最先进性能。

ABSTRACT

Multimodal magnetic resonance imaging (MRI) provides complementary information for sub-region analysis of brain tumors. Plenty of methods have been proposed for automatic brain tumor segmentation using four common MRI modalities and achieved remarkable performance. In practice, however, it is common to have one or more modalities missing due to image corruption, artifacts, acquisition protocols, allergy to contrast agents, or simply cost. In this work, we propose a novel two-stage framework for brain tumor segmentation with missing modalities. In the first stage, a multimodal masked autoencoder (M3AE) is proposed, where both random modalities (i.e., modality dropout) and random patches of the remaining modalities are masked for a reconstruction task, for self-supervised learning of robust multimodal representations against missing modalities. To this end, we name our framework M3AE. Meanwhile, we employ model inversion to optimize a representative full-modal image at marginal extra cost, which will be used to substitute for the missing modalities and boost performance during inference. Then in the second stage, a memory-efficient self distillation is proposed to distill knowledge between heterogenous missing-modal situations while fine-tuning the model for supervised segmentation. Our M3AE belongs to the 'catch-all' genre where a single model can be applied to all possible subsets of modalities, thus is economic for both training and deployment. Extensive experiments on BraTS 2018 and 2020 datasets demonstrate its superior performance to existing state-of-the-art methods with missing modalities, as well as the efficacy of its components. Our code is available at: https://github.com/ccarliu/m3ae.

研究动机与目标

• 解决因临床或技术限制导致一个或多个MRI模态（T1、T1c、T2、FLAIR）缺失时的脑肿瘤分割挑战。
• 开发一种“全兼容”框架，在无需微调的情况下泛化于所有可能的缺失模态子集。
• 通过自监督预训练学习鲁棒的多模态表征，捕捉缺失数据下的模态间与模态内依赖关系。
• 通过模型反演生成高质量、优化的替代图像以提升推理性能，用于缺失模态的补全。
• 在不引入过多内存或参数开销的前提下，高效地在异构缺失模态配置之间蒸馏语义知识。

提出的方法

• 提出一种多模态掩码自编码器（M3AE），在预训练过程中同时掩码随机选择的MRI模态子集和剩余模态中的随机空间块。
• 将M3AE的重建损失用作自监督信号，以学习对缺失数据具有鲁棒性的多模态表征。
• 应用模型反演通过反向传播M3AE损失来优化一个代表性全模态图像，使其在推理阶段可作为缺失模态的高质量替代品。
• 提出一种内存高效的自蒸馏策略，在微调阶段实现不同缺失模态配置之间的知识迁移，相比双网络联合训练显著降低参数开销。
• 采用单一编码器-解码器架构并共享权重，使模型无需重新训练即可部署于所有可能的模态组合。
• 端到端训练：首先在全模态数据上进行M3AE预训练，然后在异构缺失模态场景下进行自蒸馏微调。

实验结果

研究问题

• RQ1 自监督多模态掩码自编码器能否有效学习在缺失模态下仍具泛化能力的鲁棒表征？
• RQ2 基于模型反演的模态补全是否能提升模态缺失时的分割性能？
• RQ3 内存高效的自蒸馏能否有效在多样化缺失模态配置间迁移语义知识？
• RQ4 所提出的M3AE框架是否在所有缺失模态场景下均优于现有最先进方法的分割精度？
• RQ5 考虑到其单一模型、全兼容的设计，该框架在实际部署中是否高效且实用？

主要发现

• 在BraTS 2018测试集的所有缺失模态场景中，M3AE在平均Dice相似度分数（DSC）上均达到最高，优于四种最先进方法。
• 消融实验表明，若移除M3AE预训练或块掩码机制，性能显著下降，验证了二者在鲁棒表征学习中的必要性。
• 基于模型反演的替代方法在DSC上显著优于零填充方法，甚至优于使用训练集均值图像的基线，后者反而导致性能下降。
• 所提出的自蒸馏策略在性能上优于基线蒸馏方法，并且相比双网络联合训练，模型参数减少了约470万个。
• 在全模态验证集上，M3AE在BraTS 2018上实现了最先进性能，超越非竞赛类SOTA方法，且与竞赛优胜模型保持竞争力。
• 在各种缺失模态条件下，该框架在所有肿瘤区域（水肿、增强肿瘤、坏死核心）均保持优异性能，展现出良好的泛化能力与鲁棒性。