跳转至

Budgeted Quotient-Residual Guidance for Frozen Pocket-Conditioned Molecular Diffusion

会议: NeurIPS 2026
arXiv: 2609.31222
领域: 计算生物
关键词: 分子扩散模型、商空间残差引导、置信域预算、冷冻骨干网络、靶点口袋条件生成

一句话总结

针对口袋条件分子扩散模型中商空间先验引导梯度微弱易钝化的问题,提出预算化商残差引导(QRG),将商余向量度量水平提升至环境空间,并以冷冻采样器自身运动步长严格约束引导范数,实现免重新训练的高质量先导化合物生成与编辑。

研究背景与动机

基于蛋白口袋条件的3D分子生成与先导化合物优化是计算机辅助药物发现的核心任务。扩散模型通常在环境原子坐标空间中执行逆向去噪采样,但药物化学中大部分关键设计目标——如局部片段几何构型、母核骨架保持、连接子闭环连通度以及配体-口袋相互作用接触,本质上均定义在商空间几何特征上,即排除了全局旋转平移、原子索引置换及无关局部坐标等冗余自由度的等价类。即便底层骨干网络具备几何等变性,环境去噪坐标与商空间优化目标之间的表征不匹配依然显著存在。

在先导化合物优化流水线中,针对每一个特定化学设计目标重新训练或微调庞大的等变扩散骨干网络代价极其高昂,因此保持预训练生成骨干网络参数冻结、在推理阶段引入即插即用式引导成为极具吸引力的范式。然而,直接在环境坐标上对商空间不变能量函数施加局部梯度引导存在严重的“梯度钝化”现象:由于商特征映射的尺度与采样器的离散去噪步长脱节,商能量反向传播至环境空间的局部梯度方向虽然物理明确,数值尺度却往往过小,在有限步长的去噪漂移中几乎处于休眠状态;若采用昂贵的多步前瞻展开(rollout)来激活信号,又会带来沉重的采样耗时与算力负担。

面对这一两难困境,核心症结不在于寻找某种人为调整的绝对梯度步长,而在于如何让商空间能量梯度以采样器自身步长为尺度进行传递。核心 idea:提出预算化商残差引导(QRG),通过度量水平提升将商余向量投影为无垂直冗余的环境几何运动方向,并严格将引导校正幅度限制在冷冻采样器单步更新范数设立的置信预算内,实现“商几何定方向、采样器定尺度”的高效免重训控制。

方法详解

整体框架

预算化商残差引导在冷冻分子扩散采样器的每一个离散逆向步骤中运行。整个框架接收当前带噪分子构型与蛋白口袋环境,首先调用预训练冷冻去噪网络计算无引导的基础更新向量并确定置信预算;随后根据任务需求分别构建商空间残差目标与截面保留目标,经过水平提升消除冗余规范自由度;最后在乘积置信预算下将独立截断的双路校正增量叠加到采样步中,驱动分子构型演化。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入:带噪坐标 x_t 与口袋环境"] --> B["1. 冷冻采样器步长基准与置信预算<br/>前向计算基础位移 v_t 与范数预算 B_t"]
    B --> C["2. 商余向量度量水平提升<br/>计算商微分 J_t 并消除垂直冗余求得 h_t"]
    C --> D["3. 乘积预算截断与残差控制<br/>截面稳定与商残差双分支独立约束合成 u_t"]
    D --> E["4. 轻量化商特征代用与部署落地<br/>PredNext 预测状态或局部半径特征计算"]
    E --> F["输出:校正后的去噪坐标 x_{t-1}"]

关键设计

1. 冷冻采样器步长基准与置信预算:解决引导尺度脱节与梯度钝化

针对环境坐标空间中商能量梯度尺度过小、被离散扩散步长掩盖的失效模式,QRG 确立了以生成模型自身动态为绝对参考基准的自适应步长机制。在时间步 \(t\),设冷冻去噪网络给出的原始环境更新位移为 \(v_t = v_\theta(x_t, t)\),在度量张量 \(M_t\) 下其步长范数为 \(\|v_t\|_{M_t}\)。该机制不使用固定的外部超参数学习率,而是将允许的最大引导幅度约束在局部置信半径 \(B_t = \rho_t \|v_t\|_{M_t}\) 之内,其中 \(\rho_t\) 为相对预算比例。从动力学视角来看,当高斯逆向扩散核具备协方差 \(\tau_t M_t^{-1}\) 时,该置信椭球约束严格等价于限制逆向转移核与引导修正核之间的局部条件 KL 散度 \(\mathrm{KL}(P_t^u(\cdot \mid x_t) \parallel P_t(\cdot \mid x_t)) \le \kappa_t = \rho_t^2 \|v_t\|_{M_t}^2 / (2\tau_t)\),而在确定性采样器下则退化为局部动能预算。这确保了外部引导信号永远不会强行破坏冷冻扩散先验的生成流形,彻底消除了局部梯度尺度与采样演化尺度失配的问题。

2. 商余向量度量水平提升:消除垂直冗余方向的高效投影

为了将定义在局部商特征(如距离矩阵、口袋相互作用锚点距离)上的任务能量函数 \(E_{\mathrm{quo}}(x; y) = \ell(q(x), y)\) 准确回传至环境原子坐标,必须剔除任何不改变商特征纯属表象旋转平移的“垂直运动”。设商特征映射的雅可比微分矩阵为 \(J_x = Dq_x \in \mathbb{R}^{m \times d}\),其核空间 \(\ker J_x\) 即对应环境空间中的垂直自由度;商余向量 \(c_x = \nabla_z \ell(z, y)|_{z=q(x)}\) 通过度量张量 \(M_x\) 被提升至环境空间的度量水平向量: $\(h_x = M_x^{-1} J_x^\top c_x\)$ 该向量在 \(M_x\) 内积意义下与 \(\ker J_x\) 严格正交,保证了所求方向是单位范数约束下使商目标函数下降最快的最速下降方向。结合置信预算 \(B_t\),在活跃预算状态下的解析闭式解形式为 \(u_t^{\mathrm{quo}} = -B_t \frac{h_t}{\|h_t\|_{M_t} + \varepsilon}\)。当等变旋转平移作用在分子构型上时,该水平提升与闭式更新满足严格的等变性传递,确保了几何对称性在引导过程中的数学自洽。

3. 乘积预算截断与残差控制:解耦母核代表保留与商空间优化

在先导化合物优化(如片段生长、骨架跃迁、连接子设计)中,系统往往面临双重约束:既要保持既定分子母核或锚点原子的具体三维坐标代表(截面分支),又要对待生成区域施加商空间相互作用或几何正则(残差分支)。传统的单分支标量化混合能量求梯度再截断的方法无法独立调控两者的物理约束边界。QRG 引入乘积置信域机制,将总位移校正向量 \(u_t\) 解耦为独立的截面分支 \(h_t^{\mathrm{sec}} = M_t^{-1} \nabla_x E_{\mathrm{sec}}(\tilde{x}_t)\) 与商残差分支 \(h_t^{\mathrm{res}} = M_t^{-1} J_t^\top c_t\): $\(u_t = -\min\left\{1, \frac{\rho_s \|v_t\|_{M_t}}{\|h_t^{\mathrm{sec}}\|_{M_t} + \varepsilon}\right\} h_t^{\mathrm{sec}} - \min\left\{1, \frac{\rho_r \|v_t\|_{M_t}}{\|h_t^{\mathrm{res}}\|_{M_t} + \varepsilon}\right\} h_t^{\mathrm{res}}\)$ 数学证明表明,该乘积空间预算分配所能探索的动力学更新方向,无法通过任何固定标量化加权 \(\nabla E_{\mathrm{sec}} + \lambda \nabla E_{\mathrm{res}}\) 单独归一化重现。这种解耦赋予了推理过程灵活控制保守结构保留强度与开放构型优化力度的自由度。

4. 轻量化商特征代用与部署落地:规避昂贵回滚展开的实用方案

理论上的理想引导梯度需要通过多步扩散前瞻展开(Rollout Teacher)来评估,但在实际化学筛选中多倍的去噪前向传播完全不可接受。QRG 针对落地部署给出了两种轻量化实现方案:第一种为单步预测代用(PredNext-QRG),直接复用冷冻采样器已计算出的单步去噪估计 \(\tilde{x}_t = x_t + v_t\) 作为状态求导点,并在采样轨迹中按每隔两步稀疏触发残差校正,实现极低的额外推理开销;第二种为局部半径代用(Local-QRG),通过截断远距离无意义相互作用、仅在局部有效空间半径内计算稀疏商特征雅可比矩阵。二者彻底将前瞻展开解耦为仅用于上界分析的诊断基准,构建了高度实用的免重训推理加速流水线。

实验关键数据

主实验

在 CBGBench 评估协议下,采用冷冻 TargetDiff 作为等变分子生成骨干网络,在统一的 Seed-0 官方测试集上针对分子片段生长(Fragment Growing)、骨架跃迁(Scaffold Hopping)和连接子设计(Linker Design)三大先导化合物编辑任务展开评测。

任务 方法 有效分块数 生成有效率 (Validity) ↑ \(\Delta\) 有效率 采样耗时 (sec/sample) 引导/基准范数比 新颖性 (Novelty) ↑ 多样性 (Diversity) ↑
Fragment Base 5/5 0.815 – 16.68 0.000 1.000 0.898
Fragment PredNext-QRG 5/5 0.829 +0.014 18.28 0.192 1.000 0.899
Fragment Local-QRG 5/5 0.864 +0.049 19.91 0.265 1.000 0.896
Scaffold Base 5/5 0.664 – 17.99 0.000 1.000 0.913
Scaffold PredNext-QRG 5/5 0.683 +0.019 20.36 0.194 1.000 0.911
Scaffold Local-QRG 5/5 0.707 +0.043 20.18 0.268 1.000 0.912
Linker Base 5/5 0.681 – 18.82 0.000 1.000 0.888
Linker PredNext-QRG 5/5 0.680 -0.001 19.72 0.183 1.000 0.892
Linker Local-QRG 5/5 0.712 +0.031 21.38 0.251 1.000 0.891

消融实验

为了验证引导增益来源于真实商残差信号而非任意次级扰动或昂贵多步展开,作者在多随机种子下的分子片段生成(Fragment)任务上对比了不同引导机制与前瞻展开配置。

配置分支 评测种子数 生成有效率 (Validity) ↑ 相对运行耗时 引导/基准位移比 说明
Base 3 0.842 1.00× 0.000 冷冻无引导基准采样器
Section-only 1 0.896 1.07× – 仅保留截面母核分支,残差预算设为 0
Sham-QRG 3 0.888 1.22× 0.255 注入洗牌打乱的伪残差信号作为负对照
PredNext-QRG 3 0.896 1.10× 0.193 部署版:单步前瞻预测代用(隔步更新)
Local-QRG 3 0.905 1.22× 0.271 部署版:局部半径截断商特征代用
Rollout-QRG Teacher 3 0.898 2.15× 0.253 诊断版:显式单步完整前瞻展开(耗时倍增)

关键发现

  • 显著激活休眠信号:在受控合成玩具任务(2D 高斯、轨道点匹配与玩具分子构型)中,直接施加原始局部商梯度时成功率与未引导基准几乎完全一致(如玩具分子 0.079 vs 0.079),而施加 QRG 预算化残差引导后成功率跃升至 0.371,确证了自适应置信半径在激活休眠梯度中的核心作用。
  • 任务异质性与质量-耗时权衡:Local-QRG 展现出最优的泛化有效率提升(Fragment +0.049、Scaffold +0.043、Linker +0.031),仅引入 15%-25% 的平滑运行耗时;而 PredNext-QRG 在轻度受限的 Fragment 和 Scaffold 上效果良好,在强拓扑闭环约束的 Linker 任务上则接近中性(-0.001)。
  • 多样性与新颖性零牺牲:在所有测试任务及部署变体中,生成分子的新颖性指标稳定保持在 1.000,且化学多样性(约 0.89-0.91)与无引导基准相比完全未受破坏,排除了模式坍缩或过拟合局部口袋的假象。
  • 侧链修饰任务暴露预算迁移失配:在极度狭窄、类似旋转异构体局部堆叠的侧链修饰(Side-chain Decoration)压力测试中,默认配体预算导致有效率从 0.873 下降到 0.837;只有通过大幅缩减预算比例并执行稀疏引导(恢复至 0.869),才能缓解过控现象,证实引导预算与几何自由度强相关。

亮点与洞察

  • 商几何与动力学尺度的解耦思想:巧妙利用微分几何中的水平提升消除代表选取的垂直冗余,并将引导向量长度锚定在去噪采样器自身的动力学步长上,以极简的数学形式解决了扩散引导长期存在的步长难以调节的问题。
  • 严格对应条件 KL 散度与动能约束:揭示了置信域更新在高斯扩散核下正是局部最小条件相对熵控制的解析极值点,为经验性的启发式梯度截断赋予了坚实的随机最优控制理论底座。
  • 求真务实的严谨评估范式:不盲目宣称绝对亲和力提升,主动揭示分子对接得分(AutoDock Vina)在部分体系上的波动性与混杂性;将耗时巨大的展开版本(Rollout Teacher)归类为理论诊断工具,仅以轻量部署版本作为核心宣称,具备高度的工程诚信。

局限与展望

  • 商特征代用对复杂拓扑适用性不一:在连接子设计或高度受限的闭环几何中,简单局部半径近似难以捕捉长程构象协同,导致必须精细调整截面与残差预算比例。
  • 侧链微环境对先验预算敏感:在蛋白质侧链修饰等高密度原子碰撞场景中,预设的配体级置信预算容易过度干预几何堆积,未来需探索基于局部空间阻碍度动态退火预算的自适应机制。
  • 冷冻骨干能力的固有天花板:虽然无需重新训练极大拓展了模型适用性,但若冷冻骨干本身的化学化学价态先验极度匮乏,纯推理时引导无法根本性消除物理不可行结构的生成。

相关工作与启发

  • vs 等变扩散生成模型(TargetDiff / DiffDock):此类方法专注于在训练阶段设计复杂的 SE(3) 等变神经网络架构以习得基础分子先验;本文与之完全正交互补,专为冷冻骨干提供推理期即插即用的商几何目标控制,零重训开销。
  • vs 通用扩散引导方法(DPS / Universal Guidance):通用引导方法通常直接将环境损失的无约束梯度添加到扩散轨迹中,极易引起离散步长失稳或被微弱梯度淹没;QRG 引入商空间水平提升与步长相对预算约束,在保持等变性的同时保障了数值稳定性。
  • vs 数据对称性与轨道平均(Orbit Diffusion):Orbit Diffusion 通过在训练期做群轨道条件期望来减少方差;本文则从推理期切入,探索了对称不变性在冷冻采样时如何借助几何投影与自适应控制实现低成本高效引导。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次提出商空间度量水平提升与扩散步长相对置信预算相结合的免重训分子引导理论。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖受控玩具任务、官方基准三大核心任务、严格的消融负对照(Sham/Teacher)、侧链极限压测及靶点级异质性统计。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学推导严密自洽,概念界定清晰,对方法边界、算力账单及实验局限性论述极为诚恳客观。
  • 价值: ⭐⭐⭐⭐⭐ 为结构化药物分子生成中利用冷冻大模型进行低成本先导化合物优化提供了极具实用价值的范式。