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在脑肿瘤分割中的有效性。
本文创新点和核心思想如下:
- 风格匹配U-Net(SMU-Net)架构:这是一种新型的神经网络架构,专门为处理MRI脑肿瘤分割中的缺失模态问题而设计。它通过风格匹配机制,有效地弥补了传统方法在处理缺失模态数据时的不足。
- 内容与风格匹配机制:SMU-Net引入了一个独特的内容和风格匹配机制。该机制能够从具有全部模态的网络中提取信息特征,并将这些特征有效地转移到缺失某些模态的网络中,从而提高了分割的准确性和可靠性。
- 共同训练策略:该方法采用共同训练策略,同时训练全模态和缺失模态的网络。这种策略使得网络能够在缺少一部分数据的情况下,仍然能够有效地学习和适应,提高了模型对缺失模态的适应性和鲁棒性。
- 针对极端缺失模态情况的优化:SMU-Net特别针对极端缺失模态情况(如单一模态输入)进行了优化,显著提升了在这类挑战性情境下的性能。
- 神经风格转移的应用:文中将神经风格转移的概念应用于医学图像处理,特别是在MRI脑肿瘤分割中。这一点在之前的研究中较少见,为医学图像分析领域带来了新的视角和方法。
网络架构
-
Co-training Network:
设计两种不同的学习路径,1)以完整的模式作为输入的路径,2)输入带有缺失模式的路径。本设计的主要目标是将全模态路径中的信息蒸馏到缺失模态路径,其中协同训练策略鼓励缺失模态网络重构缺失信息。使用以下一些损失函数:
-
常规分割损失Dice Loss计算分割结果和Ground Truth之间损失: L s e g = L dice ( Y f ′ , Y ) + L dice ( Y m ′ , Y ) \mathcal{L}_{s e g}=\mathcal{L}_{\text {dice }}\left(Y_{f}^{\prime}, Y\right)+\mathcal{L}_{\text {dice }}\left(Y_{m}^{\prime}, Y\right) Lseg=Ldice (Yf′,Y)+Ldice (Ym′,Y)
-
Jensen-Shanon estimator用于算全模态路径和缺失模态路径的输出(软logits (SL))之间的差异,即上面两条独立训练路径输出结果之间的差异:
L I ( S L f , S L m ) = E P S L f S L m [ − sp ( C T ϕ ( s l f , s l m ) ) ] − E P S L f ⊗ P S L m [ sp ( C T ϕ ( s l f , s l m ) ] \mathcal{L}_{I}\left(S L_{f}, S L_{m}\right)=\mathbb{E}_{\mathbb{P}_{S L_{f} S L_{m}}}\left[-\operatorname{sp}\left(C T_{\phi}\left(s l_{f}, s l_{m}\right)\right)\right]-\mathbb{E}_{\mathbb{P}_{S L_{f}} \otimes \mathbb{P}_{S L_{m}}}\left[\operatorname{sp}\left(C T_{\phi}\left(s l_{f}, s l_{m}\right)\right]\right. LI(SLf,SLm)=EPSLfSLm[−sp(CTϕ(slf,slm))]−EPSLf⊗PSLm[sp(CTϕ(slf,slm)]
-
L1损失:为了进一步包括全局表示匹配,计算全模态路径和缺失模态路径的soft logits之间的L1距离: L L 1 ( S L f , S L m ) = ∑ i = 1 c ∣ G P ( S L f ) − G P ( S L m ) ∣ \mathcal{L}_{L 1}\left(S L_{f}, S L_{m}\right)=\sum_{i=1}^{c}\left|G P\left(S L_{f}\right)-G P\left(S L_{m}\right)\right| LL1(SLf,SLm)=∑i=1c∣GP(SLf)−GP(SLm)∣
-
一致性损失由上面的Jensen-Shanon estimator和L1加权得到:
L consistency = L I + L L 1 \mathcal{L}_{\text {consistency }}=\mathcal{L}_{I}+\mathcal{L}_{L 1} Lconsistency =LI+LL1
-
-
Style-matching Module
使用该模块克服domain shift,从而在全模态路径和缺失模态路径上恢复丢失的风格信息。对于样式表示,将卷积滤波器响应连接到浅层和深层,而对于内容表示,使用最后的卷积输出。派生的样式表示保留了有价值的纹理信息,而内容特征包含了图像的核心结构和语义特征。这部分使用了三种损失:
-
KL散度:量化两种路径风格分布之间的匹配,将缺失的模态分布对齐到全模态版本,因此,它学会了通过学习后验分布来恢复缺失的信息。 L style k l ( f s f , f s m ) = D K L ( T ( z ∣ f s f , f s m ) ∥ P ( z ∣ f s f ) ) \mathcal{L}_{\text {style }}^{k l}\left(f s_{f}, f s_{m}\right)=D_{\mathrm{KL}}\left(T\left(z \mid f s_{f}, f s_{m}\right) \| P\left(z \mid f s_{f}\right)\right) Lstyle kl(fsf,fsm)=DKL(T(z∣fsf,fsm)∥P(z∣fsf))
-
对抗损失:将缺失模态的特征分布与全模态的特征分布对齐,以重新获得缺失的信息。 L style a d v ( f s f , f s m ) = log ( 1 − D ϑ ( f s f ) ) + log ( D ϑ ( f s m ) ) \mathcal{L}_{\text {style }}^{a d v}\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) Lstyle adv(fsf,fsm)=log(1−Dϑ(fsf))+log(Dϑ(fsm))
-
MSE损失:最大化两个表示的纹理之间的相关性,计算Gram矩阵元素的均方误差,这些元素是为每条路径创建的。 L style texture ( f s f , f s m ) = ∑ l L w l 1 4 C l 2 N l 2 ∑ ( f s f l , j − f s m l , j ) 2 \mathcal{L}_{\text {style }}^{\text {texture }}\left(f s_{f}, f s_{m}\right)=\sum_{l}^{L} w_{l} \frac{1}{4 C_{l}^{2} N_{l}^{2}} \sum\left(f s_{f}^{l, j}-f s_{m}^{l, j}\right)^{2} Lstyle texture (fsf,fsm)=∑lLwl4Cl2Nl21∑(fsfl,j−fsml,j)2
-
-
Content Module
目的是捕捉模式之间共享的结构和上下文信息,从全模态中提取区分性强的信息到缺失模态。使用MSE损失:
L content ( f c f , f c m ) = 1 2 ∑ j ( f c f j − f c m j ) 2 \mathcal{L}_{\text {content }}\left(\boldsymbol{f} \boldsymbol{c}_{\boldsymbol{f}}, \boldsymbol{f} \boldsymbol{c}_{\boldsymbol{m}}\right)=\frac{1}{2} \sum_{j}\left(f c_{f_{j}}-f c_{m_{j}}\right)^{2} Lcontent (fcf,fcm)=21∑j(fcfj−fcmj)2
-
总的目标函数
L joint = λ 1 L segmentation + λ 2 L consistency + λ 3 L style + λ 4 L content \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 }} Ljoint =λ1Lsegmentation +λ2Lconsistency +λ3Lstyle +λ4Lcontent
数据集和硬件
3090单卡,BraTS 2018
性能比较