RWTD Aligns One-Step Generators with Black-Box Rewards, Lifting GenEval from 0.73 to 0.80

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

Austin Wang, Ziheng Cheng, Lexing Ying

cs.LG, cs.CV

2026-09-25

RWTD aligns one-step generators using only samples and scalar rewards, lifting SANA Sprint GenEval from 0.73 to 0.80 with better cross-reward generalization than gradients.

What problem this solves

Diffusion and flow models deliver high-fidelity images at the cost of dozens of sequential network evaluations. Distillation compresses them into one-step generators such as SANA Sprint and SDXL Turbo, which fixes speed and leaves post-training unsolved. RLHF-style alignment fails here in two ways: an implicit one-step generator offers neither a tractable likelihood nor a denoising trajectory to backprop through, and many rewards (a scorer model, a human rating) are non-differentiable. Existing one-step alignment methods demand reward gradients (FAV, DI++), stay tied to diffusion dynamics, or build specialized preference fields (DrPO). What is missing is an interface that needs nothing beyond sampling and scoring.

Method

RWTD (Reward-Weighted Transport Distillation) runs a three-step loop:

• Mixed target. The canonical alignment target is the pretrained reference distribution tilted by reward: mass reweighted by exp(β·r), with β controlling strength. RWTD detaches the target from the reference and uses (1-ρ) parts tilted current model plus ρ parts tilted reference. The current-model component absorbs high-reward regions discovered during training and speeds adaptation; the reference component keeps feeding pretrained modes back in so diversity holds.
• Feature-space optimal transport. Current and reference samples pass through a frozen visual encoder (DINOv2 for SANA Sprint, MAE for SDXL Turbo). With uniformly weighted current samples as the source and the reward-weighted mixture as the target, Sinkhorn solves entropic OT and returns how each current sample's mass splits across target particles. Feature space wins over pixel space because Euclidean pixel distance tracks neither perceptual nor semantic similarity.
• Partial transport plus fixed-point regression. Each source feature moves an η fraction of the way toward its assigned conditional mean, damping distribution jumps caused by minibatch coupling noise; the generator then regresses onto the frozen (stop-gradient) targets. Targets are rebuilt from the updated generator each round, making this fixed-point regression.

Backprop touches only the generator; gradients never flow through the reward, the weights, or the coupling, and the reward only ever returns a scalar. The theory gives the mixed fixed points in closed form: ρ=1 recovers the conventional tilted reference, while 0<ρ<1 introduces a reward-dependent rational factor that sharpens as ρ decreases.

Results

The main benchmark is GenEval (compositional correctness: position, counting, color-attribute binding):

A one-step model overtakes multi-step peers of similar scale, plus SD3.5-L (0.71) and the 4-step SANA Sprint (0.75). On SDXL Turbo, RWTD goes from 0.55 to 0.61, ahead of DrPO (0.58) and the preference-tuned LAIR-SDXL (0.59 at 50 NFEs).

The preference runs (reward: HPSv2, evaluation: PartiPrompts) show the sharpest contrast. DRaFT pushes in-domain HPSv2 to 37.02 but drops CLIP from 0.2753 to 0.2573 and out-of-domain GenEval from 0.73 to 0.62, textbook reward overoptimization. RWTD reaches only 32.45 on HPSv2, yet it is the only method that improves all five preference metrics plus GenEval over the base model (GenEval 0.75; FAV and DrPO both fall to 0.72). Qualitatively, gradient methods drift into heavy stylization while RWTD keeps realism.

Two ablations matter. Mixed targeting (ρ=0.15) trains faster and reaches 0.80 while the off-policy-only target (ρ=1) plateaus near 0.77. Dropping transport geometry (independent reward-weighted regression) raises Aesthetics slightly (6.878 vs 6.751) but hurts CLIP and PickScore, so the gain comes from OT's per-sample assignments rather than barycentric averaging.

Why it matters

Post-training a one-step generator used to be a dilemma: black-box rewards were untouchable, and gradient methods overoptimized on contact. RWTD lowers the bar to two calls, generate and score, so any scalar signal works, including non-differentiable evaluators and human ratings. The fixed-point analysis makes ρ an interpretable knob running from tilted reference to pure on-policy sharpening. Code and checkpoints are released. In scope this is an alignment-layer improvement on existing backbones, useful and incremental rather than a new generation paradigm.

Limitations

The authors state the theory is population-level and idealized; practice uses finite minibatches, entropic OT, and barycentric projections, so the propositions give intuition, not guarantees. The pure on-policy endpoint (ρ=0) has no unique fixed point and measurably destroys diversity, so ρ is a tuned hyperparameter (0.15 for GenEval). On in-domain reward depth RWTD (32.45 HPSv2) trails DRaFT (37.02); it wins on balance, not on raw optimization. Experiments cover two backbones at up to 8 GPUs, there is no human evaluation, and the cost discussion sits in an appendix. Transfer to larger backbones such as FLUX-scale models is untested.

Terms

• One-step generator: produces an image in a single network forward pass, trained by distillation or native objectives; far faster than multi-step diffusion.
• Reward tilting: reweighting a distribution's mass by exp(β·reward), the standard primitive for turning scores into target distributions.
• Optimal transport / Sinkhorn: moving probability mass between distributions at minimum total cost; entropy regularization makes it solvable by Sinkhorn iterations.
• GenEval: a text-to-image benchmark for compositional correctness (object position, counting, color-attribute binding).
• HPSv2 / PickScore: image-scoring models trained on human preference data, commonly used as rewards.
• NFEs: number of function evaluations, the count of network forward passes per generated image.

Source

• arXiv 2609.30840v1
• PDF
• Hugging Face
• GitHub · https://github.com/austin-k-wang/Reward-Weighted-Transport-Distillation

Related papers

• Muon-trained Moonlight gains 7.9 HumanEval points over its AdamW twin at 1.2T
• Parallax: a learned covariance branch cuts pretraining perplexity 17%, but only under Muon
• Muon beats Adam on ImageNet-C and typo text; ViT effective rank 6.88 vs 1.86
• Retrospective distillation lifts Qwen3-14B on ScienceWorld by 59.7% relative
• Qwen3-8B hits 41.5% exact-match tool selection, beating 402B Maverick
• DeepMind fixes two numerical bugs in softmax and cross-entropy, eliminating Slingshot loss spikes

All paper explainers