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

The paper addresses a fundamental challenge in text-to-image generation: 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. The authors note that distillation may wash out the reward gains obtained from RL, while subsequent RL can disturb a previously distilled few-step generator and degrade its structure or low-step inference stability.

The key insight is that 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. The authors take an RL-native perspective: distillation should be built directly into RL training, without a separate pre-distillation stage, external data, or auxiliary models such as fake-score networks.

The authors propose REST (Reward-Enhanced Scored-Trajectory Distillation), described as a single-stage RL-distillation co-training framework that attaches a decoupled student to an arbitrary RL teacher. The framework has three key components:

1. Decoupled Teacher-Student Co-Training: The teacher retains its original reward-optimization objective, while the student performs segment-wise imitation along the same reward-scored rollout trajectory without affecting teacher optimization. The teacher branch samples with an M-step schedule (e.g., M=20) with classifier-free guidance, producing a trajectory T = xt0, xt1,..., xtM and a terminal reward r = R(x0, c). The student is designed to run with only K ≪ M steps (e.g., K=8) at inference time. Rather than sampling a separate trajectory for the student, the method select a K-step subset of the teacher's schedule and re-index it as an imitation trajectory, computing piecewise trajectory velocity for each student step as vkgt = (xtq(k+1) − xtqk)/(tq(k+1) − tqk).

2. Advantage-Modulated Distillation (AMD): This mechanism transforms rollout advantages into signed weights over a base distillation loss. The AMD objective is formalized as: LAMD(θ) = E(sk,ak)∼D [λ(Amix + b) lbase(θ; k)], where λ = 1 is a global scale, b = 0.5 is a positive shift, and Amix is the fused advantage from multiple reward sources. The advantages are clipped group-normalized advantage of the rollout associated with prompt c normalized to [−1, 1]. The method strengthens supervision from preferred trajectories and mildly repels the student from low-reward ones.

3. EMA Regularization: The student is regularized against an EMA copy of itself using a fixed-variance Gaussian KL surrogate: LKL-EMA(θ) = Ek ∥vθ(sk) − vθema(sk)∥2, with the final student objective being Lstudent(θ) = LAMD(θ) + βLKL-EMA(θ), where β = 0.2.

The paper decomposes the AMD objective into two terms: The first term is a constant-weighted imitation prior (λb E[lbase]) and The second term is the reward-driven correction inherited from reward-weighted supervised regression (λ E[Amix lbase]). The positive shift b is crucial because the initial student branch does not yet possess reliable image-generation capability, and therefore needs a stable positive imitation signal to bootstrap its few-step generator.

AMD is described as a generic reward-aware distillation wrapper that can wrap more advanced distillation losses, such as PCM-style phase consistency losses, DMD-style distribution-matching losses, or other trajectory-level objectives.

The experiments use Stable Diffusion 3.5 Medium (SD3.5M) as the backbone with LoRA training (r=32, α=64), bf16 precision, learning rate 3×10−4, 48 prompts per group, and 24 samples per prompt. Three reward objectives are considered: GenEval (compositional correctness), OCR (visual text rendering), and PickScore (human preference alignment). For OCR and GenEval, a multi-reward setting with PickScore is used to prevent reward hacking, since although OCR or GenEval scores can become high, the generated images may collapse into reward-hacking artifacts.

The teacher follows RAM's forward-process RL setup: 20-step training solver, 40-step evaluation solver, and CFG-enabled sampling. The student is jointly trained with AMD and evaluated using the EMA student (decay 0.9) with 8 CFG-free steps.

Performance on Three Tasks: REST achieves performance comparable to the 40-step RAM teacher on both training rewards and out-of-domain DrawBench metrics, while using only a few CFG-free inference steps. Key results include:

• For GenEval: REST achieves 0.94 GenEval score with 8 steps vs. 0.95 for 40-step RAM teacher

• For OCR: REST achieves 0.96 OCR score with 8 steps vs. 0.96 for 40-step RAM teacher

• For PickScore: REST achieves 23.73 PickScore with 8 steps vs. 23.70 for 40-step RAM teacher

Comparison with RTDMD: REST improves DrawBench PickScore over RTDMD by 0.82 while requiring only one-fifth of the training iterations. Specifically, REST achieves 23.96 PickScore with 500 training iterations, while RTDMD achieves 23.14-23.71 with 2,500 iterations (1,500 warm-up + 1,000 parallel training).

Training Efficiency: Sampling rollout with reward computation accounts for 40.7% of the wall-clock time, teacher training for 33.3%, and student training for only 26.1%. The overall additional computational cost introduced by REST is below 25% over pure RL.

The progressive ablation shows:

• Full-step standard RL improves 8-step CFG-free sampling (Aesthetic: 5.15→6.42)

• Simultaneous dual-branch RL-distillation retains most gains of RL-then-distill pipeline while reducing training cost by 1/4 iterations

• AMD further strengthens reward preference and improves perceptual quality

The ablation of AMD components shows:

• Setting b=0 removes this baseline imitation pressure, making the objective dominated by aggressive negative modulation and thus potentially destabilizing early student training

• A moderate shift (b=0.5) preserves the necessary imitation prior while turning clearly low-quality trajectories into repulsive supervision

• Removing EMA regularization causes much stronger performance oscillation rather than a stable improvement

The paper validates AMD with both per-step imitation MSE and PCM-style phase consistency losses. The PCM variant surpasses 8-step RAM in our experiments, and gets similar performance with segment-velocity MSE, implying REST's compatibility with other distillation losses.

The appendix provides theoretical support showing that RAM and DiffusionNFT share an anchored, reward-shifted velocity-regression structure despite differences in their anchor choices and hyperparameters. Both methods perform forward-process velocity regression along the same reward-controlled direction vgt − vold. DiffusionNFT uses vanchor = vold and γ(A) = A/β, while RAM uses vanchor = vbase and γ(A) = ηARAM.

The authors note two main limitations: its integration with DMD2 remains unexplored (a promising direction for 4-step CFG-free inference), and our experiments instantiate the teacher branch only with RAM. However, they provide theoretical support for extending REST to DiffusionNFT and other ODE-based diffusion RL teachers.

The paper summarizes three key advantages: (1) it delivers high-quality few-step generation, matching or even surpassing the full-step inference RL-trained teacher; (2) it is plug-and-play: the decoupled teacher–student design and the generic AMD formulation make the method independent of a specific teacher RL algorithm or a particular distillation loss; (3) "it is simple and efficient: the framework avoids warm-up scheduling and delicate coordination among multiple competing objectives; it also requires no extra image rollouts, no separate distillation dataset, and no fake-score model training."

Based on the paper, here are the specific improvements I can make to AI systems:

• Improvement: Integrate reward optimization and few-step distillation into a single training phase rather than sequential stages

• Capability: The system can simultaneously align with human preferences and learn efficient inference paths, eliminating the risk of reward gains being washed out during compression or distillation destabilizing the generator

• Improvement: Implement a signed weighting mechanism that uses rollout advantages to modulate distillation loss—strengthening supervision from high-reward trajectories while mildly repelling from low-reward ones

• Capability: The system can automatically prioritize learning from preferred outputs and avoid reward-hacking artifacts, producing higher-quality images that better match compositional and aesthetic criteria

• Improvement: Attach a separate student branch to an existing RL teacher without interfering with teacher optimization, using the teacher's reward-scored trajectories as direct supervision

• Capability: The system can achieve 5× faster inference (8 steps vs. 40 steps) while matching or exceeding the teacher's performance on GenEval (0.94 vs. 0.95), OCR (0.96 vs. 0.96), and PickScore (23.73 vs. 23.70)

• Improvement: Add a fixed-variance Gaussian KL surrogate against an EMA copy of the student to stabilize training

• Capability: The system exhibits stable, monotonic improvement without performance oscillation, ensuring reliable few-step generation quality across training iterations

• Improvement: Combine multiple reward signals (e.g., GenEval + PickScore, OCR + PickScore) with clipped group-normalized advantages

• Capability: The system can simultaneously optimize for task-specific correctness (compositional, text rendering) and human preference, preventing reward hacking while maintaining high perceptual quality

• Improvement: Design AMD as a wrapper that works with various base distillation losses (per-step MSE, PCM-style phase consistency, DMD-style distribution matching)

• Capability: The system can adapt to different distillation objectives without architectural changes, making it flexible for future improvements in few-step generation

• Improvement: Eliminate warm-up phases, separate distillation datasets, extra image rollouts, and fake-score model training

• Capability: The system reduces training cost by 75% compared to sequential RL-then-distill pipelines (500 iterations vs. 2,500), while achieving better DrawBench PickScore (23.96 vs. 23.14–23.71)

• Improvement: Leverage the unified theoretical framework showing RAM and DiffusionNFT share anchored, reward-shifted velocity-regression structures

• Capability: The system can be extended to multiple RL teacher algorithms without re-engineering, supporting broader applicability across different diffusion model families

The improved AI system can: Generate high-quality images in 8 CFG-free steps that match or exceed 40-step RL-trained teachers across compositional, text-rendering, and human-preference metrics, while training 5× faster and requiring no additional data or auxiliary models—all within a single-stage, plug-and-play framework.

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.