PerturbDiff：AI预测抗癌新药效果，首次将细胞群体分布视为随机变量，刷新单细胞响应预测SOTA

PerturbDiff：将细胞群体分布视为随机变量，刷新单细胞响应预测SOTA

如何利用AI预测一款抗癌新药对特定癌细胞的影响？一个直观的思路是让AI学习同一个细胞在药物处理前后的基因表达变化规律。

然而，单细胞转录组测序（scRNA-seq）的“破坏性”构成了根本障碍：一旦对细胞状态进行观测，该细胞即被溶解。因此，我们无法在施加微扰（如药物干预、基因编辑）前后观测到同一个细胞。

面对这种本质上的非配对数据，传统AI模型往往表现不佳。它们要么强行随机配对，导致模型学到平庸的“平均”响应；要么在预测关键靶点基因时频频“脱靶”。

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

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

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

要理解PerturbDiff的成功，首先需厘清此前SOTA方法的局限。

近年来进展较快的方法（如STATE、CellFlow等）已意识到对单细胞进行强行配对和回归训练并不可靠，转而尝试进行“群体到群体”的映射。但这些方法的底层世界观存在一个关键的静态假设：给定一种细胞类型和一种药物，其微扰后的细胞分布结果是唯一且固定的。

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

如果迫使AI去学习一个“唯一标准答案”，它学到的只是所有可能性的平均值。一旦面临真正的新药或跨批次测试，其泛化能力便会急剧下降。

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

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

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

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

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

在这套框架下，一个包含数万个细胞的复杂群体，被映射为高维RKHS空间中的一个“点”。

在传统扩散模型（如图像生成的Sora或Stable Diffusion）中，AI是在有限的像素矩阵上添加高斯噪声，然后从噪声中逐步重建图像。

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

为了解决高维空间计算难题，作者通过严密的数学证明（基于一阶线性泰勒展开）指出：在真实世界的单细胞表达矩阵中，对每一个单细胞独立添加欧氏空间的高斯噪声，在数学上恰好等价于在RKHS空间中对分布点进行高斯随机游走。

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

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

在损失函数的设计上，PerturbDiff再次展现了其理论体系的自洽性。

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

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

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

这解释了为何在消融实验中，移除MMD损失仅保留均方误差（MSE）会导致模型崩溃。单细胞基因表达数据具有高度稀疏性（零值比例常高达95%以上），仅使用细胞级别的MSE会驱使模型陷入“预测全零”的次优解。而内生于泛函扩散体系的MMD损失，则迫使模型在总体分布密度、子群权重等高阶统计特征上进行对齐。这也使得PerturbDiff在预测差异表达基因（DEGs）——这一生物学界评估微扰模型的最核心指标（如AUPRC, AUROC）上，表现出了压倒性优势。

霸榜背后的生物学先验：6100万细胞“全景图”催生零样本能力

理论的优美最终转化为在实验基准上的卓越性能。

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

更令生物学家兴奋的是，它在差异表达基因相关指标上的超高精度。预测细胞群体的形态变化只是表面，新药研发真正关切的是：“这款药究竟激活或抑制了哪些关键基因通路？”在这个核心痛点上，由于PerturbDiff真正学到了微扰引起的“系统性分布偏移”，其表现远超以往所有模型。

除了模型架构，PerturbDiff的训练范式也为应对生物医药数据固有的“标注稀缺”难题提供了富有洞见的思路：边缘分布预训练。

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

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

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

总结

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

如同卷积神经网络（CNN）之于图像识别，Transformer之于文本生成，AlphaFold Evoformer之于蛋白结构预测，符合数据内在规律的归纳偏置（inductive bias）才能真正释放数据的潜力。PerturbDiff通过在“细胞群体分布的分布”这一泛函空间中训练扩散模型，正是朝着这一方向迈出的重要一步。

随着这项研究的进展，那个能够无限次试错、精准模拟药物扰动响应的终极“AI虚拟细胞”，或许离现实又近了一大步。

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

关注“鲸栖”小程序，掌握最新AI资讯