Transformer在生成细胞数据上的应用

来自：scTranslator: A pre-trained large generative model for translating single-cell transcriptome to proteome

工程地址：https://github.com/TencentAILabHealthcare/scTranslator

在scTranslator中，有3个阶段：paired bulk上监督学习，paired sc上监督学习，sc RNA上推理。前两阶段的监督信号来自配对数据，从一个模态翻译到另一个模态。在第三阶段，预训练的scTranslator可以被用于预测sc分辨率的protein丰度，而且如果存在部分配对的RNA和protein的sc数据，scTranslator还能在这个subset上进行微调。

scTranslator 论文写的十分清晰，提供的工具也非常容易使用

数据集

用于预训练的bulk数据集
TCGA（The Cancer Genome Atlas）包含63个数据集，有30种cancer type，CPTAC（Clinical Proteomic Tumor Analysis Consortium）包含7个数据集，7种cancer。最终，scTranslator的统一bulk数据集：
fig1

用于预训练的sc数据集
包含了10个独立的数据集，这些数据集包含来自各个组织的2,405,115个细胞，组织包括blood，bone marrow，liver。这些样本代表了一系列不同的情况，包括健康捐赠者和被诊断患有COVID-19、结核病、多系统炎症综合征、癫痫和肥胖的捐赠者。

用于benchmarking的sc数据集
Seurat v4 PBMC：cell-161,764，gene-23,385，protein-224
REAP-seq PBMC：cell-4,330，gene-21,005，protein-44
CITE-seq CBMC：cell-8,005，gene-16,508，protein-13
sc Pan-caner：patient-43，cancer-8，cell-65,698，gene-15,844，没有protein数据

预处理
考虑到scTranslator有2万个基因作为输入，为了保留完整的基因级信息和可解释性，scTranslator没有采用特征选择和特征提取的降维方法。对于RNA或蛋白质的计数矩阵，对每个细胞的数据进行缩放的Max-Min归一化形式化为：
eq1
其中，cell index $i=1,2,...,num_{cell}$ ，scale factor $H=1,Low=10^{-8}$ 。 $L o w$ 的目的是区分真正的0计数（0计数被缩放到 $10^{−8}$ ）和用于模型输入的 0 padding。

模型架构

scTranslator建立在Transformer架构上，可以有效利用RNA和protein组学提供的所有信息。尽管scTranslator遵循类似于Transformer的编码器到解码器框架，但它没有在解码器中使用掩码多头注意力，该注意力用于以自回归的方式聚焦输出嵌入中的信息，以防止未来位置的信息泄露。这种舍弃是基于生物学原理，即NLP中token之间的绝对位置关系对于蛋白质丰度数据是没有意义的。为了使模型能够区分不同基因，受Transformer的位置编码的启发，设计了一个 re-indexed GPE模型。为了应对超长RNA序列的挑战，使用了Transformer的线性复杂度版本Performer来代替具有二次空间和时间复杂度的常规full-rank注意力。scTranslator使用端到端训练策略构建，其中模型由预测蛋白质和ground truth之间的MSE损失监督。

Re-indexed gene position encoding (GPE)和表达值embedding
以往的研究局限于固定RNA和蛋白质的位置，只考虑相应基因的表达值。然而，这种方法忽略了从基因名称及其相互作用中获得的有价值信息。为了解决这一不足，scTranslator提出了一个re-indexed的GPE模块。具体来说，将所有基因(总计75,500个，来自最新的NCBI数据库)重新索引为离散的ID值。这里避免使用NCBI提供的原始基因ID号，因为最大ID (1,255,000) 很大，这可能导致内存冗余。GPE模块采用embedding layer，通过可学习的查找表实现，将re-indexed的Gene ID转换为连续的向量空间。该模块支持的最大词汇量为85,500，其中有10,000个保留位置用于当前研究中未涵盖的新基因。与scTranslator的其余部分一起，GPE模块被同时训练以获取所有基因的信息表示。

除了基因ID之外，如何利用RNA的表达价值也是一个挑战。NLP中的标准Transformer模型接受token的1D序列，并通过嵌入层将token映射到嵌入向量。在scTranslator编码器中，每个基因的表达被视为一个“token”，来自一个细胞或患者的所有表达值构成一个“sentence”。原始值通过线性嵌入层投影到预先设置维数的向量上。

如扩展图A7所示，GPE和表达值嵌入的结合为全面理解基因表达及其相互作用提供了可能。此外，GPE模块使scTranslator在非对齐模式下有效，增强了扩展性。

exp_fig10

扩展图A7：GPE与value embedding的工作流程。

模型backbone
scTranslator采用Transformer编码器-解码器架构，并采用多头FAVOR+机制作为backbone。编码器由 $N = 2$ 个相同的层组成，每层包括一个多头FAVOR+机制的sub-layer和一个位置全连接前馈网络的sub-layer。在每个sub-layer周围使用残差连接和层归一化。解码器也由 $N = 2$ 个相同的层组成，具有与编码器相同的sub-layer设计，以及用于表示蛋白质ID而不是RNA ID的额外GPE模块。sub-layers、embedding layers和GPE module产生的输出维度为 $d_{model} = 128$ 。考虑到目前RNA和蛋白质的覆盖范围，编码器和解码器的最大序列长度分别设置为20,000和1,000。在编码器和解码器之间使用多层感知器(MLP)进行sentence长度转换。此外，在每个子层中都加入了一个8头注意机制。整个模型由1.17亿个参数组成。

Performer中提出的FAVOR+机制在理论和实验上都被证明可以将常规注意力机制所需的二次计算复杂度降低到线性复杂度。假设 $X\in\R^{L\times d_{model}}$ 是包含位置嵌入和表达嵌入的序列长度为 $L$ 的输入嵌入。query，key，value是输入的中间表示： $Q=W_{Q}X,K=W_{K}X,V=W_{V}X$ 标准的双向dot-product注意力为（其中 $A\in\R^{L\times L}$ 为注意力矩阵）：
eq2
FAVOR+机制尝试找到一个通用kernel来映射 $Q$ 或 $K$ 到 $Q^{'}$ 或 $K^{'}$ ，以使得 $A$ 有以下近似：
eq3
映射函数的通用形式为：
eq4
其中，函数 $f_{j}:\R\rightarrow\R$ ，函数 $h:\R^{d}\rightarrow\R$ ，随机从分布 $D\in P(\R^{d})$ 中进行独立同分布采样 $w_{i}\sim D$ 。近似注意力为：
eq5

第一阶段：在配对bulk上训练
在第一阶段，scTranslator对配对的bulk RNA和protein数据进行训练。该阶段的每个数据点对应于一个tissue样本，并从同一样本中获得RNA和蛋白质的测量值。本阶段的目的是在不同组织类型中建立RNA表达水平和蛋白质丰度之间的牢固关系。这一阶段使用的数据集包括从31种不同组织类型中收集的18,227对RNA和蛋白质样本。训练参数包括学习率为 $2 × 10^{−4}$ ，batch size为8，epoch为200。

第二阶段：配对sc上持续学习
在第二阶段，通过对配对的sc RNA和protein数据的持续学习，进一步训练第一阶段预训练的scTranslator。此阶段中的每个数据点表示单个细胞，这允许模型在单个细胞级别捕获更细粒度的关系。这一阶段增强了模型从scRNA-seq数据预测蛋白质丰度的能力。这一阶段的数据集包括2,405,115对单细胞RNA和蛋白质样本。除了epoch的数量外，训练参数与阶段1保持一致。考虑到数据的庞大和异构性质，通过观察预测和真实蛋白质水平之间的余弦相似度来监测模型在预训练期间的性能。当余弦相似度超过0.85，即停止预训练。

损失函数与优化
在scTranslator的所有训练阶段，观察到的和预测的蛋白质水平之间的均方误差(MSE)被用作损失函数。这种损失函数的选择确保模型最小化预测和实际蛋白质丰度之间的差异，从而生成更准确和可靠的预测。Adam优化器用于训练。

训练时间
两阶段的预训练均在32台V100 gpu上进行，总训练时间约为1个月。

注意力用于调控推断

注意机制可以识别数据点之间的重要特征和相互作用，使其非常适合于基因表达分析中的调控推断。scTranslator实施了一个全面的可解释性分析，以确定重要的基因，以及RNA-蛋白质之间的相互作用。

编码器中的注意力得分反映了每个RNA的贡献和RNA pair的相互作用。同样，可以通过解码器的注意力矩阵来识别关键蛋白质和蛋白质之间的相互作用。此外，从编码器到解码器的注意力矩阵揭示了RNA在推断蛋白质时的重要性。计算所有样本的注意力矩阵，并将它们的和作为最终的注意力矩阵。以编码器到解码器的注意力矩阵为例，在这个矩阵中，每个值 $A (i, j)$ 表示蛋白质 $i$ 对基因 $j$ 的注意力程度。注意矩阵可以通过热图可视化。此外，在g:Profiler网站https://biit.cs.ut.ee/gprofiler上可以对排名靠前的基因进行富集分析。