MARBLE: Multi-Aspect Reward Balance for Diffusion RL
:::info Stub — Full Engineering Breakdown Coming This paper has a linked code implementation and was featured on Hugging Face Papers with 35 upvotes. A full breakdown with production viability rating, implementation notes, and honest limitations is being written. Subscribe to AI Letters → :::
| Authors | Canyu Zhao et al. |
| Year | 2026 |
| HF Upvotes | 35 |
| arXiv | 2605.06507 |
| Download | |
| Code | https://github.com/aim-uofa/MARBLE |
Abstract
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.
Engineering Breakdown
Plain English
This paper addresses a fundamental problem in diffusion model fine-tuning: when you need to optimize multiple image quality criteria simultaneously (e.g., aesthetic quality, prompt alignment, safety), naive approaches like weighted-sum rewards fail because different reward models disagree on which samples are good. MARBLE proposes a multi-aspect reward balancing technique that enables a single unified model to be trained jointly on all rewards without the heavy manual tuning required by existing sequential fine-tuning pipelines.
Key Engineering Insight
The core insight is that sample-level mismatch breaks weighted-sum aggregation—most generated samples are 'good' according to only one specialist reward model, not the aggregate, causing gradient conflicts during training. By balancing rewards at the distribution level rather than the sample level, you can align gradient directions across multiple objectives.
Why It Matters for Engineers
In production, you almost always have multiple competing objectives (visual quality, alignment with user intent, safety, diversity). Current approaches force you to either maintain separate models per objective, do expensive manual tuning of sequential stages, or accept degraded output quality. MARBLE lets you train one model end-to-end that actually respects all constraints, which directly reduces complexity and inference costs in real systems.
Research Context
Prior work on diffusion model alignment relied on single-reward RL or hand-crafted multi-stage pipelines that don't scale. This paper advances the state by identifying why naive reward combination fails mathematically and proposing a principled solution. It opens the door to treating multi-objective diffusion fine-tuning as a genuine optimization problem rather than an engineering workaround, similar to how Pareto optimization advanced multi-objective ML in other domains.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::