Arc虚拟细胞挑战入门指南

来自：https://fleetwood.dev/posts/virtual-cell-challenge

背景

Arc研究所最近公布了虚拟细胞挑战赛。参赛者需要训练一个模型，该模型能够预测在（部分）未知细胞类型中沉默某个基因所产生的效果，他们将这项任务称为上下文泛化。对于几乎没有生物学背景的机器学习工程师而言，相关的专业术语和所需的背景知识可能显得十分棘手。为了鼓励更多人参与，这里以一种更适合其他学科工程师的形式来简要介绍这项挑战赛：

• 目标：训练一个模型，以预测使用CRISPR技术沉默某个基因对细胞产生的影响。

在原子世界中做实验成本高昂、费力且容易出错。要是我们不用接触培养皿就能测试数千种候选药物，会怎么样呢？这正是虚拟细胞挑战赛的目标——构建一个模型（很可能是神经网络），它能精确模拟当我们改变某些参数时，细胞会发生什么变化。

要训练这个神经网络，我们需要数据。为了此次挑战赛，Arc 研究所整理了一个包含约 30 万个单细胞的 RNA 图谱数据集。

训练数据

训练集由一个稀疏矩阵和一些相关的元数据组成。更具体地说，我们有22万个细胞，每个细胞都有一个转录组。这个转录组是一个稀疏行向量，其中每个条目代表相应基因（即列所对应的基因）编码的RNA分子（转录本）的原始计数。在这22万个细胞中，约有3.8万个是未受扰动的，也就是说没有使用CRISPR技术沉默任何基因。我们很快就会发现，这些对照细胞至关重要。

为了更具体地理解这个数据集，我们选取一个基因——TMSB4X（该数据集中被沉默次数最多的基因），并对比在对照细胞和经扰动处理的细胞中检测到的RNA分子数量。

• 我们可以看到，与对照细胞相比，TMSB4X基因被沉默的细胞中，转录本的数量大幅减少。

建模

敏锐的人可能会疑惑，为什么不直接测量基因沉默前后的RNA分子数量——我们到底为什么需要对照细胞呢？遗憾的是，读取转录组的过程会破坏细胞，这是一个让人联想到观测者效应的问题。

这种无法测量细胞扰动前后状态的情况带来了诸多问题，因为我们不得不将一组基础状态（basal）的细胞（也称为对照细胞、未受扰动细胞）用作参考点（reference point）。即便在扰动发生之前，对照细胞和受扰动细胞也并非完全属性一致。这意味着我们现在必须从异质性引发的噪声中分离出真正的信号——即扰动带来的影响。

更正式地说，我们可以将受扰动细胞中观察到的基因表达建模为：
$${\hat{X}}_p\sim {\hat{T}}_p\left({\mathcal{D}}_{basal}\right)+H\left({\mathcal{D}}_{basal}\right)+\epsilon ,\epsilon \sim P_{\epsilon }$$

其中：

• ${\hat{X}}_p$：在受到扰动 $p$ 的细胞中观察到的基因表达测量值。
• ${\mathcal{D}}_{basal}$：未受扰动的基线细胞群体的分布。
• ${\hat{T}}_p({\mathcal{D}}_{basal})$：扰动 $p$ 对该群体产生的真实效应。
• $H\left({\mathcal{D}}_{basal}\right)$：基线群体的生物学异质性。
• $\epsilon$：特定实验的技术噪声，假设其与未受扰动的细胞状态及${\mathcal{D}}_{basal}$无关。

STATE：Arc开发的baseline

在虚拟细胞挑战赛之前，Arc研究所发布了STATE——这是他们尝试用一对基于Transformer的模型来解决该挑战的成果。它为参赛者提供了一个强有力的基准起点，因此我们将对其进行详细探讨。

STATE由两个模型组成，分别是状态转换模型（ST）和状态嵌入模型（SE）。SE旨在生成丰富的细胞语义嵌入，以努力提升跨细胞类型的泛化能力。ST是“细胞模拟器”，它接收对照细胞的转录组，或是SE生成的细胞嵌入，再加上一个代表目标扰动的独热编码向量，最终输出受扰动后的转录组。

ST

状态转换模型是一种相对简单的Transformer，以Llama backbone为架构，其作用对象如下：

1. 一组经过协变量匹配的basal cells（手动匹配确保相同批次与相同cell line）的转录组（或SE嵌入）。
2. 一组独热向量，代表每个细胞的基因扰动情况。

使用一组经过协变量匹配的对照细胞（ontrol cells ）与目标细胞（target cells）配对，应该能帮助模型识别我们预期扰动带来的实际效应。对照集张量和扰动张量均通过独立的编码器输入，这些编码器只是带有GELU激活函数的4层多层感知器（MLP）。如果直接在基因表达空间中运行（即生成完整的转录组），它们会将输出结果通过一个经过训练的解码器。

状态转换模型（ST）采用最大均值差异（Maximum Mean Discrepancy）进行训练。简单来说，该模型通过学习来最小化两个概率分布之间的差异。

SE

状态嵌入模型（SE）是一种BERT模型，通过掩码预测任务进行训练。为了更深入地理解这一点，我们首先需要稍微偏离一下主题，补充一些生物学基础知识。

基因由外显子exons（蛋白质编码区段）和内含子introns（非蛋白质编码区段）组成。如上图所示，DNA首先被转录为前体mRNA（pre-mRNA）。随后，细胞会进行可变剪接——这本质上是“对外显子进行筛选和选择”，并将所有内含子剪切掉。你可以把基因想象为一张桌子组装说明书：人们也可以通过省略某些部件，组装出一张三条腿的桌子，甚至费点功夫做出一个造型奇特的书架。这些不同的“物件”就类似于蛋白质异构体——由同一基因编码的不同蛋白质。

有了这些基础知识，我们就可以继续了解状态嵌入模型（SE）的工作原理了。请记住，SE的核心目标是生成有意义的细胞嵌入。要实现这一点，我们首先必须生成有意义的基因嵌入。

为了生成单个基因嵌入，我们首先获取目标基因所编码的所有不同蛋白质异构体（protein isoform）的氨基酸序列（例如，TMSB4X的序列……）。然后，我们将这些序列输入ESM2——这是来自FAIR（Meta人工智能研究院）的一个拥有150亿参数的蛋白质语言模型。ESM会为每个氨基酸生成一个嵌入，我们通过均值池化将这些嵌入合并，从而得到一个“转录本”（又称蛋白质异构体）嵌入。

现在我们已经得到了所有这些蛋白质异构体（protein isoform）嵌入，接下来只需对它们进行均值池化，就能得到基因嵌入。然后，我们使用一个经过训练的编码器将这些基因嵌入投影到模型维度，具体如下：
$${\tilde{g}}_j=SiLU\left(LayerNorm\left(g_j{W}_g+b_g\right)\right)$$

（注：SiLU是Sigmoid Linear Unit的缩写，即Sigmoid线性单元，是一种激活函数；LayerNorm指层归一化，是一种归一化技术）

我们现在已经得到了基因嵌入，但我们真正想要的是细胞嵌入。为实现这一点，Arc研究所将每个细胞表示为按对数倍表达量排名的前2048个基因。

然后，我们根据这2048个基因嵌入构建“细胞句子”，具体如下：
$${\tilde{c}}^{(i)}=\left[z_{cls},{\tilde{g}}_1^{(i)},{\tilde{g}}_2^{(i)},...,{\tilde{g}}_L^{(i)},z_{ds}\right]\in {\mathbb{R}}^{(L+2)\times h}$$

我们在句子中添加了[CLS] token（标记）和[DS] token（标记）。[CLS]标记最终被用作我们的“细胞嵌入”（非常类似BERT），而[DS]标记则用于“分离数据集特异性效应”。尽管基因是按对数倍表达量排序的，但Arc研究所通过以一种类似于位置嵌入的方式整合转录组，进一步强化了每个基因的表达量大小。借助一种独特的“软分箱”算法和两个多层感知器（MLP），他们创建了一些“表达编码”，然后将其添加到每个基因嵌入中。这应该能根据每个基因在转录组中的表达强度来调节其嵌入的大小。

为了训练该模型，他们在每个细胞中掩盖1280个基因，模型的任务就是对这些被掩盖的基因进行预测。所选的这1280个基因具有广泛的表达强度。对于偏好图形化理解的人来说，下图展示了“细胞句子”的构建过程。

评估

了解你的提交作品将如何被评估是成功的关键。Arc选择的三个评估指标是扰动区分度、差异表达和平均误差。由于平均误差简单易懂，正如其名，我们将不在分析中对其进行阐述。

扰动区分度

扰动区分度旨在评估你的模型能够在多大程度上揭示不同扰动之间的相对差异。为此，我们计算测试集中所有已测量的扰动转录组（即我们试图预测的真实值$y_t$，以及所有其他扰动转录组$y_p^n$）与我们预测的转录组${\hat{y}}_t$之间的曼哈顿距离。然后，我们进行如下排序：
rt=∑p≠t1{d(y^t,yp)<d(y^t,yt)}r_{t}=\sum_{p \neq t} 1\left\{d\left(\hat{y}_{t}, y_{p}\right)<d\left(\hat{y}_{t}, y_{t}\right)\right\}rt​=p=t∑​1{d(y^​t​,yp​)<d(y^​t​,yt​)}

之后，我们通过转录组的总数进行归一化：
$$PDis{c}_t=\frac{r_t}{T}$$

其中，0是完美匹配，总分是所有$PDis{c}_t$的均值，再被归一化为：
$$PDiscNorm=1-2PDisc$$

我们乘以2是因为对于随机预测而言，大约一半的结果会更接近真实值，而另一半则会更偏离真实值。

差异表达

差异表达旨在评估在真正受影响的基因中，你正确识别出显著受影响的基因的比例。首先，对于每个基因，使用带有结校正的威尔科克森秩和检验（Wilcoxon rank-sum test with tie correction）计算p值。我们会同时对预测的扰动分布和真实的扰动分布进行这一操作。

接下来，我们应用本杰明尼-霍奇伯格（Benjamini Hochberg）校正法——这本质上是通过一些统计学处理来调整p值。因为当研究20,000个基因且p值阈值设为0.05时，会预期出现1000个假阳性结果。我们将预测的差异表达基因集合记为$G_{p,pred}$，将真实的差异表达基因集合记为$G_{p,true}$。

如果我们的集合规模小于真实集合规模，就取两个集合的交集，再除以真实差异表达基因的数量，具体如下：
$$D{E}_p=\frac{G_{p,pred}\cap G_{p,true}}{n_{p,true}}$$

如果我们预测的集合规模大于真实集合规模，则选择我们预测为差异表达最显著的子集（即我们的“最有信心”的预测，记为${\tilde{G}}_{p,pred}$），取其与真实集合的交集，然后除以真实差异表达基因的数量：
$$D{E}_p=\frac{{\tilde{G}}_{p,pred}\cap G_{p,true}}{n_{p,true}}$$

对所有预测的扰动都执行此操作，然后取平均值以获得最终分数。