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

TL;DR

LOOM-CFM improves data-noise coupling in flow-based models by optimizing across minibatches, leading to faster, higher-quality image and video generation, and better high-resolution synthesis.

Contribution

The paper introduces LOOM-CFM, a novel method extending minibatch optimal transport to optimize data-noise assignments across training, enhancing inference speed and quality.

Findings

Consistent speed-quality improvements across datasets

Enhanced high-resolution synthesis capabilities

Better initialization for distillation processes

Abstract

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.