🤖 machine learning

MapPFN: Learning Causal Perturbation Maps in Context

MapPFN is a prior-data fitted network that leverages in-context learning on synthetic causal perturbation data to adaptively predict post-perturbation gene expression distributions in unseen biological contexts, outperforming existing methods in identifying differentially expressed genes.

Original authors: Marvin Sextro, Weronika Kłos, Gabriel Dernbach

Published 2026-04-22

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Marvin Sextro, Weronika Kłos, Gabriel Dernbach

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Cellular Weather" Dilemma

Imagine you are a doctor trying to figure out how a specific medicine (a "perturbation") will affect a patient's cells. In the real world, you can't test every drug on every single patient because:

It's expensive and slow: You have to grow cells, treat them, and sequence their DNA.
It's impossible to cover everyone: There are billions of possible combinations of genes and diseases. You can't run an experiment for every single scenario.

Existing computer models are like weather forecasters who only know the last 10 days of weather. If you ask them to predict the weather for a new city or a new season they've never seen, they usually guess wrong because they haven't "seen" that context before. They are stuck relying only on the specific data they were trained on.

The Solution: MapPFN (The "Super-Apprentice")

The authors created a new AI called MapPFN. Think of it not as a weather forecaster, but as a super-intelligent apprentice who has read every book in the library of "how cells work" before ever stepping into a real lab.

Here is how it works, broken down into three simple steps:

1. The "Simulated Universe" Training (Pre-training)

Instead of waiting for real experiments (which are slow), the team built a giant video game simulation of biology.

They created thousands of fake "worlds" with different rules for how genes talk to each other.
They simulated millions of fake experiments: "What happens if we knock out Gene A in World 1? What about Gene B in World 2?"
The Analogy: Imagine a chess player who has played 10 million games against a computer that generates random, impossible board setups. They haven't just memorized specific moves; they have learned the deep logic of how pieces interact in any situation.

MapPFN learns this logic. It becomes a master of "Cause and Effect" in biology, even though it has never seen a real human cell.

2. The "Context Clue" Trick (In-Context Learning)

This is the magic part. Usually, AI needs to be retrained from scratch for every new hospital or new type of cancer. MapPFN doesn't need that.

The Analogy: Imagine you are a detective. You have never been to a specific crime scene before. But, you are handed a notebook (the "Context") containing clues from three other similar crimes that just happened in that same neighborhood.
How MapPFN uses this: When you give MapPFN a new, unseen biological context (like a new patient's cells), you also give it a small set of real experimental results from that specific patient (e.g., "Here is what happened when we knocked out Gene X").
MapPFN looks at these clues, says, "Ah, I've seen similar patterns in my simulation! I know how this specific neighborhood works," and instantly predicts what will happen if you knock out Gene Y.

It adapts on the fly, without needing to retrain its brain.

3. The "Virtual Lab" (Prediction)

Once MapPFN has the clues, it acts as a crystal ball.

You ask: "What happens if we give this drug to these cells?"
MapPFN doesn't just guess a single outcome; it predicts the entire distribution of how the cells will change. It tells you, "Most cells will react this way, but a few might react differently."

Why is this a Big Deal?

Zero-Shot Learning: MapPFN can predict results for genes or diseases it has never seen in real life, simply because it learned the underlying rules from its simulation. It performs just as well as models trained on real data, but without needing the real data first.
Better than the Rest: When tested on real-world data (like melanoma and leukemia cells), MapPFN beat all the previous "state-of-the-art" models.
The "Context" is Key: The paper proved that giving the AI a few real examples (the "context") makes it much smarter than just giving it a label like "Drug A." It's the difference between telling a chef "Make a cake" vs. handing them a picture of the specific cake the customer wants and saying "Make this."

The Bottom Line

MapPFN is a "Virtual Cell" foundation model.

Think of it as a universal translator for biology. It learned the grammar of life in a simulated world, and now, when you show it a few sentences of a new language (a new biological context), it can instantly translate and predict the rest of the story. This could revolutionize drug discovery by letting scientists test millions of hypotheses in a computer before ever touching a test tube.

1. Problem Statement

The core challenge addressed is single-cell perturbation prediction: predicting the post-perturbation distribution of gene expression ( $Y^{int}$ ) given an observational distribution ( $Y^{obs}$ ) and a specific intervention (e.g., a gene knockout).

The Bottleneck: Existing methods struggle to generalize to unseen biological contexts (e.g., new cell lines or conditions) because they are typically trained on limited experimental data. They often fail to leverage new interventional evidence available at inference time to adapt their predictions.
The Data Constraint: Single-cell sequencing destroys cells, meaning pre- and post-perturbation data are unpaired. This transforms the problem into learning a map between distributions rather than individual cell pairs.
The Limitation of Current Models: Most state-of-the-art models (e.g., CPA, CellOT) condition only on covariates or the observational distribution. They lack the ability to perform in-context learning using a set of observed interventional experiments from the target context to adapt to unseen perturbations.

2. Methodology: MapPFN

The authors propose MapPFN, the first Prior-Data Fitted Network (PFN) designed for perturbation prediction. It frames the problem as in-context learning (ICL) over interventional distributions.

A. Core Concept: Prior-Data Fitted Networks (PFNs)

Instead of training a model on a specific dataset and then fine-tuning, MapPFN is meta-learned on a massive synthetic prior.

Training: The model is pre-trained on synthetic datasets generated from a distribution over Structural Causal Models (SCMs).
Inference: At test time, the model receives a "context" consisting of observational data and a few interventional experiments from the target biological context. It uses this context to predict the outcome of a query perturbation in a single forward pass (amortized inference), without needing to explicitly recover the underlying causal graph.

B. Synthetic Biological Prior

To overcome the scarcity of real perturbation data, the authors construct a synthetic prior:

Graph Generation: They sample diverse Gene Regulatory Networks (GRNs) using a preferential attachment algorithm to mimic real-world properties (sparsity, modularity, scale-free degree distribution).
Simulation: They use SERGIO, a simulator based on Stochastic Differential Equations (SDEs) with Hill functions, to generate realistic single-cell gene expression dynamics.
Intervention: Interventions are simulated as "hard" knockouts (removing a node and its incoming edges in the SCM).
Counterfactual Pairing: A key innovation is the use of a counterfactual paired prior. By keeping the random seed constant across treatments during simulation, they generate paired interventional data. This isolates causal effects from noise, significantly improving downstream performance compared to unpaired data.

C. Model Architecture

MapPFN utilizes a modified Multimodal Diffusion Transformer (MMDiT):

Input Modalities: The model processes three streams:
1. Noise: The initial noise vector for the diffusion process.
2. Cell States: The observational ( $Y^{obs}$ ) and interventional ( $Y^{int}$ ) cell data.
3. Treatments: One-hot encoded perturbation identifiers.
In-Context Mechanism: The model treats cells as tokens. It uses joint attention to allow cross-modal interactions. Crucially, it learns to differentiate between context (observed interventions) and query (the unseen intervention to be predicted) via learnable embeddings.
Objective: The model is trained using Conditional Flow Matching, learning a velocity field that transports noise to the predicted post-perturbation distribution.

3. Key Contributions

Novel Formulation: Frames perturbation prediction as a distribution mapping problem solvable via in-context learning, enabling a single pre-trained model to adapt to arbitrary biological contexts and gene sets.
MapPFN Architecture: Introduces the first PFN for perturbation prediction, leveraging synthetic pre-training to meta-learn causal perturbation maps across diverse structures.
Counterfactual Prior: Demonstrates that training on counterfactual paired data (where initial conditions are identical across treatments) significantly improves the model's ability to isolate causal effects.
Test-Time Adaptation: Shows that providing a small set of interventional experiments at inference time (context) drastically improves prediction accuracy over methods that rely solely on observational data.

4. Experimental Results

The authors evaluated MapPFN on synthetic linear SCMs and two real-world single-cell datasets (Melanoma and Leukemia).

Synthetic Benchmarks (Linear SCMs):
- MapPFN achieved the best performance across all metrics (Wasserstein distance, MMD, RMSE, and Magnitude Ratio).
- Baselines (CondOT, MFM) suffered from "identity collapse" (predicting the observational distribution unchanged), whereas MapPFN successfully learned the magnitude of the perturbation effect.
Real-World Datasets (Zero-Shot & Fine-Tuned):
- Zero-Shot: Pre-trained only on synthetic data, MapPFN recovered differentially expressed genes (DEGs) with performance on par with baselines trained from scratch on real data.
- Fine-Tuned: When fine-tuned on real data, MapPFN consistently outperformed all baselines (CPA, CondOT, MFM, CellFlow, STATE) across both datasets.
- Ablation Studies:
  - Removing the interventional context (using only observational data) degraded performance, proving the value of in-context learning.
  - Removing the counterfactual paired prior (using unpaired data) significantly reduced performance, highlighting the importance of the synthetic training strategy.
Scalability:
- Gene Sets: Through Test-Time Augmentation (TTA), MapPFN successfully scaled to predict perturbations for gene sets larger than those seen during training (e.g., 100 genes vs. 50 training genes).
- Data Volume: Performance improved monotonically as the number of cells and context experiments increased at inference time, demonstrating true in-context learning capabilities.

5. Significance and Impact

Virtual Cell Foundation Models: MapPFN represents a significant step toward "Virtual Cell" foundation models that can generalize across unseen biological contexts without retraining on massive new datasets.
Cost Reduction: By leveraging synthetic pre-training and in-context adaptation, the method reduces the need for expensive, high-throughput wet-lab experiments for every new cell line or condition.
Causal Inference: The approach bridges the gap between causal inference and generative modeling, showing that meta-learning on synthetic causal priors can effectively recover causal mechanisms in complex biological systems.
Practical Utility: The ability to predict treatment effects in unseen contexts with limited data makes this highly relevant for drug discovery and personalized medicine, where patient-specific cell lines are often unique and data-sparse.

In summary, MapPFN successfully demonstrates that a model pre-trained on diverse synthetic causal mechanisms can adapt to real-world biological perturbations via in-context learning, outperforming traditional methods that lack this flexibility.