计算病理学中转录组引导的切片表示学习.

0. TL; DR

TANGLE是一种slide和expression预训练策略。该策略的核心思想是：利用来自基因表达谱的互补信息来引导切片表示的学习。TANGLE通过模态特定的编码器处理WSI和基因表达数据，并利用对比学习来对齐这两种模态的输出。

该模型在一个包含来自人类和大鼠的肝脏、乳腺和肺部三种不同器官的数千个S+E配对样本上进行了预训练。在三个独立的测试数据集（共计7,795张WSI）上，TANGLE展示了显著优于监督学习和SSL基线方法的few-shot性能。在使用基于原型的分类和切片检索进行评估时，TANGLE也显示出比所有基线方法都优越的性能。

1. 背景介绍

Self-supervised learning (SSL) 近年来在计算病理学（CPath）领域获得了巨大的关注。SSL特别适合于建模WSI，因为它们尺寸巨大，难以直接用Vision Transformers (ViTs) 或Convolutional Neural Networks (CNNs) 处理。因此，CPath领域的主流方法是采用“分而治之”的策略：首先将WSI切分成小的patches，然后用一个预训练的网络（通常是在ImageNet或病理学特定数据集上训练的）来提取低维的patch embeddings。

除了在patch层面应用SSL，研究人员也开始探索在slide层面构建embeddings。这样得到的slide embeddings可以用于各种下游任务，如少样本分类、切片检索等，且无需或只需很少的训练。然而，用SSL构建slide embeddings仍然面临挑战：

  1. 基于patch级增强来构建slide的“视图”计算成本高昂。
  2. 在扩展到巨大的输入时，需要学习的视觉基元和不变性变得不明确。
  3. WSI内部的异质性可能会妨碍模型获得一致且强大的训练信号。

受多模态视觉-语言模型的启发，作者提出利用基因表达数据来引导切片表示学习。WSI（形态学）和基因表达（分子）为同一生物过程提供了不同但互补的视角。作者假设，用基因表达来引导切片表示学习，会比单纯使用切片自身的增强或掩码提供更强的训练信号。

因此，作者提出了TANGLE，一个基于多模态对比学习的框架，旨在通过对齐slide和expression (S+E) 的嵌入来预训练一个强大的切片表示学习模型。

2. TANGLE 方法

TANGLE (TrANscriptomics-Guided sLidE representation learning) 的框架如图所示，它由一个视觉编码器、一个基因表达编码器和一个多模态对齐模块组成。

2.1 切片编码器 (Slide encoder)

该部分的目标是将一张WSI $X_i$ 编码成一个slide embedding $h_i \in \mathbb{R}^d$。

首先将WSI分割成小的patches。然后，使用一个预训练的视觉编码器将每个patch编码成一个patch embedding。对于人类组织，作者使用了CTransPath，一个在超过32,000张WSI上训练的Swin Transformer。对于rat组织，由于缺乏现成的模型，作者从零开始训练了一个名为iBOT-Tox的ViT-Base模型。该模型在来自47,227张WSI的1500万个H&E patches上使用iBOT方法进行了训练。

在得到一个WSI的所有patch embeddings $H_i \in \mathbb{R}^{N_H \times d_H}$后，作者使用一个attention-based MIL (ABMIL) 模型来学习一个映射函数$f(H_i)$，将patch embeddings集合池化成一个单一的slide embedding $h_i$。

2.2 基因表达编码器 (Gene expression encoder)

该部分的目标是将原始的转录组测量值编码成一个expression embedding $g_i \in \mathbb{R}^d$。

首先，计算相对于对照组的log2 fold change，这代表了基因相对于正常样本的过表达或低表达程度。然后，选择log2 fold change最大的前$k$个基因（实验中$k=1000$）。

使用一个3层的MLP $\phi(\cdot)$来学习一个映射，将筛选后的基因表达谱$t_i \in \mathbb{R}^{N_G}$投影到一个与slide embedding维度相同的expression embedding $g_i$。

\[g_i = \phi(t_i): \mathbb{R}^{N_G} \to \mathbb{R}^d\]

2.3 多模态对齐 (Multimodal alignment)

作者使用一个对称的跨模态对比学习目标来对齐slide和expression编码器的嵌入空间。

对于一个包含$M$个(S+E)配对样本$(h_i, g_i)$的batch，其目标函数$L_{SymCL}$定义为：

\[L_{SymCL} = - \frac{1}{2M} \sum_{i=1}^M \log \frac{\exp(\tau h_i^T g_i)}{\sum_{j=1}^M \exp(\tau h_i^T g_j)} - \frac{1}{2M} \sum_{j=1}^M \log \frac{\exp(\tau g_j^T h_j)}{\sum_{i=1}^M \exp(\tau g_j^T h_i)}\]

该损失函数的第一项是slide-to-expression的对比损失，第二项是expression-to-slide的对比损失。每一项都旨在最大化来自同一配对的嵌入之间的点积相似度，同时通过Softmax归一化来推远与其他负样本对的相似度。

2.4 下游任务

在推理阶段，一张查询slide被送入训练好的视觉编码器和MIL模块，得到其slide embedding，然后用于各种下游任务。

作者在三个下游任务上评估了TANGLE的性能：

  1. 少样本分类 (Few-shot classification)：使用线性探测（linear probing）。
  2. 基于原型的分类 (Prototype-based classification)：通过计算查询slide与预定义原型之间的L2距离进行分类。
  3. 切片检索 (Slide retrieval)：通过计算slide embeddings之间的余弦距离来检索相似的slides。

3. 实验分析

3.1 数据集

3.2 少样本线性探测分类

利用基因表达数据作为引导信号的S+E预训练范式，比纯视觉的SSL和传统的监督MIL方法能学习到更具判别力的切片表示。

3.3 基于原型的少样本分类

该实验旨在评估TANGLE构建的slide-level原型能否用于识别特定的形态学特征。结果显示，TANGLE和TANGLE-REC（TANGLE的变体）在两个病变检测任务中都优于所有基线方法。一个惊人的发现是，仅用3个或更多样本构建原型的TANGLE，其性能就超过了在一个包含18,552张WSI的完整训练集上训练的ABMIL模型。这突显了TANGLE学习到的嵌入空间的强大判别能力，以及基于原型的方法在特定形态学检测任务上的巨大潜力。

3.4 切片检索

该任务旨在评估TANGLE能否检索出与查询slide具有相同药物、剂量和处死时间等特征的slides。TANGLE取得了最佳的检索性能，在top-10的检索结果中，平均能正确找回2.88/4张相关切片，在top-20中则能找回3.44/4张。这表明TANGLE能够捕捉到由不同药物剂量或处死时间引起的微妙形态学差异。

3.5 可解释性分析

TANGLE学习到的表示不仅具有强大的预测能力，而且在生物学上是可解释的，能够与已建立的生物学发现对齐。