通过Transformer进行单细胞多模态预测.
0. TL; DR
scMoFormer框架能够以端到端的方式进行学习，利用下游任务信息，并灵活地整合外部领域知识（如STRING蛋白质互作数据库）。scMoFormer通过为不同模态（细胞、基因、蛋白质）分别设计Transformer模块来建模模态内部的交互，同时利用跨模态聚合组件来建模模态间的关系。为了解决细胞数量庞大带来的计算复杂性问题，作者在细胞维度上采用了线性化Transformer。实验证明，scMoFormer在多个基准数据集上均取得了优越的性能，并在NeurIPS 2022的一项竞赛中取得了Top 2%的优异成绩。
1. 背景介绍
多模态单细胞技术的进步使得在同一个细胞内同时分析多种数据类型（如染色质可及性、基因表达、蛋白质水平等）成为现实。这些技术以前所未有的分辨率揭示了细胞的身份和状态，极大地促进了对多细胞生物体中基因调控网络的研究。
然而，尽管多模态单细胞数据迅速积累，其分析仍面临诸多挑战。首先，单细胞测量数据通常表现出高度的稀疏性和噪声，这使得从中提取有意义的生物学见解变得困难。其次，由于样本通常在不同批次、时间或实验条件下测量，会引入系统性的技术偏差，进一步增加了数据解释的复杂性。
当前的研究为了有效捕捉细胞和基因间的复杂相互作用，主要采用构建静态图的方法，并在此基础上应用图神经网络（GNN）进行信息提取。这些图通常基于启发式的相似性度量（如k-NN图）或利用公共数据库中的先验知识（如STRING数据库）构建。然而，这类静态图的构建过程并未利用下游任务的信息，可能导致图结构并非最优。此外，GNN自身也存在一些固有的局限性，如在深度堆叠网络层数时会出现“过平滑”（over-smoothing）和“过压缩”（over-squashing）问题，限制了其性能。
鉴于Transformer模型在捕捉对象间成对关系方面的强大能力，作者探索了利用Transformer端到端地学习细胞、基因和蛋白质之间交互图的可能性。Transformer通过自注意力机制（self-attention），让所有对象相互关注，并用可学习的注意力分数来表示它们的交互强度，从而以数据驱动的方式表征交互关系。为了将图结构中的先验知识（如生物知识图谱）与Transformer的优势相结合，Graph Transformer应运而生。它既能利用GNN从图结构中学习到的先验信息，又能通过注意力机制学习数据特有的交互，同时缓解了GNN的过平滑和过压缩问题。因此，研究（图）Transformer在单细胞分析中的潜力具有重要意义。
2. scMoFormer框架
scMoFormer框架由三个核心部分组成：多模态图构建、多模态Transformer和预测层。其整体流程如图所示。
2.1 多模态图构建
作者首先构建了一个包含细胞、基因和蛋白质三类节点的异构图。该图由四个子图组成。
- 蛋白质-蛋白质图（Protein-Protein Graph）: 为了整合先验生物学知识，作者利用STRING数据库构建蛋白质之间的交互网络。该数据库包含了多种来源（如基因组背景、实验结果、文本挖掘等）的蛋白质互作信息。作者使用所有来源的组合置信度分数作为边的权重。
- 基因-基因图（Gene-Gene Graph）: 同样利用STRING数据库构建，以保证信息来源的一致性。通过MyGene.Info等工具将STRING中的蛋白质编码（ENSP）与数据集中的基因编码（ENSG）进行对齐。
- 基因-蛋白质图（Gene-Protein Graph）: 根据分子生物学的中心法则（信息从RNA翻译为蛋白质），如果一个基因和一个蛋白质共享相同的符号（symbol），则认为该蛋白质由该基因编码，并在这两者之间建立连接。
- 细胞-基因图（Cell-Gene Graph）: 该图基于数据驱动构建。如果一个基因在某个细胞中表达，就在该细胞和基因节点之间建立连接。边的权重由细胞内标准化的基因表达计数取对数后得到，以描述细胞与基因之间的关联强度。
值得注意的是，作者没有构建静态的细胞-细胞图，而是利用Transformer的注意力机制来动态学习细胞间的关系，以避免单细胞数据噪声导致的不准确性。同时，为了防止模型过拟合，也没有直接建立细胞与目标蛋白质节点之间的连接。
最终，整个多模态异构图的邻接矩阵$A$可以表示为：
\[A = \begin{pmatrix} A_p & S^T & 0 \\ S & A_g & A_{RNA}^T \\ 0 & A_{RNA} & 0 \end{pmatrix} \in \mathbb{R}^{N \times N}\]其中，$A_p$是蛋白质-蛋白质互作图，$A_g$是基因-基因互作图，$S$是基因与蛋白质之间的编码关系矩阵，$A_{RNA}$是细胞-基因表达关系矩阵。$N$是节点总数，即$N = N_p + N_g + N_c$。
节点特征初始化：
- 细胞（Cell）: 原始的基因表达数据（RNA计数）维度高且稀疏。作者首先对其进行库大小归一化和中心化对数变换，然后使用奇异值分解（SVD）进行降维，得到初始的细胞特征$h_c^0$。
- 基因（Gene）: 基因的初始特征$h_g^0$通过其在所有细胞中表达的加权和得到，即$h_g^0 = \bar{X}_g^T \cdot h_c^0$，其中权重为标准化的基因表达计数$\bar{X}_g$。
- 蛋白质（Protein）: 由于蛋白质是预测目标，它们的特征根据其索引进行随机初始化。
2.2 多模态Transformers
为了处理异构图中的不同节点类型，scMoFormer为每种模态都设计了专门的Transformer，并通过跨模态聚合机制进行信息融合。
细胞Transformer (Cell Transformer)
为了学习细胞间的复杂关系，作者引入了Transformer。然而，标准Transformer的计算复杂度为$O(N_c^2 d)$，对于动辄数万个细胞的单细胞数据集来说是不可接受的。
为了解决这个问题，作者采用了广义可核化注意力（generalized kernelizable attention），即一种线性化的Transformer。其注意力计算可以近似为：
\[Attn(h_c) = \hat{D}^{-1}(Q'((K')^T h_c))\]其中 $\hat{D} = \text{diag}(Q’((K’)^T 1_{N_c}))$。这里的$Q’$和$K’$是经过特征映射$\phi(\cdot)$后的查询和键矩阵。这种方法将时间和空间复杂度都降低到了线性级别$O(N_c r d)$和$O(N_c r + N_c d + r d)$，使得模型能够高效处理大规模细胞数据。
基因与蛋白质Transformer (Gene and Protein Transformer)
对于基因和蛋白质，作者采用了Graph Transformer架构（参考GraphGPS），以同时利用先验知识图（来自STRING）和学习数据特有的全局关系。
一个Graph Transformer层包含两个并行部分：一个消息传递的GNN块和一个全局注意力块。
\[h^{l+1}_g = \text{MLP}^l(h^{l+1}_{g,M} + h^{l+1}_{g,T})\]其中，$h^{l+1}{g,M} = \text{GNN}^l(h^l_g, A_g)$从先验知识图中提取局部邻域信息，$h^{l+1}{g,T} = \text{Attn}^l(h^l_g)$通过全局注意力学习数据特有的全局关系。这种混合结构既利用了生物学先验，又克服了纯GNN的表达能力瓶颈。
此外，为了区分图中具有相似局部结构但位置不同的节点，作者还引入了位置编码（Positional Encoding, PE），包括：
- 拉普拉斯PE (Laplacian PE): 利用图拉普拉斯矩阵的特征向量来捕捉节点的谱信息。
- 随机游走PE (Random Walk PE): 利用节点在k步随机游走后返回自身的概率来捕捉节点的局部簇信息。
跨模态聚合 (Cross-Modality Aggregation)
为了在不同模态的Transformer之间建立联系，作者在连接不同类型节点的边上使用了消息传递GNN（具体为GraphSAGE）。信息在源节点$u_j$和目标节点$v_i$之间通过聚合（Aggregate）和更新（Update）两个步骤进行传递：
\[h^{(l+1)}_{v_i} = \text{Update}(h^l_{v_i}, \text{Aggregate}(\{h^l_{u_j}, \forall u_j \in N(v_i)\}))\]这样，来自细胞Transformer的数据驱动信息可以传递给基因和蛋白质，而来自基因和蛋白质Transformer的生物学先验知识也可以指导细胞表示的学习。
最终，每一层的节点嵌入更新可以总结为：
\[\begin{aligned} h^{l+1}_g &= \text{GT}^l_g(h^l_g, A_g) + \text{MPG}^l_{p \to g}(h^l_p, S^T) + \text{MPG}^l_{c \to g}(h^l_c, A_{RNA})\\h^{l+1}_c &= \text{Trans}^l_c(h^l_c) + \text{MPG}^l_{g \to c}(h^l_g, A^T_{RNA}) + \text{FC}^l(h^l_c)\\h^{l+1}_p &= \text{GT}^l_p(h^l_p, A_p) + \text{MPG}^l_{g \to p}(h^l_g, S) \end{aligned}\]其中GT代表Graph Transformer，Trans代表细胞Transformer，MPG代表消息传递GNN。
2.3 预测层
经过$L$层网络后，作者将所有中间层的细胞嵌入$h_c^l$进行拼接，然后通过一个全连接层得到最终的蛋白质水平预测值$\hat{X}_p$：
\[\hat{X}_p = \text{FC}^{L+1}(\text{Concat}(h^1_c, h^2_c, \dots, h^L_c))\]模型使用均方误差作为损失函数进行优化：
\[L(\hat{X}_p, X_p) = \frac{1}{N_p N_c} \sum_{i=1}^{N_p} \sum_{j=1}^{N_c} (X_{ij_p} - \hat{X}_{ij_p})^2\]3. 实验分析
作者在一系列基准数据集上对scMoFormer的性能进行了评估。
3.1 与基线模型的比较
作者将scMoFormer与四种最先进的多模态预测模型（CMAE、BABEL、scMM、ScMoGNN）在CITE和GEX2ADT两个数据集上进行了比较。评估指标包括均方根误差（RMSE）、平均绝对误差（MAE）和皮尔逊相关系数（Corr）。
scMoFormer在两个数据集上的所有三个指标都显著优于所有基线模型。这表明scMoFormer能够更准确地捕捉从基因表达到蛋白质水平的定量关系。特别是在CITE数据集上，scMoFormer的RMSE比第二好的模型ScMoGNN低了0.04，并且在Corr指标上取得了显著提升，同时表现出更低的运行方差，说明模型更稳定。
该图展示了在所有模型中平均RMSE最高的10种蛋白质上的表现。scMoFormer和ScMoGNN在这些最难预测的蛋白质上表现相对稳定，并且持续优于BABEL和CMAE，证明了其在处理挑战性预测任务上的鲁棒性。
3.2 位置编码（PE）的影响
作者探究了Laplacian PE和Random Walk PE对模型性能的影响。
PE的影响因数据集而异。在CITE数据集上，不使用PE时性能最好；而在GEX2ADT数据集上，两种PE都带来了性能提升。
作者认为这是因为CITE数据集的RNA零表达率（RNA Zero Rate）远低于GEX2ADT，数据本身包含的信息更丰富，因此不需要PE提供的额外先验知识。而对于更稀疏的GEX2ADT数据，PE提供的结构信息则能有效提升性能。此外，捕捉局部信息的Random Walk PE在两个数据集上都优于捕捉全局信息的Laplacian PE，表明在这些任务中，局部结构信息可能比全局结构更重要。
3.3 不同融合策略的比较
作者比较了三种不同的信息融合策略：并发融合（Concurrent，模型原始策略）、GNN优先融合（GNN-first）和混合融合（Mixed）。
原始的并发融合策略在两个数据集上都表现最佳。并发融合允许每种模态在与其他模态信息结合之前，先独立处理其模态内部的信息。由于不同模态对下游任务的重要性不同，这种策略能带来更好的性能。
3.4 对其他模态预测任务的处理能力
作者测试了scMoFormer在其他单细胞预测任务上的泛化能力，如ADT2GEX（蛋白质到基因）、GEX2ATAC（基因到染色质可及性）等。
scMoFormer在涉及蛋白质模态的任务（GEX2ADT、ADT2GEX）上优于ScMoGNN。scMoGNN在不涉及蛋白质模态的任务（GEX2ATAC、ATAC2GEX）上表现更好。这表明scMoFormer中为蛋白质节点设计的先验知识图和Transformer模块在处理与蛋白质相关的任务时能发挥关键作用。
3.5 可扩展性与效率
作者比较了scMoFormer与表现最好的基线模型ScMoGNN的训练效率。
scMoFormer的训练时间显著短于ScMoGNN，但消耗了更多的GPU内存。内存消耗更高是因为scMoFormer为了达到最佳性能，设置了较大的批处理大小（batch size）和隐藏维度。作者指出，通过调整这些参数可以降低内存需求。
作者从数据和模型两方面论证了其可扩展性。数据层面，可以通过对细胞进行小批量（mini-batching）处理；模型层面，核心的细胞Transformer是线性复杂度的，而GNN部分（GraphSAGE）的复杂度也与批处理大小相关，因此模型可以方便地适应大规模数据集。
3.6 消融实验
作者通过一系列消融实验来探究模型各组件的贡献。
移除细胞-细胞图和细胞-蛋白质图的合理性
在异构图中加入静态的细胞-细胞图或细胞-蛋白质图都会导致性能下降，后者下降更明显。基于基因表达构建的静态图可能因数据噪声而引入偏差信息。相比之下，由细胞Transformer动态学习的细胞交互关系更优越。将目标蛋白质水平作为边的权重会引入训练集标签信息，导致模型严重过拟合。
每个Transformer的影响
- 图(a) (CITE数据集)：基因Transformer和蛋白质Transformer的贡献大于细胞Transformer。
- 图(b) (GEX2ADT数据集)：细胞Transformer的贡献最大。
这再次与数据的稀疏性有关。GEX2ADT数据稀疏，细胞Transformer对于从数据本身挖掘信息至关重要。而CITE数据更稠密，此时通过基因和蛋白质Transformer引入外部生物学知识能带来更大收益。
如何利用先验知识
将先验知识图加入到GNN模型中（GNN-prior vs GNN）能带来轻微的性能提升。而在引入Transformer后（scMoFormer vs GNN-prior），性能得到了大幅提升。
这表明先验知识本身是有用的，但Transformer架构是有效利用这些知识并将其转化为显著性能增益的关键。
