MARBLE: Solving the Multi-Reward Alignment Tug-of-War in Diffusion RL

TL;DR

Aligning diffusion models with human preferences is rarely a single-objective task; we want images to be aesthetic, text-accurate, and prompt-faithful all at once. MARBLE (Multi-Aspect Reward BaLancE) is a new framework that moves beyond the flawed "weighted-sum" approach. By treating rewards as independent gradient vectors and harmonizing them via Quadratic Programming, MARBLE creates a single model that excels across all dimensions without the need for manual weight tuning or complex sequential schedules.

The Problem: The "Specialist Sample" Dilution

Current practices in Diffusion RL fine-tuning (like FlowGRPO or DiffusionNFT) often collapse multiple rewards into a single scalar: $R(x)=\sum w_k{R}_k(x)$.

The authors identify a critical failure mode here: the Specialist Sample phenomenon. In any given batch, an image might be a "specialist"—for example, it contains perfect text rendering (high OCR reward) but has mediocre lighting (average aesthetic reward). When these rewards are summed, the specific, high-quality signal from the OCR dimension is "diluted" by the noisy or irrelevant signals from other dimensions.

Mathematically, this leads to gradient interference. The authors found that in 80% of training mini-batches, the standard weighted-sum update direction actually opposes at least one of the individual reward gradients.

Methodology: Harmonizing the Gradients

MARBLE rejects the scalar collapse in favor of Gradient-Space Harmonization.

1. Per-Reward Advantage Decomposition

Instead of one advantage, MARBLE calculates $K$ independent advantages. Each sample is credited precisely for the dimensions where it actually succeeds.

2. The Harmonization Solver

MARBLE treats the update as a Multi-Objective Optimization problem. It finds a unified direction $d^{\ast }$ by solving: ${\alpha }^{\ast }=\arg {\min }_{\alpha \in {\Delta }^K}{\Vert {\sum }_{k=1}^K{\alpha }_k{\hat{g}}_k\Vert }^2$ This finds the "minimum-norm point" in the convex hull of normalized gradients, ensuring the update is a fair compromise that doesn't sacrifice one reward for another.

Figure 1: The MARBLE workflow: Independent reward scoring → Per-reward gradient computation → QP Harmonization → Unified Parameter Update.

3. Efficiency via Amortization

Computing $K$ gradients every step is slow. MARBLE uses an Amortized Formulation. By exploiting the affine nature of the NFT loss, the model only performs a full harmonization every $N$ steps. In between, it reuses the learned importance coefficients ($\alpha$), maintaining 97% of the speed of a standard single-reward baseline.

Experimental Battlefront: SD3.5-Medium

The researchers tested MARBLE against specialist models and sequential fine-tuning using five rewards: PickScore, HPSv2, CLIPScore, OCR, and GenEval.

Key Results:

• Simultaneous Mastery: Unlike "Specialist" models that peak in one area and fail others, MARBLE improved all five dimensions concurrently.
• Zero Manual Tuning: Unlike DiffusionNFT† (sequential), which requires a "hand-crafted stage schedule" (e.g., 800 steps of A, then 300 of B), MARBLE is fully automated.
• Pareto Superiority: MARBLE achieved the highest "Composite Score" (mean z-score across 8 metrics), proving it is a better generalist.

Table 1: Quantitative comparison showing MARBLE's superior balance across rule-based and model-based metrics.

Why It Works: Dynamic Adaptation

A fascinating insight from the ablation studies is that MARBLE’s coefficients ($\alpha$) actually track optimization difficulty. "Easy" rewards (like basic aesthetics) receive lower weights over time, while "hard" specialist rewards (like GenEval's spatial constraints) receive higher weights, ensuring the model focuses its "learning capacity" where it's needed most.

Critical Analysis & Future Outlook

Pros:

• Eliminates the "gradient tug-of-war" inherent in multi-reward RL.
• Computationally efficient enough for production-scale training.

Limitations:

• Currently validated primarily on image generation.
• Scaling to 10+ rewards might introduce new stability challenges in the QP solver.

The Takeaway: MARBLE is a significant step toward "World Models" and video diffusion, where the number of rewards (temporal consistency, physics, realism) will only grow. It proves that we don't need better weights; we need better gradient geometry.