Few-Step Boltzmann Generators via Scalable Likelihood Flow Maps
RuiKang OuYang, Hanlin Yu, Xinyue Ai, Yutong He, Nicholas M. Boffi, Pradeep Ravikumar, Jose Miguel Hernandez-Lobato, Max Simchowitz, Benjamin Kurt Miller, Omar Chehab
cs.LG, stat.ML
Submitted: 2026-06-27
License: http://creativecommons.org/licenses/by/4.0/
The gist: Recent progress in flow-based generative modeling has led to models that output high-quality samples while using only a small number of function evaluations.
Terminology
Abstract
Recent progress in flow-based generative modeling has led to models that output high-quality samples while using only a small number of function evaluations. However, at present, there is a lack of similar advances in estimating the model likelihood. In particular, most existing methods either rely on restrictive architectures that enable exact calculations, or use stochastic approximations such as Hutchinson's trace estimator that introduce substantial variance. In this work, we introduce SCAlable LikeLihood distillation of flOw maPs (SCALLOP). SCALLOP builds on the recently proposed F2D2, a likelihood flow map model that can generate samples and their densities in a small number of function evaluations. While F2D2 uses Hutchinson's estimator during training, we introduce an alternative and more scalable likelihood distillation objective that is Hutchinson-free and admits a vectorized formulation. Empirically, we demonstrate the effectiveness of SCALLOP as a Boltzmann generator in molecular science, and further validate its benefit on image datasets. SCALLOP significantly reduces both training variance and training time while consistently improving performance compared to F2D2, and is competitive with the state-of-the-art while achieving up to 10x inference speedup over the fastest baseline.
Sources
- BoltzNCE: Learning Likelihoods for Boltzmann Generation with Stochastic Interpolants and Noise Contrastive Estimation
- How to build a consistency model: Learning flow maps via self-distillation
- Neural Ordinary Differential Equations
- Free-form Flows: Make Any Architecture a Normalizing Flow
- Mean Flows for One-step Generative Modeling
- FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models
- Learning normalized image densities via dual score matching
- RNE: plug-and-play diffusion inference-time control and energy-based training
- Denoising Diffusion Probabilistic Models
- Enhancing Diffusion-Based Sampling with Molecular Collective Variables
- Improving Probabilistic Diffusion Models With Optimal Diagonal Covariance Matching
- A Diffusive Classification Loss for Learning Energy-based Generative Models
- BNEM: A Boltzmann Sampler Based on Bootstrapped Noised Energy Matching
- Consistent Sampling and Simulation: Molecular Dynamics with Energy-Based Diffusion Models
- Progressive Tempering Sampler with Diffusion
- Scalable Boltzmann Generators for equilibrium sampling of large-scale materials
- Score-Based Generative Modeling through Stochastic Differential Equations
- Scalable Equilibrium Sampling with Sequential Boltzmann Generators
- Amortized Sampling with Transferable Normalizing Flows
- Enhanced Diffusion Sampling: Efficient Rare Event Sampling and Free Energy Calculation with Diffusion Models
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: Local-Consistency Optimisation for Bayesian Neural Networks
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks