MARBLE: Equilíbrio de Recompensa Multiaspecto para RL de Difusão

Resumo

O ajuste fino por aprendizagem por reforço tornou-se a abordagem dominante para alinhar modelos de difusão com preferências humanas. No entanto, a avaliação de imagens é intrinsecamente uma tarefa multidimensional, sendo necessário otimizar múltiplos critérios de avaliação simultaneamente. As práticas existentes lidam com múltiplas recompensas treinando um modelo especialista por recompensa, otimizando uma recompensa de soma ponderada R(x)=Σₖ wₖ Rₖ(x), ou realizando ajuste fino sequencial com um cronograma de estágios definido manualmente. Essas abordagens ou falham em produzir um modelo unificado que possa ser treinado conjuntamente em todas as recompensas, ou necessitam de um pesado treinamento sequencial com ajuste manual. Descobrimos que a falha decorre do uso de uma agregação de recompensas por soma ponderada ingênua. Essa abordagem sofre de um descompasso a nível de amostra porque a maioria das trajectórias são amostras de especialistas, altamente informativas para certas dimensões de recompensa mas irrelevantes para outras; consequentemente, a soma ponderada dilui sua supervisão. Para resolver este problema, propomos o MARBLE (Multi-Aspect Reward BaLancE), um framework de otimização no espaço do gradiente que mantém estimadores de vantagem independentes para cada recompensa, calcula gradientes de política por recompensa e os harmoniza numa única direção de atualização sem ponderação manual de recompensas, resolvendo um problema de Programação Quadrática. Propomos ainda uma formulação amortizada que explora a estrutura afim da perda usada no DiffusionNFT, para reduzir o custo por passo de K+1 retropropagações para um custo próximo ao da linha de base de recompensa única, juntamente com uma suavização por média móvel exponencial (EMA) nos coeficientes de balanceamento para estabilizar as atualizações contra flutuações transitórias de um único lote. No SD3.5 Medium com cinco recompensas, o MARBLE melhora todas as cinco dimensões de recompensa simultaneamente, transforma o cosseno do gradiente da recompensa menos alinhada de negativo sob a soma ponderada em 80% dos minilotes para consistentemente positivo, e executa a 0,97X da velocidade de treinamento da linha 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.