Aligning One-Step Generative Models with Reward-Weighted Transport Distillation
cs.LG, cs.CV
Submitted: 2026-09-25
Updated: 2026-09-25
Code: https://github.com/austin-k-wang/Reward-Weighted-Transport-Distillation
License: http://creativecommons.org/licenses/by/4.0/
The gist: 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
Terminology
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.
Sources
- Rethinking the Design Space of Reinforcement Learning for Diffusion Models: On the Importance of Likelihood Estimation Beyond Loss Design
- Generative Modeling via Drifting
- Representation Distribution Matching for One-Step Visual Generation
- One-Step Generative Modeling via Wasserstein Gradient Flows
- LoRA: Low-Rank Adaptation of Large Language Models
- How to Guide Your Flow: Few-Step Alignment via Flow Map Reward Guidance
- Drifting Preference Optimization for One-Step Generative Models
- Aligning Few-Step Generative Models by Amortizing Sample-based Variational Inference
- Flow Matching for Generative Modeling
- Playground v3: Improving Text-to-Image Alignment with Deep-Fusion Large Language Models
- Smoothed Preference Optimization via ReNoise Inversion for Aligning Diffusion Models with Varied Human Preferences
- Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference
- Diff-Instruct++: Training One-step Text-to-image Generator Model to Align with Human Preferences
- DINOv2: Learning Robust Visual Features without Supervision
- A General Framework for Inference-time Scaling and Steering of Diffusion Models
- Score-Based Generative Modeling through Stochastic Differential Equations
- Consistency Models
- Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review
- Inference-Time Alignment in Diffusion Models with Reward-Guided Generation: Tutorial and Review
- Beyond Pairwise Preferences: Listwise Reward-Aware Alignment for Diffusion Models
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks