MARBLE: Multi-Aspect Beloningsbalans voor Diffusion RL

MARBLE: Multi-Aspect Reward Balance for Diffusion RL

May 7, 2026

Auteurs: Canyu Zhao, Hao Chen, Yunze Tong, Yu Qiao, Jiacheng Li, Chunhua Shen

cs.AI

Samenvatting

Reinforcement learning fine-tuning is de dominante aanpak geworden voor het afstemmen van diffusiomodellen op menselijke voorkeuren. Het beoordelen van afbeeldingen is echter van nature een multidimensionale taak, waarbij meerdere evaluatiecriteria gelijktijdig geoptimaliseerd moeten worden. Bestaande praktijken gaan om met meerdere beloningen door het trainen van één specialistisch model per beloning, het optimaliseren van een gewogen-som beloning R(x)=∑k wk Rk(x), of sequentiële fine-tuning met een handmatig vastgesteld faseschema. Deze benaderingen produceren ofwel geen verenigd model dat gezamenlijk op alle beloningen getraind kan worden, of vereisen zwaar handmatig afgestelde sequentiële training. Wij constateren dat het falen voortkomt uit het gebruik van een naïeve gewogen-som aggregatie van beloningen. Deze aanpak lijdt onder een mismatch op monsterniveau omdat de meeste rollouts specialistische monsters zijn, zeer informatief voor bepaalde beloningsdimensies maar irrelevant voor andere; bijgevolg verdunt gewogen sommatie hun supervisie. Om dit probleem aan te pakken, stellen we MARBLE (Multi-Aspect Reward BaLancE) voor, een optimalisatiekader in de gradïentruimte dat onafhankelijke advantage-schatters voor elke beloning handhaaft, beleidsgradïenten per beloning berekent en deze harmoniseert tot één enkele update-richting zonder handmatig afgestelde beloningsweging, door een kwadratisch programmeerprobleem op te lossen. We introduceren verder een geamortiseerde formulering die gebruikmaakt van de affiene structuur van het verlies gebruikt in DiffusionNFT, om de kosten per stap te reduceren van K+1 backward passes naar bijna de kosten van een enkele-beloningsbaseline, samen met EMA-afvlakking van de balanceringscoëfficiënten om updates te stabiliseren tegen tijdelijke fluctuaties per batch. Op SD3.5 Medium met vijf beloningen verbetert MARBLE alle vijf beloningsdimensies gelijktijdig, verandert de gradïentcosinus van de slechtst afgestemde beloning van negatief onder gewogen sommatie in 80% van de mini-batches naar consistent positief, en draait op 0.97X de trainingssnelheid van de baseline-training.

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.