Training-Free Refinement of Flow Matching with Divergence-based Sampling

1. Problem Statement

Flow Matching (FM) has emerged as a powerful paradigm for generative modeling, learning a marginal velocity field to transport a simple prior distribution to a target data distribution. However, a fundamental limitation arises from trajectory crossings:

• The Conflict: At intermediate states, different sample-wise conditional velocities (connecting specific source-target pairs) may point in opposing directions.
• The Consequence: The learned marginal velocity field represents the average of these conflicting velocities. When conflicts occur, this average often points toward low-density regions (between modes) rather than toward any valid mode.
• The Result: Generative trajectories passing through these ambiguous states drift into low-density areas, resulting in blurry, degraded, or low-fidelity outputs.
• Existing Limitations: Previous solutions (e.g., HRF, VRFM) attempt to resolve this by redesigning the training objective to model multi-modal velocity distributions. However, these approaches require costly retraining and architectural changes, preventing the use of powerful, pre-trained "off-the-shelf" flow models.

Core Question: Can this degradation be mitigated at inference time without retraining the model?

2. Methodology: Flow Divergence Sampler (FDS)

The authors propose Flow Divergence Sampler (FDS), a training-free, plug-and-play framework that refines intermediate states during the sampling process.

Key Theoretical Insight

The paper establishes a theoretical link between the discrepancy (the error caused by conflicting velocities) and the divergence of the marginal velocity field.

• Theorem 1: The optimal Conditional Flow Matching (CFM) residual (the mean-squared discrepancy between the marginal velocity $u_t$ and sample-wise velocities $v_t$) can be expressed as a function of the divergence of the marginal velocity field ($\nabla \cdot u_t$) and known schedule coefficients.
• Implication: High divergence in the velocity field at a specific state $x_t$ indicates a high probability of velocity conflict and low reliability. Conversely, low divergence indicates a reliable, unambiguous region.
• Surrogate: Since the true sample-wise velocities are unavailable at inference, FDS uses the divergence of the pre-trained model $u_\theta$ as a data-free surrogate to estimate local ambiguity.

Algorithmic Workflow

FDS operates as a spatial refinement step inserted before standard ODE solver steps (e.g., Euler, Heun):

1. Perturbation: At a solver state $x_t$, the algorithm generates $M$ candidate states by adding small random perturbations: $x^{(m)} = x_t + \sigma_t \xi^{(m)}$, where $\xi \sim \mathcal{N}(0, I)$.
2. Divergence Evaluation: The divergence $\hat{\delta}_t(x) = \nabla_x \cdot u_\theta(x, t)$ is computed for each candidate using Hutchinson's trace estimator (to avoid expensive second-order derivatives).
3. Selection: The candidate with the lowest divergence is selected as the refined state $\tilde{x}_t$.
4. Integration: The standard solver step is performed starting from $\tilde{x}_t$ rather than the original $x_t$.
5. Efficiency: To minimize computational overhead, the method restricts refinement to the early stages of generation ($t < 0.5$) and typically uses only one iteration ($N=1$) and one candidate ($M=1$) beyond the base state.

3. Key Contributions

1. Theoretical Characterization: Proved that the conditional mean-squared discrepancy between marginal and sample-wise velocities is directly characterized by the divergence of the marginal velocity field, enabling a data-free proxy for ambiguity.
2. Training-Free Framework: Introduced FDS, a plug-and-play sampling method that steers trajectories away from high-discrepancy regions without modifying the underlying model or requiring retraining.
3. Plug-and-Play Versatility: Demonstrated compatibility with standard solvers (Euler, Heun) and various backbones (EDM, JiT, SD3), making it applicable to existing pre-trained models immediately.

4. Experimental Results

The authors evaluated FDS across diverse datasets and tasks, comparing it against baselines with matched computational budgets (increasing the number of function evaluations, NFEs, to match FDS's wall-clock time).

• Synthetic 2D Data: FDS successfully steered samples away from low-density regions in a checkerboard distribution, significantly reducing the Wasserstein Distance (WD) compared to standard FM.
• Image Generation (CIFAR-10 & ImageNet 256x256):
  • FDS consistently improved FID (Fréchet Inception Distance) and IS (Inception Score) across different solvers and models.
  • Crucially, FDS outperformed baselines that simply increased the number of solver steps. While more steps reduce temporal integration error, they do not fix the spatial drift caused by velocity conflicts.
  • Qualitative: FDS produced sharper details (e.g., distinct sesame seeds on a burger) that high-NFE baselines failed to generate.
• Text-to-Image (SD3-Medium): Applied to the DrawBench dataset, FDS improved human preference scores (ImageReward, HPSv2) and aesthetic scores across different guidance scales (CFG).
• Inverse Problems: In deblurring and super-resolution tasks, FDS enhanced the recovery of fine-grained details (e.g., fur texture, eye pupils) compared to standard TFG baselines.
• Comparison to Training-Based Methods: On CIFAR-10, FDS applied to a standard EDM model outperformed training-based methods like HRF (Hierarchical Rectified Flow) that specifically aimed to resolve crossings, despite HRF having a similar parameter count.

5. Significance and Impact

• Bridging the Gap: FDS demonstrates that inference-time spatial refinement can be more effective than increasing temporal resolution (more solver steps) for addressing fundamental flow matching limitations.
• Efficiency: It unlocks the full potential of existing, powerful pre-trained models without the prohibitive cost of retraining or architectural redesign.
• New Perspective: The work shifts the focus from purely temporal integration errors to spatial trajectory errors, suggesting that guiding the generation path through low-divergence regions is a critical factor in high-fidelity synthesis.
• Practicality: As a lightweight, training-free module, FDS is immediately deployable in production pipelines for image, audio, and video generation.