以 Reward-Weighted Transport Distillation 对齐一步生成模型

TL;DR RWTD 是一种后训练方法，仅用生成样本和标量奖励即可对齐一步生成模型，通过混合当前与参考分布的自适应目标，在奖励适应与先验保留之间取得平衡，显著提升 GenEval 分数。

问题

问题背景

一步生成模型（如 SANA Sprint 等）通过单次网络评估生成高质量图像，大幅降低推理成本，已成为视觉生成领域的重要方向。后训练对齐（文本遵循、美学、组合能力等）是当前提升模型实用性的关键环节。

现有方法局限

• 通用隐式一步生成器既不提供 可处理的似然，也不提供 去噪轨迹，因此基于扩散模型的强化学习或轨迹对齐方法无法直接迁移。
• 许多奖励函数（美学评分、人类偏好模型）是非可微 的，无法通过反向传播直接优化。
• 已有基于样本的奖励加权方法往往依赖固定参考分布或离策略数据，难以在训练中平衡当前模型的在线探索与预训练先验保留，且对一步模型的适配不足。

为什么难 / 重要

一步模型为推理效率牺牲了 likelihood 与 denoising trajectory，导致后训练工具链出现断层。业界需要一种仅利用生成样本和标量奖励即可优化模型的方法，同时避免奖励 hacking 与灾难性遗忘。此外，由于像素空间最优传输成本高且不稳定，必须在特征空间构造有效目标分布，这增加了方法设计难度。

RWTD 通过自适应混合当前分布与参考分布的目标，并用特征空间最优传输与不动点回归实现，在理论上证明了固定点分布可插值于离策略奖励倾斜与在线策略倾斜之间，为平衡奖励适应和先验保留提供了原则化方案。

行业类比

类似自动驾驶中，仅通过试驾采样与黑盒驾驶评分来优化策略网络：无法获得中间轨迹梯度，必须依赖基于样本的奖励加权更新。

核心洞察

• RWTD 的核心创新在于构造了一个混合当前模型与参考模型各自 reward-tilted 分布的自适应目标,而不是像多数后训练方法那样仅对齐固定的 reward-tilted 参考分布。这种设计允许在训练过程中同步纳入当前模型已学到的高奖励改进,同时以参考分布作为锚点防止过度漂移,从而在奖励适应与先验知识保留之间取得原理性平衡。理论分析表明其不动点分布可在这两种倾斜模式之间插值,工程上可通过调节混合比例控制探索-利用权衡,对 one-step 生成器的持续微调尤其关键。
• 针对 one-step 隐式生成器既无可用似然也无去噪轨迹的难点,RWTD 仅依赖生成样本与标量奖励,通过特征空间最优传输和不动点回归实现蒸馏。这不同于依赖 likelihood 或 denoising 的扩散模型对齐方法,直接解决了 SANA Sprint 等 one-step 骨干网络的后训练不可微奖励问题。GenEval 从 0.73 提升至 0.80 且偏好对齐实验展示强跨奖励泛化,说明该方法能在保持构图能力的同时获得平衡提升,工程落地性强。

方法

RWTD 的输入为：一个预训练的一步生成器（如 SANA Sprint 1.6B 骨干）、一个标量奖励函数（不可微分，仅需对生成样本打分）、以及一个特征提取器（用于度量分布间距离）。训练完全基于当前生成器采样，无需真实数据或似然估计。

目标分布构造：算法每次迭代构建自适应目标。分别对 参考分布（预训练生成器分布）和 当前模型分布 进行奖励加权倾斜（reward tilting），再将两个倾斜后的分布按比例混合。参考分量采用 off-policy 方式，利用预训练模型样本计算倾斜权重；当前分量采用 on-policy 方式，利用本次迭代生成样本计算倾斜权重。混合后的分布既包含当前模型已发现的改进，又保留预训练先验锚定。

特征空间最优传输：在特征空间（如预训练视觉编码器的输出）中，通过 最优传输 求解从当前生成分布到目标分布的映射，得到每个生成样本对应的目标特征点。使用可微 Sinkhorn 迭代或熵正则化 OT，避免在像素空间直接匹配，降低维度灾难。

固定点回归：将当前生成样本的特征与其对应的目标特征构成回归对，以最小化特征空间距离为目标更新生成器参数。重复采样与更新直至收敛到固定点，使生成器分布逼近自适应目标。

与同类方法（如 RAFT、DDPO 直接对齐单一奖励倾斜参考分布）的差异在于：RWTD 通过混合 on-policy 与 off-policy 倾斜分布，在奖励适应与先验知识保留之间取得平衡，避免过度的奖励黑客或模式崩溃。

实验

实验设计

RWTD 在 SANA Sprint 1.6B 单步生成器上进行后训练，使用 GenEval 基准评估视觉生成质量。方法仅需生成样本与标量奖励，通过特征空间最优传输和固定点回归构造自适应目标分布：该目标混合 当前模型 的奖励倾斜分布与 预训练参考 的奖励倾斜分布。另设独立偏好对齐实验，测试跨奖励泛化能力。

关键发现

• GenEval 分数 从 0.73 提升到 0.80，提升幅度显著。
• 跨奖励泛化实验显示平衡的改进，同时保留组合能力。
• 理论分析表明 RWTD 的固定点分布在 off-policy 奖励倾斜（参考模型）与 on-policy 奖励倾斜（当前模型）之间插值，提供原则性权衡。

与基线对比

传统一步生成器后训练受限于不可微奖励、无似然与无去噪轨迹。RWTD 只需样本和标量奖励，工程部署门槛低。相比单纯对齐奖励倾斜参考分布，混合当前模型分布可在线吸收训练中发现的改进，并锚定预训练先验，避免奖励过拟合。

行业影响

落地场景

一步生成模型 已在电商、广告、内容平台等需要低延迟批量生成视觉素材的场景中得到应用。RWTD 允许在不依赖可微奖励或似然函数的条件下，仅用生成样本和标量奖励（如点击率、转化率、美学评分）对预训练生成器进行后训练。典型场景包括：

• 电商：商品图背景替换、多角度展示生成，用点击/购买信号优化生成效果。
• 内容平台：短视频封面、文章配图个性化生成，根据用户互动数据迭代。
• 设计工具：实时生成 logo、海报初稿，利用用户偏好评分提升设计质量。

商业价值

• 降本：一步生成本身大幅降低推理延迟与算力成本，RWTD 不改变推理架构，仅需训练阶段标量奖励，进一步降低后训练门槛。
• 增收 / 体验提升：通过奖励对齐，生成内容更符合业务指标，提高广告点击率、商品转化率；同时保持低延迟实时生成，改善交互体验。

与现有工作流接口

RWTD 可作为标准 post-training 阶段 集成进现有 ML 流水线：

1. 使用预训练一步生成器采样一批数据；
2. 调用 reward 服务（如美学模型、业务指标 API）获取标量奖励；
3. 运行特征空间最优传输与固定点回归，更新生成器权重；
4. 部署更新后的模型，无需修改推理服务架构。

该方法与 RLHF/DPO 等对齐流程互补，可替代需要可微奖励的梯度方法，降低对奖励模型结构的依赖。

具体 Use Case

• 电商平台：对商品场景图生成做点击率奖励对齐，RWTD 在不改动推理延迟的前提下提升图片吸引力，直接关联转化指标。
• 短视频平台：封面图生成模型接入用户停留时长/点赞数作为奖励，定期离线更新，实现个性化封面优化。

局限

• **计算开销与训练复杂度较高** 该方法依赖于特征空间最优传输和固定点回归，每次迭代需要求解传输计划，计算成本高于常规的奖励加权微调或直接策略优化。且需要额外预训练的特征提取器，增加了部署和训练时的资源需求，可能限制其在计算资源受限场景下的应用。
• **实验验证范围有限** 摘要仅报告了在 **SANA Sprint 1.6B** 模型上的 **GenEval** 分数提升（0.73→0.80）以及偏好对齐实验的交叉奖励泛化，未提供与现有后训练方法（如 **DPO**、**ReFL**、**DRaFT** 等）的直接定量对比，也未在更多模型架构或数据集上验证，结论的普适性有待进一步检验。
• **混合比例超参数需手动调节** 目标分布中 current 与 reference 的混合系数决定了固定点位置，理论分析表明固定点在该系数控制下插值于 off-policy 与 on-policy 之间，但论文可能未给出自动选择该系数的机制，实际训练中可能需要依赖经验调参，增加了使用门槛。