Joint Distillation for Fast Likelihood Evaluation and Sampling in Flow-based Models
Authors: Xinyue Ai
Yutong He
Albert Gu
Ruslan Salakhutdinov
J Zico Kolter
Nicholas Matthew Boffi
Max Simchowitz
Venue: ICLR 2026 Poster (2026)
Keywords: flow matching
distillation models
fast likelihood evaluation
fast sampling
generative models
Primary area: generative models
TL;DR: We propose a modular distillation framework for simultaneous fast likelihood evaluation and fast sampling for flow matching models

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.

Published as a conference paper at ICLR 2026
JOINT DISTILLATION FOR FAST LIKELIHOOD EVALU- ATION AND SAMPLING IN FLOW-BASED MODELS
Xinyue - Yutong
Gu', Ruslan Salakhutdinov', J. Zico Kolter',
Nicholas M. Boffi', Max Simchowitz! 'Carnegie Mellon University, "Peking University
yutonghe@cs. cmu. edu, axy253c7estu. edu. cn
ABSTRACT
Log-likelihood evaluation enables important capabilities in generative models, in- cluding model comparison, certain fine-tuning objectives, and many downstream applications. Yet paradoxically, some of today's best generative models diffu- sion and flow-based models - still require hundreds to thousands of neural func- tion 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 like- lihood computation entirely or still require expensive integration over full trajec- tories. We present fast flow joint distillation (F2D2), framework that simulta- neously reduces the number of NFEs required for both sampling and likelihood evaluation by two orders of magnitude. Our key insight that in continuous nor- malizing 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 single flow map. F2D2 is modular, compatible with existing flow-based few-step sampling models, and requires only additional divergence prediction head. Experiments demonstrate F2D2's capa- bility of achieving accurate log-likelihood with few-step evaluations while main- taining high sample quality, solving long-standing computational bottleneck in flow-based generative models. As an application of our approach, we propose a lightweight self-guidance method that enables 2-step McanFlow to outperform a 1024 step flow matching model with only single additional backward NFE.
INTRODUCTION
log-likelihood evaluation and likelihood based inference have long been fundamental to statistical modeling and machine learning, serving the backbone for parameter estimation (Fisher, 1922), model selection (Akaike, 1974), and hypothesis testing (Neyman & Pearson, 1933). the era of generative AI, the ability to efficiently evaluate log-likelihood (log-density) has become even more critical, as it directly enables key post-training techniques including reinforcement learning and pref- erence optimization, where likelihoods man et al., 2017; Ouyang et al., 2022; Rafailov et al., 2023; Shao et al., 2024). Beyond these ap- plications, optimizing log-likelihood also encourages generative models capture all modes of the data distribution, avoiding mode collapse that plagues adversarial approaches (Razavi et al., 2019).
important for methods like PPO, DPO and GRPO (Schul-
While likelihood evaluation is useful for modern generative modeling, the most successful gen- erative models for images and video (Rombach et al., 2022; Black Forest Labs, 2025; OpenAl, 2024; Polyak al., 2024; Google DeepMind, 2025), namely diffusion and flow matching mod- els, suffer from a critical weakness: computing likelihood requires prohibitively expensive iterative neural function evaluations (NFEs). In particular, discrete-time diffusion models like DDPM (Ho et al., 2020; Nichol & Dhariwal, 2021) require summing up variational bounds across all timesteps, which needs hundreds to thousands of forward passes to compute single likelihood. Similarly, continuous- time formulations like score SDE (Song et al., 2020) and flow matching (Lipman et al.,
Equal contribution.
Work done CMU.

---

Published as a conference paper at ICLR 2026
2022: Albergo & Vanden- 2022; Liu et al., 2022; Albergo et al., 2023) also must inte- grate the divergence along the learned (probability) flow ODE, which typically requires numerical integration with 100-1000 NFEs for accurate likelihood evaluation. While advanced solvers can significantly reduce NFEs (Karras et al., 2022), they fundamentally cannot escape the integration requirement and produce vastly inaccurate results when restricted to very few steps (s 10 NFEs). This computational burden makes many likelihood-based finetuming objectives, model comparison, and downstream applications prohibitively expensive for modern diffusion/flow matching models.
Interestingly, diffusion and flow matching models faced the same NFE bottleneck for sampling when they were first introduced, where they initially required 1000+ steps to generate a single image. Research addressing this issue has been remarkably successful, with methods that learn to skip multiple steps, either through distillation or self-consistency objectives, emerging as particularly powerful solutions (Salimans & Ho, 2022; Song et al., 2023; Kim et al., 2023; Frans et al., 2024; Geng et al., 2025; Boffi et al., 2025b;a). However, despite achieving few-step sampling, most of these methods completely lose the ability to compute likelihoods, and the few methods that preserve likelihood computation (e.g. Kim et al. (2023)) still require integrating over the entire trajectory with hundreds of NFEs, making fast likelihood evaluation impossible. Thus, while practitioners have solutions to fast sampling for diffusion and flow matching models, fast log- lig-likllihood evaluation remains an unsolved problem. Here, we show that it is possible to achieve both at the same time.
Contributions. Our key insight is that, in flow matching and continuous normalizing flows (CNFs) in general (Chen et al., 2018), computing exact likelihoods requires solving coupled ODEs: the sampling trajectory ar = uy (zx, t) and the log-density evolution ar both depending on the same learned velocity up. Since the divergence term can be viewed as another output derived from the same velocity model, we can learn to distill both the flow trajectory and its corresponding divergence computation simultaneously within a single model. By jointly optimizing for both accurate few-step sampling and log- logllikld evaluation, we can potentially achieve a model that succeeds at both tasks.
4
d Pt = div (1) t)),
Based on these insights, we propose fast flow joint distillation (F2D2), simple and modular frame- work for jointly learning fast sampling and fast log-likelihood evaluation in flow-based models. Our key idea is to leverage the flow map framework (Boffi et al., 2025a) to train single model to pre- dict both the sampling trajectory and cumulative divergence parallel using a joint self-distillation objective, ensuring both outputs learn to skip the numerous steps in training. This makes F2D2 plug-and and-play with any CNF-based few-step sampling method like shortcut models and McanFlow, and requires only a new divergence prediction head alongside the existing velocity prediction. To our knowledge, F2D2 is the first method to enable accurate few-step log-likelihood evaluation in diffussion/CNF-based generative models, solving long-standing limitation of these frameworks.
We demonstrate that our method produces both calibrated likelihoods and high quality samples with few-step NFEs on image datasets CIFAR-IO (Krizhevsky et al.) and ImageNet 64 64 (Deng al., 2009). We show that our F2D2 are compatible with and can be directly apply to pre-trained shortcut models, MeanFlow and a new distillation method we propose in this paper.
As an application of our method, we introduce maximum likelihood self- self-guidance, a lightweight test-time intervention that uses rapid likelihood evaluation to optimize over generated samples, re- quiring only an additional forward and backward pass through the model. Remarkably, we show that F2D2 with maximum likelihood self-guidance instantiated with 2-step MeanFlow achieves lower FID than 1024-step flow matching model of the same size on CIFAR-10, This proof of concept demonstrates the expanded algorithmic sandbox enabled by rapid likelihood evaluation.
2 BACKGROUND
Let Pdata denote the data distribution with samples Rdd We consider a time variable t E [0, 1] where f = 0 corresponds to a simple noise distribution Po = ~N 1) and t = 1 corresponds to the data distribution p1 Paat- We denote the marginal distribution at time t as (x)).
Flow Matching. Flow matching (Lipman et al., 2022; Albergo & Vanden- -ijjnden,, 2022; Liu et al., 2022: Albergo et al., 2023) is scalable training method for generative modeling that learns a time- dependent velocity field vu Rd x [0, 1] + Rd to transport samples from a simple noise distribution Po the data distribution P1- Along straight-line path x,, (1 (1-t)ar + lxy that linearly interpolates
2

---

Published as a conference paper at ICLR 2026
between noise sample zo Pu and a data sample æ1 P1, it models the evolution dynamic with the ordinary differential equation ODE:
Rt (0 f),
Pa
(2.1)
dt
where 8, arises from integrating the learned flow model. Since the velocity along this path is simply I1 Zo. flow matching minimizes the regression objective:
FM (0)
(2e,t) (æ1 x=)]l
(2.2)
This objective encompasses diffusion models as special case with different interpolation schemes (Gao et al., 2024). For sampling, solves the ODE from 0 to t -1 with numeri- cal solvers like Euler or dopri5, and typically requires 1000 NFEs for high quality samples.
Continuous Normalizing Flows and Likelihood Computation. A dow-matching model is spe- cial case of a continuous normalizing flow (CNF) (Chen et al., 2018), which transports data from an initial distribution To ~ Po to an estimated distribution by integrating an ODE. In the case of flow-based models, this is precisely Eq. (2.1), An advantage of the CNF formalism is the ability to explicitly compute likelihoods via the coupled system of ODEs:
d dt log [oggr;((Ee)
it
[
Up (2,,
(2.3)
div (va (2,, t)).
Above, (1,, = Tr(Va & Up (2,, t)) denotes the divergence of the velocity field Up, and pre represents the likelihood of St under Eq. (2.1), which note depends model parameter via un- Integrating backwards from = (data) to / = (noise) with initial conditions (ar1 OJT, we obtain:
lo (x11) = (00)) + J," div(vw d) = log polgoo(o)))) - f. div(va ())dl,
(2.4)
where do and the intermediate 2; are obtained by integrating the flow backward from ø1.
Likelihood evaluation is typically expensive, requiring both careful, finely-discretized integration of an ODE across time steps, and a computation of the divergence term whose exact computation (or the variance of its randomized estimator (Grathwohl et al., 2018)) scales at least linearly in ambient dimension d. Thus, likelihood evaluation is far more computationally burdensome than sampling.
Few-Step Flow-based Models. To address the computational expense of multiple ODE integrations in sampling, recent few-step flow-based models (Kim et al., 2023; Frans al., 2024; Geng al., 2025; Boffi et al., 2025a:b) learn to directly predict the outcome of integrating the ODE in Eq. (2.1) using only small number function evaluations (NFEs). These methods can be viewed as sharing a common strategy of learning to predict the flow map of the underlying ODE.
Definition 2.1 (Flow Map). Given ODE dæ, v( l)dt, the flow map $ Rd [0, 1] Rd is the solution operator that maps any state at time to its corresponding state at time s:
① (IIr,t s) = xz J. o(xr, r)dr Fs
(2.5)
After learning the flow map with network parameter 0, one can directly perform few-step sampling by discretizing the time interval [o, 1] into K steps with timesteps 0 to t1 ... tx 1, and iteratively applying the learned flow map: 6t,,, = 4e (ā, t4, 1 for = D,...,K K 1, starting from To Po- This reduces sampling from hundreds of ODE solver steps to just K NFEs (typically K < 10), as each application of directly predicts the integrated result over the interval [f;, 4441 without explicit numerical integration.
METHOD
We propose to jointly accelerate both sampling and likelihood evaluation by learning flow-map on the joint ODE system described in in Eq. (2.3). Again, Pu = (O, is the source distribution, P1 Pdata- Pt represents the marginal distribution of the interpolant 2; tæ, (1-t)) (1-t)00, v denotes the ground truth velocity and Peo is the distribution of &, under the learned flow model Eq. (2.1). Our aim is to design model which supports two key capabilities:
3

---

Published as a conference paper at ICLR 2026
1. Fast sampling: Draw 21 from trained flow model, using few number of NFEs, Ksamp < 10.
2. Fast likelihood evaluation: Evaluate the log likelihood of either model samples 1 or data samples I1 using few number of NFEs, Kn 10.
3.1 FAST FLOW JOINT DISTILLATION (F2D2): PARAMETRIZING JOINT FLOW MAP
Our key insight is that we can apply few-step flow-based models to the ODE in Eq. (2.3) which jointly parametrizes sampling and likelihood evaluation. Following Boffi al. (2025a); Frans al. (2024); Geng et al. (2025), we adopt a linear parametrization of the flow map:
(&,, Do(fe, t, s) = ân + (s () t, t)
(3.1)
where uo RRd x [0, 1] - RRd predicts the average velocity that directly transports states from time t time and ideally (xr,, -he Se olz,, T). With this parametrization, we recovers an estimate of the instantaneous velocity as ug (2,, t, t) the limit, and obtain simple conditions for valid flow maps.
Proposition 3.1 (Flow Map Conditions (Boffi al., 2025a)). operator Φ(x, t, 8) x + t)u(a, i, s) is valid flow map if and only if for all (t, 8) E [0, 112 and for all E Rd, u(x,t,t) oxx t) and any of the following conditions holds:
(a) & solves the Lagrangian equation 0,4 f, 8) = ((
(b) & solves the Eulerian equation 0,0(z, 8) V s)u 1,1) = 0.
(c) & satisfies the semigroup property ( (t),, = 4 s) <rs..
Let 2, log (1:4) R and & log Pt:o ( RR denote the log likelihood, parametrize the flow maps for the two subsystems in Eq. (2.3) as
then separately
x (Ge, 8) f; (s t)uex t,, 2 (fr, 6) 2, + (s D) (it,
(3.2)
Here uax (,2,4,8) still estimates the average velocity, and Døz (Ire (, approximates the average divergence Do (It, 1, - -
div(v(x,, along the true trajectory between i and s.
Notice that average divergence depends only on It, not zt. The fact that d, is sufficient in our parametrization follows from the joint ODE Eq. (2.3), where the evolution of the likelihood 2 is determined by the divergence of the first flow evaluated at 2t.
Therefore, denoting the joint state at time as yt (z1, 2 we can then parametrize the joint flow map using shared parameter 0 as
T x 8) (ãt, t,,
= Gt (s t)fo(ār,t,))
(3.3)
[o t, (8r, 1,))
1,
Above, the networks for ug and Dg share the same backbone with separate prediction heads for their respective components. The exact architecture is described in Appendix C.
Theoretical Justification. To justify the parameterization Eq. (3.3), we recall a property denoted the tangent condition by Boffi al. (2025a), leveraged by Kim et al. (2023); Geng et al. (2025) to recover the instantaneous velocity and divergence the s > t limit:
Lemma 3.2 (Tangent Condition). The flow map 4y (y, t, for the joint system Eq. (2.3) satisfies lim;- 0, 0y 1,8) = [(x,, 8, = (v(2, (), -div(v( T)))
We can then characterize valid joint flow maps under parametrization as the following:
Proposition 3.3 (Characterization of the Joint Flow Map). Let 4y (y.t,s) = + (s fxf t, s) denotes the dynamics for the joint sampling and likelihood system Eq. (2.3). Then, ly t, is the flow map for the joint system if and only if A 1,) Rd++1 [0, 1J2, any of the following conditions are satisfied:
satiisy f(z, s) = (v(2, t), -iv(vu(z,,t)))
4

---

Published as a conference paper at ICLR 2026
Algorithm F2D2 Flow Map Training
1: for each training step do
2:
31 ~ Pritta Ti ~ po- (1, ~ ((00,
forward-only training uses / S
3: Tt (1- 1)=r + Lx1
4: Calculate the sampling (velocity) flow map losses Lvx (0) and (.)
5:
Calculate the likelihood (divergence) flow map losses Laiv (0) and CD
6: L0) == (0) + (i) + Laiv + (0)
7: Update w.r.t. (0)
8: end for
9: return
(a) By solves the Lagrangian equation 0,0y (y, (
(b) + solves the Eulerian equation 0 (v, 8) + /
V ,, 8) (c) Φ satisfies the semigroup property (
4 r),
We provide the full analysis of the characterization in Appendix A. We refer to the family of al- gorithms which learns joint map of this characterization as fast flow joint distillation (F2D2). Notably, we can derive four separate training objectives one pair for the sampling subsystem and the other pair for the likelihood subsystem and jointly optimizing them yields a valid flow map for Eq. (2.3). The general F2D2 training objective is:
LFID2 VVM (0) (0) Laiv (0) (D)
(3.4)
where the first two terms optimize for the sampling flow map xx: : Lvw enforces the instantaneous velocity matching (i.e. the tangent condition, which is often enforced by the flow matching loss in practice), while Lu enforces one of the flow map conditions from Proposition 3.1 for Φx Similarly, the last two terms optimize for the likelihood flow map &z: Ldiv matches the instantaneous diver- gence, and LD ensures that 2 satisfies the conditions needed for the joint flow map 4y to be valid according to the conditions Proposition 3.3. We provide the pseudocode for the general recipe of F2D2 training in Algorithm 1.
3.2 INSTANTIATING F2D2 WITH SHORTCUT AND MEANFLOW
Though our method is, in principle, compatible with any flow map-based method, we instantiate our formulation for Shortcut Models (Frans et al., 2024), based on the semigroup property, and for MeanFlow (Geng et al., 2025), based on the Fulerian equation. We also provide the Lagrangian self- distillation (LSD) (Boffi et al., 2025a) instantiation based on the Lagrangian equation in Appendix A. Pscudocode for each instantiation is provided in Appendix C.
3.2.1 JOINT SHORTCUT: SHORTCUT-F2D2
Shortcut models (Frans et al., 2024) enforce the semigroup property (Proposition 3.1 (c)) by applying to the midpoint between timesteps and 8. This amounts to the shortcut self-coosistency loss:
Cu-se (0) Er,,,,
[|ue(xr, s) sg (ue (x ł, r) + up ( (
(3.5)
where (+) denotes stop- coop-gradeent. Combined with the tangent condition (at t 8),
v2 (x,
(3.6)
and training with the flow matching loss as LVM.
Lvm =o=1 [l] [[luo t,t) (x1 =o))l
(3.7)
this self-consistency loss Lu-se serves as Cu and yields a valid flow map by enforcing both the tangent condition and the semigroup property (see Corollary B. in Appendix B for details).
We convert this method to a joint self-distillation method by introducing the two additional losses: Laiv-se (0) = E [o) div(u, (,t)
(3.8)
1
LD Eics
1o(
sg (Do(2,,t,)) + Do(
(3.9)
5

---

Published as a conference paper at ICLR 2026
The first loss proceeds by analogy to Eq. (3.6), where the correct Dg is precisely the instantaneous velocity field in the Z-component, and this is precisely div(ug) by Eq. (2.3). The second loss simply the enforces the semigroup property the Z-component. Take then together, arrive at the shortcut variant of F2D2, Shortcut-F2D22 obtained by minimizing the loss
F2D2 (0) LvM-sc (0) Lu-sc Ldiv-se + (D)sc
(3.10)
We demonstrate the derivation of Shortcut-F2D2 in Appendix A and provide the pseudocode for Shortcut-F2D2 in Algorithm and 4.
3.2.2 JOINT MEANFLOW: MEANFLOW-F2DS
Alternatively, McanFlow (Geng al., 2025) learns the time-averaged velocity ug(x, t, by enfore- ing the so-called MeanFlow identity us (Zi, 5, t) v( ()++))+((() (s-t) 4 us (Ia, 8, t), which we show in Corollary B.2 in Appendix B solves the Eulerian equation (Proposition 3.1 (b)). To our knowledge, this is the first proof of this fact.
For efficiency, de (x, vecobian-vector product (JVP). The training objective then to enforce this identity by regressing the model's prediction to the target implied thereby:
can be computed as 0,u(,, s) + Vzu(x, t. (2, obtained via
LMF (0) = s0
[1 t, (
sg ((s t) + V L (os,,ts)u(( ()u(T) 1)) + v o t))
(3.11)
where v(z,,t) t) = To is the ground truth instantancous velocity of the interpolant. Importantly, Eq. (3.11) encapsulates both Lvvs and Lu since PM recovers the flow matching loss when
To extend this this method to the joint system, we introduce the additional loss
Ld (0) Do(r,, 1, sg((# (,x,,
(3.12)
V. Do i, s)u( ()) - div(uo( f, 1?)
By analogy LMP. Ldiv-MF obviates the need for explicit divergence matching term Ldiv. MeanFlow-F2D3 then amounts to training the objective
LMF FMF-F2D2( (0) = LLM + MF
(3.13)
We demonstrate the derivation of MeanFlow-F2DD in Appendix and provide the pseudocode for Shortcut-F2D2 in Algorithm 5.
3.3 PRACTICAL DESIGN CHOICES
While instantaneous velocity supervision for sampling is straightforward to obtain from data, ob- taining reliable and tractable supervision for the instantaneous divergence presents significant chal- lenges. We address these through several key practical considerations.
Parameter Sharing. Since both X and Z components of our joint flow map derive from the same underlying velocity field, learning to predict both components simultaneously fundamen- tally learning two transformations of the same underlying dynamics. As result, we efficiently parametrize the joint flow map using shared backbone network with two separate prediction heads. This parameter sharing architecture ensures both outputs are derived from consistent representation of the flow dynamics and reduces the number of parameters compared to training separate models.
Hutchinson Trace Estimator. Computing divergence terms div(u(ze, ()) requires (dd) backward passes for exact computation, making training prohibitively expensive for high-dimensional data. Following standard practice (Grathwohl et al., 2018; Lipman et al., 2022: Song et al., 2020), employ the Hutchinson trace estimator div(r) R EE-AVo,,) V.. -e] which provides unbiased estimates with only (()) computational cost per training step.
Staged Training with Warm Start. Since divergence supervision depends on having accurate velocity predictions, we adopt a staged training approach. In practice, we pre-train the sampling velocity component uy alone using existing flow map distillation techniques, which provides good initialization for joint training later. Optionally, we can also pre-train a teacher flow matching model ts that serves reliable source of divergence supervision, replacing the potentially noisy

---

Published as a conference paper at ICLR 2026
Algorithm 2 Maximum Likelihood Self-Guidance Sampling with F2D2
1: Zo ~ Pu
2: D + Do( 0,
3: LNI.I. f- - --lloggo ( -
4: 20 & Ado (NLL)
One step Adam update w.r.t. zo optimizing CNLL
5: to...., k.Emmp f- linspace(0,, 0, 1, Kaamp + 1)
6: for == 0..... Ksamp do
7:
4- uo (oa 1i,4+11)
8: Fi+1 f my + (1++1 t2)u
9: end for
10:
return 21
div(us (Ze; f, E)) with the more accurate (1 t)) during joint distillation. When using pre- trained teacher model or warm-starting, we replace the flow matching loss with a teacher matching loss, which will be described next.
To further improve training stability and performance, we propose Shorteut-Distile shortcut model variant that combines the semigroup flow map with a learned teacher instantancous velocity. Our three-stage pipeline consists of: (1) Teacher pre-training: Train Ue using standard flow matching: (2) Shortcut-Distill: Warm start 0 with the teacher param- eters and replace Ldiv-se with teacher supervision: E, [1l, t) ua (22a })] (3) Joint dis- tillation: Warm start from sampling distillation, add divergence head and train both components jointly. This approach maintains the semigroup condition while leveraging a pre-trained velocity field to ensure the joint flow map are well-aligned.
Shorteut- Chorteut-Distill-22D0
Adaptations to forward-only training. While shorteut models and MeanFlow are both consistent with the full flow-map formulation, their practical use (i.e. sampling) and the available pretrained checkpoints are forward-only: training and tuning assume & S Several implementation choices are exclusively designed for this regime (e.g., the specific parameters of MeanFlow's 's logit-normal time sampling. shortcut models' discretized time sampling schedule, and common logarithmical sine-parameterizations used in EDM/CTM-styll models). Extending training to explicitly cover both S and 2 s would require re-deriving or re-tuning these design choices and can break the plug-and-play compatibility with existing pretrained models and algorithms.
In practice, we use simple finst-order approximation to reuse a forward-only model for backward likelihood integration. Let Az(x,, t + At) At ug uz( t, + At) denote the forward predicted displacement. We approximate the backward displacement by Ar( L, At) x -Ar(a,, 6, + At) and apply the same approximation to the log-density increment. This approximation is first- order accurate in At and works well empirically at small-to-mool moderate step sizes. While there exist out-of-distribution inputs
1, in practice we find that the networks generalize these scenarios with forward-only training.
3.4 APPLICATION: MAXIMUM LIKELIHOOD SELF-GUIDANCE WITH F2D2
Now that we have access log- likelihhod computation with few NFEs, we can explore various applications. One particularly interesting one is using the one-step divergence prediction (combined with the source distribution's 's log-likelihood) as pseudo- likellhodd objective for inference-time optimization. Specifically, we optimize the initial noise 20 to improve sample quality before running the sampling procedure. This approach resembles swardd-asee initial noise optimization for one-step generation models (Eyring et al., 2024), except do not require externll reward mod- els. Instead, we obtain the guidance signal from the model's own likelihood prediction head effectively performing self-guidance at inference time improve sample quality. This maximum likelihood self-guidance sampling algorithm is described in Algorithm 2.
4 RELATED WORKS
Likelihoud computation in diffusion and flow models. While diffusion and flow-base fow-based mod- els excel at sample generation, their likelihood evaluation remains computationally expensive.

---

Published as a conference paper at ICLR 2026
Flow Matching
Shortcut Model
Shortcut- F2D2
Shortcut still
Shortcut-Distill-ee
NFE-1
NFE=2
NFE-4
NFE-8
Figure 1: CIFAR-10 generated samples from different models with different numbers of steps.
Discrete-time diffusion models compute likelihoods through variational bounds requiring hundreds of NFEs (Ho et al., 2020; Nichol & Dhariwal, 2021). Continuous formulations enable exact likeli- hood via the probability flow ODE (Song et al., 2020) but require numerical integration with 100- 1000 NFEs. Prior research has explored various techniques to improve likelihood estimation (Grath- wohl et al., 2018; Song et al., 2021), but they still require many NFEs for accurate evaluation. In parallel, normalizing flows (Rezende & Mohamed, 2015) offer exact and tractable log- lielhiiss by using specialized network architecture designs and caane-of-vvariblle formula. Many recent efforts (Zhai et al., 2024; Ho al., 2019; Chen al., 2019) aim scaling up the normalizing flow principles through better dequantization and and advance architectures.
Accelerating sampling in flow-based models. Reducing sampling costs has been major focus in diffusion and flow matching research. Advanced ODE solvers (Karras et al., 2022; Lu et al., 2022) leverage the semi-linear structure of the probability flow ODE reduce discretization error. Distillation methods (Salimans & Ho, 2022; Sauer et al., 2024) and self-distillation models (Song & Dhariwal, 2023; Zhou al., 2025) provide alternative solutions by training student models to match teacher trajectories or training to perform self- seff-boosttapping with fewer steps. In particular, consistency models (Song al., 2023) and consistency trajectory models (Kim al., 2023) learn direct mappings from any point along the trajectory to data.
Flow map-based methods. Flow maps provide a general framework to model the solution op- erator of ODEs, enabling direct prediction of integrated trajectories (Kim et al., 2023; Boffi et al., 2025a;b). Recent works exploit this structure for few-step sampling. For example, as we have shown above, shortcut models (Frans et al., 2024) imposes semigroup property to learn the flow maps while MeanFlow (Geng et al., 2025) and Align Your Flow (Sabour et al., 2025) enforce Eulerian condi- tions. While these methods successfully reduce sampling to less than 10 NFEs, they either abandon likelihood computation entirely or still require full trajectory integration for likelihood evaluation.
EXPERIMENTS
5.1 SETUPS
We empirically verify the effectiveness of our method on image datasets CIFAR-I0 (Krizhevsky et al.). ImageNet 64 x 64 (Deng et al., 2009) and CelebA-64 (Liu et al., 2015). We evaluate the sample quality using Fréchet Inception Distance (FID) (Heusel al., 2017) on 50K generated im- ages. The negative log-likelihood (NLL) is measured in bits per dimension (BPD) on the entire test set of CIFAR-10 and CelebA-64 and a randomly sampled 10K subset of the ImageNet test set. We compare our method against flow matching (Lipman et al., 2022), shortcut models (Frans et al., 2024), Lagrangian self-distillation (LSD) (Boffi et al., 2025a) and McanFlow (Geng ct al., 2025) as baselines, and augment the later three for joint distillation. All models unconditionally trained. We use 1, 2, 4, 8 Euler steps for sampling and likelihood evaluation. In order to facilitate fair com- parison to the baselines, all shorteut and MeanFlow models trained with < while LSD models are trained with the full (t. s) ~ ((0, 1J2. Implementation details can be found in Appendix C.
5.2 RESULTS
CIFAR-10 Table 1 and Figure show the quantitative and qualitative comparison CIFAR- 10 spectively, As we can observe, flow matching yields poor FID and invalid NLL estimates in few-step setting. Shorteut model and MeanFlow achieve significantly better FID and are able compute NI.I.
8

---

Published as a conference paper at ICLR 2026
Table 1: NLL and FID results on CIFAR-IO dataset with different numbers of Euler steps. The flow matching model here, which achieves BPD 3.12 2 as the NLL with 1024 steps and FID 2.60 with 200 steps, is also the teacher model we use in our Shortcut-Distill For NLL, the closer to the teacher result (3.12 BPD) the better, and for FID, the lower the better. We denote the best results in bold, the second best with underlines, the overall best results boxes and invalid predictions in gray color.
8 Steps
4 Steps
2 Steps
1 Step
Method
NLL
FID
NLL
FID
NLL
FID
NLL
FID
Flow Matching
-9.93
20.63
-24.01
64.27
-52.85
146.24
-111.19
313.54
Shortcut Model
-12.07
7.10
-28.03
9.63
-60.01
16.04
-124.15
27.28
Shortcut-Distil (Ours)
-11.42
5.01
-26.82
5.41
-57.72
7.13
-119,42
12.75
McanFlow
-9.00
4.34
-21.26
5.14
-46.63
2.84
-97.59
2.80
Shortcut-F2D2 (Ours)
3.07
8.78
3.26
10.21
2.73
15.58
0.20
27.35
Shortcu--Distil-F2DD> (Ours)
3.12
5.68
2.87
5.96
2.38
7.35
1.62
13.76
MeanFlow-F2D2 (Ours)
2.38
3.78
1.34
4.37
1.63
2.59
3.51
3.02
Table 2: Negative log- ligllike (NLL) measured in BPD and FID results on ImageNet 64 dataset with different numbers of Euler steps. The flow matching model here, which achieves BPD 3.34 as the NLL with 1024 steps and FID 13.09 with 200 steps, is also the teacher model we use in our Shortcut-Distil.. For NLL, the closer to the teacher result (3.34 BPD) the better, and for FID, the lower the better. We denote the best results in bold and invalid predictions in gray color.
8 Steps
4 Steps
2 Steps
I Step
Method
NLL
FID
NLL
FID
NLL
FID
NLL
FID
Flow Matching
-6.41
31.60
-15.87
68.55
-35.23
170.00
-74.54
363.39
Shortcut- (Ours)
-9.03
19.47
-22.30
21.73
-49.01
28.12
-102.07
42.72
Shorteut-Distill-Pe (Ours)
3.51
21.91
3.94
24.05
3.97
29.83
1.54
44.02
for their ability to recover instantaneous velocity, but their NLL values remain invalid. Incorporat- ing our proposed F2D2 brings the NLL estimations to a calibrated range close to the teacher's BPD across different settings. In particular, both Shortcut-F2D2 and Shortcut-Distill-F222 substantially improves NLL compared to plain their orignal counterparts while maintaining competitive FID, indicating that F2D2 can provide reasonable likelihood estimates without sacrificing much sam- ple quality. Finally, MeanFlow-F2D2 shows FID improvements relative to the original MeanFlow while simultaneously producing calibrated NLL, demonstrating that F2D2's potential in providing complementary training signals that are beneficial to both components.
ImageNet 64 x 64 Shown in Table 2, flow matching quickly degenerates under few-step sampling, with invalid NLL and extremely poor FID. Shortcut-Distii improves the few-step FID but still pro- duces invalid NLL. By contrast, Shortcut- Coortcut-Distil-- F2D22 achieves both competitive FID and mean- ingful likelihoods close to the teacher's 's BPD of 3.34 across all step counts. These results further confirm F2D2's ability for simultaneous fast sampling and fast likelihood evaluation.
CelebA-64 As shown in Table 3, the flow matching model degenerates under few-step Fuler sam- pling, yielding invalid likelihood estimates and worsening sample quality. L.SD) substantially im- proves the few-step FID, but still produce invalid NLL across all step counts. In contrast, LSD-F2D2 achieves both superior image quality and well-calibrated likelihoods a few steps.
5.3 MAXIMUM LIKELIHOOD SELF-GUIDANCE WITH MEANFLOW-F2D2
Figure shows the FID and qualitative comparison among different methods built upon MeanFlow using the number of forward NFE on CIFAR-10. As can observe, our F2D2 improves the model's 's inference time scaling ability. With additional self-guidance, the model not only surpasses the baseline MeanFlow performance but also outperforms 1024-step flow matching model of the same size, demonstrating the effectiveness its own likelihood predictions as valid signals to guide the sampling process toward higher-quality generations.

---

Published as a conference paper at ICLR 2026
Neanfoo
Nearton-7222
Set Cuadance
Mearfiow
5.0
Mcanfiow F202
Meanfiow
Self-Gustann.
Flow lew-matching
Wi 1024
V7S-2
4.0
3.5
3.0
V77.8
NiF
2.5
NFE
(a) FID comparison across different num- bers of forward NFEs.
(b) Example samples from various McanFlow-bascd models ing different numbers of forward NFEs.
Figure 2: Results of MeanFlow-based methods CIFAR-10.
by 101
S0IRAn mo
Table 3: Negative log-likelihood (NLL) measured
BPD and FID results on CelebA-64 dataset with different numbers of Euler steps. The flow match- ing model here achieves BPD 1.75 in 1024 steps and FID 2.48 in 200 steps. For NLL, closer to the flow matching estimate is better, for FID, lower is better. Best in bold; invalid in gray.
Methad
8 Steps
4 Stops
2 Strpp
1 Step
NLL
F10
NLLL
M10
NLL
F1D
NLL
Muccing
LSD-F2D2
1.75
1.44
Figure 3: LSD-based log- ligeliiel evaluation comparison on 2D checkerboard dataset.
5.4 2D CHECKERBOARD
In this section, we present a set of comparison of log-likelihood evaluation results on 2D checker- board, a synthetic dataset with analytically tractable ground truth log- log-likelihood. As we can observe in Figure 3, without F2D2, vanilla LSD catastrophically fails at few-step log-likelihood estimation. the other hand, our LSD-F2D2 is able to accurately recover target density distribution even with only NFE, preserving both spatial structure and density values. This directly validates that, while traditional methods fail at few-step likelihood evaluation, F2D2 enables accurate likelihood via joint flow map distillation.
6 CONCLUSION
We present fast flow joint distillation (F2D2), simple and modular framework that enables both fast sampling and fast likelihood evaluation flow-based generative models. By jointly distilling the sampling trajectory and divergence computation into unified flow map, our method simul- tancously achieves accurate likelihood evaluation and high sample quality with just few NFEs. Our experiments on CIFAR-10 and ImageNet 64 64 demonstrate that F2D2 maintains accurate likelihood estimates while preserving sample quality when applied to existing few-step methods in- cluding shortcut models and MeanFlow. The efficiency gains from F2D2 enable new algorithmic possibilities, illustrated by maximum likelihood self-guidance method, which enables a 2-step MeanFlow model to outperform a 1024-step flow matching model of the same size on CIFAR- 10. As based models continue scale, we believe that efficient likelihood evaluation alongside fast sampling will become increasingly important for enabling new training objectives, model analysis techniques, and downstream applications that require both capabilities.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
We provide proofs to our theoretical results in Appendix A and B. We also provide the imple- mentation details reproduce our algorithm and experimental results in Section 5 and Appendi C. The links to the PyTorch and JAX implementations of our algorithm can be found on our website: https://kellyyuttonm
github io/ f2d2/.
ACKNOWLEDGMENT
MS would like to recognize a TRI University 2.0 Research Partnership and a Google Robotics Research Award. We also thank Michael S. Albergo and Eric Vanden-Fijnden for the helpful dis- cussions.
REFERENCES
Hirotugu Akaike. A new look at the statistical model identification. IEEE Transactions on Automatic Control, 19(6):716-725 1-33,,,, 1974.
Michael S Albergo and Eric Sunden-Eijnden Building normalizing flows with stochastic inter- polants. arXiv preprint arXiv:2209.15577 2022.
Michael S. Albergo, Nicholas M. Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A uni- fying framework for flows and diffusions, 2023. URL 08797.
/3003.
Black Forest Labs. Flux.1 kontext: Flow matching for in-context image generation and editing in latent space. arXiv preprint arXiv:2506.15742, 2025. URL https://arxiv.orgg /bs/ 2506. 15742.
Nicholas M. Boffi, Michael S. Albergo, and Eric Vanden-Eijnden. How to build consistency model: Learning flow maps via self-distillation 2025a. URL 2505 18825.
org/ abs/
Nicholas M. Boffi, Michael S. Albergo, and Eric Vanden-Eijndee Flow map matching with stochas- tic interpolants: mathematical framework for consistency models, 2025b. URL https:
//arxiv. wwwasiiiorrr/aaae www//rggss....oo.o....e.....eeee httpaab//2.00000 07507.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Ricky TQ Chen, Jens Behrmann, David K Duvenaud, and Jörn-Henrik Jacobsen. Residual flows for invertible generative modeling. Advances neural information processing systems, 32, 2019.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hier- archical image database. In 2009 IEEE Conference Computer Vision and Pattern Recognition, pp. 245-255, 2009. doi: 10. 10110C/P..2000.52226666 528468
Luca Eyring, Shyamgopal Karthik, Karsten Roth, Alexey Dosovitskiy, and Zeynep Akata. Reno: Enhancing one-step rext-to-image models through howard-based noise optimization. Advances in Neural Information Processing Systems, 37:125487-125519 2024.
Ronald A. Fisher. On the mathematical foundations of theoretical statisties. Philosophical Transac- fions of the Royal Society of London. Series A, Containing Papers of a Mathematical or Physical Character, 222:009-668 1922.
Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel. One step diffusion via shortcut models. arXiv preprint arXiv:24100 0 2557, 2024.
Ruiqi Gao, Emiel Hoogeboom, Jonathan Heek, Valentin De Bortoli, Kevin P. Murphy, and Tim Salimans. Diffusion meets flow matching: Two sides of the same coin. 2024. URL https: //di ffusionflow www.oooowww........bb.... too..

---

Published as a conference paper at ICLR 2026
Zhengyang Geng, Mingyang Deng, Xingjian Bai, Zico Kolter, and Kaiming He. Mean flows for one-step generative modeling. arXiv preprint arXiv:2505.13347, 2025.
Google DeepMind. Veo 3: High-fidelity text- and image-to-video generation with audio. Model card, Google DeepMind, may 2025. URL http::::///// httts::/sttorgg:: googleapis com/ deepmind medi /Model Cards /Veo- M---MMoe--aaae Model-Card.pdt. Accessed 2025-09-21.
Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv:1880001166,.000 2018.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbeel. Flow++: Improving flow- based generative models with variational dequantization and architecture design. In International conference on machine learning, pp. 2722-2730. PMLR, 2019.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Tero Karras, Miia Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in neural information processing systems, 35:26365-26577, 2022.
Dongjun Kim, Chieh-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning proba- bility flow ode trajectory of diffusion. arXiv preprint arXiv:2310.02279, 2023.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 (canadian institute for advanced re- search). URL http: //www//wwwwww.....////// es teronto. wdw/ httpiiiiiiiiaaaaa.a.aaa
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamm Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0383 0.00000000,,... 2022.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face autributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: fast ode solver for diffusion probabilistic model sampling around 10 steps. Advances in neural information processing systems, 35:5775-5787, 2022.
Jerzy Neyman and Egon S. Pearson. On the problem of the most efficient tests of statistical hypothe- ses. Philosophical Transactions of the Royal Society A, 231:289-337, 1933.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International conference on machine learning, pp. 8162-8171. PMLR, 2021.
OpenAl.
Sora:
Video
generation models as world simulators.
Tech- com/indez Accessed: 2024-02-15.
nical Report, OpenAL, 2024. wow.ttiiiommmmeeeeee.m hat---oool------ooooda--aaaaaaao
URL https: wtt.sssssslllls..omm.o.oo.oooooe
Long Ouyang, Jeffrey Wu, Xu Jiang. Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to fol- low instructions with human feedback. Advances in neural information processing systems, 35: 27730-27744, 2022.
12

---

Published as a conference paper at ICLR 2026
Adam Polyak, Amit Zohar, Andrew Brown, Andros Tjandra, Animesh Sinha, Ann Lee, Apoorv Vyas, Bowen Shi, Chih Yao Ma, Ching- Yao Chuang, David Yan, Dhruv Choudhary, Dingkang Wang, Geet Sethi, Guan Pang, Haoyu Ma, Ishan Misra, Ji Hou, Jialiang Wang, Kiran Ja- gadeesh, Kunpeng Li, Luxin Zhang, Mannat Singh, Mary Williamson, Matt Le, Matthew Yu, Mitesh Kumar Singh, Peizhao Zhang, Peter Vajda, Quentin Duval, Rohit Girdhar, Roshan Sum- baly. Sai Saketh Rambhatla, Sam Tsai, Samanch Azadi, Samyak Datta, Sanyuan Chen, Sean Bell, Sharadh Ramaswamy, Shelly Sheynin, Siddharth Bhattacharya, Simran Motwani, Tao Xu, Tianhe Li, Tingbo Hou, Wei-Ning Hsu, Xi Yin. Xiaoliang Dai, Yaniv Taigman, Yaqiao Luo, Yen- Cheng Liu, Yi-Chiao Wu, Yue Zhao, Yuval Kirstain. Zecheng He, Zijian He, Albert Pumarola, Ali Thabet, Artsiom Sanakoyeu, Arun Mallya, Baishan Guo, Boris Araya, Breena Kerr, Car- leigh Wood, Ce Liu, Cen Peng. Dimitry Vengertsev, Edgar Schonfeld, Elliot Blanchard, Felix Juefei-Xu, Fraylie Nord, Jeff I Liang. John Hoffman, Jonas Kohler, Kaolin Fire, Karthik Sivaku- mar, Lawrence Chen, Licheng Yu, Luya Gao, Markos Georgopoulos, Rashel Moritz, Sara K. Sampson, Shikai Li, Simone Parmeggiani, Steve Fine, Tara Fowler, Vladan Petrovic, and Yuming Du. Movie gen: A cast of media foundation models. arXiv preprint vrXiv:2410.13722 2024. URL https: /aaiiiiiiv. org/ abs/ /0os/24/0 2410 13720. Submitted October 2024.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly reward model. Advances in neural information processing systems, 36:53728-533 53741, 2023.
Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. Advances in neural information processing systems, 32, 2019.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Interna- tional conference on machine learning, pp. 1530-1538. PMLR, 2015.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings the the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Amirmojtaba Sabour, Sanja Fidler, and Karsten Kreis. Align your flow: Scaling continuous-time flow map distillation. arXiv preprint arXiv:2506. 14603, 2025.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.005/0 2022.
Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion dis- tillation. In European Conference on Computer Vision, pp. 87-103. Springer, 2024.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06347, 2017.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, et al. Deepseekmath: Pushing the limits of mathemati- cal reasoning in open language models. arXiv preprint urXiv:2402.033000 2024.
Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. arXiv preprint arXiv:2310. 14189, 2023.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020.
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in neural information processing systems, 34:1415- 1428, 2021.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. 2023.
Shuangfei Zhai, Ruixiang Zhang, Preetum Nakkiran, David Berthelot, Jiatao Gu, Huangjie Zheng, Tianrong Chen, Miguel Angel Bautista, Navdeep Jaitly, and Josh Susskind. Normalizing flows capable generative models. arXiv preprint arXiv:2412.06329 2024.
Linqi Zhou, Stefano Ermon, and Jiaming Song. Inductive moment matching. arXiv preprint arXiv:2503.07565. 2025.
13

---

Published as a conference paper at ICLR 2026
A CHARACTERIZATION OF THE JOINT FLOW MAP
In this section, we provide analytical characterization of the joint flow map for the combined sampling and likelihood dynamics Eq. (2.3). For the simplicity of notations, we define the joint system
d
di Ut v( t), d
x(0) 2o NO
(A.1)
dt 2t = div(((())))))
20 po
where will the shorthand 2 for log (x))) Moreover, define the right-hand side of Eq. (A.1) as t) = (v(2, t), -iv--di((((;,,, (di,(t((,,)))) where = (x, 2r,2))
We then define the flow maps for the two subsystems in Eq. (A.1) as
0x(x,t,))= x (s t)u 8), Φ ==( +
(A.2)
=
((,,
In Eq. (A.2), note that the hierarchical structure in Eq. (A.1) is explicit, and that the function D only depends on æ and not on 2. We then define the joint flow map
(x(xxt,,))
(v,
Φz( 2,, =y (I,
(A.3)
u(x,t,,) D(x, t,
f(x,t,8)
We first recall the simple tangent identity denoted the tangent condition by Boffi et al. (2025a), also leveraged by (Kim et al., 2023; Frans et al., 2024; Geng et al., 2025), which allows us to recover the instantaneous velocity and divergence in the s + limit:
Lemma A.1 (Tangent Condition). The flow map (y s) for the joint system Eq. (A.1) satisfies (y,t,s) t, = gly, s) = f(z, s, s). In particular, u(x, s, s) = u(x, t) and D(x, s, s) = -div(o(x,,)) s)).
lim;- 0 (y,
Proof. The proof follows by application Lemma 2.1 from Boffi al. (2025a).
Now, given Eq. (A.3), we may now state the following proposition, which is based on an identity similar Lemma A. in reverse.
Proposition A.2 (Characterization of the Joint Flow Map). Let By (y, - (s ( ()f(r, L s) satisfy f(z, s, g(y, s) where g t) (v(z, t). diiv(v(a,, t))) denotes the dynamics for the joint sampling and likelihood system Eq. (A.1). Then, by (y. is the flow map for the joint system if and only if any of the following conditions are satisfied:
(Lagrangian condition) Zs,t satisfies the Lagrangian equation
o 0,0r(g,t,)) f(
(g,t,s) gd+1 1 10,1
(A.4)
2. (Eulerian condition) (y t, 8) satisfies the Eulerian equation
0, + by by sf =,
(y,t,s) gd+ x 10,1113. (A.5)
3. (Semigroup condition) Za,t satisfies the semigroup property
By (y,t,s) 4y ( (4)
v (g,t,r,s) gd+1 < (A.6)
Proof. The proof follows by application of Proposition 2.2 from Boffi et al. (2025a) applied to the joint system.
14

---

Published as a conference paper at ICLR 2026
Proposition A.2 gives three characterizations of the joint flow map that may be used to devise few- step flow-based training algorithms. In each case, by definition of Eq. (A.3), the X block reduces to the flow map characterizations introduced in Boffi et al. (2025b;a) for the sampling system. The second block for the likelihood dynamics is new, which we focus on now to instantiate the resulting equations.
Lagrangian likelihood equation. By inspection, Eq. (A.4) leads to the equation
Diz, t, s) = divv(e(ex (z, t, s), s)) - (2,,
(A.7)
Squaring the residual leads to the objective function
((D) = [1i((z,,,))) (-dii))o((((((((((xee))e 8)) - 1)d, (1 [2
(A.8)
At training time, because we don't have access to the true -div(e) or the true sampling flow map Φ (z,t, s), may replace them by their self-consistent estimates,
(()) = E [|] Do(s,, 1, sg (- (-div -idiviua(0e ( 4, 8),
() 4,,),, 8)) [1]]
(A.9)
where we have also placed stopgrad operator to avoid backpropagation through Jacobian- products and to control the flow of information from the teacher the student. This gives the Lagrangian likelihood self-distillation algorithm.
Notice that in the limit as s, the Lagrangian condition can cover the tangent condition.
Eulerian likelihood equation. To derive our Eulerian schemes, we first note that
x xx, by V 1202
(A.10)
(z,
Hence to compute the second component of the Eulerian equation, we must collect some simple algebraic identities,
(,0 = -D(xx,,,) + (s (2, ---000))- V (,02(2,2tt,))) = (((,))) V. V,02(z,,,t, =
(A.11)
Using the above, we find that the Eulerian relation for Y becomes
-D 1, 8) + (s - D); , t, 8) We may enforce this equation by minimizing the square residual,
t)V, , Da, t, s)u(a, t) tiv(u(x, 0. (A.12)
C(D) E [z),(t, ((s t)0,
(s D(,, (iv( t))) [2
(A.13)
At training time, we again place a sg (-) operator to avoid backpropagating through the derivatives, ((0) E [1D,(zi,t, sg ((s t)a, Dz s) (s ()), div(u(zs,, t))) [1
(A.14)
In the above, we do not have access to the ideal v(xx,() nor -div(v(x,,)) !). However, we observe that because of the placement of sg (.), resulting gradient will be linear in o(x,, 1), so that we may replace it by its Monte Carlo estimate. In practice, this reduces to conditional-Of flow matching. Second, we replace diiv(r(z,, t) by the self-consistent estimate -div(ug (2 t, t). leading to
C(0) E (o t, s) sg ((s (o t,
Do (2(1 zo) div(uo (Ze, t, t))) [1a].
(A.15)
This gives the MeanFlow-F2D2 algorithm.
Semigroup property. Last, we consider the semigroup approach. The second block is given by,
Y2,(2)= (Y,,
(A.16)
4y (y.t, t, = By (by
(A.17)
Writing this out using Eq. (A.3), we find that
z + (s (), 8) 2z 1,0
(A.18)
2 ++=-=:1(hhh.===----0------------
15

---

Published as a conference paper at ICLR 2026
Setting r = (t+ s) recovers continuous limit of shortcut models, as shown in Boffi et al. (2025a).
In this case, Eq. (A.18) becomes
s) ((0x(,))))))
(A.19)
Squaring the residual gives
1 ((D) = E [1(eent,)) ( +
(A.20)
where as in the above Eulerian and Lagrangian approaches, we have replaced the ideal flow map X by the self-consistent estimate. Again, to control the flow of information we may place stopgrad,
1 258 (Do(as,t,)) + Do (e((xx)(ttttt)))))))))))
( E (Do(
2
(A.21)
which gives Shortcut-F2DD3 algorithm.
B ADDITIONAL. THEORETICAL ANALYSIS
Corollary B.1. Shortcut models enforce semigroup property with their self-consistency loss.
Proof. The same proof as the semigroup property for by holds for 4x
Corollary B.2. MeanFlow (Geng et al., 2025) directly enforces the Eulerian condition with Mean- FLow identity.
Proof. Since u(zr,t,t) t,t) = v(z,,t) t) is implied by the linear parametrization in Equation 3.1, we can have
a ( V ()((m) 0,(xe + (s (o0-))u(rr,4450 5)) + Vn, ( (a ()))), = ()) + (s t)0,u(21, L (s
(1+
(s t) (0:u(xr, t, 8) + V t,
st u(z,,, u(x,, t)
d
(s - t)
u u(2,, u(x,, s)
de
which is exactly MeanFlow identity.
16

---

Published as a conference paper at ICLR 2026
CIFAR-10
ImageNet-64 x 64
CelebA-64
Noise embedding
Positional
Positional
Positional
Channels
128
192
128
Channels multiple
1,2,2,2
1,2,3,4
I 1,2,3,4
Attention resolution
16
32,16,8
16,8
Residual blocks
3
3
Dropout
0.13
0.1
0.0
Batch size
512
1024
256
GPUs
8 L40S
8 L40S
L40S
Iterations
150k
125k
350k
Learning Rate
le-4 (constant)
le-4 (constant)
le-2 (Sqrt decay 35k)
Warmup Steps
30k (linear warmup)
60k (linear warmup)
Precision
float32
float32
6
Optimizer
RAAda
RAdam
RAdam
EMA rate
0.9999
0.9999
0.9999
Table 4: Training hyperparameters for flow-matching.
C IMPLEMENTATION DETAILS
In this section, we provide the training details for each model experimented in this paper, and the pseudocode for each F2D2 instantiation.
Flow-natching. We use linear interpolation to generate training targets. Flow-matching models are trained on both unconditional CIFAR-10 10 and ImageNet-64 64 using the configurations in Table 4, with 200 standard Fuler steps for sampling. This model serves as the teacher model for the second distillation stage.
Shorteut Model. We reimplement the shortcut model ( (Frans et al. (2024)) for CIFAR-10. Follow- ing their method, we strictly sample 2 (i.e. forward-only training) and use discrete timesteps and set 1/128 as the smallest unit of time for approximating the ODE. We consider 8 possible shortcut lengths ranging from (1, 1/2,...... ... 1/128). Similarly, divide the batch into 3/4 for flow-matching training and 1/4 for seff-consistency training. However, instead of DiT, we use U-Net backbone with configurations in Table 4. We train for 100k iterations on CIFAR-10 and sample with 2, 4, and 8 standard Euler steps. Additionally, we follow the original paper and parametrize the model to inputs.
take v,, and
Shortcut-Distill. We use the velocity predicted by the flow-matching model as the target for the flow-mtcching loss, instead of Wasta Enoise. All other configurations remain the same in the Shortcut Model.
MeanFlow. We directly use the pre-trained model weights and parametrization provided by the offi- cial PyTorch repository to conduct the CIFAR-IO 10 MeanFlow experiments. The provided pre-trained MeanFlow model has the same size as all other models we implement for CIFAR-10 10. Because the pre-trained MeanFlow models is forward only (i.e. t S s), we follow the same convention and adapt our algorithm based on the considerations that we have elaborated Section
LSD. We follow the experimental setup of Boffi et al. (2025a). We allocate 75% of each batch to the flow matching loss and 25% to the self-distillation loss for the first 200k iterations. And then allocate 50% of each batch to the flow matching loss and 50% to the self-distillation loss for another 150k iterations. To enable both forward and backward estimation, we sample two time steps uniformly at random for the lelf-distillation loss, without enforcing any ordering constraint between them (i.e., we do not require one to be larger than the other). All LSD models are trained with the full (t, s) 4( 1
Log- lineliool estimation with vanilla flow map models. Since Flow maps can recover instanta- neous velocity through their tangent condition uo( () v(se, l), flow map models trained only
17

---

Published as a conference paper at ICLR 2026
Algorithm 3 Shortcut-F2D2 Training
1:
for each training step do
2:
Ii ~ Pdntts To ~ po. (1, -,000
forward-only training uses / S
3:
- (L 8
4:
It - (1 t + tx
5:
Ir
+ (r (r- -rr- (2), (2t, .t,))
6:
(c t [Te 1,
(0) ro)ll
7:
(us [10 L,
28 (uo r) (o( 1|2
8:
Dj f sg (div(us (Ie, (I;, t,))
9:
Ldiv-sc (0) (a (,
Dz
10:
Disc (0) f ( 8)
sg (Do(o,,(,r) + DoDDDoa r, r, 8) 1| Luss Cdiv-sc (8) + LD-sc (8)
LsC-F2D2 (8) Cv
12:
Update O w.r.t Lsc-F2D2
end for return A
14:
Algorithm Shorteut-Distill-FFF FFDt2 Training
1: for each training step do
2:
T1 ~ Pdatas Fig ~ Po. (t, L((1, 1
>> forward-onle only training uses t S
3:
r - (t st/s
4:
It - (1 t)a0 + tx1
5:
Iy < r; + (r (r- (ruo (Tt, r,
6:
LvM-sc (0) - ||ug(ze, t)
Vi (2t, t)l1
&> match with flow matching teacher us
7:
Lu.se (0) + [ 4,
2g (uo( (, uo(xr, l
8:
Dr -- sg (div((p(()))) ()))
b matching teacher divergence
9:
Caiv (0)
D,ll2
10:
sD (0) + Da( 1,
sg (Do Do(2, t,r) (o r,)) 11 su sc + LD
F2D2 ( s se
12:
Update 0 w.r.t. (22-F222
13:
end for
14:
return
for sampling can, in principle, evaluate likelihood by computing div(us t, t)) and solving for Eq 2.3 and 2.4. As a result, we use Fuler solver with this formuation to produce the NLL estimation of the baseline Shorcut, Shortcut-Distill and MeanFlow models.
F2D2 for Shortcut Model and Shorteut-Distil. For F2D2 training, we add a scalar head after the UNet decoder. The scalar head implemented first flattened, then passed through two fully connected layers with SiLU activations, and finally projected to single scalar. This head maps the spatial feature representation into scalar value, which we use to predict divergence.
MLP: the decoder's final feature map is
For CIFAR-10. 10, the hidden sizes are 128 and 64, while for ImageNet-64 64 they are 64 and 16. We linearly warm-start the model with teacher weights and train using the flow-matching loss, with the teacher model taken from either the Shortcut Model or the Shortcut Model-Distill. We also adopt the same discrete timesteps as in the Shortcut Model. To balance the four losses, we scale down divergence distillation targets by factor of 20,000 for CIFAR-10 and 300,000 for ImageNet- 64.64. We train Shortcut-Distill-Fat 64 x 64. We also train Shortcut-F2D2 for 10k iterations on CIFAR-10. We stop the training when the calibrated value of bpd is reached. All other configurations remain the same. Although derivation matches -di our implementation parametrizes to predict div(v), which yields an equivalent formulation long the training loss and the sampling process as the appropriate signs.
for 10k iterations CIFAR-10. 10k iterations ImageNet-
The pseudocode for Shorteut-F2D2 is provided in Algorithm 3 and the one for Shortcut- Shortcut-Distill-F22 Distill--2D2 provided in Algorithm 4.
F2D2 for MeanFlow. We use the same scalar head for CIFAR- 10 with Shorteut-Distill- and train for additional 50 epochs using the same configurations as in the original PyTorch imple-
18

---

Published as a conference paper at ICLR 2026
Algorithm 5 MeanFloo-FF2D Training
1: for each training step do
2:
I1 ~ Pdntts Tg ~ po- (1 ~
forward-only training uses / S
3:
Ie (1 1) + Lx1
4:
(0) (1 (,, sg ((s t) + u(z,, () f)) + v( t))l
5:
Di.a - sg( ( (1,,,
t,,, ()) div(uo(Tx, 4,1))
6:
Caiv-ur( (0) + [D t, D1.
7: LMM F2D2 (0) Lmr (0) + Ldiv (0)
8: Update 0 w.r.t. LMT F2D2
9: end for
10: return
Algorithm LSD-F2D2 Training
1: for each training step do
2:
Pdata To Po. (1, ~ ~((0)
3:
It - (1 - 1) + lxg
4:
t),
5:
LLSD (0) 10,,z, - sg (ug ( (,)) 8)]|
6:
Laiv-LSD (0) - (a L - sg (=dv 8, - (,,,,
7: CLSD-F2D2 (0) < (L)D( fLdv-LsD(O)
8:
Update O w.r.t LLSD-F2D2())
9:
end for
10:
return /
mentation. We use the pre-trained flow matching model provide the instantancous divergence supervision. Algorithm 5 shows the pseudo-code for MeanFlow-F2DD training.
F2D2 for LSD. For F2D2 training, we also add an MLP scalar head for divergence prediction and the hidden sizes are 32 and 8. We change the learning rate to le-4 and use 50% of each batch to the flow matching loss and 50% to the self-distillation loss. To balance the velocity and divergence losses, we scale down the divergence distillation targets by factor of 10,000. We train LSD-F2D2 for 60k iterations. All other configurations remain the same as LSD. Algorithm shows the pseudo- code for LSD-F2D2 training.
Maximum Likelihood Self-Guidance with F2D2. We first predict the negative likelihood using randomly sampled noise with 1-step divergence prediction and take this negative likelihood as the loss. We then update the noise with one step of Adam, using learning rate x 10-3 for 1-step sampling and 10-3 for 2-, 4-, and 8-step sampling. We then evaluate FID using 1, 2, 4, and 8 standard Euler steps for sampling.
BPD. We the same method Lipman al. (2022) to compute the BPD. To compute BPD with F2D2 models, we discretize the interval from 1 to into equal segments, while the uniform step size serves the second time input.
2D Checkerboard. We follow the experimental setup of Boffi et al. (2025a). The model is a 4-layer MLP with 512 hidden units per layer and GELU activations. Both the flow-matching baseline and the Shorteut model are trained for 150k iterations with a batch size of 100k and an initial learning rate of 10- using a square-root decay schedule after 35k steps. For the Shortcut model, each batch is divided in 3 ratio between the flow-matching loss and the self-distillation loss. For Shortcut-F2DD2. we extend the output layer with an additional dimension and initialize training from the pretrained Shortcut model, running extra 27k iterations. As for the LSD model, we train for 200k iterations with batch size of 10k while keeping other configurations unchanged. For LSD-F2D2, we add an additional MLP head (hidden size 512) to predict divergence and scale the divergence targets by factor of 0.1. In each batch, we use a 1 split between the flow- matching loss and self-distillatio loss. We train LSD-F2D2 for 200k iterations using learning rate of 1 x 10 with a square-root decay schedule starting after 50k steps. All other configurations are kept the same.
19

---

Published as a conference paper at ICLR 2026
Table 5: Sampling speed (wall-clock time and per-image latency) for different NFEs.
Dataset
Hardware
NFE
4 NFE
2 NFE
1 NFE
CIFAR-10
4xL40S
146 s (2.92 ms/ing)
83 (1.66 ms/img)
51 (1.02 ms/ing)
40 (0.80 ms/img)
ImageNet-64x64
8xL40S
289 s (5.78 ms/img)
160 (3,20 ms/img)
95 (1.90 ms/ing)
63 (1.26 ms/img)
Computation Cost and Runtime. We conduct all our training on an 8-GPU L40S node. To train typical F2D2 CIFARIO model, it takes around 1 day for the Flow-matching teacher, I day for the vanilla Shortcut or Shortcur-Distii baselines, and 5 hours for the F2D2 finetuning on top of it. For ImageNet64466 models, it takes around days for the flow matching teacher, 6 days for the vanilla shortcut or shortcut-distil model, and 5 hours for the F2D2 finetuning, For inference run time, we provide the wall-clock time generate 50k samples for each dataset using different numbers of NFEs in Table 5.
D ADDITIONAL ABLATIONS
In this section, present additional ablation study to further investigate the bevavior of our F2D2 models.
With variance around 1.8 10 each time step, the Hutchinson estimator does produce values that significantly larger in magnitude than the sampling range. Therefore, as we have mentioned in the Appendix C. we follow common deep learning practices and apply divergence target in order for the neural network to better learn the prediction. Here in Table 6, we present an ablation study on the divergence scaling factor to validate this choice of hyperparameter. As we can observe, appropriate scaling enables simultaneous calibrated likelihood prediction and high sample quality, while scaling too strong can produce invalid NLL and scaling too weak can degrade overall performance.
10-5 scaling factor the
Figure shows the self-guidance sampling algorithm with multi-step Adam optimization. As we can observe, while adding optimization steps does not further improve the FID, the performance also does not collapse. We are interested in other variants of our self-guidance algorithm to better improve the performance as exciting future work directions.
Finally, we conduct additional study on the effect that different loss components have on the sample quality well as likelihood estimation accuracy. In particular, slightly modify Equation 3. 10 to add an additional weighting scalar À to Laiv-se + LD-se and compare the performance under various weighting.
Lsc-F2D2(0) LVM-sc (0) + Lu-sc(0) + A(Laiv-sc( LD-sc(0))
(D.I)
Table 7 shows the Shortcut-Distill-Fal results with À 0.01, 0.1, 10, with to our original setting and the performance of Shortcut-Distill without F2D2 as reference. As we can observe, adding an additional loss weighting can potentially help F2D2 achieve a better sweet spot in balancing sample quality and likelihood estimation: While our original equal weighting loss (A=1) obtains the most calibrated likelihood, slightly penalizes the FID scores. Similarly, sig- nificantly down-weighting (A 0.01) further improves FID but loses the ability to produce accurate likelihood. However, when choosing an appropriate (in particular when À = 0.1), this additional weighting enable the model to produce both FID scores that match even slightly improve the original Shortcut- AhbsttlttDisstlll performance, and relatively calibrated log- likelioddoood estimation.
equivalent
The experiment in Table opens the door to many intriguing future research directions that can po- tentially further improve the performance. These include more fine- mne-grained loss weighting scheme and developing training schedules similar to the ones proposed in Frans et al. (2024) and Geng et al. (2025), where the flow matching objective and the self-consistency objective are optimized in separate portions of the training.
20

---

Published as a conference paper at ICLR 2026
Table 6: Ablation study on the divergence scaling factor.
Steps
Steps
2 Steps
1 Step
Scaling
NL.I
FID
NLL
FID
NLL
FID
NL.
FID
X 10 10-4
3.27
5.16
2,43
6.31
1.26
7.06
-2.92
13.16
x 10
3.12
5.68
2.87
5.96
2.38
7.35
1.62
13.76
x 10
5.38
15.57
5.44
25.20
5.50
10.96
5.60
23.18
Self-Gudaacc step
SullaGulcancc steps
4.0
Self-Guldiance Steps
Self-Gulcance steps
3.5
3.0
2.5
Forward NFE
Figure 4: FID results of multi step self self-guidance sampling on CIFAR10.
Table 7: Ablation study on different loss weighting on Laiv-sc + Cpisc in Shortcut-Distill-2DD
8 Steps
4 Steps
2 Steps
1 Step
Scaling
NLL
FID
NLL
FID
NLL
FID
NLL
FID
Shorteut- Distill
-11.42
5.01
-26.82
5.41
-57.72
7.13
-119.42
12.75
10
1.47
11.66
0.82
11.23
-0.07
10.11
-2.19
18.27
1
3.12
5.68
2.87
5.96
2.38
7.35
1.92
13.76
0.1
2.59
5.01
2.00
5.33
1.78
7.01
0.68
13.04
0.01
2.04
4.87
1.51
5.28
0.59
6.76
-1.68
12.70
E ADDITIONAL RESULTS
In this section, we present additional experimental results. Specifically, in Table 8 we showcase the NLL per sample error w.r.t. the teacher prediction on CIFAR10. As we can observe, the F2D2 vari- ants lower the eтror by 7-171x in comparison to the baselines. This evaluation directly verifies that our F2D2 can produce sample-level calibrated likelihood estimations that are suitable practical applications, not merely matching the summary statistics.
Figure demonstrate the training dynamics of our SeirrrutDDistil-FF22 model with the training loss curve, where we observe stable training with expected small fluctuations native uniformly random timestep selection at each iteration.
We also provide additional 2D checkerboard results in Figure 5 and 6 and qualitative image gen- eration results in Figure the generated sample density via forward integration and Figure 6 shows the data sample density via backward integration using different models. All images shown in this section are non-cherry picked results.
19,20,21,22. In particular, Figure shows
21

---

Published as a conference paper at ICLR 2026
Table NLL Mean absolute per sample error w.r.t the teacher's prediction in BPD on CIFARIO. shows that the F2D2 variants lower the error by x x in comparison to the baselines.
Method
8 steps
steps
2 steps
step
Shortcut
15.11
30.97
62.85
126.58
Shortcut-Distill
14.47
29.79
60.61
121.99
Meanflow
12.02
24.24
49.33
100.03
Shortcut-F22 F2D2
0.52
0.55
1.01
3.08
Shorteut-Distil-FF20 Disttill-E2DD2
0,41
0.50
0.90
1.69
MeanFlow-F2DD
0.80
1.78
1.51
0.60
Table 9: NLL Mcan absolute per sample error w.r. the teacher's prediction BPD on CelebA-64.
Method
steps
steps
2 steps
step
LSD LSD-FF222
8.45 0.18
16.59
34.46
71.64
0.24
0.33
0.53
Target -34
step
-32
-32
-3
-34
-34
-3
-36
-34
34
34
0
38
+0
Flow matching (Huochi, 4 steps
Shortcut (Hutchi, 4 stops
Shortcut- F202, stops
-32
-32
-3.2
-44
-44
Figure 5: Log-likelihood comparison on 2D checkerboard generated samples among different mod- els using forward integration.
Targes by pXc
flae-retciing Ikrchh
dapp
Fow-miidirg (MMn)
LSO PMAN Aapop
Figure 6: Log-likelihood comparison on 2D checkerboard ground truth dataset among different models using backward integration.
22

---

Published as a conference paper at ICLR 2026
L_F2D2
VM
div
0.5
L_2202
Lev
F202
ddv
0.0015
0.0014
0.2
2200
4200
4200
4200
6200 Iteration
Iteration
Iteration
L_
L_
Lu
0.30
L.
0.25
0.35
0.20
0.30
0.08
0.00
2200
4200
8200
#200
2200
4200
620n Iteration
R200
Figure The Shortcut-Distill-Pee Distill--222 loss curve on CIFARIO.
Flow-matching
Shortcut-Distie
Chortcut-Dissill-ii
NFE=1
NFE=2
NFE=4
NFE=8
Figure 8: Imagenet 64 x 64 unconditional generation.
23

---

Published as a conference paper at ICLR 2026
Figure 9: 8-step unconditional ImageNet 64 x 64 generation with our Shortcut- -Dittlll
24

---

Published as a conference paper at ICLR 2026
Figure 10: 2-step unconditional ImageNet 64 x 64 generation with our Shortcut- -istill..
25

---

Published as a conference paper at ICLR 2026
Figure 11: 8-step unconditional ImageNet 64 x 64 generation with our Shortcut- Shortcut-Distill-F00 Disttll--2D2..
26

---

Published as a conference paper at ICLR 2026
Figure 12: 2-step unconditional ImageNet 64 x 64 generation with our Shortcut- Shortcut-Distill-Pee Distill-F2D2.
27

---

Published as a conference paper at ICLR 2026
Figure 13: 8-step unconditional CIFAR-10 generation with our Shortcut-Din -isttill.
28

---

Published as a conference paper at ICLR 2026
Figure 14: 2-step unconditional CIFAR-10 10 generation with our Shortcut-Distin -Distill.
29

---

Published as a conference paper at ICLR 2026
Figure 15: 8-step unconditional CIFAR-10 generation with our Shortcut- SrriicuutDsttill-EEe Distill-F2D2.
30

---

Published as a conference paper at ICLR 2026
Figure 16: 2-step unconditional CIFAR-10 10 generation with our Shortcut- Sorrtuu-Distill-Peee Disttill-ttDD
31

---

Published as a conference paper at ICLR 2026
Figure 17: 8-step unconditional CIFAR-10 generation with our Shortcut-F2D2.
32

---

Published as a conference paper at ICLR 2026
Figure 18: 2-step unconditional CIFAR-10 generation with our Shortcut-F2D2.
33

---

Published as a conference paper at ICLR 2026
Figure 19: 2-step unconditional CIFAR-10 generation with McanFlow-F2D2 (Ours).
34

---

Published as a conference paper at ICLR 2026
Figure 20: 1-step unconditional CIFAR-10 generation with McanFlow-F2DD (Ours).
35

---

Published as a conference paper at ICLR 2026
Figure 21: 2-step unconditional CIFAR-10 generation with our McanFlow-F225 F2D2-Self-Guidaaeee
36

---

Published as a conference paper at ICLR 2026
Figure 22: 1-step unconditional CIFAR-10 generation with our HeanFlow-F222-- Se02-eel--Guiancee
37

---

Published as a conference paper at ICLR 2026
F LIMITATIONS AND FUTURE WORKS
In this section, we discuss the limitations and future works of our method. First, F2D2 train- ing requires careful early stopping: once the calibrated BPD value is reached, further training can potentially lead to overfitting or degraded likelihood estimation. Future work could address this through improved network architecture and auxiliary regularization. Second, on the ImageNet- 64 x 64 dataset, due to computational resources constraints, we only train with reduced model size and insufficient training iterations rather than an ideal large-scale configuration. We expect training with longer duration and larger models can further improve the performance. We also restrict our- selves to unconditional generation, which is substantially more challenging than conditional setups. Finally, the performance of our method is sensitive to divergence target scaling and its practical effectiveness in other architectures or modalities remains to be fully validated in future works.
G THE USE OF LARGE LANGUAGE MODELS (LLMS)
We use Large Language Models (LLMs) to refine and polish the manuscript. LLMs also support code debugging, but they not involved in developing the algorithms conducting the exper- iments. The authors take full responsibility for the content of manuscript.
38


Reviews:
[{"review_id": "FqoTTv958K", "reviewer": "Reviewer_tED2", "summary": "This paper introduces a technique to add log probability estimates to popular shortcut distillation techniques. They achieve this by considering the combined sample and log-density differential equation and writing down additional loss terms for the latter part in terms of the divergence of the instantaneous velocity. They specify these terms for two popular frameworks: shortcut models and MeanFlow. The authors then proceed to evaluate their method on CIFAR10 and Imagenet64, reporting NLL and FID for their method and baselines based on both MeanFlow and shortcut models. The authors also present an experiment doing self-guidance for improved FID.", "questions": "- The authors don't seem to to address the question whether the log-density still matches the log-density of the distribution. It seems to me that while the original distribution is is slow to sample from and the log-density similarly slow, at least there is theguarantee that the log-density is linked to the samples of the distribution. In this particular case, both parts (log-density and samples) are approximated, and it seems to me that there is no guarantee that these approximations match. Can the authors comment on whether the approximate log-density calculated is still a relevant number compared to the approximate distribution?\n- What makes NLL predictions \"invalid\" (Tab 1. and Tab 2. gray fields)? Is this an arbitrary threshold that the authors chose or is there something else that makes them invalid? \n- For CIFAR10, why do the authors underperform MeanFlow in some cases, and outperform it in other cases? Since the method operates on a separate head, their method should be fairly independent of performance, can the authors elaborate?\n- For CIFAR10, it seems that the MeanFlow models produce better FID as the number of samples decreases. This is contrary to what I would expect, and also contrary to what is reported in the original paper. Can the authors elaborate why this would be? Can the authors answer the same question for MeanFlow-F2D2?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper, especially up to the the experiments section is very well written, with easy to follow math that is presented with consistent notation throughout.\n- I think the topic is interesting, with potentially important untapped potential of current generative models\n- The method seems sound, and derivations (where applicable) seem correct", "weaknesses": "- My main concern with the paper is that the contribution seems somewhat limited, while it is to the best of my knowledge true that this method has not been applied before, it seems like a relatively straightforward extension of earlier shortcut methods.\n- My other concern is that the experiments present some confusing results, and for a paper targeting log-likelihood, there is a lot of focus on FID in the provided tables, and not a lot of focus on log-likelihood. Other than the small experiment with self-guidance, there are also no experiments tailored specifically to investigate the quality of the approximation of the log-likelihood\n- Continuing on the experimental evaluation: it seems strange to me to evaluate the dataset NLL to the dataset NLL of an expert, since it is comparing averages (i.e. it is a summary statistic). It would in my opinion make more sense to compare likelihoods directly on samples, and calculate a mean of errors, such as the mean squared error or mean absolute error. In the current format, individual samples could have arbitrarily large errors while the error of the mean could still be small.\n- I would have hoped to see an experiment where the log-likelihood is more of a first-class member to see the effect of downstream performance. The authors mention for example the PPO family of methods in their introduction. It could also be instructive to add a toy-experiment in 1D or 2D where the ground-truth log density function is available.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zvz5XHggQO", "reviewer": "Reviewer_5Bkt", "summary": "This paper introduces fast flow joint distillation(F2D2), which jointly learns sampling trajectory and divergence using a joint self-distillation  objective loss and a single model to enable few-step sampling and fast log-likelihood evaluation in flow-based models. It shows that the method is compatible with Shortcut and Meanflow Models and can be further distilled. The proposed method performs on image-based dataset and shows optimistic results.", "questions": "Following the weaknesses:\n\n1. In table 1, I found it is hard to interpret the results as no method performs consistently well than the others across different sampling steps, either it is existing method or the proposed methods. The results from shortcut distilling method outperforms the shortcut model sometimes. I am wondering if there is actually no pattern and the good performance happens by chance.\n2. In Shortcur-Distilll-F2D2, it only talks about the vector field training, how's divergence trained here?\n3. How are invalid NLLs computed in table 1?\n4. I am not pretty sure why we should learn the divergence term from the beginning. Presumably we can learn instantaneous/average vector fields well, and this will induce the marginal density distribution for each $\\rho_t(x)$, then it will automatically the continuity function, which means we know the divergence in the meanwhile. If this is correct, learning divergence is redundant. Please correct me if I am wrong.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. It's interesting to learn vector field and divergence by sharing the same backbone with separate prediction heads since they are related to each other.\n2. The proposed method can be extended to shortcut model and meanflow model.", "weaknesses": "1. The objective loss of shortcut F2D2 is composed of 4 components. It is not quite clear if any one of the components might dominate the whole training and make the performance better. An ablation study is required here. Also, are the 4 components equally weighted? It appears equally weighted from the equation but what if the weights are tuned.\n2. Though the idea to learn velocity field and divergence is interesting, I am wondering if it might lead to instability during training since velocity field learns the directional vector and divergence learns the expectation of the trace through Hutchinson estimator, which involves a gradient of vector field. What if they do not share the same numerical range? Again, it is better to have an ablation study here.\n3. For the experiments, no computation cost and runtime is provided.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rDUyi0dQHQ", "reviewer": "Reviewer_RL2w", "summary": "This paper targets an important problem of making FM based generative models to have fast likelihood estimation.  The proposed method is based on flow map distillation methods in the sampling space, such as short-cut model and meanflow. The authors conduct experiments on CIFAR-10 and ImageNet-64 to demonstrate that their NLL can be effectively calculated in a single step.", "questions": "- Why did the FID drop in Table 1 for meanflow with F2D2 compared to without ?  Limited network capabilities? \n- Could the author clarify why the meanflow-F2D2 is initiated from a short-cut-Distill model ?  This makes it harder to judge the effectiveness of meanflow to your likelihood distillation.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This paper focuses on an important problem and is well motivated, as the author mentioned, fast likelihood computation could open up a wide area of future work, such as in the RL community. \n\n- The writing is clear and easy to follow. \n- The proposed method is simple and straightforward without the need of adversarial training or other additional tricks. \n\n- The result in CIFAR-10 demonstrated that the proposed model can achieve state-of-the-art sampling quality in one step while having good NLL performance.", "weaknesses": "While the proposed approach sounds promising, the experiment results are lacking to fully support the paper's claim. \n\n- The ImageNet 64x64 model's FID is far from state-of-the-art. \n- MF has been reportedly by the community being not stable, adding more loss term could make it more stable for larger dataset/model. The scalability of the proposed approach is not studied. \n- One important baseline Tarflow[1] is missing given it can also do fast NLL calculation and sampling. \n\n[1] Zhai, S., et al. \"Normalizing flows are capable generative models (2025).\" URL https://arxiv. org/abs/2412.06329.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4Lm8bIvhOT", "reviewer": "Reviewer_XjHG", "summary": "The paper proposes Fast Flow Joint Distillation (F2D2), a framework that jointly distills (i) a few-step sampling flow map and (ii) a few-step log-likelihood evaluator for flow-matching models by learning a shared backbone with two heads: an average-velocity head and a cumulative-divergence head. The key idea is that sampling and likelihood evaluation requires solving a coupled ODE system that's driven by the same velocity field, making training a joint backbone neural network possible. The method is evaluated empirically with two families of few-step samplers: Shortcut-F2D2 and MeanFlow-F2D2, plus a practical Shortcut-Distill-F2D2 variant that warms up from a teacher flow. On CIFAR-10 and ImageNet-64, F2D2 is able to produce valid NLLs in a few NFEs and has competitive FID. Paper also proposes a simple maximum-likelihood self-guidance procedure that improves sample quality requiring an additional forward and backward pass.", "questions": "1. Can you provide ablations on parameter sharing vs separate backbones and provide teacher quality for Shortcut-Distill-F2D2? What is the variance of the Hutchinson-based divergence estimator and how does this affect NLL stability and bias?\n2. The experiment results show that MeanFlow-F2D2 has better FID than MeanFlow at low resolutions. MeanFlow conducted experiments on ImageNet 256x256 and the image quality is better. How does F2D2 perform at higher resolutions and will there be any qualitative failure modes when moving beyond 64×64? You mentioned early-stopping sensitivity for training in the Appendix, and this could raise some questions about scaling up training compute.\n3. For Algorithm 2, can you provide experiment results about extending self-guidance beyond one iteration (e.g., x-axis being number of iterations and y-axis being FID/NLL)? Will there be any collapse after one iteration?\n4. What is the total wall-clock time and compute resources used during training the divergence head compared to sampler and the teacher?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- Framing sampling and likelihood evaluation within the flow-map framework for distillation using a shared backbone + two heads to reduce NFE is well-motivated.\n- F2D2 is shown with semigroup Shortcut and Eulerian MeanFlow variants, demonstrating that F2D2 can plug in to different flow-map methods. \n- Writing is clear and easy to follow.\n- Hutchinson trace estimator for divergence, and a staged warm-start are sensible training choices. \n- The maximum-likelihood self-guidance algorithm is simple and provides insight into inference-time scaling for image generation.", "weaknesses": "- The experiment section says F2D2 yields NLLs “close to the teacher’s BPD,” but several entries (e.g., MeanFlow-F2D2 at 2,8 steps) show BPD values materially below the teacher on CIFAR-10. What's more, the FID for MeanFlow-F2D2 is the best among settings. This raises the question of whether or not training on sampling and training on likelihood estimation are mutually beneficial training objectives, and whether or not the approach of using a shared backbone is justified.\n- Using a 1024-step teacher as the “reference” BPD does not certify correctness of the few-step NLLs, it only checks consistency with a particular numerical property. An ablation on toy cases where the true likelihood are easily computed and comparing the F2D2 likelihood with the ground truth likelihood would better establish accuracy.\n- Results are on CIFAR-10 and ImageNet-64 with unconditional generation. Users often cares more about higher resolutions and conditional tasks. The limitations section acknowledges compute constraints, but this still leaves open how the approach scales to large-scale, conditional generation. \n- The paper argues for parameter sharing with dual heads. To validate the claim, it would be informative to see ablations on: (i) shared vs separate backbones, (ii) the dynamics of each loss term.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)