CytoWave: Perturbation-Centric Pretraining for Single-Cell Response Prediction

AI summary

Overview

Research area: Machine learning for single-cell biology — specifically, self-supervised pretraining of models that predict how a cell's gene expression changes after a genetic perturbation (e.g., CRISPR knockout).

Technical level: Advanced. The paper assumes familiarity with Transformer architectures, contrastive learning (InfoNCE), masked-reconstruction pretraining, and single-cell transcriptomics formats.

Scope (one sentence): The paper builds a 24.7-million-cell, cross-species, orthology-aligned perturbation corpus and pretrains CytoWave — a Transformer encoder–operator–decoder — in two self-supervised stages to predict single-cell genetic perturbation responses under strict dataset-held-out evaluation.

What This Paper Is About

Predicting which genes turn up or down after a gene is knocked out is central to understanding gene function and designing therapeutics, but models trained on one experiment usually fail on another because of differences in protocols, cell types, and species. This paper asks whether pretraining directly on perturbation data — rather than pretraining on unperturbed cells and adding perturbations later — produces representations that transfer to entirely unseen datasets. CytoWave answers by learning perturbations as additive shifts in a shared latent cell-state space, trained jointly with a contrastive objective that ties those latent shifts back to the perturbation identity.

Key Contributions

1. A unified cross-species perturbation corpus. The authors collect and harmonize 42 single-cell genetic perturbation datasets from 19 studies across 2 species (human and mouse), using orthology-based gene alignment and standardized metadata vocabularies, yielding approximately 24.7 million single-cell transcriptomes for pretraining and dataset-held-out benchmarking. Five datasets are held out entirely for evaluation.

2. A perturbation-centric pretraining framework. CytoWave models perturbation responses as changes in latent cellular states via a learned perturbation operator, rather than as direct gene-to-gene mappings. This differs from prior single-cell pretrained models (scGPT, scFoundation, GeneCompass, Geneformer, LPM) that pretrain on unperturbed expression and only incorporate perturbations downstream.

3. A two-stage self-supervised training scheme with bidirectional objectives. Stage 1 learns reference control cell states through masked gene-expression recovery; Stage 2 combines forward perturbation-response prediction (MSE on expression) with a reverse contrastive alignment that predicts masked perturbation embeddings from observed latent state changes using in-batch InfoNCE (temperature τ = 0.1).

4. Dataset-held-out benchmarking with a pathway-level case study. CytoWave is evaluated on five held-out datasets using ΔPCC, ΔPCC20, and MAE, plus an AUCell-based analysis over the Hallmark 50 gene sets to test whether predicted responses recover known pathway activity.

Main Findings

• Best overall on held-out datasets: CytoWave-S2 reports the strongest scores across all three metrics on all five evaluation datasets among the compared methods. Its MAE values on DatlingerBock2017, DatlingerBock2017(Stim), PapalexiSatija2021, GasperiniShendure2019, and SantinhaPlatt2023 are 1.077, 0.867, 0.791, 1.133, and 0.833 respectively, compared to 2.016, 1.996, 2.200, 2.547, and 2.057 for BaseReg and 2.007, 2.000, 2.214, 2.519, and 2.057 for TrainMean.

• Margins over CPA are large on MAE but narrow on correlation: CPA — the closest competitor — reaches MAE 1.080, 0.982, 0.997, 1.332, 1.001, and ΔPCC 0.132, 0.237, 0.251, 0.102, 0.101. CytoWave-S2's ΔPCC of 0.189, 0.279, 0.255, 0.112, 0.107 and ΔPCC20 of 0.272, 0.330, 0.305, 0.151, 0.251 are only modestly above CPA's ΔPCC20 of 0.261, 0.328, 0.304, 0.149, 0.246 — on PapalexiSatija2021 the two are effectively tied (0.305 vs 0.304).

• Stage 1 alone is already competitive: CytoWave-S1 (control-only pretraining, no perturbation-aware stage) achieves MAE 1.082, 0.891, 0.852, 1.298, 0.871 — already near or above most baselines — supporting the claim that large-scale pretraining on control cells provides a solid foundation for transcriptional state modeling.

• Stage 2 adds consistent gains: Adding perturbation-aware learning (CytoWave-S2) improves over CytoWave-S1 on every reported metric across all five datasets, with the largest MAE improvement on GasperiniShendure2019 (1.298 to 1.133) and on PapalexiSatija2021 (0.852 to 0.791).

• Metadata helps but is not the sole driver: Removing metadata tokens (CytoWave-w/o Meta) gives MAE 1.144, 0.921, 0.885, 1.317, 0.957 — worse than CytoWave-S2 on MAE, though in some ΔPCC20 cells it is marginally higher (0.313 vs 0.300 on DatlingerBock2017(Stim); 0.248 vs 0.243 on SantinhaPlatt2023). The authors' stated conclusion is that metadata contributes to performance without being the sole driver of improvements.

• Several baselines perform poorly or negatively: GEARS shows ΔPCC of 0.012, 0.002, −0.010, −0.001, 0.006 and near-zero or negative ΔPCC20 on most datasets; GenePert and scouter also produce negative ΔPCC on DatlingerBock2017 (−0.009 and −0.031).

• Pathway-level agreement: Using AUCell scores over the Hallmark 50 gene sets, predicted and ground-truth pathway activity heatmaps show consistent activation and repression patterns across the five benchmark datasets, indicating that CytoWave captures structured biological responses beyond individual gene-level predictions. The figure shows only "an illustrative subset of high-agreement pathways and perturbations."

• Evaluation caveat noted by the authors: All methods are compared under the same cell-level protocol, and the paper notes that deterministic population-level predictors may therefore appear conservative under single-cell ΔPCC evaluation.

Methodology in Plain English

Data unification. Raw datasets name genes inconsistently and mix species. The authors strip isoform suffixes from gene symbols, map human genes through HGNC and mouse genes through MGI/Ensembl GRCm39 annotations, then collapse both species into shared "orthology groups" using Ensembl Compara protein homology (release 115). Reference genomes are GRCh38 and GRCm39, Ensembl release 113. Perturbation targets go through the same pipeline, so a knockout in human and the matching knockout in mouse land on the same token; unmappable targets and genes are dropped. Each cell also gets standardized metadata tokens covering species, cell type, cell line, lineage, perturbation technology, disease state, and stimulation condition.

Model. The architecture has four parts:

• An encoder that turns a cell into a vector. Only expressed genes are kept (a sparsity trick), values are log(1+x) normalized, and each gene's expression magnitude is injected via FiLM conditioning — a small MLP produces per-gene scale and shift parameters applied to the gene embedding. Metadata tokens share one embedding table with genes but carry a separate "role" embedding. No positional encodings are used, because gene sets are unordered; the cell vector is the mean of all contextualized token states.
• A perturbation embedder that averages the embeddings of the perturbed genes, so multi-gene perturbations are handled permutation-invariantly.
• A perturbation operator: a small MLP that takes the control cell vector concatenated with the perturbation vector and outputs an additive shift to the latent state.
• A query-based decoder: to predict any gene's expression, it adds that gene's embedding to the latent state and maps the result to a number, so arbitrary gene subsets can be predicted.

Training. Stage 1 masks a random subset of expressed genes in control cells and trains the model to recover their values (MSE). Stage 2 samples a matched control cell for each perturbed cell, then does two things at once: (a) predicts post-perturbation expression for a sampled gene subset (MSE), and (b) hides part of the perturbation identity, computes the latent difference between control and perturbed cells, and asks an alignment head to recover the hidden perturbation embedding — trained with in-batch InfoNCE so the predicted embedding matches its own perturbation better than others in the batch. The two losses are summed with a weighting term λ_align.

Evaluation. Five representative datasets are excluded from pretraining, then CytoWave is finetuned on each held-out dataset's training split and tested on its test split. Results are averaged over three random seeds (0, 1, 2). Baselines span statistical predictors (Linear, TrainMean), basic ML models (BaseReg, BaseMLP), supervised perturbation models (CPA, GEARS), pretrained single-cell models (scGPT, Geneformer, LPM), embedding-transfer methods (GenePert, scouter), and a generative model (scLDM).

Not reported in the provided content: the latent dimension d, number of Transformer layers L, masking ratios ρ_x, ρ_q, ρ_u, the λ_align value, optimizer settings, and the exact train/test splits — the appendix is truncated mid-sentence at the perturbation-harmonization section.

Why This Matters

Research impact. Most single-cell foundation models are pretrained on unperturbed atlases and only see perturbations during finetuning. This paper is a direct test of the alternative: put perturbation-response prediction into the pretraining objective itself. Its dataset-held-out protocol (rather than the more common within-dataset splits) is a stricter test of generalization, and the orthology-aligned corpus is a reusable artifact for the field.

Real-world applications:

• Gene function discovery: predicting a gene's downstream transcriptional consequences without running a wet-lab screen for every target.
• Therapeutic target prioritization: simulating knockout or knockdown effects in relevant cell types before committing to experimental validation.
• Cross-species translation: because perturbations are mapped through orthology groups, a response characterized in mouse can be projected onto human contexts, and vice versa.
• Pathway and mechanism interpretation: the AUCell case study suggests predicted responses can be read out at the level of hallmark pathway activity, which is closer to how biologists reason about mechanism than per-gene error.

Industry relevance. Pharma and biotech groups running CRISPR screens face high per-experiment cost; a model that transfers across cell lines, protocols, and species could reduce the number of screens needed and help triage which perturbations are worth testing. The explicit handling of stimulation conditions and disease state in the metadata vocabulary is relevant to immunology and oncology drug discovery, where context dependence dominates.

Future Directions

• Multi-gene and combinatorial perturbations. The formulation accepts perturbation sets, but all benchmark results in the paper are for single-gene perturbations; whether the operator composes correctly for combinations is untested here.
• Closing the gap on correlation metrics. CytoWave's advantage over CPA is clear in MAE but marginal in ΔPCC and ΔPCC20 on several datasets, raising the question of whether the latent additive-shift assumption is the limiting factor or whether the metrics themselves are saturating.
• Richer perturbation encodings. The current perturbation embedder is a simple average of gene embeddings; incorporating gene–gene interaction structure (as GEARS attempts with networks) or regulatory priors could sharpen responses that depend on pathway context rather than target identity alone.
• Uncertainty and population-level modeling. The paper evaluates cell-level point predictions; the related work highlights optimal-transport, flow-based, and probabilistic approaches that model response distributions, and combining those with perturbation-centric pretraining is an open direction.

Target Audience

Primary: machine learning researchers working on single-cell foundation models, perturbation modeling, or self-supervised pretraining with contrastive and masked-reconstruction objectives; computational biologists building perturbation-response predictors.

Secondary: wet-lab and translational scientists who run Perturb-seq or CRISPR screens and want to know whether current models transfer to their own unseen datasets; bioinformatics engineers who need a curated, orthology-aligned multi-study corpus.

Less suitable for: beginners, since the paper presumes comfort with Transformer internals, FiLM conditioning, InfoNCE, and single-cell data conventions, and does not provide introductory background on Perturb-seq technology or expression normalization.

Authors’ abstract

Predicting gene expression under perturbations from single-cell data is a fundamental task for modeling cellular responses, with applications in discovering gene function and designing therapeutics. However, existing methods often struggle to generalize across datasets due to differences in experimental protocols, biological contexts, and species. We present \textbf{CytoWave}, a perturbation-centric pretraining framework for single-cell genetic perturbation response prediction. We leverage large-scale cross-species perturbation datasets unified via orthology and train CytoWave in a two-stage manner. In the first stage, we learn reference control cell states via masked gene expression recovery on unperturbed cells. In the second stage, we pretrain the model to predict perturbed gene expression given control states and perturbation signals, together with a contrastive objective that aligns latent representations of perturbation responses with perturbation signals. We perform strict dataset-held-out evaluation across multiple condition combinations and case studies of pathway activity changes, demonstrating consistent and effective performance of CytoWave in predicting single-cell genetic perturbation responses.