Generación Aumentada por Recuperación para la Predicción de Respuestas Celulares a Perturbaciones Génicas

Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation

March 7, 2026

Autores: Andrea Giuseppe Di Francesco, Andrea Rubbi, Pietro Liò

cs.AI

Resumen

Predecir cómo responden las células a las perturbaciones genéticas es fundamental para comprender la función de los genes, los mecanismos de las enfermedades y el desarrollo de terapias. Si bien los enfoques recientes de aprendizaje profundo han mostrado potencial para modelar las respuestas de las células individuales a perturbaciones, tienen dificultades para generalizar entre tipos celulares y contextos de perturbación debido a la información contextual limitada durante la generación. Presentamos PT-RAG (Generación Aumentada por Recuperación en Dos Etapas con Conciencia de Perturbación), un marco novedoso que extiende la Generación Aumentada por Recuperación más allá de las aplicaciones tradicionales de modelos de lenguaje a la biología celular. A diferencia de los sistemas RAG estándar diseñados para la recuperación de texto con LLM preentrenados, la recuperación de perturbaciones carece de métricas de similitud establecidas y requiere aprender qué constituye un contexto relevante, haciendo que la recuperación diferenciable sea esencial. PT-RAG aborda esto mediante un proceso en dos etapas: primero, recupera las perturbaciones candidatas K utilizando incrustaciones de GenePT, y luego refina la selección de forma adaptativa mediante un muestreo discreto Gumbel-Softmax condicionado tanto al estado celular como a la perturbación de entrada. Esta recuperación diferenciable consciente del tipo celular permite la optimización de extremo a extremo del objetivo de recuperación de manera conjunta con la generación. En el conjunto de datos de perturbación de un solo gen Replogle-Nadig, demostramos que PT-RAG supera tanto a STATE como a RAG básico bajo condiciones experimentales idénticas, con las mejoras más sólidas en las métricas de similitud distribucional (W_1, W_2). Cabe destacar que el fracaso drástico del RAG básico es en sí mismo un hallazgo clave: demuestra que la recuperación diferenciable y consciente del tipo celular es esencial en este dominio, y que una recuperación ingenua puede dañar activamente el rendimiento. Nuestros resultados establecen la generación aumentada por recuperación como un paradigma prometedor para modelar las respuestas celulares a la perturbación génica. El código para reproducir nuestros experimentos está disponible en https://github.com/difra100/PT-RAG_ICLR.

English

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.