【论文阅读笔记】SMU-Net: Style matching U-Net for brain tumor segmentation with missing modalities

Azad R, Khosravi N, Merhof D. SMU-Net: Style matching U-Net for brain tumor segmentation with missing modalities[C]//International Conference on Medical Imaging with Deep Learning. PMLR, 2022: 48-62.
这篇论文提出了风格匹配U-Net（SMU-Net），它采用了一种包含内容和风格匹配机制的协同训练方法。这种方法涉及将完整模态和缺失模态数据编码到潜在空间，并将其分解为风格和内容表示。风格匹配模块能够自适应地重新校准表示空间，将信息丰富和纹理特征从完整模态路径转移到缺失模态路径。此外，内容模块专注于区分性的语义特征，超越了较少信息的特征。论文通过在BraTS 2018数据集上的评估，展示了SMU-Net在脑肿瘤分割中的有效性。

本文创新点和核心思想如下：

1. 风格匹配U-Net（SMU-Net）架构：这是一种新型的神经网络架构，专门为处理MRI脑肿瘤分割中的缺失模态问题而设计。它通过风格匹配机制，有效地弥补了传统方法在处理缺失模态数据时的不足。
2. 内容与风格匹配机制：SMU-Net引入了一个独特的内容和风格匹配机制。该机制能够从具有全部模态的网络中提取信息特征，并将这些特征有效地转移到缺失某些模态的网络中，从而提高了分割的准确性和可靠性。
3. 共同训练策略：该方法采用共同训练策略，同时训练全模态和缺失模态的网络。这种策略使得网络能够在缺少一部分数据的情况下，仍然能够有效地学习和适应，提高了模型对缺失模态的适应性和鲁棒性。
4. 针对极端缺失模态情况的优化：SMU-Net特别针对极端缺失模态情况（如单一模态输入）进行了优化，显著提升了在这类挑战性情境下的性能。
5. 神经风格转移的应用：文中将神经风格转移的概念应用于医学图像处理，特别是在MRI脑肿瘤分割中。这一点在之前的研究中较少见，为医学图像分析领域带来了新的视角和方法。

网络架构

• Co-training Network：

  设计两种不同的学习路径，1)以完整的模式作为输入的路径，2)输入带有缺失模式的路径。本设计的主要目标是将全模态路径中的信息蒸馏到缺失模态路径，其中协同训练策略鼓励缺失模态网络重构缺失信息。使用以下一些损失函数：

  • 常规分割损失Dice Loss计算分割结果和Ground Truth之间损失：${\mathcal{L}}_{seg}={\mathcal{L}}_{\text{dice }}\left(Y_f^{\prime },Y\right)+{\mathcal{L}}_{\text{dice }}\left(Y_m^{\prime },Y\right)$

  • Jensen-Shanon estimator用于算全模态路径和缺失模态路径的输出(软logits (SL)）之间的差异，即上面两条独立训练路径输出结果之间的差异：

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

  • L1损失：为了进一步包括全局表示匹配，计算全模态路径和缺失模态路径的soft logits之间的L1距离：${\mathcal{L}}_{L1}\left(S{L}_f,S{L}_m\right)={\sum }_{i=1}^c\left|GP\left(S{L}_f\right)-GP\left(S{L}_m\right)\right|$

  • 一致性损失由上面的Jensen-Shanon estimator和L1加权得到:

    ${\mathcal{L}}_{\text{consistency }}={\mathcal{L}}_I+{\mathcal{L}}_{L1}$

• Style-matching Module

  使用该模块克服domain shift，从而在全模态路径和缺失模态路径上恢复丢失的风格信息。对于样式表示，将卷积滤波器响应连接到浅层和深层，而对于内容表示，使用最后的卷积输出。派生的样式表示保留了有价值的纹理信息，而内容特征包含了图像的核心结构和语义特征。这部分使用了三种损失：

  • KL散度:量化两种路径风格分布之间的匹配，将缺失的模态分布对齐到全模态版本，因此，它学会了通过学习后验分布来恢复缺失的信息。${\mathcal{L}}_{\text{style }}^{kl}\left(f{s}_f,f{s}_m\right)=D_{\mathrm{KL}}\left(T\left(z\mid f{s}_f,f{s}_m\right)\Vert P\left(z\mid f{s}_f\right)\right)$

  • 对抗损失：将缺失模态的特征分布与全模态的特征分布对齐，以重新获得缺失的信息。${\mathcal{L}}_{\text{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)$

  • MSE损失:最大化两个表示的纹理之间的相关性,计算Gram矩阵元素的均方误差，这些元素是为每条路径创建的。${\mathcal{L}}_{\text{style }}^{\text{texture }}\left(f{s}_f,f{s}_m\right)={\sum }_l^L{w}_l\frac{1}{4C_l^2{N}_l^2}\sum {\left(f{s}_f^{l,j}-f{s}_m^{l,j}\right)}^2$

​

• Content Module

  目的是捕捉模式之间共享的结构和上下文信息，从全模态中提取区分性强的信息到缺失模态。使用MSE损失：

  ${\mathcal{L}}_{\text{content }}\left(\mathbf{f}{\mathbf{c}}_{\mathbf{f}},\mathbf{f}{\mathbf{c}}_{\mathbf{m}}\right)=\frac{1}{2}{\sum }_j{\left(f{c}_{f_j}-f{c}_{m_j}\right)}^2$

• 总的目标函数

  ${\mathcal{L}}_{\text{joint }}={\lambda }_1{\mathcal{L}}_{\text{segmentation }}+{\lambda }_2{\mathcal{L}}_{\text{consistency }}+{\lambda }_3{\mathcal{L}}_{\text{style }}+{\lambda }_4{\mathcal{L}}_{\text{content }}$

数据集和硬件

3090单卡，BraTS 2018

性能比较