Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation

TL;DR

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, is introduced and establishes retrieval-augmented generation as a promising paradigm for modelling cellular responses to gene perturbation.

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.

Figures (9)

• Figure 1: Comparison of Generation, Vanilla RAG, and PT-RAG architectures.Left: Generation baseline combines cell and perturbation encodings; Vanilla RAG adds non-differentiable retrieval (dotted lines). Right: PT-RAG uses two-stage retrieval: (1) semantic similarity for $K$ candidates, (2) differentiable Gumbel-Softmax selection conditioned on $h^{ctrl}$, $h_{pert}$, $h_k^{cxt}$.
• Figure 2: Jaccard similarity matrix across cell types.
• Figure 3: The pdf $f(x) = e^{-x - e^{-x}}$ of Gumbel(0, 1).
• Figure 4: Sensitivity analysis of sparsity regularization parameter $\lambda_{\text{sparse}}$ - Core metrics. Performance metrics across different sparsity levels on HepG2 cell type with $K=32$, focusing on gene-level correlations and reconstruction accuracy.
• Figure 5: Sensitivity analysis of sparsity regularization parameter $\lambda_{\text{sparse}}$ - Distributional metrics and retrieval behavior. Additional performance metrics and retrieval patterns across different sparsity levels.