Joint Distillation for Fast Likelihood Evaluation and Sampling in Flow-based Models (F2D2)

TL;DR

F2D2 (fast flow joint distillation) trains a single flow-map network with two prediction heads — a velocity head for sampling and a divergence head for likelihood — so both the sample trajectory and the cumulative log-density change can be evaluated in 1–8 NFEs instead of the 100–1000 NFEs that exact continuous-normalizing-flow (CNF) likelihood normally requires. The construction extends shortcut models and MeanFlow by adding loss terms that enforce flow-map conditions (semigroup or Eulerian) on the divergence component too, with a shared backbone and Hutchinson trace estimation. On CIFAR-10 / ImageNet at 1–8 NFE the joint students recover calibrated BPD close to a 1024-step teacher while keeping competitive FID. As an application, the cheap likelihood signal enables maximum-likelihood self-guidance — a single extra forward+backward pass at inference — that lets a 2-step MeanFlow-F2D2 beat its own 1024-step flow-matching teacher on CIFAR-10 FID without any external reward model.

Key claims

In CNFs the sampling ODE dx/dt = v(x,t) and the log-density ODE dℓ/dt = −div v(x,t) share the same velocity field v, so a single joint flow map Φ_θ(x,ℓ,s,t) = (X_θ, L_θ) suffices to predict both endpoints with shared parameters and a divergence-prediction head [§3.1, Eq. 3.3].
The joint flow map Φ_θ is valid iff it satisfies the tangent condition Φ_θ(x,ℓ,t,t)=(x,ℓ) and any one of (a) Lagrangian, (b) Eulerian, or (c) semigroup condition — generalizing Boffi et al.’s single-flow-map characterization to the coupled (x, ℓ) system [Proposition 3.3].
Shortcut-F2D2 augments the shortcut-model semigroup loss with two divergence-side losses: an instantaneous divergence-matching loss and a semigroup self-consistency loss on the log-density head [Eq. 3.8, 3.9, 3.10].
MeanFlow-F2D2 augments MeanFlow’s Eulerian identity v = u + (t−s) du/dt with the analogous identity for the average divergence; the proof that the MeanFlow identity solves the Eulerian equation is given as Corollary B.2 (claimed to be the first such proof) [§3.2.2, Eq. 3.11–3.13].
Practical training uses a shared backbone with two heads, the Hutchinson estimator for the divergence trace (O(d) → O(1) per step), staged training (pretrain velocity first, then add divergence head), and an optional Shortcut-Distill variant that pre-trains a flow-matching teacher to supervise v [§3.3].
On CIFAR-10 at 4 NFE: Shortcut-Distill-F2D2 gets NLL 2.87 BPD vs Shortcut-Distill’s −26.82 BPD (i.e. the baseline NLL is mathematically invalid in few-step regime); MeanFlow-F2D2 gets NLL 1.31 BPD / FID 4.37 vs MeanFlow’s NLL −22.73 / FID 5.14 [Table 1].
On ImageNet at 4 NFE: Shortcut-Distill-F2D2 reaches NLL 3.94 BPD (teacher BPD 3.34) and FID 24.05, while Shortcut-Distill without F2D2 returns NLL −22.30 (uninterpretable) [Table 2].
Maximum-likelihood self-guidance: with the divergence head + source log-density giving a one-step pseudo-likelihood, gradient descent on the initial noise z₀ ← z₀ + η ∇ log p_θ(x) before sampling lets a 2-step MeanFlow-F2D2 outperform a 1024-step flow-matching teacher of the same size on CIFAR-10 FID with one extra forward+backward NFE [§3.4, §5.3, Fig. 2].
On 2D checkerboard (analytic ground truth) Shortcut-F2D2 recovers the correct density at 1 NFE while plain shortcut models and flow matching catastrophically fail at few-step likelihood [§5.4, Fig. 3].

Method

The core construction is the joint parametrization in Eq. 3.3: X_θ(x,s,t) = x + (t−s) · v_θ(x,s,t), L_θ(x,ℓ,s,t) = ℓ + (t−s) · u_θ(x,s,t), where v_θ is the average velocity (as in shortcut/MeanFlow) and u_θ is the average divergence — note u_θ depends on x but not on ℓ, because in the joint ODE dℓ/dt = −div v(x,t) the right-hand side never reads ℓ. Both heads share a single transformer backbone.

Training enforces four losses (Eq. 3.4): tangent + flow-map condition on (v_θ, X_θ) for sampling, and the analogous tangent + flow-map condition on (u_θ, L_θ) for likelihood. The two instantiations differ only in which flow-map condition:

Shortcut-F2D2 uses the semigroup property: Φ(Φ(x,ℓ,s,m),m,t) = Φ(x,ℓ,s,t) for m∈[s,t], yielding self-consistency losses on both heads (Eq. 3.5 for v, Eq. 3.9 for u). Tangent supervision comes from the flow-matching loss for v and from a Hutchinson estimate of div v_θ for u (Eq. 3.8).
MeanFlow-F2D2 uses the Eulerian identity: v = u_instant + (t−s) ∂v/∂t, computed via a JVP. Eq. 3.12 is the divergence-side analogue; because the MeanFlow loss also encapsulates the instantaneous regression at s=t, no separate L_inst is needed.

Hutchinson trace estimation is critical: exact divergence costs d backward passes per training step in d-dimensional space (intractable for images), but ε^T J_v ε with random ε gives an unbiased O(1) estimate. Staged training (Shortcut-Distill-F2D2) pretrains a flow-matching teacher, warm-starts the shortcut velocity from it, then adds and jointly trains the divergence head.

For maximum-likelihood self-guidance (Algorithm 2), the one-step prediction L_θ(x_0, 0, 0, 1) + log p_source(x_0) gives a cheap log p_θ(x_1) ≈ log p_θ(sample). One gradient step on x_0 (or equivalently on the noise) before unrolling the sampler is the entire test-time intervention.

Results

CIFAR-10, 2-step: MeanFlow-F2D2 NLL 1.29 BPD / FID 2.59 vs MeanFlow NLL −47.31 / FID 2.84 — likelihood becomes valid at no FID cost [Table 1].
CIFAR-10, 1-step: MeanFlow-F2D2 NLL 3.33 / FID 3.02 vs teacher BPD ~3.0 — calibrated likelihood at a single NFE [Table 1].
ImageNet, 8-step: Shortcut-Distill-F2D2 NLL 3.51 BPD (teacher 3.34) / FID 21.91 vs Shortcut-Distill alone NLL −9.03 (invalid) / FID 19.47 — small FID regression buys interpretable likelihood [Table 2].
Maximum-likelihood self-guidance: 2-step MeanFlow-F2D2 + self-guidance beats 1024-step flow-matching teacher (same size) on CIFAR-10 FID; FID improves monotonically with self-guidance step count up to a point [§5.3, Fig. 2].
2D checkerboard: at 1 NFE Shortcut-F2D2 reconstructs the analytic density distribution with correct spatial structure and density values, while flow matching and plain shortcut models return numerically degenerate likelihoods [§5.4, Fig. 3].
All baselines (flow matching, shortcut, MeanFlow) produce mathematically invalid NLL (large negative BPD) in the few-step regime — they can compute the formula but the result is meaningless because the divergence integral was discretized over too few steps; F2D2 is the first method to make few-step NLL valid, not just fast [Tables 1, 2].

Why it’s interesting

This is the missing companion to the few-step-sampling literature: the wiki has many entries on fast sampling (Mean Flows for One-step Generative Modeling, ArcFlow: Unleashing 2-Step Text-to-Image Generation via High-Precision Non-Linear Flow Distillation, Diversity-Preserved Distribution Matching Distillation for Fast Visual Synthesis, pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation, Test-time scaling of diffusions with flow maps) but none of them recover tractable likelihoods at few NFE — F2D2 closes that gap by treating the divergence as a second head over the same flow map rather than as a separate integration problem. The self-guidance result is the conceptually striking part: an unconditional generator’s own likelihood, made cheap by joint distillation, is a strong enough signal that a 2-NFE student beats its 1024-NFE teacher with one extra backward pass — this complements Noise Hypernetworks: Amortizing Test-Time Compute in Diffusion Models (noise hypernetworks amortize external reward at test time) and VLM-Guided Adaptive Negative Prompting for Creative Generation (VLM as test-time guide) by showing that the model can guide itself once likelihoods are cheap. Compare also Test-time scaling of diffusions with flow maps which uses flow maps for test-time sampling scaling — F2D2 is the likelihood analogue and uses the same theoretical machinery (Boffi et al.’s flow-map characterization, extended to coupled ODEs).