# MARBLE: Multi-Aspect Reward Balance for Diffusion RL

> MARBLE harmonizes per-reward policy gradients via quadratic programming, enabling a single diffusion model to improve all five reward dimensions simultaneously without hand-tuned weights.

- **Source:** [arXiv](https://arxiv.org/abs/2605.06507)
- **Published:** 2026-05-09
- **Permalink:** https://picx.dev/p/hnsXXL

## Summary

## 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)

| Method | GenEval | OCR | PickScore | CLIPScore | HPSv2.1 | Aesthetic | ImgRwd | UniRwd | Composite ↑ |
|--------|---------|-----|-----------|-----------|---------|-----------|--------|--------|-------------|
| SD3.5-M + CFG | 0.63 | 0.59 | 22.34 | 0.285 | 0.279 | 5.36 | 0.85 | 3.03 | -0.255 |
| + FlowGRPO | 0.95 | 0.66 | 22.51 | 0.293 | 0.274 | 5.32 | 1.06 | 3.18 | +0.120 |
| + DiffusionNFT † (sequential) | 0.94 | 0.91 | **23.80** | **0.293** | 0.331 | 6.01 | 1.49 | 3.49 | +1.015 |
| + DiffusionNFT ‡ (weighted-sum) | 0.92 | 0.91 | 21.53 | 0.267 | 0.300 | 6.15 | 1.16 | 3.04 | +0.184 |
| **+ MARBLE** | **0.94** | **0.96** | 22.83 | 0.286 | **0.355** | **6.59** | **1.53** | **3.52** | **+1.116** |

**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)

| Method | Relative speed | GPU memory |
|--------|---------------|------------|
| Weighted Sum (K=5, baseline) | 1.00× | 59G (1.00×) |
| **MARBLE w/ amortization (N=10)** | **0.97×** | 67G (1.14×) |
| MARBLE w/o amortization | 0.56× | 67G (1.14×) |

### Ablation Study (Table 3)

| Variant | GenEval | OCR | PickScore | CLIPScore | HPSv2.1 | Aesthetic | ImgRwd | UniRwd |
|---------|---------|-----|-----------|-----------|---------|-----------|--------|--------|
| **MARBLE (full, ρ=0.7)** | **0.93** | **0.96** | **22.62** | **0.283** | **0.355** | **6.59** | **1.52** | **3.45** |
| w/o gradient normalization | FAIL | | | | | | | |
| Fixed α=0.2 | 0.86 | 0.89 | 22.64 | 0.272 | 0.346 | 6.55 | 1.45 | 3.42 |
| Solve α every step | 0.92 | 0.92 | 21.32 | 0.267 | 0.301 | 5.89 | 1.17 | 3.04 |

**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)

| Statistic | Weighted sum | MARBLE | Δ |
|-----------|-------------|--------|---|
| Worst-reward alignment min$_k$ cos($d, g_k$) ↑ | −0.1346 | **+0.3721** | +0.5067 |
| Average alignment mean$_k$ cos($d, g_k$) ↑ | +0.4559 | +0.4014 | −0.0545 |
| Alignment imbalance var$_k$ cos($d, g_k$) ↓ | +0.1605 | **+0.0058** | −0.1548 |
| Conflict rate P(min$_k$ cos($d, g_k$) < 0) ↓ | 0.800 | **0.000** | −0.800 |

### Human Preference Study (Table S5)

| Method | Text-image alignment ↑ | Image quality ↑ |
|--------|----------------------|-----------------|
| DiffusionNFT ‡ | 3.60 | 2.79 |
| DiffusionNFT † | 4.26 | 3.58 |
| **MARBLE** | **4.63** | **4.41** |

---

## 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."

---

_Markdown view of https://picx.dev/p/hnsXXL, served by PicX — AI-generated visual whiteboard summaries of research papers._