MapPFN: Learning Causal Perturbation Maps in Context

Sextro, Marvin; Klos, Weronika; Dernbach, Gabriel

MapPFN: Learning Causal Perturbation Maps in Context

Marvin Sextro^1,2,3, Weronika Kłos^1,2, Gabriel Dernbach^1,4,2,3

¹Machine Learning Group, TU Berlin ²BIFOLD ³Aignostics ⁴Charité

Paper Code Models Datasets

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) pre-trained on a synthetic biological prior with causal interventions, decoupling pre-training from limited wet-lab data. Unlike existing methods, MapPFN uses in-context learning to map a sequence of experiments to a post-perturbation distribution, enabling a single pre-trained model to adapt to new datasets and arbitrary gene sets at inference time. Zero-shot, MapPFN identifies differentially expressed genes on par with models trained on real single-cell data, and fine-tuning further improves predictions across biological contexts.

The Problem

Perturbation effects are context-dependent
Combinatorial space infeasible to cover experimentally
Existing methods retrain and tune per dataset
No method conditions on interventional data

Our Solution

Pre-training on Synthetic Data

A prior-data fitted network (PFN) that meta-learns causal perturbation prediction.

Interventional Context Conditioning

Leverage observed interventions to improve identifiability of causal structure.

In-Context Learning (ICL)

Adapt to unseen biological contexts and gene sets at inference time.

Method

Draw causal model $\psi$ from prior $p(\psi)$, generate $\mathbf{Y}^{\text{obs}}$ and build context $\mathcal{C} = \{(t_k, \mathbf{Y}^{\text{int}}_k)\}_{k=1}^K$ by intervening with treatments $t_k \in \mathcal{T}$. MapPFN uses a Multimodal Diffusion Transformer (MMDiT) to approximate the posterior predictive distribution:

$$p(\mathbf{y}_q^{\text{int}} \mid \text{do}({\color{#D55E00} t_q}), {\color{#0072B2} \mathbf{Y}^{\text{obs}}}, {\color{#009E73} \mathcal{C}}) = \int p(\mathbf{y}_q^{\text{int}} \mid \text{do}({\color{#D55E00} t_q}), {\color{#0072B2} \mathbf{Y}^{\text{obs}}}, {\color{#E69F00} \psi}) \; p({\color{#E69F00} \psi} \mid {\color{#0072B2} \mathbf{Y}^{\text{obs}}}, {\color{#009E73} \mathcal{C}}) \; d{\color{#E69F00} \psi}$$

MapPFN overview. During pre-training, synthetic causal models are drawn to generate observational and interventional distributions. MapPFN meta-learns to map between pre- and post-perturbation distributions across many causal structures. At inference, it predicts cell-level post-perturbation distributions in one step through amortized inference.

Synthetic Biological Prior

Preferential attachment gene regulatory network (GRN) generator with sparse, directed, modular graphs
Simulation of gene expression dynamics and in silico knockouts via stochastic differential equations (SDEs)
Counterfactual prior improves downstream performance

Get Started

Fine-tune MapPFN

Download the pre-trained model and fine-tune on your own perturbation data.

Models

Download Datasets

Access the synthetic and real-world perturbation datasets used for pre-training and evaluation.

Datasets

View Source Code

Reproduce experiments, explore the codebase, and build on MapPFN.

Code

BibTeX

@inproceedings{sextro2026mappfn,
  title     = {Map{PFN}: Learning Causal Perturbation Maps in Context},
  author    = {Marvin Sextro and Weronika K{\l}os and Gabriel Dernbach},
  booktitle = {Advances in Neural Information Processing Systems},
  volume    = {39},
  year      = {2026}
}