Align Your Flow: Scaling Continuous-Time Flow Map Distillation

TL;DR

Align Your Flow (AYF, NVIDIA) is the high-performance continuous-time flow map distillation recipe that the diffusion-distillation literature keeps pointing at. The paper (i) proves analytically that consistency models accumulate error past 1–2 sampling steps and cannot be fixed by training, (ii) introduces two new continuous-time flow map losses — AYF-EMD and AYF-LMD — that generalize both continuous-time CM losses and the flow matching loss as limits, and (iii) is the first work to distill an autoguided teacher and to combine flow-map training with a light adversarial finetune that keeps diversity. AYF sets SOTA few-step generation on ImageNet 64² (FID 1.07 at 8 NFE) and 512² (FID 1.64 at 4 NFE adv, 1.70 at 4 NFE non-adv) using small networks, and beats every non-adversarial few-step T2I sampler on distilled FLUX.1 via a LoRA-only training framework.

Key claims

Consistency models are provably incompatible with multi-step sampling. For any non-optimal CM, Wasserstein-2 error to the data distribution grows past some finite step count $N$ ; with the standard $\sigma_d = 0.5$ data std the optimum is at 2 steps, matching the well-known empirical observation that CMs peak at 1–2 NFE and degrade thereafter [§3.1, Theorem 3.1, Fig. 5].
Flow maps — networks $f(x_t, t, s)$ that map any noisy state at time $t$ to the corresponding PF-ODE point at any other time $s$ — generalize denoisers ( $s = t + dt$ ), CMs ( $s = 0$ ), and rectified flow / Euler sampling ( $s = t - \Delta t$ ) in one framework [§3, §4].
AYF-EMD (Eulerian Map Distillation) fixes the starting time $r$ and asks $f(x_t, t, r)$ to remain constant as $x_t$ drifts along the PF-ODE. In the limit it reduces to the continuous-time CM loss when $r = 0$ and to the flow matching loss as $r \to t$ , so it is a strict generalization of both [§3.2, Theorem 3.2].
AYF-LMD (Lagrangian Map Distillation) instead fixes the ending time $r$ and asks the trajectory $(x_t, t) \mapsto f(x_t, t, r)$ to stay on the same PF-ODE solution; this is more stable in 2D toy experiments but produces overly smooth image samples in practice, so the image experiments all use AYF-EMD [§3.2, Theorem 3.3, Appendix ablations].
Autoguidance distillation is a free upgrade over CFG distillation. Distilling a teacher whose velocity is sharpened by autoguidance (a weaker model used to subtract low-quality directions, Karras et al. 2024a) instead of by CFG gives crisper samples without CFG’s overshooting failure mode, and is the first reported distillation of an autoguided teacher [§3.3, Eq. 3].
Short adversarial finetuning preserves diversity. Unlike adversarial training (DMD2, SiDA, CTM, ADD), running a brief GAN-style finetune on top of a pretrained AYF model — with the flow-map loss still active — sharpens samples while keeping Recall ≈ 0.65 on ImageNet 64², matching or beating prior adversarial methods that report no Recall [Tables 1–2, Fig. 8 recall plots].
Stable continuous-time training requires three knobs. AYF adopts sCM’s tangent normalization (critical), drops sCM’s adaptive weighting (no observed effect), and changes the time parameterization to $\tilde t = \ln(t/(1-t))$ to give stable $\partial /\partial t$ in the flow matching framework where $t \in [0,1]$ . A LoRA-style finetune transfers EDM2 checkpoints (different time parameterization) to the new $\tilde t$ parameterization [§3.4].
Regularized tangent warmup beats vanilla warmup. sCM ramps the tangent coefficient $r$ from 0→1 over 10k steps then holds at 1; AYF clamps $r$ at 0.9, keeping a 10% mixture of the straight-line-regularizer term throughout training, which fixes late-training destabilization [§3.4, Eq. 5].
AYF beats every non-adversarial few-step ImageNet 64² method at comparable parameter count. Without adversarial loss: AYF-S at 280M params reaches FID 3.32 (1 NFE), 1.87 (2 NFE), 1.70 (4 NFE) — vs sCD-S (same params) 3.07/2.50, and matches sCD-XXL (1.5B params) at 4× fewer parameters [Table 2].
With adversarial finetune, AYF is SOTA few-step at any NFE. ImageNet 64² FID: 1.32 (1 NFE), 1.17 (2 NFE), 1.10 (4 NFE), 1.07 (8 NFE), beating SiDA’s 1.11 and approaching the autoguided 63-NFE EDM2 teacher (1.01) [Table 1].
FLUX.1 [dev] distillation via LoRA only. Distilling FLUX.1 [dev] with an AYF-EMD LoRA reaches text-conditioned SOTA among non-adversarial few-step methods at 4 NFE, without retraining the full backbone [§5, Fig. 1].

Method

A flow map $f_\theta(x_t, t, s)$ is parameterized so that the boundary condition $f(x_t, t, t) = x_t$ holds by construction, via $f(x_t, t, s) = c_{\text{skip}}(t,s) x_t + c_{\text{out}}(t,s) F_\theta(x_t, t, s)$ with $c_{\text{skip}}(t,t) = 1$ and $c_{\text{out}}(t,t) = 0$ ; the paper picks the Euler-aligned form $c_{\text{skip}} = 1$ , $c_{\text{out}} = s - t$ . Training distills a pretrained flow-matching teacher $v_\phi(x_t, t)$ — optionally with autoguidance applied: $v^{\text{ag}}_\phi(x_t,t) = v_\phi(x_t,t) + w[v_\phi(x_t,t) - v^{\text{weak}}_\phi(x_t,t)]$ .

The AYF-EMD loss between two adjacent starting times $t, t+\Delta t$ is a squared distance between $f_\theta(x_t, t, r)$ and $f_\theta(\text{stop-grad}(\text{Euler}(x_t, t, t+\Delta t)), t+\Delta t, r)$ . Taking $\Delta t \to 0$ and stopping gradient on the bootstrapped target yields a gradient whose closed form involves the tangent $\mathrm{d} f / \mathrm{d} t$ along the PF-ODE (Theorem 3.2). The tangent decomposes into a $\partial f / \partial x_t \cdot v_\phi$ JVP term plus a $\partial f / \partial t$ time-derivative term plus the explicit $\partial f / \partial t$ inside the parameterization; the time-derivative term is the unstable piece, addressed by tangent normalization (dividing by $\|\cdot\| + c$ , with $c = 0.1$ ) and the $\tilde t = \ln(t/(1-t))$ reparameterization.

Image experiments use EDM2-S/M/L/XL/XXL student backbones initialized from the corresponding pretrained EDM2 teachers (after a short time-embedding finetune), distill from an autoguided EDM2 teacher, and run for ~1M iterations with AYF-EMD; the adversarial variant adds a short GAN finetune (a few k iterations) with the AYF loss still applied. FLUX.1 [dev] distillation runs as a LoRA on top of frozen FLUX weights, with the same AYF-EMD objective applied in latent space.

Results

ImageNet 64² (Table 1, non-adversarial): AYF-S 280M, 1 / 2 / 4 / 8 NFE → FID 3.32 / 1.87 / 1.70 / — ; vs sCD-S 3.07 / 2.50 (same params, CM family degrades past 2 NFE) and sCD-XXL 1.5B 2.28 / 1.88. AYF uniformly improves with more NFE; sCD plateaus.
ImageNet 64² (Table 1, adversarial): AYF-S + adv 1 / 2 / 4 / 8 NFE → FID 1.32 / 1.17 / 1.10 / 1.07, Recall 0.65 throughout — vs SiDA 1.11 / — / — / —, DMD2 1.28, GDD-I 1.16. AYF is the only multi-step adversarial method that keeps improving past 2 NFE.
ImageNet 512² (Table 2): AYF-S 280M 1 / 2 / 4 NFE FID 3.32 / 1.87 / 1.70 non-adv, 1.92 / 1.81 / 1.64 adv — beats every sCD size at matched params and lifts the Pareto frontier of FID-vs-FLOPs (Fig. 4): AYF-S at 4 NFE (408 GFLOPs) matches sCD-XXL at 2 NFE (1104 GFLOPs).
FLUX.1 [dev] T2I: AYF-LoRA at 4 NFE beats Phased CMs, Multistep-CD, CTM, and every other non-adversarial few-step distillation in the comparison; qualitative examples in Fig. 1 show text alignment and detail comparable to the 50-NFE teacher.
Diversity (Fig. 8): AYF + adv keeps Recall ≈ 0.65 across 1–8 NFE; the SiDA / DMD2 / CTM adversarial baselines collapse Recall to ≤ 0.60.

Why it’s interesting

This is the “AYF” the Slack message names alongside sCM and MeanFlow, and it slots directly into the Diffusion Distillation thread. AYF is the strongest continuous-time flow-map recipe filed on the wiki: unlike consistency models, the FID monotonically improves as NFE grows, which is exactly the test-time scaling property that AnyFlow: Any-Step Video Diffusion Model with On-Policy Flow Map Distillation (AnyFlow) argues consistency distillation loses on the video side. AYF and AnyFlow are sibling solutions to the same problem — drop endpoint mapping for a two-time flow-map target — applied to image (AYF) and video (AnyFlow) teachers; AnyFlow’s “Flow Map Backward Simulation” is essentially AYF-EMD with an on-policy rollout twist. The conceptual map in Learning the integral of a diffusion model (flow maps) explicitly names AYF-EMD as the canonical Eulerian-consistency instantiation, alongside MeanFlow (Lagrangian) and shortcut / progressive matching (compositionality) — Alex Z’s note (“sCM, AYF, MeanFlow are essentially the same”) is exactly Dieleman’s three consistency rules viewed from one teacher. The autoguidance angle also connects to Classifier-Free Guidance variants: AYF is the first datapoint on the wiki where distillation explicitly benefits from replacing CFG with autoguidance, a concrete production lever.