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

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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 hctrlh^{ctrl} , hperth^{pert} , and candidate embeddings hkcxth_{k}^{cxt} . The Transformer Generator is frozen in the STATE-based experiments and trained in the scGPT-based ones.
 GenerationNaïve RAGPT-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.
Figure 3
MetricSTATESTATE+GenePTNaïve RAGPT-RAG
Gene-level expression correlations
Pearson DEG ↑\uparrow0.624†0.624^{\dagger}0.6310.6310.396†0.396^{\dagger}0.633
Spearman DEG ↑\uparrow0.403†0.403^{\dagger}0.4110.4110.307†0.307^{\dagger}0.412\mathbf{0.412}
Expression reconstruction accuracy
MSE PCA50 ↓\downarrow8.438.438.428.4212.64†12.64^{\dagger}8.39\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 ( pFDR<0.01p_{\mathrm{FDR}}<0.01 ), ††\dagger\dagger ( pFDR<0.05p_{\mathrm{FDR}}<0.05 ), †††\dagger\dagger\dagger ( pFDR<0.1p_{\mathrm{FDR}}<0.1 ). Bold: best. Results with standard deviations in Table 6 .
MetricscGPTscGPT+GenePTscGPT+PT-RAG
Gene-level expression correlations
Pearson DEG ↑\uparrow0.92090.91670.9128
Spearman DEG ↑\uparrow0.9500 †0.9509 †0.9368
Expression reconstruction accuracy
MSE PCA50 ↓\downarrow18.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 ( pFDR<0.01p_{\mathrm{FDR}}<0.01 ), ††\dagger\dagger ( pFDR<0.05p_{\mathrm{FDR}}<0.05 ), †††\dagger\dagger\dagger ( pFDR<0.1p_{\mathrm{FDR}}<0.1 ). Bold: best. Results with standard deviations in Table 7 .
MetricSTATE+GenePTPT-RAG
Gene-level expression correlations
Pearson DEG ↑\uparrow0.6377\mathbf{0.6377}0.63140.6314
Spearman DEG ↑\uparrow0.4293\mathbf{0.4293}0.42240.4224
Expression reconstruction accuracy
MSE PCA50 ↓\downarrow1.1477†1.1477^{\dagger}1.1316\mathbf{1.1316}
Distributional similarity (PCA space)
Table 4: Cross-perturbation generalization. Results averaged across three folds. Significance vs. PT-RAG: †\dagger ( pFDR<0.01p_{\mathrm{FDR}}<0.01 ), ††\dagger\dagger ( pFDR<0.05p_{\mathrm{FDR}}<0.05 ), †††\dagger\dagger\dagger ( pFDR<0.1p_{\mathrm{FDR}}<0.1 ). Bold: best. Results with standard deviations in Table 8 .
MetricSTATE+GenePTPT-RAG
Gene-level expression correlations
Pearson DEG ↑\uparrow0.2985 †0.3856
Spearman DEG ↑\uparrow0.2760 †0.3564
Expression reconstruction accuracy
MSE PCA50 ↓\downarrow5.44825.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 ( pFDR<0.01p_{\mathrm{FDR}}<0.01 ), ††\dagger\dagger ( pFDR<0.05p_{\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)(0,1) : f(x)=e−x−e−xf(x)=e^{-x-e^{-x}} .
MetricSTATESTATE+GenePTNaïve RAGPT-RAG
Pearson DEG ↑\uparrow0.624±0.0480.624\pm 0.0480.631±0.0510.631\pm 0.0510.396±0.0630.396\pm 0.0630.633±0.048\mathbf{0.633\pm 0.048}
Spearman DEG ↑\uparrow0.403±0.0460.403\pm 0.0460.411±0.0520.411\pm 0.0520.307±0.0410.307\pm 0.0410.412±0.051\mathbf{0.412\pm 0.051}
MSE PCA50 ↓\downarrow8.43±0.938.43\pm 0.938.42±0.888.42\pm 0.8812.64±1.0412.64\pm 1.048.39±0.87\mathbf{8.39\pm 0.87}
W1W_{1} ↓\downarrow35.70±1.7635.70\pm 1.7635.53±1.6835.53\pm 1.6848.48±1.7148.48\pm 1.7135.41±1.62\mathbf{35.41\pm 1.62}
W2W_{2} ↓\downarrow646.1±63.6646.1\pm 63.6638.7±60.6638.7\pm 60.61189.5±83.31189.5\pm 83.3633.7±58.4\mathbf{633.7\pm 58.4}
Energy ↓\downarrow9.41±1.189.41\pm 1.189.40±1.159.40\pm 1.1514.18±1.1714.18\pm 1.179.33±1.15\mathbf{9.33\pm 1.15}
Appendix
Table 6: Cross-cell-type: full results with standard deviations. n=1,635n=1{,}635 (375 HepG2, 416 RPE1, 443 Jurkat, 401 K562).
MetricscGPTscGPT+GenePTscGPT+PT-RAG
Gene-level expression correlations
Pearson DEG ↑\uparrow0.9209 ±\pm 0.01580.9167 ±\pm 0.01490.9128 ±\pm 0.0146
Spearman DEG ↑\uparrow0.9500 ±\pm 0.03410.9509 ±\pm 0.03060.9368 ±\pm 0.0299
Expression reconstruction accuracy
MSE PCA50 ↓\downarrow18.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 ( pFDR<0.01p_{\mathrm{FDR}}<0.01 ), ††\dagger\dagger ( pFDR<0.05p_{\mathrm{FDR}}<0.05 ), †††\dagger\dagger\dagger ( pFDR<0.1p_{\mathrm{FDR}}<0.1 ). Bold: best.
MetricSTATE+GenePTPT-RAG
Gene-level expression correlations
Pearson DEG ↑\uparrow0.6377±0.0367\mathbf{0.6377\pm 0.0367}0.6314±0.03940.6314\pm 0.0394
Spearman DEG ↑\uparrow0.4293±0.0344\mathbf{0.4293\pm 0.0344}0.4224±0.03810.4224\pm 0.0381
Expression reconstruction accuracy
MSE PCA50 ↓\downarrow1.1477±0.2243†1.1477\pm 0.2243^{\dagger}1.1316±0.2150\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 ( pFDR<0.01p_{\mathrm{FDR}}<0.01 ), ††\dagger\dagger ( pFDR<0.05p_{\mathrm{FDR}}<0.05 ), †††\dagger\dagger\dagger ( pFDR<0.1p_{\mathrm{FDR}}<0.1 ). Bold: best.
MetricSTATESTATE+GenePTNaïve RAGPT-RAG
Pearson DEG ↑\uparrow0.596±0.0510.596\pm 0.0510.595±0.0550.595\pm 0.0550.351±0.0540.351\pm 0.0540.604±0.049\mathbf{0.604\pm 0.049}
Spearman DEG ↑\uparrow0.398±0.0440.398\pm 0.0440.394±0.0580.394\pm 0.0580.289±0.0360.289\pm 0.0360.401±0.049\mathbf{0.401\pm 0.049}
MSE PCA50 ↓\downarrow8.89±0.328.89\pm 0.328.86±0.328.86\pm 0.3213.35±0.7713.35\pm 0.778.83±0.31\mathbf{8.83\pm 0.31}
W1W_{1} ↓\downarrow36.30±1.3436.30\pm 1.3435.53±1.2035.53\pm 1.2048.68±1.3148.68\pm 1.3135.41±1.18\mathbf{35.41\pm 1.18}
W2W_{2} ↓\downarrow668.5±49.5668.5\pm 49.5637.7±43.3637.7\pm 43.31198.5±64.51198.5\pm 64.5631.8±42.3\mathbf{631.8\pm 42.3}
Energy ↓\downarrow9.68±0.699.68\pm 0.699.54±0.679.54\pm 0.6714.25±0.9014.25\pm 0.909.48±0.67\mathbf{9.48\pm 0.67}
Appendix
Table 9: Results on HepG2 ( n=375n=375 ).
MetricSTATESTATE+GenePTNaïve RAGPT-RAG
Pearson DEG ↑\uparrow0.621±0.0430.621\pm 0.0430.633±0.0430.633\pm 0.0430.450±0.0520.450\pm 0.0520.640±0.042\mathbf{0.640\pm 0.042}
Spearman DEG ↑\uparrow0.419±0.0350.419\pm 0.0350.432±0.0330.432\pm 0.0330.342±0.0270.342\pm 0.0270.438±0.033\mathbf{0.438\pm 0.033}
MSE PCA50 ↓\downarrow9.56±0.359.56\pm 0.359.48±0.329.48\pm 0.3213.29±0.8013.29\pm 0.809.44±0.33\mathbf{9.44\pm 0.33}
W1W_{1} ↓\downarrow36.75±1.2536.75\pm 1.2536.62±1.2436.62\pm 1.2449.31±1.4649.31\pm 1.4636.53±1.22\mathbf{36.53\pm 1.22}
W2W_{2} ↓\downarrow682.5±47.1682.5\pm 47.1676.3±46.4676.3\pm 46.41227.5±72.91227.5\pm 72.9672.8±45.8\mathbf{672.8\pm 45.8}
Energy ↓\downarrow10.17±0.8410.17\pm 0.8410.14±0.8410.14\pm 0.8414.56±0.8114.56\pm 0.8110.08±0.84\mathbf{10.08\pm 0.84}
Appendix
Table 10: Results on RPE1 ( n=416n=416 ).
MetricSTATESTATE+GenePTNaïve RAGPT-RAG
Pearson DEG ↑\uparrow0.636±0.0480.636\pm 0.0480.652±0.049\mathbf{0.652\pm 0.049}0.398±0.0560.398\pm 0.0560.649±0.0480.649\pm 0.048
Spearman DEG ↑\uparrow0.392±0.0560.392\pm 0.0560.424±0.0580.424\pm 0.0580.318±0.0330.318\pm 0.0330.424±0.054\mathbf{0.424\pm 0.054}
MSE PCA50 ↓\downarrow7.35±0.41\mathbf{7.35\pm 0.41}7.36±0.417.36\pm 0.4111.46±0.6211.46\pm 0.627.36±0.417.36\pm 0.41
W1W_{1} ↓\downarrow35.94±1.5235.94\pm 1.5236.10±1.5036.10\pm 1.5049.06±1.5049.06\pm 1.5035.82±1.48\mathbf{35.82\pm 1.48}
W2W_{2} ↓\downarrow654.4±56.7654.4\pm 56.7660.1±56.1660.1\pm 56.11219.6±75.41219.6\pm 75.4649.9±55.1\mathbf{649.9\pm 55.1}
Energy ↓\downarrow8.15±1.018.15\pm 1.018.23±1.028.23\pm 1.0213.12±1.1013.12\pm 1.108.14±1.01\mathbf{8.14\pm 1.01}
Appendix
Table 11: Results on Jurkat ( n=443n=443 ).
MetricSTATESTATE+GenePTNaïve RAGPT-RAG
Pearson DEG ↑\uparrow0.639±0.037\mathbf{0.639\pm 0.037}0.638±0.0370.638\pm 0.0370.381±0.0460.381\pm 0.0460.633±0.0410.633\pm 0.041
Spearman DEG ↑\uparrow0.406±0.040\mathbf{0.406\pm 0.040}0.391±0.0440.391\pm 0.0440.274±0.0290.274\pm 0.0290.383±0.0470.383\pm 0.047
MSE PCA50 ↓\downarrow8.02±0.33\mathbf{8.02\pm 0.33}8.07±0.328.07\pm 0.3212.61±0.5912.61\pm 0.598.04±0.338.04\pm 0.33
W1W_{1} ↓\downarrow33.77±1.2133.77\pm 1.2133.76±1.18\mathbf{33.76\pm 1.18}46.77±1.2346.77\pm 1.2333.80±1.1533.80\pm 1.15
W2W_{2} ↓\downarrow578.1±42.4578.1\pm 42.4576.9±41.3\mathbf{576.9\pm 41.3}1108.5±58.91108.5\pm 58.9577.1±40.2577.1\pm 40.2
Energy ↓\downarrow9.73±0.91\mathbf{9.73\pm 0.91}9.80±0.929.80\pm 0.9214.89±0.9714.89\pm 0.979.74±0.919.74\pm 0.91
Appendix
Table 12: Results on K562 ( n=401n=401 ).
Figure 6: Sparsity regularization sensitivity (gene-level and reconstruction metrics). HepG2, K=32K=32 .
Figure 7: Sparsity regularization sensitivity (distributional metrics and retrieval count). HepG2, K=32K=32 .
Figure 8: K=16K=16 vs. K=32K=32 (gene-level and reconstruction metrics). HepG2.
Figure 9: K=16K=16 vs. K=32K=32 (distributional metrics). HepG2.
Figure 10: Naïve RAG K∈{2,5,10,32}K\in\{2,5,10,32\} vs. PT-RAG ( K=32K=32 ). Gene-level and reconstruction. HepG2.
Figure 11: Naïve RAG K∈{2,5,10,32}K\in\{2,5,10,32\} vs. PT-RAG ( K=32K=32 ). Distributional metrics. HepG2.
ModelPearson ↑\uparrowSpearman ↑\uparrowMSE PCA50 ↓\downarrowW1W_{1} ↓\downarrowW2W_{2} ↓\downarrowEnergy ↓\downarrow
STATE+GenePT (no retrieval)0.64090.43141.30560.81660.71730.3618
Retrieval pool available at test time
PT-RAG, full pool (2,008)0.62420.41391.25560.80730.69390.3511
PT-RAG, train+val pool (802)0.62370.41351.25520.80720.69370.3511
PT-RAG, train-only pool (669)0.62370.41351.25510.80720.69360.3510
Contribution of the selector and robustness
Appendix
Table 13: Additional controls, cross-perturbation (one fold, K=32K=32 ). See the text for the description of each row.
MetricK=32K=32K=128K=128
Pearson DEG ↑\uparrow0.63260.6314
Spearman DEG ↑\uparrow0.41210.4072
MSE PCA50 ↓\downarrow8.3928.368
W1W_{1} ↓\downarrow35.41135.406
W2W_{2} ↓\downarrow633.74633.67
Energy ↓\downarrow9.3339.308
Appendix
Table 14: Candidate-set size, cross-cell-type (pooled over the four cell types, n=1,635n=1{,}635 ).
ModelParametersFLOPs/BatchFLOPs/Cell
STATE20,936,4801.67B34.9M
STATE+GenePT20,874,0161.67B34.8M
Naïve RAG21,120,2881.67B34.8M
STATE+PT-RAG20,973,6022.86B59.9M
Appendix
Table 15: Computational costs on Cross-Cell type generalization (STATE). FLOPs per batch (64 cells).
ModelParametersFLOPs/Batch (64 cells)FLOPs/Cell
STATE+GenePT20,874,0162.23B34.87M
PT-RAG K=128K=128 , sparsity =1.0=1.020,973,6027.46B116.61M
Appendix
Table 16: Computational costs on Cross-Perturbation generalization. FLOPs per batch (64 cells).
ModelParametersFLOPs/Batch (64 cells)FLOPs/Cell
scGPT51,334,6584.15T64.88B
scGPT+GenePT52,121,6024.15T64.88B
scGPT+PT-RAG ( K=128K=128 , sparsity =1.0=1.0 )53,699,5884.19T65.49B
Appendix
Table 17: Computational costs on Cross-Cell type generalization (scGPT). FLOPs per batch (64 cells).
Metricvs. STATEvs. STATE+GenePT
Pearson DEG2.44×10−8  †2.44\times 10^{-8\;\dagger}0.5900.590
Spearman DEG4.89×10−10  †4.89\times 10^{-10\;\dagger}0.7850.785
MSE PCA500.5770.5770.5900.590
W1W_{1}3.41×10−6  †3.41\times 10^{-6\;\dagger}0.082†††0.082^{\dagger\dagger\dagger}
W2W_{2}5.47×10−8  †5.47\times 10^{-8\;\dagger}0.041††0.041^{\dagger\dagger}
Energy0.082†††0.082^{\dagger\dagger\dagger}0.1860.186
Appendix
Table 18: FDR-corrected p-values for PT-RAG vs. baselines (cross-cell-type). n=1,635n=1{,}635 .

Explore similar work

CardsList

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

    Jun 10, 2026Wei Zhang, Xun Jiang, Yuesi Xi +1Single-Cell Perturbation ModelingGenomics

    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.

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

    Jun 26, 2026Dongxia Wu, Mingyu Li, Yuhui Zhang +4Single-Cell Perturbation ModelingReinforcement Learning with Verifiable Rewards

    Peking University Beijing, China · Stanford University Stanford, CA

  3. LLM-Guided Retrieval for Prediction of Molecular Perturbation Responses

    Aug 3, 2026Betty Xiong, Jan-Christian Huetter, Gabriele Scalia +2Single-Cell Perturbation ModelingDrug Discovery

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