论文阅读笔记——M3AE:缺失模态下脑肿瘤分割的多模态表征学习

​
论文地址:[2303.05302] M3AE: Multimodal Representation Learning for Brain Tumor Segmentation with Missing Modalities (arxiv.org)
论文代码: https://github.com/young-geng/m3ae_public

前置知识

MRI的四种模态(T1、T2、FLAIR)

1. T1看解刨结构
   
   从上图可以看出白质是白色的部分，灰质是灰色部分，脑脊液是黑色部分

2. T2看病变
   
   上图中的白色区域为病灶部位，由于水分积累导致局部(坏死区、增强肿瘤区及非增强肿瘤区)在T2图像中呈现高亮(白色)，所以T2图像一般适用于判断病灶所处的位置、大小(但难以清晰勾勒出由于水分积累导致的浮肿区域的大小)

3. FLAIR (磁共振成像液体衰减反转回升)
   
   抑制脑脊液的高信号(使其成像变暗)，将T2图像中与脑脊液成像邻近的病灶部位浮肿区显示清楚，更好的表现肿瘤部位周遭情况。

4. T1ce (对比增强加权T1)
   
   在做MRI之前要注射造影剂，在MRI图像中，高亮的地方血供丰富，强化显示说明血流丰富，肿瘤部位正是供血量大的地方，从而可以显示肿瘤部位的内部情况，鉴别肿瘤与非肿瘤性病变。

肿瘤异质性：

肿瘤异质性是恶性肿瘤的特征之一。肿瘤在生长过程中，经过多次分裂增殖，其子细胞呈现出分子生物学或基因方面的改变，从而导致肿瘤的生长速度、侵袭能力、对药物敏感性、预后等各方面产生差异。简而言之，肿瘤异质性指肿瘤中既有致瘤细胞亚群也有非致瘤细胞亚群。本质在于外因(环境)和内因(基因)的双重作用。

肿瘤异质性分类：

1. 肿瘤间异质性：不同肿瘤的细胞之间的基因和表型不同
2. 肿瘤内异质性(本篇文章中涉及的异质性)：相同肿瘤的不同细胞之间的基因和表型不同

肿瘤内异质性：

1. 空间异质性：相同肿瘤不同区域
2. 时间异质性：原发性肿瘤和继发性肿瘤

相似性分数 (DSCs) ：

功能：用于度量两个集合的相似度，取值范围为 [0, 1] ，值越大表示两个集合越相似，常用于判断两个样本的相似性。

性能量化指标HD95:

功能：描述两个点集之间相似度的一种量度。

DSCs与HD95比较：

DSCs对mask的内部填充比较敏感，而HD95 对分割出的边界比较敏感

摘要

主要贡献：提出了一个用于缺失模态脑肿瘤分割的两阶段框架

M3AE框架可以应用于全模态的所有可能子集，且在BraTS-2018和BraTS-2020数据集上的表现优于所有现有的先进方法

第一阶段：提出了一种多模态掩码的自动编码器(M3AE)

其中所有的随机模态(即对全模态进行随机dropout)和剩余模态的随机块(patch)被掩码用于重建任务，用于针对缺失模态鲁棒多模态表征的自监督学习。同时，使用微不足道的成本优化一个代表性的全模态图像，这将用于替代缺失的模态并提高推理过程中的性能

第二阶段：提出内存高效的自蒸馏

在微调模型进行监督分割的同时，提取异质缺失模态之间的知识。

引言：

Figure1：
( a )BraTS 2018中四种模态的示例图像，其中一种或多种模态在实际中可能缺失(e.g. T1在红框中表示缺失)
( b )对应肿瘤区域：蓝色为水肿区，黄色为增强区，绿色为坏死区
( c )部署模型的大小和平均Dice相似性分数在BraTS 2018测试拆分所有缺失模态情况，圆圈大小表示GFLOPS。与四种最新方法相比，M3AE以紧凑和高效的模型实现了最佳的性能。

创新点：

这篇文章的主要创新点在于将多模态掩码的自编码器、基于模型反演的模态补全和存储高效的自蒸馏集成在编码器-解码器的架构中。

多模态掩码的自动编码器：

其中全模态的随机子集和其余模态的随机块同时被掩码，模型利用全局和局部固有的模态间相关性以及模态内的局部语义恢复被掩码的内容。

基于模型反演的模态补全：

这是首次尝试将模型反演应用于医学图像模态补全，通过模型反演学习到具有代表性的全模态图像，作为推断过程中缺失模态的替代物，同时利用反向传播的自监督M3AE损失对替代图像进行优化。

存储高效的自蒸馏：

用于促进不同模态之间的语义一致性

实验表明：

1. 该框架对于缺失模态的鲁棒性以及对现有的catch-all(一个模型应用于多种模态组合)和专用(对每一种模态组合构建一个模型，N个模态构建2ⁿ-1个模型)方法的优越性
2. 其构件的有效性
3. 其对全模态的多模态表示学习的能力

相关工作

模态缺失下多模态脑肿瘤分割

现有方法大致分为两类

1. 一个模型用于解决一种模态缺失的问题(dedicated)：对于N种模态通常需要构建$2^N-1$个模型
   例如：
   1. 协同训练策略被使用到从全模态到缺失模态网络中提取知识。
   2. 在图像(整体语义)和像素(即网络输出)水平上将知识从多模态教师网络提取到单模态学生网络。
   3. 对抗协同训练网络(ACN):通过熵和知识对抗学习对潜在表示进行对齐，增强了全缺失模态蒸馏
   4. 样式匹配U-Net (SMU-Net)将全模态数据和缺失模态数据的共同潜在空间分解为内容和样式表示，并使用内容和样式匹配机制将全模态网络中的信息特征提取到缺失模态网络中。
   优点：性能不错(特别是在一个以上模态缺失的情况下)
   缺点：训练与部署(N个模态需要$2^N-1$个模型)的计算和内存成本都很高
   本文做法：采用单一模型适用于所有缺失模态的框架
   一组通过合成具有保真度的缺失模态来解决该问题：
   1. 生成对抗网络(GANs)：缺点难以训练用于3D图像的生成，且可能会为训练和部署带来额外的开销
   2. 生成式的模型不能完全再现通过钆元素造影剂获得的对比图像
   本文做法：不生成缺失模态的图像，而是选择用微不足道的成本优化一个通用的全模态替代图像，从而促进缺失模态的分割

2. 一个模型用于解决多种模态缺失组合的问题(catch-all)
   使用一个模型来处理所有缺失模态的情况：
   1. 通常使用模态特异性Encoder将模态嵌入到共享的潜在空间中，然后进行特征融合和进一步处理，以产生分割(Havaei 2016)
   2. 异模态变分编码器解码器(HVED)融合了多模态变分自编码器，从共享的潜在变量中重构模态，迫使形成真正共享的潜在表示。
   3. 对抗训练，使缺失模态的特征图适应完整模态的特征图
   4. 潜在相关表征学习：通过建模模态间的相关性来估计缺失模态在潜在空间中的表示
   5. 生成特征增强图像，以提供缺失模态的必要特征表示
   6. 区域感知融合网络：依靠一个区域感知融合模块，根据不同的区域自适应的从可用的图像模态进行特征融合。
   缺点：设计复杂，使用复杂交互的多个编码器(有时也有多个解码器)
   本文做法：使用单编码器-单解码器架构(本质上是一个3D U-Net)可以学习丰富的多模态表示，并同时处理异构的缺失模态情况

面向医学图像分析的自监督多模态表示学习：

虽然针对单模态医学图像有相当多的工作提出了有效的自监督表示学习，但研究者们刚刚开始探索基于多模态医学图像“多模态性”的独特前文本任务。

1. Table等人提出了一种新颖的跨模态拼图(CMJP)任务，来学习模态无关的特征嵌入，尽管它是有效的，
   缺点：CMJP是针对2D网络提出的并且没有考虑缺失模态的实际情况，而且将它扩展到3D网络或缺失模态的方法并不简单。
2. 在存在模态丢失的HVED中，正则化模态重构任务是一种有效的自监督任务，用于准备丢失的模态。
   缺点：它只关注全部模态间的相关性，而忽略了局部结构完整性，但局部结构完整性对于学习用于分割的强大表示有极大价值的。

本文做法 M3AE通过同时建模全部模态间相关性和单模态内解刨完整性来学习丰富的多模态表示，从随机丢弃的模态和随机掩码块的输入中学习。 为此M3AE对缺失的模态具有内在的鲁棒性，同时适用于精细尺度的语义学习。

知识蒸馏

最初是为了将知识从一个或多个教师网络(往往是大型的复杂网络或模型集成)压缩到一个学生网络(常常是轻量化的模型)而提出的。
对于存在缺失模态的多模态分割，有工作提出通过协同训练将全模态网络的“暗知识”迁移到缺失模态网络中。虽然取得了不错的性能，但由于双网络架构，协同训练策略为训练带来了不可忽略的内存开销。此外，每一对协同训练网络只关注完整模态和特定类型缺失模态(e.g.独自充满到T1)之间的固定相关性，未能挖掘所有不同缺失模态情况所共享的共同语义。
本文做法： 采用自蒸馏，我们的框架在单个网络中提取不同类型模态缺失情况(包括全模态的特殊情况)之间的共享语义信息 ，在消耗比先前方法更少的训练资源的同时，实现了更好地缺失模态和全模态的分割性能。

方法

Figure 2：框架概述
本文提出的框架概述如图2所示，包括一个预训练和一个微调阶段。

1. 预训练阶段：提出了一种新颖的多模态掩码的自动编码器(M3AE)，用于自监督学习对缺失模态的鲁棒表示。同时，通过反向传播M3AE(${\mathcal{L}}_{\mathrm{mse}}$)的训练损失，通过模型反演学习缺失模态的全模态替代。
2. 微调阶段：提出一种存储高效的自蒸馏策略，通过一致性损失(${\mathcal{L}}_{\mathrm{con}}$)提取不同类型的缺失模态情形之间共享的语义信息，同时微调网络使用监督损失(${\mathcal{L}}_{\mathrm{seg}}$)用于脑肿瘤分割。训练好的分割网络作为一个catch-all模型，既可以用于任何全模态的子集(即不同缺失模态情形)，也可以用于全模态。

接下来，首先详细描述一下框架中新提出的构件，包括M3AE、模型反演和自蒸馏，以及集成他们的训练和推理过程。

基于M3AE的自监督多模态表示学习

掩码自编码器(MAEs)作为可扩展的自监督视觉学习器已经被证实是成功的，其中的前件任务是在给定原始信号的部分观测值的情况下重建原始信号。受此启发，本文提出了一种面向医学图像的多模态掩码自编码器(M3AE) 。考虑一幅多模态图像$\text{x }\in {\mathbb{R}}^{N\times D\times H\times W}$，其中$W$,$H$和$D$分别表示图像的宽度、高度和深度，N为模态数。在实际应用中，N个模态中的任何一个子集都可能缺失。因此，我们采用一个随机的模态子集进行掩码来模拟真实情况，除此之外，对该模态的3D块进行随机的掩码，如在原始的自然图像MAE中。恢复作为整体被掩盖的模态需要网络利用全局模态间的相关性，而恢复被掩码的块需要同时利用模态内的结构完整性和模态间的局部相关性。因此，M3AE同时促进了解剖知识和模态间相关性的自监督学习，将重建图像和原始图像(图2中的$\hat{\mathbf{x}}$和$\mathbf{x}$)之间的均方损失作为M3AE的损失函数，记为${\mathcal{L}}_{\mathrm{mse}}$。
原始的自然图像MAE的掩码块只能从周围的上下文中推断出来，而M3AE的掩码块可以从其他模态中推断出来，这更容易补全掩码块。因此，我们在M3AE中实证设置了更高的合并掩码率87.5%(与He等人使用的75%相比)，以使自监督任务变得非平凡。

基于模型反演的模态补全

大多数现有的模态补全方法，诸如GANs来合成缺失模态的图像，导致除了分割网络外，额外的模型和相关训练和部署开销。通过模型反演，在这项工作中提出以微不足道的成本从M3AE的训练过程中合成空间和时间高效的全模态替代图像。模型反演长期以来一直被用于可解释的深度学习，以合成某些网络预测中最具代表性的图像，例如用于分类的显著图。具体来说，我们优化了一个图像${\mathrm{x}}^{\mathrm{sub}}\in {\mathbb{R}}^{N\times D\times H\times W}$当用于替换$\mathbf{x}$的掩码内容(不论是整体模态还是模态内掩码块)时，可以导致更小的重建误差：

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

式中，$S(\mathbf{x},{\mathbf{x}}^{\mathbf{sub}})$为将$\mathbf{x}$的内容替换为${\mathbf{x}}^{\mathbf{sub}}$中位置对应的掩码内容的操作(将两个图像对齐，用${\mathbf{x}}^{\mathbf{sub}}$替换需要掩码的内容)，其中的$F$为级联骨干网络$f$和回归头的重构函数，$\mathcal{R}$为正则化项，$\gamma$为权重。仿照Nguyen，我们对$\gamma =0.005$的$\mathcal{R}$使用少量的$L_2$正则化项。在这里，我们对原始的MAE进行了修改，用${\mathbf{x}}^{\mathbf{sub}}$替换被掩码的内容，并以与未被掩码的内容相同的方式进行处理，而不是丢弃它们。直觉上，为了获得更好的重建效果，最优的全模态替代图像必须捕获最具代表性的模态特异性pattern，这也有望帮助完成多模态分割的目标任务。为了实现，${\mathbf{x}}^{\mathbf{sub}}$通过反向传播进行更新，同时更新网络参数。这样，不需要引入任何额外的模块，对${\mathbf{x}}^{\mathbf{sub}}$的优化只产生微不足道的成本。优化(${\hat{\mathbf{x}}}^{\mathrm{sub}}$)的例子如图3所示。注意，对于给定的训练数据集，只学习一个(${\hat{\mathbf{x}}}^{\mathrm{sub}}$)，这可以认为是训练数据的一种多模态表示形式，适用于所有的受试者。

Figure 3：在BraTS 2018上通过模型反演优化的全模态替代图像(${\hat{\mathbf{x}}}^{\mathrm{sub}}$)实例

使用异构缺失模态自蒸馏进行肿瘤分割的模型微调

通过协同训练，从全模态到缺失模态网络的知识蒸馏在缺失模态的多模态分割中被证明是有效的，尽管由于配对网络而付出了巨大的计算开销。 受自蒸馏策略的启发，我们提出存储高效自蒸馏策略，用于在单个网络内的异构缺失模态之间进行知识蒸馏。具体来说，在每个批量中，通过模态dropout随机采样一个受试者的两种不同的缺失模态情况(包括全模态的特殊情况)作为网络输入，并通过一致性损失${\mathcal{L}}_{\mathrm{con}}$鼓励它们之间的语义特征保持一致：$${\mathcal{L}}_{\mathrm{con}}({\mathbf{x}}_0,{\mathbf{x}}_1,{\hat{\mathbf{x}}}^{\mathrm{sub}})={\mathcal{L}}_{\mathrm{mse}}({\mathbf{f}}_0,{\mathbf{f}}_1),(2)$$式中：${\mathbf{x}}_0$和${\mathbf{x}}_1$为$\mathbf{x}$的两个随机缺失模态实例，${\mathbf{f}}_0$和${\mathbf{f}}_1$$\in {\mathbb{R}}^{C\times D^{\prime }\times H^{\prime }\times W^{\prime }}$分别是从$S({\mathbf{x}}_0,{\hat{\mathbf{x}}}^{\mathrm{sub}})$和$S({\mathbf{x}}_1,{\hat{\mathbf{x}}}^{\mathrm{sub}})$中提取的对应特征图，并且$C$、$D^{\prime }$、$H^{\prime }$和$W^{\prime }$分别是特征图的通道数、深度、高度和宽度。通过Eqn实现知识的相互迁移。(2)是双向有益的，从更多的模态到更少的模态的知识转移鼓励恢复丢失的模态信息，而反向的(特别是从单模态到多模态)增强了模态特异性的特征。除此之外，由于${\mathbf{x}}_0$和${\mathbf{x}}_1$是在每个epoch中随机模态丢弃得到的，因此我们的自蒸馏是在异构缺失模态情形之间进行知识迁移，而不是像成对协同训练那样在固定的模态之间进行知识迁移。在网络的瓶颈处提取潜在空间(见图2) 。

训练和推理过程

使用两阶段(预训练和微调)训练方案。在第一阶段，对M3AE同时进行随机模态和随机块替换两种方式进行训练，同时对替换图像${\mathbf{x}}^{\mathbf{sub}}$进行优化，优化问题可以表示为：$$\underset{F,{\mathbf{x}}^{\mathrm{sub}}}{\min }{\mathcal{L}}_{\mathrm{mse}}(\mathbf{x},F(S(\mathbf{x},{\mathbf{x}}^{\mathrm{sub}})))+\gamma \mathcal{R}({\mathbf{x}}^{\mathrm{sub}}).(3)$$该阶段作为自监督预训练，在全模态替代图像$\hat{\mathbf{x}}$尽可能接近原有全模态图像$\mathbf{x}$的同时,级联骨干网络学习固有的模态间相关性和解刨学完整性。然后再第二阶段，通过随机drop0到3个模态，为每个案例实例化2个缺失模态volume${\mathbf{x}}_0$和${\mathbf{x}}_1$。同时将第一阶段使用的回归头替换为随机初始化的分割头$f_s$，优化问题现在变成了： min ⁡ f , f s , { f d } λ L c o n ( x 0 , x 1 , x ^ s u b ) + ∑ i = 0 1 L s e g ( s g t , x i , x ^ s u b ) ,  (4) \min_{f,f_s,\{f_d\}}\lambda\mathcal{L}_{\mathrm{con}}(\mathbf{x}_0,\mathbf{x}_1,\mathbf{\hat{x}}^\mathrm{sub})+\sum_{i=0}^1\mathcal{L}_{\mathrm{seg}}(\mathbf{s}^\mathrm{gt},\mathbf{x}_i,\mathbf{\hat{x}}^\mathrm{sub}),\text{ (4)} f,fs​,{fd​}min​λLcon​(x0​,x1​,x^sub)+i=0∑1​Lseg​(sgt,xi​,x^sub), (4)其中$\lambda$为权重，${\mathcal{L}}_{\mathrm{seg}}$为深度监督的分割损失。 L s e g ( s g t , x i , x ^ s u b ) = ∑ α ∈ { 1 , 1 2 , 1 4 } L ( s g t , s ^ i α ) , i ∈ { 0 , 1 } , ( 5 ) \mathcal{L}_{\mathrm{seg}}(\mathrm{s}^{\mathrm{gt}},\mathrm{x}_i,\hat{\mathrm{x}}^{\mathrm{sub}})=\sum_{\alpha\in\{1,\frac12,\frac14\}}\mathcal{L}(\mathrm{s}^{\mathrm{gt}},\hat{\mathrm{s}}_i^{\alpha}),i\in\{0,1\},(5) Lseg​(sgt,xi​,x^sub)=α∈{1,21​,41​}∑​L(sgt,s^iα​),i∈{0,1},(5)式中，$\mathcal{L}$为医学图像分割常用的Dice损失加上交叉熵损失，${\mathrm{s}}^{\mathrm{gt}}$为分割基准值，${\hat{\mathrm{s}}}_i^{\alpha }$为$S({\mathbf{x}}_i,{\hat{\mathbf{x}}}^{\mathrm{sub}})$在特定尺度$\alpha$下的网络输出，如果需要则进行上采样以匹配${\mathrm{s}}^{\mathrm{gt}}$的大小。具体来说，使用1 x 1 x 1的卷积(用$f_d^{\alpha }$表示)和三线性上采样来产生$\alpha =1/2$和$\alpha =1/4$(当$\alpha =1$时就只剩下一个分割头了)的中间预测。因此，第二阶段借助自蒸馏一致性损失，针对模态缺失的多模态分割目标任务，对网络进行了调优。至于推断，我们简单的用${\hat{\mathbf{x}}}^{\mathrm{sub}}$代替缺失的模态(如果有的话)，将得到的图像输入到训练好的模型中，通过$f_s$得到分割。

实验与结果

数据集和评估指标

数据集： BraTS 2018 和 BraTS 2020
BraTS数据集包括4个序列的多对比MRI检查：T1、T1c、T2和FLAIR。
将四个肿瘤内结构(水中、强化瘤核、坏死及无强化瘤核)被分为三个肿瘤区域进行评估：

1. 整个肿瘤，包括所有肿瘤组织
2. 肿瘤核心 ：由增强肿瘤、坏死和非增强肿瘤核心组成
3. 增强肿瘤

(训练集 : 验证集 : 测试集)：199：29：57 和 219：50：100 的比例进行拆分
使用Dice相似性系数(DSC)和HD95进行性能量化。

实现：

使用Pytorch框架进行实验
使用两块NVIDIA RTX2080Ti GPU进行训练
批量大小为2volume，即两个受试者的两volume用于预训练，一个受试者的两个随机缺失模态实例化用于微调。使用Adam优化器，初始学习率为0.0003，使用余弦衰减调度器，进行预训练(600epoch)和微调(300epoch)将${\mathbf{x}}^{\mathbf{sub}}$初始化为高斯噪声，进行普通数据增广进行训练。根据BraTS 2018验证拆分实验结果，经验设定权重$\lambda$和掩码比率分别为0.1和0.875

与SOTA比较

Table 1：使用(DSC以mean ± std.)与SOTA方法，包括U-HVED(catch-all)、ACN(dedicated)、SMU-Net(dedicated)和RFNet(catch-all)，在BraTS 2018的测试拆分上进行性能比较。当前模态和缺失模态分别用●和o表示，p<0.05通过Wilcoxon符号秩检验与本文的方法进行比较

Table 2：使用(DSC以mean ± std.)与SOTA方法，包括U-HVED(catch-all)、ACN(dedicated)、SMU-Net(dedicated)和RFNet(catch-all)，在BraTS 2020的测试拆分上进行性能比较。当前模态和缺失模态分别用●和o表示，p<0.05通过Wilcoxon符号秩检验与本文的方法进行比较
在仅使用单个训练模型的情况下，大大优于两个专用方法(ACN和SMU-Net)，相比之下，后两者对所有模态组合一共需要15个模型，这使得M3AE不仅在性能上更加优越，而且在实际训练和部署中都更加高效。

Table 3：消融实验

Table 4：在BraTS 2018在线验证集上进行全模态性能对比