Lingshu-Cell:面向虚拟细胞的生成式细胞世界模型用于转录组建模 Lingshu-Cell: A generative cellular world model for transcriptome modeling toward virtual cells
基于掩码离散扩散模型的单细胞转录组生成框架,实现细胞状态模拟和扰动响应预测
前置知识
单细胞RNA测序(scRNA-seq)
单细胞RNA测序是一种高通量技术,能够在单个细胞水平上测量基因表达。通过将每个细胞分离出来,捕获其mRNA并测序,得到细胞×基因的计数矩阵。该技术具有稀疏性(大多数基因不表达)、离散计数性质(UMI计数)、以及非顺序性(基因之间没有固有顺序)。典型的scRNA-seq数据集包含数千到数百万个细胞,每个细胞测量约18,000个基因的表达水平。
本文的核心任务是建模scRNA-seq数据,理解其稀疏、离散、非顺序的特性对于理解为什么选择掩码离散扩散模型而非传统的自回归模型或连续扩散模型至关重要。
掩码离散扩散模型(MDDM)
掩码离散扩散模型是扩散模型在离散数据上的变体。在前向过程中,它逐步将序列中的token替换为掩码token[MASK],掩码概率随时间t从0增加到1。在反向过程中,模型学习预测被掩码位置的原始token。与自回归模型不同,MDDM不需要固定的生成顺序;与连续扩散模型(如DDPM)不同,它不使用连续噪声而是直接操作离散token。训练目标是交叉熵损失:$L(\theta) = -\mathbb{E}_{t,x_0,x_t}\sum_{i=1}^L \mathbb{I}[x_i^t = M] \log p_\theta(x_i^0|x^t)$。
MDDM是Lingshu-Cell的核心架构,它天然契合scRNA-seq数据的稀疏离散特性,避免了自回归模型的基因顺序偏差和连续扩散模型与计数数据的分布不匹配问题。
分类器无关引导(Classifier-Free Guidance, CFG)
分类器无关引导是一种增强条件生成的技术,通过联合训练条件模型和无条件模型,在采样时将条件预测和无条件预测结合起来。给定条件c和控制条件cnt(如非靶向对照),引导后的logits计算为:$\tilde{a}_\theta(v|x^t, c) = a_\theta(v|x^t, cnt) + (w + 1)(a_\theta(v|x^t, c) - a_\theta(v|x^t, cnt))$,其中w是引导强度。引导权值越大,生成越偏向条件特定模式。
CFG是Lingshu-Cell在扰动预测任务中的关键技术,它能够增强扰动特异性信号,提高预测的准确性,特别是在扰动引起的转录变化相对背景变异较微弱的情况下。
研究动机
现有的单细胞转录组基础模型(如scGPT、Geneformer、scFoundation、CellFM)主要学习静态表示,而不是生成式模拟。这些模型虽然在表示学习方面取得成功,但无法生成新的细胞状态或模拟细胞对扰动的响应。现有的生成式方法(如scDiffusion、scVI)受限于连续数据假设,与scRNA-seq数据的稀疏、离散、非顺序特性不匹配。专注于扰动的方法(如STATE、CellFlow、scDFM、AlphaCell)学习从对照状态到扰动结果的直接映射,但不建模潜在的转录组状态分布或其条件动力学。这些限制使得现有方法无法实现高保真的细胞状态生成和准确的扰动响应预测,阻碍了虚拟细胞的发展。
本文的目标是本文的目标是开发一个能够学习转录组状态分布并支持在扰动条件下进行条件模拟的细胞世界模型。这个模型应该直接在离散token空间中操作,与scRNA-seq数据的稀疏、非顺序特性相匹配,不需要预先的基因选择(如高变基因过滤或表达水平排序),能够捕捉跨约18,000个基因的复杂转录组范围内的表达依赖关系,并且能够在多种组织和物种上准确重现转录组分布、标记基因表达模式和细胞亚型比例。
与已有工作不同的是,本文的独特切入角度是将细胞世界模型的概念形式化,并采用掩码离散扩散模型(MDDM)作为统一的生成框架。与现有方法不同,Lingshu-Cell不假设基因有固定的生成顺序(避免了自回归模型的偏差),不使用全局连续噪声破坏(避免了与离散计数数据的分布不匹配),而是随机掩码和预测基因表达值,以一种与基因表达数据的无序结构本质兼容的方式建模。此外,本文还引入了序列压缩、分类器无关引导和生物先验注入等技术创新,使得单个架构能够同时支持无条件生成和条件扰动预测,这是现有方法无法实现的功能组合。
核心方法
Lingshu-Cell的整体思路是将单细胞转录组数据建模为离散token序列,然后使用掩码离散扩散模型学习数据的联合分布。在前向过程中,基因表达值从原始观察状态(t = 0)逐步被掩码直到完全掩码状态(t = T)。在反向过程中,模型迭代预测被掩码的基因表达值,最终生成生物学上真实的scRNA-seq表达谱。这个掩码-预测范式使模型能够学习复杂的基因调控依赖关系,同时自然地适应基因表达谱的无序结构。模型通过将细胞类型或供体身份与扰动上下文嵌入到联合潜空间中,支持条件生成以预测全转录组表达变化。这种设计使得Lingshu-Cell能够在单个架构中实现无条件细胞生成和条件扰动预测两种功能。
Lingshu-Cell的核心创新点在于将掩码离散扩散模型(MDDM)应用于单细胞转录组数据,并引入了三个关键技术组件:序列压缩、分类器无关引导(CFG)和生物先验注入。序列压缩模块在嵌入空间中将高维基因表达序列转换为信息密度更高的更短序列,提高了建模效率同时有助于捕捉全局表达模式。分类器无关引导通过在采样时结合条件预测和无条件预测,将生成引向与扰动条件更一致的转录组状态。生物先验注入在推理时使用来自其他细胞系的扰动概要信息,为生成过程提供生物学先验信号。这些技术与MDDM架构相结合,使得Lingshu-Cell能够同时实现高保真的细胞生成和准确的扰动响应预测。
方法步骤详情
Lingshu-Cell的方法步骤包括数据预处理、离散序列表示、模型训练和推理生成。首先,将UMI计数矩阵进行量化,将每个基因的计数映射到有限的表达级别集合,量化函数为$q(x)$,将[0, 99]范围内的值保留100个bin,更高值按对数尺度分配90个bin每十年数级,超过阈值C=9999的值映射到溢出token OV。然后,每个细胞被表示为离散序列$z^{(i)} = (q(x^{(i)}_1), q(x^{(i)}_2), ..., q(x^{(i)}_G))$,作为MDDM的输入/输出空间。对于条件生成,细胞类型和扰动目标作为额外的条件token前缀到表达序列中,这些条件token在掩码过程中保持不变。模型训练使用交叉熵损失:$L(\theta) = -\mathbb{E}_{t,x_0,x_t}\sum_{i=1}^L \mathbb{I}[x_i^t = M] \log p_\theta(x_i^0|x^t)$。推理时,对于条件生成,使用分类器无关引导,引导后的logits为$\tilde{a}_\theta(v|x^t, c) = a_\theta(v|x^t, cnt) + (w + 1)(a_\theta(v|x^t, c) - a_\theta(v|x^t, cnt))$。对于生物先验注入,将先验基因集$G^\downarrow$初始化为低表达值$\mu=1$,其他位置保持掩码,采样过程中保持先验位置固定。序列压缩模块通过随机置换、分组和线性投影将18,080个基因的序列压缩到更短的序列(S=8时压缩到2,260,S=32时压缩到565),经过Transformer处理后反向扩展到原始长度。
技术新颖性
Lingshu-Cell的技术新颖性体现在多个方面。首先,它是第一个将掩码离散扩散模型成功应用于单细胞转录组数据的工作,开创了一种新的生成范式,与传统的自回归模型和连续扩散模型有本质区别。其次,引入的序列压缩模块在嵌入空间中操作,不改变token级别的建模目标,既提高了计算效率又通过多基因信号的简单线性混合减少了单基因噪声的影响。第三,分类器无关引导的应用使得模型能够在不显式训练分类器的情况下增强条件生成,这是在扰动预测任务中的创新应用。第四,生物先验注入策略利用外部细胞系的扰动概要信息,在不假设直接定量跨系统可转移性的情况下为生成过程提供方向性信号。最后,Lingshu-Cell在单个架构中同时实现了无条件生成和条件扰动预测,这是现有方法无法实现的功能整合,展示了MDDM作为统一生成框架的潜力。
实验结果
Lingshu-Cell在多个实验中取得了领先的性能。在无条件生成方面,模型在PARSE 10M PBMC数据集上训练后,生成的10,000个细胞忠实再现了五个主要PBMC谱系(T细胞、NK细胞、B细胞、单核细胞和树突状细胞)的标记基因表达模式,生成的细胞类型比例与真实数据高度一致。当扩展到200,000个细胞时,标记基因表达模式和细胞类型比例仍然与真实数据高度一致,在17个亚型的更高分辨率注释下,生成和真实数据继续紧密对齐。在五个评估指标(Pearson相关性、Spearman相关性、MMD、1-WD和iLISI)上,Lingshu-Cell取得了最低的MMD(0.0088,而scDiffusion为0.0178,scVI为0.0343),表明生成和真实表达分布之间的总体匹配最接近。在跨越九个组织和五个物种的大规模数据集上,Lingshu-Cell一致地产生了高质量样本,准确捕获了每个组织中的主要细胞类型以及组织特异性细胞类型。在基因扰动预测方面,Lingshu-Cell在Virtual Cell Challenge H1基准测试中取得了最佳平均排名(8.7),在七个评估指标中获得了最低的MAE(0.052)和最高的Pearson-$\Delta$相关性(0.306)。在细胞因子扰动预测方面,Lingshu-Cell在PARSE 10M PBMC数据集上取得了最高的平均分数,在PDS和Pearson-$\Delta$相关性方面排名第一,表明预测的转录组保留了每个细胞因子条件的独特身份,并准确捕获了细胞因子诱导的表达变化的方向和幅度。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 无条件生成(PARSE-PBMC) | MMD | 0.0088 | scDiffusion: 0.0178, scVI: 0.0343 | 相比scDiffusion降低51%,相比scVI降低74% |
| 无条件生成(PARSE-PBMC) | 1-WD | 0.0102 | scDiffusion: 0.0064, scVI: 0.0178 | 相比scVI降低43% |
| 基因扰动预测(VCC H1) | 平均排名 | 8.7 | cleopatra: 9.1, xBio: 10.7 | 相比最佳基线提升4% |
| 基因扰动预测(VCC H1) | MAE | 0.052 | cleopatra: 0.086, xBio: 0.770 | 相比最佳基线降低40% |
| 基因扰动预测(VCC H1) | Pearson-$\Delta$ | 0.306 | cleopatra: 0.203, xBio: 0.217 | 相比最佳基线提升51% |
| 细胞因子扰动预测(PARSE 10M PBMC) | 平均分数 | 领先 | PertMean, STATE, scGPT, scVI | 在所有比较方法中排名第一 |
局限与改进
作者承认的局限性包括:当前评估依赖于群体水平的分布指标(如MMD、iLISI)和伪批量相关性,这些指标不能完全评估单细胞生物学合理性或极其罕见状态的保留。更根本的是,高保真生成并不意味着生物因果性:忠实地重表达分布并不一定反映产生它们的因果调控机制,因此Lingshu-Cell目前应该被视为强大的概率假设生成工具,其预测仍然需要严格的湿实验室实验验证。此外,当前模型仅在转录组数据上操作,更完整的虚拟细胞还将整合表观基因组、蛋白质组、代谢组和空间模态。自己观察到的局限性包括:模型依赖于预定义的基因列表(18,080个基因),无法处理新基因或未知基因;量化过程可能会丢失某些表达水平的详细信息;生物先验注入依赖于外部细胞系数据,在某些扰动目标上可能不可用;模型仅建模静态的转录组状态,无法捕捉动态的细胞轨迹或时间依赖的扰动响应。
独立分析的弱点
Lingshu-Cell的第一个弱点是在稀有细胞状态建模方面可能存在不足。由于评估指标主要关注群体水平的分布相似性,模型可能无法准确生成极其罕见的细胞亚型或状态,这些状态在训练数据中出现频率很低。改进方向可以引入专门针对稀有细胞的采样策略或损失函数权重调整。第二个弱点是缺乏因果建模能力。虽然模型能够高保真地重表达分布,但不捕获基因之间的因果关系或调控网络,这限制了对机制性生物问题的应用。改进方向可以结合因果推断技术或显式建模基因调控网络。第三个弱点是生物先验注入对外部数据的依赖。对于某些扰动目标,可能没有可用的外部细胞系数据,此时生物先验注入策略无法应用。改进方向可以开发基于基因功能或通路信息的替代先验,或使用少样本学习技术从有限的外部数据中提取先验。第四个弱点是量化过程中的信息损失。虽然量化保留了约前两位有效数字,但对于极高表达值(>9999)的信息可能会丢失。改进方向可以设计更精细的量化策略,或使用基于学习的嵌入替代固定量化。
未来方向
作者提出的未来研究方向包括:将当前框架从单基因扰动扩展到更复杂的干预,包括药物诱导、多靶点和组合干预,其中剂量依赖性和时间动态可能发挥关键作用。更广泛地,设想在同一离散扩散框架内联合建模染色质可及性、蛋白质丰度和空间分辨转录组学,并融入时间结构以模拟动态细胞轨迹如分化和疾病进展。最终,最具前景的应用之一可能是闭环实验,其中模型预测指导靶向扰动,新生成的数据迭代地细化模型本身。这种框架将从静态数据拟合转向生物发现的自适应平台。基于成果可延伸的方向包括:开发更复杂的条件机制,支持多因素条件(如细胞类型+扰动+时间+药物浓度);探索更高效的条件引导技术,如层次化引导或多尺度引导;整合多模态生物学数据,如表观遗传学、蛋白质组学和代谢组学,以构建更全面的虚拟细胞模型;开发实时交互式界面,使研究人员能够交互式地探索和设计扰动实验;将模型应用于个性化医疗,根据患者特异性转录组数据预测个体化的治疗响应。
复现评估
关于复现评估,论文提供了详细的实现细节和数据来源,使得复现相对可行。模型架构和训练超参数都有详细描述,包括Transformer层数(L=13)、注意力头数(nh=10)、嵌入维度(D=640)、压缩组大小(S=8或32)等。训练配置包括优化器(AdamW)、学习率调度(余弦退火)、批量大小(256)、训练数据集(PARSE PBS、CZ CELLxGENE、VCC H1+外部数据等)。然而,论文没有明确说明代码是否开源,这是复现性的一个关键因素。数据方面,使用的数据集都是公开可用的(PARSE 10M、CZ CELLxGENE、scBaseCount、Virtual Cell Challenge等),但需要遵循各自的数据访问协议。算力方面,训练需要多个GPU(根据任务不同使用2-16个NVIDIA A800 GPU),对于资源有限的研究人员来说可能有挑战。总体而言,如果论文作者开源代码并提供预训练模型,复现实验应该是可行的,但完全从头训练模型可能需要大量计算资源。
论文图表