arXiv:cs.LG· Andrea Giuseppe Di Francesco, Andrea Rubbi, Rishabh Jain, Pietro Li\`o·· 3 小时前
PT-RAG:面向基因扰动细胞响应预测的检索增强生成框架
Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation
AI 导读
PT-RAG 是一个用于生成式细胞扰动响应预测的即插即用两阶段检索增强模块,通过 GenePT 语义检索筛选 K 个候选扰动,再用可微分 Gumbel-Softmax 选择器依据对照细胞状态自适应选择上下文。
正文
Abstract:Predicting transcriptional responses to genetic perturbations is fundamental to functional genomics and therapeutic discovery. Recent deep learning models have shown promise in single-cell perturbation response prediction, but they typically generate each response in isolation, without explicitly leveraging related perturbations. We introduce PT-RAG (Perturbation-aware Two-stage Retrieval-Augmented Generation), a plug-in retrieval-and-conditioning module for generative cellular perturbation response. PT-RAG augments an existing perturbation-response backbone with learned access to related perturbation contexts. The key challenge is that relevance is not fixed in this setting: functionally related genes may elicit different effects across cell types. PT-RAG addresses this with a two-stage retrieval mechanism: GenePT-based semantic retrieval first identifies K candidate perturbations, after which a differentiable Gumbel-Softmax selector adaptively selects retrieved contexts conditioned on the control cell state, the query perturbation, and each candidate perturbation. We evaluate PT-RAG on two backbones, a STATE-style generator used as a frozen random reservoir and a fully trained scGPT, across cross-cell-type and cross-perturbation generalization tasks. PT-RAG consistently improves distributional similarity and often overall predictive quality; for example, on scGPT cross-cell-type results, the 2-Wasserstein distance drops by 5.9% relative to scGPT alone. The code to reproduce our experiments is available at this https URL.
| Comments: | Accepted to NeurIPS 2026 main track. 34 pages, 11 figures, 18 tables |
| Subjects: | Machine Learning (cs.LG); Information Retrieval (cs.IR) |
| Cite as: | arXiv:2603.07233 [cs.LG] |
| (or arXiv:2603.07233v2 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2603.07233 arXiv-issued DOI via DataCite |
Submission history
From: Andrea Giuseppe Di Francesco [view email]
[v1]
Sat, 7 Mar 2026 14:31:27 UTC (7,267 KB)
[v2]
Thu, 8 Oct 2026 10:43:12 UTC (500 KB)
来源:arXiv:cs.LG · arxiv.org