Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation
This work is significant for researchers in cellular biology and drug discovery, as it provides a more accurate method for predicting how cells respond to genetic changes, which is crucial for understanding disease and developing therapies. It represents an incremental but important step in applying RAG to a new domain.
This paper introduces PT-RAG, a novel framework that extends Retrieval-Augmented Generation to predict cellular responses to genetic perturbations. It addresses the challenge of generalizing across cell types and perturbation contexts by using a two-stage differentiable retrieval process. PT-RAG significantly outperforms existing methods like STATE and vanilla RAG on the Replogle-Nadig dataset, showing strong gains in distributional similarity metrics ($W_1$, $W_2$).
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.