PaperlayerView exact revision on arXiv

Source-linked AI summary

Coefficients-Preserving Sampling for Reinforcement Learning with Flow Matching

Feng Wang, Zihao Yu

arXiv:2509.05952v4cs.CV

TL;DR

Flow-SDE sampling adds conspicuous noise to generated images, which can undermine reward estimation and policy learning. The paper proposes Coefficients-Preserving Sampling (CPS), inspired by DDIM, to keep latent noise consistent with the scheduler. The authors report improved reward optimization versus SDE-based approaches, while noting unresolved challenges in Flow Matching-based GRPO.

  • Problem

    Flow-SDE sampling corrupts generated images with noise artifacts, which can lead reward models to assign inaccurate scores and rankings and mislead policy updates.

    Source: §1 · p. 1

  • Method

    CPS reformulates sampling, inspired by DDIM, so the latent noise level remains consistent with the scheduler at every timestep.

    Source: §1 · p. 2

  • Results

    CPS reward optimization surpasses SDE-based approaches across tasks; aesthetic and human-preference rewards converge faster and reach higher terminal values, while detection rewards converge faster with similar final reward.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    CPS generates diverse, clean samples even under extremely high noise, and reward optimization guided by CPS outperforms SDE-based approaches on varied tasks.

    Source: §6 · p. 10

    Flow Matching-based GRPO methods still face reward hacking, multi-step credit assignment, and optimization of stochasticity arising from input noise.

    Source: §6 · p. 10

Abstract

from arXiv · show

Reinforcement Learning (RL) has recently emerged as a powerful technique for improving image and video generation in Diffusion and Flow Matching models, specifically for enhancing output quality and alignment with prompts. A critical step for applying online RL methods on Flow Matching is the introduction of stochasticity into the deterministic framework, commonly realized by Stochastic Differential Equation (SDE). Our investigation reveals a significant drawback to this approach: SDE-based sampling introduces pronounced noise artifacts in the generated images, which we found to be detrimental to the reward learning process. A rigorous theoretical analysis traces the origin of this noise to an excess of stochasticity injected during inference. To address this, we draw inspiration from Denoising Diffusion Implicit Models (DDIM) to reformulate the sampling process. Our proposed method, Coefficients-Preserving Sampling (CPS), eliminates these noise artifacts. This leads to more accurate reward modeling, ultimately enabling faster and more stable convergence for reinforcement learning-based optimizers like Flow-GRPO and Dance-GRPO. Code will be released at https://github.com/IamCreateAI/FlowCPS

1 INTRODUCTION

The introduction identifies noisy Flow-SDE samples as a threat to reward estimation and policy learning, and presents CPS as a coefficient-preserving alternative inspired by DDIM. Experiments report stronger convergence and terminal rewards for aesthetic and human-preference rewards, with faster convergence and similar final reward for detection rewards.

1 INTRODUCTION

  • Flow-SDE sampling introduces conspicuous noise artifacts that can make aesthetic-quality and human-preference reward models assign inaccurate scores and rankings, misleading policy updates.The noisy samples are the inputs used to compute rewards that guide policy updates.Source: §1 · p. 1, Figure 1 · p. 2
  • Flow-SDE injects excess noise while retaining the original ODE scheduler, so noise accumulates into a nonzero final level and visibly noisy outputs.The analysis traces the mismatch to the SDE score-function term and the Wiener-process noise level.Source: §1 · pp. 1–2
  • CPS reformulates sampling, inspired by DDIM, to keep latent noise levels consistent with the scheduler at every timestep.The method generalizes DDIM to Flow Matching and is designed to produce high-fidelity images even under high noise levels.Source: §1 · p. 2
  • Aesthetic and human-preference reward models converge faster and reach higher terminal rewards with CPS, while detection rewards converge faster with similar final reward.The experiments compare CPS against Flow-SDE sampling across multiple baseline models and reward functions.Source: §1 · p. 2
  • The paper introduces the coefficient-preserving requirement and shows that Flow-SDE fails it; it also traces excess noise to Taylor-expansion approximation errors and a 1/t source of numerical instability.The Taylor expansion used in the Flow-SDE derivation is identified as the source of the approximation errors and numerical instability.Source: §1 · p. 2
  • The authors report that CPS facilitates reward estimation and optimization, substantially outperforming Flow-SDE-based sampling.Source: §1 · p. 2

2 RELATED WORK

Related work covers reinforcement learning alignment for language and generative models, then reviews diffusion sampling methods. It positions the paper's sampling approach alongside deterministic and stochastic samplers and recent coefficient-matrix unification.

2.1 ALIGNMENT FOR LARGE LANGUAGE MODELS

2.2 ALIGNMENT FOR DIFFUSION

2.3 DIFFUSION SAMPLER

  • Diffusion samplers span deterministic DDIM (Song et al. (2021a)), ODE solvers from Euler to DPM-Solver, and concurrent coefficient-matrix unification (Zheng & Zheng (2025)).The passage describes Euler as faster but less accurate than Heun’s method, while high-order solvers reduce required sampling steps.Source: §2.3 · p. 3

3 PRELIMINARIES

The preliminaries introduce Flow Matching, Flow-GRPO and Dance-GRPO, GRPO, and DDIM as foundations for the proposed algorithm. They contrast deterministic ODE sampling with stochastic variants that inject noise.

Flow Matching

  • Flow Matching interpolates data and Gaussian-noise samples, trains a network to predict velocity, and uses a deterministic ODE sampler.The estimated velocity is denoted by the model-predicted quantity in the article.Source: §3 · p. 3

Flow-GRPO and Dance-GRPO

  • Flow-GRPO and Dance-GRPO introduce randomness into Flow Matching by converting its deterministic Flow-ODE into Flow-SDE sampling.The preliminaries relate this stochastic sampling to the need for diverse samples in reinforcement learning.Source: §3 · p. 3

GRPO

  • GRPO optimizes a policy objective after diverse images are sampled and rewards are transformed into advantages; the KL loss is described as closed form.The supplied passage also states that the mean predicted sample is implemented by removing injected noise.Source: §3 · p. 4

DDIM

  • DDIM permits stochasticity through a noise parameter: zero yields an implicit probabilistic model, while a particular value makes the forward process Markovian and the generative process DDPM.The text compares DDIM and ODE as deterministic, and DDPM and SDE as stochastic counterparts.Source: §3 · p. 4

4 ANALYSIS AND METHODS

The section defines coefficients-preserving sampling (CPS), diagnoses why Flow-SDE violates it, and presents a DDIM-inspired alternative. It also relates Flow-SDE to a Taylor approximation of Flow-CPWS and discusses limitations of that approximation.

4.1 COEFFICIENTS-PRESERVING SAMPLING

  • CPS requires the scheduler to set the sample coefficient at every timestep and the total noise level to match the scheduler.The definition assumes predicted noise is standard Gaussian and independent of newly injected noise.Source: §4.1 · p. 5
  • DDIM is an example of CPS: its sample and noise coefficients match the scheduler, so their combined noise level remains aligned for any σt.Source: §4.1 · p. 5

4.2 FLOW-SDE IS NOT COEFFICIENTS-PRESERVING SAMPLING

  • Flow-SDE adds excess noise at each timestep, violating CPS and producing noisy final images; its noise-level mismatch worsens at fewer sampling steps.The mismatch is reported for both Flow-GRPO and Dance-GRPO.Source: §4.2 · p. 5, Figure 2 · p. 6

4.3 OUR SOLUTION

  • Inspired by DDIM, the proposed formulation injects noise while satisfying CPS, with η controlling stochastic strength.The authors name the resulting sampling algorithm Flow-CPS.Source: §4.3 · p. 6
  • The method removes σt from the log-probability denominator and removes the denominator in the KL loss, shifting log-probability weight toward earlier, more diverse timesteps.The stated motivation includes avoiding division by zero or very small values at the last timestep.Source: §4.3 · pp. 6–7

4.4 DISCUSSION

  • The section also considers a Wiener-process version of the sampling function and derives an approximate SDE from Flow-CPWS using Taylor expansion.Source: §4.4 · p. 7
  • Flow-SDE is a first-order Taylor approximation of Flow-CPWS, which the paper names for its sampling function with a Wiener process.The approximate equality to Flow-SDE follows from the Taylor expansion.Source: §4.4 · p. 7
  • With distillation, sampling steps can fall to 4 or even 1, so ∆t→0 no longer holds, the fundamental reason Flow-SDE produces inaccurate noise levels; its Taylor-expansion t term yields significant error around t=0 even with many steps.Source: §4.4 · p. 8

5 EXPERIMENTS

The section evaluates Flow-CPS on verifiable and preference-reward tasks against Flow-GRPO and Dance-GRPO. It reports stronger performance on three tasks, parity on near-saturated GenEval, and faster convergence.

5.1 EXPERIMENTAL SETUP

  • Experiments compare two GRPO baselines across GenEval, Text Rendering, PickScore, and HPSv2, using 8 NVIDIA A100 GPUs.The experiments change the sampling method and log-probability while following the baselines’ settings.Source: §5 · p. 8, §5.1 · p. 8
  • GenEval and OCR provide verifiable rewards, while PickScore and HPSv2 score images using models trained on human preferences.GenEval evaluates compositional object properties; OCR compares rendered text with the prompt.Source: §5.1 · p. 8
  • The experiments omit KL loss in most tasks for training speed, but retain β = 0.001 for GenEval after omission degraded performance.The figure caption also notes significant degradation or collapse without KL loss on GenEval.Source: §5.1 · p. 9, Figure 5 · p. 9

5.2 CLEAN IMAGE SAMPLING

  • Flow-CPS produces diverse, noise-free images and higher early rewards than Flow-SDE, with a smaller train-test discrepancy.The discrepancy is attributed to deterministic inference after stochastic training.Source: §5.2 · p. 9

5.3 EXPERIMENTAL RESULTS

  • Flow-CPS outperforms Flow-GRPO and Dance-GRPO on PickScore, HPSv2, and OCR, while matching them on near-saturated GenEval.The results are reported in Tables 1–4, respectively.Source: §5.3 · p. 9
  • Flow-CPS reaches the optimal result faster; both methods perform best at η = 0.7, while neither converges properly at η = 0.1.The convergence comparison and η ablation are reported in Figures 5 and 6.Source: §5.3 · p. 10

6 CONCLUSION

The conclusion presents Coefficients-Preserving Sampling as a response to image noise in SDE-based sampling. It also notes that several challenges remain for Flow Matching-based GRPO.

6 CONCLUSION

  • CPS addresses image noise in SDE-based sampling; its analysis characterizes SDE as a first-order Taylor approximation of CPS.CPS can generate diverse, clean samples even under extremely high noise, and CPS-guided reward optimization surpasses SDE-based approaches on varied tasks.Source: §6 · p. 10
  • Flow Matching-based GRPO still faces reward hacking, multi-step credit assignment, and optimization of stochasticity from input noise.These are identified as unresolved issues for further research.Source: §6 · p. 10
  • The solution defined by Equations 13 and 15 can be implemented in 10 lines of code provided in the supplementary material.Source: §6 · p. 10

A VP-SDE IS AN APPROXIMATION OF DDPM

The appendix relates VP-SDE to DDPM through a first-order Taylor approximation. It describes the omitted higher-order terms and the reverse-process interpretation.

A VP-SDE IS AN APPROXIMATION OF DDPM

B DPM-SOLVER SERIES

The appendix checks DPM-Solver variants against the coefficient-preserving requirement and reports OCR-reward training behavior. It identifies instability in SDE-DPM-Solver++1 and variation between runs.

B DPM-SOLVER SERIES

  • SDE-DPM-Solver-1 is not coefficient-preserving because its total noise matches the requirement only when eh = 1.Its sample coefficient matches the first CPS condition, but its total noise does not generally meet the second.Source: B DPM-SOLVER SERIES · p. 13
  • SDE-DPM-Solver++1 matches CPS; Flow-CPS is presented as a DDIM special case with η retained to tune injected noise.The passage also characterizes SDE-DPM-Solver++1 as a DDIM special case.Source: B DPM-SOLVER SERIES · p. 14
  • With OCR reward, Flow-CPS becomes unstable at η ≤ 0.5, while SDE-DPM-Solver++1 has comparable sudden reward drops associated with low diversity.The equivalent η for SDE-DPM-Solver++1 varies across timesteps, from 1.0 down to about 0.43 and then back to 1.0.Source: B DPM-SOLVER SERIES · p. 14, Figure 7 · p. 14
  • SDE-DPM-Solver++1 OCR runs peak at rewards of 0.966 and 0.970, illustrating variable outcomes that the authors say harm reproducibility.The passage says multiple trials are needed to obtain a satisfactory model.Source: B DPM-SOLVER SERIES · p. 14, Table 5 · p. 14
  • For Equation 29, the ideal-versus-SDE noise error is described as ignorable when η ≤ 1 and the step count is at least four.Source: Figure 8 · p. 15

C AN ALTERNATIVE FOR THE NUMERICAL PROBLEM

A modified reverse Flow-SDE is considered as an alternative for the numerical problem. It reduces error but still does not meet CPS requirements, though it may be useful when SDE characteristics are needed.

C AN ALTERNATIVE FOR THE NUMERICAL PROBLEM

  • Setting σt = η(t−∆t) modifies the reverse Flow-SDE as a patch for the numerical problem.The passage introduces this choice as a possible patch when ∆t ≪ t−∆t and ∆t → 0.Source: C AN ALTERNATIVE FOR THE NUMERICAL… · p. 15
  • The modified formula has smaller error than the original Flow-SDE but still fails to meet CPS requirements.It may be useful when SDE characteristics are necessary.Source: C AN ALTERNATIVE FOR THE NUMERICAL… · p. 15
  • Images sampled by Equation 28 with η = 1 show no obvious noise.Source: Figure 9 · p. 15
  • The variant using denominator 2σ2t fails to converge and is omitted from the log-probability ablation figure.Source: Figure 10 · p. 16

D ABLATION ON THE LOGPROB

The numerical modification initially accelerates convergence, but final performance is comparable to the original version.

D ABLATION ON THE LOGPROB

  • The modification initially accelerates convergence, but final performance is comparable to the original version.The change addresses numerical instability from division by near-zero values in the final diffusion steps and was also applied to Flow-GRPO for an ablation study.Source: D ABLATION ON THE LOGPROB · p. 16

E QUALITATIVE RESULTS

Images optimized by PickScore and HPSv2 illustrate that higher reward scores do not necessarily correspond to better-looking images.

E QUALITATIVE RESULTS

Related papers

Loading 2509.05952v4…