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

1. Problem Statement

Context: Conditional Flow Matching (CFM) is a simulation-free method for training Continuous Normalizing Flows (CNFs), offering an efficient alternative to diffusion models for image and video generation. However, like diffusion models, CFM-based generators rely on iterative Ordinary Differential Equation (ODE) solvers, which require multiple network evaluations (NFE) during inference, leading to slow generation speeds.

The Core Issue: The speed of inference is determined by the curvature of the sampling trajectories. Straighter trajectories allow for fewer integration steps (lower NFE) while maintaining high quality. The curvature is heavily influenced by the data-noise coupling (the joint distribution $p(x, z)$ between data $x$ and noise $z$).

• Independent Coupling: Standard CFM samples $x$ and $z$ independently, resulting in highly curved, inefficient trajectories.
• Optimal Transport (OT) Coupling: Theoretically, coupling data and noise via an Optimal Transport plan minimizes trajectory curvature. However, computing exact OT for large datasets is computationally infeasible ($O(n^3)$).
• Minibatch OT Limitations: Existing approaches (e.g., BatchOT, MFM) approximate OT by solving the transport problem independently within small minibatches. While faster, this approach fails to capture global structure, leading to suboptimal assignments that do not scale well with dataset size. The assignments are discarded after each step, preventing the model from learning a consistent global mapping.

2. Methodology: LOOM-CFM

The authors propose LOOM-CFM (Looking Out Of Minibatch-CFM), an iterative algorithm that extends the scope of minibatch OT to approximate a global optimal transport plan by preserving and refining assignments across training iterations.

Key Components:

1. Iterative Global Assignment (The "Loom" Mechanism):

   • Instead of sampling noise independently for every minibatch, LOOM-CFM maintains a global assignment map $\tau$ (a permutation matrix) linking the entire dataset of $n$ data points to $n$ noise samples.
   • Algorithm Flow:
     1. Initialize $\tau_0$ (e.g., identity).
     2. Sample a minibatch of data points $\{x_{n_j}\}$ and retrieve their currently assigned noise samples $\{z_{\tau_{k-1}(n_j)}\}$.
     3. Solve a local Optimal Transport problem within this minibatch to find a local permutation $\omega_k$ that minimizes the transport cost for the current batch.
     4. Update Global Assignment: Update the global map via composition: $\tau_k = \omega_k \circ \tau_{k-1}$.
     5. Train the vector field $v_\theta$ using the updated pairs $(x_{n_j}, z_{\tau_k(n_j)})$.
   • Significance: This allows local optimal assignments to propagate and refine the global assignment over time, effectively "weaving" a better global OT plan without ever solving the full $O(n^3)$ problem.

2. Multiple Noise Caches (Preventing Overfitting):

   • A fixed assignment of noise to data can lead to overfitting, especially on smaller datasets, causing poor generalization when new noise is sampled at inference.
   • Solution: The method assigns multiple noise instances (caches) to each data point. During training, one noise sample is randomly selected from the cache for each data point.
   • Effect: This artificially increases the effective dataset size without changing the underlying data distribution, preventing overfitting to a static noise set while maintaining the benefits of the optimized coupling.

3. Convergence Analysis:

   • The authors prove that the algorithm converges to a stationary solution in a finite number of steps.
   • The final assignment has no "negative alternating cycles" of length less than the minibatch size $m$. While not guaranteed to be the global optimum, it is a significantly better approximation than standard minibatch OT.

3. Key Contributions

• LOOM-CFM Algorithm: A novel iterative method that improves CFM training by optimizing global data-noise assignments through the preservation of local minibatch OT solutions.
• Overfitting Mitigation: A computationally free technique using multiple noise caches per data point to prevent overfitting to fixed assignments, enabling robust inference with new noise samples.
• Theoretical Analysis: A formal convergence proof showing the method reaches a stationary transport plan with specific optimality properties regarding alternating cycles.
• Empirical Superiority: Demonstrated state-of-the-art performance in the speed-quality trade-off across multiple datasets (CIFAR10, ImageNet-32/64, FFHQ).

4. Experimental Results

The authors evaluated LOOM-CFM against baselines like OT-CFM, BatchOT, and Rectified Flows using the FID (Fréchet Inception Distance) metric at varying NFE (Number of Function Evaluations).

• CIFAR10:
  • LOOM-CFM reduced FID by 41% at 12 NFE compared to minibatch OT methods.
  • Achieved an FID of 4.60 at 12 NFE (vs. 7.77 for OT-CFM).
• ImageNet-32 & ImageNet-64:
  • ImageNet-32: 46% FID reduction at 12 NFE (3.89 vs. 7.18 for BatchOT).
  • ImageNet-64: 54% FID reduction at 12 NFE (8.49 vs. 18.27 for BatchOT).
• Reflow Initialization: LOOM-CFM serves as an excellent initialization for the Reflow algorithm (iterative retraining). A single Reflow step initialized with LOOM-CFM outperformed 3-step Rectified Flows, eliminating the need for multiple reflow iterations.
• High-Resolution Synthesis: Successfully applied to latent space training (FFHQ 256x256), achieving competitive results with significantly fewer NFE than previous latent flow matching methods.
• Efficiency: The method adds negligible disk overhead (storing seeds rather than full noise arrays) and converges faster in terms of epochs than BatchOT (e.g., 100 epochs vs. 575 for ImageNet-64).

5. Significance

• Bridging the Gap: LOOM-CFM effectively bridges the gap between the theoretical optimality of global OT and the practical constraints of large-scale deep learning. It achieves near-global optimality without the prohibitive computational cost.
• Inference Speed: By straightening sampling trajectories, LOOM-CFM drastically reduces the number of network evaluations required for high-quality generation, making flow-based models more viable for real-time applications.
• Simplicity and Compatibility: The method does not require auxiliary models or complex loss functions. It is a drop-in replacement for the coupling distribution in standard CFM training and is compatible with other acceleration techniques like distillation and latent space training.
• Scalability: The "noise cache" mechanism ensures the method scales effectively from small datasets (CIFAR) to large ones (ImageNet), solving the overfitting issue inherent in fixed-coupling approaches.

In summary, LOOM-CFM represents a significant advancement in generative modeling by optimizing the fundamental data-noise coupling mechanism, resulting in faster, higher-quality image synthesis with minimal computational overhead.