MEGA Hub

RL-Native Distillation: Exploiting Scored Trajectories for Few-Step Image Generation

Authors

Do you know Yuhan Li?You can claim authorship or link another user.Do you know Fangao Zeng?You can claim authorship or link another user.Do you know Sicong Kang?You can claim authorship or link another user.Do you know Mengfei Xu?You can claim authorship or link another user.Do you know Hao Zhou?You can claim authorship or link another user.Do you know Wei Li?You can claim authorship or link another user.Do you know Pipei Huang?You can claim authorship or link another user.Do you know Bingbing Ni?You can claim authorship or link another user.

Abstract

Efficient text-to-image generation requires both reinforcement-learning (RL)-based reward alignment and few-step distillation, yet these procedures are typically performed sequentially, increasing training cost and risking the loss of reward gains during compression. We instead take an RL-native perspective: diffusion RL already generates reward-scored finite-step trajectories, whose intermediate states provide a natural source of distillation supervision rather than a disposable byproduct of sampling. Based on this insight, we propose REST (Reward-Enhanced Scored-Trajectory Distillation), a single-stage RL-distillation co-training framework that attaches a decoupled student to an arbitrary RL teacher. The student learns segment-wise from the teacher's evolving rollout trajectories while leaving the original teacher optimization unchanged. To prevent uniform imitation from preserving undesirable low-reward behaviors, we further introduce Advantage-Modulated Distillation (AMD), which transforms rollout advantages into signed weights over a base distillation loss. AMD strengthens supervision from preferred trajectories and mildly repels the student from low-reward ones. The resulting framework is lightweight and plug-and-play, requires no extra image rollouts, no separate distillation dataset, and no adversarial training. Experiments on compositional generation, visual text rendering, and human-preference alignment show that REST enables few-step CFG-free inference that matches or surpasses its 40-step RL teacher, with an overall additional training cost below 25% over pure RL. REST improves DrawBench PickScore over RTDMD by 0.82 while requiring only one-fifth of the training iterations.

Community

00