单细胞预训练模型

几篇单细胞预训练基座模型文献学习笔记

Transfer learning enables predictions in network biology （2023-5）

主要方法：迁移学习策略，在 29 million 的单细胞转录组数据上进行自监督预训练，然后在下游任务中进行微调

方法特点：

• 每个细胞中的基因表达转化成 rank value encoding：对于一个细胞中的基因，先将基因的转录本 counts 除以这个细胞中总的转录本 counts，然后再除以这个基因在整个数据集细胞中的转录本 counts 的非零中位数（所有非零值的中位数），接着对这些值进行排序得到 Rank value，这种方法会将那些高表达的管家基因的 Rank 变低，而将那些低表达的转录因子 Rank 变高；模型的词汇表为 25424 个基因（能够检测到这些基因的细胞中位数为 173152），也就是说每个基因相当于一个词，基因的特征是 25424 维的词向量，这个向量中该基因的位置的值为 rank value 而其他位置为 0
• 采用 mask 基因的方法进行预训练：将 15% 的基因掩盖掉，然后预测在这些位置上的基因是什么
• 算法优化：输入长度是 2048，对于自注意力模型比较大（Bert 是 512），在 Trainer​ 中采用自定义的 Tokenizer​ 来实现动态的，分组长度填充，可以实现在预训练步骤 29 倍的速度提升，另外还使用了 Deepspeed​ 进行并行计算；先抽样 megabatch （多个 minibatch 在一起）然后按照 minibatch 的序列长度进行降序排列（内存限制会较早碰到），这样每个 minibatch 里面的长度是类似的，因此可以减少 padding 的时间

代码：https://huggingface.co/ctheodoris/Geneformer

scGPT: Towards Building a Foundation Model for Single-Cell Multi-omics Using Generative AI（2023-7）

​

主要方法：在 33 million 单细胞转录组数据上进行预训练，对下游任务进行微调，包括细胞类型注释，遗传扰动预测（包括逆向扰动预测，也就是给定细胞的表达预测在细胞上对哪个基因进行了扰动），批次矫正和多组学整合。

方法特点：

• scGPT 有三个输入（gene token，表达值，条件 token），M 是提前定义的最大输入长度：

  • 对基因使用 Gene Token 进行编码，其实就是对每个基因给一个 ID ，然后由这些 ID 组成一个词汇表

  • 为了减少不同测序方法，不同测序深度，不同批次的影响，作者开发了一种 value binning 的方法将基因的 counts 转化为相对表达值，将每个细胞中的非零表达基因的表达值排序，分成平均的 B 份，基因的值就是这个基因在哪一份中

    

  • 条件 token 是其他的元信息，比如扰动的值，也使用 ID 来表示

  • 对于基因 token 和条件 token，使用 embedding layer 将其转化为 D 维度的 embedding vector，对于 binning 表达则使用全连接层进行转化，最后基因的 embedding （h）为三个的相加（维度为 M*D）

• 对得到的 h 采用堆叠的自注意力 transformer 层进行转化，由于输入的维度 M 比较大，这里使用了 Flash-Attention 技术进行优化（github.com/Dao-AILab/fla...）；M 设定为 1200，当细胞的非零表达基因大于 1200 的时，在每个epoch 迭代中就随机选择 1200 个基因作为输入

• 细胞的表示来自 <cls>​​ token（最开始的 token）

• scGPT 在进行下游的微调任务之前将批次和组学类型也编码成 token 并进行 embedding，将这些 embedding 和 transformer 的输出进行合并，然后进行下游的微调任务。组学（模态）特征是 feature 层面的，表示这个 feature 是基因，蛋白还是 ATAC-seq 中的一个 peak 区域，批次 Token 是细胞层面的，但是也广播到单个 feature （进行复制）。

• 预训练流程：采用的是生成式预测的预训练框架（gene prompts：根据已知基因生成未知基因的表达值；cell prompts：根据细胞类型条件生成整个基因组基因的表达值）。但是和自然语言不同，细胞中的基因 Token 的顺序是可交换的，因此和自然语言模型中的预测 “下一个” token 的含义有所不同的，因此需要使用不同的 attention mask 策略（在 GPT 里面是 mask 后面的 tokens）。作者这里面开发了第一个用于非序列数据的自回归生成式预训练框架：

  ​

  上图表示 M*M 的 attention map，query 是行，key 是列，蓝色表示 mask 的位置，这个特点就是 cls​ token 和已知基因的 token 不和未知基因的 token 做注意力计算，而未知基因除了和已知基因和 cls​ token 之外只和自己的 token 做注意力计算。在训练的时候随机挑选一部分基因作为 unknown 基因，将这些基因的表达值去掉（也就是这些基因的输入只有 gene token embedding 和条件 token embedding 的和），然后利用上面的 attention mask 进行 Transformer 转化，利用得到的 embedding 预测未知基因的表达值。在 cell promot 的推理步骤中，第一步所有的基因都是未知的基因，只有第一个位置是 cell embedding，然后从预测中选择 1/K 个可信度最高的基因作为已知基因，进行逐步迭代，最终将所有基因预测（上图 B，类似语言模型中的自回归模式）。gene prompts 则是开始的时候是选择一些已知基因，剩下的步骤和 cell promot 一样。

目前预训练的框架有两种：1）预测 mask 的 token，如 BERT（使用的是 Transformer 的 encoder），2）自回归生成式的顺序预测，比如 GPT 模型（使用的是 Transformer 的 decoder）

代码：https://github.com/bowang-lab/scGPT

Large Scale Foundation Model on Single-cell Transcriptomics （2023-6）

主要方法：xTrimoGene 在 50 million 单细胞转录组数据上进行预训练，对下游任务进行微调，是目前最大参数（一亿参数），最大数据量的单细胞预训练模型。

方法特点：

• 输入维度

• 不进行基因筛选

• 模型架构：

  • Embedding module 将细胞的表达值转化成 Transformer 可以处理的 embedding 向量。这个模型的特点就是考虑了表达为 0 的基因，对于 0 表达值，直接转化成随机的初始 embedding（维度为 d）；对于非零表达值（X_i​），先通过一个带有 leakyRelu 激活函数的线性层得到中间 embedding：

    v_1 = leakyReLU(X_i*w_1)

    w_1 的维度为 1*b，这样 v_1 的维度就转变成了 1×b，接着利用另一个线性层转化成 v_2:

    v_2 = v_1 \cdot w_2+\alpha*v_1

    w_2 的维度为 b*b，\alpha 是一个可学习的参数，然后利用 softmax 函数将 v_2 进行归一化得到 v_3，也就是 v_3 中的 b 个元素加起来是 1（直接对 v_1 操作不也可以吗，为什么要中间生成一个 v_2），接着将这个 v_3 和一个初始化的 look-up 矩阵 T 相乘（维度是 b * d），相当于这个矩阵 T 中 b 个 embedding 的加权和，得到基因的值 embedding（维度就变成了 d）。除了对值进行 embedding 外，还需要对基因的身份进行 embedding，这里采用的是一般的词汇表方式，使用一个初始化的矩阵 T^G，维度为 19266 × d，最终的基因 embedding 则是值的 embedding 加上身份的 embedding；mask embedding 也是随机的初始化 embedding（维度为 d）

  • Encoder module 使用 Transformer 处理非零和非 mask 的基因表达，得到中间的 embedding，所以这个 encoder 的输入大概是总的长度的 10% 左右

  • Decoder module 输入为 Encoder module 得到的中间 embedding 和 Embedding module 得到的零值和 mask 的 embedding 进行合并，因此长度重新变成了 19266 维，然后使用 **Performer ** 层（Tranformer 的核近似变体，可以用来处理比较长的输入序列）得到 19266 * f 维度的输入，然后采用 MLP 进行转化（T 和 S 的 embedding 丢弃，T 和 S 的含义见后面）最终得到一个预测值（预测的基因表达值）

• 测序深度感知的预训练任务：作者对输入数据进行了增强，用 T 表示细胞中原始数据的 read counts 总数，S 表示增强后的细胞中 read counts 总数；增强的操作是（两）层次贝叶斯降采样过程（即第一个参数绝对要不要抽样，如果要抽样就使用 Beta-二项分布来采样）为什么要用 Beta-二项分布？ ：

  ​

  想法是，如果原始输入不变的话，模型可以学习到一个细胞中不同基因的相互关系，而如果输入是降低后的 read counts，那么模型就可以学习到不同测序深度的细胞之间的相关性；这个特性使得模型可以更好的迁移到不同的数据集上。在训练的时候，随机 mask 30% 基因的表达值，然后利用模型来预测这些基因的表达，使用的损失函数为 MSE。

代码：https://github.com/biomap-research/scFoundation

‍