Geração Aumentada por Recuperação para Prever Respostas Celulares a Perturbações Genéticas

Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation

March 7, 2026

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

cs.AI

Resumo

Prever como as células respondem a perturbações genéticas é fundamental para compreender a função dos genes, os mecanismos das doenças e o desenvolvimento de terapias. Embora abordagens recentes de aprendizagem profunda tenham mostrado potencial na modelização de respostas a perturbações em células únicas, elas lutam para generalizar entre tipos celulares e contextos de perturbação devido a informações contextuais limitadas durante a geração. Apresentamos o PT-RAG (Perturbation-aware Two-stage Retrieval-Augmented Generation), uma estrutura inovadora que estende a Geração Aumentada por Recuperação para além das aplicações tradicionais de modelos de linguagem, aplicando-a à biologia celular. Ao contrário dos sistemas RAG padrão concebidos para recuperação de texto com LLMs pré-treinados, a recuperação de perturbações carece de métricas de similaridade estabelecidas e requer a aprendizagem do que constitui um contexto relevante, tornando a recuperação diferenciável essencial. O PT-RAG aborda esta questão através de um *pipeline* em duas fases: primeiro, recupera perturbações candidatas K usando incorporações do GenePT, depois refina adaptativamente a seleção através de uma amostragem discreta de Gumbel-Softmax condicionada tanto pelo estado da célula como pela perturbação de entrada. Esta recuperação diferenciável e consciente do tipo celular permite a otimização de ponta a ponta do objetivo de recuperação em conjunto com a geração. No conjunto de dados de perturbação de gene único Replogle-Nadig, demonstramos que o PT-RAG supera tanto o STATE como o RAG padrão em condições experimentais idênticas, com os ganhos mais significativos nas métricas de similaridade distribucional (W_1, W_2). É notável que o fracasso dramático do RAG padrão é, por si só, uma descoberta fundamental: demonstra que uma recuperação diferenciável e consciente do tipo celular é essencial neste domínio, e que uma recuperação ingénua pode ativamente prejudicar o desempenho. Os nossos resultados estabelecem a geração aumentada por recuperação como um paradigma promissor para modelar respostas celulares a perturbações genéticas. O código para reproduzir as nossas experiências está disponível em 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.