Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Efficient Regression-Based Training of Normalizing Flows for Boltzmann Generators

Danyal Rehman, Oscar Davis, Jiarui Lu, Jian Tang, Michael Bronstein, Yoshua Bengio, Alexander Tong, Avishek Joey Bose

TL;DR

This work tackles the need for fast, exact likelihood evaluation in Boltzmann Generators by revisiting classical normalizing flows and proposing Regression Training of Normalizing Flows (RegFlow). RegFlow replaces maximum-likelihood training with a simple $\ell_2$ regression objective that aligns a learned invertible map with a fixed invertible target, drawn from either an offline optimal-transport map or a pretrained continuous normalizing flow (CNF); stability is enhanced via regularization and a forward-backward consistency term. The authors demonstrate that RegFlow enables training of NF architectures (e.g., affine coupling, neural spline flows) that were previously intractable under MLE for BGs and yields improvements in equilibrium sampling and targeted free energy perturbation across alanine dipeptide, tripeptide, and tetrapeptide—often with substantial inference-speedups over MLE-trained models. The results indicate RegFlow can deliver faithful samples and exact likelihoods in a computationally efficient framework, expanding the practical deployment of BGs in molecular systems while highlighting reliance on the quality of chosen targets. Overall, RegFlow offers a principled, scalable route to leverage classical NFs for scientific applications requiring reliable likelihoods and fast sampling.

Abstract

Simulation-free training frameworks have been at the forefront of the generative modelling revolution in continuous spaces, leading to large-scale diffusion and flow matching models. However, such modern generative models suffer from expensive inference, inhibiting their use in numerous scientific applications like Boltzmann Generators (BGs) for molecular conformations that require fast likelihood evaluation. In this paper, we revisit classical normalizing flows in the context of BGs that offer efficient sampling and likelihoods, but whose training via maximum likelihood is often unstable and computationally challenging. We propose Regression Training of Normalizing Flows (RegFlow), a novel and scalable regression-based training objective that bypasses the numerical instability and computational challenge of conventional maximum likelihood training in favour of a simple $\ell_2$-regression objective. Specifically, RegFlow maps prior samples under our flow to targets computed using optimal transport couplings or a pre-trained continuous normalizing flow (CNF). To enhance numerical stability, RegFlow employs effective regularization strategies such as a new forward-backward self-consistency loss that enjoys painless implementation. Empirically, we demonstrate that RegFlow unlocks a broader class of architectures that were previously intractable to train for BGs with maximum likelihood. We also show RegFlow exceeds the performance, computational cost, and stability of maximum likelihood training in equilibrium sampling in Cartesian coordinates of alanine dipeptide, tripeptide, and tetrapeptide, showcasing its potential in molecular systems.

Efficient Regression-Based Training of Normalizing Flows for Boltzmann Generators

TL;DR

This work tackles the need for fast, exact likelihood evaluation in Boltzmann Generators by revisiting classical normalizing flows and proposing Regression Training of Normalizing Flows (RegFlow). RegFlow replaces maximum-likelihood training with a simple regression objective that aligns a learned invertible map with a fixed invertible target, drawn from either an offline optimal-transport map or a pretrained continuous normalizing flow (CNF); stability is enhanced via regularization and a forward-backward consistency term. The authors demonstrate that RegFlow enables training of NF architectures (e.g., affine coupling, neural spline flows) that were previously intractable under MLE for BGs and yields improvements in equilibrium sampling and targeted free energy perturbation across alanine dipeptide, tripeptide, and tetrapeptide—often with substantial inference-speedups over MLE-trained models. The results indicate RegFlow can deliver faithful samples and exact likelihoods in a computationally efficient framework, expanding the practical deployment of BGs in molecular systems while highlighting reliance on the quality of chosen targets. Overall, RegFlow offers a principled, scalable route to leverage classical NFs for scientific applications requiring reliable likelihoods and fast sampling.

Abstract

Simulation-free training frameworks have been at the forefront of the generative modelling revolution in continuous spaces, leading to large-scale diffusion and flow matching models. However, such modern generative models suffer from expensive inference, inhibiting their use in numerous scientific applications like Boltzmann Generators (BGs) for molecular conformations that require fast likelihood evaluation. In this paper, we revisit classical normalizing flows in the context of BGs that offer efficient sampling and likelihoods, but whose training via maximum likelihood is often unstable and computationally challenging. We propose Regression Training of Normalizing Flows (RegFlow), a novel and scalable regression-based training objective that bypasses the numerical instability and computational challenge of conventional maximum likelihood training in favour of a simple -regression objective. Specifically, RegFlow maps prior samples under our flow to targets computed using optimal transport couplings or a pre-trained continuous normalizing flow (CNF). To enhance numerical stability, RegFlow employs effective regularization strategies such as a new forward-backward self-consistency loss that enjoys painless implementation. Empirically, we demonstrate that RegFlow unlocks a broader class of architectures that were previously intractable to train for BGs with maximum likelihood. We also show RegFlow exceeds the performance, computational cost, and stability of maximum likelihood training in equilibrium sampling in Cartesian coordinates of alanine dipeptide, tripeptide, and tetrapeptide, showcasing its potential in molecular systems.

Paper Structure

This paper contains 29 sections, 5 theorems, 25 equations, 13 figures, 9 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background and Preliminaries
  3. Continuous normalizing flows
  4. Normalizing flows
  5. Regression Training of Normalizing Flows
  6. Warmup: One-step generative models without likelihood
  7. Training normalizing flows using regression
  8. RegFlow targets
  9. Experiments
  10. Related work
  11. Conclusion
  12. Acknowledgements
  13. Proofs
  14. Proof of \ref{['prop:pointwise']}
  15. Proof of \ref{['prop:wasserstein_bound']}
  16. ...and 14 more sections

Key Result

Proposition 1

Suppose that $f^\star_t$ is invertible for all $t$, that $(f_t^\star)^{-1}$ is continuous for all $t$. Then, as ${\mathcal{L}}(\theta)\to 0$, it holds that $((f_t^\star)^{-1}\circ f_{t,\theta})(x) \to x$ for almost all (with respect to $p_0$) $x$.

Figures (13)

  • Figure 1: Evaluation of IMM and shortcut models with exact likelihood on the synthetic checkerboard experiment. Depictions are provided of the 2D histograms after self-normalizing importance sampling is used.
  • Figure 2: Energy distributions and Ramachandran plots for alanine dipeptide. (left to right): Energy distribution of most best MLE-trained NF; energy distribution of best RegFlow; ground truth MD data torsion angle distribution, best MLE-trained model Ramachandran plot; best RegFlow Ramachandran plot.
  • Figure 3: Left and center: Ablations demonstrating performance improvements with an increasing number of reflow samples. Right: Increasing regularization improves $\mathbb{T}$-$\mathcal{W}_2$ up to a certain point, beyond which numerical invertibility is guaranteed but the regression objective, and subsequently, sample quality, is adversely impacted.
  • Figure 4: Left: The $\beta_{\mathrm{planar}}$ and $\alpha_{\mathrm{R}}$ conformation states; Right: RegFlow's ability to learn free energy differences.
  • Figure 5: Generations of IMM trained with an iUNet with a variable number of steps.
  • ...and 8 more figures

Theorems & Definitions (7)

  • Proposition 1
  • Proposition 2
  • Proposition 2
  • Lemma 1
  • proof
  • Proposition 2
  • proof