MARBLE: Баланс многокомпонентного вознаграждения для RL с диффузией

MARBLE: Multi-Aspect Reward Balance for Diffusion RL

May 7, 2026

Авторы: Canyu Zhao, Hao Chen, Yunze Tong, Yu Qiao, Jiacheng Li, Chunhua Shen

cs.AI

Аннотация

Тонкая настройка с подкреплением стала доминирующим методом согласования диффузионных моделей с человеческими предпочтениями. Однако оценка изображений по своей природе является многомерной задачей, требующей одновременной оптимизации нескольких критериев. Существующие подходы к работе с множественными вознаграждениями включают обучение отдельной специализированной модели для каждого вознаграждения, оптимизацию взвешенной суммы R(x)=∑_k w_k R_k(x) или последовательную тонкую настройку по заранее заданному расписанию этапов. Эти методы либо не позволяют создать унифицированную модель, совместно обучаемую по всем вознаграждениям, либо требуют трудоемкой ручной настройки последовательного обучения. Мы обнаружили, что проблема кроется в использовании наивной агрегации вознаграждений методом взвешенной суммы. Этот подход страдает от несоответствия на уровне выборок: большинство траекторий являются специализированными — высокоинформативными для одних измерений вознаграждения, но нерелевантными для других; как следствие, взвешенное суммирование размывает их обучающий сигнал. Для решения этой проблемы мы предлагаем MARBLE (Multi-Aspect Reward BaLancE) — фреймворк оптимизации в пространстве градиентов, который поддерживает независимые оценки преимущества для каждого вознаграждения, вычисляет градиенты политики для каждого вознаграждения и гармонично объединяет их в единое направление обновления без ручной настройки весов вознаграждений путем решения задачи квадратичного программирования. Мы также предлагаем амортизированную формулировку, которая использует аффинную структуру функции потерь в DiffusionNFT для снижения стоимости на шаг с K+1 обратных проходов до стоимости, близкой к базовой с одним вознаграждением, в сочетании со сглаживанием по скользящему среднему коэффициентов балансировки для стабилизации обновлений против transient флуктуаций в отдельных батчах. На SD3.5 Medium с пятью вознаграждениями MARBLE улучшает все пять измерений одновременно, изменяет косинус градиента наихудше согласованного вознаграждения с отрицательного при взвешенном суммировании в 80% мини-батчей на стабильно положительный и работает со скоростью 0.97X от базового обучения.

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.