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 的几何一致性,又避开边界奇点导致的方差/支撑瞬时坍缩。
方法详解¶
整体框架¶
输入是蛋白口袋 \(P=\{(\mathbf{x}_P^{(i)},\mathbf{v}_P^{(i)})\}\) 和先验 \(\bm{\eta}_0\),输出是配体三元组 \(M=(\mathbf{x}_M,\mathbf{v}_M,\mathbf{b}_M)\)。整条 pipeline 完全在自然参数空间里运转:
- 统一表示:把 \(M\) 看作积指数族 \(p(\mathbf{x}|\bm{\eta}^{\mathbf{x}})\,p(\mathbf{v}|\bm{\eta}^{\mathbf{v}})\,p(\mathbf{b}|\bm{\eta}^{\mathbf{b}})\),其中坐标用各向同性高斯、原子/键类别用 Dirichlet,全部都是指数族成员,自然参数可以拼接成一个长向量 \(\bm{\eta}=(\bm{\eta}^{\mathbf{x}},\bm{\eta}^{\mathbf{v}},\bm{\eta}^{\mathbf{b}})\)。
- 动态端点测地线:用 \(\bm{\eta}_t=(1-t)\bm{\eta}_0+t\tilde{\bm{\eta}}_1(t)\) 构造时间演化的目标,沿 e-测地线走,路径始终留在开凸的自然参数域 \(\Omega\) 内部。
- 渐进参数细化网络:仿 BFN 与 PIF,神经网络 \(\bm{\Phi}(M_t,t,P)\) 接收当前噪声样本,直接预测终点参数 \(\hat{\bm{\eta}}_1\);训练用一阶 KL 转化为 Fisher-Rao 范数下的二次损失。
- 采样:从 \(M_0\sim p(\cdot|\bm{\eta}_0)\) 起步,每步预测 \(\hat{\bm{\eta}}_1\)、组合得到 \(\hat{\bm{\eta}}_t\),再从中重采样作为下一步输入,直到 \(t=1\) 输出分子。
关键设计¶
-
统一自然参数空间下的同步 e-测地线:
- 功能:让坐标、原子类型、键类型在同一时间表下"同步收紧",消除连续—离散模态的时序错配。
- 核心思路:对指数族 \(p(\mathbf{x}|\bm{\eta})=h(\mathbf{x})\exp(\langle\bm{\eta},\mathbf{T}(\mathbf{x})\rangle-A(\bm{\eta}))\),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\) 牵着走。
- 设计动机:先前方法把欧氏 MSE 和分类 CE 用人工权重拼起来,本质上忽略了"分布之间的距离"。Fisher-Rao 度量自动按每个分量的内在不确定性来加权监督,原文损失展开 \(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 校准的合法性。
-
动态收紧端点替换 Dirac 目标:
- 功能:避免几何奇点导致的方差瞬间归零或单纯形支撑瞬时塌缩,把可用的训练时间窗口拉满。
- 核心思路:把固定终点 \(\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\) 时端点始终在流形内部,自然参数有界,路径不再发散。
- 设计动机:作者在 §3.2 与 Fig.2 中分析,e-测地线一旦瞄准 Dirac,自然参数在端点趋于无穷,方差 \(\sigma_t^2\) 在 \(t\) 接近 1 时迅速塌到 0,几乎所有训练信号被压缩在末段;动态端点让收紧过程匀速分摊到全段时间。
-
渐进参数细化训练 + Fisher 校准 KL 损失:
- 功能:在参数空间里直接学一个"看一眼噪声态就给出终点参数"的网络,回避在样本空间显式估计速度场的难题,同时让多模态损失自动平衡。
- 核心思路:训练时采样 \(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]\);类别分量退化为 Dirichlet 的 KL,包含多元 Beta 项与 digamma 差 \(\Delta\psi_k\)。
- 设计动机:联合离散—连续的样本空间速度场不唯一,直接学很难训稳;在参数空间学终点参数让训练目标退化成熟悉的回归形式,并且权重 \(\mathbf{G}^\mathbf{c}(\bm{\eta}_t^\mathbf{c})\) 由 Fisher 信息天然给出,免去手工配比。论文还指出 SLDM 是 EGF 在静态端点 + 正则化下的特例。
损失函数 / 训练策略¶
总损失是三块 Fisher 加权 KL 的和:坐标 \(\mathcal{L}_\mathbf{x}\)、原子类型 \(\mathcal{L}_\mathbf{v}\)、键类型 \(\mathcal{L}_\mathbf{b}\),三者在 \(t\sim\mathcal{U}(0,1)\) 下求期望并相加,由于积指数族下 Fisher 矩阵分块对角,各分量天然解耦但权重彼此协调。训练采用 BFN/PIF 式的"预测终点参数 + 一步细化"范式,时间步统一在 \([0,1]\) 上均匀采样,端点收紧速度由超参 \(\lambda\) 单独控制。
实验关键数据¶
主实验¶
在 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 大幅提升非常实用。