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

2026-08-10

AI Summary

This work addresses the high training cost and reward degradation inherent in existing text-to-image generation methods that decouple reinforcement learning (RL) reward alignment from few-step distillation into separate stages. The authors propose REST, a unified framework that jointly optimizes RL and distillation in a single stage by directly leveraging intermediate states from reward-scored trajectories as distillation supervision signals during RL training. Innovatively treating scored trajectories in diffusion RL as natural sources for distillation, REST introduces Advantage-Modulated Distillation (AMD), which employs signed weights to amplify high-reward trajectories and suppress low-reward behaviors—without requiring additional data or adversarial training. Experiments demonstrate that REST enables classifier-free few-step inference matching or surpassing the performance of 40-step RL teacher models, with less than a 25% increase in training cost, a 0.82 improvement in DrawBench PickScore over RTDMD, and a fivefold reduction in training iterations.

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.

Problem

• reinforcement learning
• distillation
• text-to-image generation
• reward alignment
• few-step generation

Innovation

• RL-native distillation
• scored trajectories
• few-step generation
• advantage-modulated distillation
• co-training framework