Efficient Regression-Based Training of Normalizing Flows for Boltzmann Generators

• URL: http://arxiv.org/abs/2506.01158v2
• Date: Thu, 30 Oct 2025 07:20:59 GMT
• Title: Efficient Regression-Based Training of Normalizing Flows for Boltzmann Generators
• Authors: Danyal Rehman, Oscar Davis, Jiarui Lu, Jian Tang, Michael Bronstein, Yoshua Bengio, Alexander Tong, Avishek Joey Bose,
• Abstract summary: Boltzmann Generators (BGs) offer efficient sampling and likelihoods, but their training via maximum likelihood is often unstable and computationally challenging.<n>We propose Regression Training of Normalizing Flows (RegFlow), a novel and scalable-based training objective that bypasses the numerical instability and computational challenge of conventional maximum likelihood training.
• Score: 85.25962679349551
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• FALCON: Few-step Accurate Likelihoods for Continuous Flows [78.37361800856583]
  We propose Few-step Accurate Likelihoods for Continuous Flows (FALCON), which allows for few-step sampling with a likelihood accurate enough for importance sampling applications.<n>We show FALCON outperforms state-of-the-art normalizing flow models for molecular Boltzmann sampling and is two orders of magnitude faster than the equivalently performing CNF model.
  arXiv  Detail & Related papers  (2025-12-10T18:47:25Z)
• Flow Straighter and Faster: Efficient One-Step Generative Modeling via MeanFlow on Rectified Trajectories [14.36205662558203]
  Rectified MeanFlow is a framework that models the mean velocity field along the rectified trajectory using only a single reflow step.<n>Experiments on ImageNet at 64, 256, and 512 resolutions show that Re-MeanFlow consistently outperforms prior one-step flow distillation and Rectified Flow methods in both sample quality and training efficiency.
  arXiv  Detail & Related papers  (2025-11-28T16:50:08Z)
• On-the-Fly Fine-Tuning of Foundational Neural Network Potentials: A Bayesian Neural Network Approach [2.0700747055024284]
  Fine-tuning foundation models can reduce the amount of training data necessary to reach a desired level of accuracy.<n>A key challenge for applying this form of active learning to the fine-tuning of foundation models is how to assess the uncertainty of those models during the fine-tuning process.
  arXiv  Detail & Related papers  (2025-07-18T10:33:06Z)
• Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional [2.010573982216398]
  Current machine learning approaches use expensive, task-specific, and data-free training procedures.<n>We demonstrate our approach on varied molecular systems, obtaining diverse, physically realistic transition pathways.<n>Our method can be easily incorporated into new generative models, making it practically relevant as models continue to scale.
  arXiv  Detail & Related papers  (2025-04-25T17:17:17Z)
• TFG-Flow: Training-free Guidance in Multimodal Generative Flow [73.93071065307782]
  We introduce TFG-Flow, a training-free guidance method for multimodal generative flow.<n>TFG-Flow addresses the curse-of-dimensionality while maintaining the property of unbiased sampling in guiding discrete variables.<n>We show that TFG-Flow has great potential in drug design by generating molecules with desired properties.
  arXiv  Detail & Related papers  (2025-01-24T03:44:16Z)
• Transferable Post-training via Inverse Value Learning [83.75002867411263]
  We propose modeling changes at the logits level during post-training using a separate neural network (i.e., the value network)<n>After training this network on a small base model using demonstrations, this network can be seamlessly integrated with other pre-trained models during inference.<n>We demonstrate that the resulting value network has broad transferability across pre-trained models of different parameter sizes.
  arXiv  Detail & Related papers  (2024-10-28T13:48:43Z)
• Free-form Flows: Make Any Architecture a Normalizing Flow [8.163244519983298]
  We develop a training procedure that uses an efficient estimator for the gradient of the change of variables formula. This enables any dimension-preserving neural network to serve as a generative model through maximum likelihood training. We achieve excellent results in molecule generation benchmarks utilizing $E(n)$-equivariant networks.
  arXiv  Detail & Related papers  (2023-10-25T13:23:08Z)
• SE(3)-Stochastic Flow Matching for Protein Backbone Generation [54.951832422425454]
  We introduce FoldFlow, a series of novel generative models of increasing modeling power based on the flow-matching paradigm over $3mathrmD$ rigid motions. Our family of FoldFlowgenerative models offers several advantages over previous approaches to the generative modeling of proteins.
  arXiv  Detail & Related papers  (2023-10-03T19:24:24Z)
• Kalman Filter for Online Classification of Non-Stationary Data [101.26838049872651]
  In Online Continual Learning (OCL) a learning system receives a stream of data and sequentially performs prediction and training steps. We introduce a probabilistic Bayesian online learning model by using a neural representation and a state space model over the linear predictor weights. In experiments in multi-class classification we demonstrate the predictive ability of the model and its flexibility to capture non-stationarity.
  arXiv  Detail & Related papers  (2023-06-14T11:41:42Z)
• Training Energy-Based Normalizing Flow with Score-Matching Objectives [36.0810550035231]
  We present a new flow-based modeling approach called energy-based normalizing flow (EBFlow) We demonstrate that by optimizing EBFlow with score-matching objectives, the computation of Jacobian determinants for linear transformations can be entirely bypassed.
  arXiv  Detail & Related papers  (2023-05-24T15:54:29Z)
• Stochastic Generative Flow Networks [89.34644133901647]
  Generative Flow Networks (or GFlowNets) learn to sample complex structures through the lens of "inference as control" Existing GFlowNets can be applied only to deterministic environments, and fail in more general tasks with dynamics. This paper introduces GFlowNets, a new algorithm that extends GFlowNets to environments.
  arXiv  Detail & Related papers  (2023-02-19T03:19:40Z)
• Improving and generalizing flow-based generative models with minibatch optimal transport [90.01613198337833]
  We introduce the generalized conditional flow matching (CFM) technique for continuous normalizing flows (CNFs) CFM features a stable regression objective like that used to train the flow in diffusion models but enjoys the efficient inference of deterministic flow models. A variant of our objective is optimal transport CFM (OT-CFM), which creates simpler flows that are more stable to train and lead to faster inference.
  arXiv  Detail & Related papers  (2023-02-01T14:47:17Z)
• Normalizing flow neural networks by JKO scheme [22.320632565424745]
  We develop a neural ODE flow network called JKO-iFlow, inspired by the Jordan-Kinderleherer-Otto scheme. The proposed method stacks residual blocks one after another, allowing efficient block-wise training of the residual blocks. Experiments with synthetic and real data show that the proposed JKO-iFlow network achieves competitive performance.
  arXiv  Detail & Related papers  (2022-12-29T18:55:00Z)
• Learning GFlowNets from partial episodes for improved convergence and stability [56.99229746004125]
  Generative flow networks (GFlowNets) are algorithms for training a sequential sampler of discrete objects under an unnormalized target density. Existing training objectives for GFlowNets are either local to states or transitions, or propagate a reward signal over an entire sampling trajectory. Inspired by the TD($lambda$) algorithm in reinforcement learning, we introduce subtrajectory balance or SubTB($lambda$), a GFlowNet training objective that can learn from partial action subsequences of varying lengths.
  arXiv  Detail & Related papers  (2022-09-26T15:44:24Z)
• Deep Equilibrium Optical Flow Estimation [80.80992684796566]
  Recent state-of-the-art (SOTA) optical flow models use finite-step recurrent update operations to emulate traditional algorithms. These RNNs impose large computation and memory overheads, and are not directly trained to model such stable estimation. We propose deep equilibrium (DEQ) flow estimators, an approach that directly solves for the flow as the infinite-level fixed point of an implicit layer.
  arXiv  Detail & Related papers  (2022-04-18T17:53:44Z)
• Embedded training of neural-network sub-grid-scale turbulence models [0.0]
  The weights of a deep neural network model are optimized in conjunction with the governing flow equations to provide a model for sub-grid-scale stresses. The training is by a gradient descent method, which uses the adjoint Navier-Stokes equations to provide the end-to-end sensitivities of the model weights to the velocity fields.
  arXiv  Detail & Related papers  (2021-05-03T17:28:39Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.