scGPT：从基因 token 化到表达生成

发表于 更新于 分类于 ， ，

scGPT 在想什么？

如果我们认为 Transformer 能理解文意与文法，那我们也不妨假定 Transformer 能理解基因之间的关系。如果文字能变成 token 喂进 Transformer 训练，然后吐出文字，那基因和它的表达量也可以喂进 Transformer 训练，然后吐出未知表达量基因的表达量。

scGPT 是怎么做的？

总览

大致概念

scGPT 是一个预训练模型。换而言之就是，我们认为它对细胞内的情况了解得不错，所以它既能直接被拿来生成其他的基因，又能在输出层的位置加点东西改装一下，然后通过微调完成其他的生物学任务。

预训练流程

我们把各种各样测到的数据喂进 scGPT，然后试图在两个任务上优化它：

  1. 随机隐藏一些基因表达，让模型预测被隐藏位置的表达量
  2. 额外给出细胞状态，让模型生成全基因表达组

不同于 BERT 在预训练时同时做掩码猜词和句对匹配两个任务，scGPT 做的这两个任务是有先后顺序的。

在第一次前向传播时，模型先做 gene prompt，把已知基因表达和未知基因身份喂进 Transformer，让它预测未知表达，得到 hcellh_\text{cell}（一个 Transformer 的中间向量，可以把它类比为对整个细胞状态的摘要），进而得到 lossgene\text{loss}_\text{gene}。

然后进行第二次前向传播，把第一次得到的 hcellh_\text{cell}、已知基因表达和未知基因身份三个都喂给自己（这次还喂了 hcellh_\text{cell}，也就是 cell prompt），其中 hcellh_\text{cell} 替换掉最前面的 <cls> 向量，运行 Transformer，再次预测未知表达，进而得到 losscell\text{loss}_\text{cell}。

最后的得到 losspretrain=lossgene+losscell\text{loss}_\text{pretrain}=\text{loss}_\text{gene}+\text{loss}_\text{cell}，再进行下一步优化。

这是不是很像自己教自己呢？那会不会把自己越教越蠢呢？其实不是这样的。这个预训练流程不是让一个学生自己制定答案，而是让它写摘要后答题，再拿标准答案批改摘要和答案。

从另一个角度理解，因为训练时始终有一个真实的外部目标，所以，如果第一次得到的 hcellh_\text{cell} 很差，第二次预测就会很差，losscell\text{loss}_\text{cell} 就会变大。反向传播会通过 hcellh_\text{cell} 把误差信号传回去，调整 Transformer，使 hcellh_\text{cell} 逐渐变得更有用。

微调训练流程

对于大部分微调任务，GEP 和 GEPC 基本上都会被使用。

先根据输入获得 hjh_j（每个基因经过 Transformer 后的向量）和 hcellh_\text{cell}，然后优化两个任务：

  1. GEP，Gene Expression Prediction，遮住一些基因，让已有的基因对被遮住的做预测
  2. GEPC，GEP for Cell modeling，遮住一些基因，联合已有的基因和 hcellh_\text{cell}，对被遮住的做预测

GEP 使用被遮住表达量的基因经过 Transformer 后得到的 h被遮h_\text{被遮}，然后通过一个 MLP 把向量转化为标量，进而对被遮住的位置进行表达值预测。GEPC 只使用 hcellh_\text{cell} 和被遮住的基因自己的嵌入向量（未经过 Transformer），对被遮住的位置进行表达值预测。

GEP 更偏向于学习基因与基因之间的上下文关系，而 GEPC 更偏向于让 hcellh_\text{cell} 真的携带细胞整体状态。

对于同一个掩码集合，通常把这两个任务的 loss 加到一起，然后统一反向传播。在完整任务中，还可能继续加上 ECS、DAR、分类损失等。

输入

比如，对于一个细胞，我们获得了这样的信息：

GATA1 CD3 peak_12345
表达量 7 5 3
批次 batch-1 batch-1 batch-1
模态 RNA 蛋白质 ATAC

我们将它们先转换为模型输入：

<cls> GATA1 CD3 peak_12345
ID 0 54 763 21
表达量 0 7 5 3
批次 0 1 1 1
模态 0 2 3 1

于是得到了一个初始输入矩阵。然后取每一列，把 ID 送进基因编码器，把表达量送进表达值编码器，再把批次、模态等等条件送进条件编码器，得到三个向量，再加起来，就得到了嵌入向量。然后我们再把这些嵌入向量一行一行地排到一起，就变成了输入矩阵。形如：

(0.17−0.210.09⋯0.800.130.89−0.53⋯0.320.210.180.33⋯0.82−0.91−0.780.50⋯−0.98)\begin{pmatrix} 0.17 & -0.21 & 0.09 & \cdots & 0.80 \\ 0.13 & 0.89 & -0.53 & \cdots & 0.32 \\ 0.21 & 0.18 & 0.33 & \cdots & 0.82 \\ -0.91 & -0.78 & 0.50 & \cdots & -0.98 \\ \end{pmatrix}

其中：

  • 每一行对应一个基因
  • 每一列对应一个隐藏维度

当然，在训练时，我们会随机遮掉一些基因的表达量。这些表达量就会以 [MASK] 的形式出现并进入表达值编码器，不让真实的表达值编码进入训练。

与 Transformer 原始论文中不同的是，基因同时表达在细胞中，因此不存在先后顺序之分。也因此，scGPT 不需要进行位置编码。

实际情况中还有一些预处理的小技巧，不过我们已经理解了整体过程。小技巧可以等到之后再介绍。

输出

scGPT 可以通过部分已知的基因逐步预测全基因组的表达。它先输出置信度高的，再把原来给的和自己预测的都当作自己的输入，再进行计算与输出。

为逐轮生成，scGPT 使用了掩码注意力，当需要预测未知基因时，直接把除自己以外没有预测的基因遮盖住。这不仅仅是 [MASK] 掉表达量，甚至把这个基因存在这件事都抹掉了。换而言之，等待被预测的嵌入向量就只能看到已经有表达量的基因和被遮盖住的它自己，甚至不知道有没有其他未知基因也在被模型询问表达量。

另外，对于“置信度”的高低，文章原文并没有给出明确的公式或解释。换而言之，论文“逐轮生成”的总体思路清楚，但没有说明如何把注意力分数严格转换成经过校准的预测置信度。

置信度拓展

2019 年发表的 Mask-Predict 给出了很类似的逐步补全方法。它先推断所有被 [MASK] 的 token，然后只取置信度高的，把置信度低的重新 [MASK] 上，然后再次预测，不断循环。然而，词语是离散的，因此只需要 softmax 一下就能轻易地获取概率分布——这本身就在告诉我们模型是否犹豫。但是 scGPT 的输出是连续的，只有一个光秃秃的数字。

我个人简单想了想，认为有三种解决方法。其一是，模型不仅预测表达量，还输出标准差。为避免标准差报得太小或太大，使用高斯负对数似然替代掉普通的 MSE。推断时，优先接受标准差最小的预测。第二种方法是，输出进入某个表达量分箱的概率。这样就可以把连续问题转换为离散问题，直接套用 Mask-Predict 就好。第三种方法是，训练多个 scGPT，让它们一起进行预测，如果一起预测的方差大就说明置信度低，反之说明置信度高。

小技巧

表达量分箱

为保证稳定性，模型在预处理时会先将表达量按百分位数分箱。论文中说，分箱可以缓解测序深度和不同平台造成的数值尺度差异，保留相对表达高低关系。除此以外，我还有一个推测，就是这样做可以缓解“开始时预测的表达量存在 10% 的偏差，进而蝴蝶效应影响后方预测”的情况。

batch 和 modality 拼接

有时候，batch 和 modality 并不进入条件编码器。它们并不转化为条件向量，然后和基因向量、表达量向量加在一起，也不参与 Transformer 流程，而是在 Transformer 输出后再拼接在输出 h 的后面。这种措施常在多组学整合中看到。

这样做可以避免模型过度利用“同一模态”或“同一批次”的表面相关性，从而导致不同模态之间的基因关系被忽略。

预训练时只输入非零表达基因

这样可以显著减少序列长度和注意力计算量。完整人类基因组仍然存在于词表中，只是没有表达的基因通常不进入当前细胞的输入。

超过 1200 个非零基因时随机采样

模型的最大输入长度设置为 1200。如果一个细胞有超过 1200 个非零表达基因，则每次训练迭代随机采样 1200 个基因。这相当于对同一个细胞进行不同的局部观察，也避免了直接截断固定前 1200 个基因带来的偏差。

对扰动预测任务有专门的改动

这些改动包括：

  1. 使用 log1p 后的连续表达值，而不是分箱值；
  2. 每个基因位置增加一个二值扰动 token；
  3. 用对照细胞作为输入；
  4. 用扰动细胞作为目标；
  5. 预测所有输入基因的扰动后表达。

它把“同一个细胞遮住部分基因，预测遮住部分”问题转换为了“对照细胞加扰动，预测扰动后”问题了。