PaperlayerView exact revision on arXiv

Source-linked AI summary

Faster Inference of Flow-Based Generative Models via Improved Data-Noise Coupling

Aram Davtyan, Leello Tadesse Dadi, Volkan Cevher, Paolo Favaro

arXiv:2603.15279v1cs.LGcs.CV

TL;DR

CFM inference is slowed by repeated ODE evaluations, and its sampling efficiency depends on how data and noise are coupled. LOOM-CFM extends minibatch OT by preserving and refining assignments across training, improving the sampling speed-quality trade-off and supporting distillation and latent high-resolution synthesis.

  • Problem

    CFM sampling requires repeated evaluations, while exact OT coupling is infeasible at large scale and minibatch OT becomes less effective as datasets grow.

  • Method

    LOOM-CFM stores and iteratively refines minibatch data-noise assignments across training, using multiple noise caches to reduce overfitting to fixed assignments.

  • Results

    LOOM-CFM achieves a better sampling speed-quality trade-off than prior work across experiments, with improved few-step FID and no additional computational cost.

  • Takeaways & Limitations

    LOOM-CFM can be combined with rectified flows, improves distillation initialization, and supports latent flow matching for higher-resolution outputs.

  • Takeaways & Limitations

    The algorithm is not guaranteed to recover the optimal coupling and instead may converge to a sub-optimal transport plan.

Abstract

from arXiv · show

Conditional Flow Matching (CFM), a simulation-free method for training continuous normalizing flows, provides an efficient alternative to diffusion models for key tasks like image and video generation. The performance of CFM in solving these tasks depends on the way data is coupled with noise. A recent approach uses minibatch optimal transport (OT) to reassign noise-data pairs in each training step to streamline sampling trajectories and thus accelerate inference. However, its optimization is restricted to individual minibatches, limiting its effectiveness on large datasets. To address this shortcoming, we introduce LOOM-CFM (Looking Out Of Minibatch-CFM), a novel method to extend the scope of minibatch OT by preserving and optimizing these assignments across minibatches over training time. Our approach demonstrates consistent improvements in the sampling speed-quality trade-off across multiple datasets. LOOM-CFM also enhances distillation initialization and supports high-resolution synthesis in latent space training.

1 INTRODUCTION

Iterative generative models can produce high-quality content but require many inference evaluations, motivating methods that straighten trajectories and improve CFM data-noise coupling. LOOM-CFM extends minibatch OT by preserving and refining assignments across minibatches, with reported gains in quality, speed, distillation initialization, and high-resolution latent synthesis.

  • Motivation: Iterative methods require multiple model evaluations during inference, making generation slower than single-pass methods.The gradual transformation from noise to data necessitates repeated evaluations.
  • Motivation: Trajectory straightening is a proposed way to reduce integration steps and accelerate probability-flow ODE generation.The paper identifies enhanced training, distillation, and sampling modifications as routes toward faster generation.
  • Problem: CFM performance depends substantially on the coupling between data and noise, while exact OT is infeasible for large datasets and minibatch OT weakens as datasets grow.The coupling affects sampling-trajectory curvature, creating a scalability problem for existing OT approximations.
  • LOOM-CFM: LOOM-CFM preserves and iteratively refines minibatch noise-data assignments so local matchings communicate across minibatches, approximating a global OT plan.Multiple noise instances per data point are randomly selected during training to avoid overfitting to static assignments.
  • Results: 41%, 46%, and 54% FID reductions at 12 NFE are reported on CIFAR10, ImageNet-32, and ImageNet-64, respectively, versus minibatch OT methods.The contribution summary also reports convergence analysis and extensive evaluation against prior work.

2 BACKGROUND

The background develops CFM, explains why independent data-noise coupling can create curved ODE trajectories, and motivates OT-based coupling for faster sampling. Exact OT is computationally costly, while LOOM-CFM uses cross-minibatch assignment updates as a tractable intermediate.

  • Conditional Flow Matching: CFM trains a neural vector field so an ODE maps source noise into a valid approximation of the data distribution.The model starts from noise and obtains generated samples by numerically solving the learned ODE.
  • Faster Sampling: Lower-curvature ODE trajectories require fewer discretization steps to achieve the same accuracy, reducing neural-network evaluations during sampling.This connects trajectory geometry directly to inference cost.
  • Conditional Flow Matching: Independent coupling averages source-to-target directions and can make marginal sampling trajectories highly curved, even when pairwise interpolations are straight.In the Gaussian toy case, trajectories move toward the origin before turning back.
  • Faster Sampling: Reflow and related coupling-based methods can straighten paths, but reflow requires at least twice the training time and auxiliary approaches may add training complexity.The paper therefore focuses on changing the coupling distribution rather than adding auxiliary models.
  • Optimal Transport: Exact OT seeks a minimum-cost coupling between data and noise distributions, but solving the finite-support problem has O(n^3 log(n)) complexity and regularized alternatives remain difficult at modern scale.These costs motivate tractable approximations such as minibatch OT and the proposed cross-minibatch scheme.
  • Optimal Transport: LOOM-CFM reuses previous minibatch assignments, has complexity no greater than O(m^3), and converges to a stationary but potentially sub-optimal transport plan.Unlike averaged minibatch couplings, the method produces a deterministic coupling.

3 METHOD

LOOM-CFM extends minibatch optimal transport by preserving locally optimized data-noise assignments across training iterations, improving their approximation to a global coupling. It also uses multiple noise caches to reduce overfitting while retaining convergence guarantees and comparable computational complexity.

  • Looking Out Of The Minibatch: The assignment objective seeks a bijection between data and noise samples, but exact optimization is impractical at scale because common solvers require O(n3) time.LOOM-CFM instead applies minibatch matching while carrying assignments across iterations.
  • Looking Out Of The Minibatch: At each iteration, the method samples assigned data-noise pairs, locally updates their assignment for minibatch optimality, then trains CFM using the updated pairing.The updated assignments are saved and reused in later iterations.
  • Looking Out Of The Minibatch: LOOM-CFM preserves and iteratively refines local minibatch assignments, allowing them to influence future minibatches rather than discarding them after each step.This implicitly communicates local OT updates across minibatches and improves the global assignment approximation.
  • Convergence: LOOM-CFM converges in finitely many steps with probability 1 to an assignment whose matching has no negative alternating cycles shorter than m.The assignment costs are nonincreasing during convergence.
  • Limitations and Implementation: LOOM-CFM is not guaranteed to recover the globally optimal coupling, although it produces a deterministic coupling and is reported to improve path straightening over minibatch OT schemes.Its local assignments use the Hungarian algorithm, with minor I/O overhead offset by faster convergence in the reported ImageNet settings.
  • Multiple Noise Caches: Multiple noise caches assign several noise instances to each data point and randomly select among them during training to reduce overfitting to fixed source samples.This artificially enlarges the dataset without changing its underlying data distribution and supports using new noise instances at inference.

4 EXPERIMENTS

Experiments evaluate LOOM-CFM through ablations and comparisons on image-generation tasks. Results show improved few-step quality, straighter trajectories, better reflow initialization, and applicability to high-resolution latent-space synthesis.

  • Ablations: Training from the beginning while assignments improve performs better than waiting for a stable matching, which may overfit fixed noise-data assignments.The authors interpret progressive matching as a soft transition from independent coupling to the method’s learned coupling.
  • Ablations: Larger minibatches improve small-NFE performance, while batch size 32 matches OT-CFM at 8 and 12 NFE and is surpassed by larger LOOM-CFM batches.The authors associate larger batches with straighter sampling trajectories and attribute the comparison to the caching scheme.
  • Ablations: More noise caches improve performance, with four sufficient for CIFAR10-scale datasets; gains diminish beyond that, while one may suffice for ImageNet-scale data.Increasing caches can slow convergence and produce worse couplings, whereas sufficiently large datasets may require fewer caches.
  • Unconditional image generation: LOOM-CFM consistently improves few-step FID over minibatch OT methods and matches NFDM-OT overall, surpassing it at 12 NFE but slightly underperforming at lower NFE.Unlike compared methods requiring additional components, LOOM-CFM optimizes the original CFM objective with a modified coupling distribution.
  • Rectified flows: LOOM-CFM with reflow outperforms both 2-Rectified Flow and 3-Rectified Flow, indicating that additional reflows are unnecessary when the first reflow is well initialized.The evaluation uses one million noise-data pairs generated from the four-cache CIFAR10 model.
  • High-resolution image synthesis: In FFHQ 256 × 256 latent-space training, LOOM-CFM achieves lower FID with an order of magnitude fewer NFE than prior methods using the same architecture.The result supports direct compatibility with conventional autoencoder-based high-resolution synthesis.

5 CONCLUSIONS

LOOM-CFM improves the sampling speed-quality trade-off by recycling locally optimized data-noise matchings across training iterations. It incurs no additional computational cost and negligible disk overhead, and can compose with rectified flows.

  • LOOM-CFM recycles locally optimal minibatch matchings to improve global data-noise assignments and sampling trajectories.The method stores and reuses matchings from prior iterations rather than recomputing assignments independently.
  • LOOM-CFM achieves a better sampling speed-quality trade-off than prior work with no additional computational cost.Training incurs negligible disk-usage overhead from storing the assignments.
  • LOOM-CFM is compatible with other techniques, including rectified flows, to further enhance sampling speed.

A GAUSSIAN TO GAUSSIAN CASE

This section derives the optimal vector field for standard-normal source and target distributions. The resulting field is radial, with a time-dependent multiplier that reverses direction halfway through the interpolation.

  • The Gaussian-to-Gaussian case seeks a closed-form optimal vector field for standard-normal source and target distributions.
  • The CFM objective minimizes the squared discrepancy between the learned vector field and the displacement x − z under Gaussian sampling.
  • The derivation obtains the optimal field by applying the Euler-Lagrange equation to the objective functional.
  • The optimal vector field has the form ˆv(y, t) = y · s(t), with s(t) = (2t − 1)/(σ2 + t2 + (1 − t)2).
  • The resulting field is parallel to the line from y to the origin and changes direction halfway through t ∈ [0, 1].This analytical result matches an empirical two-dimensional neural-network fitting observation.

B ANALYSIS

LOOM-CFM can be viewed as iterative elimination of negative alternating cycles in locally selected matching subgraphs. It converges finitely to a stationary assignment, but minibatch updates do not guarantee global optimality.

  • Matching analysis: Optimal transport matching is represented as a minimum-cost perfect matching between data and noise nodes in a complete bipartite graph.The edge cost is the squared distance between each data point and noise sample.
  • Matching analysis: A matching is optimal if and only if it contains no negative M-alternating cycles.Swapping matching and nonmatching edges along a negative cycle lowers the total cost.
  • LOOM-CFM interpretation: LOOM-CFM updates the matching by solving optimal subproblems on randomly selected subgraphs, thereby eliminating negative alternating cycles of length at most m.
  • Convergence: With probability 1, LOOM-CFM converges in finitely many steps to an assignment with no negative alternating cycles shorter than m.Assignment costs are nonincreasing during the process.
  • Limitations: LOOM-CFM does not guarantee global optimality because some subproblem-based methods require a full cyclic permutation involving all points.The paper nevertheless reports more optimal couplings than minibatch OT methods and straighter trajectories in real-data experiments.

C IMPLEMENTATION DETAILS

The implementation uses an improved ADM U-Net to parameterize the learned vector field and matches prior architectures and training settings where possible. Models were trained on four Nvidia RTX 3090 GPUs, excluding ablations.

  • The implementation section documents model architectures, training parameters, and related algorithmic details.
  • LOOM-CFM parameterizes the learned vector field with an improved ADM U-Net architecture.
  • The experiments use prior-work architectures and training parameters where possible to enable fair comparisons.
  • All non-ablation models were trained on 4 Nvidia RTX 3090 GPUs.

D QUANTITATIVE RESULTS

The quantitative analysis examines convergence, minibatch OT costs, and training-time overhead. Larger batches produce better matching, while larger datasets slow convergence and assignment storage modestly increases training time.

  • Convergence: LOOM-CFM converges in roughly the same time across batch sizes despite different initial reassignment counts.Initial reassignment counts are approximately the batch size because assignments start randomly.
  • Convergence: Larger batch sizes converge to better matching, as indicated by lower minibatch OT cost.
  • Convergence: Increasing dataset size slows convergence because the algorithm must visit more minibatches to organize assignments.The minibatch OT cost changes accordingly as dataset size increases.
  • Training time: LOOM-CFM has slightly longer training time than baselines because saving and loading assignments adds disk I/O overhead.The method has comparable time complexity to OT-CFM and BatchOT because all solve the same minibatch matching problem.
  • Training time: 17.3 hours versus 13.9 hours: CIFAR10 training for 1000 epochs with batch size 128 took longer when assignments were stored.The comparison used four Nvidia RTX 3090 GPUs.

E QUALITATIVE RESULTS

The qualitative results visualize sampling paths and uncurated outputs across several datasets. CIFAR10 comparisons indicate that storing reassignments produces sharper generations that converge earlier.

  • Sampling Paths: On CIFAR10, LOOM-CFM produces sharper generations and converges earlier than the version without stored reassignments.

F ALTERNATIVE APPROACHES

The authors evaluate alternatives for preventing overfitting to fixed noise samples. Noise refreshing destabilized training, gradual noise injection yielded little improvement, and multiple noise caches provided a simpler approach.

  • Alternative approaches: Refreshing all noise samples every N epochs caused training instability because new assignments could shift network targets substantially.
  • Alternative approaches: Gradual noise injection, including making its parameter a function of t, did not produce substantial improvements.The authors leave further exploration of this approach for future work.
  • Alternative approaches: Multiple noise caches provide a straightforward alternative that artificially increases dataset size and equalizes small- and large-dataset settings.

G LIMITATIONS

The paper identifies conditional generation as a limitation for LOOM-CFM and other OT-based methods. Additional qualitative comparisons illustrate sampling behavior and generation across datasets and evaluation budgets.

  • Limitations: LOOM-CFM and other OT-based methods are not directly compatible with conditional generation, especially with complex conditioning signals.
  • Limitations: Naively conditioning a coupling-based model without adapting couplings may introduce sampling bias because label alignments can become disproportionate.
  • Qualitative comparisons: CIFAR10 trajectories compare LOOM-CFM with a model that does not store reassignments, using the same initial noise in corresponding rows.
  • Qualitative comparisons: FFHQ-256 samples are shown at 2, 4, 8, and 12 function evaluations.
Loading 2603.15279v1…