Generazione Aumentata dal Recupero per la Previsione delle Risposte Cellulari alle Perturbazioni Geniche

Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation

March 7, 2026

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

cs.AI

Abstract

Prevedere come le cellule rispondono alle perturbazioni genetiche è fondamentale per comprendere la funzione genica, i meccanismi patologici e lo sviluppo di terapie. Sebbene i recenti approcci di deep learning abbiano mostrato potenzialità nella modellizzazione delle risposte cellulari a perturbazioni a livello di singola cellula, essi faticano a generalizzare attraverso diversi tipi cellulari e contesti di perturbazione a causa di informazioni contestuali limitate durante la generazione. Introduciamo PT-RAG (Perturbation-aware Two-stage Retrieval-Augmented Generation), un framework innovativo che estende la Generazione Aumentata tramite Recupero (RAG) oltre le tradizionali applicazioni dei modelli linguistici alla biologia cellulare. A differenza dei sistemi RAG standard progettati per il recupero di testo con LLM pre-addestrati, il recupero di perturbazioni manca di metriche di similarità consolidate e richiede di apprendere cosa costituisce un contesto rilevante, rendendo essenziale un recupero differenziabile. PT-RAG affronta questa sfida attraverso una pipeline a due stadi: prima, recupera le perturbazioni candidate K utilizzando gli embedding di GenePT, poi affina selettivamente la scelta attraverso un campionamento discreto Gumbel-Softmax condizionato sia dallo stato cellulare che dalla perturbazione in input. Questo recupero differenziabile e consapevole del tipo cellulare permette l'ottimizzazione end-to-end dell'obiettivo di recupero congiuntamente alla generazione. Sul dataset di perturbazione a singolo gene Replogle-Nadig, dimostriamo che PT-RAG supera sia STATE che il RAG standard in condizioni sperimentali identiche, con i miglioramenti più significativi nelle metriche di similarità distribuzionale (W_1, W_2). È degno di nota che il fallimento eclatante del RAG standard sia di per sé un risultato chiave: dimostra che in questo ambito è essenziale un recupero differenziabile e consapevole del tipo cellulare, e che un recupero ingenuo può attivamente danneggiare le prestazioni. I nostri risultati stabiliscono la generazione aumentata tramite recupero come un paradigma promettente per modellare le risposte cellulari alla perturbazione genica. Il codice per riprodurre i nostri esperimenti è disponibile all'indirizzo 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.