论文阅读笔记——SMU-Net:面向缺失模态脑肿瘤分割的样式匹配U-Net

论文地址：https://arxiv.org/abs/2204.02961v1

前置知识：

脑胶质瘤：https://baike.baidu.com/item/%E8%84%91%E8%83%B6%E8%B4%A8%E7%98%A4/7242862
互信息:https://zhuanlan.zhihu.com/p/240676850
Gram矩阵:https://zhuanlan.zhihu.com/p/187345192

摘要：

背景：

绝大多数脑肿瘤都可以通过磁共振成像进行唯一的鉴别。
多模态MRI的好处： 每一种模态都提供人脑软组织的不同信息，整合所有这些信息对胶质瘤准确分割提供全面的数据。
MRI现存不足：容易产生伪影，这可能导致一个或多个MRI模态的缺失

方法：

本文提出了一种基于样式匹配U-Net (SMU-Net) 的MRI图像脑肿瘤分割方法。采用的协同训练方法利用内容和样式匹配机制，将全模态网络中的信息特征提取到缺失的模态网络中。
实现方法是将全模态和缺失模态数据都编码到一个潜在空间中，然后将潜在表示空间分解为样式和内容表示。
样式匹配模块：通过学习一个匹配函数来自适应性地校准表示空间，将信息和纹理特征从全模态路径转移到缺失模态路径
内容模块：surpasses信息量较少地特征，并基于判别性语义特征重新校准表示空间。

引言：

现有的解决自动分割问题的技术：

1. Pixel to pixel network (像素到像素网络) ：是第一个被引入解决自动分割问题的CNN网络之一。消除了全连接层，利用全卷积结构来保持数据的空间维度。
2. 引入U-Net模型进一步扩展的Pixel to pixel network：
3. 之后通过对U-Net模型进行注意力机制、频率重标定器、非线性跳跃连接、辅助模块等扩展提高该模型的性能。
   缺点：所有模态在训练和推理时间上的可用性都是限制性的。这一事实限制了他们在模态确实情况下的一般性。

现有解决缺失模态问题的方法：

1. 使用比分类器适应性更强的合成方法来推断测试时的缺失数据，可以通过为分类器提供数据的转换和扩充训练集来增强多模态图像分类。
2. 使用随机森林作为简单的分类器，利用合成数据增强分割
3. 利用非线性回归的随机森林合成缺失模态的方法，可以合成T2和FLAIR。
   实验表明：对数据进行合成或填充或许不会带来任何改善。此外这些方法应用于输入层，通常受限于representative能力来重建缺失的模态。因此使用深度学习来克服这个限制。

现有在分割模型中加入缺失模态的方法：

1. U-HeMIS：U-HeMIS是最早提出在分割模型中加入缺失模态的方法之一。从每个模态到嵌入空间的映射函数(正则编码器模型)被学习。然后利用均值和方差对模态间共享的公共表示空间进行建模。计算均值和方差并不一定会重建缺失的信息
2. HVED：进一步建立共同的潜在空间表示，利用独立的编码器从每个模态中提取一阶矩特征，然后通过对编码特征上的高斯分布进行建模，形成一个通用的潜在表示空间，该潜在表示空间为所有的模态提供了一个共同的表示。缺点是无法对模态独立特点进行建模，且在一个以上的模态缺失时表现不佳
3. 其他的方法：如利用两个模态共享的显著相关型，在没有完美重建缺失序列的情况下执行分割任务。类似地，一种潜在相关表示学习方法，用于解决未被视为对所有模态都完全满足的缺失模态问题。后来的方法提出分层对抗知识蒸馏网络。缺点是通常无法从缺失模态中重建纹理(样式)信息

本文提出的SMU-Net方法的贡献：

1. 使用样式-内容匹配模块对公共表示空间进行建模
2. 协同训练的方法从全模态中提取信息特征到缺失的模态网络中
3. 具有SOTA结果的端到端训练策略
4. 公开的实现源代码：https://github.com/rezazad68/smunet

提出的方法

本文提出的SMU-Net受到协同训练方法和样式迁移网络的启发，是一种用于缺失模态的脑肿瘤分割的样式匹配U-Net，该网络考虑了协同训练策略的强度，将全模态模型中有用的信息提取到缺失的模态网络中。在作者的设计中，将表示空间分解为内容和样式表示，然后利用匹配模块来重构缺失的信息，样式匹配模块学习一个非线性的参数化函数来匹配完整和缺失模态网络之间由浅到深的特征。此外，通过对互信息进行建模，内容模块surpasses信息较少的信息(得到更加丰富的特征信息)，并基于共同特征重新校准了表示空间。

Figure 1：SMU-Net采用(1)样式匹配机制，利用全模态网络重构缺失信息； (2)内容匹配模块，捕捉信息丰富的特征，超越判别性较差特征

协同训练网络

以前的技术已经引入了使用协同训练过程来训练深度学习模型。每种MRI模态都提供了有用的信息，并揭示了人脑软组织的不同特征。因此将所有获得的特征结合起来，为我们深入的研究大脑的感兴趣区域提供全面的信息。为了建立一个对缺失模态具有弹性的训练过程，我们提出的网络利用了一个耦合的学习过程。在我们的设计中，有两条独立的学习路径， (1)以完整的模态作为输入的路径(2)以缺失的模态作为输入的路径，本设计的主要目标是从全模态路径中提取信息特征到缺失模态路径，其中协同训练策略鼓励缺失模态网络重构缺失信息。按照分割任务的符号标准，定义全模态数据为 X f = { ∀ X i , i ∈ M } X_f=\{\forall X_i,i\in M\} Xf​={∀Xi​,i∈M}，缺失模态数据为 X m = { ∀ X i , i ∈ M ′ } X_m=\{\forall X_i,i\in M'\} Xm​={∀Xi​,i∈M′}，相应的真值掩码为 Y ∈ { 0 , 1 , 2 , 4 } { H × W × D } Y\in\{0,1,2,4\}^{\{H\times W\times D\}} Y∈{0,1,2,4}{H×W×D}，在文章的设定中， X i { H × W × D } X_i^{\{H\times W\times D\}} Xi{H×W×D}​表示空间维数为(H、W、D)的3D模态特征，M表示模态的集合(T1、T1c、T2 and Flair)其中$M^{\prime }\in M$，协同训练策略的目的是将全模态模型中的语义和信息特征提取到一个缺失的缺失模型中。因此，作者定义了两个并行分割模型(U-Net)，分别使用$X_f$和$X_m$独立训练和同时训练。每个网络使用预测掩码$Y^{\prime }$和真实掩码$Y$之间的Dice损失来学习分割图，${\mathcal{L}}_{seg}={\mathcal{L}}_{dice}(Y_f^{\prime },Y)+{\mathcal{L}}_{dice}(Y_m^{\prime },Y)$，接下来，作者在局部和全局两个层面上定义了跨模型的一致性损失，以确保特征从全模态蒸馏到缺失模态。为此，采用Jensen-Shanon估计量来计算完整和缺失模态路径的输出(soft Logits(SL):应用Softmax操作前的网络输出)之间的互信息。

${\mathcal{L}}_I(S{L}_f,S{L}_m)={\mathbb{E}}_{{\mathbb{P}}_{S{L}_{f}S{L}_m}}\left[-\mathrm{sp}\left(C{T}_{\varphi }(s{l}_f,s{l}_m)\right)\right]-{\mathbb{E}}_{{\mathbb{P}}_{S{L}_f}\otimes {\mathbb{P}}_{S{L}_m}}[\mathrm{sp}\left(C{T}_{\varphi }(s{l}_f,s{l}_m)\right],(1)$

式中：$\mathrm{sp}(z)=\log (1+e^z)$，且$C{T}_{\varphi }$表示参数为$\varphi$的协同训练网络。这种相互最大化函数保证了两条路径之间的分布匹配。为了进一步包含全局表示匹配，我们计算了完整和缺失模态路径的soft logits之间的L1距离：

$\begin{array}{rl}{\mathcal{L}}_{L1}(S{L}_f,S{L}_m)& =\sum _{i=1}^c|GP(S{L}_f)-GP(S{L}_m)|,(2)\end{array}$

其中GP为全局池化操作，c为模态类别数。我们的一致性损失由这两个指标构成:${\mathcal{L}}_{consistency}={\mathcal{L}}_I+{\mathcal{L}}_{L1}$。所提出的跨模态一致性损失评估了缺失模态路径和完整模态路径产生的输出之间的分布匹配，并通过考虑它们之间的相关性来提高输出一致性，并最终促进两条路径产生相同的分布。除此之外，局部和全局的层次分布保证了两条路径之间的最小分布距离，并通过提取信息特征，鼓励全模态路径监督缺失的模态路径。为了进一步确保在网络的浅层部分和瓶颈处进行特征蒸馏，我们提出了样式和内容匹配模块，其中样式匹配功能旨在使用全模态模块来重建缺失的信息，而内容匹配模块试图学习具有判别性的表示，并surpasses信息量较少的表示。

样式匹配模块

每个MRI模态是使用不同的属性和设置组合获得的，并且每个模态都暴露了人脑不同的解刨结构。在诸如纹理、对比度、饱和度等样式信息方面，获得的MRI模态可能存在差异，而这种样式差异加剧了跨完整模态和缺失模态路径的域偏移问题。因此，我们提出了一个样式匹配模块，该模块可以克服域偏移，从而在全模态和缺失模态路径上恢复丢失的样式信息。为此，我们将输入图像的特征表示分解为样式和内容表示。对于样式特征表示，我们将浅层和深层的卷积filter相应的concatenate起来，而对于内容表示，我们使用最后一个卷积的输出。导出的样式表示保留了有价值的纹理信息，而内容特征则包含了图像的核心结构和语义特征。为了执行匹配机制，我们定义了3种不同的样式匹配模块。

Figure 2：三种样式匹配模块:( a )分布匹配( b )对抗学习( c )纹理匹配

( a )分布匹配

在第一种分布匹配的设计中，从缺失模态路径和全模态路径中提取样式特征到两个对应的连续分布空间中，然后努力最小化KL散度来度量两个分布之间的匹配(图2a)

${\mathcal{L}}_{\mathrm{style}}^{kl}(f{s}_f,f{s}_m)=D_{\mathrm{KL}}(T(z\mid f{s}_f,f{s}_m)||P(z\mid f{s}_f)),(3)$

式中：$z=(\theta ,\sigma )$为连续分布的参数。通过最小化KL损失，将缺失模态的分布对齐到全模态的分布，从而通过学习后验分布来学习恢复缺失信息。

( b )对抗学习

如图2(b)所示，我们提出的第二种样式匹配机制是一个对抗学习模块，由于精确分割高度依赖于高层表示，我们想要尽力去阻止潜在空间中的特征分布错位，为此，我们使用一个以$\theta$为参数的判别器网络D来判别完整和缺失的模态样式。通过执行对抗损失，我们真正地将缺失模态的特征分布与全模态对齐，以重新捕获缺失的信息。

${\mathcal{L}}_{style}^{adv}\left(f{s}_f,f{s}_m\right)=\log \left(1-D_{\vartheta }\left(f{s}_f\right)\right)+\log \left(D_{\vartheta }\left(f{s}_m\right)\right),(4)$

( c )纹理匹配

在第三种方法中，我们通过最大化缺失模态样式表示和全模态样式表示的纹理之间的相关性来对样式匹配模块进行建模。为此，我们为每条路径创建了Gram矩阵元素的均方误差。所有层的加权MSE(wl) 损失计算如下：

${\mathcal{L}}_{\mathrm{style}}^{texture}\left(f{s}_f,f{s}_m\right)={\sum }_{l=0}^L{w}_l\frac{1}{4C_l^2{N}_l^2}{\sum }_j{\left(f{s}_f^{l,j}-f{s}_m^{l,j}\right)}^2,(5)$

其中$N_l$表示第$l$层的空间维度，$C_l$表示第$l$层的通道数，该损失函数被一致的与其他样式模块一起使用，使得缺失模态样式近似全模态样式。

内容模块

本文设计的内容模块旨在捕获模态之间共享的结构和上下文信息。为了从全模态路径中提取这种区分性信息到缺失版本，我们部署了全模态内容与缺失模态内容表示之间的MSE损失:

内容匹配模块引导缺失模态捕获判别特征，同时surpasses信息量较少的特征。当样式和内容特征被重新校准后，我们将校准过的样式特征与编码器提取的特征进行重新组合，生成新的表示，作为解码器模块的输入。我们简单的使用一个U-Net解码器来生成分割图。

Joint Objective

在训练过程中优化的整体损失函数由四项组成:

${\mathcal{L}}_{\mathrm{joint}}={\lambda }_1{\mathcal{L}}_{\mathrm{segmentation}}+{\lambda }_2{\mathcal{L}}_{\mathrm{consistency}}+{\lambda }_3{\mathcal{L}}_{\mathrm{style}}+{\lambda }_4{\mathcal{L}}_{\mathrm{content}}.(7)$

使用的$\lambda$系数来衡量每个损失对总体损失值的贡献，我们使用Adam优化以端到端的方式训练模型，学习率为$10^{-4}$，批量大小为1，训练250个epoch。该代码是在pytorch库中开发的，并在单个RTX3090GPU上执行的。

实验结果

数据集和评估指标

使用BraTS 2018数据集，该数据集包括285例患者的MRI图像，其中高级别胶质瘤210例，低级别胶质瘤75例，每例均包含上述四种模态(T1、T1c、T2、FLAIR)，BraTS 2018数据集中包含坏死、水肿、非增强性肿瘤和增强肿瘤4中组织类别的标注。尽管提供了4中不同的肿瘤标签，但他们可以很好的分为3各自区域进行评估:全肿瘤(WT)、核心肿瘤(CT)和增强肿瘤(ET)。输入大小为160 x 192 x128。对于评估，我们使用了Dice相似性系数(DSC)。在github上提供了实现代码(https://github.com/rezazad68/smunet)

比较

Table 1：使用Dice度量在BraTS 2018数据集上对所提出的SMU-Net进行性能比较，值得注意，在此次比较过程中使用的是对抗学习样式匹配模块(图中Modalities部分的●表示模态存在，o表示模态缺失)
表1展示了SMU-Net在BraTS 2018数据集上的性能。首先，为了证明SMU-Net的有效性，我们将其与著名的U-HeMIS、HVED以及最劲的ACN架构进行比较。可以观察到，当多个模态缺失时，U-HeMIS或HVED的性能大幅下降，而SMU-Net取得了更好的效果，特别是在单模态情况下，如表1所示。这一观察揭示了我们的样式-内容匹配模块对于捕获缺失信息和重新校准缺失的模态路径的重要性。此外，我们还可以观察到，在单模态场景中，本文提出的方法比SOTA方法有较大的提升(20%)，这表明了我们的协同训练策略的有效性，此外，在我们类似设置中，我们的方法比竞争对手ACN方法取得了更好的结果。为了突出我们方法的有效性，我们在图3中给出了我们的模型对单一模态案例的定性结果，可以观察到，对于每个单一模态，该方法都产生了可接受的分割结果，而这些结果在没有重建缺失信息的情况下从输入图像中是不可行的

图3：使用单一模态作为输入的SMU-Net在BraTS 2018上的视觉分割结果。绿色区域表示WT；红色为CT和蓝色为ET

消融实验