Few-Step Boltzmann Generators via Scalable Likelihood Flow Maps

2026-06-27

AI Summary

Existing flow models struggle to simultaneously achieve efficient sampling and accurate density estimation under a limited number of function evaluations, often relying on restrictive architectures or high-variance Hutchinson trace estimators. This work proposes SCALLOP, a scalable, Hutchinson-free likelihood distillation method that trains few-step flow models using a vectorized deterministic likelihood objective, substantially reducing training variance and enhancing scalability. Built upon the F2D2 framework, SCALLOP synergistically combines the strengths of flow models and Boltzmann generators. It outperforms baseline methods across molecular science and image generation tasks, offering faster training, up to 10× acceleration in inference, and state-of-the-art performance in both high-quality sample generation and density estimation.

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.

Problem

• likelihood estimation
• flow-based generative models
• Boltzmann generators
• stochastic approximation
• scalability

Innovation

• likelihood distillation
• flow-based generative models
• Hutchinson-free
• scalable inference
• Boltzmann generators