PaperlayerView exact revision on arXiv

Source-linked AI summary

Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation

Andrea Giuseppe Di Francesco, Andrea Rubbi, Pietro Liò

arXiv:2603.07233v1cs.LGcs.IR

TL;DR

Predicting cellular responses to genetic perturbations remains difficult because relevant retrieval context and similarity metrics are not predefined across cell types. PT-RAG learns cell-type-aware retrieval and improves prediction over STATE across multiple metrics, with gains over STATE+GenePT concentrated in Wasserstein distances.

  • Problem

    Cellular perturbation modeling lacks established perturbation similarity metrics and must learn what retrieval context is relevant for generating high-dimensional cell distributions.

  • Method

    PT-RAG uses two-stage differentiable retrieval, narrowing candidates with GenePT embeddings before cell-type-aware Gumbel-Softmax selection conditioned on cellular state and perturbation identity.

  • Results

    Across 1635 test perturbations, PT-RAG significantly improved over STATE across multiple metrics, while gains over STATE+GenePT concentrated in Wasserstein distances.

  • Takeaways & Limitations

    The results indicate that naive fixed retrieval is unsuitable for perturbation biology, underscoring the necessity of differentiable, cell-type-aware retrieval.

  • Takeaways & Limitations

    PT-RAG requires approximately 1.7× more FLOPs per batch than baselines and is evaluated only on single-gene perturbations.

Abstract

from arXiv · show

Predicting how cells respond to genetic perturbations is fundamental to understanding gene function, disease mechanisms, and therapeutic development. While recent deep learning approaches have shown promise in modeling single-cell perturbation responses, they struggle to generalize across cell types and perturbation contexts due to limited contextual information during generation. We introduce PT-RAG (Perturbation-aware Two-stage Retrieval-Augmented Generation), a novel framework that extends Retrieval-Augmented Generation beyond traditional language-model applications to cellular biology. Unlike standard RAG systems designed for text retrieval with pre-trained LLMs, perturbation retrieval lacks established similarity metrics and requires learning what constitutes relevant context, making differentiable retrieval essential. PT-RAG addresses this through a two-stage pipeline: first, retrieving candidate perturbations $K$ using GenePT embeddings, then adaptively refining the selection through Gumbel-Softmax discrete sampling conditioned on both the cell state and the input perturbation. This cell-type-aware differentiable retrieval enables end-to-end optimization of the retrieval objective jointly with generation. On the Replogle-Nadig single-gene perturbation dataset, we demonstrate that PT-RAG outperforms both STATE and vanilla RAG under identical experimental conditions, with the strongest gains in distributional similarity metrics ($W_1$, $W_2$). Notably, vanilla RAG's dramatic failure is itself a key finding: it demonstrates that differentiable, cell-type-aware retrieval is essential in this domain, and that naive retrieval can actively harm performance. Our results establish retrieval-augmented generation as a promising paradigm for modelling cellular responses to gene perturbation. The code to reproduce our experiments is available at https://github.com/difra100/PT-RAG_ICLR.

1 INTRODUCTION

PT-RAG extends retrieval-augmented generation to single-cell perturbation response prediction through differentiable, cell-type-aware retrieval. The framework addresses limitations of naive retrieval by learning which related perturbations improve generation for each cellular context.

  • 1 INTRODUCTION: Existing perturbation-response methods rely only on control cell state and perturbation identity, limiting prediction for novel cell types without related-intervention supervision.The combinatorial space of perturbations and cell contexts also makes comprehensive experimental characterization infeasible.
  • 1 INTRODUCTION: PT-RAG is introduced as the first retrieval-augmented generation framework for modeling cellular responses to genetic perturbations.It extends RAG beyond classical domains to cellular response generation.
  • 1 INTRODUCTION: PT-RAG first retrieves K candidates using semantic similarity in GenePT embedding space, then applies Gumbel-Softmax selection conditioned on cell state for end-to-end learning.This two-stage pipeline learns which related perturbations are informative for each cellular context.
  • 1 INTRODUCTION: Vanilla RAG with fixed retrieval dramatically underperforms all baselines, showing that naive retrieval can actively hurt performance in cellular perturbation prediction.The result underscores the necessity of differentiable, cell-type-aware retrieval.
  • 1 INTRODUCTION: 19% overlap in selected perturbations across cell types for the same query gene provides quantitative evidence that PT-RAG learns cell-type-specific retrieval patterns.The retrieval patterns therefore vary substantially across cellular contexts.

2 RELATED WORK

Related perturbation-prediction methods generally rely only on cell state and perturbation identity, whereas PT-RAG extends STATE with retrieval-augmented generation. It also brings differentiable, learned, task-conditioned retrieval to cellular-response modeling, beyond prior text-based and biological RAG applications.

  • Perturbation prediction: Existing perturbation-prediction methods generate responses from cell state and perturbation identity without using knowledge from related perturbations.
  • Perturbation prediction: PT-RAG extends STATE, which models cell populations as sequences with distributional losses, by adding retrieval-augmented generation.
  • Differentiable RAG: PT-RAG applies differentiable RAG where relevant context must be learned from scratch and depends on cell type, unlike standard fixed retrieval in text domains.
  • RAG in biology: Prior biological RAG applications retrieve textual annotations with LLMs or augment protein encoders, but do not generate cellular responses.

3 METHOD

PT-RAG predicts distributions of cellular responses to single-gene perturbations using two-stage retrieval and cell-type-aware differentiable context selection. It combines GenePT-based candidate retrieval with Gumbel-Softmax selection, enabling retrieval and generation to be optimized jointly.

  • Task formulation: The task is to predict the distribution of perturbed-cell expression profiles from control-cell expression profiles and a perturbation identifier.
  • Generation baseline: The Generation baseline encodes control cells and perturbations, adds their representations, and passes the result to a Transformer Generator.The Cell Encoder is frozen, while the Perturbation Encoder is trainable.
  • Vanilla RAG: Vanilla RAG retrieves top-K perturbations using GenePT cosine similarity and integrates them with Cross-Attention, but retrieval is non-differentiable and cell-type agnostic.It provides identical context for a perturbation across cell types, preventing the model from learning which context improves generation.
  • PT-RAG: PT-RAG first narrows the perturbation space through semantic retrieval, then adaptively selects contexts with a differentiable, cell-type-aware mechanism conditioned on cellular state.The initial stage prunes approximately 2009 perturbations to K candidates, while the second stage selects relevant candidates using Straight-Through Gumbel-Softmax.
  • PT-RAG: Candidate selection uses triplets combining cell state, target perturbation, and candidate context, so relevance depends on all three components.The selected context representations are aggregated and passed to the Transformer Generator.
  • Training objective: PT-RAG is trained with energy distance for distributional prediction and an L1 sparsity penalty weighted by λsparse = 0.1.The sparsity term encourages selective retrieval and discourages selecting all candidates.

4 EXPERIMENTS

Experiments test PT-RAG on few-shot cross-cell-type perturbation prediction against STATE, Vanilla RAG, and STATE+GenePT. PT-RAG achieves the strongest overall results, while retrieval analyses show that effective context selection must be cell-type-aware and differentiable.

  • Baselines and metrics: PT-RAG is compared with STATE, Vanilla RAG using top-K = 32 GenePT retrieval without cell-type-aware selection, and STATE+GenePT.Performance is assessed using gene-level correlations, reconstruction errors, and distributional similarity metrics.
  • Main results: Vanilla RAG underperforms STATE, with Pearson 0.293 vs 0.624 and Spearman 0.220 vs 0.403, and varying K ∈{2, 5, 10, 32} does not rescue it.The results indicate that cell-type-agnostic retrieval can introduce irrelevant context or fail to support generation.
  • Main results: PT-RAG improves over STATE in Pearson 0.633 vs 0.624, Spearman 0.412 vs 0.403, MAE 0.295 vs 0.298, W1 35.41 vs 35.70, and W2 633.7 vs 646.1.The model combines cell-type-aware selection with sparsity regularization, and the reported improvements are statistically significant; Energy distance improves marginally significantly.
  • Retrieval behavior: PT-RAG retrieves different contexts across cell types, with off-diagonal Jaccard similarity ranging from 0.185–0.196 and averaging 0.191 for the same query gene.For WARS, all cell types retrieve aminoacyl-tRNA synthetases, but the specific selected perturbations differ markedly.

5 CONCLUSION … A.1 THE GUMBEL DISTRIBUTION AND THE GUMBEL-SOFTMAX TEMPERATURE

PT-RAG introduces differentiable, cell-type-aware retrieval for cellular perturbation-response prediction, addressing the failure of naive RAG and improving performance over STATE. The appendix supplies technical background, evaluation and implementation details, extended results, related work, and statistical analyses, including the Gumbel-Softmax mechanism underlying discrete retrieval.

  • 5 CONCLUSION: PT-RAG is presented as the first retrieval-augmented generation framework for predicting cellular responses to genetic perturbations.The framework targets perturbation biology, where relevant context is not predefined and depends critically on target cell type.
  • 5 CONCLUSION: PT-RAG retrieves candidates with GenePT embeddings, then selects cell-type-aware context using Gumbel-Softmax conditioned on cellular state and perturbation identity.This two-stage differentiable pipeline enables adaptive retrieval jointly aligned with generation.
  • 5 CONCLUSION: 1189.5 vs 646.1: Vanilla RAG’s W2 performance severely degrades relative to STATE across 1635 test perturbations.The result supports the conclusion that naive retrieval can harm perturbation-response prediction.
  • 5 CONCLUSION: PT-RAG incurs approximately 1.7× more FLOPs per batch than baselines because of scoring and Gumbel-Softmax, and currently focuses on single-gene perturbations.Future extensions include combinatorial perturbations, chemical compounds, and CRISPR activation/interference.
  • APPENDIX OVERVIEW: The appendix documents mathematical preliminaries, evaluation metrics, implementation specifications, extended related work, additional experiments, statistical tests, and LLM usage.Its extended results include cell-type-specific retrieval examples, cross-cell-type evaluations, disaggregated results, and sensitivity analyses.
  • A.1 THE GUMBEL DISTRIBUTION AND THE GUMBEL-SOFTMAX TEMPERATURE: The Gumbel distribution models maxima or minima and provides the noise distribution used by Gumbel-Softmax for differentiable approximation of categorical sampling.For X ∼ Gumbel(0, 1), the appendix gives f(x) = e−x−e−x.
  • A.1 THE GUMBEL DISTRIBUTION AND THE GUMBEL-SOFTMAX TEMPERATURE: As τ →0, Gumbel-Softmax approaches one-hot categorical sampling, whereas higher τ produces softer probability distributions.The temperature τ controls the sharpness of the differentiable approximation.
  • A.1 THE GUMBEL DISTRIBUTION AND THE GUMBEL-SOFTMAX TEMPERATURE: The straight-through estimator uses hard arg max selections during the forward pass while backpropagating through soft Gumbel-Softmax probabilities.This preserves discrete selections at inference while allowing gradients to train the scoring function end to end.

B EVALUATION METRICS … C.1 MODEL ARCHITECTURE

The paper evaluates perturbation-response predictions through gene-level correlations, reconstruction errors, and population-level distributional metrics, then specifies the shared and PT-RAG-specific model architecture. These evaluations span molecular alignment, expression accuracy, and preservation of cellular heterogeneity.

  • B.1 GENE-LEVEL EXPRESSION CORRELATIONS: DEGs are identified by Welch’s t-test comparing control and perturbed populations at p < 0.05.The evaluation focuses on genes whose expression changes significantly under perturbation.
  • B.1 GENE-LEVEL EXPRESSION CORRELATIONS: Pearson correlation measures linear agreement in DEG expression changes, while Spearman correlation evaluates their rank ordering.Pearson reflects direction and relative magnitude; Spearman assesses which genes are most strongly up- or down-regulated regardless of exact magnitudes.
  • B.2 EXPRESSION RECONSTRUCTION ACCURACY: Point-wise reconstruction accuracy is measured with MSE and MAE in the original gene-expression space.MSE and MAE quantify errors across cells and genes, with MAE reported on the original scale.
  • B.2 EXPRESSION RECONSTRUCTION ACCURACY: MSEPCA50 measures mean squared error in PCA space, and lower reconstruction-metric values indicate more accurate perturbed-cell-state recovery.MSEPCA50 captures overall distributional shift in the reduced representation.
  • B.3 DISTRIBUTIONAL SIMILARITY IN LOW-DIMENSIONAL SPACE: Predicted and true populations are projected into a 50-dimensional PCA space fit on training data before distributional evaluation.This projection addresses the high dimensionality and complex distributional structure of cell populations.
  • B.3 DISTRIBUTIONAL SIMILARITY IN LOW-DIMENSIONAL SPACE: W1 and W2 measure optimal-transport distances sensitive to distribution shape and support, while MMD compares distributions with an RBF kernel using the median heuristic.W1 is the 1-Wasserstein distance and W2 is the 2-Wasserstein distance.
  • B.3 DISTRIBUTIONAL SIMILARITY IN LOW-DIMENSIONAL SPACE: Energy distance compares between- and within-distribution pairwise distances and serves as the primary training objective.The distributional metrics evaluate population-level fidelity, including heterogeneity and cellular structure beyond average accuracy.

C.2 TRAINING CONFIGURATION … E EXTENDED RESULTS

The paper specifies reproducible training, retrieval, embedding, and computational settings, then situates PT-RAG within single-cell perturbation modeling and retrieval-augmented generation, identifying its novelty as cellular-response generation.

  • C.2 TRAINING CONFIGURATION: Adam optimization uses learning rate 10^-3 and weight decay 0.0005, with validation every 2,000 steps during a maximum 50,000-step schedule.These settings define the general training configuration.
  • C.3 PT-RAG-SPECIFIC HYPERPARAMETERS: PT-RAG retrieves K = 32 candidate perturbations without a cosine-similarity threshold, then applies Gumbel-Softmax with temperature τ = 0.5 and a straight-through estimator.The first-stage candidate count balances computational efficiency and context diversity; hard forward decisions remain differentiable backward.
  • C.3 PT-RAG-SPECIFIC HYPERPARAMETERS: The primary distributional training objective is energy distance.This objective is listed among the PT-RAG-specific Gumbel-Softmax settings.
  • C.4 GENEPT EMBEDDINGS AND PERTURBATION DATABASE: GenePT represents each gene with a 1536-dimensional unit-normalized embedding derived from GPT-3.5 encodings of NCBI gene descriptions.Cosine similarity is computed after normalization.
  • C.4 GENEPT EMBEDDINGS AND PERTURBATION DATABASE: The perturbation database contains GenePT embeddings for all 2,009 Replogle training perturbations, excluding the query perturbation during retrieval to prevent information leakage.For a query perturbation p, retrieval uses P \ {p}.
  • C.5 COMPUTATIONAL RESOURCES: Experiments ran on NVIDIA A100 GPUs with 40GB memory, and PT-RAG typically required 30,000-40,000 steps and approximately 8-10 hours per target cell type.The reported convergence and runtime characterize the computational requirements.
  • D.1 SINGLE-CELL PERTURBATION RESPONSE MODELING: Prior single-cell perturbation methods learn latent shifts, compositional mappings, graph-propagated effects, or distributional transformations, with limitations including absent cell-cell modeling, predefined networks, or context dependence.The related work discusses scGen, CPA, GEARS, CellOT, temporal optimal transport, and STATE.
  • D.2 RETRIEVAL-AUGMENTED GENERATION: Standard RAG retrieves text by query similarity, but PT-RAG extends retrieval-augmented generation to cellular response modeling, where perturbation embeddings and cell-distribution outputs differ fundamentally from text.The related work also notes non-differentiable standard retrieval and prior RAG applications in other modalities and text-based single-cell biology.

E.1 CELL-TYPE-SPECIFIC RETRIEVAL EXAMPLES

PT-RAG retrieves functionally coherent perturbations while adapting selections across Jurkat, HepG2, K562, and RPE1 cell types. Examples span aminoacyl-tRNA synthetases, RNA helicases, mitochondrial and cytoplasmic ribosomal proteins, and kinases.

  • WARS: PT-RAG retrieves aminoacyl-tRNA synthetases for WARS, but the specific top-5 selections differ across all four cell types.Jurkat: EARS2, DARS, SEPSECS, CTU2, VARS; HepG2: SARS2, GART, KARS, EARS2, TARS; K562: FARSB, KARS, FARS2, EPRS, DARS; RPE1: KARS, GART, TARS, QARS, AARS2.
  • DDX27: For DDX27, PT-RAG retrieves DDX/DHX-family helicases with substantial cell-specific selection and limited overlap, including GEMIN4 in Jurkat and HepG2.Jurkat: DHX15, GEMIN4, DHX16, DDX19A, DDX23; HepG2: EDC4, DDX56, DHX36, GEMIN4, DDX1; K562: DDX11, DDX1, DDX6, DHX36, DHX33; RPE1: EDC4, DDX10, DDX56, DHX29, DDX19A.
  • TPRKB: For TPRKB, kinase and regulatory-protein retrieval differs by cell type, with proliferative K562 and RPE1 favoring cell-division-related kinases.Examples include TTK, PLK1, CDK2, and PKMYT1 across the cell types.
  • Overall pattern: Across examples, differentiable retrieval preserves family or pathway coherence while adapting perturbation selection to cellular context.The examples are presented as evidence that cell-type-aware retrieval is essential for effective perturbation prediction.

E.2 CROSS-CELL-TYPE EVALUATION WITH STANDARD DEVIATIONS

Table 2 reports complete cross-cell-type generalization results with standard deviations, averaged across four target cell types and 1,635 test perturbations. The evaluation covers HepG2, RPE1, Jurkat, and K562, with variability computed across perturbations.

  • Evaluation scope: Table 2 evaluates cross-cell-type generalization across 1,635 test perturbations from four cell types.The perturbations comprise 375 HepG2, 416 RPE1, 443 Jurkat, and 401 K562 examples.
  • Reporting format: Results are averaged over the four target cell types and reported as mean ± standard deviation across test perturbations.The standard deviations provide variability information for the cross-cell-type evaluation.
  • Reporting format: The table marks the best performance in bold and uses arrows to indicate whether higher or lower values are preferred.Upward arrows denote higher-is-better metrics, while downward arrows denote lower-is-better metrics.

E.3 PER-CELL-TYPE RESULTS

PT-RAG performs superiorly or competitively across most cell types when metrics are computed independently for each cell–perturbation pair. Its strongest per-cell-type results are reported for HepG2, RPE1, and Jurkat.

  • Per-cell-type evaluation: PT-RAG demonstrates superior or competitive performance across most cell types under per-cell-type evaluation.Metrics are computed independently for each cell–perturbation pair within each cell type.
  • Per-cell-type evaluation: PT-RAG achieves particularly strong results in HepG2, RPE1, and Jurkat cell lines.
  • HepG2: The HepG2 comparison averages results over 375 cell–perturbation samples.
  • RPE1: The RPE1 comparison averages results over 416 cell–perturbation samples.

E.4 SENSITIVITY STUDIES · E.5 COMPUTATIONAL COSTS

Sensitivity analyses show that sparsity regularization prevents retrieving nearly all perturbations, stabilizes PT-RAG performance across non-zero settings and retrieval sizes, and enables meaningful context selection. PT-RAG incurs higher computation than baselines but remains tractable while delivering substantial performance gains.

  • E.4 SENSITIVITY STUDIES: With K = 32 on HepG2, λsparse = 0 retrieves 31.949 perturbations on average and substantially degrades performance across all metrics.The sparsity loss penalizes selected contexts to reduce noisy retrieval and computational burden.
  • E.4 SENSITIVITY STUDIES: At λsparse = 0, Pearson DEG correlation is 0.134, Spearman DEG correlation is -0.025, RMSE is 0.567, MSE is 0.322, and MAE is 0.362.Distributional metrics also degrade, with W2 distance reaching 651.744 and Energy distance reaching 11.974.
  • E.4 SENSITIVITY STUDIES: For λsparse ∈{0.01, 0.10, 1.00}, retrieved perturbations decrease to 12.536, 6.612, and 4.915, while performance remains nearly identical across settings.Pearson DEG correlations range from 0.594-0.604, Spearman DEG correlations from 0.386-0.401, RMSE from 0.469-0.470, MSE from 0.220-0.221, and MAE is 0.305.
  • E.4 SENSITIVITY STUDIES: PT-RAG is stable across K = 16 and K = 32: Pearson DEG ranges are 0.578-0.604 and 0.594-0.604, while Spearman DEG ranges are 0.369-0.401 and 0.386-0.401, respectively.This robustness permits choosing retrieval size according to computational constraints without significantly sacrificing performance.
  • E.4 SENSITIVITY STUDIES: Distributional metrics vary slightly with retrieval size: K = 32 performs marginally better for W2 at lower sparsity, whereas K = 16 performs slightly better for Energy distance.At λsparse = 1.00, W2 performance converges between configurations.
  • E.4 SENSITIVITY STUDIES: Vanilla RAG’s Pearson DEG correlation improves from 0.293 at K = 2 to 0.351 at K = 32, but remains 42% below PT-RAG’s 0.604.Spearman DEG is non-monotonic, reaching 0.289 at K = 32, while RMSE and MSE reach 0.577 and 0.334 versus PT-RAG’s 0.470 and 0.221.
  • E.4 SENSITIVITY STUDIES: Even at K = 32, Vanilla RAG underperforms because its cell-type-agnostic, non-differentiable retrieval cannot learn which perturbations provide informative cellular context.PT-RAG’s differentiable selection filters misleading context and consistently outperforms all Vanilla RAG configurations across metrics.
  • E.5 COMPUTATIONAL COSTS: PT-RAG has approximately 1.7× more FLOPs per batch than baselines, 2.86B versus 1.67B, while all models have comparable parameter counts of ∼20–21M.The overhead comes from triplet construction, scoring, and Gumbel-Softmax sampling for K = 32 candidates; the absolute cost is 60M FLOPs per cell and remains tractable.

F ADDITIONAL DETAILS ON STATISTICAL TESTS

The statistical analysis reports Benjamini-Hochberg FDR-corrected Mann-Whitney U tests for PT-RAG comparisons with STATE baselines. PT-RAG shows broad significant improvements over STATE, but only one significant improvement over STATE+GenePT.

  • Test methodology: Table 8 reports Benjamini-Hochberg FDR-corrected p-values from Mann-Whitney U tests comparing PT-RAG against STATE baselines.The analysis is provided for transparency and reproducibility.
  • PT-RAG vs. STATE: pFDR < 0.001 for five metrics—Pearson DEG, Spearman DEG, MAE, W1, and W2—favoring PT-RAG over STATE.The reported values are 2.44 × 10−8, 4.89 × 10−10, 9.33 × 10−5, 3.41 × 10−6, and 5.47 × 10−8, respectively.
  • PT-RAG vs. STATE: pFDR = 5.47 × 10−8 for W2, the most significant PT-RAG improvement over STATE.The passage associates this result with enhanced modeling of cell population distributions in low-dimensional space.
  • PT-RAG vs. STATE+GenePT: pFDR = 0.041 for W2 is the only significant PT-RAG improvement over STATE+GenePT, while W1 is marginal at pFDR = 0.082.The comparison yields fewer significant improvements than the comparison against STATE.
Loading 2603.07233v1…