MapPFN: Learning Causal Perturbation Maps in Context

TL;DR

MapPFN tackles the challenge of predicting perturbation effects in unseen biological contexts by framing perturbation prediction as a context-conditioned distribution mapping. It introduces a prior-data fitted network pretrained entirely on synthetic data from structural causal models and gene regulatory networks, and uses a Multimodal Diffusion Transformer to perform in-context learning, generating $p(oldsymbol{y}^{\text{int}}_q \mid do(t_q), \boldsymbol{Y}^{\text{obs}}, \mathcal{C})$ in a single forward pass without gradient updates. Across synthetic SCMs and real Perturb-Seq data, MapPFN achieves competitive performance and robust DEG recovery, with interventional context and paired counterfactual priors providing measurable gains. The work highlights the potential of synthetic priors to enable context-adaptive virtual cell models, while noting limitations related to prior fidelity and scalability, and outlining a path toward larger-scale, context-aware perturbation prediction.

Abstract

Planning effective interventions in biological systems requires treatment-effect models that adapt to unseen biological contexts by identifying their specific underlying mechanisms. Yet single-cell perturbation datasets span only a handful of biological contexts, and existing methods cannot leverage new interventional evidence at inference time to adapt beyond their training data. To meta-learn a perturbation effect estimator, we present MapPFN, a prior-data fitted network (PFN) pretrained on synthetic data generated from a prior over causal perturbations. Given a set of experiments, MapPFN uses in-context learning to predict post-perturbation distributions, without gradient-based optimization. Despite being pretrained on in silico gene knockouts alone, MapPFN identifies differentially expressed genes, matching the performance of models trained on real single-cell data. Our code and data are available at https://github.com/marvinsxtr/MapPFN.

Figures (5)

• Figure 1: MapPFN overview. MapPFN uses in-context learning (ICL) to predict perturbation effects in unseen biological contexts. During pretraining, we draw structural causal models (SCMs) or synthetic gene regulatory networks (GRNs) $\psi$ to generate samples from the observational distribution $\mathbf{Y}^\text{obs}$ and a context set of interventional distributions $\{(t_k, \mathbf{Y}^\text{int}_k)\}_{k=1}^K$, where $t_k$ denotes a perturbation (do-intervention). Given ${\mathbf{Y}^\text{obs}}$ and the context set, MapPFN predicts post-perturbation distributions $\mathbf{Y}^\text{int}_q$ arising from unseen interventions $t_q$. During pretraining, MapPFN meta-learns how to map between pre- and post-perturbation distributions across many causal structures $\psi$ by minimizing $\mathcal{L}({\hat{\mathbf{Y}}}_q^{\text{int}}, \mathbf{Y}^\text{int}_q)$. At inference time, MapPFN predicts cell-level post-perturbation distributions $\mathbf{Y}^\text{int}_q \in \mathbb{R}^{\text{cells} \times \text{genes}}$ in one step through amortized inference, without requiring gradient-based optimization or knowledge of the underlying causal structure $\psi$.
• Figure 2: Recovery of differentially expressed genes in real single-cell data frangieh_multimodal_2021. Precision-recall curves and AUPRC for identification of $n=68$ differentially expressed genes in the held-out IFN-$\gamma$ context. For each method, we report the model with the median AUPRC across three seeds. Despite being pretrained exclusively on synthetic data, MapPFN achieves the highest AUPRC, outperforming baselines trained on real perturbations from the held-out context.
• Figure 3: Increasing the number of perturbation experiments in context improves performance. Wasserstein distance measured on the test context of the single-cell dataset for varying numbers of perturbation experiments in the context set $\mathcal{C}$. Shaded regions indicate standard deviation over three model seeds. Increasing the context size $K = |\mathcal{C}|$ improves performance of MapPFN, with diminishing returns for more than four perturbation experiments.
• Figure 4: Counterfactual paired prior improves downstream performance. Counterfactual interventional distributions using shared noise across treatments accelerate convergence and improve final performance both within the prior (dashed) and on real single-cell data (solid). Variance correlation measures the Pearson correlation between feature variances of predicted and ground-truth samples. Shaded regions indicate rolling standard deviation; curves show EMA ($\alpha=0.95$).
• Figure 5: Data split in the few-shot and zero-shot setting. Each box represents a dataset $\mathbf{Y}^\text{int}_{ij} \in \mathbb{R}^{N\times d}$ sampled from the SCM $\psi_i$ under treatment $t_j$. Green boxes are part of the training data and purple boxes are withheld for evaluation. In the few-shot setting, the training data includes interventional distributions from a subset of perturbations in the test context. In the zero-shot setting, no perturbations on the test context are available in the training data. In this setting the model has to recover perturbation effects from observational data alone. The few-shot setting is of practical interest as highlighted by the Virtual Cell Challenge roohani_virtual_2025.