Génération Augmentée par Récupération pour la Prédiction des Réponses Cellulaires aux Perturbations Génétiques

Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation

March 7, 2026

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

cs.AI

Résumé

Prédire la réponse des cellules aux perturbations génétiques est fondamental pour comprendre la fonction des gènes, les mécanismes des maladies et le développement thérapeutique. Bien que les approches récentes d'apprentissage profond aient montré des résultats prometteurs pour modéliser les réponses aux perturbations à l'échelle monocellulaire, elles peinent à généraliser leurs performances à travers différents types cellulaires et contextes de perturbation en raison d'informations contextuelles limitées lors de la génération. Nous présentons PT-RAG (Perturbation-aware Two-stage Retrieval-Augmented Generation), un nouveau cadre qui étend la Génération Augmentée par Récupération au-delà des applications traditionnelles des modèles de langage pour l'appliquer à la biologie cellulaire. Contrairement aux systèmes RAG standard conçus pour la récupération de texte avec des LLM pré-entraînés, la récupération de perturbations manque de métriques de similarité établies et nécessite d'apprendre ce qui constitue un contexte pertinent, rendant une récupération différentiable essentielle. PT-RAG résout ce problème via un pipeline en deux étapes : premièrement, la récupération de perturbations candidates K à l'aide d'embeddings GenePT, puis l'affinage adaptatif de la sélection via un échantillonnage discret Gumbel-Softmax conditionné à la fois par l'état cellulaire et la perturbation d'entrée. Cette récupération différentiable, sensible au type cellulaire, permet une optimisation de bout en bout de l'objectif de récupération conjointement avec la génération. Sur le jeu de données de perturbation génique unique Replogle-Nadig, nous démontrons que PT-RAG surpasse à la fois STATE et le RAG standard dans des conditions expérimentales identiques, avec les gains les plus importants sur les métriques de similarité distributionnelle (W_1, W_2). Il est notable que l'échec spectaculaire du RAG standard est en soi un résultat clé : il démontre qu'une récupération différentiable et sensible au type cellulaire est essentielle dans ce domaine, et qu'une récupération naïve peut activement dégrader les performances. Nos résultats établissent la génération augmentée par récupération comme un paradigme prometteur pour la modélisation des réponses cellulaires aux perturbations géniques. Le code pour reproduire nos expériences est disponible à l'adresse 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.