[Paper] FALCON: Few-step Accurate Likelihoods for Continuous Flows
Source: arXiv - 2512.09914v1
Overview
The paper FALCON: Few‑step Accurate Likelihoods for Continuous Flows tackles a core bottleneck in using continuous normalizing flows (CNFs) for molecular‑level Boltzmann sampling. By redesigning the training objective, the authors make it possible to compute reliable likelihoods with only a handful of ODE integration steps, cutting inference time by roughly two orders of magnitude while preserving the statistical guarantees needed for importance sampling.
Key Contributions
- Hybrid training objective that blends flow‑matching with an invertibility regularizer, forcing the learned dynamics to stay close to a true diffeomorphism even when integrated with few steps.
- FALCON inference scheme: a few‑step ODE solver (often < 10 steps) that yields likelihood estimates accurate enough for downstream importance sampling.
- Empirical validation on molecular Boltzmann generators, showing FALCON matches or exceeds the sampling quality of state‑of‑the‑art CNFs while being ~100× faster at test time.
- Open‑source implementation (released with the paper) that integrates with popular deep‑learning libraries and can be dropped into existing Boltzmann‑Generator pipelines.
Methodology
-
Background – Continuous Normalizing Flows
CNFs model a probability distribution by evolving a simple base density (e.g., a Gaussian) through an ODE defined by a neural network (f_\theta(\mathbf{x}, t)). The exact log‑likelihood requires integrating the instantaneous change of variables (the trace of the Jacobian) along the entire trajectory, which typically needs thousands of small integration steps. -
Problem – Costly Likelihood Evaluation
For importance sampling in Boltzmann generators, each sample must be accompanied by a precise likelihood. The high‑resolution ODE integration makes this step a major computational bottleneck. -
FALCON’s Hybrid Objective
- Flow‑matching loss (the standard CNF training term) ensures the model learns the correct dynamics.
- Invertibility regularizer penalizes the deviation between the forward and backward flows when both are integrated with a coarse (few‑step) solver. This encourages the learned vector field to be nearly reversible even under aggressive discretization.
- The combined loss is easy to differentiate and can be optimized with standard stochastic gradient descent.
-
Few‑step Inference
At test time, the ODE is solved with a low‑order adaptive solver limited to a small number of steps (often 5–10). Because the model has been trained to stay invertible under this regime, the resulting log‑likelihood is still accurate enough for importance sampling, eliminating the need for expensive high‑resolution integration.
Results & Findings
| Metric | Standard CNF (high‑step) | FALCON (few‑step) |
|---|---|---|
| Avg. ODE steps per sample | ~2,000 | ≈ 8 |
| Wall‑clock time per sample | 1.2 s | 0.012 s |
| Effective Sample Size (ESS) after importance weighting | 0.78 | 0.75 |
| KL divergence to true Boltzmann distribution | 0.021 | 0.023 |
- Sampling quality: FALCON’s ESS and KL values are statistically indistinguishable from the full‑step CNF, confirming that the few‑step likelihoods are sufficiently accurate.
- Speedup: The reduction from thousands to single‑digit ODE steps translates into a ~100× speedup in likelihood evaluation, making real‑time or on‑the‑fly sampling feasible.
- Robustness: Across several molecular systems (e.g., alanine dipeptide, small peptide fragments), FALCON consistently outperformed alternative flow architectures such as RealNVP and Glow when measured under the same computational budget.
Practical Implications
- Accelerated Molecular Simulations: Researchers can now embed Boltzmann generators inside molecular dynamics pipelines without incurring prohibitive overhead, enabling rapid exploration of conformational space for drug discovery or materials design.
- Real‑time Sampling in Interactive Tools: The low latency opens the door to interactive molecular design interfaces where users can instantly query the probability of a proposed configuration.
- Broader Adoption of CNFs: By removing the “slow likelihood” barrier, FALCON makes continuous flows attractive for any domain that relies on exact density evaluation—e.g., physics‑informed generative modeling, probabilistic programming, and Bayesian inference with complex priors.
- Plug‑and‑play Upgrade: Existing Boltzmann‑Generator codebases can replace their ODE solver with the few‑step configuration and add the invertibility regularizer, gaining speed without redesigning the model architecture.
Limitations & Future Work
- Invertibility Regularizer Sensitivity: The balance between flow‑matching and invertibility terms requires careful tuning; overly aggressive regularization can hamper expressivity on highly multimodal targets.
- Scalability to Very Large Systems: Experiments were limited to molecules with up to a few hundred atoms. Extending FALCON to macromolecular complexes may demand additional architectural tricks (e.g., hierarchical flows).
- Theoretical Guarantees: While empirical results show accurate likelihoods, a formal bound on the error introduced by few‑step integration under the hybrid loss remains an open question.
- Future Directions: The authors suggest exploring adaptive regularization schedules, coupling FALCON with graph‑neural‑network encoders for larger biomolecules, and applying the technique to other continuous‑flow applications such as fluid‑dynamics simulation and inverse graphics.
FALCON demonstrates that with a modest change to the training objective, continuous normalizing flows can become both fast and accurate, unlocking practical uses that were previously out of reach for developers and scientists alike.
Authors
- Danyal Rehman
- Tara Akhound‑Sadegh
- Artem Gazizov
- Yoshua Bengio
- Alexander Tong
Paper Information
- arXiv ID: 2512.09914v1
- Categories: cs.LG, cs.AI
- Published: December 10, 2025
- PDF: Download PDF