Joint Distillation for Fast Likelihood Evaluation and Sampling in Flow-based Models
Xinyue Ai, Yutong He, Albert Gu, Russ Salakhutdinov, Zico Kolter, Nicholas M. Boffi, Max Simchowitz
Abstract
Log-likelihood evaluation enables important capabilities in generative models, including model comparison, certain fine-tuning objectives, and many downstream applications. Yet paradoxically, some of today's best generative models -- diffusion and flow-based models -- still require hundreds to thousands of neural function evaluations (NFEs) to compute a single likelihood. While recent distillation methods have successfully accelerated sampling to just a few steps, they achieve this at the cost of likelihood tractability: existing approaches either abandon likelihood computation entirely or still require expensive integration over full trajectories. We present fast flow joint distillation (F2D2), a framework that simultaneously reduces the number of NFEs required for both sampling and likelihood evaluation by two orders of magnitude. Our key insight is that in continuous normalizing flows, the coupled ODEs for sampling and likelihood are computed from a shared underlying velocity field, allowing us to jointly distill both the sampling trajectory and cumulative divergence using a single flow map. F2D2 is modular, compatible with existing flow-based few-step sampling models, and requires only an additional divergence prediction head. Experiments demonstrate F2D2's capability of achieving accurate log-likelihood with few-step evaluations while maintaining high sample quality, solving a long-standing computational bottleneck in flow-based generative models. As an application of our approach, we propose a lightweight self-guidance method that enables a 2-step MeanFlow to outperform a 1024 step flow matching model with only a single additional backward NFE.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers3
- FALCON: Few-step Accurate Likelihoods for Continuous FlowsDanyal Rehman, Tara Akhound-Sadegh, Artem Gazizov, Yoshua Bengio et al.ICLR 2026 · 13 citations
- Compositional Planning with Jumpy World ModelsJesse Farebrother, Matteo Pirotta, Andrea Tirinzoni, Marc Bellemare et al.ICML 2026 · 1 citation
- Learning Normalized Energy Models for Linear Inverse ProblemsNicolas M Zilberstein, Santiago Segarra, Eero Simoncelli, Florentin GuthICML 2026
Builds on21
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Training language models to follow instructions with human feedbackLong Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida et al.NeurIPS 2022 · 24,707 citations
- High-Resolution Image Synthesis with Latent Diffusion ModelsRobin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser et al.CVPR 2022 · 13,123 citations
- Direct Preference Optimization: Your Language Model is Secretly a Reward ModelRafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D. Manning et al.NeurIPS 2023 · 10,924 citations
- Improved Denoising Diffusion Probabilistic ModelsAlexander Quinn Nichol, Prafulla DhariwalICML 2021 · 5,234 citations
Related papers
- Simple and Fast Distillation of Diffusion ModelsZhenyu Zhou, Defang Chen, Can Wang, Chun Chen et al.NeurIPS 2024 · 44 citations
- SwiftVideo: A Unified Framework for Few-Step Video Generation Through Trajectory-Distribution AlignmentYanxiao Sun, Jiafu Wu, Yun Cao, Chengming Xu et al.AAAI 2026 · 6 citations
- Categorical Flow MapsDaan Roos, Oscar Davis, Floor Eijkelboom, Michael Bronstein et al.ICML 2026 · 23 citations
- Mean Flow Distillation: Robust and Stable Distillation for Flow Matching ModelsAn Zhao, Shengyuan Zhang, Zhongjian Sun, Yixiang Zhou et al.ICML 2026 · 2 citations
- FastFlow: Accelerating The Generative Flow Matching Models with Bandit InferenceDivya Jyoti Bajpai, Dhruv Bhardwaj, Soumya Roy, Tejas Duseja et al.ICLR 2026 · 3 citations
