arXivDaily arXiv每日学术速递 周一至周五更新

PAPERDAILY REPORTS

阿里&上交提出REST:少步数图像生成追平40步教师,额外训练成本低于25%

arXiv 2608.09226 cs.AI · cs.CV

阿里淘天集团与上海交通大学提出REST,将图像生成的奖励对齐和少步蒸馏放进同一训练阶段。它直接复用扩散强化学习已经生成的带奖励轨迹,不再单独采样蒸馏数据。

实验中,REST以8步、无CFG的推理追平或超过40步RL教师,整体额外训练成本低于纯RL的25%。在DrawBench上,它的PickScore达23.96,比RTDMD高0.82,训练迭代数只有后者的五分之一。

RL轨迹本身就是一套带分数的蒸馏教材

常见流程先用强化学习提高人类偏好、文字渲染或组合遵循等奖励,再另起一轮将长步数模型压缩为少步模型。两阶段训练增加成本,蒸馏还可能把RL获得的偏好收益再折损掉。

REST指出,RL教师的每次采样已经产生有限步扩散轨迹,每条轨迹还有奖励和优势值。这些中间状态可以直接成为学生的分段监督,无需再让教师额外生一批图。

40步RL生成与REST少步生成效果对比

REST以更少去噪步数复现构图和视觉文字,并保留偏好对齐收益。

好轨迹多学,低奖励轨迹轻微反向推开

如果所有轨迹都用相同权重模仿,学生会把教师的低奖励行为一起学走。REST的优势调制蒸馏将优势值转成带符号权重:强化高奖励样本的拟合,对低奖励样本则产生较弱的排斥。

REST的RL教师、解耦学生和优势调制蒸馏架构

REST不改动原有RL教师优化,学生从教师正在变化的轨迹中分段学习。

学生和教师保持解耦,因此REST可以挂到不同扩散RL算法上,不必改写教师的奖励优化。方法也不需要对抗训练或单独蒸馏数据集。

文字渲染、组合生成和偏好分数同时检验

团队在组合生成、视觉文字和人类偏好三类任务上评估。DrawBench中,REST的美学分6.30、ImageReward 1.35、HPSv2 0.333和PickScore 23.96,多项指标高于同表的RTDMD版本。消融实验中,去掉优势调制后PickScore下降,表明只同时做RL和蒸馏还不够。

REST在图像文字渲染上的少步生成结果

REST在无CFG少步推理中处理招牌、广告牌和其他视觉文字。

这些成绩来自论文选定的基准、奖励模型和训练配置,不等同于所有生成题材都能在更少步数下保持相同质量。

训练中的教师不是一个固定检查点。RL分支持续根据奖励更新,学生每个阶段看到的也是教师当下最新的轨迹。这与先训完RL再锁定教师的串行流程不同:蒸馏不必等对齐完成,对齐过程生成的每批中间状态都可被使用。

优势调制中的“带符号”也有明确边界。高优势轨迹提高模仿权重,低优势轨迹只被轻微排斥,避免负权重过大让训练不稳定。论文使用平移参数控制从正向模仿到负向排斥的分界,并比较了无平移、更大平移和不使用EMA的结果。

与RTDMD对比时,REST只训练500次迭代,对比方案为1500加1000次。除PickScore外,REST的美学分为6.30、ImageReward为1.35、HPSv2为0.333,表明少步学生的收益不是只出现在单一奖励指标。它同时使用8步NFE,因此“五分之一迭代”和“少步推理”是两个不同层面的效率改善。

下一步接入更多扩散RL训练链

对训练工程而言,解耦学生还意味着教师的RL优化不需要为蒸馏目标让路。团队可以保留既有奖励、采样器和更新规则,只将轨迹及其优势值交给学生分支。这种接口比将教师、学生和判别器紧耦合在同一对抗目标中更容易迁移到现有RL代码库。

REST的实用点是将原本会被丢弃的RL中间轨迹变成蒸馏监督,学生随教师同步演化。这为高成本偏好对齐与少步推理的合并提供了更直接的接口。

REST在人类偏好提示上的生成对比

偏好对齐对比展示不同方法对提示语构图和属性的遵循。

后续可以考察更大尺寸教师、更丰富奖励组合和不同学生步数,同时测量训练成本、端到端延迟和人工质量评估,确定联合训练在产品化图像模型中的实际收益。

参考资料

https://arxiv.org/abs/2608.09226