Aligning One-Step Generative Models with Reward-Weighted Transport Distillation
Austin Wang, Ziheng Cheng, Lexing Ying
cs.LG, cs.CV
2026-09-25
字节团队提出 RWTD,只用采样和标量奖励给一步生成器做后训练,把 SANA Sprint 的 GenEval 从 0.73 升到 0.80,跨奖励泛化优于梯度法。
扩散和流模型画质好,代价是几十次串行网络前向;蒸馏出的一步生成器(SANA Sprint、SDXL Turbo)把推理压到一次前向,速度解决了,后训练没解决。RLHF 那套搬过来会卡在两处:一步生成器是隐式分布,没有可解析计算的似然,也没有可供回传的去噪轨迹;奖励又常常不可微,比如另一个模型打的分。现有的一步模型对齐方法要么要求奖励梯度(FAV、DI++),要么绑死扩散动力学,要么构造特殊的偏好场(DrPO)。缺的是一个只依赖「采样 + 打分」的通用后训练接口,RWTD 补的就是这个位置。
RWTD(Reward-Weighted Transport Distillation)每轮训练做三件事:
反向传播只更新生成器,梯度不经过奖励、权重和传输耦合。奖励从头到尾只输出标量,黑盒即可。理论侧给出混合定点的显式形式:ρ=1 退回常规倾斜参考,0<ρ<1 时多出一个随奖励放大的有理因子,ρ 越小越向高奖励模式集中。
主实验是 GenEval(考察物体位置、计数、颜色属性绑定的组合正确性 benchmark):
| 方法 | NFEs | GenEval Overall |
| SANA Sprint 1.6B(底模) | 1 | 0.73 |
| DrPO | 1 | 0.75 |
| FAV | 1 | 0.73 |
| RWTD | 1 | 0.80 |
| Playground v3(多步) | – | 0.76 |
| FLUX.1 Dev | 50 | 0.66 |
一步模型反超了同规模多步模型,也高过 SD3.5-L(0.71)和 4 步的 SANA Sprint(0.75)。SDXL Turbo 上 RWTD 从 0.55 提到 0.61,高于 DrPO 的 0.58 和偏好调优过的 LAIR-SDXL 的 0.59(后者要 50 步)。
偏好对齐一组更能看出分化:以 HPSv2 为奖励训练,在 PartiPrompts 上评估。DRaFT 把域内 HPSv2 拉到 37.02,代价是 CLIP 从 0.2753 掉到 0.2573,域外 GenEval 从 0.73 跌到 0.62,教科书式的奖励过优化。RWTD 的 HPSv2 只有 32.45,但它是唯一五项偏好指标加 GenEval 全部优于底模的方法(GenEval 0.75,FAV 和 DrPO 都掉到 0.72)。定性对比里梯度法画风明显风格化,RWTD 保住了写实感。
消融确认了两个设计。混合目标(ρ=0.15)比纯 off-policy 目标(ρ=1)收敛更快,GenEval 到 0.80 对停在 0.77 附近;去掉传输几何、独立采样目标粒子的 RWR 变体,Aesthetics 略升(6.878 对 6.751)但 CLIP 和 PickScore 变差,说明收益来自 OT 的样本级配对,不是重心平均。
一步生成器的后训练此前是个两难:黑盒奖励碰不了,梯度法一碰就过优化。RWTD 把接口降到能采样、能打分就够,任何标量信号都可用,包括不可微评测器和人工评分。对混合定点的刻画给了 ρ 一个可解释的读法:从倾斜参考到纯 on-policy 锐化的连续旋钮。对做图像生成的人,这是一条抬升一步模型组合能力又不牺牲写实性的路径,代码和 checkpoint 已开源。性质上属于对齐层的渐进改进,骨架还是原来的一步生成器,不是新的生成范式。
论文自己承认,理论是种群层面的理想化分析,实际训练用有限 batch、熵正则 OT 和重心投影,定理提供直觉而非保证。纯 on-policy 端点(ρ=0)没有唯一不动点,实验里也确认它大幅损失多样性,ρ 成了必须调的超参(GenEval 上取 0.15)。域内奖励的优化深度不及梯度法,HPSv2 32.45 对 DRaFT 的 37.02,RWTD 赢在平衡,要极限拉分时梯度法仍更强。实验只覆盖 SDXL Turbo 和 SANA Sprint 1.6B 两个骨架,最大 8 卡规模,没有人类评测,计算成本讨论放在附录。能否迁移到 FLUX 级别的更大骨架,论文没有验证。