MARBLE : Équilibrage de Récompense Multi-Aspect pour l'Apprentissage par Renforcement avec Diffusion

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

Résumé

Le fine-tuning par apprentissage par renforcement est devenu l'approche dominante pour aligner les modèles de diffusion sur les préférences humaines. Cependant, l'évaluation d'images est intrinsèquement une tâche multidimensionnelle, nécessitant l'optimisation simultanée de multiples critères d'évaluation. Les pratiques existantes traitent les récompenses multiples en entraînant un modèle spécialiste par récompense, en optimisant une récompense par somme pondérée R(x)=∑k wk Rk(x), ou en effectuant un fine-tuning séquentiel avec un calendrier d'étapes défini manuellement. Ces approches échouent soit à produire un modèle unifié pouvant être entraîné conjointement sur toutes les récompenses, soit nécessitent un entraînement séquentiel intensif et manuellement réglé. Nous constatons que cet échec provient de l'utilisation d'une agrégation naïve des récompenses par somme pondérée. Cette approche souffre d'un décalage au niveau des échantillons car la plupart des épisodes sont des échantillons spécialistes, très informatifs pour certaines dimensions de récompense mais non pertinents pour d'autres ; par conséquent, la sommation pondérée dilue leur supervision. Pour résoudre ce problème, nous proposons MARBLE (Multi-Aspect Reward BaLancE), un cadre d'optimisation dans l'espace des gradients qui maintient des estimateurs d'avantage indépendants pour chaque récompense, calcule les gradients de politique par récompense, et les harmonise en une seule direction de mise à jour sans pondération manuelle des récompenses, en résolvant un problème de programmation quadratique. Nous proposons en outre une formulation amortie qui exploite la structure affine de la perte utilisée dans DiffusionNFT, pour réduire le coût par étape de K+1 passes arrière à un coût proche de celui de base à récompense unique, couplée à un lissage EMA sur les coefficients d'équilibrage pour stabiliser les mises à jour contre les fluctuations transitoires par lot. Sur SD3.5 Medium avec cinq récompenses, MARBLE améliore simultanément les cinq dimensions de récompense, transforme le cosinus du gradient de la récompense la moins alignée (négatif sous somme pondérée dans 80 % des mini-lots) en une valeur constamment positive, et s'exécute à une vitesse d'entraînement équivalente à 0,97 fois celle de l'entraînement de base.

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.