Jeffreys Flow: Robust Boltzmann Generators for Rare Event Sampling via Parallel Tempering Distillation
作者: Guang Lin, Christian Moya, Di Qi, Xuda Ye
分类: cs.LG, math.NA, physics.comp-ph, stat.ML
发布日期: 2026-04-07
💡 一句话要点
提出Jeffreys Flow，通过并行回火蒸馏解决玻尔兹曼生成器中的模式崩塌问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 玻尔兹曼生成器 Jeffreys Flow 并行回火 模式崩塌 罕见事件采样
📋 核心要点
- 玻尔兹曼生成器在复杂物理系统采样中易发生模式崩塌，导致无法覆盖所有模式。
- Jeffreys Flow通过Jeffreys散度蒸馏并行回火数据，平衡局部精度和全局覆盖，抑制模式崩塌。
- 实验证明Jeffreys Flow在非凸问题上具有可扩展性和准确性，并能校正随机梯度偏差。
📝 摘要（中文）
对具有粗糙能量面的物理系统进行采样受到罕见事件和亚稳态陷阱的阻碍。虽然玻尔兹曼生成器已经提供了一种解决方案，但它们对反向Kullback-Leibler散度的依赖经常导致灾难性的模式崩塌，从而遗漏多模态分布中的特定模式。本文介绍了一种鲁棒的生成框架Jeffreys Flow，它通过使用对称的Jeffreys散度从并行回火轨迹中提炼经验采样数据来缓解这种失败。这种公式有效地平衡了局部目标寻找精度与全局模式覆盖。我们表明，最小化Jeffreys散度可以抑制模式崩塌，并通过经验参考数据的提炼来结构性地纠正固有的不准确性。我们在高度非凸多维基准上展示了该框架的可扩展性和准确性，包括系统地校正副本交换随机梯度朗之万动力学中的随机梯度偏差，以及大规模加速量子热态路径积分蒙特卡罗中的精确重要性采样。
🔬 方法详解
问题定义：论文旨在解决玻尔兹曼生成器在对具有复杂能量景观的物理系统进行采样时，由于依赖反向KL散度而导致的模式崩塌问题。现有方法无法有效覆盖多模态分布中的所有模式，导致采样结果不准确。
核心思路：论文的核心思路是使用Jeffreys散度作为训练目标，并结合并行回火（Parallel Tempering）策略进行数据蒸馏。Jeffreys散度是一种对称散度，能够更好地平衡局部精度和全局模式覆盖，从而避免模式崩塌。通过从并行回火轨迹中提取经验数据，可以纠正玻尔兹曼生成器固有的不准确性。
技术框架：Jeffreys Flow框架主要包含以下几个阶段：1) 使用并行回火方法生成多个温度下的采样数据；2) 使用Jeffreys散度作为损失函数训练玻尔兹曼生成器，目标是拟合并行回火生成的经验分布；3) 通过迭代优化，使玻尔兹曼生成器能够准确地捕捉目标分布的各个模式。
关键创新：最重要的技术创新点在于使用Jeffreys散度替代传统的反向KL散度。Jeffreys散度的对称性使其在训练过程中能够同时关注生成分布和目标分布，从而避免了模式崩塌问题。此外，结合并行回火进行数据蒸馏，能够有效地利用不同温度下的采样信息，提高生成模型的鲁棒性和准确性。
关键设计：论文中关键的设计包括：1) Jeffreys散度的具体计算公式；2) 并行回火中温度的选择策略；3) 玻尔兹曼生成器的网络结构，例如可以使用神经网络来参数化能量函数；4) 优化算法的选择，例如可以使用Adam等优化器进行训练。
🖼️ 关键图片
📊 实验亮点
实验结果表明，Jeffreys Flow在高度非凸多维基准测试中表现出优异的性能，能够有效地抑制模式崩塌，并准确地捕捉目标分布的各个模式。例如，在副本交换随机梯度朗之万动力学中，Jeffreys Flow能够系统地校正随机梯度偏差。此外，Jeffreys Flow还能够大规模加速量子热态路径积分蒙特卡罗中的精确重要性采样。
🎯 应用场景
该研究成果可应用于材料科学、药物发现等领域，用于加速复杂物理系统的模拟和采样。通过更准确地模拟分子动力学、蛋白质折叠等过程，可以帮助研究人员发现新的材料和药物，并优化现有产品的性能。此外，该方法还可以应用于机器学习中的生成模型训练，提高生成模型的质量和多样性。
📄 摘要（原文）
Sampling physical systems with rough energy landscapes is hindered by rare events and metastable trapping. While Boltzmann generators already offer a solution, their reliance on the reverse Kullback--Leibler divergence frequently induces catastrophic mode collapse, missing specific modes in multi-modal distributions. Here, we introduce the Jeffreys Flow, a robust generative framework that mitigates this failure by distilling empirical sampling data from Parallel Tempering trajectories using the symmetric Jeffreys divergence. This formulation effectively balances local target-seeking precision with global modes coverage. We show that minimizing Jeffreys divergence suppresses mode collapse and structurally corrects inherent inaccuracies via distillation of the empirical reference data. We demonstrate the framework's scalability and accuracy on highly non-convex multidimensional benchmarks, including the systematic correction of stochastic gradient biases in Replica Exchange Stochastic Gradient Langevin Dynamics and the massive acceleration of exact importance sampling in Path Integral Monte Carlo for quantum thermal states.