Retrieval-augmented Generation voor het Voorspellen van Cellulaire Reacties op Genperturbatie

Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation

March 7, 2026

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

cs.AI

Samenvatting

Het voorspellen van hoe cellen reageren op genetische perturbaties is fundamenteel voor het begrip van genfunctie, ziekte-mechanismen en de ontwikkeling van therapieën. Hoewel recente deep learning-benaderingen veelbelovend zijn gebleken in het modelleren van responsies op perturbaties in enkele cellen, hebben zij moeite met generaliseren over verschillende celtypen en perturbatiecontexten vanwege beperkte contextuele informatie tijdens de generatie. Wij introduceren PT-RAG (Perturbation-aware Two-stage Retrieval-Augmented Generation), een nieuw raamwerk dat Retrieval-Augmented Generation uitbreidt voorbij traditionele taalmodellentoepassingen naar de cellulaire biologie. In tegenstelling tot standaard RAG-systemen die zijn ontworpen voor tekstretrieval met vooraf getrainde LLM's, ontbreekt het bij perturbatieretrieval aan gevestigde similariteitsmetrieken en vereist het leren wat relevante context vormt, waardoor differentieerbare retrieval essentieel is. PT-RAG lost dit op via een tweestappenpijplijn: eerst worden kandidaat-perturbaties K opgehaald met behulp van GenePT-embeddingen, vervolgens wordt de selectie adaptief verfijnd door middel van Gumbel-Softmax discrete sampling, geconditioneerd op zowel de celstatus als de inputperturbatie. Deze differentieerbare retrieval, die rekening houdt met het celtype, maakt end-to-end optimalisatie van het retrieval-doel mogelijk, gezamenlijk met de generatie. Op de Replogle-Nadig single-gene perturbatie-dataset tonen wij aan dat PT-RAG beter presteert dan zowel STATE als standaard RAG onder identieke experimentele omstandigheden, met de sterkste verbeteringen in distributionele similariteitsmetrieken (W_1, W_2). Opmerkelijk is dat het dramatische falen van standaard RAG zelf een belangrijke bevinding is: het toont aan dat differentieerbare retrieval, die rekening houdt met het celtype, essentieel is in dit domein, en dat naïeve retrieval de prestaties actief kan schaden. Onze resultaten vestigen retrieval-augmented generation als een veelbelovend paradigma voor het modelleren van cellulaire responsies op genperturbatie. De code om onze experimenten te reproduceren is beschikbaar op 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.