DiffusionNFT: Online Diffusion Reinforcement with Forward Process
Authors: Kaiwen Zheng*, Huayu Chen*, Haotian Ye, Haoxiang Wang, Qinsheng Zhang, Kai Jiang, Hang Su, Stefano Ermon, Jun Zhu, Ming-Yu Liu Affiliations: Tsinghua University, NVIDIA, Stanford University(* 同等贡献, Jun Zhu 为通讯作者) arXiv: 2509.16117 Project Page: research.nvidia.com/labs/dir/DiffusionNFT GitHub: NVlabs/DiffusionNFT Venue: ICLR 2026 Oral
1. Motivation(研究动机)
在大语言模型(LLM)的后训练中,在线强化学习(RL)已成为对齐与推理能力提升的核心手段。然而,将相似的成功复制到扩散模型上仍然面临根本性挑战:
-
前向不一致性(Forward Inconsistency):现有方法(如 FlowGRPO、DanceGRPO)完全聚焦于反向采样过程进行策略优化,破坏了扩散模型前向过程的遵从性。这带来模型退化为级联高斯分布的风险——学到的模型不再对应一个合法的前向扩散过程。
-
求解器限制(Solver Restriction):数据收集依赖一阶 SDE 采样器(如 Euler-Maruyama),无法充分利用高阶 ODE 求解器或默认的 Flow 模型求解器,限制了生成效率和质量。
-
CFG 集成复杂(Complicated CFG Integration):扩散模型大量依赖 Classifier-Free Guidance(CFG),这要求同时训练条件与无条件模型。现有 RL 实践通常将 CFG 纳入后训练,导致复杂的双模型优化方案,显著增加计算成本。
-
似然估计偏差(Likelihood Estimation Bias):基于策略梯度的方法依赖精确对数似然,而扩散模型的似然只能通过代价高昂的 ODE 数值积分或 SDE 的变分下界来近似,引入系统性估计偏差。
核心问题:能否在正向扩散过程(而非逆向过程)上进行扩散模型的强化学习?
2. Idea(核心思想)
DiffusionNFT 提出了一种全新的在线 RL 范式——Diffusion Negative-aware FineTuning(DiffusionNFT),其核心思想是:
- 在正向过程上做策略优化:通过 Flow Matching 目标直接优化前向扩散过程,而非逆向采样过程,天然保证前向一致性。
- 对比正负样本定义改进方向:将收集到的生成图像按奖励信号划分为”正样本”(高奖励)和”负样本”(低奖励),定义两者之间的隐式改进方向 ,并将该方向编码进 Flow Matching 目标中。
- 隐式参数化单模型同时拟合正负策略:通过一个隐式参数化技巧,单个模型 同时表达正策略 和负策略 ,无需训练两个独立模型。
- 无需似然估计、无需 CFG:整个训练只需干净图像(clean images),不需要存储完整采样轨迹,也不需要无条件模型,实现完全 CFG-free 的强化学习。
- 天然离策略(off-policy):训练策略 和数据收集策略 解耦,无需重要性采样,通过 EMA 软更新维持采样策略稳定性。
该方法本质上是 NLP 中 Negative-aware FineTuning(NFT)范式(Chen et al., 2025c)在扩散域的推广。
3. Method(方法)
3.1 整体框架
Figure 2 解读:图中对比了两种扩散模型强化学习范式。上方为逆向过程 RL(GRPO):将反向 SDE 过程离散化,存储完整采样轨迹 ,利用相邻步骤间的转移概率 和 构建 GRPO/PPO 目标,需要一阶 SDE 采样器且存储整个轨迹。下方为正向过程 RL(NFT,本文方法):使用任意 Black-box Solver 生成干净图像 ,然后对其加噪得到 (正向过程),只需存储最终干净图像。NFT 计算当前策略速度 与参考速度 ,配合奖励 ,通过 Flow Matching 目标进行优化。
3.2 问题设置与分布划分
在线 RL 设置:给定预训练扩散策略 和提示词数据集 。每次迭代采样 张图像 ,用标量奖励函数 评估每张图像(代表最优概率 ,Levine 2018)。
分布划分:按奖励将数据随机分为正样本集 和负样本集 ,对应两个隐式分布:
可以证明 ,即正策略优于旧策略优于负策略。
3.3 改进方向定理
Figure 3 解读:图示改进方向 的几何含义。空间中分布着高奖励区域 (蓝色)和低奖励区域 (红色)。当前策略速度场 指向混合方向,而最优策略 沿着 方向偏转,从负样本区域移向正样本区域。待优化的 需要逼近 。
定理 3.1(改进方向):考虑扩散模型三元组 ,其对应速度场之间的方向差正比于:
其中标量系数:
训练目标 定义为:
其中 是引导强度(guidance strength)的倒数超参数。
3.4 隐式参数化与训练目标
Figure 4 解读:DiffusionNFT 的完整训练架构图。左侧:以文本条件 (如”a cute dog”)为输入,使用旧策略 采样 张干净图像 ,并评估奖励 。中间:对图像加噪得到 ,送入模型同时计算两路隐式速度:上路(,正样本分支)计算隐式正策略速度 ,损失为 ;下路(,负样本分支)计算隐式负策略速度 ,损失为 。右侧:以奖励 和 加权求和,通过 操作梯度下降。整个过程只训练单一策略 ,无需单独学习 CFG 模型。
定理 3.2(策略优化):考虑训练目标:
其中隐式正策略和负策略速度定义为:
在无限数据和无限模型容量下,Eq.(5) 的最优解满足:
这说明训练目标的最优解精确对应改进方向 ,保证了策略改进的理论依据。
3.5 与 CFG 的关系
DiffusionNFT 将 CFG 解释为强化学习引导方向 的一种离线估计(Eq.(4)),其中无条件和有条件模型分别对应负策略和正策略信号。因此,DiffusionNFT 通过强化学习训练,将 CFG 的功能内化进单个条件模型,实现了真正的 CFG-free 生成:训练时不需要无条件模型,推理时也不需要 CFG。
3.6 核心设计选择
最优性奖励(Optimality Reward)
将连续奖励 转化为 的最优性概率:
其中 是归一化因子(如全局奖励标准差)。该转换借鉴 GRPO 中的奖励归一化实践。
采样策略软更新(Soft Update of Sampling Policy)
离策略性质允许将采样策略 与训练策略 解耦,通过指数移动平均(EMA)进行软更新:
其中 为迭代编号,(单奖励实验)。完全在策略()收敛快但不稳定,完全离策略()稳定但收敛极慢。
自适应损失权重(Adaptive Loss Weighting)
受扩散蒸馏方法 DMD(Yin et al., 2024)启发,将时间相关权重 替换为自归一化的 回归形式:
其中 为 stop-gradient 操作, 为对应的 预测器(rectified flow 下 )。
CFG-Free 优化
初始化时仅使用条件模型(不用 CFG),起点性能虽低,但通过强化学习迅速超越 CFG 基线,验证了 CFG 功能可被 RL 后训练习得。
3.7 完整算法
Algorithm 1: DiffusionNFT 完整训练流程
# 输入:预训练扩散策略 v_ref,原始奖励函数 r_raw,提示词数据集 {c}
# 初始化:数据收集策略 v_old <- v_ref,训练策略 v_theta <- v_ref,数据缓冲区 D = {}
for i in range(num_iterations):
# ---- Rollout Step: 数据收集 ----
for c in sampled_prompts:
# 1. 使用任意黑盒 Solver(如 DPM2-ODE)生成 K 张干净图像
x0_1_to_K = black_box_solver(v_old, c, K=24)
# 2. 计算原始奖励
r_raw_1_to_K = reward_fn(x0_1_to_K)
# 3. 奖励归一化为最优性概率 r in [0,1]
r_norm = r_raw - mean(r_raw)
r = 0.5 + 0.5 * clip(r_norm / Z_c, -1, 1)
# 4. 存入缓冲区
D.append((c, x0_1_to_K, r_1_to_K))
# ---- Gradient Step: 策略优化 ----
for (c, x0, r) in mini_batch(D):
# 正向加噪:从正向过程采样 (x_t, v)
t ~ Uniform(0, 1)
eps ~ N(0, I)
x_t = alpha_t * x0 + sigma_t * eps # 正向过程
v = dot_alpha_t * x0 + dot_sigma_t * eps # 目标速度
# 计算隐式正策略速度
v_theta_plus = (1 - beta) * v_old(x_t, c, t) + beta * v_theta(x_t, c, t)
# 计算隐式负策略速度
v_theta_minus = (1 + beta) * v_old(x_t, c, t) - beta * v_theta(x_t, c, t)
# 自适应权重(基于 x0 预测误差)
x0_pred = x_t - t * v_theta(x_t, c, t) # rectified flow 下的 x0 预测
w = |x0_pred - x0|^2 / sg(mean(|x0_pred - x0|))
# NFT 损失 (Eq. 5)
loss = mean(w * (r * ||v_theta_plus - v||^2 + (1-r) * ||v_theta_minus - v||^2))
# 梯度下降
theta <- theta - lambda * grad_theta(loss)
# ---- Online Update: 软更新采样策略 ----
eta_i = min(0.001 * i, 0.5)
theta_old <- eta_i * theta_old + (1 - eta_i) * theta
D <- {} # 清空缓冲区Algorithm 2: 奖励归一化与最优性概率转换
def compute_optimality_reward(r_raw_list, Z_c=None):
"""
将原始奖励转换为最优性概率 r in [0,1]
Args:
r_raw_list: shape [K] 每个 prompt 采样的 K 个原始奖励
Z_c: 归一化因子,默认为全局 std(来自 all-gather 跨进程)
"""
# 跨所有进程 all-gather 奖励(分布式)
r_raw_all = all_gather(r_raw_list)
# 计算均值(按 prompt 分组)
r_mean = r_raw_all.mean() # 或 per-prompt 均值
# 归一化
r_norm = r_raw_all - r_mean
# 计算 Z_c(全局 std)
if Z_c is None:
Z_c = r_raw_all.std() + 1e-4
# 转换为 [0,1]
r = 0.5 + 0.5 * clip(r_norm / Z_c, -1.0, 1.0)
return rAlgorithm 3: EMA 软更新
def soft_update_sampling_policy(theta_old, theta, global_step, eta_max=0.5):
"""
EMA 软更新采样策略参数
Args:
theta_old: 采样策略参数(LoRA adapter "old")
theta: 训练策略参数(LoRA adapter "default")
global_step: 当前全局训练步数
eta_max: EMA 衰减上限
"""
# 计算当前 eta
eta_i = min(0.001 * global_step, eta_max)
# EMA 更新(对每个 LoRA 参数)
for param_old, param_new in zip(theta_old, theta):
param_old.data = eta_i * param_old.data + (1 - eta_i) * param_new.data.clone()
return theta_old3.8 代码与论文对应关系
| 论文公式/概念 | 代码位置 | 代码实现 |
|---|---|---|
| Eq.(5) NFT 损失 | scripts/train_nft_sd3.py training loop | `loss = r * |
| 隐式正策略 | scripts/train_nft_sd3.py | positive_prediction = (1-beta)*old_pred + beta*forward_pred |
| 隐式负策略 | scripts/train_nft_sd3.py | negative_prediction = (1+beta)*old_pred - beta*forward_pred |
| 奖励归一化 Eq. | scripts/train_nft_sd3.py reward processing | r = 0.5 + 0.5 * clip(r_norm / Z_c, -1, 1) |
| EMA 软更新 | flow_grpo/ema.py + train loop | ema.step(params, global_step) 内部 decay = min(0.001*step, eta_max) |
| 自适应权重 | scripts/train_nft_sd3.py | `weight = |
| 正向加噪 | scripts/train_nft_sd3.py line 9 | x_t = alpha_t * x0 + sigma_t * eps |
| 目标速度 | scripts/train_nft_sd3.py line 9 | v = dot_alpha * x0 + dot_sigma * eps |
| LoRA 配置 | config/nft.py | lora_rank=32, lora_alpha=64 |
| 数据收集策略 (黑盒求解器) | config/nft.py | sample.solver = "dpm2", sample.num_steps = 40 |
4. Experimental Setup(实验设置)
基础模型与分辨率
- 基础模型:SD3.5-Medium(2.5B 参数,Esser et al., 2024), 分辨率
- 参数高效微调:LoRA,
- 学习率:
奖励模型
基于规则的奖励(Rule-Based):
- GenEval(Ghosh et al., 2023):评估组合图像生成(compositional generation)能力
- OCR:评估视觉文字渲染能力
基于模型的奖励(Model-Based):
- PickScore(Kirstain et al., 2023):人类图像偏好
- ClipScore(Hessel et al., 2021):图文对齐
- HPSv2.1(Wu et al., 2023):人类偏好分数
- Aesthetics(Schuhmann, 2022):图像美学质量
- ImageReward(Xu et al., 2023):综合人类偏好
- UnifiedReward(Wang et al., 2025):多模态统一奖励
提示词数据集
- GenEval/OCR:使用 FlowGRPO 的对应训练/测试集
- 其余奖励:在 Pick-a-Pic(Kirstain et al., 2023)上训练,在 DrawBench(Saharia et al., 2022)上评估
训练与评估配置
- 每轮 epoch:48 组(groups),每组大小
- 单奖励对比实验:采样步数 10(公平比较),,
- 多奖励联合训练:采样步数 40(最佳视觉质量),,(OCR)
- 评估:40 步一阶 ODE 求解器(DPM2)
- OCR 特殊配置: 以稳定训练
多奖励联合训练策略
三阶段课程式训练:
- Pick-a-Pic 数据集上训练 PickScore、ClipScore、HPSv2.1(800 轮 → 200 轮 → 200 轮)
- GenEval 数据集上叠加上述 3 个奖励 + GenEval(300 轮 → 200 轮)
- OCR 数据集上叠加所有奖励 + OCR(100 轮 → 100 轮)
5. Experimental Results(实验结果)
5.1 多奖励联合训练(CFG-Free vs SOTA)
Figure 1b 解读:雷达图展示 DiffusionNFT(CFG-free)在 8 个维度上与多个基线的全面对比。DiffusionNFT(红色)在 PickScore、ClipScore、HPSv2.1、GenEval、OCR 上均达到最优或次优,在 Aesthetic、ImageReward、UnifiedReward 上也显著超越 SD3.5-M(无 CFG)基线,整体轮廓面积最大。特别值得注意的是,DiffusionNFT(CFG-free)的雷达图面积超过了 SD3.5-M(w/ CFG)和 FLUX.1-Dev(12B 参数)。
Table 1 完整结果(评估分辨率 ,†为官方 checkpoint,‡为 ):
| 模型 | Iter | GenEval | OCR | PickScore | ClipScore | HPSv2.1 | Aesthetic | ImgRwd | UniRwd |
|---|---|---|---|---|---|---|---|---|---|
| SD-XL‡ | — | 0.55 | 0.14 | 22.42 | 0.287 | 0.280 | 5.60 | 0.76 | 2.93 |
| SD3.5-L† | — | 0.71 | 0.68 | 22.91 | 0.289 | 0.288 | 5.50 | 0.96 | 3.25 |
| FLUX.1-Dev | — | 0.66 | 0.59 | 22.84 | 0.295 | 0.274 | 5.71 | 0.96 | 3.27 |
| SD3.5-M (w/o CFG) | — | 0.24 | 0.12 | 20.51 | 0.237 | 0.204 | 5.13 | -0.58 | 2.02 |
| + CFG | — | 0.63 | 0.59 | 22.34 | 0.285 | 0.279 | 5.36 | 0.85 | 3.03 |
| + FlowGRPO† | >5k | 0.66 | 0.92 | 22.51 | 0.293 | 0.274 | 5.32 | 1.06 | 3.18 |
| 2k | 0.54 | 0.68 | 22.41 | 0.290 | 0.280 | 5.90 | 1.29 | 3.37 | |
| 4k | 0.54 | 0.68 | 23.50 | 0.280 | 0.316 | 5.90 | 1.29 | 3.37 | |
| + Ours | 1.7k | 0.94 | 0.91 | 23.80 | 0.293 | 0.331 | 6.01 | 1.49 | 3.49 |
DiffusionNFT 在 1.7k 迭代时即在 7/8 指标上取得最优,同时完全不使用 CFG。
5.2 单奖励头对头对比(效率优势)
Figure 1a 解读:GenEval 任务上 DiffusionNFT vs FlowGRPO 的 GPU 时间对比曲线。DiffusionNFT(橙色)在约 100 GPU 小时内即达到 0.98 的 GenEval 分数,而 FlowGRPO(蓝色)在超过 2500 GPU 小时才达到 0.95。两者效率差距约为 25×。DiffusionNFT 的起点(~0.24,无 CFG)明显低于 FlowGRPO 的起点(~0.63,有 CFG),但最终大幅超越。
Figure 6a 解读:OCR 任务上的头对头对比。DiffusionNFT 在约 100 GPU 小时内达到 0.96+ OCR 分数,FlowGRPO 需约 800 GPU 小时才达到相同水平,效率提升约 24×。
Figure 6b 解读:PickScore 任务上的头对头对比。DiffusionNFT 以约 8× 的效率优势达到相同 PickScore 水平。
Figure 6c 解读:HPSv2.1 任务上的头对头对比,DiffusionNFT 收敛更快,效率提升约 5×,且最终分数更高。
5.3 定性比较
Figure 5 解读:三个任务(GenEval 组合生成、OCR 文字渲染、DrawBench 通用图像生成)上的定性对比。第一行 SD3.5-M 基线经常无法遵循颜色/组合指令(如”blue pizza and yellow baseball glove”中颜色混淆);第二行 FlowGRPO(带 CFG)有所改善但仍有错误;第三行 DiffusionNFT(无 CFG)生成质量最佳,正确渲染颜色、文字和空间关系,视觉质量也更自然。
5.4 消融研究
Figure 7a 解读:不同扩散采样器对 GenEval 的影响。ODE 求解器(一阶/二阶)均显著优于 SDE 采样器,验证了 DiffusionNFT 可充分利用任意高质量求解器的优势(FlowGRPO 只能使用 SDE)。二阶 ODE 与一阶 ODE 在 GenEval 上效果相当。
Figure 8 解读:不同 EMA 软更新策略对 GenEval 的影响。完全在策略()收敛极快但训练崩溃;(接近离策略)收敛过慢;渐进式 和 均取得良好平衡,前者更稳定。
Figure 9a 解读:不同时间权重策略对 GenEval 的影响。逆策略(如 )导致训练崩溃; 和 表现尚可;自适应权重(Adaptive)始终优于所有手动调节方案,且无需人工设定 。
Figure 10 解读:引导强度 的影响。 过小导致收敛极慢; 过大引起训练不稳定; 在稳定性和收敛速度之间取得最佳平衡。实践中单奖励实验取 ,多奖励实验取 。
关键消融结论:
- 负损失不可缺少:去除负策略损失( 项),奖励几乎立即崩溃(与 LLM 中 RFT 仍有效的现象不同,扩散 RL 中负反馈至关重要)
- ODE 采样器更优:ODE 样本质量更高,有利于奖励评估和训练稳定性
- 自适应权重最优:无需人工调节时间权重,自适应方案一致优于启发式选择
- 渐进 EMA 效果最佳:兼顾早期快速收敛与后期训练稳定性