Self-Improvement Imitation with Biologically Guided Search for Protein Design Under Oracle Budgets¶
会议: NeurIPS 2026
arXiv: 2605.26690
OpenReview: https://openreview.net/forum?id=iQpv99SOhF
代码: https://github.com/grimmlab/SILO.git
领域: 计算生物
关键词: 蛋白质设计, 自我改进模仿学习, 随机集束搜索, 丙氨酸扫描, 主动学习
一句话总结¶
针对严格预言机(Oracle)评估预算下的蛋白质序列优化,本文提出了 SILO 框架,通过两阶段分层编辑策略、不放回增量随机集束搜索与融入丙氨酸扫描适应度得分的主动学习选择,并仅依据真实评测优质轨迹进行轨迹级监督模仿微调,避免了强化学习在噪声代理环境中的价值崩塌。
研究背景与动机¶
在蛋白质工程中,通过计算方法对初始蛋白质序列进行定点突变以优化热稳定性、结合亲和力或催化活性等目标适应度(Fitness),是生物医药与酶工程的核心课题。然而,蛋白质序列的组合搜索空间极度庞大(长度为 \(L\) 的序列对应 \(20^L\) 种可能),且突变之间存在高度非线性的上位效应(Epistasis),导致适应度景观(Fitness Landscape)极其崎岖稀疏。更关键的是,高精度的湿实验测量或昂贵的结构评估通常面临严苛的预算限制,要求算法在极少数的主动学习评估轮次内,既要高效探索有潜力的突变空间,又要保证每次采样的有效性。
现有的强化学习(如 PPO)与离线生成模型(如 GFlowNets、离散扩散模型)在该设置下面临显著瓶颈。一方面,基于价值函数的优化方法高度依赖代理模型(Proxy)预测作为奖励,当代理模型存在近似误差或分布偏移时,价值估计极易在分布外候选上产生灾难性的误差累积与过度拟合;而非策略生成模型若不更新参数则无法针对具体任务持续积累有效知识。另一方面,常规突变提议往往是位置不可知的(Position-Agnostic),盲目修改保守位点极易破坏蛋白质的空间折叠与核心功能。
本文的切入角度是:借鉴神经组合优化中的自我改进模仿学习(Self-Improvement Learning, SIL)思想,完全绕开脆弱的价值函数拟合,直接使用每轮经真实 Oracle 评测验证的高分序列作为自生成的“伪专家”示范。同时,在突变探索阶段引入结合丙氨酸扫描(Alanine Scanning)先验的结构化采样,确保所探索的突变位点具备生物学合理性。核心 idea:将蛋白质序列编辑建模为分层决策过程,利用不放回增量随机集束搜索配合丙氨酸扫描生物先验筛选候选,并在严格 Oracle 预算下通过对高适应度轨迹的自模仿交叉熵监督迭代更新策略。
方法详解¶
整体框架¶
SILO 将固定 Oracle 预算下的蛋白质序列设计构建为一个主动学习循环系统。系统由四个核心步骤闭环运转:首先,利用当前已获得的 Oracle 标注数据集训练代理模型集成(由 3 个 1D-CNN 构成的 Proxy Ensemble);其次,基于上一轮搜索到的全局最优序列作为起点,冻结的 ESM Cambrian 提取序列演化表征,由轻量 Transformer 策略网络通过分层动作采样生成候选突变,并使用增量无放回随机集束搜索(SBS)探索多样化轨迹;第三,代理集成计算候选序列的 UCB 预估值,并结合丙氨酸扫描突变评分(AFS)预测编辑位点的耐受性,贪心地挑选出 Top-\(K\) 个最具潜力且生物学扰动小的序列提交给真实 Oracle 评测;最后,从本轮 Oracle 实际测得的高分候选及其编辑轨迹中提取动作对,以交叉熵损失监督策略网络进行轨迹级模仿更新,引导下一轮采样概率质量聚焦于高收益突变空间。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["当前标注数据集 D_n-1 与最佳起点 x_start"] --> B["代理集成训练与置信度建模"]
B --> C["分层编辑策略生成"]
C --> D["增量随机集束搜索采样"]
D --> E["丙氨酸扫描生物先验打分与 UCB 候选筛选"]
E --> F["Oracle 评测与数据集扩充"]
F --> G["轨迹级自我改进模仿学习"]
G -->|更新策略网络参数| C
关键设计¶
1. 分层编辑策略网络与马尔可夫决策过程:将组合突变空间解耦为位置与氨基酸两阶段决策
传统的序列生成模型直接在 \(L \times |\mathcal{V}|\) 的联合动作空间上采样,动作空间巨大且难以捕捉残基之间的位置敏感性。SILO 将候选设计建模为步长为 \(d_{\max}\) 的马尔可夫决策过程(MDP),每步编辑拆解为双层动作:第 0 层动作 \(a_0^{(t)} \in \{1, \ldots, L\}\) 负责选择待突变位置;第 1 层动作 \(a_1^{(t)} \in \mathcal{V}\)(其中 \(|\mathcal{V}|=20\))负责在该选定位置下选择替换的目标氨基酸。在模型实现上,输入序列先输入包含 3 亿参数的冻结 ESM Cambrian 蛋白质语言模型提取进化表征,再输入 2 层带 FlashAttention 的轻量 Transformer 决策网络,输出维度为 \(D=512\) 的隐状态矩阵 \(H \in \mathbb{R}^{L \times D}\)。两个策略头分别产生位置 Logits \(h_0 \in \mathbb{R}^L\) 与残基 Logits \(h_1 \in \mathbb{R}^{L \times |\mathcal{V}|}\)。这种分层因式分解有效降低了联合搜索复杂度,实现了参数高效共享。
2. 增量随机集束搜索采样:无放回 Gumbel 扰动防止采样坍缩并增强探索多样性
标准贪心集束搜索在策略分布较为尖锐时容易陷入单一局部极值,导致生成的编辑序列同质化严重。SILO 引入不放回增量随机集束搜索(Incremental Stochastic Beam Search, SBS)。在构建部分编辑轨迹 \(a_{0:t}\) 时,对策略的对数概率加入 Gumbel 噪声扰动: $\(\tilde{G}(a_{0:t+1}) = \log \pi_\theta(a_{0:t+1}) + \epsilon, \quad \epsilon \sim \text{Gumbel}(0, 1)\)$ 集束保留扰动得分最高的 \(\beta\) 个候选。在增量设置中,每当一条完整轨迹采样完成,其对应的概率质量便会从策略分布中扣除并对剩余概率重新归一化。这确保了后续采样严格在未被覆盖的概率空间中进行,在不增加额外网络调用的前提下,大幅增强了突变轨迹的全局覆盖度与多样性。
3. 融入丙氨酸扫描先验的主动学习候选选择:结合 UCB 不确定性与功能容忍度评分
代理模型在远离已知数据的未知区域容易给出盲目的高估分,盲目信任代理分数常导致选择已被破坏功能的突变体。受实验生物学中丙氨酸扫描突变技术的启发,SILO 设计了综合考虑预期性能与位点耐受性的复合评分函数。对于突变方案选定的位置集合 \(M\),构造对应的丙氨酸替换变体 \(A(x)\)(将 \(M\) 中位点均换为中性氨基酸丙氨酸),并使用代理集成同时评估候选序列 \(x\) 与中性变体 \(A(x)\): $\(S(x) = \text{UCB}_{\text{proxy}}(x) + \text{UCB}_{\text{AFS}}(x)\)$ 其中 \(\text{UCB}_{\text{proxy}}(x) = \mu_\phi(x) + \gamma_1 \sigma_\phi(x)\) 代表序列本身的上置信界(探索潜在高适应度);而 \(\text{UCB}_{\text{AFS}}(x) = \mu_\phi(A(x)) + \gamma_2 \sigma_\phi(A(x))\) 则度量突变位点对中性替换的耐受程度。若变体 \(A(x)\) 的预测值同样较高,表明突变发生在功能容忍区域而非关键活性中心,降低了误损主干功能的风险。算法在采样池中贪心选取 \(S(x)\) 最高的 Top-\(K\) 个序列提交真实 Oracle 评估。
4. 轨迹级自我改进模仿学习:绕过价值函数崩溃的真实表现监督微调
在有限且带有噪声反馈的环境中拟合 Critic 或 Value 网络极易因外推误差产生价值崩塌。SILO 完全抛弃了强化学习中的价值估计,直接将经过真实 Oracle 评估验证的最高适应度序列集(BESTFOUND,取本轮评估中的 Top-\(P\) 序列)视为自生成的专家示范。对于这些优质序列所对应的编辑轨迹,提取历史动作上下文与后续动作对 \((h^{(t)}, a^{(t)})\),使用标准的自回归交叉熵损失更新策略网络: $\(\mathcal{L}(\theta) = -\frac{1}{m} \sum_{(h^{(t)}, a^{(t)})} \log \pi_\theta(a^{(t)} \mid h^{(t)})\)$ 通过在每轮主动学习中对经 Oracle 认证的优质突变轨迹执行行为克隆,策略网络逐渐将概率质量浓缩在真正产生高增益的序列子空间中,实现了鲁棒且单调递增的闭环自我进化。
实验关键数据¶
主实验¶
实验在 8 个基准蛋白质适应度景观上展开,严格限制主动学习轮次 \(N=10\)、每轮 Oracle 查询预算 \(K=128\)(总预算仅 1280 次评估),每轮突变预算 \(d_{\max} \in \{1, 2, 3\}\)。基线涵盖进化算法(PEX、AdaLead)、离线/保守强化学习(GFN-AL-\(\delta\)CS)、基于语言模型定向进化的 MLDE 以及基于扩散模型推理制导的 SOTA 方法 ProSpero。所有基线均基于相同环境与超参重新复现(Mean \(\pm\) Std,5 次独立运行)。
| 数据集 (起始适应度 \(O(x_{\text{start}})\)) | 指标 (Max Fitness) | 本文 (SILO) | 最优基线 (Method) | 提升 |
|---|---|---|---|---|
| AAV (0.500) | 最大适应度 | 0.747 ± 0.014 | 0.697 ± 0.012 (AdaLead) | +0.050 |
| AMIE (0.224) | 最大适应度 | 0.262 ± 0.004 | 0.245 ± 0.004 (PEX) | +0.017 |
| E4B (7.743) | 最大适应度 | 8.169 ± 0.052 | 8.143 ± 0.021 (ProSpero) | +0.026 |
| GFP (3.572) | 最大适应度 | 3.619 ± 0.002 | 3.615 ± 0.004 (ProSpero) | +0.004 |
| LGK (0.020) | 最大适应度 | 0.045 ± 0.003 | 0.043 ± 0.001 (ProSpero) | +0.002 |
| Pab1 (0.840) | 最大适应度 | 1.894 ± 0.348 | 1.799 ± 0.324 (PEX) | +0.095 |
| TEM (1.229) | 最大适应度 | 1.233 ± 0.003 | 1.233 ± 0.002 (PEX) | 持平 / 领先 |
| UBE2I (2.978) | 最大适应度 | 2.997 ± 0.003 | 2.991 ± 0.002 (PEX/ProSpero) | +0.006 |
在前 100 篇高质量候选的平均适应度(Mean Fitness of Top-100)评测中,SILO 同样在全部 8 个景观中取得第一(例如 AAV 上达 \(0.727 \pm 0.014\),大幅超越 AdaLead 的 \(0.661\) 与 ProSpero 的 \(0.632\);AMIE 上达 \(0.259 \pm 0.005\),显著超越基线的 \(0.234\)),证明其提升不仅在于极个别离群高分序列,更显著提升了高价值候选群体的整体质量。
消融实验¶
为了量化结构化采样、生物先验打分与策略自模仿微调各自的贡献,作者在 AAV、GFP 与 Pab1 上进行了严格的模块消融对比(10 轮主动学习后的最大适应度,5 次独立运行):
| 配置 | AAV (Max Fitness) | GFP (Max Fitness) | Pab1 (Max Fitness) | 说明 |
|---|---|---|---|---|
| SILO (完整框架) | 0.747 ± 0.014 | 3.619 ± 0.002 | 1.894 ± 0.348 | 完整模型:SBS采样 + AFS打分 + 迭代自模仿更新 |
| Frozen policy + SBS | 0.706 ± 0.013 | 3.615 ± 0.001 | 1.530 ± 0.268 | 冻结随机初始化策略不训练,仅靠 SBS 与 AFS 搜索 |
| SBS + w/o AFS | 0.694 ± 0.019 | 3.615 ± 0.001 | 1.558 ± 0.134 | 移除丙氨酸扫描先验,仅用代理 UCB 挑选候选 |
| BS + AFS | 0.587 ± 0.028 | 3.602 ± 0.012 | 1.022 ± 0.146 | 将随机集束搜索(SBS)替换为确定性集束搜索(BS) |
| Random baseline | 0.530 ± 0.015 | 3.600 ± 0.010 | 1.018 ± 0.103 | 无模型引导的随机均匀突变 |
关键发现¶
- 随机集束搜索是探索多样性的核心支柱:将 SBS 替换为确定性 BS 后性能发生剧烈滑坡(AAV 从 0.747 骤降至 0.587,Pab1 从 1.894 骤降至 1.022),表明在崎岖离散景观中,基于 Gumbel 扰动的无放回遍历能有效打破贪心局部陷阱。
- 生物学先验(AFS)具有关键正则化作用:在 AAV 和 Pab1 上,去掉 AFS 后最大适应度分别下降 0.053 和 0.336,证实将丙氨酸扫描作为加性排名信号能够有效阻断代理模型对敏感功能位点突变的虚假高估。
- 自模仿学习提供持续提升动力:即便冻结策略网络仅依靠 SBS + AFS 也能达到较强性能(如 GFP 上达 3.615),但在复杂任务(AAV、Pab1)上,通过 Oracle 高分轨迹进行自回归模仿更新带来了决定性提升(AAV 从 0.706 提升至 0.747,Pab1 从 1.530 跃升至 1.894)。
- 小样本与高噪声环境下的极强鲁棒性:在仅保留 10% 初始标注数据的低数据压力测试(UBE2I/TEM)以及添加高斯噪声的人工破坏代理测试(AMIE/E4B)中,传统 RL 基线因代理漂移而严重崩塌,而 SILO 因策略仅拟合真实 Oracle 优质轨迹,表现极其稳健。
亮点与洞察¶
- 将神经组合优化的自我模仿范式成功迁移至生物计算:传统的自模仿学习假设评估成本低廉且可无限调用,SILO 创新性地在采样生成与 Oracle 评测之间插入了结合生物先验的 UCB 代理筛选层,完美适配了严苛预算约束。
- 软性加性利用丙氨酸扫描信号:以往工作将丙氨酸扫描作为硬性阈值过滤器(容易误杀边界候选),SILO 将其作为连续的加性排序项融进 UCB,平滑量化了位点破坏风险。
- 彻底根治基于代理的强化学习价值崩塌难题:不依赖 Critic 网络或 Bellman 误差传递,以行为克隆的监督范式直接拟合高分轨迹,在有噪声的离线代理环境中展现出极佳的工程稳定性。
局限与展望¶
- 自模仿收敛导致的局部搜索局限:自模仿天然会使策略网络的采样分布向已发现的高分邻域急剧坍缩(多样性 Diversity 指标低于无约束扩散基线 ProSpero),若早期不幸跌入次优吸引盆,容易发生自我强化锁定。
- 变异步长受限于局部阶数:当前每轮突变预算限制为 1-3 个点突变,对于需要大规模结构重排或高阶协同突变的全新功能蛋白设计,局部累积策略可能难以跨越深层适应度峡谷。
- 完全依赖计算预言机(In Silico Oracle):所有实验均基于计算代理或已有高通量筛选数据集的模型插值评测,尚未进行湿实验真实验证,真实生物体系中的批次效应与物理化学不可测性有待进一步检验。
相关工作与启发¶
- vs ProSpero: ProSpero 依赖冻结的大规模扩散生成模型(EvoDiff)并在推理阶段施加硬性引导,模型参数在优化过程中完全不更新;SILO 则采用轻量级策略网络,随主动学习轮次通过 Oracle 轨迹实现自进化微调,更善于在目标靶点邻域进行深度开发。
- vs GFN-AL-\(\delta\)CS: GFlowNets 依赖离线数据集与保守掩码去噪训练生成流网络,当初始数据稀缺时泛化能力骤降;SILO 的模仿学习直接锚定主动探索中 Oracle 揭示的高分轨迹,在低数据测试中稳定性显著占优。
- vs DyNA-PPO / EvoPlay: 基于价值网络的强化学习在面对有噪声的代理模型时容易出现价值过高估计与策略崩溃;SILO 采用纯监督的下一动作交叉熵拟合,在代理强噪声扰动下依然保持高鲁棒性。
评分¶
- 新颖性: ⭐⭐⭐⭐☆ (将自改进模仿学习引入严格预算的蛋白质设计,分层动作解耦与加性 AFS 融合构思精巧)
- 实验充分度: ⭐⭐⭐⭐⭐ (覆盖 8 个真实适应度景观,涵盖低数据与高噪声压力测试,消融实验设计严密)
- 写作质量: ⭐⭐⭐⭐⭐ (逻辑链条清晰,图表完备,数学公式与生物学机制阐释严谨)
- 价值: ⭐⭐⭐⭐⭐ (为严苛实验预算下的蛋白质定向进化与主动学习序列设计提供了高稳定、易复现的新范式)