Few-Step Boltzmann Generators via Scalable Likelihood Flow Maps

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.

• The paper introduces SCALLOP, a scalable likelihood distillation method that eliminates high-variance Hutchinson estimators through conditional divergence matching.
• The paper demonstrates up to 10x inference speedup, improved training stability, and enhanced effective sample size across both molecular and high-dimensional image benchmarks.
• The paper unifies likelihood estimation and sample generation via a shared SE(3)-symmetric architecture, enabling efficient, density-aware generative modeling in complex domains.

Scalable Few-Step Likelihood Estimation via SCALLOP

Introduction and Motivation

The development of expressive, computationally tractable generative models equipped with accurate likelihood estimates remains a central challenge for applications spanning molecular modeling and high-dimensional data generation. While recent advances in continuous-time generative models, such as CNFs and distilled flow maps, have enabled rapid sample generation, the absence of scalable, low-variance likelihood estimation persists as the main obstacle to deployment in scientific domains where density estimates are essential (e.g., for free energy differences and importance sampling). Traditional approaches rely on either constrained normalizing flow architectures—limiting flexibility for complex data—or high-variance stochastic divergence estimators (e.g., Hutchinson's estimator), rendering many models impractical for large-scale systems or real-world scientific applications.

Methodology and Theoretical Contributions

The paper introduces SCALLOP, a scalable likelihood distillation method for flow maps that obviates the need for Hutchinson’s trace estimator. SCALLOP builds on the F2D2 likelihood flow map framework [Ai et al., 2026], but introduces a conditional divergence matching objective derived using the chain rule and the continuity equation. This objective enables direct, divergence-free estimation of the log-density change along learned flows, circumventing the main variance and computational costs associated with Hessian- or stochastic-trace-based approaches.

The essential innovation is the Conditional Divergence Matching (CDM) loss, formulated for coupled stochastic interpolant paths, which reduces to exact divergence matching under appropriate marginalization. The paper further proposes a vectorized form (CDM-v), regressing the full diagonal of the flow’s Jacobian, thereby leveraging a higher-bandwidth training signal and promoting data efficiency in high dimensions. The implementation amortizes flow and likelihood prediction into a shared, SE(3)-symmetric neural architecture, employing DiT transformers for both molecular and image data.

Empirical Results

Molecular Systems: Boltzmann Generator Regime

Empirical validation on alanine peptide systems (ALA-2 to ALA-6) demonstrates that SCALLOP consistently achieves test set likelihoods, effective sample size (ESS), and torsional angle Wasserstein-2 distances (T-W2) that are superior or comparable to state-of-the-art flow and importance-sampling-based Boltzmann Generators such as FALCON [Rehman et al., 2026] and SBG [Tan et al., 2026a]. Of particular note:

• Up to 10x Inference Speedup: Likelihood evaluation is realized via a single forward pass, in contrast to FALCON, which incurs cubic complexity in data dimension.
• Greater Training Stability and Data Efficiency: The variance of the training signal for likelihood estimation is reduced by up to 100x compared to F2D2, and wall-clock training converges faster.
• Global Conformation Fidelity: SCALLOP attains the lowest T-W2 on ALA-3 and ALA-6, confirming effective recovery of macroscopic metastable states.
• ESS and Mode Robustness: Consistently high ESS across tasks reflects reduced mode collapse and higher overlap with target Boltzmann densities.

High-Dimensional Images: Consistency Testing

On CelebA-64, SCALLOP matches or exceeds F2D2 in bits-per-dimension (BPD) for negative log-likelihood and FID for sample quality at all numbers of Euler evaluation steps (1, 2, 4, 8). The per-sample BPD error is systematically lower, and training remains more stable, with gradient norms and loss variance markedly reduced. Loss computation consumes only 83% of the wall-clock time required by F2D2.

Practical and Theoretical Implications

SCALLOP provides a practical route to scalable, likelihood-equipped, few-step generative modeling in both molecular and high-dimensional image domains. The elimination of the Hutchinson estimator’s variance and the integration of vectorized objectives are immediately impactful for large-scale scientific modeling, where data dimension and physical symmetries must be accommodated. By consolidating likelihood and sample generation into a unified, highly parallelizable architecture, SCALLOP resolves the typical inference-time bottlenecks encountered by flow-based Boltzmann generators and supports applications such as SNIS-based free energy estimation.

The theoretical insight of conditional divergence matching prompts a re-examination of training objectives for generative ODEs, potentially informing new approaches for trace-free, higher-order score matching and density-ratio learning.

Limitations and Future Directions

The framework assumes that the learned velocity field closely tracks the ideal transport; the impact of model misspecification and resulting bias in likelihood estimates remains to be quantified. While generation in 1 step is possible, accurate modeling of high-dimensional systems typically requires about 10 steps. Further, global metric improvements over FALCON are not universal—energy-based Wasserstein distances still evidence outlier sensitivity—posing open questions regarding the trade-off between local and global structure fidelity.

Future work should assess the effect of velocity learning bias, extend SCALLOP to stochastic interpolant generalizations, and evaluate integration with advanced SMC-based samplers for enhanced molecular modeling.

Conclusion

SCALLOP establishes a new computational and statistical regime for flow-based generative modeling by enabling efficient, stable, and scalable estimation of both samples and their likelihoods in a few steps. Its adoption will facilitate equilibrium sampling and density-aware generation in high-dimensional scientific and machine learning applications that were previously dominated by the computational burden of likelihood evaluation (2606.29110).