Differentiable Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation¶

会议: NeurIPS 2026
arXiv: 2603.07233
OpenReview: https://openreview.net/forum?id=UbCFFQ9TmZ
代码: https://github.com/difra100/PT-RAG_NIPS
领域: 计算生物
关键词: 基因扰动响应预测, 单细胞转录组, 检索增强生成, 可微子集选择, Gumbel-Softmax

一句话总结¶

提出了首个面向单细胞基因扰动预测的可微端到端检索增强生成框架 PT-RAG，通过两阶段可微筛选将细胞状态自适应扰动上下文注入生成模型，显著提升分布保真度与几何结构泛化能力。

研究背景与动机¶

利用深度学习模型从单细胞转录组（Perturb-seq）中预测细胞对基因敲除或过表达的转录响应，是功能基因组学与靶向药物发现的关键基石。然而，潜在基因扰动与细胞类型的组合空间呈现指数级爆炸，实验全覆盖极度受限。现存主流模型（如潜空间模型 scGen/CPA、最优传输方法 CellOT、序列生成模型 STATE、条件流匹配 CellFlow 等）普遍采用“隔离预测”范式：它们仅依赖对照组细胞状态与目标扰动标识独立生成响应，模型无法在预测时动态借用训练集中已知相关扰动的上下文信息。

这一局限导致两类高频现实场景性能瓶颈突出：其一是跨细胞类型迁移预测（部分扰动在目标细胞系中未被观测，但在其他细胞系中已测得），其二是全新未测基因的冷启动泛化（完全未被实验扰动的基因）。自然语言处理中的检索增强生成（RAG）通过引入非参数化外部知识取得了巨大成功，但直接将 RAG 迁移至细胞扰动建模面临两大本质障碍：第一，生物领域缺乏固定的语义相关性度量基准，仅凭基因文本描述的嵌入（如 GenePT）仅能捕捉功能先验，无法反映细胞内在的调控网络；第二，朴素检索对细胞背景不敏感，相同的基因扰动在 T 细胞、肝细胞或白血病细胞中往往激发截然不同的转录程序。朴素检索若无视细胞上下文直接注入相似基因，不仅无法提供帮助，反而会引入严重的生物噪声。

针对这一矛盾，本文提出关键洞察：在细胞扰动生物学中，外部知识的检索不仅必须具备语义先验作为初筛，更必须通过端到端梯度的引导，让模型根据目标细胞的即时表达状态、查询扰动以及候选扰动特征进行联合打分与可微过滤。核心 idea：构建扰动感知的两阶段可微检索增强生成框架（PT-RAG），先利用 GenePT 语义嵌入召回候选扰动子集，再通过 Straight-Through Gumbel-Softmax 构建由细胞状态联合条件化的可微二值门控机制，端到端筛选出对特定细胞背景真正有益的上下文信息。

方法详解¶

整体框架¶

PT-RAG 作为一个通用的即插即用检索增强模块，能够无缝嫁接于不同的细胞生成骨干网络（如 STATE 架构的固定随机水库网络，或预训练的 scGPT 转换器）。系统接收对照组细胞表达量 $x_{\text{ctrl}} \in \mathbb{R}^G$（覆盖 $G=2000$ 个高可变基因）以及查询扰动标识 $p_{\text{pert}} \in \mathcal{P}$，目标是生成扰动后的单细胞转录分布 $\hat{x}_{\text{pert}}$。预测流程显式划分为两阶段检索与生成条件化：

首先，在第一阶段语义初筛中，系统利用预计算的 GenePT 特征库进行余弦相似度粗筛，在排除查询扰动自身后提取 Top-$K$ 个功能相关的候选扰动集合；随后在第二阶段可微精选过程中，编码器分别提取对照细胞表征、查询扰动表征以及各候选扰动的隐层向量，将三者拼接为三元组，通过打分 MLP 与 Straight-Through Gumbel-Softmax 离散采样做出独立的二值入选决策；最后，经过投影的网络聚合选中的上下文向量，与主流程表征融合后驱动生成解码器输出扰动后表达。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入：对照组表达 x_ctrl 与查询扰动 p_pert"] --> S1["基于 GenePT 的语义初筛候选池<br/>排除查询项后依余弦相似度选出 Top-K 扰动"]
    S1 --> S2["三元组表征与细胞状态感知打分<br/>拼接 h_ctrl, h_pert, h_cxt 并过打分 MLP"]
    S2 --> S3["Gumbel-Softmax 可微独立子集选择<br/>Straight-Through 离散门控与稀疏正则化"]
    S3 --> S4["上下文特征聚合与骨干生成解码<br/>投影聚合上下文融入解码器生成扰动表达"]
    S4 --> Out["输出：扰动后细胞转录分布 x_pert"]

关键设计¶

1. 基于 GenePT 的语义初筛候选池：在无真实扰动响应时高效压缩搜索空间

传统的单细胞扰动模型使用 one-hot 编码，既丢失了基因间的生物学联系，又完全无法泛化至未见过的基因。为了赋予检索模块通用生物学语义先验，PT-RAG 借助 GenePT 基于 NCBI 文本摘要与 GPT-3.5 提取的稠密嵌入向量 $h_g^{\text{gpt}} \in \mathbb{R}^{1536}$，为包含 2,009 个候选基因的扰动库 $\mathcal{D}$ 建立单位归一化的特征索引。对于输入的查询扰动 $p_{\text{pert}}$，第一阶段通过计算余弦相似度直接非可微地检索最接近的 $K$ 个邻居候选： $$R_{p_{\text{pert}}} = \text{TopK}\left(h_{p_{\text{pert}}}^{\text{gpt}}, \mathcal{D} \setminus \{h_{p_{\text{pert}}}^{\text{gpt}}\}\right) = \{p^{(1)}, p^{(2)}, \dots, p^{(K)}\}$$ 强制排除查询扰动自身能从物理上阻断自环信息泄露。这一初筛机制避免了在大规模基因库（成千上万个基因）上直接进行全图可微计算的算力不可承受性，在保留功能相关基因子集的同时，将庞大的离散搜索空间压缩至小规模候选集 $K$（通常设定为 16 到 128）。

2. 三元组表征与细胞状态感知打分：解决跨细胞类型生物学效应异质性

同一个基因突变或敲除在不同细胞微环境中可能触发完全异质的信号通路反应。朴素 RAG（Naïve RAG）只凭借查询基因自身的相似度强行利用 Cross-Attention 注入候选，完全忽略了靶向细胞的生理背景，因而引入了大量无关甚至有害的干扰。PT-RAG 的关键设计在于构建包含三个维度的联合三元组向量：由可微线性层从对照组转录谱提取的细胞表征 $h_{\text{ctrl}} \in \mathbb{R}^{d_h}$、由单层扰动编码器投影的查询扰动隐向量 $h_{\text{pert}} \in \mathbb{R}^{d_h}$、以及由同一扰动编码器转换的各候选特征 $h_{k}^{\text{cxt}} = \text{PertEncoder}(h_{p^{(k)}}^{\text{gpt}}) \in \mathbb{R}^{d_h}$。

针对第 $k$ 个候选扰动，模型将其显式拼接并层归一化后送入两层多层感知机进行打分： $$c_k = [h_{\text{ctrl}} ;\, h_{\text{pert}} ;\, h_{k}^{\text{cxt}}] \in \mathbb{R}^{3d_h}$$ $$s_k = \text{MLP}_{\text{score}}(\text{LayerNorm}(c_k)) \in \mathbb{R}^2$$ 这里输出的二维 logits 分别对应于当前细胞背景下针对该候选的“排除（exclude）”与“纳入（include）”倾向，确保检索决策不仅依赖基因文本相关度，更深度耦合了受试细胞的具体转录状态。

3. Gumbel-Softmax 可微独立子集选择：支持多候选并行采样的端到端离散门控

为了让网络不仅能够自主学习“选谁不选谁”，还能支持多重扰动上下文的同时纳入（或在无有效上下文时全部剔除），PT-RAG 放弃了在 $K$ 个候选上做单个互斥 Softmax 分布的做法，而是将选择建模为 $K$ 个独立的二元决策。系统采用 Straight-Through Gumbel-Softmax 重参数化技术，在温度参数 $\tau$ 下对每个候选独立采样软概率，前向传播时采用硬阈值（arg max）离散化为 $\{0, 1\}$，反向传播时将梯度通过连续松弛项回传： $$w_k = \text{GumbelSoftmax}(s_k, \tau)[\text{include}] \in \{0, 1\}$$ 通过 Straight-Through 操作 $w_k = w_k^{\text{hard}} + (w_k^{\text{soft}} - \text{stop\_gradient}(w_k^{\text{soft}}))$，优化信号能够顺畅穿透离散采样层更新打分网络及扰动编码器。最终入选的扰动三元组经过单层线性映射与二值门控系数加权求和，聚合成全局上下文表征 $z = \sum_{k=1}^K w_k \cdot \text{MLP}_{\text{proj}}(c_k)$。

4. 结构稀疏性正则化机制：防止模式崩溃至全选状态

若仅在下游利用生成损失指导筛选，打分网络极易发生退化，演变为将全部 $K$ 个候选不加筛选全盘吸收的“总是全选（always-retrieve-all）”平凡解，导致重蹈朴素 RAG 噪声泛滥的覆辙。PT-RAG 在总目标函数中设计了显式的掩码稀疏惩罚项： $$\mathcal{L} = \mathcal{L}_{\text{dist}} + \lambda_{\text{sparse}} \mathcal{L}_{\text{sparse}}, \quad \text{其中 } \mathcal{L}_{\text{sparse}} = \frac{1}{K}\sum_{k=1}^K w_k$$ 其中主任务损失 $\mathcal{L}_{\text{dist}}$ 在基于 STATE 的水库架构中采用经验能量距离（Energy Distance）$\text{Energy}(\hat{x}_{\text{pert}}, x_{\text{pert}})$，在 scGPT 骨干中采用均方误差重建损失。通过设定超参数 $\lambda_{\text{sparse}}$（如 0.1 或 1.0），模型受到严格驱动，仅保留在梯度反馈中能够显著降低转录分布误差的必要扰动，迫使入选上下文保持高度精简与针对性。

损失函数 / 训练策略¶

模型优化器采用 Adam。对于基于 STATE 的模型，学习率设为 $10^{-3}$，批大小为 64，至多训练 50,000 步，每 2,000 步评估一次验证集。在此配置下，Transformer 生成骨干（8 层、隐藏层 128、12 注意力头）与输出映射随机初始化后永久冻结作为水库，仅训练 CellEncoder、PertEncoder、打分与投影网络以及最终的 Readout MLP（全量 20.97M 参数中仅 5.21M 为可训练参数）。针对 scGPT 架构，采用公开预训练检查点微调所有参数，学习率设为 $5 \times 10^{-5}$，训练 25,000 步。

实验关键数据¶

主实验¶

在 Replogle-Nadig 大规模 Perturb-seq 数据集（包含 K562、Jurkat、RPE1、HepG2 四大细胞系、2,009 个单基因扰动、共 643,413 个单细胞转录组）上进行全面评测。主实验设定包含两类严苛协议： 1. 跨细胞类型泛化（Cross-Cell-Type）：轮流将一个细胞系作为目标测试域，其余细胞系全量训练，目标细胞系仅提供约 2/3 扰动进行少样本拟合，在剩余未见过的 1,635 个扰动任务上评测。 2. 跨扰动空间泛化（Cross-Perturbation）：全细胞系可见，但 60% 的基因扰动被严格隔离作为未见测试集（冷启动推断）。

表 1 汇总了基于 STATE 水库骨干的跨细胞类型泛化平均指标（共 1,635 个测试扰动）：

指标	STATE (One-hot, 无检索)	STATE+GenePT (语义, 无检索)	Naïve RAG (K=32, 交叉注意力)	PT-RAG (本文方法)	相对最优提升/降幅
Pearson DEG ↑	0.624	0.631	0.396	0.633	+1.4% (vs STATE)
Spearman DEG ↑	0.403	0.411	0.307	0.412	+2.2% (vs STATE)
$\text{MSE}_{\text{PCA50}}$ ↓	8.43	8.42	12.64	8.39	-0.5% (显著优于 Naïve)
$W_1$ (Wasserstein-1) ↓	35.70	35.53	48.48	35.41	-0.8% (pFDR < 1e-5)
$W_2$ (Wasserstein-2) ↓	646.1	638.7	1189.5	633.7	-1.9% (vs STATE, pFDR < 1e-7)
Energy Distance ↓	9.41	9.40	14.18	9.33	-0.9% (vs STATE)

表 2 展示了跨骨干可迁移性，在完全微调的 scGPT 骨干上的跨细胞类型泛化结果：

指标	scGPT (One-hot 基线)	scGPT+GenePT (语义基线)	scGPT+PT-RAG (本文方法)	相对基线提升
Pearson DEG ↑	0.9209	0.9167	0.9128	略有下降 (-0.008)
Spearman DEG ↑	0.9500	0.9509	0.9368	略有下降 (-0.013)
$\text{MSE}_{\text{PCA50}}$ ↓	18.6440	19.8907	17.3212	-7.1% (优于 scGPT)
$W_1$ 距离 ↓	33.7992	34.6746	32.7828	-3.0% (优于 scGPT)
$W_2$ 距离 ↓	578.6169	609.7348	544.6977	-5.9% (vs scGPT, -10.7% vs GenePT)
Energy 距离 ↓	23.5341	24.3923	22.3410	-5.1% (优于 scGPT)

表 3 汇总了 3 折交叉扰动（Cross-Perturbation，60% 未见扰动测试）严苛冷启动实验表现：

评估指标	STATE+GenePT (最强非检索基线)	PT-RAG (本文方法)	统计学显著性 (pFDR)
Pearson DEG ↑	0.6377 ± 0.0367	0.6314 ± 0.0394	未达显著差异
Spearman DEG ↑	0.4293 ± 0.0344	0.4224 ± 0.0381	未达显著差异
$\text{MSE}_{\text{PCA50}}$ ↓	1.1477 ± 0.2243	1.1316 ± 0.2150	$p < 0.01$ (†)
$W_1$ 距离 ↓	0.8043 ± 0.0582	0.8016 ± 0.0578	$p < 0.01$ (†)
$W_2$ 距离 ↓	0.6436 ± 0.1134	0.6367 ± 0.1093	$p < 0.01$ (†)
Energy 距离 ↓	0.3385 ± 0.0459	0.3357 ± 0.0442	$p < 0.01$ (†)

消融与分析实验¶

为进一步剖析模块贡献与参数敏感性，作者在 HepG2 数据集以及单折交叉扰动任务上进行了深度消融（见表 4）。

表 4：关键机制与检索池设计消融实验（交叉扰动单折基准，K=32）：

实验配置 / 变体	Pearson DEG ↑	Spearman DEG ↑	$\text{MSE}_{\text{PCA50}}$ ↓	$W_1$ 距离 ↓	$W_2$ 距离 ↓	Energy 距离 ↓
STATE+GenePT (无检索基准)	0.6409	0.4314	1.3056	0.8166	0.7173	0.3618
PT-RAG 全集测试池 (2,008 个扰动)	0.6242	0.4139	1.2556	0.8073	0.6939	0.3511
PT-RAG 仅训练+验证池 (802 个扰动)	0.6237	0.4135	1.2552	0.8072	0.6937	0.3511
PT-RAG 仅训练池 (669 个扰动)	0.6237	0.4135	1.2551	0.8072	0.6936	0.3510
冻结 PertEncoder（仅训练打分+投影器）	0.6432	0.4284	1.2880	0.8157	0.7104	0.3628
鲁棒性：注入 60% 强噪候选扰动	0.6274	0.4171	1.2675	0.8096	0.6990	0.3541

关键发现¶

朴素 RAG 产生灾难性性能劣化：固定余弦相似度 + Cross-Attention 的 Naïve RAG 在所有指标上均严重崩溃（Pearson 暴跌至 0.396，$W_2$ 激增至 1189.5）。对 $K \in \{2, 5, 10, 32\}$ 的连续消融证明扩大候选池并未缓解劣势，证明未与细胞状态可微对齐的生物信息检索只会引入致命噪声。
检索增益的核心在于“群体几何流形修正”：无论在 STATE 还是 scGPT 上，PT-RAG 均大幅刷新了 $W_1$、$W_2$ 及 Energy 距离等流形几何与高阶分布距离，在 scGPT 架构下将 $W_2$ 距离断崖式压缩了 5.9%（相比 scGPT+GenePT 则降低 10.7%），但在基因层级的相关系数指标上并未带来暴涨。这证实检索上下文的核心功能是修正多细胞响应的群体几何分布。
真正习得了细胞特异性的检索行为：在测试集共享的 33 个核心基因跨细胞系 Jaccard 相似度分析中，非对角线相似度均值仅为 0.191（即对于完全相同的查询基因，不同细胞系之间选取的 Top-10 扰动重合度仅约 19%），强有力地证实打分网络根据细胞内特定转录微环境实施了个性化上下文筛选。

亮点与洞察¶

证明了生物大模型中 RAG 必须具备“可微筛选”：在自然语言处理中由于语义空间良好对齐，非可微离散检索常能奏效；而本文从反面揭示了生物计算的本质差异——缺乏可微打分与细胞状态条件化，RAG 会对生成模型造成不可逆的毒害。
巧借 Straight-Through Gumbel-Softmax 构建并行子集门控：跳出传统检索选 Top-1 的框架，将 $K$ 个候选转化为独立的 Bernoulli 式二值决策，配合稀疏正则项 $\mathcal{L}_{\text{sparse}}$，优雅解决了模型全盘吸收导致的过拟合难题。
模块化插件式设计展现跨骨干普适性：既能在轻量级随机水库（STATE Reservoir）中以仅 5M 训练参数击败重度调优基线，又能在 50M+ 规模的单细胞大模型 scGPT 上保持无缝迁移并全面优化流形指标。

局限与展望¶

第一阶段初筛仍为细胞非特异性：当前的初筛依靠 GenePT 文本向量的余弦相似度完成，如果某种在目标细胞系中起决定性代偿作用的扰动未能进入 GenePT 的前 $K$ 个语义近邻，第二阶段的可微模块将无从检索。
计算复杂度随扰动空间扩大而上升：在跨扰动大池推断中，模型必须针对 2,000+ 个扰动计算三元组相关度，在 64 细胞批次下推断开销相比非检索基准有显著增加（FLOPs/Cell 升高约 3 倍）。
仅聚焦于单基因扰动：未拓展至更加复杂的组合基因敲除（如双基因、多基因协同效应）及跨化学分子扰动场景，多扰动相互作用下的动态检索空间建模仍有待攻克。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次打通单细胞转录生成与端到端可微检索，提出针对生物异质性的条件化子集选择机制。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖跨细胞系、跨扰动冷启动、跨骨干（STATE/scGPT）以及跨技术平台（CRISPRa），数据完备。
写作质量: ⭐⭐⭐⭐⭐ 动机阐述深刻，反向印证实验严密，理论公式与可视化分析逻辑闭环。
价值: ⭐⭐⭐⭐⭐ 为功能基因组学与计算虚拟细胞设计提供了全新的系统性知识注入范式。