Differentiable Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation

The gist: PT-RAG introduces Perturbation-aware Two-stage RetrievalAugmented Generation, a plug-in retrieval-and-conditioning module for generative cellular perturbation response, which augments an existing perturbationresponse backbone with learned access to related perturbation contexts; this framework consistently improves distributional similarity and often overall predictive quality in predicting transcriptional responses to genetic perturbations <ref:2603.07233#pg11>.

How it works

PT-RAG addresses the limitation where existing deep learning models typically generate each response in isolation, without explicitly leveraging related perturbations. The key challenge addressed is that relevance is not fixed in this setting, as functionally related genes may elicit different effects across cell types. PT-RAG tackles this with a twostage 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.

The two stages of PT-RAG are as follows:

1. Stage 1: Semantic Candidate Retrieval, which is identical to Naïve RAG and non-differentiably selects K candidates via cosine similarity in GenePT embedding space. The query perturbation is excluded from the pool to prevent information leakage.

2. Stage 2: Differentiable Cell-Type-Aware Selection, where each retrieved candidate p(k) is encoded as h cxt k, and a triplet capturing the three-way relationship between cell state, target perturbation, and candidate is formed. This triplet ck is scored with a two-layer MLP to output logits sk, which are then used in a Straight-Through GumbelSoftmax to make K independent binary decisions wk. Selected contexts are then aggregated via z = sum k wk · MLPproj(ck) and decoded by the TransformerGenerator.

Key Components and Backbones

PT-RAG is designed to be a plug-and-play retrieval module rather than a backbone-specific architecture, as demonstrated by its cross-cell-type transfer experiments on the scGPT perturbation generator. The framework can be applied to two backbones:

** A STATE-style generator used as a frozen random reservoir in the STATE experiments. In this case, the context enters at the cell level as the generator input. The training objective for STATE-based models is L = Ldist + λsparse Lsparse, where Ldist = Energy(ˆx pert, xpert) and λsparse penalizes dense masks to prevent mode collapse. In the scGPT-based models, a meansquared-error reconstruction loss is used in place of Ldist. The scGPT+PT-RAG model adds the context vector to the transformer output at every gene position, on top of h pert, before the expression decoder. The parameter count for PT-RAG is 5.36M trainable parameters. The computational cost comparison shows that PT-RAG requires substantially more FLOPs because it has to score and filter retrieved contexts over the full pool of roughly 2,009 perturbations. The scGPT transfer experiment shows that PT-RAG recovers performance where GenePT alone can degrade scGPT, and it outperforms both baselines on most metrics. 5.4 Cell-Type-Specific Retrieval Patterns**

To validate the cell-state awareness, the paper analyzes the overlap in selected perturbations across cell types using Jaccard similarity. For each of the 33 genes common to all four cell types in the cross-cell-type test set, only ∼19% of selected perturbations overlap between cell types for the same query gene. This low overlap confirms that PT-RAG actively conditions on cell state rather than relying solely on the query perturbation—a capability that Naïve RAG fundamentally lacks. The examples show functional coherence, where retrieved genes are biologically plausible and cell-type-specific differences are maintained across contexts.

Experimental Evaluation

The model is evaluated on the Replogle-Nadig Perturb-seq dataset, using two complementary generalization settings:

1. Cross-Cell-Type Generalization, where each cell type serves as a held-out target in turn. This tests generalization across biological contexts.

2. Cross-Perturbation Generalization, which is the more direct test of retrieval's value when a query perturbation has no direct training signal. This setting tests generalization across the perturbation space itself.

The results consistently show that Naïve RAG actively hurts performance, substantially underperforming STATE across all metrics in the cross-cell-type setting. PT-RAG demonstrates statistically significant improvements over STATE in gene-level correlations and distributional similarity (W1, W2). In the cross-perturbation setting, PT-RAG consistently improves most metrics, with significant gains on reconstruction and distributional geometry. The strongest pattern is again distributional: PT-RAG yields lower Wasserstein and Energy distances, indicating better recovery of population-level response structure for unseen genes.

Conclusion

Taken together, the results conclude that PT-RAG is effective when retrieval is learned, cell-aware, and transferred as a module rather than baked into one backbone. Learned retrieval beats Naïve RAG and improves on the retrieval-free baseline, with the strongest gains appearing where population geometry matters. The scGPT experiment reinforces the plug-and-play framing, because the same retrieval rule can be attached to another generator without changing the retrieval logic. The cross-perturbation results then show how the same method behaves under a harder split in which direct perturbation supervision is much thinner. Taken together, we conclude that PT-RAG is effective when retrieval is learned, cell-aware, and transferred as a module rather than baked into one backbone. The limitations noted include the extra computation introduced by retrieval and the fact that retrieval examples are biologically suggestive rather than causal. The current study is also limited to single-gene perturbations, which keeps the biological setting focused but leaves combinatorial perturbations and other modalities that would require a completely new analysis on context modelling for future work. The retrieval pool at test time contains the embeddings of all perturbations, including those never seen in training, which is regarded as part of the protocol rather than an advantage of the method. The current study is also limited to single-gene perturbations, which keeps the biological setting focused but leaves combinatorial perturbations and other modalities that would require a completely new analysis on context modelling for future work. The retrieval examples are biologically suggestive rather than causal, so they should be read as model behavior rather than mechanistic proof.

**How it works (Cont.

1. Bold header: Differentiable Retrieval-Augmented Generation (PT-RAG) for Perturbation Prediction. This module augments existing perturbation response backbones with learned access to related perturbation contexts, allowing the system to leverage knowledge from related perturbations already observed in training, which is crucial when predictions are generated solely from the control cell state and perturbation identity, without leveraging knowledge about related perturbations.

2. Bold header: Cell-Type-Aware Selection via Gumbel-Softmax. The system uses a twostage retrieval mechanism where a differentiable Gumbel-Softmax selector adaptively selects retrieved contexts conditioned on the control cell state, the query perturbation, and each candidate perturbation, enabling the system to learn cell-type-specific retrieval patterns as evidenced by a low Jaccard similarity of only ~19% perturbation overlap across cell types for the same query gene.

3. Bold header: Robustness to Unseen Perturbations via Cross-Perturbation Generalization. By employing a cross-perturbation protocol where 60% of perturbations are held out for testing, the improved system can predict responses to genes never experimentally perturbed, demonstrating that PT-RAG consistently improves most metrics, with significant gains on reconstruction and distributional geometry in this challenging cold-start regime.

4. Bold header: Backbone Agnosticism and Transferability. PT-RAG is designed as a plug-and-play retrieval module rather than a backbone-specific architecture, allowing the same retrieval rule to be attached to another generator without changing the retrieval logic, specifically demonstrated by its successful transfer to the scGPT perturbation generator.

5. Bold header: Improved Distributional Fidelity in Unseen Contexts. The system is optimized for predicting the distribution of perturbed cells by reducing metrics like Wasserstein distances (W2) and Energy distance, showing that retrieved GenePT embeddings capture substantial functional relationships leading to a significant drop in W2 on scGPT results.

6. Bold header: Learned Context Filtering over Naive Retrieval. The system proves that na¨ıve retrieval actively harms performance, contrasting it with its ability to improve generation quality, showing that retrieved contexts have to be filtered by a differentiable, cell-state-aware selector rather than relying on non-differentiable methods.

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.