MapPFN: Learning Causal Perturbation Maps in Context

MapPFN is a prior-data fitted network (PFN) that uses in-context learning to predict perturbation effects in unseen biological contexts.

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.

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 gene sets. Our results show that combining synthetic pre-training with fine-tuning on real single-cell data improves perturbation predictions across biological contexts.

Setup

A Docker image and devcontainer configuration are provided with all dependencies:

docker run --rm -it --gpus all -v .:/srv/repo ghcr.io/marvinsxtr/mappfn:latest bash

VSCode & Slurm

Use the Remote - Containers extension to open the devcontainer locally, or connect to a remote tunnel by replacing bash with code tunnel.

This setup also works with Apptainer on Slurm clusters. See the ml-project-template for instructions.

WandB logging (optional)

Create a .env file in the repository root:

WANDB_API_KEY=your_api_key
WANDB_ENTITY=your_entity
WANDB_PROJECT=your_project_name

Data

Download pre-trained weights and datasets from Hugging Face:

python -m map_pfn.scripts.download_data

Add --include-prior to also download the synthetic prior (27 GB) used for pre-training and fine-tuning.

Preprocessing & generation

Preprocess single-cell datasets:

for seed in 10 20 30 40 50; do
  for dataset in frangieh papalexi; do
    python -m map_pfn.scripts.process_sc_data --dataset $dataset --seed $seed
  done
done

Generate synthetic datasets:

python map_pfn/scripts/generate_data.py cfg=linear   # Linear SCMs
python map_pfn/scripts/generate_data.py cfg=sergio    # Biological prior

Inference

from map_pfn.eval.evaluate import load_model

trainer, module, datamodule = load_model(
    method="map_pfn",
    checkpoint_path="checkpoints/model.ckpt",
    dataset_path="datasets/single_cell/frangieh/split_10.h5ad",
)
preds = trainer.predict(module, datamodule=datamodule)

Fine-tuning

Fine-tune from a pre-trained checkpoint:

python map_pfn/scripts/train.py \
    cfg=map_pfn_rna \
    cfg/datamodule=frangieh_finetune \
    cfg.load_checkpoint=checkpoints/model.ckpt \
    cfg.globals.num_steps=600000 \
    cfg.trainer.val_check_interval=500 \
    cfg.trainer.callbacks.2.max_steps=3000 \
    cfg/wandb=base

Pre-training

Train MapPFN from scratch:

python map_pfn/scripts/train.py cfg=map_pfn_rna

Configuration

This project uses hydra-zen for configuration. Display all available options:

python map_pfn/scripts/train.py --help
python map_pfn/scripts/generate_data.py --help

Repository Structure

MapPFN/
├── map_pfn/
│   ├── configs/         # Hydra-zen configuration
│   ├── data/            # Datasets and data generation
│   ├── models/          # MapPFN and MMDiT architecture
│   ├── eval/            # Evaluation metrics
│   ├── loss/            # Loss functions (CFM)
│   ├── scripts/         # Training and data generation
│   ├── train/           # Training utilities
│   └── utils/           # Helpers
├── baselines/
│   ├── condot/          # Conditional Optimal Transport
│   └── metafm/          # Meta Flow Matching
└── datasets/            # Generated datasets (gitignored)

Citation

@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}
}

Contributing

If you have any feedback, questions, or ideas, please open an issue or reach out via email.