一步生成模型的目标很明确:让模型从噪声到数据只经过一次前向计算,以降低生成过程的时间复杂度。但当用户希望生成结果满足某种奖励目标时,如何在不破坏模型分布的前提下完成微调,仍然缺少成熟方法。arXiv 近日发表的这项工作,将问题放到最优传输框架下,提出通过 Wasserstein Gradient Flow(WGF)进行奖励引导微调。

从奖励优化到分布演化

论文没有把一步生成器简单看成一个逐样本修正器,而是关注生成数据整体分布如何随奖励目标变化。WGF 被用于描述概率空间中的平滑、可控分布演化:模型不是直接把输出强行推向高奖励区域,而是沿着一条连续路径更新生成分布。

这一视角对应了一个实际问题。只追求奖励,可能会让模型过度利用奖励函数的漏洞,或者把样本集中到少数模式,出现奖励 hacking 和 mode collapse。论文提出的方法希望通过更平滑的分布更新,兼顾奖励对齐与生成分布的稳定变化。

不需要奖励梯度

方法的一个关键特点是训练过程不要求奖励函数可微。论文给出的实践方案既能处理不可微奖励,也能处理可微奖励,这使它可以覆盖一些难以直接反向传播的评价目标。摘要列出的实验包括二维合成数据、CIFAR-10,以及 ImageNet 256×256,并使用了 JPEG 可压缩性、类别概率、黑白属性和 CLIP 对齐等不同奖励。

从结果描述看,在这些实验设置中,该方法相较基线取得了更好的奖励对齐效果。这里的重点不只是“能不能提高奖励”,还在于用统一的训练方式处理不同类型的奖励函数。对于一步生成模型来说,这种兼容性可能比针对某个单一任务的技巧更有价值。

我的判断

这项工作的价值主要在于补上了一步生成模型奖励引导微调这一方向的方法空白,并提供了一个具有明确数学视角的训练框架。WGF 强调分布级、平滑的更新,对缓解奖励投机和模式坍塌提出了合理路径,但摘要并未给出具体提升幅度,也没有说明在更大规模任务或更多真实应用中的代价。

因此,现阶段更适合把它看作一种有潜力的训练方法,而不是一步生成模型奖励优化的通用答案。它尤其适合奖励函数不可微、需要比较稳定分布变化的场景;至于训练成本、超参数敏感性,以及对奖励设计本身的依赖,还需要结合论文完整实验进一步判断。