MARBLE: Equilibrio de Recompensa Multi-Aspecto para RL de Difusión

MARBLE: Multi-Aspect Reward Balance for Diffusion RL

May 7, 2026

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

cs.AI

Resumen

El ajuste fino por aprendizaje por refuerzo se ha convertido en el enfoque dominante para alinear los modelos de difusión con las preferencias humanas. Sin embargo, evaluar imágenes es intrínsecamente una tarea multidimensional que requiere optimizar múltiples criterios de evaluación simultáneamente. Las prácticas existentes manejan las múltiples recompensas entrenando un modelo especialista por recompensa, optimizando una recompensa de suma ponderada R(x)=Σₖ wₖ Rₖ(x), o realizando ajuste fino secuencial con una programación de etapas manual. Estos enfoques o bien no logran producir un modelo unificado que pueda entrenarse conjuntamente en todas las recompensas, o bien requieren un costoso entrenamiento secuencial con ajuste manual. Encontramos que el fallo proviene del uso de una agregación de recompensas por suma ponderada ingenua. Este enfoque sufre de un desajuste a nivel de muestra porque la mayoría de las trayectorias son muestras especializadas, muy informativas para ciertas dimensiones de recompensa pero irrelevantes para otras; en consecuencia, la suma ponderada diluye su supervisión. Para abordar este problema, proponemos MARBLE (Multi-Aspect Reward BaLancE), un marco de optimización en el espacio del gradiente que mantiene estimadores de ventaja independientes para cada recompensa, calcula gradientes de política por recompensa y los armoniza en una única dirección de actualización sin ponderación manual de recompensas, resolviendo un problema de Programación Cuadrática. Además, proponemos una formulación amortizada que explota la estructura afín de la pérdida utilizada en DiffusionNFT, para reducir el coste por paso de K+1 pasos hacia atrás a un coste cercano al de la línea base de recompensa única, junto con un suavizado por media móvil exponencial (EMA) en los coeficientes de balanceo para estabilizar las actualizaciones contra fluctuaciones transitorias de lotes individuales. En SD3.5 Medium con cinco recompensas, MARBLE mejora las cinco dimensiones de recompensa simultáneamente, convierte el coseno del gradiente de la recompensa peor alineada de negativo bajo suma ponderada en el 80% de los minilotes a consistentemente positivo, y se ejecuta a 0.97X la velocidad de entrenamiento de la línea 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.