跳转至

EvoEGF-Mol: Evolving Exponential Geodesic Flow for Structure-based Drug Design

会议: ICML 2026
arXiv: 2601.22466
代码: https://github.com/BLEACH366/EvoEGF-Mol (有)
领域: 科学计算 / 分子生成 / 结构基药物设计
关键词: 信息几何, 指数测地线, Fisher-Rao 度量, 流匹配, SBDD

一句话总结

EvoEGF-Mol 把 SBDD 的连续坐标与离散原子/键类型放到同一个指数族自然参数空间里,用动态收紧的目标分布替代奇异的 Dirac 端点,沿着 Fisher-Rao 几何下的指数测地线同步演化,在 CrossDock 上把 PoseBusters 通过率推到 93.4%,逼近参考分子水平。

研究背景与动机

领域现状:结构基药物设计 (SBDD) 想根据蛋白口袋 \(P\) 生成具有结合活性的小分子配体 \(M=(\mathbf{x}_M,\mathbf{v}_M,\mathbf{b}_M)\),包括三维原子坐标、原子类型与键类型三类异构变量。主流方法已从早期的自回归 (AR、Pocket2Mol、PocketFlow) 转向扩散与流匹配范式 (TargetDiff、DecompDiff、FLOWR、DynamicFlow、ECloudGen),并出现统一概率框架 (MolCRAFT、MolPilot)。

现有痛点:几乎所有方法都为连续坐标和离散类别各自单独设计概率路径——前者走欧式空间的高斯加噪,后者走类别概率单纯形上的离散调度。这种"分而治之"会导致模态时序错配(modality mismatch):几何坐标已经被驱赶到近似收敛,而原子身份还含糊不清,破坏了药物分子里几何—化学的强耦合。

核心矛盾:异构变量没有统一的"距离"。坐标的高斯方差与类别的 Dirichlet 浓度衡量的是不同空间里的不确定性,强行用两套损失加权拼起来,需要靠人工调权,本质上违背了分布间的内在几何。

本文目标:(1) 给坐标、原子类型、键类型一个统一概率对象的描述;(2) 在这个对象上构造几何上合理的概率路径;(3) 避免 Dirac 端点带来的瞬时坍缩,保证训练信号在整段 \(t\in[0,1]\) 上有效。

切入角度:信息几何告诉我们,指数族分布在 Fisher-Rao 度量与指数连接下的指数测地线 (e-geodesic) 恰好对应自然参数 \(\bm{\eta}\) 上的线性插值。如果把高斯坐标和 Dirichlet 类别都视作指数族的"积",它们就共享同一种线性时间表,天然消除模态时序错配。

核心 idea:把分子表示成"高斯×Dirichlet×Dirichlet"的复合指数族分布,沿着 e-测地线演化,并把固定 Dirac 端点替换成"随时间逐步收紧"的动态端点,从而既保留 Fisher-Rao 的几何一致性,又避开边界奇点导致的方差/支撑瞬时坍缩。

方法详解

整体框架

EvoEGF-Mol 要解决的是"连续坐标和离散原子/键类型用两套概率路径、节奏对不齐"的老问题。它的做法是把分子三元组 \(M=(\mathbf{x}_M,\mathbf{v}_M,\mathbf{b}_M)\) 整体看成一个"高斯×Dirichlet×Dirichlet"的积指数族分布,把所有变量的状态压进同一个自然参数向量 \(\bm{\eta}\),再让这个向量沿 Fisher-Rao 几何下的指数测地线从先验 \(\bm{\eta}_0\) 演化到目标。网络只需"看一眼当前噪声分子就预测终点参数",配合动态收紧的目标分布避开端点奇异,最终在蛋白口袋 \(P\) 条件下输出配体。

关键设计

1. 统一自然参数空间下的同步 e-测地线:让坐标与类别按同一节拍收紧

痛点在于以往把坐标的欧式高斯加噪和类别单纯形上的离散调度分开设计,几何已经收敛了原子身份还含糊,破坏了几何—化学耦合。EvoEGF 的关键观察是:对任意指数族 \(p(\mathbf{x}|\bm{\eta})=h(\mathbf{x})\exp(\langle\bm{\eta},\mathbf{T}(\mathbf{x})\rangle-A(\bm{\eta}))\),Fisher-Rao 度量与指数连接下的 e-测地线恰好等价于自然参数上的线性插值 \(\bm{\eta}_t=(1-t)\bm{\eta}_0+t\bm{\eta}_1\)。于是把各向同性高斯的两个自然参数 \(\sigma_t^{-2}\bm{\mu}_t\)\(-\tfrac{1}{2}\sigma_t^{-2}\) 和 Dirichlet 的 \(\bm{\eta}=\bm{\alpha}-\mathbf{1}\) 全部线性插值,异构变量就被同一个时间 \(t\) 牵着同步收紧,从根上消掉模态时序错配。这种几何而非人工配权的好处在损失展开里很直观:一阶 KL 退化成 \(D_{\mathrm{KL}}\approx \tfrac{1}{2}\sum_{\mathbf{c}}(\bm{\xi}_t^\mathbf{c})^\top \mathbf{G}^\mathbf{c}(\bm{\eta}_t^\mathbf{c})\bm{\xi}_t^\mathbf{c}\),每个分量的监督权重直接由 Fisher 信息矩阵 \(\mathbf{G}^\mathbf{c}\) 给出,不再需要手调坐标 MSE 与分类 CE 的相对比例。

2. 动态收紧端点替换 Dirac 目标:把训练信号摊到整段时间轴

直接让 e-测地线瞄准 Dirac 终点会出大问题:自然参数在端点趋于无穷,坐标方差 \(\sigma_t^2\)\(t\to1\) 时迅速塌到 0,类别单纯形支撑瞬间坍缩,结果几乎所有有效训练信号都被挤在末段(作者在 §3.2 与 Fig.2 中正是这样分析的)。解法是把固定终点 \(\bm{\eta}_1\) 换成时间相关的 \(\tilde{\bm{\eta}}_1(t)\),用平滑超参 \(\lambda\) 控制收紧速度——坐标侧设 \(\tilde{\sigma}_1(t)=\lambda(1-t)\),类别侧设 \(\tilde{\bm{\alpha}}_1(t)=(1-\lambda(1-t))\mathbf{e}_k+\lambda(1-t)\tfrac{1}{K}\mathbf{1}_K\)。这样只要 \(t<1\),端点就始终落在开凸的自然参数域 \(\Omega\) 内部、参数有界、路径不发散,收紧过程被匀速分摊到 \(t\in[0,1]\) 全段,训练窗口被拉满。

3. 渐进参数细化训练 + Fisher 校准 KL 损失:在参数空间里把目标退化成回归

联合离散—连续的样本空间速度场不唯一、直接学很难训稳,所以 EvoEGF 仿 BFN/PIF,让网络 \(\bm{\Phi}(M_t,t,P)\) 接收当前噪声样本直接预测终点参数 \(\hat{\bm{\eta}}_1\),而不是显式估计速度场。训练时采样 \(t\sim\mathcal{U}(0,1)\) 算出 \(\bm{\eta}_t\) 与噪声样本 \(M_t\),网络预测 \(\hat{\bm{\eta}}_1\) 后重建 \(\hat{\bm{\eta}}_{t+\Delta t}\),用它与真实演化的一阶 KL 差 \(\bm{\xi}_t\) 做监督。坐标分量于是化简为加权 MSE \(\mathcal{L}_\mathbf{x}=\mathbb{E}[\tfrac{t^2\sigma_t^2}{2\tilde{\sigma}_1^4(t)}\|\mathbf{x}^*-\hat{\mathbf{x}}\|^2]\),类别分量退化为含多元 Beta 项与 digamma 差 \(\Delta\psi_k\) 的 Dirichlet KL。由于积指数族的 Fisher 矩阵分块对角,权重 \(\mathbf{G}^\mathbf{c}(\bm{\eta}_t^\mathbf{c})\) 天然给出、各分量解耦却又彼此协调,免去跨模态手工配比——论文还顺带证明了 SLDM 只是 EGF 在静态端点 + 正则化下的一个特例。

损失函数 / 训练策略

总损失是坐标 \(\mathcal{L}_\mathbf{x}\)、原子类型 \(\mathcal{L}_\mathbf{v}\)、键类型 \(\mathcal{L}_\mathbf{b}\) 三块 Fisher 加权 KL 在 \(t\sim\mathcal{U}(0,1)\) 下的期望之和;积指数族下 Fisher 矩阵分块对角,三个分量天然解耦但权重彼此协调。整体延续 BFN/PIF 式"预测终点参数 + 一步细化"范式,时间步统一在 \([0,1]\) 均匀采样,端点收紧速度由超参 \(\lambda\) 单独控制。采样时从 \(M_0\sim p(\cdot|\bm{\eta}_0)\) 起步,每步预测 \(\hat{\bm{\eta}}_1\)、组合出 \(\hat{\bm{\eta}}_t\) 再重采样作为下一步输入,直到 \(t=1\) 输出分子。

实验关键数据

主实验

在 CrossDock 上对比统一框架与 SOTA 扩散/自回归基线,PoseBusters 通过率与多种 Vina 打分、应变能、连接率、QED、SA、Clash Ratio 共同评估几何与化学合理性。

数据集 指标 本文 EvoEGF-Mol 之前最佳 (MolCRAFT) 提升
CrossDock PB-Valid (↑) 93.4% 84.6% +8.8 pp
CrossDock Connected (↑) 98.6% 96.7% +1.9 pp
CrossDock Strain (Med., ↓) 25.96 195 -86.7%
CrossDock Vina Min (Avg., ↓) -6.98 -7.21 略低 0.23
CrossDock SA (↑) 0.75 0.67 +0.08
CrossDock Clash Ratio (↓) 0.24 0.26 -0.02

PoseBusters 通过率 93.4% 已经接近参考分子集本身的 95.0%,应变能中位数从 MolCRAFT 的 195 一路压到 25.96,说明几何上的"物理合理性"提升远比单纯 Vina 分数更显著。

消融与拓展实验

评测套件 / 配置 关键指标 说明
CrossDock vs Dirac 端点 EGF (Fig.2) 训练窗口宽度 静态端点导致方差/支撑瞬时坍缩,训练信号集中末段;动态端点平滑展开整段时间轴
MolGenBench (In) Pass Rate / Hit Recovery (↑) EvoEGF-Mol 在 In/In(RM.)/Not 三种蛋白划分下均取得 top-2 的命中率与片段恢复率
与 SLDM 关系 形式化对比 (Appx. E) 证明 SLDM 是 EGF 在正则化静态端点下的特例,EvoEGF 是更广义的动态解
Fisher 校准 vs 人工加权 多模态平衡 KL 二次展开给出每个分量天然权重 \(\mathbf{G}^\mathbf{c}\),免去交叉模态调权

关键发现

  • 几何先验 (e-测地线) 比手工设计的概率路径更有助于"物理合理"的分子:应变能与 clash 都大幅下降,说明模型生成的构象更贴近真实化学。
  • 端点动态化是关键:相同 e-测地线骨架下,把 Dirac 改成 \(\lambda\)-收紧目标即可显著缓解方差坍缩,无需修改网络架构。
  • 在 MolGenBench 的真实任务上,scaffold recovery 的 hit rate 提升说明该框架不只是在 CrossDock 这种受限基准上过拟合,对实际药物候选检索也有价值。

亮点与洞察

  • 信息几何视角是真正"统一"连续—离散生成的有力工具:自然参数空间的线性插值同时把高斯方差与 Dirichlet 浓度按同一节拍拉到目标,远比"两套损失加权"优雅。
  • 动态端点是一招通用的"奇异性疗法":任何想瞄准 Dirac 的生成式流都会遇到端点发散,把端点改成时间相关的收紧分布几乎可以无痛迁移到其它指数族生成任务 (语言、点云、属性图)。
  • 渐进参数细化范式让训练目标退化成熟悉的回归 + Dirichlet KL,与 BFN/PIF 体系无缝衔接,工程实现成本低。
  • 论文对 SLDM 的"特例化"分析很值得借鉴:把已有方法纳入更广的几何框架可以同时澄清各自的边界条件,未来研究者可以照葫芦画瓢分析 MolPilot/FLOWR 是不是 EGF 族另一些特例。

局限与展望

  • 框架目前固定使用各向同性高斯 + Dirichlet 这一具体指数族选择,更复杂的分布族 (如混合高斯、von Mises) 是否同样稳定还未验证。
  • 收紧速度 \(\lambda\) 是全局常数,分子的不同部位 (骨架 vs 取代基) 可能需要不同收紧速率,未来可探索 \(\lambda\) 随位点自适应。
  • 实验只覆盖两个口袋数据集 (CrossDock、MolGenBench),对脂质、多肽、共价结合等更难的药物场景缺乏验证。
  • 推理仍需要多步迭代,相对一次性的回归模型推理成本较高,能否结合一致性模型 (RCM) 类思想做少步采样值得探索。

相关工作与启发

  • vs MolCRAFT / MolPilot (BFN/VOS 系):他们仍是欧式 + 类别两套调度,需要 VOS 这种额外的噪声时间表对齐;EvoEGF 用统一指数族直接消除时序错配,且参数空间训练范式延续 BFN 思想,工程开销相近。
  • vs FLOWR / DynamicFlow (流匹配系):他们把连续 OT 流和离散 categorical FM 拼接;EvoEGF 给出几何上更内在的"积指数族 + e-测地线"路径,并通过 Fisher 信息自动平衡多模态。
  • vs Fisher-Flow / SFM / E-Geodesic FM (信息几何系):先前工作把信息几何用于纯离散或单纯形数据,本文首次把它推广到"连续—离散混合"的分子结构,并通过动态端点解决了原生 e-测地线对 Dirac 端点的奇异性。
  • vs SLDM:作者证明 SLDM 是 EGF 在静态正则化端点下的特例,提供了将多种近期"端到端直线扩散"方法纳入统一几何框架的视角。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 第一次把 SBDD 的连续—离散变量真正放到同一个 Fisher-Rao 流形上演化,并系统化"动态端点"解决奇异性。
  • 实验充分度: ⭐⭐⭐⭐ CrossDock 与 MolGenBench 双重验证,但缺少与更多 2025 年新方法 (如 ECloudGen、DynamicFlow) 的同条件对比。
  • 写作质量: ⭐⭐⭐⭐ 信息几何推导清晰,把每个分布族的自然参数写得很显式,附录补足边界奇点与 SLDM 关系。
  • 价值: ⭐⭐⭐⭐⭐ 几何统一框架与动态端点都是可以扩展到其它生成式问题的通用工具,对药物发现而言 PoseBusters 大幅提升非常实用。