From Navigation to Refinement: Revealing the Two-Stage Nature of Flow-based Diffusion Models through Oracle Velocity

AI summary

Overview

Research area: Machine learning / generative modeling — specifically the training dynamics of flow-based diffusion models (flow matching, rectified flow), and the balance between memorization and generalization.

Technical level: Advanced. The paper derives a closed-form expression for the marginal velocity field and reasons about Gaussian mixtures, posterior weights, and probability flow dynamics, though the high-level conclusions are stated plainly.

Scope: A theoretical-plus-empirical study that derives a closed-form "oracle velocity" for flow matching and uses it to show that the effective training target splits into a navigation stage (near the prior) and a refinement stage (near the data), then uses this lens to explain practical techniques such as timestep shifting, classifier-free guidance intervals, and latent space choice.

What This Paper Is About

Flow-based diffusion models are trained to predict a velocity field, but it is hard to say what the model is actually being asked to learn at different points along the trajectory, and why it sometimes reproduces training images (memorization) and sometimes invents new ones (generalization). The authors show that under a Gaussian prior with a finite dataset and linear interpolation, the otherwise intractable marginal velocity field has an exact closed form, which can be computed and inspected at any point in the sample space. Analyzing this oracle reveals that training is not one uniform task but two distinct ones.

Key Contributions

1. A closed-form oracle velocity (Theorem 2.1). Under a Gaussian prior, a finite dataset of N samples, and the rectified-flow linear interpolation, the oracle velocity u*_t(x_t, t) = E[u_t(x_t | x_1) | x_t] is written as A_t · Σ_i γ_i(x_t, t) x_1^(i) + B_t x_t, where γ_i are normalized Gaussian posterior weights. This makes the FM target exactly computable.

2. Identification of the two-stage training target. The oracle target is shown to be a mixture of data modes in an early interval (navigation) and dominated by a single nearest data sample later (refinement), with the transition around t ≈ 0.1 on ImageNet-scale latent data.

3. A stage-level explanation of memorization versus generalization. Using mixed sampling that starts with the oracle and switches to a trained model at a threshold t_switch, the authors attribute generalization to the navigation stage and memorization to the refinement stage.

4. Explanations of three practical techniques. The framework is used to interpret timestep-shifted schedules, CFG interval selection, and latent space design choices, with supporting gFID ablations.

Main Findings

• The target splits into two regimes. On ImageNet 256² resolution data with latent dimensionality D ∈ {4096, 8192} and N ≈ 1400, the MSE between the oracle target and the CFM target (x_1 − x_0) is concentrated in t ∈ [0.0, 0.1]; the two targets align closely for t > 0.1. The average top-1 posterior weight γ_i(x_t, t) saturates to 1 beyond t = 0.1, indicating collapse onto a single dominant data point.

• Dimensionality and sample size set the transition point. Higher data dimensionality accelerates saturation of the top-1 posterior weight because squared distances scale with D and enter the exponent divided by 2σ_t²; a larger sample size delays the transition. The paper states this means the effective training target naturally varies across datasets even under the same rectified flow objective.

• Model behavior transitions slightly later than the target. Visualizing LightningDiT-XL/1 intermediate predictions (a single Euler step from t' ∈ [0, 1] to t = 1), global layout stabilizes around t ≈ 0.2, lagging the training-target shift of t ≈ 0.1. The authors suggest the model needs temporal margin to rectify accumulated prediction errors.

• Generalization comes from navigation, memorization from refinement. When the oracle is used up to t_switch ∈ [0.0, 0.2], the strong prior corruption prevents recovery of the training trajectory and outputs are novel. When t_switch ∈ (0.2, 1.0], the model largely replicates training trajectories. Running the oracle for the whole trajectory deterministically retrieves a training sample. The paper also reports that memorization from scratch is highly unlikely when the dataset size substantially exceeds the model's effective memorization capacity, whereas resuming from refinement-stage trajectories readily triggers memorization.

• Refinement transfers to unseen latents. Combining Gaussian priors with ImageNet validation latents and resuming at t_resume ≥ 0.2, the model preserves the global layout while improvising fine details.

• Learning difficulty differs by stage. Training MSE against the conditional target and the class-conditioned oracle target diverges mainly in the navigation stage; the curves align in the refinement stage. Oracle loss is near-perfect as t → 0 (trivial class-mean prediction), rises through navigation, bumps at the transition, dips in mid-refinement, and rises again as t → 1. Across model sizes, all models perform nearly identically in navigation, while larger models and longer training produce lower oracle loss in refinement. Most models were trained for 100 epochs on ImageNet; one model marked with an asterisk was trained for 800 epochs (8× training compute).

• Mixed-model generation confirms capacity asymmetry. With LightningDiT Base and XL and 25 uniform sampling steps per sub-interval (no CFG), gFID@50K for [0.0, 0.1] + [0.1, 1.0] is 2.94 (XL/XL), 3.71 (Base/XL), 11.26 (XL/Base), 12.45 (Base/Base). For [0.0, 0.2] + [0.2, 1.0] it is 2.60 (XL/XL), 4.47 (Base/XL), 9.24 (XL/Base), 12.01 (Base/Base). Swapping the Stage 1 model for a weaker one costs little; swapping Stage 2 costs a lot.

• Timestep shifting: modest extra navigation steps help most. With LightningDiT-B/1, NFE = 50, no CFG, gFID@50K by shift factor: s = 4.0 (6% of t ∈ [0.0, 0.2]) 18.89; s = 2.0 (12%) 14.82; s = 1.0 uniform (22%) 12.99; s = 0.7 (28%) 12.46; s = 0.5 (34%) 12.23; s = 0.3 (46%) 12.66; s = 0.1 (72%) 19.91. The best gFID is at s = 0.5.

• CFG interval: best results span early-to-mid refinement. With LightningDiT-B/1 and CFG factor ω = 2.5, gFID@50K is 12.99 with no CFG and 10.79 for [0.0, 1.0]. Best values: [0.1, 0.6] 2.80, [0.1, 0.5] 2.82, [0.1, 0.7] 2.86, [0.1, 0.8] 2.97, [0.1, 0.3] 3.54, [0.1, 0.4] 4.16. Single 0.1-wide intervals degrade as they move later: [0.1, 0.2] 5.21, [0.2, 0.3] 7.70, [0.3, 0.4] 9.39, [0.4, 0.5] 10.39, [0.5, 0.6] 10.59, [0.6, 0.7] 11.06, [0.7, 0.8] 11.38, [0.8, 0.9] 11.64, [0.9, 1.0] 12.20. Applying CFG at the very earliest steps is harmful.

• Latent space structure matters. Comparing VA-VAE (aligned with semantic structure via DINO-based VF loss) and SD-VAE (trained purely for low-level reconstruction), both show a gradual oracle loss increase in navigation, but refinement differs: VA-VAE converges smoothly toward a parabola-like shape, while SD-VAE is flatter with noticeable waviness. The paper notes absolute magnitudes are not directly comparable.

• Oracle-supervised training works but does not scale. The authors report that training directly on the oracle target is feasible and converges similarly to CFM, but is impractical at scale because of the computational and I/O overhead of evaluating the oracle over the full dataset.

Methodology in Plain English

Instead of treating the flow matching target as unknowable, the authors impose three mild assumptions: the prior is Gaussian, the data is a finite set of samples, and the interpolation between prior and data is linear (rectified flow). Under these, the path marginal becomes a Gaussian mixture, and applying Bayes' rule lets them write the expected velocity given the noisy sample x_t in closed form. That expression is essentially a soft nearest-neighbor weighted average over data points, plus a term linear in x_t.

Because this oracle can be computed exactly, they can (a) measure how far the standard per-sample training target (x_1 − x_0) deviates from the true target at each timestep, (b) track how the posterior weight concentrates on a single data point as time progresses, and (c) run controlled sampling experiments where the oracle drives early steps and a real trained model (LightningDiT-XL/1, LightningDiT-B/1, and in the appendix Flux.1 [dev]) drives later steps. They also visualize intermediate one-step predictions of a trained model to locate the point where behavior shifts, measure training loss against both the conditional and oracle targets across timesteps and model sizes, and run gFID ablations over timestep shift factors and CFG intervals.

Why This Matters

Impact on research. The paper supplies a concrete, computable object (the oracle velocity) for reasoning about what diffusion training targets actually are at each timestep, extending prior closed-form analyses of diffusion formulations to rectified flow and using it specifically to dissect memorization versus generalization at ImageNet scale — a regime where the authors note earlier small-scale studies are less informative. It also argues that navigation quality is a bottleneck that scaling model capacity and compute does not fix, and that standard FID gives little insight into intermediate layout fidelity.

Real-world applications:

• Text-to-image and video generation pipelines, where timestep schedules, guidance intervals, and latent encoders are tuned heuristically today.
• Dataset and training-budget planning, since the stage transition depends on data dimensionality and sample size and therefore shifts across resolutions and corpus sizes.
• Generative model auditing and copyright/privacy review, where understanding when a model reproduces training data versus when it invents new content is directly relevant.
• Generative model post-processing or editing workflows that resume sampling midway, e.g. latent-based refinement of an existing image.

Industry relevance. Timestep shifting, CFG interval selection, and VAE latent choice are core knobs in production diffusion systems. The paper offers an explanation for why the commonly used non-uniform schedules and partial-guidance intervals work, and where their optimal settings sit relative to the navigation-to-refinement transition. The finding that refinement — not navigation — consumes additional model capacity and training compute is directly relevant to compute allocation decisions.

Future Directions

• Improving navigation capability. The authors state that navigation performance under FM/CFM remains largely unchanged after scaling model capacity and training compute, and that it is unclear whether the near-prior prediction task is simply trivial or whether the per-sample CFM target provides insufficient supervision for larger models.
• New metrics for intermediate prediction quality. The paper argues that generation FID provides little insight into early-stage layout fidelity and that new measures are needed.
• Stage-aware training schedules. The closed-form oracle is said to identify a theoretical stage-shift point that could allow more training steps to be allocated to refinement for faster convergence; the appendix explores stage-aware training timestep sampling.
• Scaling behavior. Because the stage transition depends on data dimensionality and sample size, the authors suggest that scaling along these axes reshapes navigation-refinement dynamics, and that state-of-the-art systems scale data, capacity, and compute simultaneously.

Target Audience

Researchers and engineers working on diffusion and flow-based generative models who want a principled account of training dynamics — particularly those tuning sampling schedules, guidance intervals, or latent space designs, and those studying memorization versus generalization in large-scale generative models. Readers need comfort with probability flow ODEs, Gaussian mixtures, and classifier-free guidance to follow the derivations, though the stage-level conclusions and practical recommendations are stated accessibly.

Authors’ abstract

Flow-based diffusion models have emerged as a leading paradigm for training generative models across images and videos. However, their memorization-generalization behavior remains poorly understood. In this work, we revisit the flow matching (FM) objective and study its marginal velocity field, which admits a closed-form expression, allowing exact computation of the oracle FM target. Analyzing this oracle velocity field reveals that flow-based diffusion models inherently formulate a two-stage training target: an early stage guided by a mixture of data modes, and a later stage dominated by the nearest data sample. The two-stage objective leads to distinct learning behaviors: the early navigation stage generalizes across data modes to form global layouts, whereas the later refinement stage increasingly memorizes fine-grained details. Leveraging these insights, we explain the effectiveness of practical techniques such as timestep-shifted schedules, classifier-free guidance intervals, and latent space design choices. Our study deepens the understanding of diffusion model training dynamics and offers principles for guiding future architectural and algorithmic improvements. Our project page is available at: https://maps-research.github.io/from-navigation-to-refinement/.