Contextual Flow Matching: Adaptive Step Selection in Flow Models for Efficient Visual Generation

Flow Matching enables high-quality visual generation via continuous-time dynamics, but inference remains costly due to multiple sequential function evaluations. Existing acceleration methods reduce the number of function evaluations but often introduce additional training overhead, degrade quality, or fail to account for input-dependent variability. We propose an inference-time method, CoFlow, that adaptively selects the step counts for each generation based on the prompt features. Our context-aware CoFlow is trained online with an unsupervised reward that balances efficiency and fidelity. Our method is plug-and-play, requiring no retraining of the generative model. It generalizes to image and video generation, achieving over $2.5\times$ speedup while preserving perceptual and semantic quality. We also provide theoretical insights into the link between adaptive step allocation and discretization error. The anonymized source code is available at \url{https://anonymous.4open.science/r/Contextual_Flow_Matching-9970}.