Robust Inference-Time Steering of Protein Diffusion Models via Embedding Optimization¶

会议: NeurIPS 2026
arXiv: 2602.05285
OpenReview: https://openreview.net/forum?id=DUpxIDDWMK
代码: https://github.com/rs-station/embedopt
领域: 计算生物
关键词: 蛋白质扩散模型、推理时引导、条件嵌入优化、冷冻电镜密度拟合、贝叶斯后验采样

一句话总结¶

针对传统扩散后验采样在先验低密度区易失稳和破坏几何结构的难题，EmbedOpt 提出在推理过程中直接优化多序列比对条件嵌入以动态调整结构先验，实现了超两个数量级学习率的鲁棒性以及 4 倍采样步数压缩。

研究背景与动机¶

生物大分子具有天然的动态构象变化特征，解析其在不同功能状态下的三维构象是揭示生命机制的核心。然而，从冷冻电镜（cryo-EM）或化学交联等实验测量数据中反演结构是一个极具挑战的高维逆问题。传统基于物理力场和最大熵原理的逐步采样方法计算代价昂贵且依赖大量手工精细调优。近年来，以 AlphaFold 3 和 Boltz-1 为代表的蛋白质序列到结构生成扩散模型提供了强大的数据驱动先验，使研究者能够以贝叶斯后验采样的视角，通过实验似然引导生成目标构象。

然而，现有方法（如 Diffusion Posterior Sampling, DPS）普遍基于固定的结构先验，在去噪过程中直接将似然梯度加到含噪的三维原子坐标上。当目标构象处于预训练先验的低概率密度区域（如大尺度结构域重排或亚稳态）时，这种坐标空间引导必须施加极大的似然权重才能将样本拉向测量数据，这极易引发数值爆炸并破坏局部的立体化学几何，导致结构散架且对超参数极其敏感。

本文的切入角度是：与其在固定先验下对原子坐标施加脆弱的外部扰动，不如借助条件嵌入本身富含的共进化与残基间相互作用信息直接形变先验分布。核心 idea：提出 EmbedOpt 推理时引导框架，在单次扩散去噪轨迹中直接对条件序列/配对嵌入进行替代奖励单步梯度上升，将结构先验平滑对齐到实验观测，从而在保持立体化学约束的同时实现稳定鲁棒的逆问题求解。

方法详解¶

整体框架¶

EmbedOpt 面向以 AlphaFold 3 / Protenix 为代表的蛋白质条件扩散骨架。其整体输入包括目标氨基酸序列及其多序列比对（MSA）特征，由此计算出初始单残基嵌入 \(s\) 和残基对嵌入 \(z\)（统称条件嵌入 \(c=\{s, z\}\)）；以及由实验测量构成的可微奖励函数 \(R(\hat{x}_0) \propto \log p(y \mid \hat{x}_0)\)。在反向去噪采样过程中，算法从高斯噪声坐标 \(x_T\) 和初始条件嵌入 \(c_T = c\) 出发，在每个时间步交替进行：去噪预测、基于替代奖励计算关于嵌入的梯度、RMS 归一化梯度上升更新嵌入，以及依据更新后的嵌入推进扩散坐标更新。

整个流程的流向图如下所示：

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    InCoords["含噪坐标 x_t<br/>与当前嵌入 c_t"] --> Pred["去噪均值预测<br/>x^0 = x^theta(x_t, c_t, sigma_t)"]
    Pred --> RewGrad["可微实验奖励评估<br/>与嵌入梯度反向传播"]
    RewGrad --> RMSNorm["RMS 梯度归一化<br/>与单步上升更新 c_{t-1}"]
    RMSNorm --> CondUpdate["更新后先验步进<br/>推进坐标更新 x_{t-1}"]
    CondUpdate --> Out["最终松弛输出 x_0<br/>高奖励且保立体化学"]

关键设计¶

1. 替代奖励驱动的动态条件嵌入优化：避免全局反传并对齐结构先验

直接优化全局目标 \(\max_c R(x_0(x_T, c))\) 需要对整条扩散微分方程轨迹进行梯度反向传播，显存开销巨大且数值极不稳定。EmbedOpt 针对这一痛点，引入单步局部替代奖励 \(F(x_t, c_t, \sigma_t) := R(\hat{x}_\theta(x_t, c_t, \sigma_t))\)，即以当前时间步 Tweedie 去噪预测结构 \(\hat{x}_0\) 处的实验奖励作为局部代理目标。在每步采样中仅需通过单次网络前向与反向传播获取梯度 \(g_{c_t} = \nabla_{c_t} R(\hat{x}_0)\)。这一机制使得优化直接作用于残基自注意力与配对表征空间，动态改变先验能量面，使得扩散模型在生成过程中自发流向符合实验约束的构象，而非生硬拉扯原子坐标。

2. 残基与配对嵌入 RMS 梯度归一化：消除多尺度参数敏感性

在 AlphaFold 3 架构中，单残基嵌入 \(s \in \mathbb{R}^{N_{\text{res}} \times C_s}\) 与配对嵌入 \(z \in \mathbb{R}^{N_{\text{res}} \times N_{\text{res}} \times C_z}\) 的维度和梯度范数存在巨大尺度差异。若直接使用固定学习率或未经缩放的梯度更新，极易导致某一分支欠优化或梯度爆炸。EmbedOpt 针对此痛点，分别对单残基梯度和配对梯度计算均方根（RMS）归一化：

\[\bar{g}_{s_t} = \frac{\nabla_{s_t} R(\hat{x}_0)}{\sqrt{\frac{1}{d_s} \sum_{i=1}^{d_s} (\nabla_{s_t} R(\hat{x}_0)_i)^2}}, \quad \bar{g}_{z_t} = \frac{\nabla_{z_t} R(\hat{x}_0)}{\sqrt{\frac{1}{d_z} \sum_{i=1}^{d_z} (\nabla_{z_t} R(\hat{x}_0)_i)^2}}\]

随后统一使用恒定标量学习率 \(\alpha\) 进行上升：\(s_{t-1} = s_t + \alpha \bar{g}_{s_t}\)，\(z_{t-1} = z_t + \alpha \bar{g}_{z_t}\)。这一机制彻底消除了手动设计复杂学习率衰减调度的负担，使得单一基础学习率即可跨体系和跨实验模态保持稳定。

3. 信赖域理论单调性保证与几何预条件：从理论上阐明鲁棒性优势

传统 DPS 的坐标更新将奖励梯度通过雅可比转置 \(J_{x_t}^\top \nabla_{\hat{x}_0} R(\hat{x}_0)\) 映射回含噪坐标空间，强耦合于去噪器对坐标的局部灵敏度。EmbedOpt 经一阶泰勒展开可等价表达为在坐标更新中引入了基于嵌入空间几何的预条件算子 \(J_{c_t} J_{c_t}^\top \nabla_{\hat{x}_0} R(\hat{x}_0)\)，使得坐标更新方向始终落在奖励梯度张成的子空间内，而幅值和各向异性由嵌入流形几何进行正则化。理论分析（命题 1 与命题 2）进一步证明，在局部光滑与有界假设下，当学习率 \(\alpha_t \le \alpha_{\max}\) 且噪声步长 \(\Delta \sigma_t \le \Delta_{\max}\) 落在信赖域内时，单步替代奖励满足单调非递减：

\[F(x_{t-1}, c_{t-1}, \sigma_{t-1}) \ge F(x_t, c_t, \sigma_t) + \frac{\alpha_t}{2} \|g_{c_t}\|^2 - \Delta \sigma_t \left[ \frac{G_x}{\sigma_t} \|\hat{x}_\theta(x_t, c_{t-1}, \sigma_t) - x_t\| + G_\sigma \right]\]

这从数学上解释了为何 EmbedOpt 的优化轨迹呈现光滑单调上升，且对大扰动具备天然容忍度。

损失函数 / 训练策略¶

作为完全免训练（training-free）的推理时优化算法，EmbedOpt 不需要对底层扩散主干进行微调。在推理阶段，EmbedOpt 兼容 AlphaFold 3 / Protenix 的随机采样机制（在初始高噪声步注入适度随机噪声并采用步长缩放因子 \(\eta = 1.5\)）。生成的最终全原子结构随后经过 AMBER 力场下的局部能量最小化松弛（Energy Relaxation），在仅造成微弱实验奖励衰减的前提下进一步消除微小立体化学冲突。

实验关键数据¶

主实验¶

在 77 个多样化蛋白质系统的合成冷冻电镜密度拟合基准（5.0 Å 分辨率，SFC_Torch 严格模拟静电散射势）以及真实 CryoBoltz 实验基准（6 个真实实验冷冻电镜目标）上，对比未引导先验（Prior）、同架构坐标引导基线（DPS）和已发表 SOTA 基线（CryoBoltz）。

数据集 / 评测基准	评估指标	Prior (Protenix)	DPS (Protenix)	EmbedOpt (本文)	CryoBoltz (Boltz-1)
合成 Cryo-EM (77体系, 中位数)	拟合相关系数 CC mask ↑	0.64	0.80	0.89	--
合成 Cryo-EM (77体系, 中位数)	MolProbity 立体化学分数 ↓	0.70	0.71	0.77	--
真实 Cryo-EM (6目标, 中位数)	拟合相关系数 CC mask ↑	0.58	0.74	0.81	0.76
真实 Cryo-EM (6目标, 中位数)	MolProbity 立体化学分数 ↓	0.82	0.85	0.84	1.95
稀疏距离约束 (24体系, 50步)	约束满足率 (% passed) ↑	25.0%	15.0%	75.0%	--

消融实验¶

在合成冷冻电镜基准上评估能量松弛前后以及不同学习率区间对模型表现在中位数 [四分位距 IQR] 上的消融影响：

配置 / 学习率区间	未松弛 CC mask ↑	松弛后 CC mask ↑	未松弛 MolProbity ↓	松弛后 MolProbity ↓	简要说明
Prior (未引导基线)	0.68 [0.47, 0.82]	0.64 [0.44, 0.79]	1.54 [1.42, 1.69]	0.70 [0.60, 0.83]	先验分布无法自发贴合实验测量
DPS (适中 LR: 0.01~0.5)	0.84 [0.69, 0.92]	0.80 [0.66, 0.89]	1.52 [1.37, 1.76]	0.71 [0.59, 0.88]	窄有效窗口，过度增大权重易崩溃
DPS (高 LR: \(\ge 0.5\))	0.30 [0.00, 0.49]	0.01 [-0.01, 0.05]	4.48 [4.34, 4.66]	3.27 [2.72, 3.97]	坐标严重脱轨，61.3% 样本构象退化无法解析
EmbedOpt (适中 LR: 0.01~0.5)	0.95 [0.86, 0.98]	0.89 [0.79, 0.94]	1.92 [1.60, 2.43]	0.77 [0.64, 0.90]	显著优于 DPS，松弛后几何与拟合俱佳
EmbedOpt (高 LR: \(\ge 0.5\))	0.98 [0.93, 0.99]	0.89 [0.73, 0.94]	3.60 [2.77, 4.35]	1.37 [0.89, 2.26]	仅 0.2% 异常，松弛后仍保持良好物理几何

关键发现¶

大尺度结构域重排优势显著：对于先验相关性低（难度系数 \(>0.4\)）的挑战性体系（如 8H1I），DPS 容易陷于局部极小（\(CC=0.58\)），而 EmbedOpt 成功重定向结构域并将相关系数大幅提升至 \(0.93\)。
超宽超参数平台期：在学习率横跨两个数量级（\(\alpha \in [0.01, 1.0]\)）的范围内，EmbedOpt 始终保持优异的拟合度与物理合理性，彻底规避了 DPS 必须精细网格搜索的调参陷阱。
卓越的时间步效率：在稀疏距离约束测试中，将扩散步数由 200 步压缩至 50 步时，DPS 彻底崩溃，而 EmbedOpt 仍能满足 75% 以上的约束且 MolProbity 分数保持在 2.0 以下，带来 4 倍推理加速。
真实实验图表现：在 6 个真实实验冷冻电镜体系上，EmbedOpt 在 5 个体系上全面超越了 CryoBoltz，且生成的构象无 reward-fitting 畸变，立体化学指标（MolProbity 分数）显著优于 CryoBoltz（0.84 vs 1.95）。

亮点与洞察¶

将嵌入空间确立为扩散引导的新正交轴：过去生成式逆问题几乎全部集中于坐标/像素空间梯度引导，本文证明蛋白质多序列比对嵌入编码了深厚共进化流形，优化先验本身比在外层扭曲坐标更为本质且平滑。
RMS 梯度归一化的高效解耦：简单优雅的单残基与残基对 RMS 归一化策略，无需人工设计随时间步衰减的启发式权重因子，大幅增强了工程实用性。
广阔的跨任务迁移潜力：该嵌入优化范式可无缝推广至核酸-蛋白质复合物预测、配体诱导契合对接（induced-fit docking）以及单细胞/小角 X 射线散射（SAXS）等多种结构生物物理实验约束引导中。

局限与展望¶

流形覆盖范围的先决约束：若目标构象完全超出预训练模型所涵盖的构象流形（例如在无 MSA 深度输入或极端突变体系中），单纯在嵌入空间做局部梯度上升无法凭空生成未知状态。
确定性局部极小陷阱：在极少数靶点（如 8F2R）上，DPS 坐标空间的随机噪声扰动能够碰巧跳出局部势阱，而 EmbedOpt 纯梯度上升可能在局部鞍点处梯度消失而停滞。未来可结合二者优势，构建“粗粒度嵌入导向 + 细粒度坐标随机微调”的混合采样框架。
显存与计算开销：由于残基对嵌入梯度具有 \(\mathcal{O}(N_{\text{res}}^2)\) 的显存规模，并需要额外一次前向预测，单步计算时间约为 DPS 的 1.6~1.8 倍。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次提出优化条件嵌入作为蛋白质扩散模型推理时引导的新正交轴，并给出信赖域理论证明。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖稀疏距离约束、77 个合成冷冻电镜体系及 6 个真实高难度实验靶点，并细致进行了步数、学习率及 MSA 深度消融。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨自洽，算法图解与对比实验逻辑明晰，详尽分析了失败模式与适用边界。
价值: ⭐⭐⭐⭐⭐ 为结构生物学与实验生物物理中的反演重构难题提供了高鲁棒、低超参数开销的实用新范式。