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

Aligning One-Step Generative Models with Reward-Weighted Transport Distillation. By Austin Wang and two others, from ByteDance Seed and UC Berkeley. Published on the twenty-fifth of September. Its code repository is worked out: about four hundred source files, with a readme.

One-step generators produce high-quality images with a single network evaluation. Post-training them is another matter. A general implicit generator offers neither a tractable likelihood nor a denoising trajectory to work with, and many of the rewards one would want to align to are not differentiable. Most of the usual handles are missing. Methods that need likelihoods have nothing to compute, methods that need a trajectory have no steps to follow, and methods that backpropagate through the reward have no gradient to use.

Reward-Weighted Transport Distillation, RWTD, asks for very little: generated samples, and a scalar reward for each. The usual move is to align to the reference distribution tilted by the reward. RWTD aims instead at an adaptive target that mixes two separately tilted distributions. One is the current model, which brings in improvements discovered during training. The other is the reference, which keeps the target anchored to the pretrained generator.

That target is realized through optimal transport in feature space, combined with fixed-point regression. The theory explains what the mixture buys. The fixed-point distributions of RWTD interpolate between two familiar extremes: off-policy reward tilting of the reference, and on-policy tilting of the current model. Sitting between them is the authors' principled way of balancing adaptation to the reward against keeping what the model already knew.

On the one-step SANA Sprint backbone, at one point six billion parameters, RWTD raises the GenEval score from zero point seven three to zero point eight zero, which the authors describe as a substantial improvement. A separate set of preference alignment experiments shows strong cross-reward generalization. Improvements stay balanced across rewards, and the model's compositional capabilities are preserved, which is the retention half of the balance the theory describes.

A generator without likelihoods or trajectories, aligned to rewards without gradients. The method gets there by asking only for samples and scores, and by refusing to choose between chasing the reward and staying near where it began: its target is built to sit somewhere between the two.