1亿参数的细胞大模型来了!登Nature子刊,清华大学团队发布scFoundation:对2万基因同时建模

数据派THU

共 7490字,需浏览 15分钟

 · 2024-07-05

   

本文约4800字,建议阅读9分钟

清华大学自动化系生命基础模型实验室和电子系/AIR 合作开展研究,构建了拥有 1 亿参数的 scFoundation 细胞大模型。


近年来,大规模预训练模型正在引领新一轮人工智能浪潮。「大模型」通过从大规模、多来源的数据中提取深层次规律,进而能够作为「基础模型」服务不同领域的多样化任务。例如,大语言模型通过学习大量文本数据,掌握了理解和识别语言的能力,革新了自然语言处理领域的范式。


同理,在生命科学领域,生物体也有其「底层语言」——细胞是人体的基本结构和功能单位,如果将 DAN、RNA、蛋白质和基因表达值比作「词语」,其组合在一起就连成了「细胞」这个句子。所以,如果能够基于细胞「语言」开发人工智能细胞大模型,将有望为生命科学和医学提供全新的研究范式和革命性的研究工具。


然而,当前训练大规模单细胞数据主要存在以下三点挑战:


  • 基因表达预训练数据需要涵盖不同状态和类型的细胞景观,目前大多数单细胞 RNA 测序 (scRNA-seq) 数据组织松散,全面完整的数据库仍然缺失;

  • 在训练过程中,传统的 Transformer 难以处理近 2 万个蛋白质编码基因构成的「句子」;

  • 不同技术和实验室的 scRNA-seq 数据在测序深度上存在差异,这妨碍了模型学习统一且有意义的细胞和基因表示。


为了解决这些挑战,清华大学自动化系生命基础模型实验室主任张学工教授、电子系/AIR 马剑竹教授和百图生科宋乐博士合作开展研究,于 2024 年 6 月在 Nature Methods 上发表了题为「Large-scale foundation model on single-cell transcriptomics」的研究论文。


论文介绍了名为 scFoundation 的细胞大模型,能够同时处理约 2 万个基因。作为基础模型,它在细胞测序深度增强、细胞药物响应预测和细胞扰动预测等多种生物医学下游任务中,表现出卓越的性能提升,提供了人工智能在单细胞研究中的新范式。

研究亮点:

  • scFoundation 细胞大模型基于 5 千万个细胞的基因表达数据进行训练,拥有 1 亿参数,能够同时处理约 2 万个基因;

  • 该模型采用非对称设计减少计算和内存挑战;

  • 该模型为基因网络推断和转录因子识别提供了新的研究思路。


论文地址:
https://www.nature.com/articles/s41592-024-02305-7 


开源项目「awesome-ai4s」汇集了百余篇 AI4S 论文解读,还提供海量数据集与工具:

https://github.com/hyperai/awesome-ai4s


数据集:构建全面的单细胞数据集


研究人员通过收集所有公开可用的单细胞资源数据,构建了一个全面的单细胞数据集。其中包括 Gene Expression Omnibus (GEO),Single Cell Portal,HCA,人类基因组计划 (hECA),深度整合人类单细胞组学数据 (DISCO),欧洲分子生物学实验室-欧洲生物信息学研究所数据库 (EMBL-EBI) 等。


* GEO 下载地址:https://www.ncbi.nlm.nih.gov/geo/ 
* Single Cell Portal 下载地址:
https://singlecell.broadinstitute.org/single_cell
* HCA 下载地址:https://data.humancellatlas.org/ 
* EMBL-EBI 下载地址:https://www.ebi.ac.uk/


研究人员将所有数据与 HUGO 基因命名委员会确定的 19,264 个蛋白质编码和常见线粒体基因的基因列表对齐。经过数据质量控制,获得了超过 5 千万个用于预训练的人类 scRNA-seq 数据。



研究收集了 5 千万个单细胞基因表达谱


丰富的数据来源构建了富含生物模式的预训练数据集。在解剖学上,它跨越了 100 多种组织类型,涵盖了各种疾病、肿瘤和正常状态,如上图所示,几乎包括了所有已知的人类细胞类型和状态。


模型架构:构建 1 亿参数的 scFoundation 模型


研究人员开发的 scFoundation 模型具有约 1 亿参数,其参数规模、基因覆盖率和数据规模在单细胞领域均「名列前茅」。


模型设计


研究人员开发了 xTrimoGene 模型作为 scFoundation 的 backbone 模型,这是一个可扩展的基于 Transformer 架构的模型,包括一个向量模块 (embedding module) 和一个非对称的编码器-解码器结构。


其中,向量模块将连续的基因表达标量值转换为可学习的高维向量,确保完全保留原始表达值;编码器以非零 (nonzero) 和非掩码 (nonmasked) 的表达基因为输入,使用了 vanilla transformer 块并具有较大的参数量;解码器以所有基因为输入,使用了 performer 块并具有相对较小的参数量。


与其他架构相比,这种非对称设计减少了计算和内存挑战。数据显示,该模块在保持相同参数规模的情况下,所需的计算量仅为传统语言模型 Transformer 的 3.4%。


预训练任务


研究人员设计了一种名为 RDA (read-depth-aware) 建模的预训练任务,这是掩码语言模型的扩展,考虑了大规模数据中测序深度的高方差。


预训练过程概要


在 RDA 建模中,模型根据细胞的上下文基因预测被掩码的基因表达。研究人员将总计数视为一个细胞的测序深度,并定义了两个总计数指标:T(目标)和 S(来源),分别表示原始样本和输入样本的总计数。研究人员随机掩码输入样本中的零表达和非零表达基因,并记录它们的索引。


然后,模型使用掩码输入样本和两个指标来预测原始样本在掩码索引处的表达值。这使得预训练模型不仅能够捕捉细胞内的基因关系,还能协调不同测序深度的细胞。在推断过程中,研究人员将细胞的原始基因表达输入预训练模型,并设置 T 高于其总计数 S,以生成增强测序深度的基因表达值。


简单来说,RDA 能够对测序深度进行降采样,使得模型在预训练阶段除了完成传统的掩码恢复任务外,还能够由低质量细胞恢复高质量细胞的基因表达信息。



预训练为下游任务提供便利


预训练完成后,研究人员进一步将 scFoundation 模型应用于多个下游任务。scFoundation 编码器的输出被汇总为细胞级向量,用于细胞级任务,包括 (数据集内和跨数据集的) 聚类、批量和单细胞水平的药物反应预测以及细胞类型注释。scFoundation 解码器的输出是基因级上下文向量,用于基因级任务,如扰动预测和基因模块推断。


研究结果:scFoundation 模型具备卓越性能


在实际应用中,scFoundation 模型支持「开箱即用」和「微调」两种模式。在「开箱即用」模式下,得益于其独特的预训练任务,该模型能直接用于提升细胞数据的质量,在不需要进一步调整的情况下便可达到或超越现有方法的效果。此外,用户可以利用 scFoundation 提取细胞的预训练表征,该表征可以用于识别细胞类型特异基因模块和转录因子,并可广泛应用于下游任务中。


可扩展的无微调测序深度增强模型


研究人员分别训练了参数量为 3M、10M 和 100M 的三个模型,并记录了它们在验证数据集上的损失。



不同参数大小和 FLOPs 下的训练损失
绿色曲线代表 3M 模型,橙色曲线代表 10M 模型,红色曲线代表 100M 模型


随着模型参数和浮点运算次数 (FLOPs) 的增加,验证数据集上的损失呈现出幂律下降。然后,研究人员估算了各种规模的 xTrimoGene 架构模型的性能,并与 scVI 进行比较。如上图所示,具有 1 亿参数的 scFoundation 模型在所有模型中表现最佳。研究人员进一步在细胞类型注释任务上评估了 3 个模型,观察到随着模型规模的增加,性能有所提升。


研究人员在从验证数据集中随机抽取的 1 万个细胞的独立测试数据上评估了这种能力,其将总计数降采样到原始数据的 1%、5%、10% 和 20%,生成了 4 个具有不同总计数变化的数据集。对于每个数据集,其使用未微调的 scFoundation,测量了预测值与实际非零基因表达之间的平均绝对误差 (MAE)、平均相对误差 (MRE) 和皮尔逊相关系数 (PCC)。


对未知数据集上读深度增强性能的评估
使用 MRE 和 PCC 来评估恢复的基因表达性能,较低的 MRE 和较高的 PCC 表明性能更佳


如上图所示,即使在降采样率低于 10% 的情况下,scFoundation 的 MAE 和 MRE 也显著减少了一半。这些结果展示了 scFoundation 在总计数极低的情况下增强基因表达的能力。


下游任务——癌症药物反应预测任务


癌症药物反应 (CDRs) 旨在研究肿瘤细胞在药物干预下的反应,计算预测 CDR 对于指导抗癌药物设计和理解癌症生物学至关重要。在本研究中,研究人员将 scFoundation 与 CDR 预测方法 DeepCDR 结合,预测多个细胞系数据中药物的半数抑制浓度 IC50 值,以验证 scFoundation 能否为基于单细胞训练的整体基因表达数据提供有用的嵌入信息。


研究人员评估了基于 scFoundation 的结果与基于基因表达的结果在多个药物和癌症细胞系上的表现,结果显示大多数药物和所有癌症类型使用 scFoundation 嵌入都取得了更高的皮尔逊相关系数 (PCC),如下图所示:



使用 scFoundation 嵌入进行药物反应预测


* 注:皮尔逊相关系数是一种衡量变量之间线性关系强弱的统计量,取值范围在 -1 到 1 之间。相关系数接近 1,表明两个变量之间存在完全正向的线性关系;接近 -1,说明存在完全负向的线性关系;接近 0,表示两个变量之间没有线性关系。


这说明,尽管 scFoundation 在单细胞转录组数据上进行了预训练,但学习到的基因关系可以转移到整体水平的表达数据上,生成压缩的向量,促进更准确的 IC50 预测。因此 scFoundation 在扩展对癌症生物学中药物反应的理解,以及指导更有效的抗癌治疗设计中具备一定潜力。


下游任务——单细胞药物反应分类任务


在单细胞水平上推断药物敏感性有助于识别表现出不同药物耐药特征的特定细胞亚型,从而提供有关潜在机制和新疗法的宝贵见解。因此,研究人员将 scFoundation 应用于单细胞药物反应分类这一关键任务,基于一个称为 SCAD 的下游模型。


研究人员重点研究了四种药物 (索拉非尼、NVP-TAE684、PLX4720 和足叶乙甙),它们在原始研究中表现出较低 AUC (Area Under Curve) 值。其将基于 scFoundation 的模型与以所有基因表达值作为输入的基线 SCAD 模型进行了比较,结果显示,基于 scFoundation 的模型在所有药物的 AUC 值上,都取得了更高的分数,对于 NVP-TAE684 和索拉非尼尤为显著,AUC 值增加了超过 0.2,如下图所示。


基于 scFoundation 细胞嵌入的单细胞药物反应分类任务

注:AUC 可以用来衡量模型的性能,AUC 的取值范围是 0 到 1,值越大代表模型分类性能越好。


这些结果验证了使用 scFoundation 嵌入具有捕捉药物敏感性生物标志物信号的潜力。


下游任务——细胞扰动预测任务


理解细胞对扰动的响应对于生物医学应用和药物设计至关重要,有助于识别不同细胞类型间的基因-基因相互作用和潜在的药物靶点。研究人员将 scFoundation 与一个高级模型 GEARS 结合起来,以预测单细胞分辨率的扰动响应,并计算了前后基因表达轮廓差异显著的前 20 个差异表达 (DE) 基因的平均均方误差 (MSE),作为评估标准。


结果显示,与原始 GEARS 基线模型相比,基于 scFoundation 的模型取得了更低的 MSE 值。下图展示了双基因扰动 ETS2 + CEBPE 的前 20 个基因表达变化:



组合扰动 (ETS2 + CEBPE) 后,前 20 个 DE 基因的预测基因表达量高于对照


这些结果表明,通过提取单个细胞的基因表征来构建特定的基因共表达网络,scFoundation 成功捕捉了不同条件下的细胞和基因表征,显著提高了单/双扰动预测的准确度。


综上所述,scFoundation 模型为建立细胞预训练大模型的模型架构、训练框架,和下游示范应用体系都提供了新的思路和方法,为生物医学任务的学习提供了基础功能,拓展了单细胞领域基础模型的边界。


探究性能更优的生命科学大模型


全球领先的人工智能企业 DeepMind 首席执行官兼创始人 Demis Hassabis 曾经表示:「从最基本的层面来看,生物学可以被看作是一个非常复杂且动态的信息处理系统。就像数学被证明是物理学的正确描述语言一样,生物学可能会成为完美的人工智能应用领域。」


然而,传统的人工智能方法需要大量的标记数据才能作出准确的预测。但在生命科学领域,高质量的标记数据往往供不应求。想要在更少数据的基础上构建更准确的下游任务模型,就意味着底层基础模型需要具有更好的表征或通用能力。因此,越来越多研究人员开始致力于设计更优的生物领域垂直大模型。


2023 年 5 月,多伦多大学研究团队发布首个单细胞生物学基础大型语言模型 scGPT,其在超 1 千万个细胞上进行预训练,该模型可实现跨各种下游任务的迁移学习。同年 7 月,该团队进一步尝试对超过 33,00 万个细胞进行生成预训练来更新 scGPT,结果表明,scGPT 可以有效地提炼有关基因和细胞的关键生物学见解,并在各种下游任务中实现先进的性能,包括 multi-batch 整合、多组学整合、细胞类型注释、遗传扰动预测和基因网络推断等。


该研究以「scGPT: toward building a foundation model for single-cell multi-omics using generative AI」为题,发表于 Nature Methods 上。
* 论文链接:https://www.nature.com/articles/s41592-024-02201-0


2023 年 9 月,由中国科学院多学科交叉研究团队组成的「指南针联盟」 (Xcompass Consortium) 成功构建了世界首个跨物种生命基础大模型——GeneCompass。该模型集成了人和小鼠超过 1.26 亿个单细胞的转录组数据、融合了包括启动子序列和基因共表达关系等四种先验知识、基础模型参数量达到 1.3 亿,实现了对基因表达调控规律的全景式学习理解,同时支持细胞状态变化预测及多种生命过程的精准分析。


该研究以「GeneCompass: Deciphering Universal Gene Regulatory Mechanisms with Knowledge-Informed Cross-Species Foundation Mode」为题,发布在 bioRxiv 上。
* 论文链接:https://www.biorxiv.org/content/10.1101/2023.09.26.559542v1


2023 年 10 月,全球制药巨头赛诺菲宣布与 BioMap 百图生科达成大型战略合作,双方将基于百图生科的生命科学大模型 (Life Science AI Foundation Model),共同开发用于生物治疗药物发现的前沿模型。


展望未来,将大语言模型所拥有远超人类想象的复杂理解能力和创新生成能力,应用在更加复杂的生命 「自然语言」上,将有望真正改变生命科学的研究范式。


参考资料:
1.https://www.jiqizhixin.com/articles/2023-9-29
2.https://www.tsinghua.edu.cn/info/1175/112118.htm
3.https://hope.huanqiu.com/article/4FYZxnpu88J
4.https://www.jiqizhixin.com/articles/2023-7-5-26


编辑:于腾凯

校对:林亦霖

浏览 61
点赞
评论
收藏
分享

手机扫一扫分享

举报
评论
图片
表情
推荐
点赞
评论
收藏
分享

手机扫一扫分享

举报