Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation
Andrea Giuseppe Di Francesco, Andrea Rubbi, Pietro Liò
2026-03-10
Summary
This paper introduces a new method, PT-RAG, for predicting how cells will change when scientists alter their genes. It's about using artificial intelligence to understand what happens inside cells when we tweak their genetic code, which is important for understanding diseases and developing new treatments.
What's the problem?
Currently, AI models struggle to accurately predict cellular responses to gene changes across different types of cells and different kinds of genetic alterations. Existing methods don't effectively consider the specific context of the cell and the change being made, leading to inaccurate predictions. Simply grabbing information from a database and feeding it to the AI doesn't work well because it's hard to determine what information is actually relevant to the specific cell and genetic change.
What's the solution?
The researchers developed PT-RAG, which stands for Perturbation-aware Two-stage Retrieval-Augmented Generation. It works in two steps: first, it finds potentially relevant genetic changes using a special way of representing genes called GenePT embeddings. Then, it refines this selection, focusing on the changes that are most relevant to the specific cell type and the initial genetic change. This process is 'differentiable,' meaning the AI can learn and improve how it selects the most useful information. Essentially, it's a smarter way of finding and using relevant data to make predictions.
Why it matters?
This work shows that using a carefully designed information retrieval system, specifically one that understands the context of cells and genetic changes, is crucial for accurately predicting cellular responses. It demonstrates that simply using a standard information retrieval method can actually make predictions *worse*. This establishes a new, promising approach – retrieval-augmented generation – for modeling how cells react to genetic alterations, which could significantly advance our understanding of biology and medicine.
Abstract
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.