One-step generative models improved by mixing reward and prior knowledge
Aligning One-Step Generative Models with Reward-Weighted Transport Distillation
Machine LearningComputer Vision and Pattern Recognition
Summary
Generating high-quality images with a single step is hard to improve after training because the usual approaches don’t work well with these models or with rewards that can’t be easily measured. The authors propose a method called Reward-Weighted Transport Distillation (RWTD) that adjusts the model using only samples and simple reward scores. RWTD balances learning from new rewards with keeping what the model already knows by mixing old and new goals in a smart way. This approach helps the model get better at creating images while still retaining its original skills.
What this means in practice
- •For visual ai developers: Enhance pretrained image generation models by balancing new reward improvements with prior knowledge retention.
- •For media content producers: Use improved one-step generative models to create higher-quality visual content with fewer computational steps.
Authors
Austin Wang, Ziheng Cheng, Lexing Ying
Abstract
One-step generators enable high-quality visual generation with a single network evaluation, but their post-training is difficult: general implicit generators provide neither tractable likelihoods nor denoising trajectories, and many rewards are non-differentiable. We introduce Reward-Weighted Transport Distillation (RWTD), a post-training method that requires only generated samples and scalar reward evaluations. Rather than aligning solely to the conventional reward-tilted reference distribution, RWTD constructs an adaptive target that mixes separately tilted current and reference distributions. The current component incorporates improvements discovered during training, while the reference component anchors the target to the pretrained generator. RWTD realizes this target through feature-space optimal transport and fixed-point regression. Theoretical analysis shows that the fixed-point distributions of RWTD interpolate between off-policy reward tilting of the reference and on-policy tilting of the current model, providing a principled approach to balancing reward adaptation with retention of prior knowledge. Empirically, RWTD substantially improves the GenEval score of the one-step SANA Sprint 1.6B backbone from 0.73 to 0.80, while separate preference alignment experiments demonstrate strong cross-reward generalization that yields balanced improvements and preservation of compositional capabilities.