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

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 experimentally characterised responses to 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. Through 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, Wasserstein distance drops by 10.7%. The code to reproduce our experiments is available at https://anonymous.4open.science/r/PT-RAG_NIPS-13D8/.