MARBLE: 拡散RLのためのマルチアスペクト報酬バランス

要旨

強化学習によるファインチューニングは、拡散モデルを人間の選好に合わせるための主流手法となっている。しかし、画像評価は本質的に多次元的な課題であり、複数の評価基準を同時に最適化する必要がある。既存手法では、複数の報酬を扱う際に、報酬ごとに特化モデルを訓練する、重み付き和報酬 R(x)=Σ_k w_k R_k(x) を最適化する、または手作りの段階的スケジュールで順次ファインチューニングするといった方法が取られてきた。これらのアプローチは、すべての報酬で共同訓練可能な統一モデルを生成できないか、手動調整が煩雑な逐次訓練を必要とする。我々は、この問題の原因が単純な重み付き和報酬の集約にあることを明らかにした。この手法はサンプルレベルのミスマッチに悩まされる。なぜなら大半のロールアウトが特定の報酬次元には情報豊富だが他の次元には無関係な特化サンプルであり、重み付き和がこれらの監督信号を希釈してしまうためである。この問題を解決するため、我々はMARBLE（Multi-Aspect Reward BaLancE）を提案する。これは勾配空間最適化フレームワークであり、各報酬に対して独立したアドバンテージ推定量を維持し、報酬ごとの方策勾配を計算した後、二次計画問題を解くことで手動調整なしで単一の更新方向に調和させる。さらに、DiffusionNFTで用いられる損失関数のアフィン構造を利用した償却定式化を提案し、ステップ当たりの計算コストをK+1回の逆伝播から単一報酬ベースラインに近いコストに削減するとともに、平衡係数にEMA平滑化を適用して一時的なバッチ変動による更新の不安定性を抑制する。5つの報酬を用いたSD3.5 Mediumでの実験では、MARBLEは5つの報酬次元すべてを同時に改善し、最悪の調整報酬における勾配コサイン類似度を、重み付き和ではミニバッチの80%で負だった値から一貫して正に転換させ、ベースライン訓練の0.97倍の速度で動作することを実証した。

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.

MARBLE: 拡散RLのためのマルチアスペクト報酬バランス

MARBLE: Multi-Aspect Reward Balance for Diffusion RL

要旨

Support