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

📄 arXiv: 2609.30840v1 📥 PDF

作者: Austin Wang, Ziheng Cheng, Lexing Ying

分类: cs.LG, cs.CV

发布日期: 2026-09-25

💡 一句话要点

提出奖励加权传输蒸馏方法以优化一步生成模型

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 一步生成模型 奖励加权 传输蒸馏 后期训练 生成对抗网络 计算机视觉 强化学习

📋 核心要点

  1. 现有的一步生成模型在后期训练中面临困难，尤其是在处理不可微的奖励函数和缺乏可处理的似然性时。
  2. 本文提出的RWTD方法通过仅依赖生成样本和奖励评估，构建适应性目标，结合当前和参考分布的倾斜，来优化生成模型。
  3. RWTD在实验证明中显著提升了GenEval得分，并在偏好对齐实验中展示了良好的跨奖励泛化能力，保持了生成模型的组合能力。

📝 摘要（中文）

一步生成器能够通过单次网络评估实现高质量的视觉生成，但其后期训练面临挑战：一般的隐式生成器既没有可处理的似然性，也没有去噪轨迹，许多奖励函数是不可微的。本文提出了奖励加权传输蒸馏（RWTD）方法，仅需生成样本和标量奖励评估。RWTD通过构建适应性目标，结合当前和参考分布的倾斜，来实现这一目标。理论分析表明，RWTD的固定点分布在参考模型的离策略奖励倾斜和当前模型的在策略倾斜之间进行插值，提供了一种平衡奖励适应与保留先前知识的原则性方法。实验证明，RWTD显著提升了一步SANA Sprint 1.6B主干的GenEval得分，从0.73提高至0.80，同时在偏好对齐实验中展示了强大的跨奖励泛化能力。

🔬 方法详解

问题定义：本文旨在解决一步生成模型在后期训练中面临的挑战，尤其是隐式生成器缺乏可处理的似然性和去噪轨迹，以及许多奖励函数的不可微性问题。

核心思路：RWTD方法通过构建一个适应性目标，结合当前模型和参考模型的奖励倾斜，来实现更有效的后期训练。这种设计旨在同时保留先前知识和适应新的奖励信号。

技术框架：RWTD的整体架构包括生成样本的收集、奖励评估、目标构建和特征空间的最优传输。主要模块包括当前和参考分布的倾斜混合，以及通过固定点回归实现目标的过程。

关键创新：RWTD的核心创新在于其适应性目标的构建方式，通过结合当前和参考分布的倾斜，提供了一种新的平衡奖励适应与知识保留的方法。这与传统的奖励倾斜方法有本质区别。

关键设计：RWTD在参数设置上采用了特征空间的最优传输技术，并通过固定点回归来实现目标的优化。损失函数设计上，强调了当前和参考分布的结合，以确保模型在训练过程中能够有效地适应新的奖励信号。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

RWTD方法在实验证明中显著提升了一步SANA Sprint 1.6B主干的GenEval得分，从0.73提高至0.80，展示了其在后期训练中的有效性。此外，偏好对齐实验表明RWTD具有强大的跨奖励泛化能力，能够在多个任务中保持生成模型的组合能力。

🎯 应用场景

该研究的潜在应用领域包括计算机视觉、生成对抗网络和强化学习等。RWTD方法能够提升生成模型的性能，具有广泛的实际价值，尤其是在需要高质量生成的任务中，如图像生成、视频生成等。未来，该方法可能会推动生成模型在更多复杂场景中的应用。

📄 摘要（原文）

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.