MapPFN: Learning Causal Perturbation Maps in Context¶
会议: NeurIPS 2026
arXiv: 2601.21092
OpenReview: https://openreview.net/forum?id=skdhOj5NkC
代码: https://marvinsxtr.github.io/MapPFN
领域: 因果推理
关键词: 单细胞微扰预测, 上下文学习, 先验数据拟合网络 (PFN), 因果干预, 条件流匹配
一句话总结¶
MapPFN 是首个用于单细胞因果微扰预测的先验数据拟合网络，通过在合成生物学因果先验上元学习观测与干预分布映射，在测试期仅凭序列化干预上下文即可零样本自适应全新生物学背景与任意基因集。
研究背景与动机¶
通过单细胞微扰测序（如 Perturb-seq）研究基因敲除后细胞群体的转录组响应，是破译基因调控网络（GRN）和发现药物靶点的金标准实验手段。然而，测序技术的破坏性本质导致实验无法观测同一个细胞在干预前后的成对状态，使得干预效果估计本质上成为了未配对分布之间的最优传输或生成建模问题。现有的虚拟细胞模型与微扰预测方法（如 CPA、CellOT、CellFlow、Meta Flow Matching 和 STATE）高度依赖真实湿实验数据训练，但现实中即便是规模最大的公开微扰数据集也仅覆盖数十种细胞系，导致模型难以泛化到未见过的生物学背景中。
更关键的矛盾在于，同一基因在不同细胞系中因调控机制的异质性会产生截然不同的干预响应。现有方法在推理阶段仅能接收未干预的对照组分布或类别协变量，完全缺乏利用测试时新干预实验证据进行即时上下文适应的能力；同时，现有架构被硬性绑定在固定的基因输入空间上，面对新基因集必须从头重新训练。
本文的切入角度是打破对湿实验数据的直接依赖，将基于合成先验元学习贝叶斯推断的先验数据拟合网络（PFN）从一维标量推断扩展到高维细胞分布映射，并通过上下文学习（In-Context Learning）在推理期无缝吸收干预线索。核心 idea：将微扰预测建模为包含干预上下文的分布映射问题，构建基于 Hill 非线性微分方程的合成生物学因果先验，利用多模态扩散 Transformer 在合成基因敲除数据上元学习后验预测分布，实现无需重训即可单步自适应未见生物学背景与任意基因集。
方法详解¶
整体框架¶
MapPFN 的输入包含目标生物学背景下的未干预观测细胞样本 \(\mathbf{Y}^{\text{obs}}\in\mathbb{R}^{N\times d}\)、包含 \(K\) 组已知微扰实验的干预上下文集合 \(\mathcal{C}=\{(t_k, \mathbf{Y}_k^{\text{int}})\}_{k=1}^K\)（其中每个 \(t_k\) 为单基因敲除干预，\(\mathbf{Y}_k^{\text{int}}\in\mathbb{R}^{M\times d}\) 为干预后细胞表达谱分布），以及一个待预测的目标干预查询 \(t_q\)。模型的目标是在不显式重建调控网络图结构的前提下，单步生成该目标干预下的细胞表达谱后验分布 \(p(\mathbf{y}_q^{\text{int}}\mid\text{do}(t_q), \mathbf{Y}^{\text{obs}}, \mathcal{C})\)。
整个系统分为两个阶段：预训练阶段在大规模采样的合成因果图（线性 SCM 或无标度网络 SERGIO 模拟器）上模拟基因敲除过程，生成海量合成先验任务 \((\mathbf{Y}^{\text{obs}}, \mathcal{C}, t_q, \mathbf{Y}_q^{\text{int}})\)，利用条件流匹配（Conditional Flow Matching）训练多模态扩散 Transformer；推理阶段直接将真实单细胞实验作为上下文输入，通过单次前向传播并求解常微分方程（ODE）实现后验细胞分布的分摊推断。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}, 'subGraphTitleMargin': {'top': 8, 'bottom': 16}}}%%
flowchart TD
A["输入数据<br/>观测细胞 Y_obs + 实验上下文 C + 查询干预 t_q"] --> B["多模态流匹配架构<br/>三路模态流独立投影 + 联合自注意力"]
B --> C["合成生物学因果先验<br/>无标度 GRN 拓扑 + SERGIO Hill 非线性动力学"]
C --> D["反事实配对预训练<br/>恒定随机初值模拟 + 消除非配对环境噪声"]
D --> E["测试期测试时增强<br/>重叠子集采样 + 跨子集细胞池化统计检验"]
E --> F["输出结果<br/>后验扰动细胞分布 Y_q_int 与差异表达基因"]关键设计¶
1. 多模态流匹配架构：无序细胞集合的联合注意力与连续速度场建模
针对单细胞测序数据中细胞排列完全无序且无配对关系的特性，传统时序或位置编码会导致错误的归纳偏差。MapPFN 采用三路分支的多模态扩散 Transformer（MMDiT）架构，将高斯噪声流、细胞表达谱状态流和干预条件流解耦为独立的参数空间，并在自注意力层中实现跨模态联合交互。模型彻底移除了正弦绝对位置编码，完全依托自注意力机制的置换不变性，并引入可学习的语义嵌入以显式区分三种核心元信息：上下文中的不同微扰实验标识、观测对照组与干预实验组、以及查询干预与上下文样本。在此基础上，网络采用条件流匹配目标进行训练，在插值轨迹 \(\mathbf{Y}_\tau = (1-\tau)\mathbf{Y}_0 + \tau\mathbf{Y}_q^{\text{int}}\) 上学习回归从标准正态分布向量场指向真实干预后分布的连续速度场： $\(\mathcal{L}_{\text{CFM}}(\theta) = \mathbb{E}_{\tau, \mathbf{Y}_0, \mathbf{Y}_q^{\text{int}}}\left\|v_\tau^\theta(\mathbf{Y}_\tau \mid t_q, \mathbf{Y}^{\text{obs}}, \mathcal{C}) - (\mathbf{Y}_q^{\text{int}} - \mathbf{Y}_0)\right\|_F^2\)$ 在推理阶段，通过数值求解器（Dopri5）求解速度场常微分方程，配合无分类器引导（Classifier-Free Guidance, CFG 权重设为 2.0），在单步前向流中直接生成高质量的目标扰动细胞分布。
2. 合成生物学因果先验：非线性调控动力学与可辨识性保证
真实单细胞干预数据极度匮乏，直接在真实数据上元学习会导致严重的过拟合与模式坍塌。MapPFN 构建了一套高保真合成生物学先验生成器，首先通过优先连接算法（Preferential Attachment）采样具有生物学真实属性的无标度基因调控网络，精确模拟真实生物网络中的稀疏度、非对称出入度以及模块化结构，并强制消除循环依赖及设立主调控因子（Master Regulators）。随后，先验生成器调用基于化学朗之万方程的单细胞随机微分方程模拟器 SERGIO，利用非线性 Hill 动力学函数建模转录因子与靶基因之间的饱和调节效应，并严格对齐 10x Chromium 测序技术的文库大小与 Dropout 噪声分布。根据因果理论，纯观测数据只能确定图结构的马尔可夫等价类，而包含 \(K\) 个单基因硬敲除（do-intervention）的干预上下文能够显著缩减等价类空间，为模型提供超越纯观测方法的结构可辨识性保证。
3. 反事实配对预训练：剥离非配对测量噪声以聚焦因果效应机制
真实测序破坏性导致实验只能获取非配对的细胞快照，这种由初始状态采样带来的随机差异会严重干扰模型对微弱因果调控信号的学习。MapPFN 在合成预训练中引入反事实配对机制：在同一个 GRN 结构下模拟不同干预条件（包括对照组与敲除组）时，固定随机微分方程模拟器的底层随机数种子，使得细胞表达状态在对照组和各个干预组之间具有精确的一对一反事实对应关系。这一设计将流匹配的学习目标从"补偿未配对样本分布的初始噪声方差"转化为"纯粹捕捉干预算子引起的内在因果跃迁机制"，显著加速了模型特征方差相关性的收敛速率，使预训练特征方差与真实值的相关系数从非配对基线的 0.6 大幅提升至 0.8 以上。
4. 测试期测试时增强：突破固定维度限制的外推缩放机制
传统的微扰预测模型由于网络第一层线性权重与固定基因列表死绑定，若输入基因集或目标维度发生改变，模型必须彻底重新训练。MapPFN 的输入仅与单细胞调控机制在上下文中的相对表达模式交互，预训练设置虽然针对 50 个基因子集，但在推理期可通过测试时增强（Test-Time Augmentation, TTA）灵活扩展到任意规模基因集。针对包含 100 个基因的大型靶标系统，MapPFN 在推理时随机采样 10 次包含 50 个基因的重叠子集分别进行前向预测，将各子集生成的细胞级后验分布在基因维度进行跨子集池化聚合，最后执行单基因 Wilcoxon 秩和检验识别差异表达基因（DEG）。TTA 机制不仅消除了重新训练模型的算力与数据开销，还通过多子集集成有效压低了预测方差。
损失函数 / 训练策略¶
模型采用带有条件 Dropout（丢弃概率 \(p=0.2\)）的条件流匹配目标，时间标量采样自对数正态分布 \(\tau\sim\text{LogitNormal}(0, 1)\)。优化器使用 AdamW，配合预热-稳定-衰减（Warmup-Stable-Decay, WSD）学习率调度，峰值学习率设为 \(10^{-4}\)，维持指数移动平均（EMA，衰减率 0.999）权重用于最终推理。在线性 SCM 先验上预训练 50k 步，在合成生物学先验上预训练 400k 步。针对真实数据微调时，使用仅需 3,000 步的轻量级微调方案（单 GPU 上耗时约 10 分钟，学习率 \(5\times 10^{-5}\)）。
实验关键数据¶
主实验¶
论文分别在受控线性结构因果模型（SCM，20 个节点）和两个具有显著生物学异质性的真实单细胞 CRISPR 敲除数据集上开展评估：黑素瘤数据集（Melanoma，218k 细胞，248 种免疫逃逸相关基因敲除，以 IFN-\(\gamma\) 处理组作为未知留出背景）和白血病 THP-1 细胞系数据集（Leukemia，20k 细胞，26 种免疫检查点调控基因敲除）。评估指标包括 Wasserstein 距离（\(\text{W}_2\)）、最大均值差异（\(\text{MMD}\times 10^{-3}\)）、均方根误差（\(\text{RMSE}\)）、微扰辨别得分（\(\text{PDS}\)，越低越好）、幅度比率（\(\text{MR}\)，\(\text{MR}=d(\mathbf{Y}^{\text{obs}}, \hat{\mathbf{Y}}^{\text{int}})/d(\mathbf{Y}^{\text{obs}}, \mathbf{Y}^{\text{int}})\)，1.0 为最优，0 为恒等坍塌）以及差异表达基因恢复精准率-召回率曲线下面积（\(\text{AUPRC}\)）。
表 1：真实单细胞 CRISPR 敲除数据集主实验对比（Melanoma 与 Leukemia）
| 数据集 | 方法 | \(\text{W}_2\) ↓ | \(\text{MMD} (\times 10^{-3})\) ↓ | \(\text{RMSE}\) ↓ | \(\text{PDS}\) ↓ | \(\text{MR}\) | \(\text{AUPRC}\) ↑ |
|---|---|---|---|---|---|---|---|
| Melanoma | CPA [11] | 15.57±0.10 | 140.09±0.35 | 0.13±0.00 | 0.49±0.01 | 0.68±0.01 | 0.04±0.00 |
| CondOT [46] | 22.09±0.39 | 7.11±0.12 | 0.10±0.00 | 0.06±0.01 | 0.05±0.00 | 0.34±0.05 | |
| MFM [13] | 20.99±0.14 | 7.28±0.13 | 0.10±0.00 | 0.09±0.02 | 0.13±0.00 | 0.28±0.04 | |
| CellFlow [12] | 22.27±0.59 | 7.16±0.17 | 0.10±0.00 | 0.41±0.01 | 0.01±0.00 | 0.10±0.02 | |
| STATE [14] | 20.52±0.07 | 7.82±0.09 | 0.08±0.00 | 0.07±0.02 | 0.94±0.00 | 0.33±0.04 | |
| MapPFN (纯合成预训练/零样本) | 22.75±0.16 | 10.07±0.19 | 0.13±0.00 | 0.17±0.01 | 1.00±0.01 | 0.34±0.02 | |
| MapPFN (微调) | 21.38±0.12 | 7.84±0.14 | 0.10±0.00 | 0.03±0.01 | 0.99±0.00 | 0.38±0.03 | |
| Identity (下界) | 22.91±0.18 | 7.90±0.17 | 0.11±0.00 | 0.51±0.02 | 0.00±0.00 | 0.04±0.01 | |
| Observed (上界) | 8.54±0.10 | 2.59±0.07 | 0.04±0.00 | 0.00±0.00 | 1.00±0.00 | 0.64±0.06 | |
| Leukemia | CPA [11] | 12.41±0.11 | 78.74±1.27 | 0.17±0.00 | 0.50±0.02 | 0.61±0.01 | 0.15±0.01 |
| CondOT [46] | 17.92±0.37 | 26.51±0.68 | 0.27±0.01 | 0.54±0.04 | 0.19±0.00 | 0.14±0.01 | |
| MFM [13] | 41.71±0.26 | 105.64±0.73 | 0.71±0.00 | 0.51±0.02 | 1.67±0.01 | 0.16±0.01 | |
| CellFlow [12] | 16.87±0.37 | 14.55±0.46 | 0.17±0.00 | 0.50±0.01 | 0.02±0.00 | 0.16±0.01 | |
| STATE [14] | 15.27±0.15 | 15.28±0.44 | 0.17±0.00 | 0.47±0.03 | 0.83±0.00 | 0.17±0.01 | |
| MapPFN (纯合成预训练/零样本) | 44.42±0.23 | 191.88±1.46 | 0.78±0.00 | 0.49±0.01 | 2.56±0.02 | 0.16±0.01 | |
| MapPFN (微调) | 16.32±0.11 | 12.24±0.58 | 0.15±0.00 | 0.42±0.03 | 0.91±0.01 | 0.18±0.01 | |
| Identity (下界) | 16.72±0.18 | 13.17±0.44 | 0.14±0.01 | 0.49±0.02 | 0.00±0.00 | 0.18±0.01 | |
| Observed (上界) | 10.29±0.12 | 2.59±0.04 | 0.04±0.00 | 0.03±0.02 | 0.99±0.01 | 0.91±0.04 |
消融实验¶
论文在 Melanoma 数据集上系统消融了先验设计中的核心机制（包括反事实配对先验与干预上下文）：
表 2：核心模块消融实验对比（Melanoma 数据集）
| 配置 | \(\text{W}_2\) ↓ | \(\text{MMD} (\times 10^{-3})\) ↓ | \(\text{RMSE}\) ↓ | \(\text{PDS}\) ↓ | \(\text{MR}\) | \(\text{AUPRC}\) ↑ | 说明 |
|---|---|---|---|---|---|---|---|
| MapPFN (预训练基准) | 22.75±0.16 | 10.07±0.19 | 0.13±0.00 | 0.17±0.01 | 1.00±0.01 | 0.34±0.02 | 完整预训练模型 |
| − 反事实配对先验 | 24.44±0.31 | 21.84±1.28 | 0.23±0.01 | 0.20±0.02 | 1.14±0.01 | 0.21±0.03 | 移除配对，替换为非配对合成先验，AUPRC 下降 38.2% |
| − 干预上下文 (\(\mathcal{C}=\emptyset\)) | 23.78±0.16 | 15.88±0.19 | 0.20±0.00 | 0.20±0.01 | 1.10±0.01 | 0.13±0.02 | 仅利用纯观测数据推断，AUPRC 下降 61.8% |
| MapPFN (随机初始化从头训练) | 21.23±0.11 | 51.78±0.80 | 0.24±0.00 | 0.18±0.02 | 1.02±0.01 | 0.12±0.02 | 无合成预训练直接在真实数据训练，MMD 激增 5 倍 |
关键发现¶
- 干预上下文对识别因果效应具有决定性贡献：移除测试期干预上下文（\(\mathcal{C}=\emptyset\)）使差异表达基因检测的 AUPRC 从 0.34 骤降至 0.13（跌幅达 61.8%），充分证明纯观测数据无法充分确定因果拓扑，而上下文中的序列化扰动实验显著增强了机制的可辨识性。
- 基线模型普遍存在恒等坍塌（Identity Collapse）：通过幅度比率（Magnitude Ratio, MR）指标可清晰发现，CondOT、CellFlow 和 MFM 在真实数据上的 MR 普遍低于 0.15（CellFlow 甚至低至 0.01），实质上仅退化为复读对照组分布；而 MapPFN 能够精准恢复 0.99~1.00 的真实干预幅度。
- 合成预训练赋予强大的泛化先验：纯粹在合成生物学先验上预训练的 MapPFN 在零样本设定下即达到 0.34 的 AUPRC，追平了在真实数据上从头训练的 SOTA 模型 STATE（0.33）；经过 10 分钟快速微调后进一步提升至 0.38，全面超越所有真实数据基准。
亮点与洞察¶
- 将 PFN 推向高维无配对连续分布生成：巧妙地将先验数据拟合网络从传统表格数据或一维因果效应估计，推广至以单细胞群为代表的高维无序分布流匹配，解决了生物学数据破坏性测序带来的无配对测量瓶颈。
- 解耦昂贵湿实验与先验模型预训练：摆脱了对百万人体细胞微扰大图谱的数据依赖，证明仅靠结合图论无标度网络与 Hill 方程动力学的合成先验，即可元学习到高度可迁移的转录组调控规律。
- 测试期任意基因集外推能力：通过结合子集采样与池化统计检验的测试时增强（TTA），彻底打破了固定神经网络输入维度的硬编码束缚，展示了作为通用虚拟细胞基础模型（Virtual Cell Foundation Model）的技术雏形。
局限与展望¶
- 先验对硬敲除的假设限制：当前合成生物学先验基于纯基因硬敲除（Hard Knockout）模拟，难以直接原生建模基于 CRISPRi 的弱抑制软敲减（Soft Knockdown）或表观遗传过表达。
- 组合扰动与化学药物扰动尚未覆盖：目前的实验设定仅限于单基因敲除，无法直接扩展到多基因协同互作或化学小分子靶向药物的剂量响应曲线预测。
- 高维系统模拟的算力开销：随着基因数量向全转录组（20,000+ 基因）规模扩展，基于常微分/随机微分方程的合成数据生成时间和流匹配训练代价将显著上升。
相关工作与启发¶
- vs CPA / CellFlow / CellOT: 传统基于最优传输或流匹配的生成式模型需要依赖特定细胞系的类别标签或协变量，在未见过的全新生物学背景中无法根据少量探索性干预自适应；MapPFN 通过上下文学习实现了纯数据驱动的免重训即时泛化。
- vs Meta Flow Matching (MFM) / STATE: MFM 和 STATE 引入了对观测对照组分布的条件注入，但依然无法利用干预实验带来的因果约束，极易发生向均值坍塌；MapPFN 首次实现了对干预分布集合的显式联合注意力和机制识别。
- vs Do-PFN / CausalPFN: 先前的因果 PFN 仅能处理一维连续变量的处理效应（ATE/CATE）推断，且仅依赖观测数据；MapPFN 首次实现了高维多变量细胞表达谱分布级因果推断与干预引导。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次将先验数据拟合网络（PFN）引入单细胞微扰分布预测，提出了干预上下文学习与基于 Hill 非线性动力学的合成生物先验。
- 实验充分度: ⭐⭐⭐⭐☆ 包含从受控线性 SCM 到两个真实单细胞大图谱的系统评测，设计了细致的先验配对与上下文消融，略显遗憾的是基因维度受限于 50-100 维。
- 写作质量: ⭐⭐⭐⭐⭐ 数学形式化规范，问题陈述与图表设计清晰，逻辑推演扎实流畅。
- 价值: ⭐⭐⭐⭐⭐ 为虚拟细胞系统的研发提供了"合成生物学先验元学习 + 推理期上下文自适应"的全新范式，对计算生物学和因果表征学习具有重要启发。