论文
MARBLE:扩散 RL 的多方面奖励平衡
MARBLE: Multi-Aspect Reward Balance for Diffusion RL
摘要
强化学习 微调 已成为使扩散模型与人类偏好保持一致的主要方法。然而,评估图像本质上是一个多维度的任务,需要同时优化多个评估标准。现有的实践通过为每个奖励训练一个专家模型、优化加权和奖励 $R(x)=\sum_k w_k R_k(x)$ 或使用手工制作的阶段时间表依次 微调 来处理多种奖励。这些方法要么无法产生可以对所有奖励进行联合训练的统一模型,要么需要大量手动调整的顺序训练。我们发现失败源于使用朴素的加权和奖励聚合。这种方法存在样本级别不匹配的问题,因为大多数采样轨迹都是专业样本,对于某些奖励维度来说信息量很大,但与其他维度无关;因此,加权求和削弱了他们的监督。为了解决这个问题,我们提出了 MARBLE(多方面奖励 BaLancE),这是一种梯度空间优化框架,它为每个奖励维护独立的优势估计器,计算每个奖励的策略梯度,并通过解决二次规划问题将它们协调为单个更新方向,而无需手动调整奖励权重。我们进一步提出了一种摊销公式,利用 DiffusionNFT 中使用的损失的仿射结构,将每步成本从 K+1 向后传递降低到接近单奖励基线成本,同时对平衡系数进行 EMA 平滑,以稳定更新以应对瞬态单批次波动。在具有五个奖励的 SD3.5 Medium 上,MARBLE 同时改进了所有五个奖励维度,将 80% 小批量中最差对齐奖励的梯度余弦从加权求和下的负值变为一致的正值,并以基线训练的 0.97 倍的训练速度运行。