Source-linked AI summary
Flow Map Distillation Without Data
Shangyuan Tong, Nanye Ma, Saining Xie, Tommi Jaakkola
TL;DR
Flow-map distillation commonly relies on external data, although those samples may not represent the teacher’s intermediate states. FreeFlow predicts and corrects the teacher’s path using prior samples alone. On ImageNet, it achieves state-of-the-art FID scores at one sampling step.
Problem
Conventional distillation trains on dataset-derived intermediate states that may differ from the states along the teacher’s generating paths.
Source: §3.1 · p. 3Method
FreeFlow samples from the prior and combines prediction of the teacher’s path with correction using the student’s noising-flow velocity.
Source: §4.1 · p. 5, §4.2 · p. 6Results
FID 1.45 on ImageNet 256×256 and 1.49 on ImageNet 512×512, with 1-NFE, significantly surpasses data-based baselines.
Source: §1 · p. 2Takeaways & Limitations
The experiments support high-fidelity flow-map distillation without an external dataset, avoiding Teacher-Data Mismatch by construction.
Source: §1 · p. 2, §6 · p. 8The comprehensive correction design for the full predicted trajectory was not helpful in the experiment settings considered.
Source: §4.2 · p. 6
Abstract
from arXiv · showhide
State-of-the-art flow models achieve remarkable quality but require slow, iterative sampling. To accelerate this, flow maps can be distilled from pre-trained teachers, a procedure that conventionally requires sampling from an external dataset. We argue that this data-dependency introduces a fundamental risk of Teacher-Data Mismatch, as a static dataset may provide an incomplete or even misaligned representation of the teacher's full generative capabilities. This leads us to question whether this reliance on data is truly necessary for successful flow map distillation. In this work, we explore a data-free alternative that samples only from the prior distribution, a distribution the teacher is guaranteed to follow by construction, thereby circumventing the mismatch risk entirely. To demonstrate the practical viability of this philosophy, we introduce a principled framework that learns to predict the teacher's sampling path while actively correcting for its own compounding errors to ensure high fidelity. Our approach surpasses all data-based counterparts and establishes a new state-of-the-art by a significant margin. Specifically, distilling from SiT-XL/2+REPA, our method reaches an impressive FID of 1.45 on ImageNet 256x256, and 1.49 on ImageNet 512x512, both with only 1 sampling step. We hope our work establishes a more robust paradigm for accelerating generative models and motivates the broader adoption of flow map distillation without data.
1. Introduction
Flow maps can accelerate ODE-based generation, but conventional distillation depends on external data that may misrepresent a teacher’s generative capabilities. FreeFlow instead learns from the prior and combines prediction with correction; on ImageNet it reaches state-of-the-art FID scores with one sampling step.
1. Introduction
- Flow maps accelerate ODE-based generation by learning the solution operator directly and taking large jumps instead of iteratively solving the ODE.Distilling from pre-trained teachers offers a modular route to compressing models built with advanced training and post-training techniques.Source: §1 · p. 1
- Conventional data-based distillation uses external-dataset samples, which may not represent the teacher’s full generative distribution and can create Teacher-Data Mismatch.This can occur when teachers generalize beyond their training set, shift after post-hoc fine-tuning, or have unavailable proprietary training data.Source: §1 · pp. 1–2
- Sampling only from the prior avoids mismatch by construction because the prior is the shared, guaranteed-aligned origin of the teacher’s generation.The prior is also the termination point of the noising process.Source: §1 · p. 2
- FreeFlow predicts the teacher’s sampling path from a prior sample and integration interval, aligning its generating velocity with the teacher’s instantaneous velocity.A correction mechanism also aligns the student’s noising velocity with the teacher to address compounding prediction errors.Source: §1 · p. 2
- FID 1.45 on ImageNet 256×256 and 1.49 on ImageNet 512×512, with 1-NFE, establish FreeFlow’s state-of-the-art performance and surpass all data-based baselines.The teacher is SiT-XL/2+REPA [100].Source: §1 · p. 2
2. Preliminaries
Flow and diffusion models reverse a noising process that transports data toward an easy-to-sample prior. Sampling integrates the flow ODE backward, while flow maps approximate its solution operator directly to reduce iterative evaluations.
2. Preliminaries
- Diffusion and flow models reverse a noising process from data distribution p to prior π; the paper uses linear interpolants xt=(1−t)x+tz and conditional velocity x−z.The interpolants combine a data sample x with a prior sample z.Source: §2 · p. 2
- Sampling integrates the flow ODE backward, and standard generation evaluates ϕu(z,1,0); numerical solvers may require dozens or hundreds of NFEs.The generating flow ϕu is equipped with velocity field u, and the integration interval is t−s.Source: §2 · p. 2
- Flow maps approximate the generating flow ϕu directly rather than the instantaneous velocity u, enabling students to reproduce the teacher’s process with fewer NFEs.In MeanFlow, Fθ represents the average velocity traveled over the path; flow-map distillation uses a pre-trained flow model as teacher.Source: §2 · pp. 2–3
- With the teacher fixed, stronger data augmentation increases Teacher-Data Mismatch and degrades student performance.Source: Figure 2 · p. 3
3. With or Without Data?
Flow map distillation aims to reproduce a teacher’s generative process with fewer function evaluations, but conventional data-noising states may not match the teacher’s actual paths. The section motivates prior-only distillation as an alternative.
3.1. The Risk of Teacher-Data Mismatch
- Teacher-Data Mismatch occurs when data-noising states differ from teacher-path states, so training on the former may leave parts of the teacher’s behavior unlearned.Even a perfectly converged student is not guaranteed to reproduce the teacher’s outputs when trained on trajectories anchored to a static dataset.Source: §3.1 · p. 3
- Stronger data augmentation increases the discrepancy between teacher-generating and data-noising distributions and degrades the learned flow map’s performance.The controlled experiment deliberately misaligns the training distributions through data augmentations.Source: §3.1 · p. 3
- Teacher generalization or classifier-free guidance, post-hoc fine-tuning, and unavailable proprietary training data are practical settings in which dataset and teacher trajectories can diverge.When teacher training data is unavailable, a proxy dataset may create severe mismatch.Source: §3.1 · p. 3
3.2. Towards a Data-Free Alternative
- Directly sampling teacher trajectories for distillation requires costly on-the-fly integration, while offline precomputation scales poorly and can be intractable for complex tasks.The proposed direct-sampling remedy starts from the prior and integrates the teacher to a random time.Source: §3.2 · p. 4
- At t = 1, the teacher-generating and data-noising distributions coincide at the prior, making it a guaranteed on-distribution source for the teacher.The distributions may diverge for t in [0, 1), but both equal the prior at t = 1.Source: §3.2 · p. 4
- The paper argues that flow map distillation can instead use an objective governed only by prior samples, avoiding reliance on an external dataset.This data-free alternative is motivated by the prior’s guaranteed alignment with the teacher’s generative process.Source: §3.2 · p. 4
4. Flow Map Distillation Without Data
The method anchors flow-map distillation at the prior, predicts the teacher’s trajectory through velocity alignment, and corrects accumulated prediction errors with a prior-only objective.
4. Flow Map Distillation Without Data
- Prior-only sampling fixes the start state at t=1, so the method differentiates the optimal condition with respect to the predicted end time instead of perturbing the start point.Source: §4 · p. 4
4.1. Predict With Generating Flows
- The student learns to align its generating velocity with the teacher’s velocity at its own predicted state, extending outward from the prior.Source: §4.1 · p. 5
- Prediction errors can compound as the student queries the teacher at its own deviating state, progressively diverging from the true trajectory.Source: §4.1 · p. 5, Figure 4 · p. 5
4.2. Correct With Noising Flows
- The correction objective aligns the student’s noising velocity with the teacher’s velocity and uses only prior samples, avoiding Teacher-Data Mismatch.Source: Figure 5 · p. 5, §4.2 · p. 6
- Although full-trajectory correction would be comprehensive, the authors did not find it helpful in the experimental settings considered.Source: §4.2 · p. 6
5. Experiments
Experiments show FreeFlow's data-free distillation outperforms prior methods on ImageNet. Paired prediction and correction support high-fidelity generation, while a one-step proxy enables inference-time scaling.
5.1. Design Decisions
- Higher-noise emphasis improves alignment between vN and u, while dropping higher noise levels under uniform r sampling worsens results.The authors motivate emphasizing high noise because differences in probability flux accumulate over time.Source: §5.1 · p. 6, Table 1 · p. 6
- Correction-objective guidance differs from sampling guidance; a more aggressive guidance interval performs better, and SiT-B/2 does not benefit from γ = 2.The correction objective drops guidance at high noise levels.Source: §5.1 · pp. 6–7, Table 1 · p. 6
- Training splits batches 75% for prediction and 25% for correction, then adaptively balances correction gradients; performance remains robust across α.Correction is slightly more expensive computationally.Source: §5.1 · p. 7, Figure 6 · p. 7
- Stronger decay on ΔvG,u improves joint training, which the authors hypothesize dampens conflicts between prediction and correction signals.The weighting applies a power-law decay to the prediction-objective gradient.Source: §5.1 · pp. 6–7
5.2. Main Results
- FID 1.45 at 256×256 and 1.49 at 512×512 let FreeFlow surpass prior proposals when distilling SiT-XL/2+REPA, using 1-NFE.Evaluation uses FID-50K on ImageNet at both resolutions.Source: §5 · p. 6, §5.2 · p. 7, Table 2 · p. 8
- Combining prediction and correction outperforms either objective alone: prediction constructs the path while correction stabilizes compounding errors.Prediction alone plateaus at suboptimal fidelity; correction alone leads to gradual mode collapse and performance degradation.Source: §5.2 · pp. 7–8, Figure 7 · p. 7
- With a total budget of 80 NFEs, proxy-guided noise search outperforms standard teacher classifier-free guidance sampling at 128 NFEs.The one-step student searches for transferable noise candidates, which the teacher then uses for generation.Source: Figure 8 · p. 7, §5.2 · p. 8
6. Conclusion
The conclusion frames external-dataset dependence as an unnecessary vulnerability in flow-map distillation. It presents prior-only predictor-corrector distillation as a robust alternative supported by strong empirical results.
6. Conclusion
- A static dataset can be an unreliable proxy for a teacher’s full generative capabilities, creating Teacher-Data Mismatch.Source: §6 · p. 8
- The proposed predictor-corrector framework samples only from the prior, resolving the mismatch by construction.Source: §6 · p. 8
- The paper reports strong empirical results that establish the practical viability of data-free flow-map distillation.Source: §6 · p. 8
A.1. More on the Prediction Objective
The prediction objective trains the student’s flow map to follow the teacher’s velocity field, with discrete-time training approximating that condition through finite differences. Its error analysis separates discretization error from model approximation error, while the correction objective makes performance less sensitive to discretization choices.
A.1. More on the Prediction Objective
- At optimality, the student’s path derivative matches the teacher velocity at every prior sample and time; with the stated regularity conditions, this identifies the teacher’s flow map.The criterion also requires fθ(z, 0) = z.Source: §A.1 · p. 13
- Finite differences over N time intervals approximate the student’s generating velocity, yielding a discrete-time objective that mimics Euler integration without requiring efficient JVP implementations.The objective approaches its continuous-time counterpart as the interval size tends to zero.Source: §A.1 · pp. 13–14
- The total trajectory error combines network approximation error with solver discretization error, and this total error controls the Wasserstein distance between student and teacher distributions.For local truncation error O((δ_i,i+1)^(p+1)), the global discretization error is O((δ_max)^p).Source: §A.1 · p. 14
- More discretization steps can degrade performance as model approximation error accumulates, but the correction signal makes results robust to step count and solver precision.The correction objective’s robustness to discretization steps and solver type is also noted in Figure 9.Source: §A.1 · p. 14, Figure 9 · p. 14
- A confident-region warmup adds noise to predicted states before evaluating the teacher velocity, with the transition level reduced from 1 to 0 over 10K steps.After warmup, the authors report no training instability with a reasonable δ sampling distribution.Source: §A.1 · pp. 14–15
A.2. More on the Correction Objective
The correction objective is derived as an alignment between the student’s and teacher’s noising velocities. Its training signal depends on estimating the student velocity with an auxiliary model and on suitable learning-rate and sampling choices.
A.2. More on the Correction Objective
- The correction objective can be expressed as aligning the student’s marginal noising velocity with the teacher’s, using the student’s predicted sample and a noising interpolation.The derivation connects score functions to marginal velocities and yields Eq. (11).Source: §A.2 · p. 15
- The auxiliary model gψ estimates the student’s noising-flow velocity; a higher learning rate for gψ improves this signal, while correction training uses 3 × 10−5 for fθ versus 1 × 10−4 in SiT.A two-time-scale update is mentioned but not explored because of its overhead.Source: §A.2 · p. 15, Table 4 · p. 15
- The authors sample r more often at higher noise levels and note that velocity differences should be on a consistent scale across r for sampling to control their contributions.For the EDM parameterization, the magnitude of the difference varies with cout(r).Source: §A.2 · pp. 15–16
- An additional weighting proportional to 0.5r / √(0.5² + r²) heavily downplays contributions from small r.The passage describes this weighting in connection with prior work’s treatment of Eq. (11).Source: §A.2 · p. 16
A.3. More on How to Combine Both
The section reports that prediction and correction work synergistically and describes a schedule for balancing their training signals. The schedule delays and warms up α before its later decay, with its value interpreted in light of the batch split.
A.3. More on How to Combine Both
- For REPA distillation, α is set to zero for 10K steps, linearly warmed up for another 10K, then follows an inverse-square-root decay toward αref.The authors report that larger α tends to learn faster early, while smaller α tends to converge better later.Source: §A.3 · p. 16
- With the 75–25 prediction–correction batch split, α = 0.3 corresponds to a correction-gradient contribution around 10% of the prediction-gradient contribution.Source: §A.3 · p. 16
- At 1.5M iterations, prediction alone still lagged and correction alone had suffered mode collapse, while the combined objectives achieved performance neither attained in isolation.The comparison uses SiT-XL/2 distillations evaluated by FID and Inception Score.Source: §A.3 · p. 16
B. Implementation Details
The implementation uses a DiT-based model with additional conditioning for jump duration and guidance strength, and evaluates reported models through a common GPU-based pipeline. The section also identifies the combined objectives’ empirical synergy and points to detailed configurations for the main results.
B. Implementation Details
- The reported comparison attributes performance to the synergy of prediction and correction objectives, which together achieve performance neither reaches alone.Source: Table 5 · p. 16
- Training uses VAE latent space and a DiT-based architecture with 2×2 patches; scalar embeddings for time, class, jump duration, and guidance strength are combined and fed to each block.The auxiliary gψ also takes time, class, and guidance strength as inputs.Source: B. Implementation Details · p. 16
- Models are trained with an internal JAX codebase on TPU, while reported XL results are converted to PyTorch and evaluated on GPU using ADM’s evaluation suite.The implementation configurations for reported entries are listed in Table 6.Source: B. Implementation Details · p. 17, Table 6 · p. 17
- Because the model is trained over γ ∼ U(1, 2), inference can sweep guidance strengths efficiently; the paper reports the best FID and supplies the full curves.The curves are identified in Figure 10.Source: B. Implementation Details · p. 17
C. Related Work
Prior work accelerates generative models by distilling trajectories or matching distributions. FreeFlow contributes a data-free trajectory-distillation algorithm and distribution-matching design choices through a predictor-corrector framework.
C. Related Work
- Unlike BOOT, which frames data-free distillation as a practical response to storage and privacy burdens, FreeFlow presents it as necessary for distributional fidelity.The paper also distinguishes its approach from BOOT’s signal-ODE parameterization, whose boundary condition requires a separate loss.Source: C. Related Work · p. 18
- FreeFlow’s average-velocity objective satisfies the boundary condition by design and avoids costly partial-derivative gradients used in broader ODE formulations [81].The paper attributes improved training efficiency to this perspective and an identity derived in Eq. (7).Source: C. Related Work · p. 18
- An auxiliary correction objective addresses accumulated prediction errors, and the paper reports performance beyond the then-current state of the art.The related-work discussion also positions the framework as combining strengths of trajectory- and distribution-level methods.Source: C. Related Work · p. 18
- FreeFlow spans trajectory- and distribution-level acceleration with data-free trajectory distillation and a predictor-corrector framework intended to retain quality and diversity.Trajectory methods compress ODE integration, while distribution-level methods match teacher marginals without necessarily following its trajectory.Source: C. Related Work · p. 18
D. Additional Visual Results
The appendix presents additional uncurated FreeFlow-XL/2 samples at 512×512 resolution with 1-NFE. The authors state that training used the SiT-XL/2+REPA teacher without querying ImageNet samples.
D. Additional Visual Results
- FreeFlow-XL/2 generates additional uncurated 512×512 samples with 1-NFE, using the SiT-XL/2+REPA teacher without querying ImageNet during training.Figures 11–18 contain these samples.Source: D. Additional Visual Results · p. 18
- The displayed class labels include robin, Siberian husky, barn, and desktop computer.These correspond to Figures 11–14.Source: Figure 11 · p. 19, Figure 12 · p. 19, Figure 13 · p. 20
- The displayed class labels also include dogsled, fire truck, lakeside, and volcano.These correspond to Figures 15–18.Source: Figure 15 · p. 21, Figure 16 · p. 21, Figure 17 · p. 22
Related papers
This paper cites
- Attention Is All You Need2017
- ADAM: A METHOD FOR STOCHASTIC OPTIMIZATION2014
- Score-Based Generative Modeling through Stochastic Differential Equations2020
- Classifier-Free Diffusion Guidance2022
- Diffusion Models Beat GANs on Image Synthesis2021
- Deep Unsupervised Learning using Nonequilibrium Thermodynamics2015