通过奖励倾斜分布匹配强化少步生成器 Reinforcing Few-step Generators via Reward-Tilted Distribution Matching
两阶段框架RTDMD统一DMD与RL，在4步内实现少步文生图SOTA。
前置知识
扩散模型与流匹配
通过连续概率路径 $p_t$ 将数据分布 $p_0$ 与高斯先验 $p_1$ 连接的生成模型。流匹配采用线性路径 $x_t=(1-t)x_0+t\epsilon$，需要数十步迭代采样，部署成本高。
本文方法建立在流匹配之上并对速度场进行蒸馏，少步生成是核心目标，因此必须先理解原始多步模型的概率路径与采样过程。
分布匹配蒸馏 DMD
Yin等人提出的少步蒸馏方法，引入假分数模型 $s_\phi$ 逼近学生分布的分数，通过最小化 $D_{KL}(p_\theta\|p_\psi)$ 的反向KL让学生输出逼近教师分布。
RTDMD直接以DMD为基础并扩展为两阶段框架，理解DMD中假分数模型的训练目标是看懂Stage 1的前提。
GRPO 策略梯度
Shao等人提出的分组相对策略优化，对每个prompt采样 $N$ 条轨迹，用组内归一化优势 $A_i=(r_i-\bar r)/std$ 替代value网络，降低方差。
Stage 2的SubGRPO是GRPO在少步流生成器上的变体，必须先理解GRPO的分组归一化思想才能体会子集采样与共享噪声的方差缩减原理。
一致性模型正则化
Daras等人提出的方法，要求模型在相邻时间步的预测保持自洽，即 $\hat x(x_{t'},t')=\mathbb E[\hat x(\tilde x_{t''},t'')]$，用耦合约束降低估计方差。
AC-DMD的关键就是将一致性正则引入到假分数模型训练中，使其在generator分布持续漂移下仍能稳定跟踪。
研究动机
现有文生图扩散/流模型（如SD3、FLUX.2）虽然质量高，但需要50–100步采样，无法满足交互式内容创作、端侧部署、实时视觉系统等延迟敏感场景。少步蒸馏方法DMD系列可将步数压缩到4步，但教师分布本身并不对齐人类偏好，对高奖励与低奖励样本赋同等概率。此外，在少步生成中两个核心难题未解决：(1) 中间潜变量本身含噪，假分数模型不得不在带噪中间态上训练，且生成器每步更新都让目标分布漂移，冷启动蒸馏信号不可靠；(2) 奖励优化需要兼顾采样的混合性质——中间 $K-1$ 步为随机高斯转移，末步 $t_K=0$ 为确定性映射，纯GRPO或纯反传都会丢失一边的优化信号。
本文的目标是本文提出统一框架RTDMD，目标是在仅用4步推理的前提下，把分布匹配蒸馏与奖励引导的强化学习耦合到同一个KL目标中，既保留教师分布的生成先验又让少步生成器对齐人类偏好。具体地，希望在SD3-M、SD3.5-M与FLUX.2 4B上，4步生成的CLIPScore、Aesthetic、PickScore、HPSv2、ImageReward等指标全部达到新的SOTA，并能在多数指标上反超50步FLUX.2 9B教师。
与已有工作不同的是，此前两类工作各走极端：纯蒸馏派（DMD/DMD2/TDM）只拟合教师分布不考虑奖励；纯RL派（DMDR/RDM/TDM-R1/GDMD）只对随机步做GRPO或只对末步做反传，二者缺乏统一目标。本文独特切入点是先把KL到奖励倾斜教师分布的梯度严格分解为分布匹配项与奖励最大化项，再用AC-DMD稳定冷启动、混合策略梯度（SubGRPO+末步反传）联合优化奖励，使两个原本割裂的目标首次在同一目标函数下被一起优化，并提供方差缩减与冷启动稳定的工程方案。
核心方法
RTDMD把少步流生成器训练拆成冷启动与强化两阶段，建在KL目标 $\min_\theta D_{KL}(p_\theta\|\tilde p_\psi)$ 上。先定义奖励倾斜分布 $\tilde p_\psi(x)\propto p_\psi(x)\exp(\beta r(x))$，由 $\log\tilde p_\psi=\log p_\psi+\beta r-\log Z$ 把KL梯度拆为分布匹配项与 $-\beta\mathbb E_{p_\theta}[r]$ 奖励最大化项。Stage 1 用CPS（$\eta=0.9$）的AC-DMD在带噪中间态做子区间分布匹配并加一致性正则；Stage 2 对奖励项做混合策略梯度，对 $K-1$ 个随机中间步用SubGRPO、对末步用路径式反传。最终梯度为 $\nabla_\theta\mathcal L_{total}=\nabla_\theta\mathcal L_{AC-DMD}-\beta\nabla_\theta\mathcal L_{SubGRPO}+\nabla_\theta\mathcal L_{det}$。
与已有方法相比，本文的本质区别有三点：(1) 在KL层面首次把奖励倾斜统一进DMD的目标推导，使RL不再是外挂后处理；(2) 针对带噪中间态的假分数训练提出一致性正则，解决fake score跟不上漂移generator分布的难题；(3) 引入SubGRPO让非选中步共享噪声，把GRPO的方差来源从全部 $K-1$ 步压缩到 $M$ 步，并用混合策略梯度把此前被忽视的确定性末步纳入优化，整体在统一目标下联合处理分布匹配与奖励最大化。
方法步骤详情
Stage 1 AC-DMD：(a) CPS调度 $K=4$、$\eta=0.9$，每步注入 $\sigma_k=t_k\sin(\eta\pi/2)$ 高斯噪声；(b) 在 $[t_k,1]$ 重噪声化 $\hat x_{t_k}$ 为 $x_t^{(k)}=\alpha_k(t)\hat x_{t_k}+\sigma_k(t)\epsilon$，最小化 $\mathcal L_{gen}^{(k)}=\mathbb E_t\lambda_k(t) D_{KL}(p_{\theta,t}^{(k)}\|p_{\psi,t})$；(c) 假分数 $s_\phi$ 同子区间DSM训练并加一致性正则 $\gamma=0.01$。Stage 2：(d) 对prompt采 $N$ 条轨迹算组内优势 $A_i$；(e) SubGRPO采子集 $|S|=M=2$，未选中步共享噪声；(f) 末步直接反传奖励 $\nabla_\theta\mathcal L_{det}=\frac1N\sum_i \nabla_{\hat x_0}r^\top \partial_\theta G_\theta$。
技术新颖性
技术新颖性体现在四个层面：第一，理论贡献上给出 $D_{KL}(p_\theta\|\tilde p_\psi)$ 自然分解为分布匹配+奖励最大化的严格推导，让两阶段训练不再是经验拼凑；第二，方法贡献上首次把一致性正则移植到DMD假分数训练中，针对中间带噪latent设计新的 $\mathcal L_{cons}^{(k)}$；第三，算法贡献上提出SubGRPO的共享噪声方差缩减机制，相比MixGRPO更明确地把方差归因到选中的 $M$ 步；第四，工程贡献上首次把整套两阶段框架同时在SD3-M、SD3.5-M与FLUX.2 4B三种不同规模/架构的模型上验证并刷新SOTA。
实验结果
实验在SD3-M $1024^2$、SD3.5-M与FLUX.2 4B上对比GDMD/DMDR/RDM/TDM-R1/Hyper-SD等SOTA。SD3-M 4 NFE上RTDMD全部5项指标第一：CLIPScore 0.3161（vs SD3-M w/ CFG 100 NFE 0.2936），Aesthetic 5.9642（比100步CFG教师高+0.39），PickScore 22.86，HPSv2 0.3211，ImageReward 1.3024（比100步教师+0.23，用25×更少NFE）。FLUX.2 4B上拿下9项中7项SOTA：ImageReward 1.3712 vs 50步4B 0.8538与9B 1.0021，HPSv3 15.58 vs 9B 13.30，4B模型在多数指标上反超9B 50步教师。消融显示 $\gamma=0.01$ 比A-DMD PickScore再涨+1.56；SubGRPO $(M=2)$ 相对朴素GRPO PickScore由23.31升至23.45；加末步反传后进一步升至23.57。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| SD3-M 文生图 1024² 4 NFE | CLIPScore | 0.3161 | SD3-M w/ CFG 100 NFE 0.2936 / RDM 4 NFE 0.2936 | +0.0225（vs SD3-M w/CFG），+0.0225（vs RDM） |
| SD3-M 文生图 1024² 4 NFE | Aesthetic Score | 5.9642 | SD3-M w/ CFG 100 NFE 5.5711 / RDM 5.8769 | +0.3931（vs 100步CFG教师），+0.0873（vs RDM） |
| SD3-M 文生图 1024² 4 NFE | PickScore | 22.8593 | SD3-M w/ CFG 100 NFE 22.3236 / RDM 22.5783 | +0.5357（vs 100步CFG教师），+0.2810（vs RDM） |
| SD3-M 文生图 1024² 4 NFE | ImageReward | 1.3024 | SD3-M w/ CFG 100 NFE 1.0759 / GDMD 1.2702 | +0.2265（vs 100步CFG教师），+0.0322（vs GDMD） |
| FLUX.2 4B 4 NFE | ImageReward | 1.3712 | FLUX.2 4B 50 NFE 0.8538 / FLUX.2 9B 50 NFE 1.0021 | +0.5174（vs FLUX.2 4B 50步），+0.3691（vs FLUX.2 9B 50步） |
| FLUX.2 4B 4 NFE | HPSv3 | 15.5772 | FLUX.2 9B 50 NFE 13.2955 / Z-Image-Turbo 6B 12.9136 | +2.2817（vs FLUX.2 9B 50步） |
局限与改进
作者在结论中明确指出局限：本文聚焦静态文生图任务，没有把两阶段框架推广到视频合成与图像编辑。在视频任务中时间一致性、跨帧奖励设计与长轨迹信用分配会带来新难题；在图像编辑任务中reward既要保证编辑指令遵循又要保持原图身份，对reward-tilted分布的偏置方向也更敏感。此外观察到虽然RTDMD在大多数指标上反超教师，但在GenEval2（0.2755 vs FLUX.2 9B 0.3570）和OCR（0.6858 vs 0.7566）上仍落后，作者承认这部分差距主要来自基座模型能力差异，方法本身的相对提升幅度仍可观。另一个隐含局限是 $\eta=0.9$、$\gamma=0.01$、$M=2$ 等超参都是经验选择，对不同base model的迁移性尚未系统研究。
独立分析的弱点
独立分析有三处可改进：(1) 假分数模型的一致性正则用单步近似的reverse transition，对极小子区间 $[t',t'']$ 的方差缩减有效，但当 $K=2$ 时子区间几乎没有重叠，正则项退化为自洽约束，此时SubGRPO与一致性正则的贡献会重叠，建议引入多跳transition提高利用率；(2) 奖励项与分布匹配项共享 $\beta$ 单标量，无法对不同prompt自适应，prompt难度差异大时会出现简单prompt过优化、困难prompt欠优化，可考虑prompt条件化的 $\beta$ 或双层优化；(3) 实验只对比了4 NFE，缺少2步与8步的步数敏感度曲线，论文中虽讨论了 $\eta$ 影响但未量化步数-质量的帕累托前沿，建议补充2/4/8 NFE的横向对比。
未来方向
作者明确指出未来要扩展到视频合成与图像编辑；基于成果还可延伸三个方向：(a) 把奖励倾斜KL分解推广到非可微奖励或人类反馈奖励（本文奖励需可微以做末步反传）；(b) SubGRPO的子集采样策略目前是均匀随机，可以用学习得到的credit assignment给更关键步更高被选中概率；(c) AC-DMD的假分数一致性正则本质上是一种自监督约束，可考虑与rejection sampling或人类偏好数据结合，让fake score同时逼近人类判断的“好分布”而非纯学生分布。
复现评估
作者声明代码与模型开源在 https://github.com/Harahan/RTDMD，但截至论文撰写尚未公开具体权重。训练数据为t2i-2M、ShareGPT-4o-Image子集与DrawBench评测集，均为公开数据集；超参与训练曲线在App. J详细给出。算力需求较高：8或16块NVIDIA H20 GPU，SD3-M两阶段共约 $1.5\times10^3+1\times10^3$ 次迭代，每迭代需采48组、组大小24的轨迹（FLUX.2 4B为64组），单次实验估计在数十到上百H20·小时量级。复现难点主要在于LoRA $r=64,\alpha=32$ 与CPS $\eta=0.9$ 的组合、fake score与generator的交替频率、以及多奖励加权 $\beta$ 的选择，整体复现难度偏高但有完整算法伪代码支撑。
论文图表
