MoME: 癌症生存预测的多模态混合专家.

0. TL; DR

整合Whole Slide Images (WSIs)(全切片图像)和genomic data(基因组数据)进行生存分析,作者提出了两个核心创新:

  1. 偏向渐进式编码范式 (Biased Progressive Encoding, BPE):该范式在编码一种模态时,将另一种模态作为参考,从而实现编码和融合的同时进行。通过多次交替迭代,这种方法能够逐步减少跨模态的差异,促进互补信息的交互,实现更深层次的融合。
  2. 多模态专家混合层 (Mixture of Multimodal Experts, MoME):考虑到不同患者的关键生物标志物可能存在于不同模态中,作者提出了MoME层。该层包含多个专门的“专家”,能够动态地选择最适合当前样本的专家进行处理。这些专家以不同程度整合参考模态的信息,使得模型能够在编码过程中实现对不同模态的平衡或偏向性关注。

作者在包括TCGA-BLCA, TCGA-UCECTCGA-LUAD在内的多个数据集上进行了广泛实验。结果表明,与现有state-of-the-art方法相比,该方法取得了优越的性能。

1. 背景介绍

利用WSI和基因组数据进行生存分析,对于泛癌种预后至关重要。这项任务的关键在于如何有效利用两种模态的信息,例如,检测图像-组学(image-omic)生物标志物,以及探索组织病理学图像中肿瘤微环境与基因组数据中共表达之间的交互。

然而,这项任务面临两大挑战:

  1. 显著的异质性:组织病理学图像和基因组数据在数据性质和预处理方法上存在巨大差异。
  2. 复杂的交互:两种模态都包含丰富的信息,但只有一小部分是相互关联且对生存预测有用的。

现有的方法,如基于cross-modality attention(跨模态注意力)或co-attention的方法,通常是在对每个模态进行独立编码后,才进行一次性的特征融合。作者认为,考虑到任务的复杂性和模态间的巨大差异,这种“浅层”的融合策略是不够的。

因此,作者提出了一个Biased Progressive Encoding (BPE) 范式,它在编码一个模态时,就将另一个模态作为参考,从而实现编码与融合的同步进行,并允许更深层次的交互。

此外,对于不同的患者,决定其生存预后的关键特征可能存在于不同的模态中。这要求模型具有灵活性,能够根据具体情况自适应地调整对不同模态的关注度。为此,作者进一步提出了一个Mixture of Multimodal Experts (MoME) 层,它能够动态地选择最合适的“专家”来处理数据。

2. 方法介绍

2.1 问题定义

遵循Attention-Based Multiple Instance Learning (AB-MIL) 框架,将WSI分割成一个patch包,并使用预训练网络提取特征。WSI patch特征表示为$P \in \mathbb{R}^{n_p \times d}$。

根据生物学功能分为6个类别,并通过一个全连接层得到基因组特征$G \in \mathbb{R}^{n_g \times d}$。

问题定义为:给定输入对$(P, G)$,估计风险函数$h(t)$,并最终得到一个序数风险值。

2.2 偏向渐进式编码范式 (Biased Progressive Encoding, BPE)

BPE的核心思想是,在编码一种模态时,将另一种模态作为“参考”,以提取更相关的信息,并逐步减少模态间的差异。

编码过程是交替进行的。假设$F_1$和$F_2$分别代表两种模态的特征,在第$i$次迭代中:

  1. 首先,以$F_2^{(i)}$为参考,对$F_1^{(i)}$进行编码,得到$F_1^{(i+1)}$。
  2. 然后,以更新后的$F_1^{(i+1)}$为参考,对$F_2^{(i)}$进行编码,得到$F_2^{(i+1)}$。
\[F_1^{(i+1)} = \text{BPE}_{2i}(F_1^{(i)}, F_2^{(i)})\\ F_2^{(i+1)} = \text{BPE}_{2i+1}(F_2^{(i)}, F_1^{(i+1)})\]

其中,$\text{BPE}(\cdot, \cdot)$函数由作者提出的MoME层实现。

2.3 多模态专家混合层 (Mixture of Multimodal Experts, MoME)

MoME层源于传统的Mixture of Experts (MoE),但作者对其进行了创新,使其在样本和层级上进行操作,而非传统的token级。

MoME层包含两部分:一个门控网络和一个多模态专家池。

2.3.1 门控网络 (Gating Network)

门控网络是一个轻量级的网络,用于根据输入的两种模态特征,决定选择哪个专家来处理当前的数据。

\[\text{logits} = W \cdot (\text{mean}(\text{GELU}(\text{RMSNorm}(W_1 F_1)))) + W \cdot (\text{mean}(\text{GELU}(\text{RMSNorm}(W_2 F_2))))\]

与传统的MoE使用加权和不同,MoME强制只选择一个最合适的专家,这有助于降低计算成本。

2.3.2 多模态专家池 (Multimodal Expert Pool)

作者设计了四种不同的专家,它们以不同程度整合参考模态的信息,从而实现对不同模态的偏向性关注。

  1. TransFusion (TF):该专家基于self-attention,通过将两种模态的特征拼接后送入一个self-attention模块,实现两种模态之间信息的完全交换,从而最大化地利用参考模态。 $\text{TF}(F_1, F_2) = \text{SA}([F_1, F_2])[:n_1, :]$
  2. Bottleneck TransFusion (BTF):该专家引入了bottleneck features(瓶颈特征)作为桥梁,避免了两种模态特征之间的直接交互。参考模态的利用程度低于TF。 $\text{BTF}(F_1, F_2) = \text{SA}_1(F_1, \text{SA}_2(B, F_2)[:n_b, :])[:n_1, :]$
  3. SNNFusion (SF):该专家主要为基因组数据设计,它包含两个Self-Normalizing Networks (SNNs)。对参考模态的利用程度更低。 $\text{SF}(F_1, F_2) = \text{SNN}_1(\text{RMSNorm}(F_1)) + \text{mean}(\text{SNN}_2(\text{RMSNorm}(F_2)))$
  4. DropF2Fusion (DF):该专家完全忽略参考模态$F_2$,仅对$F_1$进行处理。当单一模态的信息已经足够准确时,该专家非常有用。 $\text{DF}(F_1, F_2) = F_1$

作者通过数学公式推导指出,Co-Attention (CA) 实际上是Self-Attention (SA) 的一个子集。在对拼接的多模态特征进行SA时,其注意力矩阵中包含了模态内自注意力(SA(F1), SA(F2))和模态间协同注意力(CA(F1,F2), CA(F2,F1))的所有情况。因此,作者认为CASA的一个次优替代品,并选择基于SA来设计专家。

3. 实验分析

3.1 实验设置

使用了来自TCGA的三个数据集:BLCA(膀胱癌), UCEC(子宫内膜癌), LUAD(肺腺癌)。

与一系列单模态和多模态的SOTA方法进行比较,包括SNN, TransMIL, MCAT, Porpoise, MOTCAT, CMTA等。评估指标使用Concordance index (C-index)

3.2 比较结果

MoME在所有三个数据集上都一致地取得了最佳性能。与之前表现最好的方法相比,MoMEBLCA, UCEC, LUAD上的C-index分别提升了0.4%, 1.4%, 1.5%,总体性能提升了1.4%。

一个有趣的观察是,尽管结构简单,但仅基于基因组数据的方法(如SNN)的性能优于基于WSI的方法,这突显了基因组数据在生存分析中的重要性。多模态方法一致地优于单模态方法,证明了多模态方法的有效性和整合多模态的必要性。

MoME的优越性能表明,其渐进式编码和动态专家选择机制能够有效处理异构数据并捕捉复杂的模态间交互,适用于通用的生存分析场景。

3.3 消融研究

作者通过逐一移除不同的专家来评估它们各自的有效性。在LUAD上,包含所有专家的完整MoME取得了最佳性能。在UCEC上,移除了SNNFusion专家的MoME表现最好。这两个数据集上最优的专家组合不同,这表明针对不同数据集,模型需要关注的重点(特定模态或模态间交互)是不同的。这进一步证明了MoME能够动态选择不同专家来适应不同样本的必要性。

当瓶颈特征数量为2时,模型在两个数据集上取得了最佳的平衡和总体性能。当瓶颈特征数量为8时,模型在UCEC上表现最好,但在LUAD上表现最差。这也从侧面反映了不同数据集对模型组件的偏好不同。