Precise: SDE-Consistent Stochastic Sampling for RL Post-Training of Flow-Matching Models

Paper: arXiv:2605.23522 Code status: 代码搜索未找到开源实现;论文 PDF 声明的 https://github.com/Tencent-Hunyuan/Precise 当前通过 GitHub API / repo search 均不可访问(404 / 0 hits),因此本文笔记不设置 github_ref

1. Motivation (研究动机)

现有 image / text-to-image flow-matching model 已经可以用 RL post-training 提升 prompt alignment 和 perceptual quality,但在线 RL 需要一个“可采样、可探索、可计算 policy ratio”的 stochastic rollout policy。Flow matching 原本常用 deterministic probability-flow ODE 采样;把它直接换成 reverse-time SDE 后,sampler 本身就成为 policy 的一部分:它决定每一步注入多少随机性,也决定小步数 NFE 下这条 stochastic trajectory 是否仍忠实于目标 flow-matching marginal。

本文指出当前 sampler 的核心瓶颈不是单纯“噪声太大/太小”,而是两个耦合问题:

  • exploration schedule 控制随机探索强度,也同时放大 score term;如果 late denoising 时仍注入过强随机性,模型需要额外纠错,细节容易被破坏;如果 early denoising 太保守,又缺少探索。
  • finite-step discretization:RL 训练通常只用 10–30 NFE,而不是无限小步长。即使连续时间 reverse SDE 在 exact score 下保持 marginal,小步数 Euler-Maruyama 也可能引入额外 discretization noise。
  • posterior-mean collapse:CPS 这类 coefficient-preserving sampler 能约束噪声系数,但在非退化数据分布上只使用 posterior mean,会丢掉 的 residual uncertainty,导致 marginal covariance contraction。

论文要解决的具体问题是:为 flow-matching model 的 GRPO-style RL post-training 设计一个 stochastic sampler,使其既能给 RL 足够探索,又在少量 NFE 下尽量保持 reverse-SDE 一致性,从而减少训练不稳定、sample quality 下降和 wall-clock 浪费。这个问题值得研究,是因为 sampler 不是“训练细节”,而是 RL policy 的转移核;sampler 偏差会直接改变 reward optimization 看到的 trajectory distribution。

2. Idea (核心思想)

Precise 的核心 insight 是:flow-matching RL 的 stochastic sampler 应同时在 logSNR 空间控制“额外 denoising burden”,并用一个 SDE-consistent finite-step transition 消除 Euler 小步数下的 excess noise。它不是只调大/调小 ,而是把 exploration schedule 和 finite-step transition 作为同一个 reverse SDE policy 的两个部分来设计。

关键创新有两点。第一,作者用 logSNR decomposition 推出 ,让 score-induced progress 与 velocity-induced progress 的比例 在时间上保持常数。第二,在一阶 reverse step 内假设 clean-latent posterior mean 近似冻结,得到闭式 transition:新噪声的反馈不再被 Euler 忽略,transition 直接匹配 frozen-mean SDE 的 exact solution。

与 Flow-GRPO / Dance-GRPO 的区别在于,后者使用 Euler-style discretization,fresh noise 注入后仍把 velocity/score 冻结在旧状态,少步数时会产生 excess discretization noise;与 CPS 的区别在于,CPS 只保留 coefficient constraint ,但没有修复 posterior-mean 引起的 covariance contraction。Precise 试图同时覆盖这两个 failure mode。

3. Method (方法)

3.1 Overall framework:把 sampler 当作 RL policy 的转移核

Flow matching 的 forward path 写作: 模型学习 velocity 。为了做 RL post-training,Flow-GRPO / Dance-GRPO 将 deterministic ODE 替换为 reverse-time SDE: 其中 控制 exploration, 是 Brownian motion。连续时间、exact score 的理想情况下,任意 都能保持 flow-matching marginal;但 RL 实际训练用 imperfect denoiser 和小 NFE,因此 与 discretization 都会改变 reward optimization 的行为。

Figure 1 解读:左图把 sampler design 拆成两个轴:exploration-stability balance 与 SDE consistency。Dance-GRPO 的 constant noise schedule 在探索/稳定上不够自适应;Flow-GRPO 的 Euler transition 有 discretization inconsistency;CPS 提升了 coefficient consistency 但仍有 marginal bias。右图的 double-ring toy example 展示了这种差异:Flow-GRPO 会把分布打得过散,CPS 会让 mass 偏向 inner ring,而 Precise 目标是让 stochastic rollout 更接近目标 marginal。

直觉上,RL 需要 sampler 注入随机性来探索 reward landscape,但每一份随机性都会让 denoiser 多承担一次“把 noisy latent 拉回 clean manifold”的负担。Precise 的做法不是让所有 timestep 用同一个噪声强度,而是问:在当前 logSNR 位置,模型正常 ODE velocity 本来会推进多少 denoising?score term 又额外要求模型推进多少?如果二者比例在全程可控,sampler 就既能探索,又不会在 late denoising 破坏细节。

3.2 Exploration schedule:用 logSNR 控制额外 denoising burden

作者用 logSNR 度量 denoising progress。对一个 infinitesimal reverse step ,velocity、score、stochastic 三项对 logSNR 的一阶贡献为: score 与 stochastic noise 在 net first-order progress 上互相抵消,所以总体 denoising progress 仍由 velocity term 主导;但 score term 并不是“免费”的,因为在被 re-randomization 抵消前,它要求模型额外预测更强的 denoising direction。如果这个额外 burden 远大于 normal velocity progress,模型误差会被放大。

因此 Precise 固定比例: 并得到 exploration schedule: 这里 是一个全局强度参数。论文实验中 Precise 默认 ,接近 ,表示 score-induced progress 与 normal velocity progress 同量级;这比 constant 更符合 flow-matching path 的 logSNR geometry。

3.3 为什么现有 finite-step sampler 会偏

Euler / Flow-GRPO / Dance-GRPO:fresh noise feedback 被忽略

Euler-Maruyama 从 的一步写作: $ \mathbf{z}_{t’}^{\mathrm{Euler}}

\mathbf{z}_t+ \left(-\mathbf{u}t+\frac{1}{2}\varepsilon_t^2\nabla{\mathbf{z}}\log p_t(\mathbf{z}_t)\right)\Delta t +\varepsilon_t\sqrt{\Delta t},\mathbf{w}, \Delta t=t-t’t$ 时刻;在 deterministic ODE 里这通常只是普通 discretization error,但 stochastic step 注入的 fresh noise 会立刻改变 latent state,进而改变后续 velocity/score。Euler 忽略这个 feedback。

在 point-mass data toy case 中,Euler step 会变成: 因为 独立,最终噪声系数为 只要 就严格大于目标 ,即使模型 perfect、数据分布最简单,也会产生 excess noise。

CPS:coefficient-preserving 但 posterior uncertainty 被压缩

CPS transition 的一般形式为: 这个条件对 point-mass model 是必要的,但对真实非退化数据分布不充分。原因是 只保留 posterior mean,丢掉 的 residual uncertainty。局部例子里,CPS 的 covariance contraction 正好是: 这解释了为什么 CPS 在 double-ring toy distribution 上会持续偏向 inner ring:它不是噪声系数错了,而是 posterior mean 把多模态 clean latent ambiguity 压成了均值。

3.4 Precise transition:frozen clean-latent posterior mean 下的闭式解

Precise 的关键近似是:在一个 finite reverse step 内,clean-latent posterior mean 近似不变: 这个近似在 point-mass data 中 exact,在实际模型中是 local approximation:当 denoiser 对 clean latent 已经比较确定时,新噪声可以显著移动 noisy latent,但 clean-latent prediction 只会轻微变化。作者在 appendix 中用 posterior covariance 控制 posterior mean 的 Jacobian,形式化这个 intuition。

在该假设下,reverse SDE 可以写成线性 SDE: Theorem 给出 exact transition: 其中 代入 logSNR schedule 后,有: 结合 最终 transition 变成: 这条公式很像 CPS 的 coefficient-preserving 形式,但 来自 frozen-mean SDE 的 exact solution,而不是手工分配 。它保留了 总噪声尺度,同时让新噪声的影响以 SDE-consistent 的方式进入下一步。

Figure 2 解读:图中检验 frozen posterior-mean approximation 的局部稳定性。作者在 SD3.5-M 上沿 uniform logSNR grid 回放 deterministic denoising trajectory,并对相邻 logSNR anchors 做 forward-style renoising;曲线显示 clean-latent prediction 的变化相对受控,支持“一步内 近似冻结”的局部假设,但这也说明 Precise 的理论保证依赖 step size 与 denoiser 稳定性。

3.5 Paper-level pseudocode(无 released code,可复现公式级逻辑)

代码搜索未找到开源实现;下面伪代码只根据论文公式与实验描述重构 sampler 逻辑,不声称对应 Tencent-Hunyuan/Precise 的实际文件或实现细节。

import torch
 
 
def logsnr_exploration(t: torch.Tensor, eta: float = 1.5, eps: float = 1e-6):
    t = t.clamp(eps, 1.0 - eps)
    return eta * torch.sqrt(t / (1.0 - t))
import torch
 
 
def precise_transition(z_t, t, t_next, z0_pred, eta: float = 1.5):
    """One reverse step t -> t_next for flow-matching RL rollouts."""
    eps_hat = (z_t - (1.0 - t) * z0_pred) / t
    rho = (t_next * (1.0 - t) / (t * (1.0 - t_next))).pow((eta ** 2) / 2.0)
    noise = torch.randn_like(z_t)
    z_next = (
        (1.0 - t_next) * z0_pred
        + t_next * rho * eps_hat
        + t_next * torch.sqrt(torch.clamp(1.0 - rho ** 2, min=0.0)) * noise
    )
    return z_next
import torch
 
 
def rollout_with_precise(model, prompt_embeds, time_grid, eta: float = 1.5):
    z = torch.randn_like(prompt_embeds.latent_template)
    log_probs = []
    for t, t_next in zip(time_grid[:-1], time_grid[1:]):
        velocity = model(z, t, prompt_embeds)
        z0_pred = z - t * velocity  # from u_t = epsilon - z0 and z_t=(1-t)z0+t epsilon
        z_next = precise_transition(z, t, t_next, z0_pred, eta=eta)
        # RL implementation would record transition density / policy ratio here.
        log_probs.append(None)
        z = z_next
    return z, log_probs

3.6 Code search / paper mapping status

Code reference: 无可访问 released code;GitHub 搜索 "Precise" "SDE-Consistent""2605.23522""Tencent-Hunyuan/Precise" 均为 0 hits,https://api.github.com/repos/Tencent-Hunyuan/Precise 返回 404(2026-05-25 检查)。

Paper conceptPaper anchorSource file / class status
logSNR exploration scheduleEq. 未找到可访问 source;无法映射到 launch/config
SDE-consistent finite-step transitionTheorem exact transition + Eq. 未找到可访问 source;伪代码基于论文公式
GRPO rollout policyPreliminaries RL post-training, Flow-GRPO-style group trajectories未找到可访问 source;论文称基于 Flow-GRPO codebase
Training configsSD3.5-M / FLUX.2 protocols in Experiments未找到可访问 launch script;本文只记录论文报告值

论文公式与 released code 实现差异:未发现可访问 released code,无法核对 paper/code gap;当前 blocker 是论文声明的 GitHub repo 未公开或返回 404。

4. Experimental Setup (实验设置)

4.1 Tasks, prompts, rewards

  • Rule-based reward:GenEval,用 COCO categories 上的 synthetic object-centric prompt templates;论文未详细说明本实验使用的 GenEval prompt 总数。
  • Model-based rewards:训练 reward 包括 PickScore、CLIPScore、HPSv2.1;evaluation 额外报告 ImageReward、UnifiedReward v2、Aesthetics,用来覆盖 reward alignment、text-image matching 与 perceptual quality。
  • Prompt dataset:PickScore / CLIPScore / HPSv2.1 prompts 来自 Pick-a-Pic training captions,过滤为 unique captions 且至少 6 个 whitespace-separated words;held-out evaluation split 为 2,048 shuffled captions。训练 captions 的具体数量论文未详细说明。
  • Prompt splits:所有实验使用 Flow-GRPO codebase 分发的 splits。

4.2 Baselines and protocols

主要 baseline 是 Dance-GRPO、Flow-GRPO、CPS;主实验还把 DiffusionNFT 作为 30-NFE reference。所有方法用于 GRPO-based reward optimization of flow-matching models。

SD3.5-M main protocol

  • Backbone:Stable Diffusion 3.5 Medium, text-to-image flow-matching model。
  • Fine-tuning:LoRA ;使用 GRPO-Guard。
  • CFG / KL:training 和 evaluation 均无 classifier-free guidance;training 无 KL loss。
  • NFE:training 与 evaluation 使用相同 NFE。
  • Protocols:10-NFE / 3000 iterations;30-NFE / 1000 iterations;两者约 3 天。
  • Exploration scalar :Dance-GRPO 0.3,Flow-GRPO 0.7,CPS 0.7,Precise 1.5。
  • Hardware:所有 runs 使用 NVIDIA H20 GPUs,每张 H20 约 80GB GPU memory。
  • Reporting:final checkpoint;可用时报告 3 个 evaluation seeds 的 mean ± std。Dance-GRPO 在 30-NFE 中不稳定,论文报告 iteration 600 的最高 checkpoint。

FLUX.2 Klein transfer protocol

  • Backbone:FLUX.2 Klein 4B Base,
  • Fine-tuning:LoRA
  • Protocol:20-NFE / 1000 iterations。
  • Hardware / time NVIDIA H20 GPUs,约 1 天。

由于 released code 不可访问,上述 training-config numbers 均来自论文 Experiments section,而不是 actual launch script;无法验证是否存在代码默认值与论文报告值不一致的问题。

5. Experimental Results (实验结果)

5.1 SD3.5-M main results

ProtocolMethodPickScoreCLIPScoreHPSv2.1ImageRewardUnifiedReward v2Aesthetics
10-NFEDance-GRPO22.428±0.0070.942±0.0000.340±0.0001.350±0.0100.647±0.0006.620±0.007
10-NFEFlow-GRPO23.242±0.0031.018±0.0010.377±0.0001.627±0.0040.652±0.0006.530±0.003
10-NFECPS23.670±0.0021.026±0.0010.389±0.0001.607±0.0050.650±0.0006.667±0.003
10-NFEPrecise23.745±0.0041.038±0.0010.391±0.0001.615±0.0040.654±0.0006.652±0.003
30-NFEDance-GRPO22.237±0.0040.868±0.0010.330±0.0001.082±0.0080.652±0.0016.568±0.005
30-NFEFlow-GRPO23.227±0.0060.967±0.0010.374±0.0001.528±0.0020.656±0.0006.736±0.004
30-NFECPS23.288±0.0070.994±0.0010.379±0.0001.542±0.0010.652±0.0006.693±0.002
30-NFEPrecise23.421±0.0061.002±0.0010.379±0.0001.548±0.0040.658±0.0006.766±0.003
30-NFE ref.DiffusionNFT23.3491.0100.3631.5730.6606.494

Figure 3 解读:10-NFE 和 30-NFE training curves 都显示 Precise 在三个 in-domain training rewards 上更快上升,并且最终 PickScore / CLIPScore / HPSv2.1 基本领先。论文把 prior samplers 的 best value 当作 target,Precise 在 10-NFE 下少用 14.2%–22.1% iterations,在 30-NFE 下少用 13.1%–53.2% iterations;由于 sampler overhead 可忽略,这基本转化为 wall-clock saving。

结果解释:Dance-GRPO 的 constant noise schedule 被作者认为 early denoising 太保守、late denoising 太激进,导致 layout variation 少且细节破坏,甚至在 30-NFE 后期不稳定;Precise 与 Flow-GRPO 的差距主要来自去除 excess discretization noise;Precise 与 CPS 的较小差距则体现了 posterior uncertainty / covariance bias 的影响。

5.2 Task-level validation

MethodGenEvalPickScore
Flow-GRPO0.96924.283
CPS0.97224.594
Precise0.98024.668

这组 head-to-head 在 10-NFE / 3000-iteration protocol 下隔离一个 rule-based reward 和一个 model-based reward。Precise 同时在 GenEval 与 PickScore 上最高,说明它不是只对某一个 learned reward model 过拟合。

5.3 FLUX.2 Klein transfer

MethodPickScoreCLIPScoreHPSv2.1ImageRewardUnifiedReward v2Aesthetics
Flow-GRPO22.591±0.0060.905±0.0010.355±0.0001.266±0.0010.654±0.0006.676±0.005
CPS22.486±0.0050.898±0.0010.362±0.0001.334±0.0100.649±0.0016.513±0.000
Precise22.791±0.0100.930±0.0010.366±0.0001.410±0.0060.653±0.0006.677±0.007

Figure 4 解读:在 FLUX.2 Klein 20-NFE protocol 下,Precise 的 training curves 仍然上升更快、最终更高;最终 checkpoint 在 6 个 metrics 中 5 个第一,UnifiedReward v2 只比最佳 Flow-GRPO 低 0.001。以 prior samplers 的最佳值为 target,Precise 在 PickScore、CLIPScore、HPSv2.1 上少用 18.4%–46.3% iterations。

5.4 Ablations:两个组件都有效

AblationPickScoreCLIPScoreHPSv2.1ImageRewardUnifiedReward v2Aesthetics
Dance-GRPO22.4280.9420.3401.3500.6476.620
+ SDE-consistent discretization+0.826+0.043+0.037+0.151+0.003+0.008
+ logSNR exploration schedule+0.491+0.053+0.014+0.114+0.004+0.024

这张 ablation 说明 finite-step transition 和 exploration schedule 不是互相替代的 trick。先把 Dance-GRPO 的 Euler transition 换成 SDE-consistent discretization,所有指标都有明显增益;再把 constant schedule 换成 logSNR-derived schedule,又在所有指标上继续增加。

Figure 5 解读:左侧 NFE ablation 表明 Precise 在低步数预算下尤其有价值;右侧 sweep 使用 ,在 上性能稳定,极端小值缺探索、极端大值会带来 denoising instability,但下降是 graceful 的。

5.5 Appendix evidence:CPS bias 与 qualitative samples

Figure 6 解读:double-ring distribution 的目标 outer-ring mass 是 0.5,但 CPS 即使增加 NFE 仍偏向 inner ring。这支持作者关于 covariance contraction 的分析:只 preserving coefficients 不足以恢复多模态 posterior uncertainty。

Figure 7 解读:10-NFE qualitative examples 展示了 Precise 在低步数 RL rollout 下的视觉质量和 prompt alignment。结合主表可看出,低 NFE 下 sampler bias 更容易显性影响 reward optimization,因此 Precise 的 SDE-consistency 对 10-NFE 特别关键。

Figure 8 解读:30-NFE qualitative examples 对应更长 denoising trajectory。论文结果显示,即使 NFE 提高,Dance-GRPO 仍可能不稳定,CPS 仍可能有 marginal bias;Precise 的优势不是“少步数专用”,而是从 sampler transition 层面减少 stochastic rollout 的系统偏差。

5.6 Limitations and takeaways

作者明确的主要限制在于理论近似:Precise 的 closed-form transition 依赖 one-step frozen posterior-mean assumption。该假设在 point-mass 情况 exact,在真实 model 中由 posterior covariance / local stability 支撑,但如果 denoiser 的 clean-latent prediction 对 fresh noise 非常敏感,近似误差会变大。

实验层面的限制是代码仓库当前不可访问,外部读者无法从 released config 验证 learning rate、batch size、scheduler、policy-ratio 计算和 GRPO-Guard 细节;论文也未详细列出 Pick-a-Pic 过滤后训练 captions 的精确数量。尽管如此,论文的主要证据相当一致:Precise 在 SD3.5-M、FLUX.2 Klein、rule-based / model-based reward、NFE / ablation 上都显示了更稳的 reward optimization。