用于鲁棒多模态生存预测的条件潜在微分变分自编码器.
0. TL; DR
LD-CVAE (Conditional Latent Differentiation Variational AutoEncoder) 用于实现即使在基因组数据缺失情况下的鲁棒多模态生存预测。作者开发了LD-VAE，它通过学习基因组的和功能特异性的后验分布，使用product-of-experts技术来整合基因组后验和图像后验，以估计联合潜在分布，来生成具有多样化功能的基因组嵌入。
作者在五个不同的癌症数据集上测试了该方法的有效性，实验结果表明，无论是在完整模态还是缺失模态的场景下，该方法都显示出优越的性能。
1. 背景介绍
生存分析是评估癌症治疗和精准医疗的有效解决方案。病理图像和基因组数据是两种重要的生物标志物来源。病理图像提供了丰富的视觉信息，如组织学特征和细胞组织结构，被认为是癌症预后的金标准。而基因组数据则揭示了癌症潜在的分子机制。
近年来，研究人员开始探索将这两种模态结合起来进行癌症预后预测。现有方法如MCAT和MOTCat等，通过co-attention（协同注意力）transformer框架取得了不错的进展。这些研究表明，结合基因组和病理图像能比单一标志物更好地对癌症患者进行风险分层。
然而，一个关键的现实问题被大多数研究忽略了：数据不完整性。采集基因组数据的成本远高于获取病理图像，这使得基因组数据在测试样本中可能无法获得，从而导致多模态模型的性能下降。
一种常见的解决思路是利用一种可用模态（如病理图像）来重建缺失模态（如基因组）的特征。但这种策略在实践中面临两大挑战：
- WSI的有效表示：WSI尺寸巨大，其中只有一小部分区域与基因型信息密切相关。如何从中学习到有效的表示来重建基因组特征是一个难题。
- 生成多样化的基因组嵌入：不同基因通常具有不同的生物学功能。在一个统一的生成框架中，如何解决这种差异性，生成具有不同功能类别的基因组嵌入，是另一个挑战。
基于此，作者提出了LD-CVAE，一个即使在基因组数据缺失的情况下，也能从WSI中重建具有多样化生物学功能的基因组嵌入，从而实现鲁棒多模态生存预测的模型。
2. LD-CVAE 方法
LD-CVAE的整体框架如图所示。它主要由三个部分组成：多模态包表示、Conditional Latent Differentiation VAE以及多模态融合与预测。
2.1 数据构建与表示
- 病理学特征：使用预训练的CTransPath（一个Swin Transformer）作为特征提取器$\phi_p(\cdot)$，将WSI中的每个patch编码为一个特征向量，形成病理学特征集$Y^{(k)}$。
- 基因组学特征：根据生物学功能将基因组数据分为6个类别。然后，使用一个全连接层$\phi_g(\cdot)$提取每个类别的特征，形成基因组学特征集$X^{(k)}$。
2.2 条件变分自编码器 (Conditional VAE)
模型的核心目标是在给定病理学特征$Y$的条件下，生成目标基因组学特征$X$。Conditional VAE (CVAE) 的目标是最大化证据下界（ELBO），其损失函数可以表示为：
\[L_{ELBO} = -E_{q_\phi(z|X,Y)} \log p_\theta(X|z,Y) + \beta KL[q_\phi(z|X,Y), p(z)]\]其中，$q_\phi(z|X,Y)$是后验分布，$p_\theta(X|z,Y)$是解码器生成的似然。
2.3 病理学VIB Transformer (VIB-Trans)
为了从巨大的WSI中学习压缩且有用的表示，作者提出了VIB-Trans。它基于变分信息瓶颈（Variational Information Bottleneck, VIB）理论，旨在学习一个既能最大程度压缩原始病理学特征$Y$，又能最大程度保留对生存预测目标$T$有用信息的潜变量$z_Y$。
VIB-Trans的核心是一个Transformer编码器。它将一个WSI的所有patch特征作为输入，并引入两个可学习的tokens，$\mu_{token}$和$\Sigma_{token}$，用于学习整个WSI的全局信息，并最终输出后验分布$p(z_Y|Y)$的参数$\mu_Y$和$\Sigma_Y$。
2.4 条件潜在微分变分自编码器 (LD-CVAE)
2.4.1 潜在微分VAE (LD-VAE)
为了解决生成具有不同功能类别的基因组特征的挑战，作者提出了LD-VAE。对于$N$个功能类别的基因组特征${x_1, …, x_N}$，作者假设它们在给定一个共同的潜变量$z_X$时是条件独立的。但直接从$z_X$生成所有的$x_i$会影响生成特征的多样性。因此，作者引入了一个“latent differentiation”过程。
为每个基因组特征$x_i$分配一个独特的潜变量$z_i$，并假设存在一个马尔可夫链$x_i \leftrightarrow z_i \leftrightarrow z_X$。这意味着，可以从共同的基因组后验$p(z_X|X)$中“微分”出每个功能特异性的后验$p(z_i|X)$。
\[p(z_i|X) \approx E_{z_X \sim q(z_X|X)} p(z_i|z_X)\]LD-VAE使用一个transformer编码器从所有基因组特征中学习共同的后验分布参数$\mu_X$和$\Sigma_X$。然后，通过一组映射器$\psi_i$将共同的潜变量$z_X$转换为每个功能特异性的潜变量$z_i$，再由各自的解码器$\theta_i$重建出$x_i$。
2.4.2 联合病理-基因组分布学习
为了估计LD-CVAE中的联合后验分布$q(z|X,Y)$，作者采用了Product-of-Experts (PoE) 技术。PoE假设联合后验可以由各个模态的边际后验的乘积来近似：
\[p(z|X,Y) \propto \frac{p(z|X)p(z|Y)}{p(z)} \approx \tilde{q}(z|X)\tilde{q}(z|Y)p(z)\]由于高斯分布的乘积仍然是高斯分布，因此可以以闭合形式计算出联合后验的参数。为了增强病理学特征和基因组学特征之间的一致性，作者引入了一个对齐损失，旨在最小化两个边际潜分布$q(z|X)$和$q(z|Y)$之间的Wasserstein距离。
\[L_{Align} = ||\mu_X - \mu_Y||^2 + ||\Sigma_X - \Sigma_Y||^2\]2.5 多模态融合与预测
在得到病理学特征$Y$和由LD-CVAE重建的基因组学特征$\tilde{X}$后，作者使用一个Genomic-Guided Co-Attention Layer来建模两者之间的成对交互，并引导对生存相关WSI实例的选择。
最终的损失函数由三部分组成：生存预测损失$L_{surv}$、VIB-Trans损失$L_{VIB-Trans}$和LD-CVAE损失$L_{LD-CVAE}$。
3. 实验分析
3.1 实验设置
使用了TCGA中的五个癌症数据集：BLCA, BRCA, GBMLGG, LUAD, UCEC。
评估指标C-index和Kaplan-Meier (KM) 分析。与多种单模态（MLP, SNN, TransMIL等）和多模态（Porpoise, MCAT, MOTCat, CMTA, SurvPath, PIBD）的SOTA方法进行比较。
3.2 在完整模态下的比较
多模态方法的性能普遍优于单模态方法，证明了整合基因组学和组织学的优势。与所有比较的多模态方法相比，作者提出的方法在5个数据集中的4个上取得了最高的C-index，总体性能提升了1.0%-3.9%。这些结果清晰地证明了作者的方法能够更好地学习基因组表示并融合多模态数据，以用于癌症生存预测。
3.3 在缺失模态下的比较
作者将模型在完整数据上进行训练，然后在仅有病理图像、基因组数据缺失的场景下进行测试。与VAE, GAN, MVAE, SMIL, ShaSpec等专门用于处理缺失模态的方法相比，作者的方法在所有数据集上都取得了更高的C-index值。这突显了作者的方法在处理基因组数据缺失的生存预测任务上的强大能力。
3.4 消融研究
- VIB-Trans的影响：与使用TransMIL作为基线WSI编码器相比，VIB-Trans在完整和缺失数据场景下分别带来了1.8%和1.9%的总体性能提升，证明了其在有效表示WSI方面的优势。
- LD-VAE的影响：与使用普通VAE相比，LD-VAE在完整和缺失数据场景下分别带来了1.2%和0.9%的性能提升，验证了其在生成具有功能差异性的基因组表示方面的优势。
- 对齐损失的影响：移除对齐损失后，模型性能下降。这表明对齐损失有助于提升多模态整合的一致性，从而增强联合后验的估计。
3.5 患者分层与可视化
作者使用模型预测的风险中位数将患者分为高风险（红色）和低风险（绿色）组。在完整模态和缺失基因组数据两种场景下，模型都能在所有五个癌症队列中显著地区分高/低风险患者（p-value均很低）。这进一步清晰地证明了作者的方法生成的基因组特征也与生存信息相关，从而有助于预测癌症患者的临床结局。
作者可视化了由真实基因组特征和生成的基因组特征计算出的co-attention权重。结果显示，两者产生了相似的co-attention分数可视化结果，并且它们各自对应的top-5高相关性patches也具有相似的视觉外观。这再次表明，作者的方法能够生成与病理图像具有紧密联系的基因组特征。
