FORT: Forward-Only Regression Training of Normalizing Flows

2025-06-01

AI Summary

Traditional normalizing flow models rely on computationally expensive Jacobian determinant evaluations, limiting their applicability to scientific tasks such as molecular conformational sampling. To address this, we propose a forward regression training paradigm that abandons inverse simulation and maximum-likelihood objectives, introducing instead a novel single-step ℓ₂ regression loss. This loss unifies optimal transport and continuous normalizing flow (CNF) distillation pretraining. The approach significantly improves training stability, scalability, and architectural generality. On equilibrium conformational sampling of alanine di-, tri-, and tetrapeptides, our model efficiently generates high-fidelity Cartesian coordinates, outperforming conventional training methods in both sample quality and likelihood estimation accuracy. Our work establishes a new paradigm for efficient generative modeling in scientific computing.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Natural Language Processing: Generation
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: User modeling and simulation for interactive and conversational systems
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Simulation-free training frameworks have been at the forefront of the generative modelling revolution in continuous spaces, leading to neural dynamical systems that encompass modern large-scale diffusion and flow matching models. Despite the scalability of training, the generation of high-quality samples and their corresponding likelihood under the model requires expensive numerical simulation -- inhibiting adoption in numerous scientific applications such as equilibrium sampling of molecular systems. In this paper, we revisit classical normalizing flows as one-step generative models with exact likelihoods and propose a novel, scalable training objective that does not require computing the expensive change of variable formula used in conventional maximum likelihood training. We propose Forward-Only Regression Training (FORT), a simple $ell_2$-regression objective that maps prior samples under our flow to specifically chosen targets. We demonstrate that FORT supports a wide class of targets, such as optimal transport targets and targets from pre-trained continuous-time normalizing flows (CNF). We further demonstrate that by using CNF targets, our one-step flows allow for larger-scale training that exceeds the performance and stability of maximum likelihood training, while unlocking a broader class of architectures that were previously challenging to train. Empirically, we elucidate that our trained flows can perform equilibrium conformation sampling in Cartesian coordinates of alanine dipeptide, alanine tripeptide, and alanine tetrapeptide.

Problem

• Eliminates expensive simulation in normalizing flow training
• Proposes scalable one-step generative models with exact likelihoods
• Enables equilibrium sampling for molecular systems efficiently

Innovation

• Forward-Only Regression Training (FORT) method
• Avoids expensive change of variable computation
• Uses ell_2-regression for prior sample mapping