Few-Step Cofolding with All-Atom Flow Maps¶
会议: NeurIPS 2026
arXiv: 2606.08375
代码: https://github.com/genesistherapeutics/decaf
领域: 计算生物
关键词: 全原子共折叠, 流映射, 模型蒸馏, 推理时搜索, 蛋白质配体结合
一句话总结¶
本文提出全原子去噪共折叠流映射框架 DeCAF，通过将预训练扩散模型蒸馏为直接在 \(\sigma\) 空间参数化的两时刻去噪流映射，并在推理阶段引入前瞻引导搜索 DeCAF-SEARCH，仅需数步推理即可生成高质量、物理有效的全原子蛋白质-配体共折叠复合物结构。
研究背景与动机¶
基于全原子欧几里得坐标的扩散生成模型（如 AlphaFold 3、Boltz-1 等）已成为预测蛋白质三维结构与蛋白质-配体结合构象（共折叠，Cofolding）的主流范式。这类方法在捕捉生物大分子全局拓扑、主链几何和相对位姿方面展现出卓越能力。然而与传统图像生成不同，生物大分子建模对局部立体化学具有极其苛刻的物理约束——键长键角微小偏差、侧链排布应变或局部原子立体碰撞都会导致结构物理有效性彻底崩溃。为满足精细局部结构的精度要求，现有扩散模型在采样时必须依赖数百步（通常 \(\mathcal{O}(200)\) 次神经网络前向评估 NFE）的细粒度数值积分，高昂的时间与计算成本严重制约了高通量药物筛选与下游结构生物学管线的实际部署。
由于生物物理约束极其严苛，研究界逐渐将计算重心转向推理阶段的搜索与对齐，例如利用立体化学势能指导的序贯蒙特卡洛（SMC）、费曼-卡茨（FK）引导以及扩散蒙特卡洛树搜索（Diffusion MCTS）。但这些测试期引导方法在每一步粒子演化时，都必须将处于高噪声状态的中间结构还原或评估其代理物理奖励。现存的核心矛盾在于：生物物理有效性奖励函数只在无噪声的干净三维结构上有意义且具有非光滑特性，在带噪状态直接计算梯度极不稳定；而多粒子分支在多步扩散轨迹上的反复展开，会令本就沉重的推理计算成本呈乘数级暴增。
本文的切入角度是：若能绕过微元数值积分，直接学习连接任意扩散噪声水平的状态跳转算子（Flow Map），并以无噪端点预测的形式显式保留刚体变换对齐能力，便能兼得超少步采样与高保真端点奖励评估。核心 idea：将预训练全原子扩散共折叠模型蒸馏为基于 \(\sigma\) 噪声空间和两时刻去噪器参数化的全原子流映射 DeCAF，在端点损失中自然保留 \(\mathrm{SE(3)}\) 刚体对齐，并借助其精准的无噪前瞻能力构建推理时引导搜索 DeCAF-SEARCH，在压缩近一个数量级推理步数的同时实现更优的物理有效性与结合精度。
方法详解¶
整体框架¶
DeCAF 整体旨在将预训练的 EDM 风格扩散模型（如 Boltz-1 或 Pearl）蒸馏为支持跳步采样的全原子流映射，并配合专为流映射设计的推理期奖励引导搜索。整个框架分为训练与推理两阶段：训练阶段，针对 EDM 非线性时间调度在时间偏导处容易产生数值溢出的痛点，直接在噪声标准差 \(\sigma\) 空间重新参数化流映射，利用雅可比-向量积（JVP）计算瞬时速度并重构预测端点 \(\hat{x}_{\text{tgt}}^{\text{DeCAF}}\)，与教师模型的无噪预测通过 Kabsch 算法求解最优 \(\mathrm{SE(3)}\) 刚体对齐后计算端点均方误差损失；推理阶段，由支持确定性与随机性插值的 \(\gamma\)-采样器执行少量步数（如 2-10 步）采样，同时基于流映射在任意噪声步均能一步前瞻干净结构 \(\hat{x}_0\) 的特性，在干净空间评估物理奖励并做梯度修正或粒子树搜索。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["初始高斯噪声构象<br/>x_sigma_N ~ N(0, sigma_N^2 I)"] --> B["1. sigma-空间流映射重参数化<br/>跳过时间偏导，在噪声标度跨度计算平均速度"]
B --> C["2. 两时刻去噪端点参数化与 SE(3) 对齐<br/>Kabsch 算法刚体对齐并优化端点损失"]
C --> D["3. 随机 gamma-采样器<br/>在确定性 ODE 步与随机重加噪间灵活平衡"]
D --> E["4. DeCAF-SEARCH 推理时前瞻引导<br/>干净构象前瞻 + 物理奖励修正 + SMC/MCTS"]
E --> F["高保真物理有效共折叠复合物构象<br/>RMSD < 2Å & 满足 PoseBusters"]关键设计¶
1. \(\sigma\)-空间流映射重参数化：根除扩散时间导数奇异性
在连续时间流映射或一致性模型蒸馏中，标准做法是在归一化时间 \(t \in [0, 1]\) 下建立状态演化方程。然而现代蛋白质扩散模型普遍采用 EDM 噪声调度 \(\sigma(t)\)，其映射高度非线性。若在时间域计算流映射一致性或欧拉平均流目标，根据链式法则 \(\frac{\partial x}{\partial t} = \frac{\partial x}{\partial \sigma} \frac{\partial \sigma}{\partial t}\)，在时间端点附近导数 \(|\frac{\partial \sigma}{\partial t}|\) 会产生极大数值，引发严重的数值不稳定甚至梯度爆炸。
针对该痛点，DeCAF 提出直接将流映射建立在噪声标准差 \(\sigma\) 空间。定义概率流常微分方程（PF-ODE）沿 \(\sigma\) 轴的瞬时速度为 \(v_\sigma(x_\sigma) \triangleq \frac{dx_\sigma}{d\sigma} = \frac{x_\sigma - D_\sigma(x_\sigma)}{\sigma}\)，其中 \(D_\sigma(x_\sigma)\) 为 EDM 去噪器。对于两个噪声等级 \(\rho > \sigma\)，定义平均速度 \(u_{\rho, \sigma}(x_\rho) = \frac{1}{\rho - \sigma} \int_\rho^\sigma v_{\bar{\sigma}}(x_{\bar{\sigma}}) d\bar{\sigma}\)，跳步更新由状态映射算子直接给出： $\(X_{\rho, \sigma}(x_\rho) = x_\rho - (\rho - \sigma) u_{\rho, \sigma}(x_\rho)\)$ 通过直接在 \(\sigma\) 坐标系下定义流映射，彻底消除了对时间变量求导所产生的放大因子，为全原子坐标尺度下的稳定蒸馏奠定了数学基础。
2. 两时刻去噪端点参数化与 \(\mathrm{SE(3)}\) 对齐：保障大分子对称性与训练低方差
如果直接采用速度回归损失（如 MeanFlow 方案直接逼近瞬时速度向量），在蛋白质全原子三维坐标上会导致训练严重发散。原因在于蛋白质空间结构具有全局旋转与平移不变性，AlphaFold 3 系列扩散模型能够稳定收敛的关键在于将预测结构与真实结构进行 \(\mathrm{SE(3)}\) 加权刚体对齐。若直接在速度场上做平移去除，会损失一个自由度，无法直接施加刚体对齐。
为解决这一冲突，DeCAF 将双时间流映射参数化为端点去噪器。利用雅可比-向量积（JVP）计算得到带阻断梯度的瞬时 \(\sigma\)-速度估算值 \(V(x_\rho, \rho, \sigma)\)： $\(V(x_\rho, \rho, \sigma) \triangleq u_{\rho, \sigma}(x_\rho) + (\rho - \sigma) \cdot \mathrm{sg}\left(\mathrm{jvp}\left(u_{\rho, \sigma}(x_\rho), (x_\rho, \rho, \sigma), (v_\rho^{\text{EDM}}, 1, 0)\right)\right)\)$ 进而构造端点预测 \(\hat{x}_{\text{tgt}}^{\text{DeCAF}} \triangleq D(x_\rho, \rho, \sigma) = x_\rho - \rho \cdot V(x_\rho, \rho, \sigma)\)。以此为基础，构建支持 Kabsch 算法刚体对齐的端点损失函数： $\(\mathcal{L} = \mathbb{E}_{x_0, x_\rho, x_\sigma} \left[ \frac{1}{\sigma^2} \min_{g \in \mathrm{SE(3)}} \left\| \hat{x}_{\text{tgt}}^{\text{DeCAF}} - \mathrm{sg}(\zeta(g) \circ \hat{x}_0^{\text{EDM}}) \right\|^2 \right]\)$ 该设计使学生模型自然继承了 EDM 预训练中通过刚体对齐降低损失方差的特性，同时对角采样（\(\rho = \sigma\)）直接对齐教师得分网络，非对角采样则学会大步幅轨迹跳跃。
3. 随机 \(\gamma\)-采样器：平衡流映射外推轨迹与探索多样性
确定性常微分方程数值求解在极端少步（如 2-5 步）时容易累积截断误差，且无法通过重复采样探索构象空间中的多峰分布。DeCAF 设计了带控制参数 \(\gamma \in [0, 1]\) 的随机采样机制。
在每个噪声递减区间 \([\sigma_n, \sigma_{n-1}]\)，首先通过流映射向前演化到收缩噪声水平 \(\tilde{\sigma}_{n-1} = \sqrt{1 - \gamma^2} \sigma_{n-1}\)，获得预测状态 \(x_{\tilde{\sigma}_{n-1}} \leftarrow X(x_{\sigma_n}, \sigma_n, \tilde{\sigma}_{n-1})\)；随后显式注入高斯随机噪声完成重加噪：\(x_{\sigma_{n-1}} \leftarrow x_{\tilde{\sigma}_{n-1}} + \gamma \sigma_{n-1} \epsilon\)（其中 \(\epsilon \sim \mathcal{N}(0, I)\)）。当 \(\gamma = 0\) 时退化为严格的确定性流映射轨迹跳转；当 \(\gamma > 0\) 时，引入受控随机扰动，既能纠正大步长跳转带来的局部几何畸变，又为后续基于粒子的引导搜索提供了必要的多样性。
4. DeCAF-SEARCH：基于干净空间前瞻的高保真物理奖励引导搜索
传统的无分类器引导或带噪得分引导难以直接作用于非光滑的生物物理立体化学惩罚（如原子位阻重叠、键长键角失真）。带噪构象 \(x_\sigma\) 上的原子普遍处于非物理的弥散状态，在此计算能量梯度的信噪比极低。
DeCAF 的核心优势在于流映射天然具备一步到干净构象的前瞻能力：\(\hat{x}_0 = X(x_\sigma, \sigma, 0) = x_\sigma - \sigma \cdot u_{\sigma, 0}(x_\sigma)\)。在此基础上设计的 DeCAF-SEARCH 具备两重关键机制： - 干净空间物理梯度上升（MC-GRAD）：在干净结构 \(\hat{x}_0\) 处直接评估物理势能奖励 \(R(\hat{x}_0)\)，计算解析梯度并沿梯度方向微调原子位置：\(\hat{x}_0 \leftarrow \hat{x}_0 + \frac{\beta}{L} \sum_{l=1}^L w_l \nabla R(\hat{x}_0^l)\)，再将优化后的结构通过重加噪投影回中间噪声层继演化。 - 粒子树与重采样引导（FK / MCTS）：维护多个候选粒子，在不同噪声水平下前瞻干净端点并评估奖励。在低算力预算下采用费曼-卡茨（FK）顺序重采样，依据奖励权重向高潜力粒子分配计算资源；在充足算力下则扩展为蒙特卡洛树搜索（MCTS），结合 UCT 上置信界限选择回溯更新，以极少步数遍历搜索最优低能结合位姿。
损失函数 / 训练策略¶
训练以开源 Boltz-1 权重为骨干初始化并冻结大部分主干模块，仅将单时刻时间条件模块替换为来自 FastGen 的双时刻（\(\rho, \sigma\)）融合条件模块作为蒸馏头。训练数据使用 2021-09-30 之前发布的 RCSB PDB 晶体结构（分辨率 \(\le 9.0\) Å）。使用 64 张 NVIDIA H200 GPU 训练 100 个 epoch，单卡 batch size 为 2，有效 batch size 为 128，扩散重采样重数为 16。优化器采用 Muon（动量 0.95，Nesterov 加速），对 1D 参数辅以 AdamW。学习率采用线性预热（1,000 步）至峰值 \(1.8 \times 10^{-3}\)。噪声等级采样采用 Karras 调度（\(\sigma_{\min}=4 \times 10^{-4}, \sigma_{\max}=160, \rho=7\)），包含 10% 对角样本（\(\rho=\sigma\)）与 90% 非对角样本。损失函数加入缩放归一化因子 \(\frac{1}{\sqrt{\mathcal{L} + 10^{-6}}}\) 稳定梯度更新。
实验关键数据¶
主实验¶
在严格测试基准 Runs N' Poses（RnP，702 个 2023-06-01 之后发布的独立测试构象）上，评估在极低 NFE 预算下 DeCAF-SEARCH 相较于 Boltz-1x 的性能表现。
| 预算 (NFE) | 方法 | PB-Valid (%) ↑ | 成功率 (RMSD<2Å & 有效) (%) ↑ | lDDT-PLI ↑ |
|---|---|---|---|---|
| \(\le 10\) NFEs (2 steps) | Boltz-1x (default) | 0.0 | 0.0 | 0.0 |
| \(\le 10\) NFEs (2 steps) | Boltz-1x (tuned) | 21.8 | 17.2 | 45.8 |
| \(\le 10\) NFEs (2 steps) | DeCAF-SEARCH (Ours) | 24.0 | 19.6 | 46.4 |
| \(\le 20\) NFEs (5 steps) | Boltz-1x (tuned) | 75.9 | 50.3 | 59.8 |
| \(\le 20\) NFEs (5 steps) | DeCAF-SEARCH (Ours) | 83.5 | 57.4 | 66.5 |
| \(\le 40\) NFEs (10 steps) | Boltz-1x (tuned) | 86.6 | 55.4 | 64.2 |
| \(\le 40\) NFEs (10 steps) | DeCAF-SEARCH (Ours) | 91.6 | 61.8 | 67.2 |
| \(\le 50\) NFEs (12 steps) | Boltz-1x (tuned) | 86.3 | 56.5 | 63.7 |
| \(\le 50\) NFEs (12 steps) | DeCAF-SEARCH (Ours) | 92.7 | 62.6 | 67.1 |
| \(\le 160\) NFEs | Boltz-1x (default, 40 steps) | 88.5 | 58.7 | 67.6 |
| \(\le 160\) NFEs | DeCAF-SEARCH (MC-GRAD, 15 steps) | 92.1 | 63.7 | 69.2 |
| 800 NFEs (Full budget) | Boltz-1x (ref., 200 steps) | 95.9 | 64.9 | 69.9 |
| 800 NFEs (Full budget) | DeCAF-SEARCH (MCTS, 10 steps) | 95.5 | 65.0 | 69.5 |
| 参考前沿模型 | AlphaFold 3 (Full budget) | 73.9 | 60.7 | 80.9 |
| 参考前沿模型 | Pearl (Full budget) | 96.1 | 72.1 | 81.9 |
在 post-2023 极具挑战的子集（\(n=196\)）上，将当前最强共折叠模型 Pearl 蒸馏为 DeCAF-Pearl，并在 40 NFE（仅 10 步）下与完整算力前沿模型对比：
| 评估阶段 | 模型 | PB-Valid (%) (best@1 / best@5) | 成功率 (%) (best@1 / best@5) | lDDT-PLI (best@1 / best@5) |
|---|---|---|---|---|
| 全预算基线 | AlphaFold 3 | 66.3 / 86.2 | 53.5 / 69.9 | 76.1 / 79.3 |
| 全预算基线 | Chai-1 | 71.2 / 90.3 | 49.5 / 63.3 | 69.9 / 72.9 |
| 全预算基线 | Boltz-1x | 93.7 / 95.9 | 62.2 / 66.3 | 71.3 / 73.2 |
| 全预算基线 | Pearl-2026.1.dev (Teacher) | 96.7 / 98.5 | 67.8 / 77.0 | 73.0 / 78.9 |
| 40 NFE 严格限制 | Pearl-2026.1.dev (10 steps) | 83.9 / 97.4 | 52.0 / 75.0 | 60.6 / 78.7 |
| 40 NFE 严格限制 | DeCAF-Boltz | 89.5 / 98.0 | 55.5 / 66.8 | 62.3 / 67.7 |
| 40 NFE 严格限制 | DeCAF-Pearl (Ours) | 95.1 / 99.5 | 64.8 / 76.0 | 71.1 / 77.2 |
消融实验¶
在 PoseBusters 基准上验证参数化方式的关键作用（统一在 30 NFE 预算下评测）：
| 参数化设计 | PB-Valid (%) ↑ | RMSD < 2Å (%) ↑ | 说明 |
|---|---|---|---|
| \(x_0\)-aligned (DeCAF 本文设计) | 90.4 | 75.6 | 端点参数化 + Kabsch 刚体对齐，完整实现立体化学稳定 |
| 纯速度回归 (velocity prediction) | 9.2 | 49.1 | 缺少刚体对齐，全原子构象严重形变，物理有效率暴跌 81.2% |
| 一致性蒸馏 (Consistency Distillation) | 0.0 | 49.5 | 传统 CD 目标无法收敛，产生完全无效的立体碰撞结构 |
采样器随机性参数 \(\gamma\) 扫描消融（PoseBusters, 10 步 DeCAF-SEARCH FK, \(P=3\)）：
| 随机度 \(\gamma\) | RMSD < 2Å (%) ↑ | PB-Valid (%) ↑ | 成功率 (%) ↑ | 说明 |
|---|---|---|---|---|
| \(\gamma = 0.3\) | 75.3 | 89.2 | 65.2 | 随机扰动偏低 |
| \(\gamma = 0.5\) | 75.6 | 90.3 | 65.9 | 最优折中：保留 ODE 轨迹方向同时保证多样性探索 |
| \(\gamma = 0.7\) | 72.8 | 90.3 | 63.8 | 随机扰动过大引发位姿发散 |
| \(\gamma = 1.0\) | 70.3 | 90.7 | 61.3 | 纯随机扩散方差过大，成功率下降 4.6% |
关键发现¶
- 低步数下刚体对齐是全原子蒸馏的生死线：消融表明，常规图像域成功的纯速度回归和标准一致性蒸馏在生物大分子场景完全失效（PB-Valid 分别仅为 9.2% 与 0.0%），唯有本文提出的带 Kabsch 算法刚体对齐的 \(x_0\)-端点去噪损失才能维持局部化学几何构象的正确性。
- 5-20 倍推理加速且不伤精度：DeCAF-SEARCH 在 PoseBusters 上仅需 30 NFE 即匹配了 Boltz-1x 在 600 NFE 时的精度与有效性帕累托前沿；DeCAF-Pearl 在 40 NFE 下达到了完整教师模型（\(\ge 200\) NFE）98.7% 的 best@5 成功率，同时在物理有效性（99.5% vs 98.5%）上实现反超。
- 算力自适应引导范式：在极低算力（\(\le 30\) NFE）下，基于粒子重采样的 FK 引导最具成本效益；在适度算力（100-250 NFE）下，蒙特卡洛平均梯度（MC-GRAD）表现最佳；而在高算力场景下，MCTS 树搜索能最彻底地挖掘高奖励结合构象。
亮点与洞察¶
- \(\sigma\)-空间速度重参数化规避了 EDM 调度的时间导数奇异点：巧妙绕过传统扩散流映射对 \(t\) 变量求导导致的数值爆炸，使全原子流映射在极宽噪声区间（\(\sigma \in [4\times 10^{-4}, 160]\)）能够平稳蒸馏收敛。
- 端点前瞻彻底解耦了中间噪声态与非光滑生物物理势能评估：流映射单步准确预测无噪几何端点的能力，完美解决了中间噪声态物理势能无法评估的困境，使干净空间的化学梯度反传与引导重采集成为了可能。
- 低 NFE 样本与高 NFE 样本构成了互补性而非简单包含：分析表明高算力 MCTS 成功的靶点并非低算力 FK 的纯超集，揭示出粗粒度轨迹探索与精细能量优化的正交性，为未来研发自适应动态计算分配机制提供了重要参考。
局限与展望¶
- 物理奖励模型本身的缺陷导致系统性误差：由于 DeCAF-SEARCH 继承了 Boltz-1x 的立体化学势能函数，某些特定构象缺陷（如 \(sp^2\) 杂化双键平面性不足）依然存在；此外，PoseBusters 评测中对金属离子配位等场景存在基于中性原子范德华半径的误报 clash 现象。
- 动态算力分配与端到端联合微调尚待探索：当前流映射蒸馏与推理引导搜索是两阶段解耦的，未来若将流映射直接与特定结合亲和力或能量打分函数联合训练（RL / Direct Alignment），有望进一步缩减搜索开销。
- 多分子体系与大尺寸复合物推广：当前验证主要集中在单蛋白-单配体共折叠任务，扩展至核酸、大尺寸多链大分子机器及共价结合抑制剂的少步生成仍需进一步验证。
相关工作与启发¶
- vs AlphaFold 3 / Boltz-1(x) / Pearl：传统全原子共折叠模型依赖 200 步左右的连续布朗运动数值反向积分，单次采样耗时极长；DeCAF 将其蒸馏为直接跨步演化的流映射，利用无噪端点前瞻将推理评估次数压缩 5-20 倍，使超大规模虚拟文库共折叠筛选成为现实。
- vs Consistency Models (CM) / MeanFlow：经典一致性模型与 MeanFlow 主要面向图像生成，采用直接速度预测或单一时间变量；DeCAF 针对生物分子的空间对称性与 EDM 尺度特性，首创 \(\sigma\)-空间重参数化与显式 \(\mathrm{SE(3)}\) 刚体对齐端点损失，填补了全原子大分子少步流映射的技术空白。
- vs Diffusion MCTS / FK-Steering：已有扩散搜索方法在带噪特征空间逼近能量梯度，受噪声干扰严重且粒子计算乘积极大；DeCAF-SEARCH 将搜索建立在流映射确切的一步端点预测上，结合 \(\gamma\)-采样器的随机重噪，实现了更高信噪比的快速推理对齐。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次将流映射引入全原子生物大分子共折叠领域，巧妙解决 \(\sigma\) 空间导数与刚体对齐矛盾。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 Runs N' Poses 与 PoseBusters 严苛评测，横向对比 AF3、Pearl 等顶级前沿模型，消融实验严谨详实。
- 写作质量: ⭐⭐⭐⭐⭐ 数学推导清晰，理论动机与生化物理场景结合紧密，实验设计层层递进。
- 价值: ⭐⭐⭐⭐⭐ 带来 5-20 倍推理加速并维持原子级精度的共折叠能力，将极大推动基于结构的计算机辅助药物设计与虚拟筛选落地。