用于千兆像素全切片图像生存预测的多模态协同注意力Transformer.

0. TL; DR

Multimodal Co-Attention Transformer (MCAT) 框架学习了一个WSI和基因组特征之间的可解释的、密集的共注意力映射。核心贡献在于一种跨模态注意力机制Genomic-Guided Co-Attention (GCA)，使用基因组特征作为“查询”，从庞大的WSI bag中识别出信息丰富的实例。这不仅实现了多模态交互的可视化，还显著降低了WSI bag的空间复杂度。

作者在五个不同的癌症数据集（共计4,730张WSI，6700万个patches）上进行了实验。结果表明，MCAT的性能始终优于现有SOTA方法。

1. 背景介绍

尽管深度学习已在许多领域取得革命性突破，但将其应用于计算病理学中的数G像素whole-slide imaging (WSI) 仍然面临巨大障碍，尤其是在癌症预后等监督学习任务中。

与自然图像分类不同，病理WSI的特点是尺寸巨大，可达150,000 x 150,000像素。标签的稀疏性或全局性：标签可能只与图像中极小一部分区域相关，如微转移灶的检测；也可能依赖于整个图像的组成和各成分间的复杂交互，如基质、肿瘤团、免疫细胞等。

生存结局预测就属于后者，它需要理解肿瘤微环境中异构的视觉概念，这超出了传统multiple instance learning (MIL) 方法的能力。MIL将WSI视为一个由patches（实例）组成的bag（包），它能解决“大海捞针”问题，但难以建模实例间的复杂、长程交互。

此外，生存预测本质上是一个多模态学习任务。临床上，病理学家会结合histology（组织学）和genomics（基因组学）信息进行风险评估。然而，在计算模型中融合这两种数据存在巨大的数据异质性鸿沟（data heterogeneity gap）：WSI被表示为包含数万个图像patches的bag，而基因组特征通常是表格化的属性。这导致许多方法只能采用晚期融合，限制了对多模态交互的学习。

为了应对这些挑战，作者提出了Multimodal Co-Attention Transformer (MCAT) 框架。该框架旨在：

  1. 实现可解释的早/中期融合：通过co-attention机制，学习组织学特征和基因组特征之间的直接关联。
  2. 解决WSI的计算复杂性：利用基因组信息来指导WSI的特征聚合，从而降低其表示的复杂度。
  3. 提升特征聚合能力：引入Transformer架构来建模WSI中实例间的长程依赖关系。

2. MCAT 框架

MCAT框架的核心思想是将WSI和genomic数据都看作是bags（实例包），然后通过一个基因组引导的共注意力（Genomic-Guided Co-Attention）层来捕捉它们之间的多模态交互，并最终利用set-based MIL Transformers进行特征聚合和生存预测。

WSI和Genomic Bag的构建

作者首先将问题形式化为Multiple Instance Learning (MIL)。一个bag $X = {x_1, …, x_M}$ 是一个包含$M$个实例的无序集合。MIL的目标是学习一个置换不变的函数$F(X)$来预测整个bag的标签。

1. WSI Bag构建

首先，从WSI中分割出所有组织区域，并将其切割成不重叠的256x256的patches。将这些patches输入一个在ImageNet上预训练的ResNet-50 CNN，提取出$d_k$维的特征嵌入$h$。

对于一个患者的所有WSIs中的$M$个patches，将提取出的特征嵌入打包成一个WSI bag: $H_{bag} \in \mathbb{R}^{M \times d_k}$。在实践中，$M$的平均值约为15,231，最大可达230,000。

2. Genomic Bag构建

传统的基因组特征（如突变状态、RNA-Seq丰度）是$1 \times 1$的标量属性，缺乏描述其生物学功能的语义信息。为了获得类似词嵌入的、更具表达力的特征表示，作者首先将基因根据其生物学功能（如Tumor Suppression, Oncogenesis等）分为$N$个不同的集合。然后，对于每个基因集合，使用一个fully-connected (FC) 层（参数化为$\phi_g$）将其编码为一个$d_k$维的基因组嵌入$g$。最终，得到$N$个基因组嵌入，打包成一个Genomic bag: $G_{bag} \in \mathbb{R}^{N \times d_k}$。在实现中，$N=6$。

Genomic-Guided Co-Attention (GCA) Layer

GCA层是MCAT的核心，它利用Genomic bag ($G_{bag}$) 来引导对WSI bag ($H_{bag}$) 的特征聚合。其机制类似于Transformer中的标准注意力机制。

GCA将$G_{bag}$作为Query ($Q$)，将$H_{bag}$作为Key ($K$) 和 Value ($V$)，计算得到一个聚合后的WSI表示$\tilde{H}$：

\[\text{CoAttn}_{G \to H}(G, H) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V = A_{\text{coattn}}V \to \tilde{H}\]

其中，$A_{\text{coattn}} \in \mathbb{R}^{N \times M}$是co-attention（共注意力）矩阵。

对于Genomic bag中的每一个基因组嵌入$g_n$（代表一种生物学功能），GCA层会计算$H_{bag}$中所有patch嵌入$h_m$与之的相似度，得到一行注意力权重$[a_{n1}, …, a_{nM}]$。然后，用这行权重对$H_{bag}$进行加权求和，得到一个新的WSI级别的特征嵌入$\tilde{h}_n$。这个$\tilde{h}_n$可以被看作是一个基因引导的视觉概念，它主要聚合了那些在表型上与基因功能$g_n$相关的patches。

GCA层的一个关键作用是，它将一个包含$M$个实例的巨大WSI bag ($H_{bag} \in \mathbb{R}^{M \times d_k}$) 聚合（或压缩）成一个只包含$N$个视觉概念的小bag ($\tilde{H} \in \mathbb{R}^{N \times d_v}$)，其中$N \ll M$。这使得后续应用计算复杂度为$O(N^2)$的self-attention（自注意力）层成为可能。

基于集合的MIL Transformer与生存预测

作者将Transformer架构引入到MIL框架中，用于对聚合后的WSI表示$\tilde{H}$和Genomic bag $G_{bag}$进行进一步的特征聚合。

作者证明了Transformer中的self-attention层是一个permutation-equivariant（置换等变）的集合函数，而position-wise的前馈网络和残差连接也保持了置换不变性。因此，可以将多个Transformer编码器块堆叠起来，用于在全局池化之前对bag中的元素进行复杂的交互建模。

在Transformer编码器之后，作者使用了一个全局注意力池化函数，自适应地计算bag中所有嵌入的加权和，得到最终的bag级特征表示$h^{(L)}$和$g^{(L)}$。

将图像和基因组的最终bag级特征向量进行拼接，再通过几个FC层得到最终的共享表示$h_{\text{final}}$，用于生存风险预测。模型的训练使用了标准的Cox生存损失函数，以及一个用于离散化生存时间的损失函数。

3. 实验分析

作者在TCGA的五个大型癌症数据集上进行了实验，包括BLCA、BRCA、GBMLGG、LUAD和UCEC，共计4,730张WSI和6700万个patches。

3.1 与SOTA方法的比较

作者将MCAT与多种先进的MIL方法进行了比较，包括Deep Sets、Attention MIL和DeepAttnMISL。为了公平比较，作者还为这些基线方法实现了两种多模态融合机制：concatenation（拼接）和bilinear pooling（双线性池化）。

与仅使用基因组（SNN）或仅使用WSI（Attention MIL, DeepAttnMISL）的单模态基线相比，MCAT在所有基准测试中都取得了优越的性能，总体C-Index性能分别提升了23.9%，6.35%和12.4%。

与采用late fusion（拼接或双线性池化）的多模态基线相比，MCAT在总体C-Index上实现了3.0%至6.87%的性能提升。在5个癌症类型中的4个上，MCAT都取得了最高的C-Index。

在GBMLGG（胶质母细胞瘤和低级别胶质瘤）这个异质性极强的数据集上，MCAT取得了所有模型中最高的C-Index（0.817），并能很好地区分高低风险病例。

3.2 消融研究

作者通过消融实验评估了Transformer层在MIL任务中的作用。结果显示，与仅使用全局注意力池化相比，加入Transformer层（$T_H$或$T_G$）能够持续提升总体C-Index性能。这表明Transformer能够正确地建模在两种模态内部对癌症生存具有预后意义的成对特征交互。

3.3 注意力可视化

作者通过可视化GCA层计算出的共注意力权重，来分析WSI和基因组之间的特征交互。

作者发现，基因组引导的WSI嵌入能够反映许多已知的基因型-表型（genotype-phenotype）关系。

这些可视化结果不仅为模型的预测提供了可解释性，还揭示了histology形态学特征是如何与特定的基因功能相关联的。