Exactly invertible flow maps improve image generation and sampling quality
Normalizing Flow Maps
Machine Learning
Summary
Many computer programs that create images or data often struggle to calculate exactly how likely their creations are, which limits understanding and control. The authors introduce a new kind of model called Normalizing Flow Maps that combines fast generation with the ability to exactly compute likelihoods. These models learn how data smoothly changes over time in a way that makes it easier to both create new samples and evaluate them accurately. Tests on image generation and molecular sampling tasks show these models produce better results than similar previous methods.
What this means in practice
- •For image generation engineers: Create higher-quality conditional images with efficiently computable exact likelihoods for better model training and evaluation.
- •For computational chemists: Generate molecular configurations at equilibrium states with improved sampling accuracy using Normalizing Flow Maps in Boltzmann Generators.
Authors
Francesco M. Ruscio, Timothy James Hitge, Fred Zhangzhi Peng, Danyal Rehman, Zander W. Blasingame, Nicholas M. Boffi, Michael M. Bronstein, Alexander Tong, Avishek Joey Bose
Abstract
Scalable approaches to learning neural dynamical systems have placed diffusion, flow matching, and few-step generators in flow maps at the forefront of generative modeling. Such state-of-the-art models have limited access to explicit model densities despite evolving from classical generative modeling approaches that leveraged tractable likelihoods, e.g., Normalizing Flows and Autoregressive models. In this paper, we combine these threads of few-step generation and tractable sample likelihoods and introduce Normalizing Flow Maps (NFM), a class of exactly invertible pushforward maps that generalizes normalizing flows and flow maps. More precisely, NFMs learn the integrator of an underlying dynamical system, but crucially, they provide efficient exact likelihood evaluation by construction. To enable scalable training of NFMs, we derive two complementary axes of consistency inherent to NFMs that yield (self) distillation objectives: (1) pointwise matching, which unlocks $\ell_2$-based flow-map style regression losses, and (2) distributional matching, which matches pushforward distributions to the marginals of a prescribed forward process. On conditional image generation on CIFAR-10 and ImageNet-64, NFMs achieve state-of-the-art sample quality among parameter-matched normalizing flows, while supporting likelihood-based RLHF. For equilibrium sampling in Boltzmann Generators, NFMs outperform prior state-of-the-art normalizing flows on both single-system and transferable peptide benchmarks.