解决扩散模型强化学习中多维度奖励冲突问题,实现高效统一优化。
MARBLE: Multi-Aspect Reward Balance for Diffusion RL

- 通过梯度空间优化,为每项奖励独立计算优势值并融合更新方向。
- 在五维奖励下同时提升所有指标,使最差奖励的梯度方向从负转正。
- 无需手动调参,训练速度接近单奖励基线,适合多目标生成任务。
强化学习微调已成为对齐扩散模型与人类偏好的主流方法。然而,图像评估本质上是多维度任务,需同时优化多个评价标准。现有方法通常为每个奖励训练专用模型、使用加权和奖励 $R(x)=\ sum_k w_k R_k(x)$,或采用人工设计的分阶段微调流程。这些方法要么无法联合训练统一模型,要么需要大量手动调参的序列训练。我们发现,问题根源在于简单的加权和聚合方式:多数轨迹为特定奖励的专家样本,对其他奖励无信息量,导致加权求和稀释了监督信号。为此,我们提出 MARBLE(Multi-Aspect Reward BaLancE),一种梯度空间优化框架,为每项奖励维护独立的优势估计器,计算各奖励策略梯度,并通过求解二次规划问题将它们融合为单一更新方向,无需人工调参权重。此外,我们提出一种近似形式,利用 DiffusionNFT 中损失的仿射结构,将每步开销从 K+1 次反向传播降至接近单奖励基线水平,并引入 EMA 平滑平衡系数以稳定因单批次波动带来的更新。在 SD3.5 Medium 上使用五项奖励进行测试,MARBLE 同时提升全部五个奖励维度,使最差奖励的梯度余弦角度在 80% 的小批次中由负转正,且训练速度保持在基线的 0.97 倍。
原文摘要 · Abstract (English)
Reinforcement learning fine-tuning has become the dominant approach for aligning diffusion models with human preferences. However, assessing images is intrinsically a multi-dimensional task, and multiple evaluation criteria need to be optimized simultaneously. Existing practice deal with multiple rewards by training one specialist model per reward, optimizing a weighted-sum reward $R(x)=\sum_k w_k R_k(x)$, or sequentially fine-tuning with a hand-crafted stage schedule. These approaches either fail to produce a unified model that can be jointly trained on all rewards or necessitates heavy manually tuned sequential training. We find that the failure stems from using a naive weighted-sum reward aggregation. This approach suffers from a sample-level mismatch because most rollouts are specialist samples, highly informative for certain reward dimensions but irrelevant for others; consequently, weighted summation dilutes their supervision. To address this issue, we propose MARBLE (Multi-Aspect Reward BaLancE), a gradient-space optimization framework that maintains independent advantage estimators for each reward, computes per-reward policy gradients, and harmonizes them into a single update direction without manually-tuned reward weighting, by solving a Quadratic Programming problem. We further propose an amortized formulation that exploits the affine structure of the loss used in DiffusionNFT, to reduce the per-step cost from K+1 backward passes to near single-reward baseline cost, together with EMA smoothing on the balancing coefficients to stabilize updates against transient single-batch fluctuations. On SD3.5 Medium with five rewards, MARBLE improves all five reward dimensions simultaneously, turns the worst-aligned reward's gradient cosine from negative under weighted summation in 80% of mini-batches to consistently positive, and runs at 0.97X the training speed of baseline training.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。