MARBLE: 확산 RL을 위한 다중 측면 보상 균형

초록

강화 학습 미세 조정은 확산 모델을 인간의 선호도에 맞추는 데 있어 지배적인 접근법으로 자리 잡았습니다. 그러나 이미지 평가는 본질적으로 다차원적인 작업이며, 여러 평가 기준을 동시에 최적화해야 합니다. 기존 방법론은 보상마다 별도의 전문가 모델을 훈련하거나, 가중합 보상 R(x)=∑_k w_k R_k(x)을 최적화하거나, 수작업으로 설계된 단계별 일정에 따라 순차적으로 미세 조정하는 방식으로 다중 보상을 처리합니다. 이러한 접근법들은 모든 보상에 대해 공동으로 훈련될 수 있는 통합 모델을 생산하지 못하거나, 많은 수작업 조정이 필요한 순차적 훈련을 필요로 합니다. 우리는 이러한 실패의 근본 원인이 단순한 가중합 보상 집계 방식에 있다고 판단했습니다. 이 방식은 샘플 수준의 불일치 문제를 겪는데, 대부분의 롤아웃은 특정 보상 차원에서는 매우 유익하지만 다른 차원에서는 관련성이 낮은 전문가 샘플이기 때문입니다. 결과적으로 가중합 계산은 이들의 지도 신호를 희석시킵니다. 이 문제를 해결하기 위해 우리는 MARBLE(Multi-Aspect Reward BaLancE)을 제안합니다. MARBLE은 각 보상에 대해 독립적인 어드밴티지 추정기를 유지하고, 보상별 정책 그래디언트를 계산한 후, 2차 계획법 문제를 해결하여 수동 조정된 보상 가중치 없이도 이들을 단일 업데이트 방향으로 조화시키는 그래디언트 공간 최적화 프레임워크입니다. 또한 우리는 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