基于 Wasserstein 梯度流的一步生成模型奖励引导微调

Reward-guided Fine-Tuning of One-Step Generative Models via Wasserstein Gradient Flow

arXiv: 2608.29647v1

论文信息

标题: Reward-guided Fine-Tuning of One-Step Generative Models via Wasserstein Gradient Flow

作者: Hoseong Hwang, Woorim Han, Joungin Chun, et al.

发布日期: 2026-08-30

arXiv ID: 2608.29647v1

PDF 链接: 下载 PDF

3 分钟速览

  • 研究问题：一步生成模型虽然能通过单次前向映射显著降低推理成本，但缺少适用于其结构的奖励引导微调方法；现有多步扩散模型微调依赖中间去噪状态，无法直接迁移到一步模型。
  • 核心方法：将奖励引导微调建模为 Wasserstein 梯度流（WGF），通过 JKO 方案离散化，推导半对偶形式的目标，并用重要性采样估计奖励加权目标分布，从而无需奖励梯度即可训练。
  • 关键结果：在 ImageNet 256×256 的四个奖励任务中，本文方法在三个任务上取得最高奖励，并在所有四个任务上的美学分数都高于 DPO（表 3）；在 CIFAR-10 JPEG 不可压缩性任务中，奖励从预训练的 1.30 提升到 1.62，但 FID 从 2.90 上升至 55.9（表 1）。
  • 主要局限：作者明确说明该方法尚未在 prompt 条件文本到图像生成这一重要应用中评估（论文第 6 节）；理论收敛保证依赖于奖励加权目标强对数凹的假设，而作者指出实际微调任务通常不满足该条件。
  • 适合读者：关注生成模型微调、一步生成模型、最优传输、奖励对齐与黑盒奖励优化的研究者和工程师，尤其是需要处理不可微奖励指标的团队。

论文背景和研究动机

扩散模型在图像生成质量上表现突出，但其推理需要求解概率流 ODE 或执行数十到数百步去噪迭代，计算成本高。为缓解这一问题，近年出现了一类一步生成模型，如 Consistency Trajectory Model、MeanFlow 等，它们直接学习从噪声分布到数据分布的映射 T0T_0，生成过程只需一次前向传播（论文第 2.1 节）。

与推理加速同步发展的是奖励微调：让预训练生成模型在图像美学、可压缩性、类别对齐、文本-图像对齐等目标下被进一步优化。然而，已有奖励微调方法大多针对多步扩散或流模型设计。它们依赖采样轨迹中的中间状态构造损失；一步生成器不产生这些中间状态，且其生成路径通常为线性，不同于多步模型假设的非线性轨迹。因此，如果直接把 DDPO、DPOK 等方法用于一步模型，结构上不可行。

该论文把一步生成器的奖励微调问题放在最优传输视角下。预训练模型生成分布 T0#μ≈νT_0\#\mu\approx\nu，微调目标是让该分布向奖励加权分布 νr\nu_r 演化。论文选择 Wasserstein 梯度流作为概率空间中的动力学框架，而不是在扩散时间步上做优化。这使更新直接在分布层面进行，并借由 Wasserstein 距离引入对大幅分布偏移的惩罚，从而在训练稳定性和奖励对齐之间保持平衡。

核心方法和技术细节

论文首先定义奖励加权目标分布：

νr(y)=exp⁡(βr(y))Mβν(y),Mβ=∫exp⁡(βr(y)) dν(y),\nu_r(y)=\frac{\exp(\beta r(y))}{M_\beta}\nu(y),\quad M_\beta=\int \exp(\beta r(y))\,d\nu(y),

其中 rr 是奖励函数，β\beta 控制奖励强度。目标是最小化当前分布 ρ\rho 与 νr\nu_r 之间的 ff-散度 Df(ρ∣νr)D_f(\rho\mid\nu_r)。直接求解连续 WGF 困难，因此采用 JKO 方案离散化：

νk+1=arg⁡min⁡ρ[12hW22(νk,ρ)+Df(ρ∣νr)].\nu_{k+1}=\arg\min_{\rho}\left[\frac{1}{2h}W_2^2(\nu_k,\rho)+D_f(\rho\mid\nu_r)\right].

这一步可以理解为：在保持与当前分布 νk\nu_k 的水分子的距离不太远的前提下，向奖励加权目标方向移动。W22W_2^2 项充当近端正则，使更新更平滑。

为了让 JKO 步骤可训练，论文使用半对偶形式。对于代价 ch(x,y)=12h∥x−y∥22c_h(x,y)=\frac{1}{2h}\|x-y\|_2^2，通过 cc-变换引入势函数 vv，并把更新映射 ΔTk\Delta T_k 隐式定义为

ΔTk(x)=arg⁡min⁡y(ch(x,y)−v(y)).\Delta T_k(x)=\arg\min_y\left(c_h(x,y)-v(y)\right).

最终单个 JKO 步骤的目标可以写成一类对抗形式：生成器 ΔTk\Delta T_k 最小化运输代价与势函数的差值，判别器 vv 同时拟合来自 νk\nu_k 和 νr\nu_r 的样本。论文进一步用累积变换 Tk=ΔTk∘⋯∘ΔT0T_k=\Delta T_k\circ\cdots\circ\Delta T_0 替换逐步模拟，避免每次更新都沿着整个 JKO 轨迹采样，只需从初始噪声分布 μ\mu 采样。

一个关键实现问题是 νr\nu_r 不可直接采样。论文通过重要性采样解决：

∫f∘(v(y)) dνr(y)=∫f∘(v(y))exp⁡(βr(y))Mβ dν(y).\int f^\circ(v(y))\,d\nu_r(y)=\int f^\circ(v(y))\frac{\exp(\beta r(y))}{M_\beta}\,d\nu(y).

这意味着训练样本从预训练分布 ν\nu 抽取，仅用标量权重 exp⁡(βr(y))/Mβ\exp(\beta r(y))/M_\beta 加权。奖励函数不需要对生成器输出求导，因此可以处理 JPEG 压缩文件大小等不可微奖励。论文还指出固定参考分布 ν\nu 会与当前生成器逐渐失配，因此训练中动态更新 ν←νk\nu\leftarrow\nu_k（论文第 3.2 节）。最终算法交替更新生成器 TθT_\theta 与判别器 vϕv_\phi。

创新点和贡献

论文给出三个主要贡献。第一，它提出了首个基于 WGF 和最优传输的一步生成模型奖励引导微调框架。此前一步模型微调方法缺乏系统性理论工具；该方法不是简单地把多步模型损失裁剪到一步，而是重新定义概率分布演化问题。

第二，训练算法不依赖奖励梯度。与需要奖励反传的 ReFL 等方法不同，本文通过重要性采样把奖励转化为样本权重，因此既可以处理可微奖励，也可以处理不可微奖励。论文在 CIFAR-10 的 JPEG 不可压缩性和可压缩性任务上验证了这一点，这些奖励来自离散黑盒压缩算法，没有可用梯度。

第三，WGF 和 JKO 方案赋予了更新过程稳定性和保守性。W22W_2^2 近端项抑制大幅分布漂移，缓解奖励破解和模式坍塌。在 8-Gaussians 合成实验（图 1）中，本文方法逐步分离出八个目标模式；而 SFT、RWR、DPO 在该实验设置下出现模式平均或合并为连续环的现象。论文还给出一个理论动机：当 νr\nu_r 强对数凹时，WGF 在 W2W_2 距离下指数收敛（Remark 1）。不过作者明确说明实际奖励加权目标通常不满足该条件，所以该结果只作为理论动机，不算实用性保证。

实验结果分析

在 2D 合成数据上，论文通过可视化展示训练动态：本文方法在 5k 到 25k 迭代间稳定地分离出八个模式；SFT、RWR、DPO 则难以避免结构坍塌（图 1）。

CIFAR-10 实验覆盖不可微的 JPEG 不可压缩性、可压缩性，以及可微的类别概率奖励（Class 5）。在不可压缩性任务上，本文方法达到奖励 1.62，FID 55.9；预训练模型为奖励 1.30、FID 2.90（表 1）。论文指出 DPO 虽然得到奖励 1.36，但图像出现严重噪声和模糊伪影；SFT 和 RWR 的不可压缩性奖励反而低于预训练。在 Class 5 任务上，本文方法奖励从预训练的 -7.71 提升到 -0.87，但 FID 升至 67.70；DPO 奖励为 -8.18，几乎未改善（表 1）。此外，在表 2 的与多步 DDPM 方法对比中，本文在奖励约 1.40 水平下 FID 为 9.32，而 20 步 DDPM 上的 q-Learning 和 DPOK 分别为 89.79 和 138.89。不过论文明确说明两种方法的预训练模型 FID 不同，不能直接视为同一基线上的对比。

ImageNet 256×256 实验是更高维度的验证。四个奖励任务中，本文方法在不可压缩性、Black-and-White、CLIP-red 三个任务上取得最高奖励，并且在所有四个任务上的美学分数都高于 DPO（表 3）。在 FID 上，本文在 Black-and-White 和 CLIP-red 任务中比 DPO 高；论文作者推测，这是因为奖励加权目标本身离数据分布更远，强化奖励自然伴随 FID 上升，不能仅凭 FID 评估质量。

消融实验显示 β\beta 在 CIFAR-10 不可压缩性任务上从 1 到 20 单调提升奖励，但 FID 也从 5.43 增加到 56.96（表 2 部分结果）。在固定 ν\nu 与更新 νk\nu_k 的选择上，更新策略在相同 FID 下获得更高奖励，在相同奖励下 FID 更低；固定策略在约 3000 迭代时出现 FID 尖峰（图 5）。

实践建议

对于需要在实时或低成本推理环境中部署生成模型的团队，本文方法提供了一条相对实用的微调路径。它不要求奖励函数可微，因此可以围绕真实业务指标——例如压缩后文件大小、安全过滤分数、用户点击或偏好评分——直接构建奖励并微调一步生成器。

在工程实施时，有几点值得注意。第一，β\beta 在本文实验中被证明是奖励-保真度的关键旋钮：CIFAR-10 不可压缩性任务中，β\beta 越高奖励越高，FID 也越高（表 2）。实践者应基于自身场景做小范围扫描。第二，论文在消融中发现固定参考分布会导致 FID 尖峰，因此推荐在线更新参考分布 νk\nu_k。第三，本文框架尚未在 prompt 条件文本到图像生成中验证，若业务依赖类条件或文本条件生成，需要自行验证并可能需要扩展条件输入机制。

总体而言，该方法是奖励微调工具箱中的一个新选择，尤其适合不可微奖励、黑盒指标和一步推理场景。但其结论主要来自 CIFAR-10 和 ImageNet 类条件生成的实验设置，迁移到其他数据分布或更高分辨率任务时，仍需谨慎评估奖励上升与保真度下降之间的权衡。