FewStepRL: Fine-tuning Few-step Diffusion Models with Reinforcement Learning
Haitao Lin, Minsi Ren, Yang Li, Xingtong Ge, Zheyuan Hu, Peiyan Hu, Tailin Wu
cs.LG, cs.CV
2026-11-06
五种扩散强化学习目标被改到少步平均速度模型上，混合奖励的留出 GenEval2 里适配比硬套原公式最高多 0.013。
扩散模型做强化学习后训练，每一轮都要按当前策略采样。多步采样器常见 32 或 64 次网络评估，rollout 占掉大部分时间和显存，视频更甚。少步模型把生成收到 1 到 4 次前向，学的是一段区间上的平均速度。MeanFlow 直接拟合这个平均，DMD 和 AnyFlow 把教师的传输蒸馏成一次能跳很远的 flow map。平均速度是轨迹的割线，轨迹弯的地方，它和瞬时速度就分开了。
FlowGRPO、AWM、DiffusionNFT、DiffusionOPSD、ReFL 按瞬时速度或多步转移来写。平均速度填进去之后，对哪一次求梯度、参考策略在哪一端评估、哪些 (t, r) 被训练，都没定下来。Figure 1 从同一个少步检查点比较，三族的留出 GenEval2 都是硬套更低。
线性加噪下，速度是一次仿射搬运：水平 t 的状态减去 (t−r) 乘以速度，落到水平 r。速度的均方误差于是等于搬运后两点的距离乘一个时间系数。密度、速度、奖励三种目标就可以放在同一个干净样本上比：预测该靠近谁、躲开谁。默认看干净端，转移和 flow map 目标另标 r。
梯度是预测对参数的雅可比乘两股力。自回归拉向锚点，锚通常是旧策略在同一状态的预测，奖励因此推不远。奖励增广沿奖励样本和基线锚的连线走，符号看组内优势：高于组均值就吸引，低于组均值就推开。原始分高、但跑输同组的样本，仍会被推开。到这一步，换算法就是换锚和权重。AWM 的自回归权重是 0，另用速度惩罚贴住旧策略。DiffusionNFT 的回归锚和基线锚重合，奖励系数按优势截断。
少步网络输出平均速度 u(x, t, r)。u 形式只在 r=0 跳到干净端，采样要用的其他 r 没有梯度。v 形式先恢复瞬时速度。MeanFlow 在旧策略上做有限差分，估计平均速度随 t 的变化，修正不回传，每次查询多两次前向。DMD 用辅助网络 vφ 以流匹配拟合旧策略的 rollout，每个外层迭代 5 步，vφ 只当基线锚。带梯度的预测必须来自学生，否则奖励进不了参数。
外层循环是固定的：旧策略采样打分，按需拟合 vφ，内层查询累加损失，更新一次，再用 EMA 刷新旧策略。换锚点就得到 FewStepGRPO、AWM、DiffusionNFT、DiffusionOPSD 和 ReFL 的少步版。OPSD 要求奖励对干净样本可微，沿归一化奖励梯度造出上升点和下降点再拟合。ReFL 只留一步上升、关掉自回归，混合权重恒为 1，奖励打在旧策略的预测上。
图像从同一个 AnyFlow 蒸馏的 SD3.5-Medium 出发。Figure 1 画混合奖励下的留出 GenEval2（Soft-TIFA），训到 400 步。虚线是平均速度硬套原目标，实线是改写，标注为改写减去硬套。
| 方法 | 指标 | 结果 |
| FewStepGRPO 对 FlowGRPO | 留出 GenEval2，最佳与第 400 步 | +0.005 |
| MeanFlowAWM 对 AWM | 同上 | +0.003 |
| MeanFlowNFT 对 DiffusionNFT | 最佳检查点 | +0.005 |
| MeanFlowNFT 对 DiffusionNFT | 第 400 步 | +0.013 |
前两族停在千分位，NFT 到终点拉开到 0.013。三组都是改写更高。
奖励可微且梯度算得起时，奖励梯度方法最有效，单独 ReFL 已是强基线。不可微时用组优势，图像上另有 OCR 奖励；DiffusionNFT 风格往往较少迭代就把奖励做高，单次有时更贵。DMD 变体图像上每次迭代更快，视频上更慢，也更常带来留出退化。视频底座是 AnyFlow 蒸馏的 Wan2.1-T2V-1.3B。推理步加到训练的 2 到 4 倍，总体指标没有系统性收益，个别样本肉眼更好。
少步模型省下的采样，会在几十步 rollout 的后训练里还回去。五种目标共用采样、打分和评测，新变体主要改锚点。奖励可微、反传便宜，就用奖励梯度，ReFL 不必上成更重的 OPSD。分数不可微，就用 NFT 的组优势，并检查 DMD 有没有用速度换来留出退化。推理时把步数加密，不要默认总分会涨。
0.013 的 GenEval2 说明硬套有损，也说明改写是修补。少步学生已经在手、后训练又被 rollout 卡住，这套拆法才用得上。
Figure 1 只覆盖混合奖励和同一个蒸馏检查点。+0.003 到 +0.013 说明硬套会吃亏，撑不起「所有少步后训练都该换这套框架」。实线是 v 形式、虚线是 u 形式，多出来的分有多少来自整段 (t, r) 都有梯度，有多少来自奖励改写，没有拆开。
ReFL 在框架里是脱离后的代理：奖励导数取在旧策略预测上，再平移到锚点。和直接对当前预测反传共线，要 EMA 衰减为 0，数值相等还要系数对齐。
自回归被写成防奖励投机的手段，文中没有单独的投机测试。MeanFlow 每次查询多两次前向，DMD 多养一个跟着 rollout 的网络。少步省下的成本被这两项吃掉多少，只有图像更快、视频更慢，没有统一墙钟。实验停在 AnyFlow 蒸馏的 SD3.5-Medium 和 Wan2.1-T2V-1.3B，一致性模型和对抗蒸馏没有覆盖。