🤖 machine learning

Efficient Training of Boltzmann Generators Using Off-Policy Log-Dispersion Regularization

本文介绍了一种离策对数离散化正则化（off-policy log-dispersion regularization, LDR），这是一种通过利用目标能量标签来正则化能量景观，且无需额外在策样本的新型框架，旨在提升玻尔兹曼生成器的数据效率和性能。

原作者： Henrik Schopmans, Christopher von Klitzing, Pascal Friederich

发布于 2026-08-06

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Henrik Schopmans, Christopher von Klitzing, Pascal Friederich

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图预测天气，但你观察的不是云朵，而是试图弄清楚一滴水中一个微小的、看不见的分子是如何扭动和旋转的。在计算科学的世界里，这是一个巨大的谜题。分子就像复杂的、高维度的迷宫，它们可能会陷入深谷（稳定的形状），也可能会爬上高山（不稳定的形状）。为了理解药物如何发挥作用或蛋白质如何折叠，科学家需要绘制出分子可能采取的所有形状的图谱，并根据这些形状出现的可能性进行加权。这被称为从“玻尔兹曼分布”（Boltzmann distribution）中进行采样。

问题在于，分子如此复杂，以至于传统的方法就像是通过一步一挪来绘制整个大陆的地图。你可能会花费数年时间只是绕着一座山转圈，永远看不到山的另一边。为了加速这一过程，科学家们使用了“玻尔兹曼生成器”（Boltzmann generators）。把它们想象成超级聪明的、由人工智能驱动的传送门。与其一步步行走，它们学会了地图，并能瞬间“传送”到分子可以采取的任何有效形状。但问题在于，教导这些传送门是非常昂贵的。每当 AI 猜出一个形状时，计算机都必须计算其能量，这就像是在向一个非常缓慢、非常昂贵的先知询问一张票的价格。如果 AI 需要向先知提问数百万次才能学会地图，那么整个过程就会变得太慢而无法使用。

这篇论文介绍了一种聪明的技巧，通过一种叫做**离策对数离散正则化（Off-Policy Log-Dispersion Regularization, LDR）**的方法来更快地教导这些 AI 传送门。作者 Henrik Schopmans、Christopher von Klitzing 和 Pascal Friederich 意识到，在教导 AI 时，我们忽略了一个巨大的线索：能量值本身。通常情况下，AI 只是被告知：“这里是分子采取的一些形状；学习去模仿它们。”但论文提出了一种更好的方法：“这里是一些形状，以及每个形状对应的精确能量成本。确保你的猜测与这些成本相匹配。”

他们提出了一种新的“正则化项”（regularizer），它就像是一个给 AI 提供额外反馈的教练。想象一下你正在学习画一幅风景画。一个标准的老师可能会只说：“在这里画一棵树。”但这个新的教练会说：“在这里画一棵树，并确保阴影与山的高度完全匹配。”如果你的画作相对于真实的能量高度来说太暗或太亮，教练就会轻轻推你一下，让你修正画作的形状。这就是这个“推力”——即 LDR。它不需要 AI 向先知请求新的、昂贵的能量计算；它只是利用了数据中已经存在的能量标签。

结果非常令人印象深刻。在测试中，作者发现使用这种新的教练方法，可以让 AI 以更少的昂贵能量计算次数来学习地图。例如，在对一种名为丙氨酸二肽（alanine dipeptide）的分子进行训练时，新方法仅用 100 万个样本就达到了高质量的结果，而旧方法需要 500 万个样本才能达到类似的质量。这实现了五倍的效率提升。在针对更大的分子（丙氨酸六肽，alanine hexapeptide）进行更具挑战性的测试中，改进程度非常显著，AI 学习稳定地图所需的能量评估次数比之前的最佳方法减少了 10 倍。

该论文还表明，即使数据是“有偏的”（biased）——也就是说，AI 只被展示了分子世界的某个微小、奇特的角落——这个技巧仍然有效。通常情况下，这会让 AI 感到困惑，但 LDR 教练帮助它理清了其余部分的地图。作者在简单的 2D 形状和复杂的 3D 分子模型上都进行了测试，在每种情况下，AI 都学得更快，犯的错误也更少。他们甚至展示了当 AI 在没有看到任何真实分子形状的情况下，仅依靠能量计算进行学习时，该方法依然奏效。

简而言之，这篇论文并没有发明一种新的分子类型或新的物理定律。相反，它发明了一种更好的方法来教导 AI 如何阅读现有的地图。通过将能量标签作为一种“形状正则化项”，作者表明我们可以将模拟分子的过程效率提高多达十倍。这意味着在未来，科学家们可能能够更快地设计新药或理解复杂的生物过程，因为进行繁重计算的计算机不必如此频繁地向昂贵的先知寻求帮助。该方法具有灵活性，适用于不同类型的数据和不同类型的 AI 模型，这表明它可能成为任何试图模拟微观世界的人的标准工具。

技术摘要：使用离策对数离散正则化实现玻尔兹曼生成器的高效训练

1. 问题陈述

从未归一化概率密度（特别是物理系统的玻尔兹曼分布）中进行采样是计算物理和化学中的一个基本挑战。目标分布定义为 $p_X(x) = \tilde{p}_X(x)/Z$ ，其中 $\tilde{p}_X(x) = \exp(-E(x)/k_B T)$ 是未归一化密度， $E(x)$ 是能量函数， $Z$ 是难以计算的归一化常数。

虽然经典的分子动力学（MD）和马尔可夫链蒙特卡洛（MCMC）方法在渐近意义上是正确的，但由于在高维或能量复杂的系统中存在长相关轨迹，它们往往面临混合缓慢和计算成本高的问题。玻尔兹曼生成器（Boltzmann Generators, BGs）提供了一种生成式建模的替代方案，能够实现从平衡分布中进行独立采样。然而，其实际应用受到数据效率的限制。标准的训练目标通常依赖于：

来自目标分布的平衡样本（通常通过 MD 生成，成本昂贵）。
变分设置中大量的能量评估。

现有方法通常无法充分利用在数据集生成过程中可获得的靶向能量标签，或者需要额外的在策（on-policy）样本和能量评估才能有效地整合这些信息。

2. 方法论：离策对数离散正则化 (LDR)

作者提出了**离策对数离散正则化（Off-Policy Log-Dispersion Regularization, LDR）**框架，该框架能够在不要求额外的在策样本或额外能量评估的情况下，将目标能量信息整合到玻尔兹曼生成器的训练中。

核心概念

该方法基于传统用于在策变分训练中的对数方差目标。作者定义了一个函数 $f_\theta(x)$ ，表示负对数模型密度与负对数目标密度（未归一化重要性权重）之间的差异：
$f_\theta(x) = -\log q_\theta^X(x) - (-\log \tilde{p}_X(x)) = -\log q_\theta^X(x) - \frac{E(x)}{k_B T}$
在最优解处，即 $q_\theta^X = p_X$ 时， $f_\theta(x)$ 应为常数 ( $\log Z$ )。

对数离散目标

作者将对数方差目标推广为一族对数离散目标：
$\mathcal{L}_\theta^{LD} = \mathbb{E}_{r_X} \left[ \rho(f_\theta(x) - c_\theta) \right], \quad c_\theta = \mathbb{E}_{r_X}[f_\theta(x)]$
其中 $\rho(u) \geq 0$ 是离散度度量（例如 $|u|^p$ ）， $r_X$ 是参考分布。

LDR-L2 ( $p=2$ )： 最小化方差（二阶中心矩），恢复标准的对数方差目标。
LDR-L1 ( $p=1$ )： 最小化一阶绝对中心矩。作者假设该变体对能量异常值（outliers）的敏感度较低，并能提供更稳定的优化，这在高维能量景观中具有重尾性质的重要性权重时至关重要。

离策正则化策略

一个关键的理论洞察（附录 A）是，如果参考分布 $r_X$ （例如一个固定的数据集）在配置空间上没有全支撑（full support），那么仅最小化对数离散目标不是一种散度；它无法约束数据流形之外的模型，从而导致归一化错误。

为了解决这个问题，LDR 被作为一种正则项应用于标准的数据驱动目标（例如 Forward KL、流匹配或分数匹配）旁边：
$\mathcal{L}_\theta = \lambda_{data}\mathcal{L}_\theta^{data} + \lambda_{LD}\mathcal{L}_\theta^{LD}$

数据驱动项确保全局归一化并收敛到正确的目标分布。
LDR 项作为一个形状正则化器，通过最小化数据流形上对数重要性权重的离散度，使提议分布（proposal distribution）与目标能量景观对齐。

这种方法允许使用任意参考分布，包括偏置模拟数据或辅助能量标记数据集，而无需在训练期间进行额外的目标能量评估。

3. 主要贡献

向离策设置的适配： 作者将对数方差目标从在策变分训练适配到离策设置，将其作为固定能量标记数据集上的正则项使用。
向对数离散的泛化： 他们将目标推广到一系列对数离散损失，引入了 LDR-L1 变体，与标准的 LDR-L2 相比，它在应对异常值方面提供了更好的稳定性。
广泛的适用性： 该框架在三种不同的训练机制中均表现出有效性：
- 在无偏平衡数据上训练。
- 在偏置模拟数据上训练（利用重要性采样和混合数据集）。
- 纯变分训练，无需访问目标样本（仅使用目标能量评估）。
与最先进技术的集成： 该方法成功集成到受限质量传输（Constrained Mass Transport, CMT）中——这是当前变分训练的最先进技术——显著提高了其效率。

4. 实验结果

作者在包括 2D 高斯混合模型 (GMM)、丙氨酸二肽 ( $d=60$ ) 和丙氨酸六肽 ( $d=180$ ) 在内的多个基准测试上评估了 LDR。性能使用负对数似然 (NLL) 和有效样本量 (ESS) 进行衡量。

关键发现：

无偏数据： 在无偏数据集上，LDR 显著改善了 NLL 和 ESS。对于丙氨酸二肽，LDR 在 $1 \times 10^6$ 个样本下的性能超过了使用 $5 \times 10^6$ 个样本的非正则化训练，表明数据效率提升了 5 倍。
偏置数据： 当在偏置数据集（通过短轨迹生成）上训练时，将 LDR 应用于重采样后的平衡数据集 ( $D_{IS}$ ) 和原始偏置数据集均取得了实质性收益。LDR 使仅使用 $1 \times 10^5$ 个重要性样本加上初始偏置数据集构建即可训练出高性能模型，而传统的非正则化基准则需要超过 10 倍的重要性样本。
变分训练（无目标样本）： 在使用 CMT 的纯变分设置中，LDR 将所需的靶向能量评估次数减少了约一个数量级。对于丙氨酸二肽，CMT + LDR 在仅使用 $1 \times 10^7$ 次评估的情况下，达到了与标准 CMT 使用 $1 \times 10^8$ 次评估相当的性能。对于丙氨酸六肽，LDR 在 vanilla CMT 失效的机制中防止了模式崩溃（实现了 >30% 的 ESS，而 vanilla CMT <6%）。
笛卡尔坐标： 当直接在笛卡尔坐标上使用 TarFlow 进行训练时，上述增益得以复制，证实了其收益并不依赖于内部坐标表示。
稳定性： 在高维系统且具有重尾重要性权重的情况下，LDR-L1 通常比 LDR-L2 提供更稳定的训练动力学（更低的梯度范数方差）。

5. 意义与主张

本文声称 LDR 提供了一个简单而强大的框架，用于将目标能量信息整合到玻尔兹曼生成器训练中。其主要意义在于：

数据效率： 它在多种不同设置下一致地提高了最终性能和数据效率，通常将所需样本数或能量评估次数减少了一个数量级。
灵活性： 它具有广泛的适用性，支持无偏、偏置和纯变分训练，且不需要额外的在策样本或额外的能量评估。
互补性： 作者指出 LDR 与现有的方差缩减技术（如使用目标梯度的路径梯度法）是互补的，因为 LDR 利用的是目标能量值（零阶信息）。

作者强调，虽然 LDR 并没有消除对数据集大小的依赖，但它通过提高每个样本所提取的信息量，显著减轻了对大规模数据集的需求。该框架旨在易于集成到现有的归一化流（normalizing flows）及潜在的扩散生成器训练流水线中。