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

arXiv: https://arxiv.org/abs/2603.07233 · HTML: https://mlsanity.com/paper/2603.07233/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Medical and biomedical RAG, one of 15 directions in Retrieval-augmented generation.
- Builds on 8 of its 35 references that are in the map.
- First posted on arXiv 7 Mar 2026.

## Placed in

- Retrieval-augmented generation / [Medical and biomedical RAG](https://mlsanity.com/rag/medical-and-biomedical-rag/)

## Builds on (most cited first, 8 of 8)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [From Local to Global: A Graph RAG Approach to Query-Focused Summarization](https://mlsanity.com/paper/2404.16130/) (2024)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [MPNet: Masked and Permuted Pre-training for Language Understanding](https://mlsanity.com/paper/2004.09297/) (2020)

## Abstract (arXiv)

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.