Velocity Scaling in Flow Matching

TL;DR

Velocity scaling multiplies a trained flow-matching velocity field $v_\theta(x_t, t)$ by a time-dependent gain $\gamma(t)$ at inference. Prior work motivated this by claiming MSE training systematically underestimates velocity magnitude. Saied & Fleuret show that story is wrong: MSE does not create a magnitude deficit. Instead, coarse integration causes population time lag — the sampled state at model time $t$ resembles training states from an earlier time — and scaling the velocity lengthens each step so samples catch up. A tiny probe estimates the current lag and picks the gain that minimizes it; pushing slightly past the lag-correction optimum sharpens the sample covariance for further quality gains, trading precision for recall. On SiT-XL/2 at NFE 25, lag correction cuts FID from 28.0 to 12.2, and the extra push drops it to 8.7 [§abstract, Fig 5 in thread].

Key claims

MSE training of flow-matching velocity fields does not create a velocity-magnitude deficit, contrary to the motivation offered by prior velocity-scaling work [§abstract].
The actual effect velocity scaling corrects is population time lag: with coarse integration, the sample at model time $t$ belongs to the distribution of an earlier training time; multiplying the velocity by $\gamma > 1$ lengthens each step so the sample catches up [§abstract, tweet 2/8].
A small probe network can estimate which training-time a given sample belongs to; selecting $\gamma(t)$ to minimize this probed lag improves generation quality without decoding, inspired by the “Blessings of Dimensionality” line of work [tweet 3/8].
The optimal gain depends jointly on the velocity field and the number of integration steps. With more NFE the required correction shrinks toward $\gamma = 1$ ; a well-trained small model (CIFAR-10) barely lags at all so any gain other than 1 hurts [tweet 4/8].
Pushing the gain further than lag correction demands sharpens the sample covariance, improving FID beyond what catch-up alone gives: on SiT-XL/2 at NFE 25, lag correction gives FID 12.2 and the extra push $\kappa$ gives FID 8.7, with the same $\kappa$ tuned at 25 steps transferring near-optimally to 50/100/250 steps [tweet 5/8, tweet 6/8].
The extra sharpening trades precision for recall: samples become more faithful but less diverse [tweet 6/8].
Velocity scaling composes with CFG and is partially redundant with it. CFG already sharpens the sample covariance, so stronger CFG requires less velocity scaling, and under strong guidance scaling the velocity down (gain < 1) can improve FID [tweet 7/8].
Across architectures and model sizes, measuring population time lag and using it to select a gain reduces ImageNet-256 FID from 28.0 to 12.2 at NFE 25 without guidance [§abstract].

Method

Flow-matching sampling integrates $\dot x_t = v_\theta(x_t, t)$ from noise to data over $N$ discrete steps. The paper’s claim is that at low $N$ , the discrete sample at step $k$ (nominal time $t_k$ ) distributionally resembles the training distribution at a strictly earlier time $t_k - \delta_k$ . The authors call $\delta_k$ the population time lag.

Two interventions address it. (1) Lag correction: scale the velocity by $\gamma(t) > 1$ so each discrete step covers more ground, pushing the sample forward in distribution time. The gain is chosen so a small probe (predicting which training-time a state belongs to) is centered on the current nominal time. (2) Extra sharpening: push $\gamma$ further than lag correction requires, which empirically sharpens the sample covariance, improving fidelity at the cost of coverage. The extra gain $\kappa$ is tuned once at 25 steps and reused.

The construction is purely at inference — no retraining — and is one multiplication per step, so it stacks on top of any pretrained flow-matching model.

Results

SiT-XL/2 on ImageNet 256×256 at NFE 25, no guidance: FID 28.0 (no scaling) → 12.2 (lag correction) → 8.7 (lag correction + extra sharpening $\kappa$ ) [§abstract, tweet 5/8].
The sharpening gain $\kappa$ tuned at 25 steps reaches near-best FID at 50, 100, and 250 steps too [tweet 5/8].
The CIFAR-10 model studied barely lags; any gain other than 1 hurts FID — the needed gain is a property of both the model and the sampling regime, not a universal constant [tweet 4/8].
Under CFG, the required velocity scale drops; with strong CFG, $\gamma < 1$ (scaling the velocity down) can improve FID — i.e. CFG and velocity scaling partly do the same job [tweet 7/8].

Why it’s interesting

Velocity scaling has been floating around as an empirical trick with a mechanistic story that this paper argues is wrong. Reframing it as population time lag correction makes it the inference-time sibling of the training-time CFG reframings on Classifier-Free Guidance variants (SoFlow, Ostris, CAFM) — all four treat a popular inference knob as a patch for a specific calibration failure, but each attributes the failure to a different source (training noise distribution, training objective, generalization quality, discretization). The redundancy with CFG — showing that the two sharpen covariance similarly, so stronger CFG needs less velocity scaling, and sometimes velocity de-scaling — is the kind of load-bearing interaction that lets you reason about what’s actually doing the work at inference. It also connects to Few-Step Generation (from scratch) and Diffusion Distillation: a one-multiplication fix that recovers 15.8 FID points at NFE 25 is directly competitive with distillation effort at the same step count, without a student model.