arXiv:2502.06061cs.LGcs.AI2025-02ICLR被引 42

用在线奖励加权与熵正则化,让生成模型高效对齐任意奖励函数。

Online Reward-Weighted Fine-Tuning of Flow Matching with Wasserstein Regularization

  • 在流匹配中引入在线奖励加权,无需奖励梯度或过滤数据集。
  • 通过W2正则化防止策略坍缩,在高奖励区与多样性间实现平衡。
  • 适合需要可控生成质量与多样性的图像/文本生成任务研究者。

近期强化学习在微调基于扩散的生成模型方面取得显著进展,但将连续流生成模型与任意用户定义的奖励函数对齐仍具挑战性,主要源于过优化导致的策略坍缩以及连续时间流中似然计算成本过高。本文提出一种简单易用且理论严谨的强化学习微调方法,称为在线奖励加权条件流匹配与Wasserstein-2正则化(ORW-CFM-W2)。该方法将强化学习融入流匹配框架,可在不依赖奖励梯度或过滤数据集的情况下,对任意奖励函数进行微调。通过引入在线奖励加权机制,引导模型聚焦数据流形中的高奖励区域。为防止策略坍缩并保持多样性,结合了Wasserstein-2(W2)距离正则化,并在流匹配中推导出其可计算上界,有效平衡了策略优化中的探索与利用。我们提供了理论分析,揭示方法的收敛性与诱导数据分布特性,建立了与带KL正则化的传统强化学习算法的联系,深化了对方法内在机制和学习行为的理解。大量实验涵盖目标图像生成、图像压缩与文本-图像对齐任务,结果表明该方法实现了最优策略收敛,同时支持奖励最大化与多样性保持之间的可控权衡。

原文摘要 · Abstract (English)

Recent advancements in reinforcement learning (RL) have achieved great success in fine-tuning diffusion-based generative models. However, fine-tuning continuous flow-based generative models to align with arbitrary user-defined reward functions remains challenging, particularly due to issues such as policy collapse from overoptimization and the prohibitively high computational cost of likelihoods in continuous-time flows. In this paper, we propose an easy-to-use and theoretically sound RL fine-tuning method, which we term Online Reward-Weighted Conditional Flow Matching with Wasserstein-2 Regularization (ORW-CFM-W2). Our method integrates RL into the flow matching framework to fine-tune generative models with arbitrary reward functions, without relying on gradients of rewards or filtered datasets. By introducing an online reward-weighting mechanism, our approach guides the model to prioritize high-reward regions in the data manifold. To prevent policy collapse and maintain diversity, we incorporate Wasserstein-2 (W2) distance regularization into our method and derive a tractable upper bound for it in flow matching, effectively balancing exploration and exploitation of policy optimization. We provide theoretical analyses to demonstrate the convergence properties and induced data distributions of our method, establishing connections with traditional RL algorithms featuring Kullback-Leibler (KL) regularization and offering a more comprehensive understanding of the underlying mechanisms and learning behavior of our approach. Extensive experiments on tasks including target image generation, image compression, and text-image alignment demonstrate the effectiveness of our method, where our method achieves optimal policy convergence while allowing controllable trade-offs between reward maximization and diversity preservation.

生成模型强化学习流匹配奖励对齐

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。