多模态癌症生存分析的结构化预后事件建模.

0. TL; DR

SlotSPE 是一个基于slot的结构化预后事件建模框架。受factorial coding（因子编码）原理的启发，SlotSPE使用slot attention（槽注意力）机制将每个患者的高维多模态输入压缩成一组紧凑的、模态特定的、相互区别的slots。每个slot被视为一个潜在的预后事件的编码。引入一个Mixture-of-Experts解码器为每个患者选择性地激活最具预测性的slots。通过一个跨模态重建任务，该框架利用组学衍生的slots从组织学图像中预测基因表达。这不仅加强了模态间的生物学对齐，也提高了在基因组数据缺失时的鲁棒性。

作者在十个癌症基准数据集上进行了广泛的实验。结果显示，SlotSPE在10个队列中的8个上优于现有方法，总体性能提升了2.9%。即使在基因组数据缺失的情况下，它也保持了鲁棒性，并通过结构化的事件分解提供了显著改善的可解释性。

1. 背景介绍

癌症生存分析旨在评估患者的死亡风险或生存时间，是实现个性化预后和精准肿瘤学的关键工具。Whole-slide images (WSIs)（全切片图像）和genomics（基因组学）是两种广泛使用的互补数据源：WSI揭示了癌症在组织中的形态学表现，而genomics则揭示了其潜在的分子驱动因素。

然而，有效且高效地进行多模态学习以进行生存分析面临一个根本性挑战：尽管患者的结局通常由少数关键事件驱动，但WSI和基因组信号本身却极其复杂和高维。这些关键事件表现为有意义的结构化关联，例如，在组织病理学中，免疫活性和肿瘤侵袭性反映在肿瘤浸润淋巴细胞和癌巢的密度和邻近关系上；在基因组学中，肿瘤促进效应源于PI3K/AKT和MAPK等通路的共激活。

但是，从海量输入中学习和发现这些与事件相关的结构化模式非常困难。现有方法，特别是基于原型的模型（如PIBD, MMP），存在一些局限性。它们要么使用固定的原型，缺乏对个体患者的适应性；要么原型本身在训练中不可学习，导致次优的端到端优化。

本文中，作者提出了SlotSPE，一个基于slot的结构化预后事件建模框架。受factorial coding思想的启发，该方法将高维输入压缩成一组紧凑的语义slots，每个slot对应一个潜在的预后事件。与固定的原型不同，这些slots是为每个患者动态实例化的，从而允许模型在保持简洁表示的同时，捕捉个性化的预后模式。

2. SlotSPE 模型

SlotSPE的框架如图所示。其核心是使用slot attention将组织学和基因组学特征压缩成slots，然后通过选择性激活和跨模态交互来建模生存风险。

2.1 基于Slot的信息压缩 (Slot-based Information Compression)

作者使用slot attention模块将WSI和基因组学两种模态的高维实例集（bag）压缩成一个紧凑的、包含$S$个slots的表示。

Slot Attention通过$T$次迭代的cross-attention，一组可学习的slots $S^{(\tau)}$与输入特征$X$进行交互和更新。在每次迭代中，每个输入特征$x_j$的注意力权重在所有slots上进行softmax归一化，这强制slots之间产生竞争，从而实现结构化的分解。

\[\begin{aligned} \alpha_j^{(\tau)} &= \text{softmax}_{S^{(\tau)}} \left( \frac{Q(S^{(\tau)}) K(x_j)}{\sqrt{d_{slot}}} \right)\\ u_k^{(\tau)} &= \sum_{j=1}^M \alpha_{kj}^{(\tau)} V(x_j)\\ s_k^{(\tau+1)} &= \text{RNN}(s_k^{(\tau)}, u_k^{(\tau)})\\ s_k^{(\tau+1)} &\leftarrow s_k^{(\tau+1)} + \text{MLP}(s_k^{(\tau+1)}) \end{aligned}\]

这种方法将输入从$\mathbb{R}^{M \times d}$压缩到$\mathbb{R}^{S \times d_{slot}}$，其中$S \ll M$。

如图(a)所示，为了避免slots的冗余，作者引入了一个Mixture-of-Experts (MoE) 风格的解码器。一个轻量级的门控函数$\phi$为每个slot $s_k$预测一个保留分数$r_k$。使用Gumbel-Top-K技巧选择得分最高的$K$个slots，生成一个可微的$K$-hot掩码$\hat{G}$。最终的患者级别预测$y$是所有被选中slots的预测结果$\ell_k$的加权和$y = \sum_{k=1}^S w_k \ell_k$。其中，权重$w_k$由$r_k$经过softmax和掩码$\hat{G}$计算得到。

为了防止未被激活的slots变成无用的“null slots”，作者引入了一个重建目标$L_{recon}$，要求slots能够重建出原始的模态特征。

2.2 具有生物学引导的多模态Slots交互

如图(b)所示，作者设计了一个跨模态重建任务，以编码生物学先验知识（即分子事件导致组织病理学表型）。

将组织病理学patch特征$X_h$输入到由组学衍生的slots $S_g$初始化的slot attention模块中，得到跨模态的slot表示$\tilde{S}_{g \to h}$。然后，要求这些跨模态slots能够重建出原始的基因组学特征$X_g$。

\[\tilde{X}_g = R(Q_g, \tilde{S}_{g \to h})\\ L_{recon}^{g \to h} = ||\tilde{X}_g - X_g||_2^2\]

这个任务不仅加强了模态间的对齐，还使得模型在基因组数据缺失时，能够从组织学图像中推断出基因组特征，从而提高了模型的鲁棒性。

为了捕捉多模态依赖关系，作者在模态内和模态间两个层面并行地对slots进行交互建模。

总的训练目标$L$是生存预测损失$L_{surv}$和重建损失$L_{recon}$的加权和。在测试时，如果所有模态都可用，则不需要重建。只有当基因组数据缺失时，才会激活跨模态重建模块作为一种填补机制。

3. 实验分析

3.1 实验设置

在来自TCGA的十个公共癌症队列上进行了评估。与三类SOTA方法进行了比较：(1) 单模态方法（基因组学和组织学）；(2) 多模态方法；(3) 基于原型的方法。

评估指标使用C-index、Kaplan-Meier曲线、log-rank检验和restricted mean survival time (RMST)。

3.2 与SOTA方法的比较

SlotSPE在十个癌症类型中的总体C-index达到了0.721，显著优于所有比较方法，比第二好的方法高出2.9%。即使在单模态设置下（即仅使用基因组学或仅使用WSI），SlotSPE的性能也优于相应的单模态基线，证明了其框架的普适性。

与基于原型的方法（PIBD, MMP）相比，SlotSPE的性能更优，这表明其动态捕捉患者特异性事件的能力强于静态原型。通过动态捕捉结构化的预后事件并进行跨模态对齐，SlotSPE成为了一个强大且可靠的生存预测框架。

3.3 鲁棒性与风险分层

作者评估了在基因组数据缺失的情况下的模型性能。SlotSPE保持了强大的性能，在某些队列中（如HNSC, STAD），其缺失模态变体的性能甚至能持平或超过完整模态版本。

与专门为缺失数据设计的LD-CVAE相比，尽管SlotSPE的性能下降幅度稍大，但其绝对性能在大多数数据集上仍然更优或相当。这得益于其跨模态重建任务，使得模型能够从组织学中推断基因组信息，从而在基因组数据缺失时提供鲁棒的预测。

KM曲线显示，SlotSPE能够将患者显著地分层为高风险和低风险组。与第二好的方法LD-CVAE相比，SlotSPE在大多数队列中都取得了更大的RMST差异和更小的RMST比率，表明其分层效果更清晰、更具临床意义。

3.4 消融研究与效率分析

移除选择性slot激活、slots正则化、跨模态重建或迭代式cross-attention中的任何一个，都会导致性能一致下降。这证明了SlotSPE的每个设计组件都对最终性能至关重要。

作者对slot数量、迭代次数、top-K选择等超参数进行了详细的消融研究，验证了当前参数选择的合理性。

作者比较了SlotSPE与四个最强基线在推理和训练时的内存和时间开销。SlotSPE在取得最高C-index的同时，保持了较低的运行时间和内存使用。它在效率和预测性能之间取得了最佳的平衡。

3.5 可解释性分析

SlotSPE能够建立事件级别的跨模态对齐，并揭示通路-形态学之间的对应关系，为生物学上合理的解释和新的科学假设提供了可能。