Few-step generation 通常靠 distillation,把几十步的 denoising trajectory 压缩到极少几步。它能换来速度,却也把 student 的目标锁成“复现 teacher”:teacher 的局限自然成了 student 的上限。单独加入 terminal reward 又很难告诉压缩后的每一步该怎么改。 RATS 把两者接起来:用 sigma-aligned multi-horizon matching 提供中间轨迹 guidance,再用 reward-aware gate 比较 teacher 与 student 的相对 reward。Teacher 更好时就多跟一点;student 追平或超过时便放松约束。最终 inference 仍只运行 few-step student,没有额外测试开销。
1. 两种常见方案,各自卡在一端
Static distillation 给出了清晰的中间监督,但 student 只能逼近既定 teacher。Reward-only optimization 允许 student 朝偏好目标继续提升,却只在完整 rollout 结束后看到分数;当 sampling 被压到极少 step,每一步承担的生成责任更重,terminal reward 的 credit 更难准确传回去。 RATS 处理的不是单纯的加速问题,而是“高效生成如何继续做 preference alignment”。
2. Multi-horizon matching:先把两条轨迹对齐
Teacher 和 student 的 step 数不同,不能按序号直接比较。RATS 按 noise level 对齐若干关键 horizon,让 few-step student 在相近的 denoising 阶段参考 multi-step teacher 的 latent prediction。 这些中间锚点比单个终点更有信息,也比强行匹配整条轨迹更贴合压缩后的采样日程。Teacher 采用 EMA 更新,只在训练阶段参与。
3. Reward-aware gate:决定此刻该听 teacher 多少
轨迹匹配并非固定强度。RATS 比较同一 prompt 和 noise 下 teacher、student 的 reward:teacher 得分更高时,gate 增强 trajectory shaping;student 已经追平或超过时,gate 自动减弱 teacher constraint,让 reward objective 继续推动 student。 这种机制保留了 distillation 的稳定先验,也避免 teacher guidance 在训练后期变成阻力。
4. Key Insight:teacher 更适合作为条件约束
传统 distillation 默认 teacher 始终正确,RATS 则把“是否值得模仿”变成一个随训练动态变化的问题。真正有价值的不是复制 teacher 的每一步,而是提取其中与当前 preference objective 一致的知识。 这也给 efficient generation 一个更一般的启发:压缩模型不必永远追赶固定上限;只要 guidance 能感知目标与能力差距,student 就可以在保持效率的同时继续优化质量。
English Summary
RATS targets the teacher-bounded nature of few-step generation. Static distillation provides useful intermediate supervision but limits the student to imitation, while reward-only training offers freedom to improve but supplies only delayed terminal guidance over a highly compressed trajectory.
Core Idea
Sigma-aligned multi-horizon matching connects key stages of the few-step student and multi-step EMA teacher. A reward-aware gate then adjusts shaping strength from their relative reward: stronger teacher performance increases guidance, whereas a competitive student is allowed to deviate.
Why This Helps
The student receives informative trajectory priors when they are useful without being permanently constrained by the teacher. The teacher is used only during training, so deployment keeps the original few-step inference cost.
Practical Takeaway
In preference-aligned distillation, a teacher is most useful as a conditional source of knowledge rather than a fixed target. Guidance should weaken once imitation stops serving the reward objective.