【论文阅读笔记】M3AE: Multimodal Representation Learning for Brain Tumor Segmentation with Missing Modalitie

Liu H, Wei D, Lu D, et al. M3AE: Multimodal Representation Learning for Brain Tumor Segmentation with Missing Modalities[J]. arXiv preprint arXiv:2303.05302, 2023. 【开源】

本文的核心思想是介绍了一个名为M3AE（Multimodal Masked Autoencoders）的新框架，用于处理在磁共振成像（MRI）中脑肿瘤分割任务时遇到的缺失模态问题。M3AE框架包含三个创新组件：

1. 多模态掩蔽自编码器：这是一种自监督学习方法，用于在缺失某些模态的情况下学习稳健的表示。这意味着即使在某些MRI模态数据不可用的情况下，该系统也能有效地识别和分割脑肿瘤。
2. 模型反演的模态补全：这个组件用于生成缺失模态的替代图像，帮助改善肿瘤分割的准确性。这种方法通过优化生成的图像来补充缺失的信息，而不是简单地用平均图像或零填充。
3. 异质模态间的自我蒸馏：这是一种内存高效的策略，用于在不同缺失模态情况下进行知识转移。它通过使用单一网络模型来处理所有可能的模态缺失组合，从而减少了所需的参数数量。

这个框架在两个公共基准数据集上进行了广泛的实验，结果表明M3AE在处理缺失模态的脑肿瘤分割任务方面设立了新的技术标准，并且在多模态表示学习方面表现出色。此外，通过消融实验验证了M3AE的三个组件的有效性。未来，作者计划将M3AE应用于完全不同的模态（例如MRI和CT）以及超出BraTS基准的其他基准。

框架组成部分解释

• Self-Supervised Multimodal Representation Learning via M3AE

  对模态的随机子集进行采样，以模拟真实情况，除此之外，我们还对剩余模态的3D patch进行随机掩蔽，就像原始的自然图像MAE一样。从整体上恢复被掩蔽的模态需要网络利用全局的模态间相关性，而恢复被掩蔽的patch需要利用模态内结构完整性和局部的模态间相关性。作者认为：原始的自然图像MAE与论文M3AE的一个显著区别是，前者的掩蔽patch只能从周围的像素中推断出来，而后者的掩蔽patch可以从其他模式中额外推断出来，因此预期更容易。因此在M3AE中设置了更高的组合掩蔽率87.5%(相比何凯明MAE的75%)。

• Model Inversion based Modality Completion:模型反演的概念：模型反演通常指的是从模型的输出反向推导出输入的过程。在这个上下文中，它涉及到利用已经训练好的模型来生成或估计缺失的模态数据。生成缺失模态：在脑肿瘤分割的场景中，MRI图像可能包含多种不同的模态（如T1、T2、FLAIR等），而在某些情况下，某些模态可能缺失。Model Inversion based Modality Completion的目的是生成缺失模态的合成版本，以便于模型能够使用完整的多模态数据进行更准确的分割。优化过程：这个过程涉及到优化一个输入图像，使其在经过模型处理后，能够尽可能接近于已知的目标输出。换句话说，它试图找到一个合成的模态图像，当这个图像被输入到模型中时，模型的输出与实际使用完整模态数据时的输出尽可能相似。补充缺失信息：通过这种方式，Model Inversion based Modality Completion能够为缺失的模态提供一个有效的替代品，这有助于模型更好地理解和处理整个数据集，即使其中一些模态信息缺失。

​ ${\hat{\mathbf{x}}}^{\mathrm{sub}}=\arg {\min }_{{\mathbf{x}}^{\mathrm{sub}}}{\mathcal{L}}_{\mathrm{mse}}\left(\mathbf{x},F\left(S\left(\mathbf{x},{\mathbf{x}}^{\mathrm{sub}}\right)\right)\right)+\gamma \mathcal{R}\left({\mathbf{x}}^{\mathrm{sub}}\right),$

​ 其中：$S\left(\mathbf{x},{\mathbf{x}}^{\mathrm{sub}}\right)$是用$x_{sub}$中位置对应的内容替换$x$的屏蔽内容，$F$为骨干网$F$与回归头级联的重构函数，$R$为正则项，$\gamma$为权值。文章对原始的MAE 做了一个修改，用$x_{sub}$替换被屏蔽的内容，并按照与非屏蔽内容相同的方式处理它们，而不是丢弃它们。这样做的直觉是，为了产生更好的重构，最佳替代必须捕获最具代表性的特定模式模式，这也有望帮助多模态分割的目标任务。实现上，$x_{sub}$通过反向传播更新，同时更新网络参数。这样就不需要引入额外的模块，$x_{sub}$的优化只会产生边际成本。

• Fine-Tune with Heterogeneous Missing-Modal Self Distillation for Tumor Segmentation：在每个batch中，通过modality dropout随机抽取一个被试的两种不同的失模态情况(包括全模态的特殊情况)作为网络输入，以一致性损失$L_{con}$鼓励它们之间语义特征一致:

  ${\mathcal{L}}_{\text{con }}\left({\mathbf{x}}_0,{\mathbf{x}}_1,{\hat{\mathbf{x}}}^{\mathrm{sub}}\right)={\mathcal{L}}_{\mathrm{mse}}\left({\mathbf{f}}_0,{\mathbf{f}}_1\right)$

​ 其中$x_0$和$x_1$是$x$的两个随机缺失模态实例，$f_0$和$f_1$分别从$S\left({\mathbf{x}}_0,{\hat{\mathbf{x}}}^{\mathrm{sub}}\right)$和$S\left({\mathbf{x}}_1,{\hat{\mathbf{x}}}^{\mathrm{sub}}\right)$提取对应的特征映射。通过上述公式实现知识的相互迁移。从多模态到少模态的知识迁移促进了缺失模态信息的恢复，反之(特别是从单模态到多模态)则增强了模态特异性特征。此外，由于$x_0$和$x_1$在每个epoch都是通过随机模态dropout获得的，所以自蒸馏在异构缺失模态情况之间传递知识，而不是像成对共同训练那样在固定的情况之间传递知识。

数据集

BraTS 2018和BraTS 2020

训练条件

2080Ti *2 PyTorch1.7

模型效果

缺点

两阶段独立训练，模型训练时只是考虑两个模态缺失的情况