面向基因扰动细胞响应预测的检索增强生成框架 Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation

完整方法按以下顺序执行。第一步，构建扰动数据库 D：对训练集中全部 $|\mathcal{P}|\approx 2009$ 个扰动预先计算 GenePT 嵌入 $h_g^{pt} \in \mathbb{R}^{1536}$，存为矩阵；查询时用余弦相似度检索并剔除查询自身。第二步，编码控制细胞和扰动身份：xctrl 送入冻结的预训练 Cell Encoder（arcinstitute/SE-600M）得到 hctrl∈R^128；查询扰动 $h_g^{pt}$ 经单层可训练 MLP PertEncoder 投影到 hpert∈R^128。第三步，第一阶段语义检索：在 GenePT 空间做 top-K 余弦检索 $R_{pert} = \mathrm{TOPK}(h_g^{pt}, \mathcal{P})$，得到 K=32 个候选扰动索引 p(1)...p(K)；这一步高效但不可微，作用是把候选池从 2009 缩到 32。第四步，第二阶段可微评分：对每个候选 k 构造三元组 $c_k = [h_{ctrl}; h_{pert}; h_{cxt}^k]$（其中 $h_{cxt}^k = \mathrm{PertEncoder}(h_g^{pt,p(k)})$），经过 LayerNorm 后送入两层 MLPscore 输出二元 logits $s_k \in \mathbb{R}^2$（exclude / include）。第五步，Straight-Through Gumbel-Softmax 采样：在前向传播取 hard 0/1 决策 $w_k = \mathrm{GumbelSoftmax}(s_k, \tau=0.5)[\mathrm{include}]$，反向传播使用 soft 梯度的直通估计器，使检索选择可参与端到端优化。第六步，上下文聚合：每个三元组再经 MLPproj 投影为 h'k，加权求和得到 $z = \sum_{k=1}^{K} w_k h'_k$（只保留被选中的扰动）。第七步，细胞响应生成：把 z 送入基于 Llama 的 Transformer Generator 输出扰动后细胞群体 $\hat{x}_{pert}$。第八步，训练：总损失 $\mathcal{L} = \mathcal{L}_{dist} + \lambda_{sparse} \mathcal{L}_{sparse}$，其中分布损失采用 Energy Distance，稀疏损失 $\mathcal{L}_{sparse} = \frac{1}{K} \sum_k w_k$ 鼓励少选，使用 Adam 优化器（lr=1e-3, weight decay=5e-4）训练最多 50000 步，每 2000 步验证一次。

技术新颖性体现在三个层面。架构层面，本文首次将 RAG 引入单细胞扰动预测，构建了「非文本检索内容 + 非文本生成目标 + 无预训练检索器」的可微 RAG 管线，论文明确声明这是 to the best of our knowledge 的首次尝试。优化层面，作者没有直接套用 NLP 中可微 RAG 的做法（如 Gumbel-Top-K、期望效用最大化），而是针对「选择/不选一个候选」这一二元决策重新设计了 Straight-Through Gumbel-Softmax 评分器，并将 L1 稀疏惩罚作为关键组件——消融研究证明没有稀疏惩罚，模型会陷入「全选所有候选」的退化模式，性能从 Pearson 0.604 骤降到 0.134。生物学洞察层面，作者通过 Jaccard 相似度（平均仅 0.191，即 19% 重叠）和具体例子（WARS、DDRX27、MRPL39、RPS2、TPRKB）证明学到的检索模式不仅在数学上细胞类型相关，更与真实的 tRNA 合成、核糖体生物合成、细胞周期调控等通路偏好一致，这是单纯套用 NLP RAG 无法得到的结果。

通过仔细审视方法与实验，我识别出四个值得改进的弱点。弱点一，GenePT 嵌入的「语义 ≠ 细胞效应」问题没有完全解决：GenePT 基于 NCBI 描述文本编码，对一个基因的功能描述在不同细胞类型下其实是一样的，因此第一阶段检索给出的 K=32 候选在所有细胞类型中完全相同；这意味着细胞类型感知的负担全部落在第二阶段 Gumbel-Softmax 上——如果 K=32 中没有真正适合当前细胞类型的候选，第二阶段也「巧妇难为无米之炊」。改进方向：可以引入细胞状态条件化的嵌入（如 scGPT、Geneformer 这类细胞上下文基础模型）来动态调整第一阶段检索池。弱点二，K=32 是硬编码的超参数，作者虽然做了 K=16/32 消融但都是「训练后」的回看，没有自适应机制；如果查询扰动是个相对孤立的基因（无功能近邻），K=32 候选可能质量普遍不高；如果查询扰动属于高度保守通路，可能 K=8 就够了。改进方向：可以加一个「候选质量预测器」动态决定 K，或在训练中加课程学习让模型从 K=8 渐进到 K=32。弱点三，第二阶段的 MLPscore 只用了简单的 [hctrl; hpert; hcktx] 拼接和 LayerNorm，没有利用扰动间的二阶关系（如两两候选之间的关系），可能漏掉「候选之间互补 vs 冗余」的信息。改进方向：可以引入 cross-attention 让候选之间互相通信，类似 Set Transformer 的思路。弱点四，计算开销不均衡：当前 PT-RAG 对所有扰动都做完整的两阶段检索，但实际上对 K562 这类训练充分的细胞类型，检索的边际收益可能很小；可以加入「何时跳过检索」的早退机制，对高置信度预测直接走 baseline 路径。

作者在结论部分明确提出了三个未来方向：（1）把 PT-RAG 拓展到组合扰动（多基因同时敲除）、化学化合物扰动和 CRISPRa/i 激活/抑制场景，因为现实药物筛选和合成生物学往往涉及多模态扰动；（2）探索 GraphRAG，把基因调控网络（如 STRING、GeneMANIA 数据库）作为图结构先验融入检索，让候选扰动不仅是「语义近」还是「调控网络近」；（3）对学到的选择模式做系统化生物分析，把 PT-RAG 视为「基因功能关系发现工具」补充现有 pathway 数据库。基于本文成果可延伸的方向还包括：把 PT-RAG 的可微检索思想应用到单细胞多组学（RNA+ATAC+protein 联合预测）；用更先进的离散-连续混合优化（如 REINFORCE 替代 Gumbel-Softmax）减少训练方差；以及把 19% Jaccard 重叠率本身作为下游任务——预测「某基因在两个细胞类型间检索模式的相似度」可以作为新的细胞类型相似度度量。

复现性整体较好但有一定门槛。代码已开源在 https://github.com/difra100/PT-RAG_ICLR，包含完整的训练脚本和评估流程。数据完全公开：Replogle 原始数据来自 Replogle et al. 2022 (Cell) 和 Nadig et al. 2024 (bioRxiv)，GenePT 嵌入也可公开下载，扰动数据库基于这二者交集的 2009 个扰动。预训练模型方面使用了 HuggingFace 上的 arcinstitute/SE-600M 作为冻结的 Cell Encoder，需要先下载。算力需求：单 NVIDIA A100 40GB，每个目标细胞类型训练到收敛需要 30k-40k 步，约 8-10 小时；4 个细胞类型共需约 32-40 小时 A100 时。超参数方面作者给出了完整配置（Adam, lr=1e-3, weight decay=5e-4, batch size 64, 序列长度 64 cells, K=32, τ=0.5, λsparse=0.1）。难度评估：中等偏上，主要工作量在（a）搭建 GenePT 检索流水线、Cell Encoder 嵌入和扰动数据库，（b）实现 Straight-Through Gumbel-Softmax 的反向传播（PyTorch 有现成的 torch.nn.functional.gumbel_softmax 但需注意 hard=True 的实现细节），（c）Mann-Whitney U + Benjamini-Hochberg FDR 统计检验的复现。论文明确标注 LLM 仅用于文本润色而非科学构思，降低了「方法是否真的来自 LLM」的疑虑。