PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Divya Jyoti Bajpai, Arun Verma, Manjesh Kumar Hanawal

arXiv:2610.03202v1cs.CVcs.AI

TL;DR

Flow Matching generation remains costly because inference requires sequential model evaluations, and input-dependent needs complicate choosing a fixed step budget. COFLOW selects budgets online from prompt features using a reward tied to trajectory dynamics. Across image, editing, and video tasks, it reduces computation while maintaining competitive quality, with reported acceleration above 2.5×.

  • Problem

    Flow Matching inference requires multiple sequential model evaluations, while generation complexity and computational needs vary across prompts.

    Source: §1 · pp. 1–2

  • Method

    COFLOW uses an online neural contextual bandit to select prompt-dependent step budgets, with an unsupervised trajectory-variation surrogate balancing computation and refinement.

    Source: §1 · p. 2, §3.2 · p. 4

  • Results

    Across image generation, editing, and video synthesis, COFLOW reduces computation while maintaining competitive quality; experiments report inference acceleration (>2.5×).

    Source: §1 · p. 2, §5 · p. 9

  • Takeaways & Limitations

    COFLOW adapts existing solver trajectories without modifying the pretrained generative model and applies across multiple visual-generation tasks.

    Source: §1 · p. 2, §4.1 · p. 9

    COFLOW requires a brief exploration phase, and adaptation depends on the expressiveness of its intentionally simple prompt features.

    Source: §4.1 · p. 9

Abstract

from arXiv · show

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 COFLOW, an inference-time method that adaptively selects the step counts each generation based on the prompt features. Our context-aware COFLOW is trained online with an unsupervised reward that balances inference efficiency and generation fidelity. Our method is plug-and-play, requiring no retraining of the underlying generative model. It generalizes to image and video generation, achieving over 2.5x speedup while preserving perceptual and semantic quality. We further provide a theoretical analysis establishing an O(1/K) forward-Euler discretization error bound under standard regularity conditions.

1 Introduction

Flow Matching offers high-quality visual generation, but sequential integration makes inference costly. COFLOW adapts step budgets to prompt-specific needs, aiming to reduce computation while preserving generation quality.

1 Introduction

  • Sequential numerical integration requires multiple model forward passes, making inference latency a central challenge as models, resolutions, and temporal horizons scale.The serial computation is especially costly for practical deployments.Source: §1 · p. 1
  • Existing acceleration methods can add training costs or rely on expensive step-importance estimates whose errors may degrade generation quality.Distillation, step truncation, and consistency-based training reduce function evaluations but can require teacher models or large datasets.Source: §1 · pp. 1–2
  • Different prompts can require different step budgets; AdaDiff adapts to this variability but requires costly policy pretraining and may be sensitive to distribution shifts.Figure 1 illustrates different behavior across integration budgets for different inputs.Source: §1 · p. 2, Figure 1 · p. 2
  • COFLOW uses a contextual bandit to predict each generation’s step budget online from prompt-specific features and a reward balancing computation with trajectory refinement.Its reward is motivated by the observation that complex prompts induce larger velocity variations across timesteps.Source: §1 · p. 2
  • Across image and video benchmarks, COFLOW reports inference acceleration above 2.5× while maintaining quality, and the paper provides an O(1/K) discretization-error bound.The framework is described as task-agnostic and applicable to image generation, image editing, and video generation.Source: §1 · p. 2

2 Related Work

Prior work accelerates flow-based generation through retraining or changes to inference-time sampling. COFLOW is positioned against methods that require extra training or rely on heuristic signals and pretrained policies.

2 Related Work

  • Flow Matching [1] [18] [19] [20] is presented as an alternative to diffusion models [3] [21] [22] [23], supporting applications from image inversion and editing to video synthesis.Related advances include multimodal extensions [27] [28], large-scale flow models such as FLUX [19], and pyramidal flow matching [29].Source: §2 · pp. 2–3
  • Retraining-based approaches, including InstaFlow [31], LeDiFlow [32], and Diff2Flow [33], reduce sampling steps but require additional training, data, or specialized objectives.Other approaches modify objectives to straighten trajectories or combine adversarial and consistency objectives [10] [11] [34].Source: §2 · p. 3
  • Inference-time methods such as TeaCache [13], FastFlow [15], and FlowCast [14] use trajectory signals or speculative decoding, while AdaDiff trains a policy using generations at multiple step budgets.The paper contrasts COFLOW’s online trajectory-signal approach with AdaDiff’s offline data and evaluation overhead; heuristic inference-time methods may introduce inconsistencies or degrade quality.Source: §2 · p. 3

3 Methodology

COFLOW builds on Flow Matching’s continuous flow trajectories by selecting an input-dependent integration budget online. A neural contextual bandit uses prompt features and a trajectory-variation reward to allocate computation, while the pretrained model and solver trajectory remain unchanged.

3.1 Flow Matching Overview

  • Flow Matching learns a time-dependent velocity field that transports source-distribution samples toward the target; forward Euler approximates its trajectory in K steps.The step count trades computational cost against discretization error and generation quality.Source: §3.1 · p. 3
  • COFLOW’s adaptive budget responds to input-dependent trajectory variation, since fixed budgets can be conservative for some inputs and insufficient for others.The selected budget is applied to the corresponding full solver trajectory.Source: §3.1 · p. 3, Figure 2 · p. 4

3.2 Algorithm Design

  • COFLOW selects a step budget online with a neural contextual bandit, then generates the full trajectory using that budget.An initial ε-greedy warm-up is followed by choices based on reward estimates and confidence bounds; observed rewards update the model.Source: §3.2 · p. 4, §3.3 · pp. 5–6
  • The unsupervised reward favors lower relative velocity variation while penalizing additional steps, but remains a surrogate rather than a direct measure of generation quality.The variation statistic reflects velocity changes along the trajectory induced by a candidate budget.Source: §3.2 · p. 4
  • Prompt context combines semantic embeddings, entity count, and mean word length to guide adaptive step allocation.The reward model predicts expected reward for context–budget pairs and selects using predicted reward plus an uncertainty term.Source: §3.2 · p. 5

3.3 Algorithm: COFLOW

  • Forward-Euler discretization error is bounded by O(1/α), with an input-dependent constant determined by velocity-field variation along the exact trajectory.The observed relative-velocity-variation statistic is a surrogate motivated by this dependence, not a direct estimate of the error or its constant.Source: §3.3 · p. 6

4 Experiments

COFLOW evaluates image generation, image editing, and video generation, comparing input-adaptive step selection with full-generation and acceleration baselines. Across these settings, it reduces inference computation while maintaining competitive output quality.

4 Experiments

  • Evaluation uses 553 GenEval prompts, 606 GEdit instructions, and 80 VBench prompts sampled across 16 dimensions, with BAGEL, FLUX, and HunYuanVideo models.The image-generation models also include PeRFlow and StableDiffusion; BAGEL and Flux-Kontext are used for editing.Source: §4 · p. 7
  • The comparison baselines span full generation, TeaCache, InstaFlow, PeRFlow, FlowCast, FastFlow, and AdaDiff.They cover full sampling, training-based acceleration, and inference-time approaches.Source: §4 · p. 7
  • COFLOW predicts an input-adaptive step budget with a neural contextual bandit, then runs the full trajectory; its online policy uses prompt features and replay-buffer updates.The context combines handcrafted complexity features and projected semantic embeddings; training updates the reward model periodically.Source: §4 · p. 7

4.1 Results

  • 0.77 overall score at 17.7 NFE vs 0.78 at 50 NFE on BAGEL; 0.65 at 19.5 NFE vs 0.66 at 50 NFE on FLUX.The reported comparison is on GenEval; the table covers generation quality and efficiency metrics.Source: Table 1 · p. 8, §4.1 · p. 9
  • Across BAGEL, FLUX, and Step-1X-Edit, COFLOW reduces the integration budget while maintaining competitive semantic consistency and perceptual quality in image editing.Edit quality is assessed with G_SC, G_PQ, and G_O, using GPT-4.1 as an automatic judge.Source: Figure 3 · p. 8, §4.1 · p. 9
  • In video generation, COFLOW reduces the integration budget while maintaining competitive performance on VBench and frame-level BRISQUE quality.VBench evaluates motion dynamics, temporal consistency, and scene composition.Source: Figure 4 · p. 8, §4.1 · p. 9
  • COFLOW requires a brief exploration phase, and adaptation depends on the expressiveness of its intentionally simple prompt features.Richer features could improve adaptation quality but add computational overhead; reported results include exploration samples.Source: §4.1 · p. 9

5 Conclusion

COFLOW selects input-dependent step budgets online for flow-based generation without retraining the generative model. Across image generation, editing, and video synthesis, it reduces computation while maintaining competitive quality, and the paper analyzes discretization error under regularity conditions.

5 Conclusion

  • COFLOW uses lightweight prompt representations and a trajectory-variation surrogate in a neural contextual bandit to select step budgets without retraining the generative model.Source: §5 · p. 9
  • Across image generation, image editing, and video synthesis, COFLOW reduces inference computation while maintaining competitive semantic and perceptual quality across multiple backbones.Source: §5 · p. 9
  • The paper establishes forward-Euler discretization-error dependence of O(1/K) under standard regularity conditions and reports that velocity variation tracks temporal-discretization sensitivity.Source: §5 · p. 9

A.1 Proof of Theorem 1

The proof uses forward-Euler updates and regularity assumptions to analyze numerical integration error. The section also illustrates trajectory-level velocity variation across prompts and integration budgets.

A.1 Proof of Theorem 1

  • For input x_i and budget α, the proof uses a uniform grid t_k = kh with h = 1/α; the initial error is zero because both trajectories start at the same state.Source: §A.1 · p. 13
  • The local truncation error is expressed through the difference between exact and Euler updates, using the fundamental theorem of calculus and the chain rule along the exact trajectory.Source: §A.1 · p. 13
  • Velocity-change profiles vary across prompts and integration budgets, illustrating the trajectory-level signal used by COFLOW.Source: Figure 6 · p. 14
  • The error analysis assumes a velocity field Lipschitz in the state variable and a uniform bound on its derivative along the exact trajectory; the latter constant captures input-dependent flow variation.Source: §A.1 · pp. 13–14

A.2 Expanded Results on the GenEval Dataset

The expanded evaluations examine GenEval category performance, contextual-bandit overhead, and combining COFLOW with existing inference-time accelerators.

A.2 Expanded Results on the GenEval Dataset

  • COFLOW retains competitive GenEval performance across prompt categories while using substantially fewer function evaluations than full-step inference.The reported categories span object composition, counting, attribute binding, and spatial relations.Source: §A.2 · pp. 14–15

A.3 Computational Overhead of the Contextual Bandit

  • Bandit decision latency is <1 second, and neural-bandit training takes ∼3–4 seconds in this setup.The decision module is a small MLP with ≈30k parameters.Source: §A.3 · p. 15, Table 2 · p. 15

A.4 Complementarity of Our Approach

  • Combining COFLOW with TeaCache, FastFlow, and FlowCast further reduces runtime versus each accelerator alone, though gains and quality effects vary by method.COFLOW selects the integration budget, and the existing acceleration method is then applied during generation without modifying the underlying generative model.Source: §A.4 · p. 15

A.5 Qualitative Analysis

The examples compare how AdaDiff and COFLOW allocate steps across prompts and the resulting visual outputs.

A.5 Qualitative Analysis

  • COFLOW assigns intermediate step budgets across inputs, whereas AdaDiff often chooses a relatively small or maximum budget.The figure contrasts the approaches' selected budgets and qualitative outputs.Source: §A.5 · p. 15, Figure 7 · p. 16
  • Across the examples, COFLOW reduces the integration budget while retaining visually comparable outputs.The examples illustrate different computational allocations for individual inputs.Source: §A.5 · p. 15, Figure 7 · p. 16

A.6 L1-Relative Velocity Change as a Surrogate for Trajectory Variation

The section evaluates L1-relative velocity change as a trajectory-level signal associated with discretization sensitivity, while distinguishing it from prompt-complexity ranking.

A.6 L1-Relative Velocity Change as a Surrogate for Trajectory Variation

  • ˆv is a model-internal L1-relative velocity-change measure along a trajectory, not a direct estimate of discretization error or perceptual quality.It requires neither a reference solution nor an external quality model.Source: §A.6 · p. 16
  • 0.924 mean within-prompt Spearman correlation links ˆv with reference deviation across budgets; pooled correlation is ρ = 0.822 (p < 10^-16).The association also appears with perceptual LPIPS distance (Spearman ρ = 0.676, p < 10^-11).Source: §A.6 · p. 16
  • ˆv tracks within-trajectory discretization sensitivity, not independent prompt-difficulty ranking; fixed-budget cross-prompt correlation was 0.255 in this pilot.Both ˆv and reference deviation vary systematically with integration budget.Source: §A.6 · p. 16

A.7 Context Matters

The ablations examine contextual information, context representations, online learning behavior, and uncertainty-aware exploration in adaptive step selection.

A.7 Context Matters

  • Contextual UCB raises overall GenEval score from 0.74 to 0.77 while reducing average NFE from 20.2 to 17.7 versus non-contextual UCB; Thompson Sampling reaches 0.76 at 17.5 NFE.The comparison tests contextual and non-contextual policies and an alternative exploration strategy.Source: §A.7 · p. 17
  • The full context representation reduces mean NFE from 39.75 to 37.67 and maximum-budget selections from 65.4% to 47.9% versus no context.Handcrafted features achieve 37.79 NFE, while MiniLM alone yields 39.58 NFE.Source: §A.7 · p. 18
  • Mean NFE falls from 43.2 in rounds 1–10 to 37.4 in rounds 81–100, while mean surrogate reward rises from 0.520 to 0.580.Across prompt permutations, final mean NFE is 37.8 ± 1.1 and mean surrogate reward is 0.579 ± 0.009.Source: §A.7 · p. 18
  • NeuralUCB achieves higher mean surrogate reward (0.582 vs. 0.541) and lower mean NFE (37.5 vs. 39.6) than the greedy variant.It also selects the maximum step budget less frequently (49.7% vs. 63.4%) and has higher action entropy (1.54 vs. 1.21).Source: §A.7 · p. 18

A.8 Direct Step Reduction vs. Trajectory Manipulation

In the shown examples, direct step reduction retains better visual quality than trajectory truncation at a comparable computational budget. COFLOW selects an input-dependent budget and executes its corresponding solver schedule.

A.8 Direct Step Reduction vs. Trajectory Manipulation

  • Direct reduction constructs a schedule for the reduced integration budget, while trajectory truncation skips intermediate evaluations along a higher-step trajectory.Source: §A.8 · p. 18
  • In the shown examples, direct step reduction retains better visual quality than trajectory truncation at a comparable computational budget.Source: §A.8 · p. 18
  • COFLOW selects an input-dependent step budget and executes the corresponding solver schedule rather than pruning a high-step trajectory.Source: §A.8 · p. 19

A.9 Performance Across Solver Families

Across solver families, COFLOW reduces function evaluations while maintaining competitive semantic alignment and perceptual quality, with mild solver-dependent sensitivity. Context features combine sentence embeddings and handcrafted prompt properties.

A.9 Performance Across Solver Families

  • Approximately 2.5× fewer function evaluations across solver families, with competitive GenEval semantic alignment and IQA perceptual quality.Source: §A.9 · p. 19
  • DPM++ variants show slightly larger GenEval degradation, while Heun demonstrates greater robustness under reduced discretization.Source: §A.9 · p. 19

B Context Features

  • The context combines MiniLM sentence embeddings with lightweight handcrafted features capturing semantic, compositional, and structural prompt properties.Handcrafted cues include word count, lexical complexity, entity and attribute densities, numerical constraints, and punctuation density; feature computation adds only a few milliseconds per prompt.Source: B Context Features · pp. 19–20
  • Table 10 evaluates adaptive COFLOW on PeRFlow models trained for 5, 8, and 10 steps, reporting competitive performance with fewer NFEs.Source: Table 10 · p. 20
  • All context features are normalized to [0, 1] to stabilize scaling and prevent any single feature from dominating the representation.Source: B Context Features · p. 20

B.1 Choice of Context Representation as MiniLM

This section describes COFLOW’s prompt-context representation—MiniLM semantic embeddings combined with handcrafted complexity features, presented as a lightweight, model-agnostic alternative to the native text encoder—and notes that an LLM was used only for editing and formatting, not as an important, original, or non-standard component of the core method.

B.1 Choice of Context Representation as MiniLM

  • COFLOW combines all-MiniLM-L6-v2 sentence embeddings with handcrafted complexity features to form a lightweight, expressive prompt representation.The embeddings add higher-level semantic information to lexical and structural cues while remaining inexpensive to compute for the contextual decision module.Source: §B.1 · p. 20
  • The authors avoid the native prompt encoder because its high-dimensional representations increase bandit input size, computational cost, and training instability.They argue the external pretrained encoder is reusable across generative backbones, including BAGEL and FLUX, without modification.Source: §B.1 · p. 21

NeurIPS Paper Checklist

  • The checklist reports an anonymized source-code repository and hyperparameter settings in Table 12, and says inference-time evaluation requires no train splits.It states the repository link is in the abstract and experimental details are also provided in Section 4.Source: NeurIPS Paper Checklist · pp. 25–26
  • The authors report no error bars, citing standard benchmark metrics and qualitative comparisons as their evaluation practice.The checklist answers [No] to reporting error bars or other statistical-significance information.Source: NeurIPS Paper Checklist · p. 26
  • The authors state that the work has no societal impact because it focuses on inference-time speedup of existing methods.The checklist also says no additional safeguards are required because the method is an inference-time speedup method.Source: NeurIPS Paper Checklist · pp. 27–28
  • The authors mark the LLM-use declaration as not applicable because the LLM was used only for editing and formatting.They state this use does not affect the research’s core methodology, scientific rigor, or originality.Source: NeurIPS Paper Checklist · p. 30
  • The checklist defines N/A to mean that LLMs are not important, original, or non-standard components of the core method.Source: NeurIPS Paper Checklist · p. 30

Related papers

Loading 2610.03202v1…