anthology-of-the-sota

LIT-791: Improving the Training of Rectified Flows

7 min read

Properties5
StatusActive
Tagsgenerative-modeling, few-step-generation, flows-and-transport, inference-optimization, training-optimization
ExtendsLIT-636 — Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow
Compared againstLIT-075 — Elucidating the Design Space of Diffusion-Based Generative Models, LIT-093 — Consistency Models, LIT-067 — Progressive Distillation for Fast Sampling of Diffusion Models, LIT-643 — One-step Diffusion with Distribution Matching Distillation
Cited as “source” byTHEORY-123 — A flow's sampling ODE curves because the noise-data pairs it was trained on cross and the learned velocity averages them, so a coupling whose straight paths do not cross gives a straight flow

Graph View

Lee, Lin and Fanti, Carnegie Mellon University and Microsoft Research (2024), NeurIPS 2024 — ARXIV-2405.20320. Read at v2 (8 Oct 2024), main text, Appendices A–G and the checklist; v1 is 30 May 2024.

Key takeaways

  • One reflow round, and why (§3, Fig. 2). For two 1-rectified-flow pairs to cross, one noise must equal another plus a scaled data difference. Such a noise is atypical: its norm sits off the Gaussian annulus and it is autocorrelated (Fig. 2c–d, Fig. 7). So crossings are rare, the optimal 2-rectified flow is nearly straight, and further rounds only add error. The checklist says this “is intuitive, and we do not prove it”. Prior work is described as needing 3-rectified flow for good one-step results.
  • The timestep distribution should be U-shaped for reflow (§4.1, Table 1). The 2-rectified-flow loss is largest near both ends of [0, 1] (Fig. 3), the opposite of a 1-rectified flow, so pₜ(u) ∝ e^{4u} + e^{−4u}. Against uniform, one-step FID goes 7.14 → 5.17 on CIFAR-10, 12.39 → 9.03 on AFHQ-64 and 8.84 → 6.81 on FFHQ-64.
  • The loss can be perceptual (§4.2, Table 1). If the optimum is nearly the posterior anyway, any premetric has the same minimizer. Pseudo-Huber alone is neutral at large batch (5.24 against 5.17 on CIFAR-10). LPIPS on the predicted clean image, weighted toward high t, is the largest single gain: 5.24 → 3.38 on CIFAR-10, 8.20 → 4.11 on AFHQ, 7.06 → 5.21 on FFHQ.
  • Initialize from a diffusion model (§4.3, Prop. 1, App. C). A VP or VE denoiser gives the rectified-flow posterior after a time and scale change. EDM initialization plus batch 512 takes CIFAR-10 from 12.21 to 7.14. A short fine-tune on (real data, inverted noise) pairs takes 3.38 → 3.07 (§4.4, config G).
  • Headline numbers (Tables 3–4). CIFAR-10: 3.07 at one step, 2.40 at two. ImageNet-64: 4.31 and 3.64, against CD 6.20/4.70, PD 15.39/8.95 and iCT 4.02/3.20. Its own tables also list DMD at 3.77 (CIFAR-10) and 2.62 (ImageNet-64), and CTM+GAN at 1.98 and 1.92.
  • Forward passes (§5.2, Table 6). On ImageNet, pair generation (395M) plus training (1,433.6M) is 1,828.6M passes, against 5,734.4M for CD and 2,867.2M for CT at the same iteration count.

Where the hedges are

Per DP-010:

  • “A single iteration of the Reflow algorithm … is sufficient” (abstract). The argument is a geometric intuition with toy evidence, not a measurement of 2-RF curvature on the trained models. It is also a claim about the optimum: the paper’s own premise is that the trained 2-RF falls short of it.
  • “Without additional Reflow or distillation stages” (§1). The training set is (noise, EDM output) pairs from a 79-step teacher, and the loss puts LPIPS weight where the model predicts data from noise. That is close to a distillation of EDM’s ODE map, so this is less a reflow-instead-of-distil result than a single stage doing both.
  • “Outperforms the state-of-the-art distillation methods” (abstract) holds against CD, PD, DFNO, TRACT and BOOT. Two distillers in the paper’s own Table 4 beat it on ImageNet-64: DMD (2.62) and CTM+GAN (1.92). The text says it closes the gap with iCT at “4.01 vs 4.31”; Table 4 gives 4.02.
  • “Without harming performance at the higher NFE setting” (§6) is not measured against the pre-reflow model. Fig. 4 shows FID falling with NFE, but ImageNet-64 at 8 NFE is 3.32 (Fig. 15) against EDM’s 2.30 at 79 (Table 4). No matched-NFE comparison with EDM is given.
  • The cost comparison omits the comparator that matters. Table 6 sets reflow against CD and CT. A direct LPIPS distillation on the same 5M pairs would cost the same as reflow and is not run.
  • One FID run per configuration (checklist item 7). Several of Table 1’s gaps are below the 2% seed variation SOTA-307 reports: 5.17 against 5.24 (adding Huber, C to D), 3.42 against 3.38 and 4.13 against 4.11 (the 1/t weighting).

Which comparisons are like for like

  • Table 1 is a controlled ablation within the paper. Same data, iteration count (800K, or 1M on AFHQ and FFHQ) and “best FID” reporting per config. Config A is Liu et al.’s published number, not a rerun.
  • Against Rectified Flow’s 2-RF (12.21) and 2-RF+distill (4.85): not like for like. Initialization (EDM, not a trained 1-RF), batch (512 against 128), loss (LPIPS-Huber), timestep distribution and real-data fine-tuning all differ.
  • Tables 3–4 copy baseline numbers from their papers. ImageNet uses batch 2048 for 700K iterations on 64 V100s for about 9 days, with config E and no real data.

Standing in the anthology

It extends Rectified Flow (LIT-636): it keeps reflow and the straight-path objective and changes how the reflow stage is trained, cutting 2-rectified flow’s one-step CIFAR-10 FID from Rectified Flow’s 12.21 to 3.07. Rectified Flow’s own distilled 1-RF (6.18) and distilled 2-RF (4.85) are rows in its Table 3, and it beats both without a distillation stage. The setup differences above mean that is not a test of their pipeline.

Its teacher is EDM (LIT-075). Every 2-rectified flow after config A is initialized from EDM’s denoiser via Prop. 1 and trained on pairs from EDM’s Heun sampler. EDM is also the many-step reference in Tables 3 and 4 (1.97 on CIFAR-10, 2.30 on ImageNet-64). The one-step student does not reach it.

The distillers it measures itself against are Consistency Models (LIT-093, CD 3.55 and 6.20 at one step), progressive distillation (LIT-067, 9.12 and 15.39) and DMD (LIT-643, 3.77 and 2.62), all as copied numbers. It beats the first two. DMD beats it on ImageNet-64, by distribution matching with no reflow.

On SOTA-392 it refines rather than tests. SOTA-392 says to distil a rectified flow for one step and treat reflow as optional. This paper says to reflow once and skip the distillation stage. It has no arm that distils the same EDM teacher without reflow at the same budget, so it cannot say which ordering wins. It does not meet promote_when: there is no distillation-alone comparison and no matched many-step comparison. What it adds is three refinements. If you reflow, one round is enough, with InstaFlow (LIT-796) also finding a second round unhelpful at text-to-image scale. The reflow stage should be trained like a distillation, with a perceptual loss and the endpoints weighted. And it reports a one-step number from reflow alone that the best distribution-matching distiller still beats in its own table.

On SOTA-266 it marks a boundary for the logit-normal half. That practice samples timesteps from a logit-normal, which concentrates on the middle of [0, 1]. This paper argues the reverse for a reflowed model, whose hard tasks are at the ends, and measures U-shaped against uniform for the reflow stage. It does not run a logit-normal. The practice’s distribution is for the base flow, and this is evidence it does not carry to reflow.

It supports SOTA-206 in kind. A 2-rectified flow stays an ODE, and its FID improves with every added step (Fig. 4). §1 makes that its main argument against distillation. It also inverts real images in 2–8 NFE, with lower reconstruction error than EDM at 16 (Fig. 5).

Filed without a NOTE: the takeaways come from one full reading of v2, main text, Appendices A–G and the checklist. Table values were read from the text layer. Fig. 4’s curves are read only where captions give values.

Backlinks