anthology-of-the-sota

SOTA-444: When a flow will be sampled with a coarse ODE solver, pair noise and data within each minibatch by exact optimal transport

4 min read

Properties6
StatusProposed
Consensusemerging
Tagsgenerative-modeling, flows-and-transport, training-optimization
SourceLIT-789 — Multisample Flow Matching: Straightening Flows with Minibatch Couplings, LIT-798 — Improving and generalizing flow-based generative models with minibatch optimal transport
Introduced byLIT-798 — Improving and generalizing flow-based generative models with minibatch optimal transport
Explained 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

Source

Pooladian, Ben-Hamu, Domingo-Enrich, Amos, Lipman and Chen (2023), LIT-789 — Multisample Flow Matching, Tables 1, 2, 6 and 9.

Tong, Fatras, Malkin, Huguet, Zhang, Rector-Brooks, Wolf and Bengio (2023), LIT-798 — OT-CFM, Table 2, Table 5 and Fig. 3.

What to do

Flow matching usually pairs each data point with an independent noise draw. Instead, for each training batch:

  1. Draw k noise samples and k data samples.
  2. Solve the exact assignment problem between them under squared Euclidean cost (a k × k linear assignment, the discrete OT plan).
  3. Train the ordinary straight-path flow-matching loss on the matched pairs.

Use exact batch OT, not the cheaper variants. Of the couplings Multisample Flow Matching tried, only exact BatchOT beat independent pairing on adaptive-solver FID at both resolutions (below).

This changes nothing at sampling time. What it buys is a learned ODE that is straighter, so a coarse Euler or midpoint solve loses less. “Coarse” here means roughly ten to thirty steps. It is not a route to one- or two-step sampling.

Evidence

The marginals are untouched by construction. Any coupling whose marginals are the noise and data distributions leaves the flow-matching gradient equal to the marginal one (LIT-789, Lemma 3.1 and §4; LIT-798, Thms. 3.1–3.2). As the batch grows, the batch-OT objective’s optimum and the paths’ straightness go to the OT values (LIT-789, Thm. 4.2; LIT-798, Prop. 3.4).

Multisample Flow Matching (LIT-789). Pixel-space, unconditional ImageNet. Architecture, optimizer, epochs and evaluation are shared across couplings (App. E.1).

  • Euler steps needed to reach a target FID (Table 1, read off Fig. 3’s curves):

    ImageNet-32, FID 10ImageNet-64, FID 20
    Independent (CondOT)2029
    BatchOT1412
    Diffusion baselines≥ 40≥ 40

  • Likelihood is unchanged: 3.58 and 3.27 bits per dimension for both couplings. Adaptive-solver FID improves, 5.04 → 4.68 on ImageNet-32 and 13.93 → 12.37 on ImageNet-64 (Table 6).

  • Samples from the same noise change less as the step count drops (Table 3).

  • The pairing costs 0.8% more time per iteration on ImageNet-32 and 3.9% on ImageNet-64 (Table 9).

OT-CFM (LIT-798), from a different group. In 2-D, normalized path energy is 0.018–0.087 with minibatch OT against 0.222–2.738 with independent pairs, over five seeds. Two and three rounds of reflow reach 0.069–0.149 and 0.055–0.129, so one OT-CFM training run does about as well as several rounds of reflow or better (Table 2). The OT batch can be small, and path energy plateaus near 64 samples (Fig. D.2). On CIFAR-10, Fig. 3 shows OT-CFM ahead of independent pairing at low Euler step counts after 400K steps, as curves only.

Conditions

  • The image gain is modest, and the numbers are one run each. OT-CFM’s CIFAR-10 table (LIT-798, Table 5) has 4.443 against 4.461 at 100 Euler steps. Independent pairing is ahead at 1,000 steps (3.643 against 3.741). Its “more efficient training” is per step: wall-clock time to convergence in 2-D is longer with OT (Table D.1). In Multisample Flow Matching, the cheaper couplings lose to independent pairing on ImageNet-32 adaptive FID: Stable 5.79, Heuristic 5.29 and BatchEOT 6.14 against 5.04 (Table 6). Several gaps there are under one FID point.
  • Unconditional, pixel space, 64 px and below. Neither paper pairs under a class or text condition. Pairing across a batch makes each noise depend on the other images in that batch, so the noise given a condition is no longer the source distribution, and neither paper shows the conditional marginals survive. Pairing within each condition is the obvious repair and is untested.
  • The minibatch plan is an approximation whose error grows with dimension (LIT-798, §6). The theorems are about the exact plan or the limit of large batches. Nothing bounds the error at image dimension and realistic batch size.
  • It does not replace a few-step student. Four-step Euler FID with BatchOT is 38.86 on ImageNet-32 and 56.75 on ImageNet-64 (LIT-789, Tables 7–8). For one to four steps, distillation or a consistency-type model is the route. SOTA-392 is about that case.
  • What it straightens is the marginal, not the conditional path. SOTA-266 already says to use the straight conditional path. THEORY-106 notes that a straight conditional path does not give a straight marginal ODE. The coupling is the lever that straightens the marginal with the conditional path held fixed. In OT-CFM’s CIFAR-10 runs that lever moves FID less than the choice of path does (LIT-798, Table 5). Use this on top of SOTA-266, not instead of it.

Known implementations

  • torchcfm, the OT-CFM authors’ library.

Backlinks