Summary (Overview)

• Problem: Multi-reward RL fine-tuning of diffusion models suffers from the specialist sample phenomenon — most rollouts are informative for only a subset of reward dimensions, and naive weighted-sum aggregation dilutes their supervision, leading to anti-aligned updates (negative worst-reward gradient cosine in 80% of mini-batches).
• Proposed Solution: MARBLE (Multi-Aspect Reward Balance), a gradient-space optimization framework that maintains independent advantage estimators per reward, computes per-reward policy gradients, and harmonizes them into a single update direction via Quadratic Programming (QP) — without manually-tuned reward weights.
• Key Innovation: An amortized formulation exploiting the affine structure of the DiffusionNFT loss reduces per-step cost from $K+1$ backward passes to near single-reward baseline cost ($0.97\times$ training speed), plus EMA smoothing on balancing coefficients for stability.
• Results: On SD3.5 Medium with five rewards, MARBLE improves all five reward dimensions simultaneously in a single model, achieves the highest Composite score across 8 metrics, and turns worst-reward gradient alignment from negative to consistently positive.
• First work to address reward balancing in multi-reward diffusion RL, bridging gradient harmonization techniques from multi-task learning into the diffusion RL setting.

---

Introduction and Theoretical Foundation

Background and Motivation

Reinforcement Learning (RL) fine-tuning has become the dominant paradigm for aligning diffusion models with human preferences. However, image quality is inherently multi-dimensional — a high-quality image must simultaneously exhibit:

• Aesthetic appeal
• Faithfulness to the text prompt
• Fine-grained correctness (accurate text rendering, coherent object placement)

Existing approaches fall into two categories, both with significant limitations:

1. One model per reward (e.g., FlowGRPO specialists): Requires maintaining multiple models and cannot generalize across reward dimensions.
2. Sequential multi-reward training (e.g., DiffusionNFT): Produces a single model but requires hand-crafted stage schedules (e.g., "800 iterations on reward 1, then 300 on reward 2...") and suffers from forgetting previously acquired rewards.

The Specialist Sample Phenomenon

The paper identifies a fundamental failure mode of scalar reward aggregation:

"Many rollouts are informative for only a part of reward dimensions and uninformative or even inapplicable for the rest. For example, an image of a cat carries no signal for OCR-related rewards, and a generation with strong text rendering may be only average aesthetically."

Under $R(x) = \sum_k w_k R_k(x)$, the value of such a sample is diluted by unrelated dimensions, and the resulting advantage no longer reflects the dimension on which the sample is genuinely useful. Figure 2 shows that high advantages are concentrated on source-specific rewards (OCR, GenEval), with few samples achieving positive rewards across all dimensions.

Theoretical Foundation: DiffusionNFT

The paper builds on DiffusionNFT, which optimizes:

$$\max_\theta \mathbb{E}_{x \sim \pi_\theta}[R(x)] - \beta_{KL} \cdot D_{KL}(\pi_\theta \| \pi_{ref})$$

The NFT loss for a sample $x$ with advantage $A(x)$ interpolates between positive and negative terms:

$$\ell(\theta; x, t) = r \cdot L^+(\theta) + (1-r) \cdot L^-(\theta)$$

where $r = \text{clamp}\left(\frac{1}{2} + \frac{A(x)}{2A_{max}}, 0, 1\right)$, $L^+(\theta) = \|v^+_\theta - v\|^2$, and $L^-(\theta) = \|v^-_\theta - v\|^2$.

Key structural property: $L^+$ and $L^-$ depend only on $\theta$ and the current sample, and are independent of the advantage value. The advantage affects the loss only through the affine mapping to $r$.

---

Methodology

1. Per-Reward Advantage Decomposition

MARBLE maintains an independent advantage estimator for each reward $R_k$, normalized within prompt groups:

$$A_k(x) = \frac{R_k(x) - \mu_k(\text{prompt})}{\sigma_k(\text{prompt}) + \varepsilon}$$

Each $A_k$ yields a separate interpolation coefficient $r_k \in [0,1]$, defining a reward-specific NFT loss $\ell_k$. The corresponding policy gradient is:

$$g_k = \nabla_\theta \frac{1}{NT} \sum_{i=1}^{N} \sum_{t=1}^{T} \ell_k(\theta; x_i, t)$$

2. Gradient Normalization and Harmonization

To remove scale disparities between rewards, each gradient is normalized:

$$\hat{g}_k = g_k / \|g_k\|$$

The unified update direction is computed by solving a convex Quadratic Program (QP):

$$\alpha^* = \arg\min_{\alpha \in \Delta_K} \left\| \sum_{k=1}^{K} \alpha_k \hat{g}_k \right\|^2$$

where $\Delta_K = \{\alpha \in \mathbb{R}^K_{\geq 0} : \sum_k \alpha_k = 1\}$ is the probability simplex. The solution $d^* = \sum_{k=1}^{K} \alpha^*_k \hat{g}_k$ is the minimum-norm point in the convex hull of normalized gradients.

3. Rescaling and KL-Decoupled Update

Since $d^*$ is computed from unit-normalized gradients, its magnitude is restored:

$$d_{final} = d^* \cdot \bar{n}, \quad \bar{n} = \frac{1}{K} \sum_{k=1}^{K} \|g_k\|$$

The final update combines the rescaled reward gradient with KL regularization:

$$\theta \leftarrow \theta - \eta \left( d_{final} + \beta_{KL} \cdot \nabla_\theta D_{KL}(\pi_\theta \| \pi_{ref}) \right)$$

4. Amortized Gradient Harmonization

Proposition 1 (Scalarization Equivalence): Let $\alpha \in \Delta_K$ and let $A_1, \ldots, A_K$ be per-reward advantages with $|A_k| < A_{max}$ for all $k$ and $\sum_k \alpha_k A_k < A_{max}$. Define the combined advantage $\bar{A} = \sum_{k=1}^{K} \alpha_k A_k$. Then:

$$\nabla_\theta \ell(\theta; \bar{A}) = \sum_{k=1}^{K} \alpha_k \nabla_\theta \ell_k(\theta; A_k)$$

Proof sketch: Since $r_k = \frac{A_k}{2A_{max}} + \frac{1}{2}$, we have $\sum_k \alpha_k r_k = \frac{\bar{A}}{2A_{max}} + \frac{1}{2} = \bar{r}$, and the simplex constraint $\sum_k \alpha_k = 1$ preserves the constant offset.

Amortized procedure: Every $N$ steps, run full harmonization to refresh $\alpha^*$ from normalized gradients. During the intervening $N-1$ steps, form $\bar{A} = \sum_k \alpha^*_k A_k$ using cached coefficients and perform only one reward backward pass. This reduces average per-step cost from $(K+1)\times$ to $(K+N)/N\times$ that of a single-reward baseline.

5. Coefficient Smoothing (EMA)

To prevent transient single-batch fluctuations from suppressing a reward for an entire amortization window:

$$\bar{\alpha}_t = \rho \bar{\alpha}_{t-1} + (1-\rho) \alpha^*_t$$

where $\rho = 0.7$ is the EMA decay. Since both $\bar{\alpha}_{t-1}$ and $\alpha^*_t$ lie on the probability simplex, their convex combination remains a valid simplex vector.

---

Empirical Validation / Results

Experimental Setup

• Base model: Stable Diffusion 3.5 Medium, fine-tuned with LoRA adapters (rank 32, alpha 64)
• Optimizer: AdamW, learning rate $3 \times 10^{-4}$
• Five training rewards: PickScore, HPSv2, CLIPScore (general), OCR accuracy, GenEval (specialist)
• Four held-out metrics: Aesthetic Score, ImageReward, UniReward (for transfer assessment)
• Hardware: 16 NVIDIA H200 GPUs

Main Results (Table 1)

Key findings:

• MARBLE improves all five training rewards simultaneously in a single model
• MARBLE ranks first on all four held-out quality metrics (HPSv2.1, Aesthetic, ImageReward, UniReward)
• Weighted-sum baseline (‡) performs substantially worse on specialist objectives
• Sequential training (†) requires extensive manual scheduling and suffers from forgetting

Training Efficiency (Table 2)

Ablation Study (Table 3)

Key ablation findings:

• Without gradient normalization: Optimization fails entirely (degenerate/numerically unstable coefficients)
• Fixed uniform α: Imbalanced convergence — general rewards improve quickly, specialist objectives remain under-optimized
• Solving α every step: Higher cost and high-frequency variation in update direction, negatively affecting stability

Update-Direction Harmony Diagnostics (Table S1)

Human Preference Study (Table S5)

---

Theoretical and Practical Implications

Theoretical Contributions

1. Formal characterization of the specialist sample problem: The paper provides quantitative evidence (80% of mini-batches have anti-aligned worst-reward gradients under weighted-sum) for why scalar reward aggregation fails when reward signals are sample-sparse.

2. Scalarization equivalence (Proposition 1): A novel theoretical result showing that convex combinations of per-reward NFT gradients can be recovered exactly by a single backward pass using a combined advantage — enabling efficient amortized training.

3. Gradient-space vs. loss-space balancing: The paper demonstrates that resolving reward interactions in gradient space (rather than loss space) preserves reward-specific supervision and avoids dilution.

Practical Implications

1. Eliminates manual tuning: MARBLE requires no hand-crafted stage schedules, reward orderings, or per-stage hyperparameters — a significant practical advantage over sequential training.

2. Near-baseline training cost: The amortized formulation achieves 0.97× training speed of the weighted-sum baseline, making gradient-space balancing practical at scale.

3. Single unified model: MARBLE produces one model that excels across all reward dimensions, avoiding the need to maintain multiple specialist models.

4. Adaptive coefficient interpretation: Learned coefficients reflect optimization difficulty rather than reward values — easier rewards (e.g., HPSv2) receive coefficients below the uniform baseline (0.2), while harder specialist rewards (e.g., GenEval) receive larger coefficients (~0.3).

---

Conclusion

Main Takeaways

MARBLE is the first multi-reward balancing method for diffusion model RL fine-tuning. It addresses the fundamental limitation of scalar reward aggregation by:

1. Preserving reward-specific supervision through per-reward advantage decomposition
2. Harmonizing gradients via QP-based minimum-norm solution with normalization
3. Amortizing computation via the affine structure of the NFT loss (Proposition 1)
4. Stabilizing coefficients via EMA smoothing

The method simultaneously improves all five reward dimensions in a single model, matches or exceeds both sequential and weighted-sum baselines, and achieves near-baseline training speed.

Limitations and Future Directions

1. Extension to video generation: Video models require optimizing heterogeneous objectives (temporal consistency, motion realism, physical plausibility), making reward balancing even more critical.

2. Scaling to larger reward sets: More rewards create tighter optimization and efficiency challenges.

3. Generative world models: Require not only high-quality generation but also coherent dynamics and plausible long-horizon evolution — a promising application area for MARBLE.

"We believe that MARBLE provides an important step toward scalable multi-objective alignment for future generative models."