Aligning One-Step Generative Models with Reward-Weighted Transport Distillation

2026-09-25

AI Summary

This study addresses the challenge of post-training alignment for one-step generators, which lack computable likelihoods and denoising trajectories. To overcome this, the authors propose reward-weighted transport distillation, a method that relies solely on generated samples and scalar rewards. By constructing an adaptive target that interpolates between the current and reference distributions, it effectively balances reward adaptation with prior knowledge preservation. Technically, the approach integrates optimal transport in feature space, fixed-point regression, and theoretical analysis of reward tilting to achieve efficient optimization. Experimental results demonstrate that SANA Sprint improves its GenEval score from 0.73 to 0.80 while exhibiting strong cross-reward generalization capabilities.

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.

Problem

• one-step generative models
• post-training alignment
• implicit generators
• non-differentiable rewards
• reward-weighted distillation

Innovation

• Reward-Weighted Transport Distillation
• One-Step Generative Models
• Optimal Transport
• Fixed-Point Regression
• Preference Alignment