Mode Coverage in Normalizing Flow Boltzmann Generators via Log-Ratio Variation
Abstract
Normalizing flow Boltzmann generators retain a tractable pushforward density, but training with forward KL depends on target samples that may be biased or omit modes. As a result, a flow can miss target mass while its observed importance weights give a high effective sample size. We introduce the log-ratio variation Xω, the mean absolute pairwise difference of the target-to-pushforward log-density ratio under a weighting measure ω, and use it to define KLXX, a new loss function. Two log-ratio variations are added to the forward KL (denoted by the two X's): one weighted by the target to improve accuracy, the other by a mixture of quench and temper samples with pushforward samples to search candidate modes. We derive the Fisher--Rao gradient flow of KLXX, where both variations contribute nonpositive dissipation, and a fixed-surrogate error bound for KLXX. We use KLXX in an adaptive-staging Boltzmann generator, with importance reweighting at every stage. We bound the sampling error of its inference scheme when the stage weights are essentially bounded, and prove it asymptotically unbiased in the sample size. In the numerical tests, KLXX improves mode coverage over forward KL. It also improves the generator's per-stage diagnostics against the loss that built the schedule. The observables the generator recovers are close to independent references. The log-ratio variations thus supply information that the forward KL loss usually omits.
Related
- Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows
- Boltzmann generators for amorphous particle systems
- Few-Step Boltzmann Generators via Scalable Likelihood Flow Maps
- Recovering Weak Signals with Normalizing Flows
- Neural Sampling with Reweighted Normalizing Flows via the Wasserstein--Fisher--Rao JKO Scheme
BibTeX
@misc{feng2026mode,
title = {{Mode Coverage in Normalizing Flow Boltzmann Generators via Log-Ratio Variation}},
author = {Qi Feng and Rongjie Lai and Di Qi and Xuda Ye},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.09473},
eprint = {2609.09473},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.09473},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Mode Coverage in Normalizing Flow Boltzmann Generators via Log-Ratio Variation
AU - Qi Feng
AU - Rongjie Lai
AU - Di Qi
AU - Xuda Ye
PY - 2026
DA - 2026/09/08
PB - arXiv
DO - 10.48550/arxiv.2609.09473
UR - https://arxiv.org/abs/2609.09473
AB - Normalizing flow Boltzmann generators retain a tractable pushforward density, but training with forward KL depends on target samples that may be biased or omit modes. As a result, a flow can miss target mass while its observed importance weights give a high effective sample size. We introduce the log-ratio variation X_ω, the mean absolute pairwise difference of the target-to-pushforward log-density ratio under a weighting measure ω, and use it to define KLXX, a new loss function. Two log-ratio variations are added to the forward KL (denoted by the two X's): one weighted by the target to improve accuracy, the other by a mixture of quench and temper samples with pushforward samples to search candidate modes. We derive the Fisher--Rao gradient flow of KLXX, where both variations contribute nonpositive dissipation, and a fixed-surrogate error bound for KLXX. We use KLXX in an adaptive-staging Boltzmann generator, with importance reweighting at every stage. We bound the sampling error of its inference scheme when the stage weights are essentially bounded, and prove it asymptotically unbiased in the sample size. In the numerical tests, KLXX improves mode coverage over forward KL. It also improves the generator's per-stage diagnostics against the loss that built the schedule. The observables the generator recovers are close to independent references. The log-ratio variations thus supply information that the forward KL loss usually omits.
ER -