用于预测细胞对基因扰动响应的检索增强生成

预测细胞如何响应遗传扰动是理解基因功能、疾病机制和治疗方法的基础。虽然最近的深度学习方法在模拟单细胞扰动响应方面显示出前景，但由于生成过程中上下文信息有限，它们难以跨细胞类型和扰动上下文进行泛化。我们引入了PT-RAG（扰动感知两阶段检索增强生成），一种新颖的框架，它将检索增强生成从传统的语言模型应用扩展到细胞生物学。与为预训练LLM设计的文本检索标准RAG系统不同，扰动检索缺乏既定的相似性度量，需要学习什么构成相关上下文，这使得可微检索变得至关重要。PT-RAG通过一个两阶段流程解决了这个问题：首先，使用GenePT嵌入检索候选扰动$K$，然后通过基于细胞状态和输入扰动条件的Gumbel-Softmax离散采样自适应地优化选择。这种细胞类型感知的可微检索使得检索目标与生成能够进行端到端的联合优化。在Replogle-Nadig单基因扰动数据集上，我们证明了PT-RAG在相同实验条件下优于STATE和vanilla RAG，在分布相似性度量（$W_1$、$W_2$）上增益最大。值得注意的是，vanilla RAG的显著失败本身就是一个关键发现：它证明了可微的、细胞类型感知的检索在该领域是必不可少的，而朴素的检索实际上会损害性能。我们的结果确立了检索增强生成作为模拟细胞对基因扰动响应的一种有前景的范式。重现我们实验的代码可在https://github.com/difra100/PT-RAG_ICLR获取。