Retrieval-Augmentierte Generierung zur Vorhersage zellulärer Reaktionen auf Genperturbationen
Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation
March 7, 2026
Autoren: Andrea Giuseppe Di Francesco, Andrea Rubbi, Pietro Liò
cs.AI
Zusammenfassung
Die Vorhersage, wie Zellen auf genetische Perturbationen reagieren, ist grundlegend für das Verständnis von Genfunktion, Krankheitsmechanismen und der Entwicklung von Therapien. Während neuere Deep-Learning-Ansätze vielversprechende Ergebnisse bei der Modellierung von Einzelzell-Perturbationsantworten gezeigt haben, leiden sie unter mangelnder Generalisierbarkeit über verschiedene Zelltypen und Perturbationskontexte hinweg, bedingt durch begrenzte kontextuelle Informationen während der Generierung. Wir stellen PT-RAG (Perturbation-aware Two-stage Retrieval-Augmented Generation) vor, ein neuartiges Framework, das Retrieval-Augmented Generation über traditionelle Sprachmodell-Anwendungen hinaus auf die Zellbiologie erweitert. Im Gegensatz zu Standard-RAG-Systemen, die für Textretrieval mit vortrainierten LLMs konzipiert sind, fehlen für das Perturbations-Retrieval etablierte Ähnlichkeitsmetriken, und es muss erlernt werden, was einen relevanten Kontext ausmacht, was differentierbares Retrieval essenziell macht. PT-RAG adressiert dies durch eine zweistufige Pipeline: Zuerst werden Kandidaten-Perturbationen K unter Verwendung von GenePT-Embeddings abgerufen, dann wird die Auswahl adaptiv durch Gumbel-Softmax-Diskretabtastung verfeinert, die sowohl vom Zellzustand als auch von der Eingabe-Perturbation abhängt. Dieser zelltypsensitive, differentierbare Retrieval-Ansatz ermöglicht eine end-to-end-Optimierung des Retrieval-Ziels gemeinsam mit der Generierung. Auf dem Replogle-Nadig-Datensatz für Einzelgen-Perturbationen demonstrieren wir, dass PT-RAG unter identischen experimentellen Bedingungen sowohl STATE als auch einfaches RAG übertrifft, mit den stärksten Verbesserungen bei den Verteilungsähnlichkeitsmetriken (W_1, W_2). Bemerkenswerterweise ist das dramatische Scheitern des einfachen RAG an sich eine zentrale Erkenntnis: Es zeigt, dass differentierbares, zelltypsensitives Retrieval in diesem Bereich unerlässlich ist und dass naives Retrieval die Leistung aktiv verschlechtern kann. Unsere Ergebnisse etablieren Retrieval-Augmented Generation als ein vielversprechendes Paradigma für die Modellierung zellulärer Antworten auf Genperturbationen. Der Code zur Reproduktion unserer Experimente ist verfügbar unter 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.