PaperlayerView exact revision on arXiv

Source-linked AI summary

Global Transport Couplings for Classifier-Free Guided Flows

Katarina Petrović, Zander W. Blasingame, Danyal Rehman, İsmail İlkan Ceylan, Michael Bronstein, Stephen Y. Zhang, Lazar Atanackovic, Alexander Tong

arXiv:2610.07555v1cs.LG

TL;DR

Conditional OT is difficult to scale across large or continuous condition spaces, and it remains an open question how couplings should be chosen for conditional generation sampled with CFG. The paper introduces Global Transport, a condition-agnostic OT coupling, and reports improved guided generation across studied settings despite worse unguided performance. The results support evaluating couplings under the guided flow used at inference.

  • Problem

    Condition-specific OT can be impractical for large or continuous conditioning spaces, and it is unclear what coupling suits conditional generation under CFG.

    Source: §1 · p. 2

  • Method

    Global Transport computes one global OT assignment without considering condition labels and adds no changes to the model architecture, sampler, or guidance rule.

    Source: §1 · p. 2, §3 · p. 4

  • Results

    GT improves guided generation across the studied settings, although its condition-specific source mismatch can worsen unguided conditional performance.

    Source: §1 · p. 2, §7 · p. 11

  • Takeaways & Limitations

    Unguided performance is not a reliable basis for choosing a coupling when CFG is used at inference.

    Source: §7 · p. 11

    The prediction-gap bound does not control learned fields along CFG trajectories or guarantee a hierarchy of generation quality.

    Source: §7 · p. 11

Abstract

from arXiv · show

Optimal-transport couplings have been shown to reduce training variance in unconditional flow models, but their role in conditional generation remains unclear. A natural approach constructs separate couplings for each condition, but this is impractical for large or continuous conditioning spaces found in modern image foundation models. We introduce Global Transport (GT), a global class-agnostic optimal-transport coupling, computed without class labels. GT can associate different conditions with different regions of the source noise, and consequently worsens performance without guidance. However, when combined with classifier-free guidance (CFG), GT consistently improves generation across domains, model scales, and sampling budgets. This reversal suggests that couplings for conditional flows should be evaluated both empirically and theoretically under the guided flow used at inference, rather than on unguided generation. We evaluate GT over both discrete class and continuous text conditioned image generation across model scales, and investigate how coupling choice alters guided trajectories. These results identify coupling design in the guided flow setting as a simple training time axis to improve performance without modifying existing architectures, samplers, or guidance mechanisms.

1 Introduction

Conditional OT can be impractical for large or continuous condition spaces, while CFG means inference uses both conditional and unconditional velocity fields. The paper introduces Global Transport and reports that its unguided disadvantage reverses under CFG, alongside reduced prediction gaps and yo-yo behavior.

1 Introduction

  • Class-conditional OT requires large batches per condition, making it impractical for many classes and infeasible for continuous conditions such as text embeddings.The approach has seen limited adoption in conditional generation.Source: §1 · p. 2
  • Because CFG combines conditional and unconditional velocity fields, the paper asks whether a coupling that considers both fields together can improve guided generation.Class-conditional OT improves the two fields separately.Source: §1 · p. 2
  • GT computes one global OT assignment without condition labels and requires no changes to the model, sampler, or CFG rule, but performs worse in unguided conditional sampling.Under CFG, the performance ordering reverses and GT improves guided generation across the studied settings.Source: §1 · p. 2
  • GT reduces the prediction gap between conditional and unconditional flows and mitigates the yo-yo effect along guided trajectories as guidance strength increases.The yo-yo effect involves contraction, re-expansion, and overshooting.Source: §1 · p. 2
  • GT outperforms independent and class-conditional couplings across class-conditioned image, text-conditioned image, and single-cell generation settings, model scales, and sampling budgets.The paper reports FID 1.91 on ImageNet-256 as one example.Source: §1 · p. 2

2 Background and preliminaries

Flow matching trains a velocity field to transport samples from a source distribution to data, and conditional flow matching uses a coupling to make its training objective tractable. Optimal-transport couplings replace independent pairing, while class-conditional OT faces scaling limits; CFG combines conditional and marginal fields during sampling.

2 Background and preliminaries

  • Flow matching learns a velocity field that evolves samples from a source distribution p0 to a target data distribution p1.Once trained, samples are generated by drawing from p0 and numerically integrating the field's ODE.Source: §2 · p. 3
  • Conditional flow matching introduces a coupling and tractable conditional paths, then trains the model by regression on their conditional velocity field.This replaces the unavailable closed-form marginal velocity field in the training objective.Source: §2 · p. 3
  • CFG combines marginal and conditional velocity fields using guidance weight w, with a null conditioning token used during training to estimate the marginal field.The conditional signal is dropped on some training examples.Source: §2 · p. 3
  • OT-CFM replaces independent noise-data pairing with approximations to the Euclidean OT coupling, whose straighter trajectories can ease few-step integration.The cited work reports modest but consistent gains in the unconditional setting.Source: §2 · p. 3
  • Class-conditional OT requires separate OT problems for each class, which becomes infeasible with many classes and cannot handle continuous conditions such as text embeddings.ImageNet has 1,000 classes; proposed smaller-batch approximations have unknown degradation and limited adoption.Source: §2 · pp. 3–4

3 Global Transport

Global Transport applies optimal transport without preserving condition labels, using a training-time assignment step rather than changes to the model or inference procedure. This can mismatch the source distribution associated with a condition, while guided trajectories show less yo-yo behavior under GT.

3 Global Transport

  • GT computes an approximate OT coupling without constraining assignments to preserve class labels, adding a training-time assignment step but no architecture or inference-algorithm changes.The paper describes both minibatch and semi-discrete approximations.Source: §3 · p. 4
  • GT uses the OT coupling to match noise to nearby data points while ignoring labels, then applies the conditional flow-matching objective.The condition is replaced by a null token with probability p_uncond and otherwise retained.Source: §3 · p. 4
  • For minibatch OT, squared Euclidean costs require O(B^2) memory and exact Hungarian assignment requires O(B^3) time, usually under 2% overhead in modern workflows.Large batches can make this approach expensive; regularized or semi-discrete methods may be more appropriate.Source: §3 · p. 4
  • The semi-discrete implementation first optimizes a semi-dual potential over discrete datapoints, then uses it to calculate couplings for continuous noise samples.The preprocessing step is expensive.Source: §3 · p. 4
  • Ignoring conditions can pair different conditions with different regions of the source prior, so the conditional source distribution need not equal the unconditional prior.Unguided sampling from the prior is not guaranteed to reach p1|c; that guarantee requires sampling from the intractable conditional source distribution.Source: §3 · p. 4
  • Independent pairings show stronger contraction and overshoot as guidance rises, whereas GT exhibits less pronounced yo-yo behavior over the tested guidance scales.The yo-yo effect is contraction followed by re-expansion toward class modes.Source: Figure 1 · p. 5

4 Coupling Choice With Classifier-Free Guidance

The section compares coupling behavior under guided and unguided sampling, then examines guided trajectories and prediction-gap bounds. GT shows less yo-yo behavior across tested guidance scales, while the bounds apply to training interpolants and do not establish that smaller gaps cause better generation.

4 Coupling Choice With Classifier-Free Guidance

  • GT may be poor for unguided conditional sampling (w = 1), but its behavior changes under guidance (w > 1), motivating evaluation of both fields as used at inference.Source: §4 · p. 5

4.1 Geometry of guided trajectories

  • Across tested guidance scales, independent trajectories contract then re-expand and may overshoot, whereas GT shows less contraction and overshoot and is stable.Yw(t) describes radial position along guided-sampler trajectories, unlike the distribution of training interpolants.Source: §4.1 · p. 5, Figure 2 · p. 6

4.2 Coupling controls the prediction gap

  • Lower transport cost tightens an upper bound on average prediction gap along training interpolants; the bound ranks GT tighter than C2OT (β > 0), then independent coupling.Finite mini-batch OT retains a non-zero prediction gap while regularizing the velocity fields.Source: §4.2 · pp. 6–7
  • The bound applies along linear training interpolants, not guided-inference states, and does not establish that smaller prediction gaps cause better generation.Source: §4.2 · pp. 6–7

5 Experiments

The experiments test Global Transport (GT) for guided class- and text-conditioned image generation and single-cell generation. Across these settings, they examine performance across coupling strategies, guidance scales, model sizes, and sampling budgets.

5.1 Class- and text- conditioned image generation on ImageNet-256

  • FID improves from 2.06 with standard SiT-XL/2 trained using independent couplings to 1.91 with SiT-XL/2 trained using GT.Table 1 compares FID against baselines and includes curated SiT-XL/2 + GT samples.Source: §5.1 · p. 7, Table 1 · p. 7
  • Across model sizes and NFEs, GT couplings deliver the strongest reported ImageNet-256 performance after sweeping guidance strength for each setting.GT also achieves a better overall FID minimum and can tolerate higher guidance scales before performance degrades.Source: Table 2 · p. 8, §5.1 · p. 8
  • On SiT-B/2, applying GT only during distillation matches pretraining with GT throughout, supporting its use as a distillation strategy for few-step generators.The study distills SiT-B/2 and SiT-L/2 teachers into flow maps, with students distilled using GT.Source: §5.1 · p. 8, Table 3 · p. 9
  • With the adaptive dopri5 solver, GT attains the lowest FID and FDDINOv2 overall, while semi-discrete GT requires the fewest function evaluations and leads at high guidance scales.Table 4 reports FID, FDDINOv2, and NFE for SiT-B/2 across coupling plans and guidance scales.Source: §5.1 · p. 8, Table 4 · p. 9
  • For text-conditioned ImageNet-256 generation, GT achieves lower FID and FDDINOv2 at higher guidance scales.Figure 5 presents both metrics across guidance scales and text-conditioned samples at w = 4.Source: Figure 5 · p. 9, §5.1 · p. 10

5.2 Single-cell Experiments

  • Across three single-cell datasets, GT ranks best on three of six metrics and second best on the other three against four comparison models.The comparison includes c-CFGen, its linear interpolant variant, scDiffusion, and scVI.Source: Table 5 · p. 10, §5.2 · p. 10

6 Related Work

The related work situates optimal transport in generative modeling and biological applications, then reviews flow maps and classifier-free guidance as approaches used in generation.

Optimal Transport for Generative Models

  • Optimal transport has been used in unconditional and class-conditional generation, including scalable semi-discrete and re-flow-style approaches, as well as biological trajectory inference and measure-to-measure transport.The cited work spans generative modeling and biological applications.Source: §6 · p. 10

Flow Maps

  • Flow maps are presented as an efficient route to one- and few-step generators, either distilled or trained from scratch, for image and video generation.Source: §6 · p. 10

Classifier-Free Guidance for Flow Models

  • Classifier-free guidance is widely used with diffusion and flow-matching models, with later work exploring interval tuning, velocity projection, weight schedules, and weaker-checkpoint guidance.Source: §6 · p. 10

7 Conclusion

The conclusion presents Global Transport (GT) as a condition-agnostic coupling that improves guided generation in the studied settings, despite degrading unguided conditional generation. It cautions that the prediction-gap bound does not guarantee a hierarchy of generation quality.

7 Conclusion

  • GT improves guided generation across all studied settings, although its source-distribution mismatch may degrade unguided conditional generation.GT requires no change to the model architecture, sampler, or guidance rule.Source: §7 · p. 11
  • The prediction-gap bound neither controls learned fields along CFG trajectories nor guarantees a resulting hierarchy of generation quality.The authors also leave open how coupling choice interacts with other guidance and post-training methods, and whether gains transfer beyond the studied domains.Source: §7 · p. 11

A Synthetic Example

The appendix illustrates coupling behavior on a four-class 40-GMM, then develops prediction-gap and transport-cost analyses. Its synthetic results contrast coupling-dependent velocity fields and guided trajectories.

A Synthetic Example

  • The synthetic example uses a two-dimensional 40-component Gaussian mixture with four classes and equal-weight components.The component means are drawn uniformly from [−40, 40]².Source: A Synthetic Example · p. 17
  • GT retains the 40-GMM modes at guidance scale w = 20, while independent pairing shows contraction followed by re-expansion near the conditioned class.The independent field forms a sink near the class origin before reorganizing at t = 0.5; under GT, the unconditional field forms at t = 0 and remains stable.Source: A Synthetic Example · pp. 17–18, Figure 7 · p. 18, Figure 8 · p. 19

B.1 Relationship between coupling and prediction gap

  • The prediction-gap analysis defines the unconditional field as Eπ[u | x = xt], where u = x1 − x0 along the coupling-induced path.The appendix also defines coupling cost and excess quadratic transport cost as transport quantities.Source: §B.1 · p. 20

B.2 Proof of Proposition 4.2: Hierarchy of Coupling Costs

  • Proposition B.3 establishes an ordering of quadratic transport costs across global and condition-aware coupling choices, including independent and hard class-matching limits.The proof considers the hard class-matching limit and combines its parts into a complete chain of inequalities.Source: §B.2 · pp. 21–23

B.3 Coupling controls flow curvature

A Gaussian plug-in model provides an analytically tractable proxy for minibatch OT. In this analysis, acceleration decreases with batch size for plug-in OT, remains order one for independent flow matching, and vanishes for population OT.

B.3 Coupling controls flow curvature

  • Pointwise and total material acceleration scale as O(n^-1) for Gaussian plug-in minibatch OT, remain O(1) for independent FM, and vanish for population OT.The corresponding expected plug-in acceleration energy decays as n^-2 at fixed dimension, while independent-coupling energy remains order one in n.Source: §B.3 · pp. 23, 28
  • The analysis uses a Gaussian plug-in model: it fits minibatch means and covariances, couples the fitted Gaussians, then marginalizes their conditional flow-matching fields.This is an analytically tractable proxy; exact analysis of the experimental discrete Hungarian coupling is described as likely more involved.Source: §B.3 · p. 23
  • Independent flow matching exhibits order-one contraction–expansion despite identical endpoint distributions, whereas population OT has zero velocity and time derivative when Σ = I.Gaussian plug-in minibatching retains residual motion whose magnitude and Eulerian time derivative decay as 1/n.Source: §B.3 · p. 24
  • At t = 1/2, the Gaussian plug-in field is exactly zero for every n because exchanging the identically distributed fitted batches preserves the midpoint and reverses the displacement.Source: §B.3 · p. 26
  • The plug-in material acceleration agrees with the Eulerian derivative to order 1/n in the standardized case because the convective term is order 1/n^2.Source: §B.3 · p. 27
  • For every fixed nonzero starting point, independent FM has the largest of the three acceleration energies for sufficiently large n, but this does not establish finite-n monotonicity for plug-in OT.The leading plug-in acceleration energy is nonincreasing in n and decays as n^-2 in the controlled large-n expansion.Source: §B.3 · p. 28

C.1 Implementation Details

Implementation details document the training setups and coupling overhead; with guidance-interval tuning, GT achieves the strongest results across three SiT scales.

C.1 Implementation Details

  • ImageNet-256 training configurations and optimization hyperparameters are reported across SiT model scales.The implementation details are organized in Tables 6 and 7.Source: §C.1 · p. 29, Table 6 · p. 29, Table 7 · p. 29
  • GT adds a modest optimizer-step cost, slightly higher than class-conditional OT because it is not restricted per class.The compute comparison is for SiT-B/2 training on ImageNet-256.Source: §C.1 · p. 30, Table 8 · p. 30
  • Class-conditional OT uses β to prevent cross-class coupling, with βm set to 10 after a sweep.The sweep is shown in Figure 10.Source: §C.1 · p. 30, Figure 10 · p. 30
  • Each ImageNet-1k image is paired with a VisualLayer caption encoded as a normalized 1024-dimensional embedding.The captions are encoded once using a frozen DFN5B CLIP ViT-H/14 text tower.Source: §C.1 · p. 30

C.2 Additional results on ImageNet-256

  • With fixed noise and class, GT keeps the generated object stable across guidance scales, unlike independent and class-conditional couplings.Under the independent coupling, increasing guidance can change composition and pose.Source: §C.2 · p. 30
  • GT achieves the strongest FID results overall with guidance-interval tuning across B/2, L/2, and XL/2 scales.Table 10 reports FID across model scales and couplings under the full and tuned guidance intervals.Source: §C.2 · p. 30, Table 9 · p. 31

D Single-cell Datasets

The appendix describes three single-cell datasets and their held-out evaluation setup, alongside classifier-guidance context and additional image-generation samples.

D Single-cell Datasets

  • The single-cell datasets cover 2,638 PBMC3K cells across 8 types, 18,213 dentate-gyrus cells across 14 types, and 584,944 HLCA cells across 50 types.The HLCA cells span 486 individuals and 49 datasets.Source: D Single-cell Datasets · p. 31
  • The single-cell evaluation uses held-out test data and compares an independent-coupling CFGen-linear baseline with GT transport.CFGen splits data into 90% training and 10% validation; its pipeline encodes cells with an RNA autoencoder before latent-space flow training.Source: D Single-cell Datasets · p. 31
  • For fixed noise and class, GT keeps object characteristics stable across guidance scales, while independent and class-conditional samples change composition or pose at higher scales.Figure 11 illustrates this comparison with SiT-XL/2 samples for class 339 sorrel.Source: Figure 11 · p. 32

E.1 Classifier and classifier-free guidance

  • Classifier-free guidance constructs a guided field without the additional classifier training required by classifier guidance.Classifier guidance requires training p_t(c | x), which the appendix identifies as a substantial added compute cost.Source: §E.1 · p. 33

Related papers

Loading 2610.07555v1…