MARBLE: Multi-Aspect Reward Balance for Diffusion RL

Zusammenfassung

Reinforcement Learning Fine-Tuning hat sich als vorherrschender Ansatz zur Ausrichtung von Diffusionsmodellen an menschlichen Präferenzen etabliert. Die Bewertung von Bildern ist jedoch von Natur aus eine mehrdimensionale Aufgabe, bei der mehrere Bewertungskriterien gleichzeitig optimiert werden müssen. Bisherige Verfahren behandeln mehrere Belohnungen entweder durch das Training eines spezialisierten Modells pro Belohnung, durch die Optimierung einer gewichteten Summenbelohnung R(x)=∑ₖ wₖ Rₖ(x) oder durch sequenzielles Fine-Tuning mit einem manuell erstellten Stufenplan. Diese Ansätze führen entweder nicht zu einem vereinheitlichten Modell, das für alle Belohnungen gemeinsam trainiert werden kann, oder erfordern einen aufwändigen, manuell abgestimmten sequenziellen Trainingsprozess. Wir stellen fest, dass das Scheitern auf der Verwendung einer naiven gewichteten Summenbelohnung zur Aggregation beruht. Dieser Ansatz leidet unter einem Stichproben-Mismatch, da die meisten Rollouts spezialisierte Stichproben sind – hochgradig informativ für bestimmte Belohnungsdimensionen, aber irrelevant für andere; folglich verdünnt die gewichtete Summierung deren Aufsichtsinformation. Um dieses Problem zu adressieren, schlagen wir MARBLE (Multi-Aspect Reward BaLancE) vor, ein Optimierungsframework im Gradientenraum, das unabhängige Advantage-Schätzer für jede Belohnung verwaltet, policy-Gradienten pro Belohnung berechnet und diese durch Lösen eines quadratischen Optimierungsproblems zu einer einzigen Aktualisierungsrichtung harmonisiert, ohne manuell abgestimmte Belohnungsgewichtung. Wir schlagen weiterhin eine amortisierte Formulierung vor, die die affine Struktur des in DiffusionNFT verwendeten Verlusts ausnutzt, um die Kosten pro Schritt von K+1 Backward-Passes auf nahezu die Kosten eines Single-Reward-Baselines zu reduzieren, kombiniert mit EMA-Glättung der Ausgleichskoeffizienten, um die Aktualisierungen gegen vorübergehende Single-Batch-Schwankungen zu stabilisieren. Bei SD3.5 Medium mit fünf Belohnungen verbessert MARBLE alle fünf Belohnungsdimensionen gleichzeitig, wandelt den Kosinus des Gradienten der am schlechtesten ausgerichteten Belohnung von negativ unter gewichteter Summierung in 80 % der Mini-Batches zu konsistent positiv und läuft mit 0,97X der Trainingsgeschwindigkeit des Baseline-Trainings.

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.