FALCON: Few-step Accurate Likelihoods for Continuous Flows

AI summary

Overview

Research area: Machine learning for scientific computing — specifically Boltzmann Generators (BGs) for molecular equilibrium sampling, continuous normalizing flows, flow matching, and few-step generative models.

Technical level: Advanced. The paper assumes familiarity with normalizing flows, the instantaneous change-of-variables formula, the Jacobian trace estimator, importance sampling, and molecular statistical mechanics.

Scope: The paper introduces FALCON, a hybrid-trained flow map that makes few-step molecular sampling invertible so that likelihoods are cheap and accurate enough for self-normalized importance sampling.

What This Paper Is About

Sampling molecular configurations from a Boltzmann distribution is hard because classical methods (molecular dynamics, MCMC) get trapped in local energy minima. Boltzmann Generators solve this by training a generative model on biased data and then correcting its samples with self-normalized importance sampling (SNIS). The catch is that the best-performing generators, continuous normalizing flows, need thousands of function evaluations to compute a likelihood accurate enough for that correction. FALCON's goal is to keep the flexible, simulation-free training of flow models while producing likelihoods from only a handful of steps.

Key Contributions

1. FALCON, a continuous flow-based generative model for Boltzmann sampling that is invertible, trainable with a regression loss, supports free-form architectures, and enables both few-step generation and efficient likelihood evaluation. It uses a hybrid objective combining a flow-matching loss, an average-velocity regression loss, and a cycle-consistency (invertibility) regularizer.
2. A softly equivariant continuous flow architecture — a standard diffusion transformer (DiT) with an added time-embedding head, plus data augmentation for soft SO(3) rotation equivariance and mean subtraction for translation invariance.
3. A two-orders-of-magnitude speedup over CNF-based Boltzmann Generators for equivalent performance.
4. State-of-the-art results over existing normalizing-flow Boltzmann Generators across all reported metrics, even when FALCON is given 250× fewer samples than the SBG baseline.

Main Findings

• Few-step likelihoods become usable. FALCON's average-velocity flow map is not guaranteed invertible before its training objective is exactly minimized, so the authors add an invertibility loss. Proposition 2 shows that minimizing this loss alone is sufficient for a valid change-of-variables computation, even without exactly reproducing the continuous-time flow.

• Invertibility is numerically verified. The authors train an auxiliary network to invert a frozen FALCON flow and report invertibility errors on the order of 10⁻⁴, the same order of magnitude as discrete and continuous NFs.

• Scalability beyond alanine dipeptide. ECNF++ performs exceptionally well on ESS and 𝕋-𝒲₂ for alanine dipeptide but fails to scale to larger molecules. On tri-alanine, alanine tetrapeptide, and hexa-alanine, FALCON substantially outperforms ECNF++ across all metrics. For example, on hexa-alanine FALCON reports ESS 0.060 ± 0.017, ℰ-𝒲₂ 0.892 ± 0.311, and 𝕋-𝒲₂ 1.256 ± 0.132, versus ECNF++ at ESS 0.006 ± 0.001, ℰ-𝒲₂ 10.668 ± 0.285, and 𝕋-𝒲₂ 1.902 ± 0.055.

• On alanine dipeptide, FALCON is not best on every metric. FALCON reports the lowest ℰ-𝒲₂ among the listed baselines at 0.225 ± 0.104, but its ESS is 0.067 ± 0.013, below ECNF++ at 0.275 ± 0.010, and its 𝕋-𝒲₂ is 0.402 ± 0.021, above both ECNF++ (0.189 ± 0.019) and FALCON-A (0.180 ± 0.005).

• Beating discrete flows on sample quality. SBG outperforms all previously reported methods on the larger systems, but FALCON outperforms SBG across all metrics in a few steps. On hexa-alanine, SBG IS reports ESS 0.034 ± 0.015, ℰ-𝒲₂ 1.021 ± 0.239, and 𝕋-𝒲₂ 1.431 ± 0.085.

• More samples only partly close the gap. Even with 5×10⁶ samples — 250× more than those used to evaluate FALCON — SBG's ℰ-𝒲₂ remains significantly worse than that of a 4-step FALCON flow.

• Cumulative training + inference time. Over 10⁴ samples on one NVIDIA L40S with batch size 1024, FALCON reports 7.65 (alanine dipeptide), 11.45 (tri-alanine), 18.84 (alanine tetrapeptide), and 25.76 (hexa-alanine), compared to ECNF++ at 12.52, 19.59, 32.17, and 137.4, SBG at 16.83, 24.67, 41.67, and 57.50, and a DiT CNF at 9.56, 17.54, 24.10, and 82.10.

• Adaptive-step FALCON trades speed for accuracy. A Dopri5-integrated FALCON achieves better metrics at 257 (ALDP), 265 (AL3), 200 (AL4), and 207 (AL6) function evaluations, while the few-step FALCON uses 4, 8, 8, and 16 evaluations respectively — and even in that few-step regime it still outperforms every method considered on the larger systems.

• Regularization trade-off. Weak invertibility regularization leads to poor invertibility and degraded performance; strong regularization enforces invertibility at the cost of reduced sample quality. The regularization constant is fixed to 10.0 for all experiments unless stated otherwise.

• Scheduler choice matters. In an 8-step alanine dipeptide ablation, the EDM scheduler substantially outperforms all other schedulers, consistent with observations in the diffusion literature, and is used for all reported results.

• Loss choice. FALCON is evaluated with a MeanFlow-based loss and an additivity loss (FALCON-A). The MeanFlow-based loss is the default because it performed best or second best on all metrics for the larger systems.

Methodology in Plain English

1. Frame the target. Learn a generative model whose samples can be re-weighted into the true Boltzmann distribution. The re-weighting (SNIS) requires evaluating the model density p_θ(x) for every sample, so that density must be both cheap and accurate.

2. Diagnose the problem with flow maps. Existing few-step flow models learn an average velocity that transports samples over a time interval. At the exact optimum, this map is invertible (Proposition 1), but in practice it does not match the true continuous flow, so invertibility — and therefore the standard change-of-variables formula — is not guaranteed.

3. Relax the requirement. The authors observe they do not need the flow map to equal the continuous flow; they only need it to be invertible. That yields a cycle-consistency regularizer: map forward from time s to t, map back from t to s, and penalize the squared difference from the starting point.

4. Combine three losses. A flow-matching term, an average-velocity regression term (implemented efficiently with a single Jacobian-vector product via forward-mode autodiff), and the invertibility term. A sign parameterization u_θ(x_s, s, t) = sign(t − s) h_θ(x_s, s, t) handles the discontinuity when s = t, since the model is the first to require flow maps in both the forward and backward directions.

5. Scale the architecture. Because few-step inference is cheap, the authors can use a diffusion transformer with a time-embedding head instead of the small equivariant networks used by prior BGs, which are described as being up to 2.3 million parameters.

6. Evaluate as a Boltzmann Generator. Train on biased data from implicit-solvent MD with the amber-14 force field and test on a held-out unbiased dataset, using SNIS and force-field energies for log-likelihoods and re-sampling. Compare against four discrete flow baselines (SE(3)-EACF, RegFlow, SBG IS, SBG SMC) and three continuous flows (ECNF, ECNF++, BoltzNCE), with all continuous flows integrated using Dormand–Prince 4(5) at atol = rtol = 10⁻⁵. Report ESS, ℰ-𝒲₂, and 𝕋-𝒲₂ over 10⁴ samples on three seeds as mean ± standard deviation.

Why This Matters

Impact on research. Continuous normalizing flows are expressive and stable to train but are hamstrung by likelihood costs that require many integration steps. FALCON shows that invertibility can be encouraged directly by the training objective rather than assumed at convergence, which lets researchers keep free-form, scalable architectures while retaining the exact-likelihood requirement that makes importance sampling statistically consistent.

Real-world applications:

• Characterizing biomolecular function by sampling equilibrium conformations of peptides and proteins.
• Accelerating drug design through faster exploration of molecular configurations.
• Discovering novel materials where equilibrium sampling underpins observable properties.
• Computing thermodynamic observables such as free energies and heat capacities, which the paper identifies as downstream quantities governed by Boltzmann sampling.

Industry relevance. Inference cost is often the binding constraint on deploying Boltzmann Generators. A method that reaches equivalent performance with two orders of magnitude less inference time, and that is faster than discrete normalizing flows on cumulative training plus inference, lowers the compute barrier for large-scale molecular sampling workflows. The 250× sample-efficiency advantage over SBG is relevant where generating large sample sets is itself expensive.

Future Directions

• Closing the remaining metric gaps. FALCON is not the best method on ESS or 𝕋-𝒲₂ for alanine dipeptide, and its E-W2 degrades relative to FALCON-Dopri5 on the larger systems; understanding and narrowing these gaps is an open question.
• Tuning the invertibility regularization. The paper fixes the regularization constant to 10.0 and shows a clear trade-off between invertibility and sample quality. A principled way to schedule or adapt this weight is not resolved.
• Scheduler design for few-step regimes. The EDM scheduler was chosen after an 8-step ablation on a single system, and the paper notes that variance of the flow field is higher near the target distribution; extending this analysis to other systems and step counts remains open.
• Scaling to larger systems and architectures. FALCON's cheap inference is what makes the DiT architecture feasible. Whether the approach continues to pay off at protein-scale systems, and what larger or more expressive backbones it enables, is left open.

Target Audience

Researchers and practitioners working on generative modeling for the physical sciences — especially those building Boltzmann Generators, normalizing flows, or flow-matching models for molecular systems. It is also relevant to machine learning researchers interested in few-step generation who need likelihoods rather than samples alone, and to computational chemists or drug-discovery teams weighing the inference cost of equilibrium sampling methods. Readers without a background in continuous normalizing flows and importance sampling will find the mathematical sections demanding.

Authors’ 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.