在线掩码自编码器混合图卷积网络用于多模态癌症生存预测.

0. TL; DR

本文提出了一种名为HGCN（Hybrid Graph Convolutional Network）的网络，该网络配备了一个online masked autoencoder（在线掩码自编码器）范式，用于鲁棒的多模态癌症生存预测。

将患者的多模态数据建模为灵活且可解释的多模态图，HGCN整合了graph convolutional networks (GCNs) 和hypergraph convolutional network (HCN) 的优势，通过节点消息传递和超边混合机制来促进模态内和模态间的交互。为了应对临床场景中缺失模态的问题，作者将一个online masked autoencoder范式融入HGCN。该范式能有效捕捉模态间的内在依赖关系，并在模型推理时无缝地生成缺失的hyperedges（超边）。

作者在来自TCGA的六个癌症队列上进行了广泛的实验和分析。结果表明，无论是在完整模态还是缺失模态的设置下，作者的方法都显著优于state-of-the-arts方法。

1. 背景介绍

预测癌症患者的生存风险对于临床治疗和监测决策至关重要。传统的生存分析模型（如Cox模型）多依赖于人口统计学和临床协变量，但这些方法未能全面地表征患者的疾病病理和遗传因素，从而限制了生存预测的效力。

近年来，随着医学影像和基因组学技术的发展，生存预测研究进入了一个新的多模态时代。许多基于特征的融合技术被提出来融合不同模态的特征表示。然而，大多数现有方法存在以下问题：

尽管已有方法（如零填充、因子分解表示、自编码器）被提出来解决数据不完整问题，但由于医疗数据的维度差异巨大且归因困难，将这些方法直接应用于多模态癌症生存预测是相当困难的。

因此，作者提出了HGCN，一个配备了online masked autoencoder的混合图卷积网络，旨在有效利用多模态数据的互补信息，并鲁棒地应对模态缺失的场景。

2. HGCN 方法

该框架的核心思想是将患者的多模态数据建模为图结构，通过HGCN促进模态内和模态间的交互，并通过一个在线的masked autoencoder来处理模态缺失的情况。

2.1 多模态图的构建 (Construction of Multimodal Graphs)

作者将一个患者的多模态数据（病理学、临床、基因组学）分别建模为灵活的图结构。

这种图表示方法具有更灵活、更能挖掘上下文、能处理缺失特征以及更具可解释性等优点。

2.2 混合图神经网络 (Hybrid Graph Neural Network, HGCN)

HGCN包含两个核心部分：用于模态内交互的GCN层和用于模态间交互的超图卷积层。

2.2.1 使用GCN层进行模态内交互

挖掘每个模态内部的特有信息。为每个模态的图表示（$G_p, G_c, G_g$）分别应用一个独立的GCN层（作者采用了GraphSAGE），以实现节点间的消息传递。

\[G_m^{intra} = \sigma(\text{GCN}(G_m^{in}))\]

其中，$m \in {p, c, g}$，$G_m^{intra}$是经过模态内交互后生成的图表示。

2.2.2 使用超边混合进行模态间交互

在GCN层之后，作者将每个模态的所有节点视为被一条hyperedge（超边）连接。通过对一个模态的所有节点特征进行attention pooling，可以提取出代表该模态高阶信息的hyperedge表示$H_m$。

\[H_m = \sum_{n=1}^N \text{Softmax}(\text{MLP}_1(V_m^n)) \odot V_m^n\]

为了促进模态间的交互，作者设计了一个hyperedge mix模块。该模块通过两个MLP层，分别在token-wise和channel-wise维度上对所有模态的hyperedges $H \in \mathbb{R}^{B \times M \times C}$ 进行混合。

\[H^{inter} := H^T + \text{MLP}_2((\text{LayerNorm}(H))^T)\\ H^{inter} := (H^{inter})^T + \text{MLP}_3(\text{LayerNorm}((H^{inter})^T))\]

最后，将经过混合的、包含了多模态信息的hyperedge $H_m^{inter}$广播加回到相应的模态内图表示$G_m^{intra}$上，得到最终的输出多模态图$G_m^{out}$。

\[G_m^{out} = G_m^{intra} \oplus H_m^{inter}\]

2.3 在线掩码自编码器 (Online Masked Autoencoder, Online MAE)

为了在模型推理时处理模态缺失的问题，作者设计了一个online MAE模块。

MAE是一种自监督学习方法，通过从可见部分预测被掩码的部分来学习有意义的表示。作者将其思想引入到多模态学习中，用于重建缺失模态的hyperedge。

与将MAE作为预训练任务不同，作者将其与HGCN的生存预测任务进行端到端的联合训练。在训练过程中，随机地掩码掉一部分模态（如基因组学）对应的hyperedge。将部分可见的hyperedges送入一个基于Transformer的非对称autoencoder中。编码器只处理可见的hyperedges，而解码器则在可见hyperedges的编码和代表缺失位置的mask tokens上进行重建。重建损失$L_{MAE}$仅在被掩码的hyperedges上计算。

在推理时，如果遇到有模态缺失的样本，训练好的MAE可以直接从可用的模态hyperedges中生成缺失的hyperedges，从而使后续的模态间交互能够正常进行。

2.4 训练与推理策略

总损失是MAE的重建损失和每个可用模态的Cox损失的加权和。

\[L_{Total} = \epsilon_o L_{MAE} + \sum_{m \in M} \epsilon_m L_{Cox}^m\]

在推理时，对于一个给定的样本，首先计算其所有可用模态的图表示和hyperedges。然后，使用MAE补全缺失的hyperedges。接着，通过HGCN进行模态内和模态间交互。最后，对所有可用模态的预测结果进行平均，得到最终的生存预测。

3. 实验分析

3.1 实验设置

数据集使用了来自TCGA的六个癌症队列：KIRC, LIHC, ESCA, LUSC, LUAD, UCEC。

评估指标使用concordance index (CI)、Brier Score (BS) 和Kaplan-Meier (KM) 分析（及Logrank检验）。

与单模态方法（Graph based MIL, Cox model）和四种SOTA多模态方法（GSCNN, MultiSurv, Metric learning, Outer product, MCAT）进行比较。

3.2 在完整模态下的实验

在CI指标上，HGCN在所有六个癌症队列中都取得了最高的平均值，一致地优于所有SOTA和基线模型。在BS指标上，HGCN在五个队列中取得了最低（最好）的值，在UCEC上也取得了具有竞争力的结果。这些结果证明了HGCN在完整数据设置下的临床价值和有效性。

作者将HGCN与在每个数据集上表现最好的SOTA方法进行了KM分析的比较。在所有六个癌症队列中，HGCN的p-value都低于相应的SOTA方法，表明其对高/低风险患者的分层能力更强。HGCN在患者分层方面同样表现出色。

3.3 可解释性分析

得益于基于图的多模态数据建模，HGCN框架提供了丰富的可解释性。

3.4 在缺失模态下的实验

作者比较了HGCN与Zero padding, MFM, Autoencoder等处理缺失数据的方法。在所有六种缺失模态的设置下（例如，仅有临床记录、临床+病理等），HGCN的性能都一致地优于所有基线方法。这证明了HGCN的online MAE范式能够有效处理缺失模态，并保持了模型的鲁棒性。

作者可视化了在不同缺失模态设置下，MAE重建的hyperedge值与在完整模态下生成的hyperedge值的关系。结果显示，两者之间存在很强的相关性（Pearson相关系数$\rho > 0.8$）。这表明online MAE能够有效地捕捉模态间的内在依赖关系，并准确地重建缺失模态的高阶信息，从而提升了模型的鲁棒性。