基于代理似然估计器的可扩展推理时间退火方法 Scalable Inference-Time Annealing with Surrogate Likelihood Estimators
SITA使用EBM代理似然实现高效分子玻尔兹曼分布采样
前置知识
玻尔兹曼分布
玻尔兹曼分布 pi(x) = Z^{-1} exp(-E(x)/k_BT) 描述了分子系统在热平衡状态下各构象的概率分布，其中 E(x) 是能量函数，T 是温度，Z 是配分函数。采样玻尔兹曼分布是计算化学和生物物理的核心任务，可以得到自由能、结合亲和力等热力学可观测量。由于分子能量函数高度非凸且粗糙，传统采样方法容易陷入局部极小值。
本文的目标就是高效采样玻尔兹曼分布，这是理解分子系统热力学性质的基础。
流匹配
流匹配通过学习速度场 v_t(x) 将简单基分布（如高斯）变换到目标分布。给定配对样本 (x_0 ~ rho_0, x_1 ~ rho_1)，插值定义为 I_t = alpha_t x_0 + beta_t x_1，速度场通过回归损失 L_v(hat{v}) = E[||hat{v}_t(I_t) - dot{I}_t||^2] 学习。与扩散模型相比，流匹配避免了反向过程的复杂评分网络，采样通过积分生成ODE完成。
SITA使用流匹配作为生成模型的主干，因为它的连续时间特性便于实现温度退火。
退火重要性采样
退火重要性采样通过一系列中间分布从简单分布逐渐过渡到目标分布。在温度退火中，定义温度阶梯 {T_k} 从高温到低温，每一步使用重要性权重重采样来逼近下一个温度的分布。关键挑战是计算重要性权重需要评估两个分布的似然，对于流匹配模型这涉及计算速度场的散度，成本随维度增长而变得不可行。
SITA的核心创新就是用代理似然估计器来避免昂贵的散度计算，使退火适用于大分子系统。
能量基模型
能量基模型通过能量函数 U(x) 定义概率分布 p(x) proportional to exp(-U(x))，不需要归一化常数。训练EBM是历史难题，但BoltzNCE将评分匹配与噪声对比估计结合，使其变得实用。训练分两阶段：先从流生成样本，然后用这些样本训练EBM学习能量函数，损失包含评分匹配项 L_SM 和InfoNCE项 L_InfoNCE。
SITA使用EBM作为流生成样本的代理似然估计器，使重要性重采样变得可行。
研究动机
现有分子玻尔兹曼采样方法面临严重挑战。传统MCMC和分子动力学模拟需要飞秒级时间步长，产生高度相关的样本，且对分子系统的微小修改需要重新运行模拟，不适合高通量分析。扩散模型如PITA（Akhound-Sadegh et al., 2025）在推理时间退火中需要计算速度场散度来估计重要性权重，基于Feynman-Kac公式的自归一化重要性采样要求在完整反向积分路径上评估散度，计算开销随着自由度数量的增长变得不可接受，严重限制了在大分子系统上的应用。
本文的目标是本文的目标是开发一种可扩展的推理时间退火方法SITA，用于训练连续流模型来恢复目标玻尔兹曼分布，同时避免计算向量场散度的高昂成本，在多个基准分子系统上达到最先进性能。具体而言，SITA将流匹配模型与能量基模型代理似然估计器结合，通过温度退火从高温分布逐步转移到室温平衡目标分布，使用代理似然实现重要性加权重采样，避免基于雅可比的重新加权的 prohibitively expensive 成本。
与已有工作不同的是，本文的独特切入点在于首次将推理时间退火应用于连续流匹配模型，并通过一类新的代理似然估计器BoltzNCE绕过基于路径积分的SNIS估计器。与PITA需要计算速度场散度不同，SITA利用EBM作为流生成样本的代理似然，使重要性权重计算变得可扩展。虽然引入了代理估计器的偏差，但作者通过实验证明，尽管存在这种内在偏差，EBM导出的重要性权重在基准分子系统上仍能提供优越性能。
核心方法
SITA的整体思路是结合流匹配生成模型和能量基模型代理似然估计器，通过温度退火阶梯逐步训练流模型逼近目标玻尔兹曼分布。直觉上，高温玻尔兹曼分布更平坦、更容易采样，从高温分布开始，每一步使用当前流模型生成样本，用EBM估计这些样本的似然，通过重要性重采样获得低温分布的样本，然后用这些样本微调流模型，重复这个过程直到达到目标温度。技术路线上，SITA维护两个模型：流模型 hat{v}_theta 和EBM hat{U}_phi，在整个退火过程中复用，只在每步重新初始化优化器进行微调。
SITA的核心创新点有两个：一是首次实现了连续流模型的推理时间退火，通过观察流模型可以隐式编码温度信息，在推理时通过调整基分布方差实现温度传递，无需修改模型架构或强制保持体积不变；二是将BoltzNCE风格的代理似然估计器集成到温度退火阶梯中，使重要性加权传输从高温玻尔兹曼分布到室温平衡目标分布成为可能，而这个维度下直接基于雅可比的重新加权是不可行的。与PITA依赖路径积分的SNIS估计器不同，SITA使用EBM计算重要性权重 tilde{w}(hat{x}_1) = exp(-(E(hat{x}_1)/k_BT_{k+1} - hat{U}_phi(hat{x}_1)))，避免了昂贵的散度计算。
方法步骤详情
SITA的退火引导过程包含四个步骤。首先，退火流：从基分布 N(0, T_{k+1}/T_k * I) 采样 x_0，通过积分生成ODE生成样本 hat{x}_1 ~ hat{rho}_{T_{k+1}}。其次，微调EBM：使用流生成的样本，通过BoltzNCE损失 L_BoltzNCE(hat{U}_phi) = L_SM(hat{U}) + L_InfoNCE(hat{U}) 微调EBM N_EBM 个epoch，使其更好地逼近 hat{rho}_{T_{k+1}} 的似然。第三，重要性采样：使用新拟合的EBM计算每个样本的重要性权重 log tilde{w}(hat{x}_1) leftarrow -(E(hat{x}_1)/k_BT_{k+1} - hat{U}_phi(hat{x}_1))，在99分位数处截断，根据自归一化权重重采样 {hat{x}_1} 获得温度 T_{k+1} 的新数据集。最后，微调流：使用重要性加权样本，通过流匹配损失 L_v(hat{v}_theta) 微调流模型 N_Flow 个epoch，使其更好地逼近退火后的目标分布。整个过程重复直到达到最终温度 T_K。
技术新颖性
SITA的技术新颖性体现在三个方面。一是实现了流模型的温度可操纵性，观察到训练在温度 T_high 上的流模型可以将密度升到 kappa = T_high/T_low 次方，通过推理时将基分布方差按 kappa^{-1} 缩放即可生成低温分布样本，无需显式温度条件。二是将代理似然估计器引入退火重要性采样，虽然引入了偏差，但证明其在大分子系统上仍有效。三是完整的端到端退火引导框架，两个模型在整个过程中复用，每步只重新初始化优化器，效率高。与现有方法相比，SITA避免了扩散模型的模式坍塌问题，也不需要像Boltzmann Generators那样需要平衡数据集进行最大似然训练。
实验结果
SITA在两个分子基准系统上达到最先进性能。在丙氨酸二肽（ADP）上，SITA在Rama-KL指标上达到 0.517 plus-minus 0.013，显著优于PITA的 4.773 plus-minus 0.460，在Energy-W2指标上达到 0.939 plus-minus 0.079，优于PITA的 1.615 plus-minus 0.053。能量评估次数为 5.08 times 10^7，与PITA的 5.01 times 10^7 相当。可视化确认SITA捕获了所有主要构象基，匹配参考能量分布。在丙氨酸三肽（ATP）上，SITA在几乎所有指标上超越所有方法：Rama-KL 0.361 plus-minus 0.025（PITA 1.209 plus-minus 0.144），Energy-W1 1.933 plus-minus 0.298（PITA 2.567 plus-minus 0.108），Energy-W2 2.054 plus-minus 0.268（PITA 2.592 plus-minus 0.107）。Remarkably，SITA不需要对生成的样本进行松弛，而PITA和TA-BG必须在目标温度运行短MD细化才能保持竞争力。IMH细化进一步改善性能，在50步后Rama-KL降至 0.313 plus-minus 0.020，Energy-W1降至 0.996 plus-minus 0.029，但代价是额外的能量评估。TICA评估显示，采用均匀下采样时SITA在ATP上TICA-W1达到 0.179 plus-minus 0.012，TICA-W2达到 0.546 plus-minus 0.035，显著优于基于首帧下采样的 0.311 plus-minus 0.014 和 0.943 plus-minus 0.021，说明原始PITA论文的TICA评估存在系统性偏差。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 丙氨酸二肽玻尔兹曼分布采样@300K | Rama-KL (↓) | 0.517 ± 0.013 | PITA: 4.773 ± 0.460 | 89.2% 相对提升 |
| 丙氨酸二肽玻尔兹曼分布采样@300K | Energy-W2 (↓) | 0.939 ± 0.079 | PITA: 1.615 ± 0.053 | 41.9% 相对提升 |
| 丙氨酸三肽玻尔兹曼分布采样@300K | Rama-KL (↓) | 0.361 ± 0.025 | PITA: 1.209 ± 0.144 | 70.1% 相对提升 |
| 丙氨酸三肽玻尔兹曼分布采样@300K | Energy-W1 (↓) | 1.933 ± 0.298 | PITA (w/o relaxation): 86.270 ± 0.294 | 97.8% 相对提升 |
局限与改进
作者承认SITA使用代理似然估计器引入了内在偏差，因为重要性权重估计收敛到倾斜分布 tilde{pi}(x) proportional to rho_theta(x)q_phi(x)/pi(x)，只有当 rho_theta = q_phi 时才能精确恢复目标分布。然而，作者通过实验证明尽管存在这种偏差，EBM导出的重要性权重在基准分子系统上仍提供优越性能。此外，IMH细化虽然改善性能但引入额外计算成本，50步IMH需要 5 times 10^5 次能量评估。我观察到有效样本大小（ESS）相对较低，ADP在300K时ESS为0.268，ATP在300K时ESS为0.065，表明重要性权重分布有长尾，重采样可能导致样本多样性损失。T-W2指标上SITA表现一般，说明准确恢复模式种群仍具挑战性。
独立分析的弱点
SITA的独立弱点包括：一是代理似然偏差可能在大分子系统上更严重，因为EBM学习高维能量函数更困难，偏差会累积到最终分布；二是ESS较低，ADP在退火步骤中ESS从0.131到0.268，ATP仅从0.045到0.074，意味着大部分样本被少数高权重样本主导，重采样后多样性损失明显；三是T-W2指标表现一般，说明构象种群的准确恢复仍是挑战，这可能因为代理似然引入的系统性偏差影响了模式权重的正确性。改进方向包括：(1) 改进EBM架构使其更准确估计似然，可以使用更强大的GNN或引入物理约束；(2) 采用更智能的重采样策略，如分层重采样或残差重采样来保留多样性；(3) 联合优化流和EBM，而不是交替优化，减少偏差累积；(4) 引入正则化项鼓励更均匀的重要性权重分布。
未来方向
作者提出的未来方向包括架构优化、跨系统可迁移性和更大分子系统的应用。基于成果可延伸的方向包括：(1) 将SITA应用于更复杂的分子系统，如蛋白质配体复合物或膜蛋白，验证其可扩展性；(2) 探索替代代理似然估计器，如归一化流或变分自编码器，比较偏差-效率权衡；(3) 研究温度阶梯的自动优化策略，根据ESS动态调整步长；(4) 将SITA与其他采样技术结合，如并行回火或 Hamiltonian Monte Carlo，进一步提升采样效率；(5) 扩展到多目标采样，如同时采样不同温度的系综用于计算热力学性质；(6) 探索SITA在非平衡系统或时间依赖系统中的应用，如构象转变路径采样。
复现评估
SITA的复现性评估如下：代码已在GitHub开源（https://github.com/countrsignal/sita.git），这是一个积极信号。实验在单块NVIDIA L40 GPU（48GB GDDR6）上运行，硬件要求合理。预训练使用500个epoch，批次大小512，流学习率1e-3，EBM学习率5e-4。退火引导使用相同的优化器配置但学习率降为5e-4。对于ADP和ATP，退火阶梯为[755.95K, 555.52K, 408.24K, 300.00K]，每步EBM微调400个epoch，流微调200个epoch。每个退火步骤生成20万个样本用于EBM微调，然后生成10万个样本用于重要性重采样。评估使用3个随机种子，每个种子生成1万个样本。算力需求较高但可接受，主要成本是初始MD训练数据生成（5 times 10^7 次能量评估），退火引导仅增加少量评估。总体而言，复现难度中等，主要挑战可能是EBM训练的不稳定性和长运行时间。
论文图表
