CAT-Flow: Curvature-Adaptive sTeps for Flow Matching

Flow matching has solidified its position as a leading framework for generative modeling, powering state-of-the-art systems such as FLUX and Stable Diffusion 3.5. However, the iterative nature of its ODE-based sampling process creates a fundamental efficiency bottleneck: the quality of generated samples is highly sensitive to the choice of step-sizes, and current models typically require 20 to 30 steps to reach satisfactory quality. To this end, we propose two lightweight, training-free, and efficient algorithms, CAT-OV and CAT-OT, which adapt step-sizes at inference time. This is made possible by establishing a novel and provable connection between flow matching sampling and gradient flow. Without requiring additional neural function evaluations, CAT-OT estimates curvature over time via a finite-difference approximation of the time-derivative of the vector field, while CAT-OV estimates curvature as a variance of past velocities via a gradient vector field. Under suitable conditions, both methods have truncation error bounds of constant order. Empirically, across four text-to-image flow matching models, CAT-OV and CAT-OT match the image quality of existing step-size heuristics with up to 40% fewer generation steps and negligible added wall-clock time, with the largest gains in metric scores in the low-step regime. We also validate our methods on audio generation, showing efficiency gains that generalize beyond image generation, and demonstrate compatibility with higher-order solvers. Beyond adapting step-sizes, our gradient flow connection opens a broader avenue for transferring optimization techniques to flow matching inference, with proof-of-concept demonstrations including momentum and gradient normalization extensions.