FALCON: Few-step Accurate Likelihoods for Continuous Flows

Abstract: Scalable sampling of molecular states in thermodynamic equilibrium is a long-standing challenge in statistical physics. Boltzmann Generators tackle this problem by pairing a generative model, capable of exact likelihood computation, with importance sampling to obtain consistent samples under the target distribution. Current Boltzmann Generators primarily use continuous normalizing flows (CNFs) trained with flow matching for efficient training of powerful models. However, likelihood calculation for these models is extremely costly, requiring thousands of function evaluations per sample, severely limiting their adoption. In this work, we propose Few-step Accurate Likelihoods for Continuous Flows (FALCON), a method which allows for few-step sampling with a likelihood accurate enough for importance sampling applications by introducing a hybrid training objective that encourages invertibility. We show FALCON outperforms state-of-the-art normalizing flow models for molecular Boltzmann sampling and is two orders of magnitude faster than the equivalently performing CNF model.

• The paper introduces FALCON, a method for amortized Boltzmann sampling using few-step invertible flow maps that cut computational cost by up to 100x compared to traditional CNFs.
• The method employs a hybrid loss combining regression and cycle-consistency to ensure both sample quality and numerical invertibility for efficient likelihood computation.
• Empirical evaluations demonstrate FALCON’s superior performance in capturing molecular energy distributions and peptide conformational diversity over existing discrete normalizing flows.

Few-step Accurate Likelihoods for Continuous Flows: FALCON

Introduction and Context

Boltzmann Generators (BGs) offer an efficient strategy for sampling from molecular Boltzmann distributions—a core challenge in computational statistical physics underlying structure-function elucidation, free energy computation, and molecular discovery. Traditional BG frameworks leverage continuous normalizing flows (CNFs), which are invertible, expressive, and enable exact likelihood computation. However, the computational cost associated with CNF-based likelihood evaluation is prohibitive, often requiring hundreds to thousands of ODE function evaluations per sample due to discretization and Jacobian trace computations. This cost severely impedes the scalability of BGs to high-dimensional molecular systems.

This paper proposes FALCON (Few-step Accurate Likelihoods for Continuous Flows) (2512.09914), a method for amortized Boltzmann sampling that achieves competitive or superior sample quality in two orders of magnitude fewer steps compared to conventional CNFs. FALCON achieves this by leveraging a hybrid loss that directly regularizes for both regression performance and invertibility in few-step discrete flow maps, facilitating both efficient sample generation and tractable likelihoods—properties necessary for downstream importance weighting in physical observables estimation.

Methodology: Few-step Invertible Flow Maps

FALCON is built on the insight that while flow-matching and consistency-model approaches achieve compelling one- or few-step generative performance, they do not guarantee invertible maps or accurate likelihoods prior to optimum convergence, which is generally unattainable in high dimensions with limited data. The method introduces a hybrid loss:

• A regression loss targeting accurate velocity estimation for sample quality.
• A cycle-consistency (invertibility) loss penalizing deviations from numerical invertibility even far from convergence.

Formally, FALCON parameterizes a flow map $X_u(x_s, s, t) = x_s + (t-s) u_\theta(x_s, s, t)$ where $u_\theta$ learns the average velocity along the flow. The cycle-consistency loss

$$\mathcal{L}_{\text{inv}} = \mathbb{E}_{s, t, x_s} \| x_s - X_u(X_u(x_s, s, t), t, s)\|^2$$

encourages invertibility of $X_u$ (via its own parameterization) across arbitrary start and end points $(s,t)$. Under mild conditions, minimizing this loss ensures invertibility everywhere and guarantees that the log density at any time point can be computed efficiently using the change-of-variables formula—a property crucial for functionality as a Boltzmann Generator.

Notably, FALCON relaxes the requirement that the learned flow map precisely follows the underlying continuous-time ODE solutions of the reference vector field $v$. Instead, it suffices for $X_u$ to be invertible, affording much greater architectural freedom and enabling the scaling of model complexity.

Empirical Results

Scalability and Performance

FALCON demonstrates substantial gains in computational efficiency and sample quality across standard molecular benchmarks: alanine dipeptide, tri-alanine, alanine tetrapeptide, and hexa-alanine. The results highlight two central claims:

1. Superior scalability over CNFs: FALCON achieves comparable or better sampling quality in as few as 4–16 steps, compared to 200–270 steps required by CNFs, yielding a speedup of approximately 100x. This scalability enables the use of more expressive architectures such as attention-based DiT backbones previously untenable under CNF bottlenecks.
2. Improved sample quality over discrete NFs: Even one-step discrete normalizing flows (e.g., SBG, TarFlow) cannot match the global and local metrics achieved by FALCON, even when allowed 250-fold higher evaluation budgets. FALCON’s reweighted samples more closely match reference energy and torsion-angle statistics.

These improvements are clearly reflected in the empirical distribution of sampled energies before and after importance reweighting, accurately covering the support of the target molecular dynamics (MD) reference for increasing system complexity.

Figure 1: True MD energy distribution with best FALCON unweighted and re-sampled proposals for alanine dipeptide (left), tri-alanine (center left), alanine tetrapeptide (center right), and hexa-alanine (right).

Likelihood-Quality versus Step Count

The impact of the number of inference steps on the closeness of sample energies to the ground-truth distribution is quantified. FALCON provides a consistent interpolation between accuracy and efficiency, with empirical evidence that even with minimal steps, reweighted proposals approach the correct Boltzmann energy statistics necessary for practical applications.

Figure 2: Improved proposal and re-weighted sample energies with increased steps for alanine dipeptide.

Structural and Angular Accuracy

FALCON captures the critical conformational diversity in peptide torsions, as established via Ramachandran plots on held-out test data, showing accurate recovery even for modes under-sampled in the training set.

Figure 3: Left: Training data for alanine dipeptide; Right: Test data for alanine dipeptide.

Figure 4: Left: Test data for alanine dipeptide; Right: FALCON's angular predictions for alanine dipeptide.

Analogous results extend to larger peptides (Figures 4–6, 8–10), verifying FALCON’s robustness against data bias and its capacity for unbiased mode recovery and correct weighting.

Ablations, Invertibility, and Inference Scheduling

The invertibility-promoting regularization is essential for balancing sample quality with precise likelihood computation. Extensive ablations show that insufficient regularization leads to degenerate flows, while excessive regularization impairs generative quality.

Inference scheduling (e.g., EDM, cosine, Chebyshev) significantly affects performance in the few-step regime, with empirical results favoring EDM-style step placement—consistent with recent findings in diffusion models.

Theoretical and Practical Implications

FALCON’s theoretical guarantee—the ability to compute exact likelihoods for invertible maps—positions it as a legitimate alternative to ODE-based CNFs in importance sampling, with much-reduced computational overhead. The relaxation from needing the true flow map to any invertible map enables flexible architectural choices and scaling.

Practically, this unlocks new capabilities for scalable BGs in physical sciences, facilitating accurate, efficient estimation of thermodynamic observables and rare-event modeling in high dimensions. The methods extend readily to any generative tasks requiring both fast sample generation and tractable likelihoods, including Bayesian inference, variational learning, and scientific simulation.

The paper also notes the open problem of efficient, guaranteed one-step invertible flow learning—a prospect for further AI research in flows, likelihood-based modeling, and scientific generative models.

Conclusion

FALCON delivers a significant advance in amortized Boltzmann sampling by combining efficient few-step flow maps with invertibility-inducing losses, thereby dramatically reducing the inference cost of likelihood-based generative models without sacrificing sample quality. By bridging regression-driven sample efficiency and the exactness required for scientific applications, this approach lays a foundation for scalable, accurate molecular modeling, and suggests promising directions for further AI development in fast-likelihood generative modeling (2512.09914).