Few-Step Boltzmann Generators via Scalable Likelihood Flow Maps
Abstract Overview
This paper introduces SCALLOP, a training method for likelihood flow maps that aims to generate samples and estimate their likelihoods in only a few function evaluations. The method is designed as an alternative to prior approaches such as F2D2, which rely on Hutchinson’s trace estimator for divergence estimation during training and therefore suffer from high variance and extra computational cost. SCALLOP replaces that component with conditional divergence matching and a vectorized formulation, so likelihood-related quantities can be distilled without Hutchinson-based estimation. The paper studies the method on molecular Boltzmann generation tasks and on image generation with CelebA-64, focusing on training stability, runtime, and likelihood accuracy.
Novelty
The main novelty is a Hutchinson-free likelihood distillation objective for flow-map models, based on conditional divergence matching, together with a vectorized divergence target that predicts Jacobian diagonal terms rather than only a scalar divergence. This extends few-step flow-map generation to more scalable joint sample-and-likelihood estimation while avoiding the usual stochastic trace-estimation bottleneck.
Results
Across alanine molecular systems, SCALLOP consistently improves over F2D2 on the reported metrics and is competitive with strong Boltzmann-generator baselines, while achieving the best torsional-angle Wasserstein scores on ALA-3 and ALA-6. The paper reports substantially lower training variance, faster convergence, about 10% faster training than F2D2, and much faster inference than Jacobian-based baselines, including up to 10× speedup over the fastest baseline and 100× over FALCON in molecular settings. On CelebA-64, SCALLOP achieves sample quality and likelihood metrics comparable to F2D2 while consistently reducing per-sample BPD error and improving training stability.
Key Points
- SCALLOP replaces Hutchinson-based divergence estimation in likelihood flow-map training with conditional divergence matching, reducing variance and removing a key training bottleneck.
- The method uses a vectorized objective that learns Jacobian diagonal entries, enabling efficient reconstruction of divergence and supporting a shared few-step framework for sampling and likelihood evaluation.
- Empirically, SCALLOP improves on F2D2 in molecular and image experiments, remains competitive with state-of-the-art Boltzmann generators, and offers markedly better runtime characteristics for likelihood-aware inference.