用Diffusion构建「AI虚拟细胞」，14项指标霸榜！Mila唐建团队破解单细胞「破坏性」测序难题

PerturbDiff团队 投稿 
量子位 | 公众号 QbitAI

如果你想用AI预测一款抗癌新药对某种癌细胞的影响，最符合直觉的训练方式是什么？

你的第一直觉可能是——给AI看同一个细胞“吃药前”和“吃药后”的基因表达状态，让它学习其中的变换规律。

然而，单细胞转录组测序（scRNA-seq）的“破坏性”为这一愿景设置了一道物理屏障：单细胞测序（scRNA-seq）是破坏性的。一旦你观测了细胞的状态，这个细胞就被溶解了。你无法在施加微扰（Perturbation，如药物干预、基因编辑）前后观测到同一个细胞。

面对这种本质上非配对（Unpaired）的数据，传统的AI模型往往败下阵来。要么强行随机配对导致学出平庸的“平均脸”，要么在预测靶点基因时频频“脱靶”。

近日，此前已在蛋白性质预测与设计、AI虚拟细胞领域发表多篇重磅论文的Mila（魁北克人工智能研究所）唐建团队发表了一篇新作PerturbDiff。他们跳出了前人纠结于“单细胞”还是“单分布”的认知盲区，首次将“细胞群体的概率分布”视作一个可以进行加噪、去噪的随机变量。

该模型在包含上亿细胞的全球最大药物微扰图谱（Tahoe100M）等多项权威基准测试中，以碾压之势刷新了单细胞响应预测的SOTA记录，完成了AI在构建高保真“虚拟细胞（Virtual Cells）”道路上的一次重大突破。

认知陷阱：微扰响应真的是一个“静态分布”吗？

如果要理解PerturbDiff为什么能赢，首先要明白此前的SOTA错在哪。

近年来进步最快的方法（如STATE、CellFlow等）已经意识到单细胞强行配对、并进行回归训练是不靠谱的，于是它们开始做“群体对群体”的映射。但在这些算法底层的世界观里，存在一个致命的静态假设：给定一种细胞和一种药物，它微扰后的细胞分布结果是唯一且固定的。

唐建团队敏锐地指出，生命系统从来不是精密的钟表。培养皿里的微环境波动、细胞周期的细微差异、甚至是实验的批次效应，这些“看不见的隐变量”会导致同一组实验在不同时间做，最终的细胞分布态都有微妙的偏移。

如果你逼着AI去死记硬背一个“唯一标准答案”，它学到的只是所有可能性的平均值。一旦面临真正的新药或跨批次测试，泛化能力就会雪崩。

PerturbDiff的破局点在于完成了一次极其深刻的概念“升维”：它不再把AI的建模/生成的目标设定为单个“细胞”，甚至不是“某一个细胞分布”，而是“细胞分布的分布”——将整个微扰后的细胞种群视为一个在分布空间中取值的“随机变量”。既然生物微扰的结果是混沌且多变的，那就直接用生成式模型去拟合这种多变性。

泛函空间上的扩散：如何向一个“分布”加噪？

确立了非静态分布的建模思想，真正的工程灾难才刚刚开始。

如何用深度学习网络对一个无形无相的“概率分布族”进行生成式建模？

PerturbDiff引入了一个极具美感的数学工具：再生核希尔伯特空间（RKHS，𝓗ₖ）与核均值嵌入（Kernel Mean Embedding, KME）。

在这套空间里，一个包含了几万只小白鼠（单细胞）的复杂种群，被高维压缩成了一个“点”。

在传统扩散模型（如Sora或是Stable Diffusion）中，AI是在有限的像素矩阵里注入高斯噪声，像捏泥人一样从噪点中还原出图像。

而PerturbDiff是直接在这个极其高维的函数空间里，对代表“整个细胞种群分布”的那个点，进行连续的马尔可夫扩散演化（即泛函扩散，Functional Diffusion）。

为了解决高维空间无法直接计算的问题，作者通过严密的数学证明（一阶线性泰勒展开）指出：在真实世界的单细胞表达矩阵里，给每一个单细胞独立注入欧氏噪声，在数学上刚好完美等价于那个高维函数空间里的高斯随机游走。

这一证明在理论与工程之间架起了一座桥梁，让极其抽象的泛函扩散过程得以通过最基础的张量加噪来实现。

MMD的内生性：从“正则化惩罚”到“底层原理解析”

在探讨损失函数时，PerturbDiff再次展现了其理论体系的自洽。

近年来，最大均值差异（MMD）常被各类单细胞对齐模型（如STATE）用作损失函数。但在过去的工作中，MMD更多只是一种“外挂”的启发式正则化惩罚项（heuristic regularizer），用于拉近组间距离。而在PerturbDiff的框架下，MMD是内生且不可避免的。

根据作者在推导RKHS空间中的反向去噪目标（Denoising Objective）时，其对数似然的变分下界自然导出了对于分布点μ的距离度量：

由于RKHS内积与经验核对齐的固有几何性质，这一项在数学上绝对等价于原单细胞空间中真实分布与预测分布之间的MMD的平方。

这解释了为什么PerturbDiff消融实验中，去掉MMD仅保留MSE会导致模型崩溃。

单细胞基因表达高度稀疏（Zero-inflation常高达95%以上），单细胞级别的MSE会让模型陷入“预测全0”的次优解，而内生于泛函扩散体系的MMD则迫使模型在总体分布密度、子群权重等高阶统计特征上进行对齐。这也使其在预测“差异表达基因（DEGs）”——这一生物学界评判微扰模型的最核心指标（如AUPRC, AUROC）上，表现出了压倒性的优势。

霸榜背后的生物学先验：6100万细胞“全底图”催生零样本奇迹

理论的优美最终转化为在实验榜单上的统治力。

在包含超1亿单细胞、1100种药物的Tahoe100M（全球最大药物响应基准）和PBMC（免疫信号基准）中，PerturbDiff在14项评估指标中实现了全面霸榜。

更令生物学家兴奋的是，它在差异表达基因（DE-metrics，如AUPRC、AUROC）上的超高精度。预测一群细胞长什么样只是表面功夫，新药研发真正关切的是：“这款药到底开启或关闭了哪些关键基因？”在这个核心痛点上，PerturbDiff因为真正学到了微扰的“系统性偏移”，其表现远超以往所有模型。

在模型架构之外，PerturbDiff的训练范式为应对生物医药数据特有的“标注稀缺（Low-Data Regime）”难题提供了一个富有洞见的思路：边缘分布预训练（Marginal Pretraining）。

微扰实验（尤其是CRISPR筛选或大规模药物测试）是昂贵的，覆盖的细胞类型极其有限。作者没有局限于有监督的微扰数据，而是利用CellxGene数据集中6100万无干预（Unperturbed）的单细胞转录组数据，对其扩散模型进行了无条件流形的预训练。

这种看似与特定微扰任务无关的预训练，不仅大幅提升了模型在极低样本量（如Replogle数据集，每种微扰仅数百个细胞）下的微调性能，更在预训练检查点上观察到了显著的零样本（Zero-shot）预测能力。这其中蕴含着深刻的生物学洞察：生物微扰并不是将细胞推向基因表达空间中任意的混沌区域，微扰诱导的转移轨迹，实际上部分对齐/嵌套于自然界本身已存在的细胞状态流形之中。

6100万细胞图谱绘制的正是这张宏大的“全底图”，理解了底图的地形，寻找微扰发生的轨迹自然事半功倍。

总结

PerturbDiff不仅在权威基准（PBMC，Tahoe100M）的十余项泛化指标上刷新了SOTA，更重要的是，它促使我们重新审视计算表型建模的理论根基。

如同卷积神经网络（CNN）之于图像识别，Transformer之于文本生成，AlphaFold Evoformer之于蛋白结构预测，符合数据内在规律的归纳偏置（inductive bias）才能真正发挥数据的全部价值。PerturbDiff通过在“细胞分布的分布”这一泛函空间训练扩散模型，朝着这个方向迈进了一小步。

而随着它的到来，那个可以无限次试错、能够精准模拟扰动响应的终极“AI虚拟细胞”，或许离我们又近了一大步。

项目主页：
https://katarinayuan.github.io/PerturbDiff-ProjectPage/