Scalable Inference-Time Annealing with Surrogate Likelihood Estimators

TL;DR SITA 通过基于能量的替代似然估计器实现可扩展的推理时退火，无需计算代价高昂的散度项，在分子玻尔兹曼采样中取得了最优性能。

问题

问题背景

在计算化学与生物物理中，高效采样玻尔兹曼分布（Boltzmann distribution） 是理解分子平衡态性质（如自由能、构象系综）的核心任务，直接影响药物设计、材料筛选等下游应用。传统分子动力学（MD）模拟需在飞秒级时间步长下长时间演化，计算成本极高，常成为瓶颈。

现有方法局限

近年，扩散模型（diffusion models） 被引入分子生成，结合温度阶梯退火与重要性采样的策略展现出潜力：通过逐步降低温度并利用推理时生成的样本重新训练模型，以避免直接 MD 采样。然而，现有方案需计算得分场散度（divergence of score field, ∇·s_θ）来估计重要性权重，该操作涉及二阶导数，在高维构象空间（如数百个原子）中内存与时间开销随自由度指数增长，导致方法无法扩展至丙氨酸三肽（Alanine Tripeptide） 这类中等尺寸分子。即使采用随机近似（如 Hutchinson 迹估计器）也会引入方差与额外计算，失去实用价值。

为什么这个问题难且重要

散度计算的本质困难在于：

• 维度灾难：得分函数定义在 3N 维相空间，直接求散度需遍历所有维度，对 GPU 显存极不友好。
• 精度-效率权衡：避开散度的 likelihood-free 方法（如拒绝采样）往往在高维拒绝率过高，而提高训练集质量又依赖准确的权重估计。 业界（包括制药、蛋白设计）迫切需要一种线性或近线性复杂度的替代方案，能在可接受的计算预算下生成无偏玻尔兹曼样本。SITA 的突破在于用基于能量的模型（EBM） 提供代理似然，将重要性权重估计转化为一次模型前向传播，完全规避散度瓶颈，同时保持流模型采样的高效性。

行业类比

这一思路类似在大语言模型 RLHF 中用奖励模型替代人类在线反馈以降低标注成本：用轻量代理模型近似昂贵的目标量，使得原本难以扩展的推理时优化变得可行。

核心洞察

• 用流模型和能量模型替代得分场散度计算, 消除了推理时退火的可扩展性瓶颈。 先前扩散模型沿温度阶梯微调需要计算得分场散度来估计重要性权重，该计算复杂度随系统尺寸增长而急剧升高，限制了其在较大分子系统的应用。SITA 通过使用流匹配模型结合**能量模型**产生的**替代似然估计器**(如 BoltzNCE)，无需显式计算散度即可进行重要性采样，从而实现了低温度样本的高效生成，为大规模生物分子模拟提供了实用路径。
• 在温度阶梯上重训练流模型而非从头采样，利用能量模型快速提供似然值，显著降低了生成高保真构象的总体成本。 传统模拟方法依赖长时间物理模拟，计算成本高昂。SITA 的策略是预训练一个流模型学习较高温度的分布，然后通过**推理时退火**逐步降低温度，每个温度阶段仅需少量重训练，并使用**能量模型**估算的替代似然进行重要性加权来调整分布。这相当于在已有知识基础上逐步精化，避免了每次退火都从头开始采样，大幅提升数据效率，同时保持了样本质量。

方法

方法总览：推理时温度退火与代理似然引导的流模型

SITA 解决的核心问题是：在不依赖昂贵 分数散度计算 的前提下，沿温度阶梯生成分子 Boltzmann 分布 样本。方法围绕 预训练流模型、能量模型（EBM）代理似然 和 重要性采样重训练 三个模块展开。

输入与初始化

• 高温度流模型：先在高温 (T_0) 下用 Flow Matching（或 Stochastic Interpolants 框架）训练一个连续流模型，该模型能高效探索构象空间，但样本对应的是平坦分布。
• 能量模型（BoltzNCE）：训练一个独立的 EBM，作为快速近似势能函数，提供 代理似然 (p_{\text{EBM}}(x))。它比直接计算分子力场快得多，且可微分。

推理时退火的关键模块

1. 温度操控的概率流 ODE
   流模型本身具备温度可操控性：通过调整 ODE 中噪声调度或条件信号，可直接从同一个模型生成指定温度 (T) 的样本，而无需重新训练。这为逐步降温提供基础。

2. 代理似然驱动的重要性采样
   在目标温度 (T < T_0) 下，从当前流模型采样一批候选构象 (x)，用 EBM 计算代理似然 (p_{\text{EBM}}(x))，并估计重要性权重 (w(x) \propto \exp(-(E(x) - E_0))/T)（其中 (E) 由 EBM 近似）。这些权重衡量候选构象在目标 Boltzmann 分布下的相对概率。

3. 流模型微调（重训练）
   以重要性权重为采样概率，对候选构象重新采样，再以这些加权样本作为目标分布，用 Flow Matching 损失 微调流模型。此过程将模型从当前温度分布逐渐推向更低温度的 Boltzmann 分布。

输出与迭代

• 微调后的流模型成为新的起点，重复步骤 1‑3，逐步降温至目标温度。最终输出一个在目标低温下准确生成 Boltzmann 分布的流模型。
• 整个过程 无需计算扩散模型的分数散度（如 ℓ₂ 散度或 Stein 散度），因为流模型直接用概率流 ODE 采样，重要性权重仅依赖代理似然的点评估。

与同类方法的差异

与 基于扩散模型 的退火方法（如 Diffusion Tempering）相比，SITA 用 流模型 + EBM 代理 完全规避了分数散度计算，大幅降低计算开销，从而可扩展至更大分子体系。

实验

实验设计

本文在两个经典分子体系 Alanine Dipeptide (ADP) 和 Alanine Tripeptide (ATP) 上验证 SITA 的有效性。实验对比了基于连续归一化流（CNF）的退火框架与扩散模型基线，着重评估温度阶梯退火与替代似然引导的采样质量。流程包含：

1. 在高温数据上预训练流模型；
2. 沿温度阶梯迭代微调，每一步用 EBM (Graphormer) 提供替代似然权重进行重要性采样；
3. 可选的 独立 Metropolis-Hastings (M-H) 精炼步骤，利用同一替代似然进一步提升样本保真度。 评估维度涵盖有效样本量（ESS）、能量误差、扭转角分布拟合（Ramachandran KL 散度）以及时间滞后独立成分分析（TICA）下的动力学一致性。

关键发现

• 无需散度项即可实现退火：替代似然估计器成功替代了以往扩散采样中必需的 score field 散度计算，消除了面向更大分子体系的计算瓶颈。
• SOTA 采样质量：在 ADP 和 ATP 上，SITA 生成的构象系综在能量和几何分布上与长时分子动力学（MD）参考高度吻合，TICA 投影的 Wasserstein 距离显著优于基线。
• 精炼模块即插即用：附加的 M-H 步骤能进一步降低残余偏差，且几乎不增加额外计算成本，表明替代似然可作为高质量提议分布的直接判别器。

基线对比深度解读

相较于 扩散模型推理时退火（需显式计算时间反向过程的散度），SITA 利用了流模型的可逆性与显式温度操控特性，将退火过程转化为概率流 ODE 的温度参数化。这一设计不仅避开了高维散度估计的不稳定性，还使得重要性权重可通过一次前向 EBM 评估快速获得，使得退火迭代的实际计算量主要由数据生成（采样）决定，而非权重的复杂微分。

与 正常化流 的同类工作相比，SITA 引入的替代似然引导本质是 bootstrap 策略：用粗粒度的 EBM 近似真实玻尔兹曼分布，在生成-重加权循环中既校正了流模型的系统误差，又免除了对昂贵物理能量函数的反复调用。这种 替代似然 + 温度操控 的组合，为面向通用分子的可扩展采样提供了清晰范式，其单次推理成本可控，适合作为下游任务（如自由能计算）的前端。

行业影响

SITA 将基于流的生成模型与能量基模型(EBM) 结合，在推理时进行温度退火，无需昂贵的散度计算即可高效采样分子构象。该技术直接服务于计算化学与生物物理，但其可扩展的采样范式对工业级 AI 应用有直接启示。

落地场景

• AI 药物研发平台: 在虚拟筛选、自由能微扰(FEP)和构象集成中，SITA 可替代传统分子动力学(MD)采样的部分计算，加速先导化合物优化。
• 分子模拟云服务: 云计算厂商可将 SITA 封装为微服务，提供按需玻尔兹曼采样，降低中小型药企的算力门槛。
• 材料科学辅助设计: 用于聚合物或催化剂的构象空间探索，缩短新材料开发周期。

商业价值

• 降本: 直接减少 MD 模拟的 GPU/CPU 机时消耗。以 FEP 计算为例，传统方法需数万节点小时，SITA 可在单卡推理秒级生成多样构象，成本降低一个量级。
• 增收: 加速药物研发管线中的构象分析环节，使更多候选分子进入评价，提高临床前成功率；对 CRO/CDMO 而言，可增加通量。
• 体验提升: 在交互式分子设计工具（如 PyMOL、Schrödinger Maestro 插件）中集成 SITA，让药剂师实时看到低能构象，优化设计体验。

与现有产品 / 工作流的接口

SITA 以流匹配 (flow matching) 框架训练，与 Equivariant GVP 编码器无缝衔接，输出为连续概率路径样本。可集成方式包括：

1. Python 库 / SDK: 提供的代码库（GitHub） 可直接加载预训练权重，通过 sample() API 生成构象。
2. 容器化服务: 打包为 Docker 镜像，通过 REST API 或 gRPC 供下游调用。
3. 分子模拟软件对接: 生成构象写入 PDB 文件，供 Amber、OpenMM 进一步精修或分析；亦可将 SITA 作为 Metropolis-Hastings 提议分布，与现有采样器协同。

具体落地用例

用例 1: 增强 FEP 工作流
在 Schrödinger FEP+ 或开源 FEP 工具中，用 SITA 生成配体扰动前后的构象集合，替代短时 MD 退火。由于 SITA 能精确采样目标温度分布，可减少自由能估计的方差，使客户能用更少的模拟时间获得可靠的结合自由能预测。

用例 2: AI 驱动 SaaS 构象采样器
某云端药物设计平台集成 SITA 为“一键退火”功能：用户上传小分子 SMILES 或蛋白质 PDB，选择温度 ladder（如 500K→300K），平台返回低能构象系综及 Ramachandran 分布，用于下游药效团建模或对接。该功能可对标现有收费的构象搜索工具（如 ConfGen），但速度更快、延展更高。

局限

• **系统规模验证不足**：论文仅在**丙氨酸二肽(ADP)**和**丙氨酸三肽(ATP)**这类小体系上展示了效果，未在更大蛋白质或实际药物分子上实验。尽管声称“scalable”，但**EBM代理似然**在高维构象空间的误差可能放大，且**流模型**的容量能否扩展到数千个原子仍存疑。对实际工程应用而言，小分子 benchmark 的 SOTA 不一定能迁移到有产业价值的体系上。
• **依赖 EBM 的准确性**：SITA 使用**能量基模型(EBM)**作为代理似然估计器来引导重要性采样，而 EBM 本身训练困难，可能模式崩塌或概率密度估计不准确。如果 EBM 未能覆盖重要构象区域，退火过程会遗漏关键模态，导致最终采样偏差。论文未讨论 EBM 质量对结果敏感度的影响，这对生产化部署是一个隐患。
• **对比基线范围有限**：实验主要与基于扩散模型的 **FAB** 和 **Timewarp** 等少数生成式方法比较，缺少与传统增强采样方法（如元动力学、副本交换）或更近期的**温度阶梯强化学习**方法的全面对比。同时未报告在更高温度下的效能，难以判断方法是否在整个退火路径上都优于现有方案。