Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation

Andrea Giuseppe Di Francesco; Andrea Rubbi; Rishabh Jain; Pietro Liò

cs.LGMar 7, 2026

Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation

Authors: Andrea Giuseppe Di Francesco, Andrea Rubbi, Rishabh Jain, Pietro Liò

Organizations: Sapienza University of Rome, Rome, Italy · ISTI-CNR, Institute of Information Science and Technologies, Pisa, Italy · University of Cambridge, Cambridge, United Kingdom · Wellcome Sanger Institute, Cambridge, United Kingdom

Abstract

Predicting transcriptional responses to genetic perturbations is fundamental to functional genomics and therapeutic discovery. Recent deep learning models have shown promise in single-cell perturbation response prediction, but they typically generate each response in isolation, without explicitly leveraging related perturbations. We introduce PT-RAG (Perturbation-aware Two-stage Retrieval-Augmented Generation), a plug-in retrieval-and-conditioning module for generative cellular perturbation response. PT-RAG augments an existing perturbation-response backbone with learned access to related perturbation contexts. The key challenge is that relevance is not fixed in this setting: functionally related genes may elicit different effects across cell types. PT-RAG addresses this with a two-stage retrieval mechanism: GenePT-based semantic retrieval first identifies K candidate perturbations, after which a differentiable Gumbel-Softmax selector adaptively selects retrieved contexts conditioned on the control cell state, the query perturbation, and each candidate perturbation. We evaluate PT-RAG on two backbones, a STATE-style generator used as a frozen random reservoir and a fully trained scGPT, across cross-cell-type and cross-perturbation generalization tasks. PT-RAG consistently improves distributional similarity and often overall predictive quality; for example, on scGPT cross-cell-type results, the 2-Wasserstein distance drops by 5.9% relative to scGPT alone. The code to reproduce our experiments is available at https://github.com/difra100/PT-RAG_NIPS.

Figures & tables

Figure 1: Comparison of architectures. Left : Generation baseline and Naïve RAG (non-differentiable retrieval, dotted lines; frozen components indicated by ice cube icon). Right : PT-RAG uses two-stage differentiable retrieval with Gumbel-Softmax selection conditioned on $h^{ctrl}$ , $h^{pert}$ , and candidate embeddings $h_{k}^{cxt}$ . The Transformer Generator is frozen in the STATE-based experiments and trained in the scGPT-based ones.

	Generation	Naïve RAG	PT-RAG
Retrieval	✗	✓	✓
Cell-type-aware Retrieval	✗	✗	✓

Table 1: Method comparison. All methods except PT-RAG use non-differentiable or no retrieval; PT-RAG uses two-stage differentiable selection.

Gene-level expression correlations
Metric	STATE	STATE+GenePT	Naïve RAG	PT-RAG
Pearson DEG $\uparrow$	$0.624^{\dagger}$	$0.631$	$0.396^{\dagger}$	0.633
Spearman DEG $\uparrow$	$0.403^{\dagger}$	$0.411$	$0.307^{\dagger}$	$\mathbf{0.412}$
Expression reconstruction accuracy
MSE PCA50 $\downarrow$	$8.43$	$8.42$	$12.64^{\dagger}$	$\mathbf{8.39}$
Distributional similarity (PCA space)

Table 2: Cross-cell-type generalization. Mean values across 1,635 test perturbations from four cell types. Significance vs. PT-RAG:

\dagger

(

p_{\mathrm{FDR}}<0.01

\dagger\dagger

(

p_{\mathrm{FDR}}<0.05

\dagger\dagger\dagger

(

p_{\mathrm{FDR}}<0.1

). Bold: best. Results with standard deviations in Table 6 .

Gene-level expression correlations
Metric	scGPT	scGPT+GenePT	scGPT+PT-RAG
Pearson DEG $\uparrow$	0.9209	0.9167	0.9128
Spearman DEG $\uparrow$	0.9500 †	0.9509 †	0.9368
Expression reconstruction accuracy
MSE PCA50 $\downarrow$	18.6440 †	19.8907 †	17.3212
Distributional similarity (PCA space)

Table 3: Cross-cell-type generalization: Results on scGPT. Mean values across 1,635 test perturbations from four cell types. Significance vs. PT-RAG:

\dagger

(

p_{\mathrm{FDR}}<0.01

\dagger\dagger

(

p_{\mathrm{FDR}}<0.05

\dagger\dagger\dagger

(

p_{\mathrm{FDR}}<0.1

). Bold: best. Results with standard deviations in Table 7 .

Gene-level expression correlations
Metric	STATE+GenePT	PT-RAG
Pearson DEG $\uparrow$	$\mathbf{0.6377}$	$0.6314$
Spearman DEG $\uparrow$	$\mathbf{0.4293}$	$0.4224$
Expression reconstruction accuracy
MSE PCA50 $\downarrow$	$1.1477^{\dagger}$	$\mathbf{1.1316}$
Distributional similarity (PCA space)

Table 4: Cross-perturbation generalization. Results averaged across three folds. Significance vs. PT-RAG:

\dagger

(

p_{\mathrm{FDR}}<0.01

\dagger\dagger

(

p_{\mathrm{FDR}}<0.05

\dagger\dagger\dagger

(

p_{\mathrm{FDR}}<0.1

). Bold: best. Results with standard deviations in Table 8 .

Gene-level expression correlations
Metric	STATE+GenePT	PT-RAG
Pearson DEG $\uparrow$	0.2985 †	0.3856
Spearman DEG $\uparrow$	0.2760 †	0.3564
Expression reconstruction accuracy
MSE PCA50 $\downarrow$	5.4482	5.2871
Distributional similarity (PCA space)

Table 5: Cross-perturbation generalization on Norman et al. (CRISPRa). Mean values across the 63 test perturbations of one fold. Significance vs. PT-RAG:

\dagger

(

p_{\mathrm{FDR}}<0.01

\dagger\dagger

(

p_{\mathrm{FDR}}<0.05

). Bold : best.

Figure 4: Jaccard similarity of top-10 retrieved perturbations across cell types. Off-diagonal values (range 0.185–0.196, mean 0.191) confirm that only $\sim$ 19% of selected perturbations overlap between cell types for the same query gene.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: PDF of Gumbel $(0,1)$ : $f(x)=e^{-x-e^{-x}}$ .

Metric	STATE	STATE+GenePT	Naïve RAG	PT-RAG
Pearson DEG $\uparrow$	$0.624\pm 0.048$	$0.631\pm 0.051$	$0.396\pm 0.063$	$\mathbf{0.633\pm 0.048}$
Spearman DEG $\uparrow$	$0.403\pm 0.046$	$0.411\pm 0.052$	$0.307\pm 0.041$	$\mathbf{0.412\pm 0.051}$
MSE PCA50 $\downarrow$	$8.43\pm 0.93$	$8.42\pm 0.88$	$12.64\pm 1.04$	$\mathbf{8.39\pm 0.87}$
$W_{1}$ $\downarrow$	$35.70\pm 1.76$	$35.53\pm 1.68$	$48.48\pm 1.71$	$\mathbf{35.41\pm 1.62}$
$W_{2}$ $\downarrow$	$646.1\pm 63.6$	$638.7\pm 60.6$	$1189.5\pm 83.3$	$\mathbf{633.7\pm 58.4}$
Energy $\downarrow$	$9.41\pm 1.18$	$9.40\pm 1.15$	$14.18\pm 1.17$	$\mathbf{9.33\pm 1.15}$

Appendix

Table 6: Cross-cell-type: full results with standard deviations.

n=1{,}635

(375 HepG2, 416 RPE1, 443 Jurkat, 401 K562).

Gene-level expression correlations
Metric	scGPT	scGPT+GenePT	scGPT+PT-RAG
Pearson DEG $\uparrow$	0.9209 $\pm$ 0.0158	0.9167 $\pm$ 0.0149	0.9128 $\pm$ 0.0146
Spearman DEG $\uparrow$	0.9500 $\pm$ 0.0341	0.9509 $\pm$ 0.0306	0.9368 $\pm$ 0.0299
Expression reconstruction accuracy
MSE PCA50 $\downarrow$	18.6440 $\pm$ 2.0080 †	19.8907 $\pm$ 2.8476 †	17.3212 $\pm$ 2.1105
Distributional similarity (PCA space)

Appendix

Table 7: Cross-cell-type generalization on scGPT: full results with standard deviations. Mean values across 1,635 test perturbations from four cell types. Significance vs. PT-RAG:

\dagger

(

p_{\mathrm{FDR}}<0.01

\dagger\dagger

(

p_{\mathrm{FDR}}<0.05

\dagger\dagger\dagger

(

p_{\mathrm{FDR}}<0.1

). Bold: best.

Gene-level expression correlations
Metric	STATE+GenePT	PT-RAG
Pearson DEG $\uparrow$	$\mathbf{0.6377\pm 0.0367}$	$0.6314\pm 0.0394$
Spearman DEG $\uparrow$	$\mathbf{0.4293\pm 0.0344}$	$0.4224\pm 0.0381$
Expression reconstruction accuracy
MSE PCA50 $\downarrow$	$1.1477\pm 0.2243^{\dagger}$	$\mathbf{1.1316\pm 0.2150}$
Distributional similarity (PCA space)

Appendix

Table 8: Cross-perturbation generalization: full results with standard deviations. Results averaged across three folds. Significance vs. PT-RAG:

\dagger

(

p_{\mathrm{FDR}}<0.01

\dagger\dagger

(

p_{\mathrm{FDR}}<0.05

\dagger\dagger\dagger

(

p_{\mathrm{FDR}}<0.1

). Bold: best.

Metric	STATE	STATE+GenePT	Naïve RAG	PT-RAG
Pearson DEG $\uparrow$	$0.596\pm 0.051$	$0.595\pm 0.055$	$0.351\pm 0.054$	$\mathbf{0.604\pm 0.049}$
Spearman DEG $\uparrow$	$0.398\pm 0.044$	$0.394\pm 0.058$	$0.289\pm 0.036$	$\mathbf{0.401\pm 0.049}$
MSE PCA50 $\downarrow$	$8.89\pm 0.32$	$8.86\pm 0.32$	$13.35\pm 0.77$	$\mathbf{8.83\pm 0.31}$
$W_{1}$ $\downarrow$	$36.30\pm 1.34$	$35.53\pm 1.20$	$48.68\pm 1.31$	$\mathbf{35.41\pm 1.18}$
$W_{2}$ $\downarrow$	$668.5\pm 49.5$	$637.7\pm 43.3$	$1198.5\pm 64.5$	$\mathbf{631.8\pm 42.3}$
Energy $\downarrow$	$9.68\pm 0.69$	$9.54\pm 0.67$	$14.25\pm 0.90$	$\mathbf{9.48\pm 0.67}$

Appendix

Table 9: Results on HepG2 (

n=375

Metric	STATE	STATE+GenePT	Naïve RAG	PT-RAG
Pearson DEG $\uparrow$	$0.621\pm 0.043$	$0.633\pm 0.043$	$0.450\pm 0.052$	$\mathbf{0.640\pm 0.042}$
Spearman DEG $\uparrow$	$0.419\pm 0.035$	$0.432\pm 0.033$	$0.342\pm 0.027$	$\mathbf{0.438\pm 0.033}$
MSE PCA50 $\downarrow$	$9.56\pm 0.35$	$9.48\pm 0.32$	$13.29\pm 0.80$	$\mathbf{9.44\pm 0.33}$
$W_{1}$ $\downarrow$	$36.75\pm 1.25$	$36.62\pm 1.24$	$49.31\pm 1.46$	$\mathbf{36.53\pm 1.22}$
$W_{2}$ $\downarrow$	$682.5\pm 47.1$	$676.3\pm 46.4$	$1227.5\pm 72.9$	$\mathbf{672.8\pm 45.8}$
Energy $\downarrow$	$10.17\pm 0.84$	$10.14\pm 0.84$	$14.56\pm 0.81$	$\mathbf{10.08\pm 0.84}$

Appendix

Table 10: Results on RPE1 (

n=416

Metric	STATE	STATE+GenePT	Naïve RAG	PT-RAG
Pearson DEG $\uparrow$	$0.636\pm 0.048$	$\mathbf{0.652\pm 0.049}$	$0.398\pm 0.056$	$0.649\pm 0.048$
Spearman DEG $\uparrow$	$0.392\pm 0.056$	$0.424\pm 0.058$	$0.318\pm 0.033$	$\mathbf{0.424\pm 0.054}$
MSE PCA50 $\downarrow$	$\mathbf{7.35\pm 0.41}$	$7.36\pm 0.41$	$11.46\pm 0.62$	$7.36\pm 0.41$
$W_{1}$ $\downarrow$	$35.94\pm 1.52$	$36.10\pm 1.50$	$49.06\pm 1.50$	$\mathbf{35.82\pm 1.48}$
$W_{2}$ $\downarrow$	$654.4\pm 56.7$	$660.1\pm 56.1$	$1219.6\pm 75.4$	$\mathbf{649.9\pm 55.1}$
Energy $\downarrow$	$8.15\pm 1.01$	$8.23\pm 1.02$	$13.12\pm 1.10$	$\mathbf{8.14\pm 1.01}$

Appendix

Table 11: Results on Jurkat (

n=443

Metric	STATE	STATE+GenePT	Naïve RAG	PT-RAG
Pearson DEG $\uparrow$	$\mathbf{0.639\pm 0.037}$	$0.638\pm 0.037$	$0.381\pm 0.046$	$0.633\pm 0.041$
Spearman DEG $\uparrow$	$\mathbf{0.406\pm 0.040}$	$0.391\pm 0.044$	$0.274\pm 0.029$	$0.383\pm 0.047$
MSE PCA50 $\downarrow$	$\mathbf{8.02\pm 0.33}$	$8.07\pm 0.32$	$12.61\pm 0.59$	$8.04\pm 0.33$
$W_{1}$ $\downarrow$	$33.77\pm 1.21$	$\mathbf{33.76\pm 1.18}$	$46.77\pm 1.23$	$33.80\pm 1.15$
$W_{2}$ $\downarrow$	$578.1\pm 42.4$	$\mathbf{576.9\pm 41.3}$	$1108.5\pm 58.9$	$577.1\pm 40.2$
Energy $\downarrow$	$\mathbf{9.73\pm 0.91}$	$9.80\pm 0.92$	$14.89\pm 0.97$	$9.74\pm 0.91$

Appendix

Table 12: Results on K562 (

n=401

Panel of Figure 6 — Figure 6: Sparsity regularization sensitivity (gene-level and reconstruction metrics). HepG2, $K=32$ .

Panel of Figure 7 — Figure 7: Sparsity regularization sensitivity (distributional metrics and retrieval count). HepG2, $K=32$ .

Panel of Figure 8 — Figure 8: $K=16$ vs. $K=32$ (gene-level and reconstruction metrics). HepG2.

Figure 9:

K=16

vs.

K=32

(distributional metrics). HepG2.

Panel of Figure 10 — Figure 10: Naïve RAG $K\in\{2,5,10,32\}$ vs. PT-RAG ( $K=32$ ). Gene-level and reconstruction. HepG2.

Panel of Figure 11 — Figure 11: Naïve RAG $K\in\{2,5,10,32\}$ vs. PT-RAG ( $K=32$ ). Distributional metrics. HepG2.

Retrieval pool available at test time
Model	Pearson $\uparrow$	Spearman $\uparrow$	MSE PCA50 $\downarrow$	$W_{1}$ $\downarrow$	$W_{2}$ $\downarrow$	Energy $\downarrow$
STATE+GenePT (no retrieval)	0.6409	0.4314	1.3056	0.8166	0.7173	0.3618
PT-RAG, full pool (2,008)	0.6242	0.4139	1.2556	0.8073	0.6939	0.3511
PT-RAG, train+val pool (802)	0.6237	0.4135	1.2552	0.8072	0.6937	0.3511
PT-RAG, train-only pool (669)	0.6237	0.4135	1.2551	0.8072	0.6936	0.3510
Contribution of the selector and robustness

Appendix

Table 13: Additional controls, cross-perturbation (one fold,

K=32

). See the text for the description of each row.

Metric	$K=32$	$K=128$
Pearson DEG $\uparrow$	0.6326	0.6314
Spearman DEG $\uparrow$	0.4121	0.4072
MSE PCA50 $\downarrow$	8.392	8.368
$W_{1}$ $\downarrow$	35.411	35.406
$W_{2}$ $\downarrow$	633.74	633.67
Energy $\downarrow$	9.333	9.308

Appendix

Table 14: Candidate-set size, cross-cell-type (pooled over the four cell types,

n=1{,}635

Model	Parameters	FLOPs/Batch	FLOPs/Cell
STATE	20,936,480	1.67B	34.9M
STATE+GenePT	20,874,016	1.67B	34.8M
Naïve RAG	21,120,288	1.67B	34.8M
STATE+PT-RAG	20,973,602	2.86B	59.9M

Appendix

Table 15: Computational costs on Cross-Cell type generalization (STATE). FLOPs per batch (64 cells).

Model	Parameters	FLOPs/Batch (64 cells)	FLOPs/Cell
STATE+GenePT	20,874,016	2.23B	34.87M
PT-RAG $K=128$ , sparsity $=1.0$	20,973,602	7.46B	116.61M

Appendix

Table 16: Computational costs on Cross-Perturbation generalization. FLOPs per batch (64 cells).

Model	Parameters	FLOPs/Batch (64 cells)	FLOPs/Cell
scGPT	51,334,658	4.15T	64.88B
scGPT+GenePT	52,121,602	4.15T	64.88B
scGPT+PT-RAG ( $K=128$ , sparsity $=1.0$ )	53,699,588	4.19T	65.49B

Appendix

Table 17: Computational costs on Cross-Cell type generalization (scGPT). FLOPs per batch (64 cells).

Metric	vs. STATE	vs. STATE+GenePT
Pearson DEG	$2.44\times 10^{-8\;\dagger}$	$0.590$
Spearman DEG	$4.89\times 10^{-10\;\dagger}$	$0.785$
MSE PCA50	$0.577$	$0.590$
$W_{1}$	$3.41\times 10^{-6\;\dagger}$	$0.082^{\dagger\dagger\dagger}$
$W_{2}$	$5.47\times 10^{-8\;\dagger}$	$0.041^{\dagger\dagger}$
Energy	$0.082^{\dagger\dagger\dagger}$	$0.186$

Appendix

Table 18: FDR-corrected p-values for PT-RAG vs. baselines (cross-cell-type).

n=1{,}635

Explore similar work

Jun 10, 2026q-bio.GN

CisTransCell: Single-Cell Perturbation Prediction via Gene Function, Regulatory Control, and Cellular Context

Predicting cellular transcriptional responses to genetic perturbations is a central problem in single-cell biology, especially in the zero-shot setting where the perturbed gene or gene combination is unseen during training. A major difficulty is that perturbation effects are not determined by expression state alone: they depend on how the perturbed gene product influences other genes and proteins, how those downstream factors act on cis-regulatory elements, and which regulatory programs are active in the current cell state. To better capture this biological complexity, we propose CisTransCell, a cell-conditioned multi-modal framework for single-cell perturbation prediction that augments each gene with two complementary priors: a regulatory-sequence prior that captures how the gene is controlled, and a coding-sequence prior that captures what the gene product does. By integrating these priors with cellular expression state, CisTransCell models perturbation response as a cascade from gene function to regulatory control to downstream transcriptional change. Experiments on benchmark single-cell perturbation datasets show that CisTransCell achieves strong performance in zero-shot perturbation prediction.

Wei Zhang, Xun Jiang, Yuesi Xi +1

1L3S Research Center, Leibniz Universität Hannover, Germany · School of Clinical Medicine & Laboratory Medicine, Jiangsu University, China · Institute for Information Processing (tnt), Leibniz Universität Hannover, Germany.

Single-Cell Perturbation Modeling Genomics Single-Cell Transcriptomics Molecular Representation Learning

Ask PDF

Jun 26, 2026cs.LG

PerturbCellRL: Verifier-Guided Reinforcement Learning for Single-Cell Perturbation Prediction

Single-cell perturbation models can reduce costly wet-lab screening by predicting how cells respond transcriptionally to interventions. While recent generative models improve population-level prediction, individual generated cells are not explicitly checked for biological consistency. We introduce PerturbCellRL, a reinforcement learning (RL) framework that post-trains a pretrained single-cell transcriptomic generator using a suite of cell-level verifiers as rewards. These verifiers define four rewards: Pearson top-k similarity, RMSE top-k proximity, DE Spearman, and Pathway activity. The Pathway activity verifier rewards cells whose pathway responses match known perturbation biology. We evaluate PerturbCellRL on multiple genetic and chemical perturbation benchmarks. Across these benchmarks, PerturbCellRL improves over the pretrained flow-matching generator on reward-aligned evaluation metrics and a held-out evaluation metric. Moreover, PerturbCellRL remains competitive with state-of-the-art methods on population-level metrics. Together, these results frame trustworthy single-cell prediction as verifier-guided generative alignment, moving beyond matching expression distributions toward predictions whose single-cell perturbation effects are explicitly checked for biological consistency.

Dongxia Wu, Mingyu Li, Yuhui Zhang +4

Peking University Beijing, China · Stanford University Stanford, CA

Single-Cell Perturbation Modeling Reinforcement Learning with Verifiable Rewards RL for Generative Models Single-Cell Transcriptomics

Ask PDF

Aug 3, 2026cs.LG

LLM-Guided Retrieval for Prediction of Molecular Perturbation Responses

Predicting transcriptomic responses to small-molecule perturbations across cell lines is central to drug discovery, but exhaustive profiling of drug-cell combinations is infeasible. We frame molecular perturbation prediction as retrieve-and-aggregate: approximate an unmeasured drug's response in a cell line by aggregating measured responses of a small set of biologically related compounds. We propose LLM-Guided Retrieval (LGR), where a large language model (LLM) ranks candidate neighbor drugs (restricted to those profiled in the target cell line); after which a fixed mean aggregator combines their observed expression deltas to form the prediction. We evaluate on the Tahoe-100M single-cell perturbation atlas under unseen-drug, unseen-cell-line, and open-world regimes. LGR consistently improves over drug mean, ChemCPA, and chemistry-based kNN baselines, with the strongest gains for unseen cell-line generalization, where it achieves higher correlation and lower error than mean baselines. Across settings, LGR improves directional (sign) accuracy of gene regulation, indicating better recovery of biologically meaningful perturbation effects even when magnitude-based metrics are similar. These results suggest that retrieval quality, rather than predictor complexity, is a key driver of zero-shot molecular perturbation prediction, and that LLMs can provide a useful biological prior when used as constrained retrieval modules.

Betty Xiong, Jan-Christian Huetter, Gabriele Scalia +2

Department of Biomedical Data Science Stanford University Stanford, CA 94305, USA · Biology Research & AI Development, Genentech DNA Way, South San Francisco, CA 94080, USA

Single-Cell Perturbation Modeling Drug Discovery

Ask PDF