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

会议: NeurIPS 2026
arXiv: 2605.23522
OpenReview: https://openreview.net/forum?id=GDousLQnYE
代码: https://github.com/Tencent-Hunyuan/Precise
领域: 图像生成
关键词: 流匹配、强化学习后训练、随机采样器、SDE一致性、探索调度

一句话总结¶

针对流匹配模型在线强化学习后训练中采样器探索不稳定与离散化偏差痛点，提出 Precise 采样器，通过对数信噪比（logSNR）平衡探索调度与冻结干净隐变量后验均值的 SDE 一致闭式步进，消除多余离散化噪声并节省 13.1%–53.2% 训练耗时。

研究背景与动机¶

强化学习（RL）后训练已成为提升扩散与流匹配生成模型提示词对齐度（prompt alignment）和感知质量的关键范式。在将在线强化学习（如 Flow-GRPO 与 Dance-GRPO 等基于组相对策略优化的方法）拓展至连续流匹配生成器时，必须将原本确定性的逆时间常微分方程（ODE）轨迹转化为随机策略，通常通过引入反向随机微分方程（SDE）来实现。在此框架下，随机采样器直接决定了策略的动作概率分布、prompt 内部优势估计所需的探索多样性，以及模型做出速度和得分预测时的去噪轨迹动态。为了规避迭代生成带来的巨大计算瓶颈，在线 RL 的 rollout 步数通常被严格压缩至极低的函数评估次数（NFE，如 6 至 30 步），使得在低步数预算下实现高质量、忠实于去噪几何的随机探索成为强化学习扩展的核心前提。

然而，现有随机采样器在低 NFE 设定下暴露出了严重的探索与离散化矛盾。一方面，探索噪声调度的设定极难权衡：若注入随机扰动过小，同组轨迹缺乏奖励方差导致策略优化迟缓；若扰动过大，去噪轨迹失稳，破坏生成保真度。另一方面，粗粒度离散化使得采样过程与理想反向 SDE 产生严重偏离。例如，Flow-GRPO 采用的标准欧拉-马山（Euler-Maruyama）离散化在单步内冻结速度与得分，忽略了注入噪声对潜在隐变量的实时反馈，导致每一步都引入超出理论方差的多余噪声，使模型被迫外推至训练分布之外；而系数保持采样（CPS）虽通过保留名义信号与噪声系数缓解了点质量数据的过量噪声，却因忽略条件后验均值以外的残余不确定性而引入协方差收缩，导致边缘分布出现系统性偏差。

为了从根本上摆脱现有采样器的次优折中，需要将采样器设计解构为互相关联的两个维度：其一，设计与局部去噪几何自适应的时间变化探索调度，在模型能可靠纠偏的阶段增强探索，在敏感阶段控制扰动；其二，在有限步离散化时保持与连续 SDE 的严格一致性。核心 idea：通过对数信噪比（logSNR）一阶分解导出恒定得分-速度推进比的探索调度，并基于“冻结干净隐变量后验均值”的局部线性化假设推导出严格 SDE 一致的闭式转移转移方程，在完全消除多余离散化噪声的同时兼顾充足探索与生成稳定性。

方法详解¶

整体框架¶

Precise 针对流匹配模型的在线 RL 后训练设计，将随机采样器显式视为可学习强化学习策略的核心构成部分。其整体输入为先验标准高斯噪声隐变量 \(\mathbf{z}_1 \sim \mathcal{N}(\mathbf{0}, \mathbf{I})\) 以及文本提示条件，输出为对齐高质量人类偏好分布的图像隐变量 \(\mathbf{z}_0\)。整个流程贯穿反向时间调度 \(t \in [1, 0]\)，在极有限的离散步数（如 10 步或 20 步）内执行迭代逆向推理。

整个采样流程在每一步包含三个紧密耦合的阶段：首先，根据当前时刻 \(t\) 的对数信噪比局部几何，按闭式公式计算自适应噪声强度 \(\varepsilon_t\)；接着，基于当前噪声隐变量 \(\mathbf{z}_t\) 调用流匹配网络预测速度场，进而显式换算出当前时刻干净隐变量的后验均值估计 \(\hat{\mathbf{z}}_0(t)\) 与后验噪声估计 \(\hat{\boldsymbol{\epsilon}}(t)\)；最后，利用冻结后验均值假设下求得的精准 SDE 转移算子，解析地生成下一步隐变量 \(\mathbf{z}_{t'}\)。该设计杜绝了一阶欧拉截断带来的方差漂移，确保在线 RL 策略梯度的精确计算。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入: 初始高斯噪声 z_1 与 文本条件"] --> B["阶段1: logSNR 自适应探索调度<br/>按局部几何动态计算噪声强度 ε_t"]
    B --> C["阶段2: 冻结后验均值局部锚定<br/>网络预测速度并估计 z0_hat(t) 与 eps_hat(t)"]
    C --> D["阶段3: SDE 一致闭式转移<br/>解析求解线性化 SDE 更新到 z_t'"]
    D -->|未到达 t=0| B
    D -->|到达 t=0| E["输出: 对齐人类偏好的干净图像潜变量 z_0"]

关键设计¶

1. logSNR 自适应探索调度：平衡探索多样性与去噪稳定性

标准反向 SDE 在连续极限下由漂移项和扩散项构成：\(\mathrm{d}\mathbf{z}_t = \left(\mathbf{u}_t - \frac{1}{2}\varepsilon_t^2 \nabla_{\mathbf{z}}\log p_t(\mathbf{z}_t)\right)\mathrm{d}t + \varepsilon_t \mathrm{d}\mathbf{w}_t\)。为量化各分量在无穷小时间步 \(\Delta t\) 内对去噪过程的实质推进，作者引入对数信噪比 \(\lambda_t \triangleq \log(\alpha_t^2 / \sigma_t^2) = \log((1-t)^2 / t^2)\) 进行一阶渐近展开。通过数学推导证明，速度项贡献的一阶信噪比净增量为 \(\Delta\lambda_{\mathrm{vel}} = \frac{2\Delta t}{t(1-t)} + o(\Delta t)\)，而得分项的增益 \(\Delta\lambda_{\mathrm{sco}} = \frac{\varepsilon_t^2 \Delta t}{t^2}\) 恰好被随机噪声项引起的退化 \(\Delta\lambda_{\mathrm{sto}} = -\frac{\varepsilon_t^2 \Delta t}{t^2}\) 在一阶上完全抵消。尽管得分项与噪声项的期望抵消，但在单步离散化时，过大的得分项会迫使模型产生远超正常 ODE 速度场的额外去噪负担，从而放大神经网络固有的预测误差。为此，Precise 提出将得分推进量与基础速度推进量的比值 \(R(t) \triangleq \Delta\lambda_{\mathrm{sco}} / \Delta\lambda_{\mathrm{vel}}\) 约束为与时间无关的常数 \(\eta^2 / 2\)。通过联立解出全局自适应噪声强度：

\[\varepsilon_t = \eta \sqrt{\frac{t}{1-t}}\]

该调度在去噪初期（\(t \to 1\)，\(1-t\) 较小）注入适量噪声以激发构图层面的广泛探索，而在去噪晚期（\(t \to 0\)，\(t\) 较小）平滑衰减噪声注入，避免损坏精细高频纹理与造成训练坍塌。

2. 冻结后验均值局部锚定：捕捉注入噪声的动态反馈

传统欧拉采样器之所以会在小步数下产生多余方差，根源在于它在 \([t', t]\) 步长内将速度 \(\mathbf{u}_t\) 和得分 \(\nabla_{\mathbf{z}}\log p_t\) 均冻结在起始点 \(t\)，完全忽略了新注入高斯扰动对内部隐变量状态改变所带来的瞬时反馈。由于对干净隐变量的后验均值估计 \(\hat{\mathbf{z}}_0(t) \triangleq \mathbb{E}[\mathbf{z}_0 \mid \mathbf{z}_t]\) 本质上是一个条件去噪期望，其各向同性扰动敏感度受后验协方差 \(\operatorname{Cov}(\mathbf{z}_0 \mid \mathbf{z}_t)\) 严格控制。在去噪轨迹上，该后验均值相对于带噪状态的变化速度极其缓慢。Precise 提出在单步时间微元 \([t', t]\) 内冻结干净隐变量的后验均值，即设 \(\hat{\mathbf{z}}_0(s) \approx \hat{\mathbf{z}}_0(t)\)（\(s \in [t', t]\)）。这使得原本高度非线性的反向 SDE 在局部完全转化为时变参数的线性高斯 SDE：

\[\mathrm{d}\mathbf{z}_s = \left(\frac{\mathbf{z}_s - \hat{\mathbf{z}}_0(t)}{s} + \frac{\varepsilon_s^2}{2s^2}\bigl(\mathbf{z}_s - (1-s)\hat{\mathbf{z}}_0(t)\bigr)\right)\mathrm{d}s + \varepsilon_s \mathrm{d}\mathbf{w}_s\]

这一转化不仅精确反映了噪声注入对漂移向量的即时线性牵引，还为后续解析积分奠定了严格的数学基础。

3. SDE 一致闭式转移：杜绝多余方差并保持真实边缘分布

基于冻结后验均值构建的线性 SDE，定义反向时间变量 \(\tau = t - s\) 并引入标准化残差向量 \(\mathbf{r}_\tau \triangleq \frac{\mathbf{z}_{s(\tau)} - (1-s(\tau))\hat{\mathbf{z}}_0(t)}{s(\tau)}\)，可将其严格解构为布朗运动驱动的标准奥恩斯坦-乌伦贝克（Ornstein-Uhlenbeck）型过程。利用伊藤等距同构（Itô isometry）积分后，作者证明了该系统具有闭式精确高斯解。结合流匹配特定的自适应调度 \(\varepsilon_s\)，最终导出无多余离散化误差的转移公式：

\[\mathbf{z}_{t'} = (1-t')\hat{\mathbf{z}}_0(t) + t' \rho(t', t) \hat{\boldsymbol{\epsilon}}(t) + t' \sqrt{1 - \rho(t', t)^2} \,\mathbf{w}, \quad \mathbf{w} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})\]

其中时间衰减相关系数 \(\rho(t', t)\) 具有优美简洁的解析形式：

\[\rho(t', t) \triangleq \left(\frac{t'(1-t)}{t(1-t')}\right)^{\eta^2 / 2}\]

该转移形式与 CPS 在结构上同属系数保持型，但存在本质区别：CPS 依赖启发式无噪声玩具假定，直接用后验均值坍缩残余不确定性导致边缘协方差收缩；而 Precise 的衰减权重 \(\rho(t', t)\) 是反向扩散过程经伊藤积分得到的解析解，严格还原了连续 SDE 的边缘方差扩散，在消除欧拉过量噪声的同时杜绝了边缘分布畸变。

损失函数 / 训练策略¶

在策略梯度优化阶段，Precise 采用基于组相对策略优化（GRPO）的目标函数。针对同一 prompt \(c\)，采样生成一组 \(G\) 条轨迹 \(\{\tau_i\}_{i=1}^G\)，计算其标量奖励 \(R(\tau_i)\) 并进行组内归一化计算优势函数 \(A_i = (R(\tau_i) - \mu_R) / \sigma_R\)。由于 Precise 的一步转移概率分布严格服从高斯分布：

\[q(\mathbf{z}_{t'} \mid \mathbf{z}_t, c) = \mathcal{N}\left(\mathbf{z}_{t'}; \; (1-t')\hat{\mathbf{z}}_0(t) + t'\rho(t',t)\hat{\boldsymbol{\epsilon}}(t), \; (t')^2(1-\rho(t',t)^2)\mathbf{I}\right)\]

其转移对数似然 \(\log \pi_\theta(\mathbf{z}_{t'} \mid \mathbf{z}_t, c)\) 具有显式二次型解析式，可无损计算重要性采样权重比率 \(r_t(\theta) = \frac{\pi_\theta(\mathbf{z}_{t'} \mid \mathbf{z}_t, c)}{\pi_{\theta_{\text{old}}}(\mathbf{z}_{t'} \mid \mathbf{z}_t, c)}\)，并通过 GRPO 经典的截断代理损失执行高效的 LoRA 参数更新。

实验关键数据¶

主实验¶

主实验在 Stable Diffusion 3.5 Medium（SD3.5-M）基础模型上进行，输入分辨率为 \(512 \times 512\)，采用 8 卡 NVIDIA H20 集群测试 10-NFE/3000 次迭代与 30-NFE/1000 次迭代两种训练协议。下表展示了最终 checkpoint 在训练奖励与未参与训练的泛化指标（Held-out metrics）上的评估结果（三随机种子均值与标准差）。

协议 / 方法	PickScore ↑	CLIPScore ↑	HPSv2.1 ↑	ImageReward ↑	UnifiedReward v2 ↑	Aesthetics ↑
10-NFE 协议
Dance-GRPO	22.428±0.007	0.942±0.000	0.340±0.000	1.350±0.010	0.647±0.000	6.620±0.007
Flow-GRPO	23.242±0.003	1.018±0.001	0.377±0.000	1.627±0.004	0.652±0.000	6.530±0.003
CPS	23.670±0.002	1.026±0.001	0.389±0.000	1.607±0.005	0.650±0.000	6.667±0.003
Precise (本文)	23.745±0.004	1.038±0.001	0.391±0.000	1.615±0.004	0.654±0.000	6.652±0.003
30-NFE 协议
Dance-GRPO (iter 600)	22.237±0.004	0.868±0.001	0.330±0.000	1.082±0.008	0.652±0.001	6.568±0.005
Flow-GRPO	23.227±0.006	0.967±0.001	0.374±0.000	1.528±0.002	0.656±0.000	6.736±0.004
CPS	23.288±0.007	0.994±0.001	0.379±0.000	1.542±0.001	0.652±0.000	6.693±0.002
Precise (本文)	23.421±0.006	1.002±0.001	0.379±0.000	1.548±0.004	0.658±0.000	6.766±0.003
DiffusionNFT (25-NFE/3d)	23.349	1.010	0.363	1.573	0.660	6.494

在更现代的 FLUX.2 Klein 4B Base 模型（20-NFE/1000 次迭代）上，Precise 同样展现出极强的迁移性：PickScore 达到 22.791（显著高于 Flow-GRPO 的 22.591 与 CPS 的 22.486），ImageReward 达到 1.410（Flow-GRPO 为 1.266，CPS 为 1.334）。

消融实验¶

在 SD3.5-M 的 10-NFE/3000 次迭代协议下，作者系统验证了“SDE 一致离散化转移”与“logSNR 自适应探索调度”两个维度的阶梯式贡献。

配置	PickScore ↑	CLIPScore ↑	HPSv2.1 ↑	ImageReward ↑	UnifiedReward v2 ↑	Aesthetics ↑	说明
Dance-GRPO 基线	22.428	0.942	0.340	1.350	0.647	6.620	恒定调度 + 欧拉离散化
+ SDE 一致离散化转移	+0.826	+0.043	+0.037	+0.151	+0.003	+0.008	替换为闭式转移，消除多余噪声
+ logSNR 探索调度 (Precise)	+0.491	+0.053	+0.014	+0.114	+0.004	+0.024	引入动态 logSNR 调度，达到最优

关键发现¶

收敛速度与训练墙钟时间大幅节省：达到已有基线采样器的最高对齐得分目标，Precise 在 10-NFE 协议下减少了 14.2%–22.1% 的训练迭代，在 30-NFE 协议下更是减少了 13.1%–53.2% 的迭代步数。由于采样器本身的解析计算开销完全可忽略不计，迭代步数的缩减直接 1:1 转化为实际 GPU 训练时间的节省。
极端低步数下的抗退化鲁棒性：在 \(N \in \{6, 10, 20, 30\}\) 的 NFE 敏感性测试中，随着步数缩减至极端的 6-NFE，Flow-GRPO 因累积的过量离散化噪声导致生成崩塌、得分断崖式下跌，而 Precise 凭借严格的 SDE 一致性依然保持领先且平稳的 PickScore。
超参数宽容度：噪声强度参数在 \(\eta \in [1.0, 2.0]\) 宽幅区间内均表现出极高鲁棒性，仅在 \(\eta=0.5\)（探索不足导致策略改善慢）和 \(\eta=2.5\)（噪声过大引起轨迹失稳）时产生温和下降。

亮点与洞察¶

将采样器从推理超参提升为强化学习策略核心：指出在在线流匹配 RL 中，随机转移概率直接构成策略似然，从而将探索调度与离散化保真度确立为直接影响策略梯度的关键设计。
巧借“去噪任务稳定性”解决微元积分难题：没有强行数值求解非线性 SDE，而是敏锐抓住“去噪模型的干净后验均值在各向同性扰动下变动极缓”的统计物理特性，通过局部冻结构造出精妙的闭式解。
通用性与即插即用特性：Precise 独立于上层强化学习算法（无论是 GRPO、PPO 还是 DPO 包装），可直接无缝替换任意流匹配生成系统的采样引擎。

局限与展望¶

极端非凸流场下的局部平稳假设边界：虽然实验证明后验均值雅可比矩阵受局部协方差有效约束，但在多模态强冲突的生成分叉区，冻结均值假设可能带来微小的近似误差。
跨模态生成场景的拓展验证：当前验证主要集中于文本到图像（SD3.5 与 FLUX.2），未来可进一步拓展至文本到视频生成系统（如 HunyuanVideo 等长序列视频流匹配模型）的在线 RL 后训练中。
高阶积分器的结合潜力：可探索将冻结后验均值思想与二阶多步求解器结合，进一步压缩极端单步/两步推理下的探索对齐误差。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从 logSNR 理论分解与冻结后验均值线性化视角推导出严格一致的流匹配随机采样器。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 SD3.5 与 FLUX.2 两大前沿模型，提供多维度对齐评测、详尽的阶梯消融与参数鲁棒性曲线。
写作质量: ⭐⭐⭐⭐⭐ 理论推导严密流畅，从引理分解到定理求解一气呵成，图表与实验动机高度自洽。
价值: ⭐⭐⭐⭐⭐ 直接解决流匹配在线 RL 训练耗时长、小步数易坍塌的产业级瓶颈，开源即用价值极高。