PaperlayerView exact revision on arXiv

Source-linked AI summary

MARBLE: Multi-Aspect Reward Balance for Diffusion RL

Canyu Zhao, Hao Chen, Yunze Tong, Yu Qiao, Jiacheng Li, Chunhua Shen

arXiv:2605.06507v1cs.CVcs.LG

TL;DR

Multi-reward diffusion RL lacks an effective way to jointly optimize diverse image-quality criteria without separate models or laborious sequential schedules. MARBLE balances per-reward policy gradients in gradient space, improving all five rewards simultaneously while keeping training at 0.97X baseline speed.

  • Problem

    Existing diffusion RL methods lack a unified approach for jointly optimizing multiple image-quality rewards without separate specialists or manually tuned sequential fine-tuning.

  • Method

    MARBLE decomposes rewards into independent advantages, harmonizes their policy gradients through quadratic programming, and amortizes balancing with affine-loss structure and EMA smoothing.

  • Results

    MARBLE simultaneously improves all five rewards with one model, while weighted-sum aggregation is anti-aligned with at least one reward gradient in 80% of mini-batches.

  • Takeaways & Limitations

    MARBLE provides a single-model strategy for balancing multiple reward dimensions in diffusion RL while keeping amortized training cost near the single-reward baseline.

  • Takeaways & Limitations

    The full harmonization procedure requires K+1 backward passes and can suffer from single-mini-batch coefficient fluctuations that produce later visual artifacts.

Abstract

from arXiv · show

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.

1 Zhejiang University 2 HiThink 3 Zhejiang University of Technology · 1 Introduction

The introduction identifies specialist-sample mismatch as the cause of naive weighted-sum reward aggregation failing in multi-reward diffusion RL. MARBLE addresses this through per-reward gradient balancing, amortized computation, and EMA smoothing, improving all rewards with one model.

  • 1 Introduction: Diffusion RL alignment is inherently multi-dimensional, requiring simultaneous optimization of aesthetics, text faithfulness, and fine-grained correctness.
  • 1 Zhejiang University 2 HiThink 3 Zhejiang University of Technology: The work was conducted during an internship at HiThink.
  • 1 Introduction: Sequential multi-reward training can forget earlier rewards and requires substantial manual tuning of handcrafted stage schedules.
  • 1 Introduction: Naive weighted-sum aggregation fails because specialist samples are informative for only some reward dimensions, diluting relevant supervision with unrelated signals.
  • 1 Introduction: MARBLE maintains independent per-reward advantages, computes normalized per-reward policy gradients, and harmonizes them into one update direction.
  • 1 Introduction: 80% of mini-batches have a weighted-sum update direction anti-aligned with at least one reward gradient on SD3.5 Medium.
  • 1 Introduction: MARBLE reduces multi-reward training cost to near a single-reward baseline through DiffusionNFT loss amortization and stabilizes balancing weights with EMA smoothing.
  • 1 Introduction: MARBLE simultaneously improves all rewards with a single model, providing a foundation for scalable multi-objective alignment.

2 Related Work

Diffusion models evolved toward efficient, high-quality generation through latent representations, scaling, rectified flow, and transformer architectures. Related work covers diffusion RL preference alignment and gradient-level methods for multi-objective interference, while identifying limited diffusion-RL research on multi-reward alignment.

  • Diffusion Models: Diffusion models became dominant for high-fidelity image generation, with latent diffusion enabling efficient high-resolution synthesis and later scaling combining rectified flow with transformers.The passage also notes that diffusion models have expanded beyond text-to-image generation.
  • Diffusion RL: Diffusion RL aligns models with human preferences by optimizing expected reward while keeping the diffusion policy close to a pretrained reference model.Reward models evaluate generated samples, and early work primarily used policy-gradient methods.
  • Gradient Harmonization: Multi-task learning addresses inter-task gradient interference through strategies such as minimum-norm convex-hull gradients, destructive-component projection, and worst-case-oriented optimization.Interference can make one update improve some objectives while harming others.
  • Multi-Reward Alignment: Multi-reward alignment has grown in large language models, but corresponding diffusion-RL research remains limited; MARBLE adapts gradient harmonization with per-reward advantages and scale-aware balancing.Its design is tailored to the diffusion training objective.
  • MARBLE: MARBLE computes separate per-reward policy gradients and finds a common descent direction that balances reward objectives before updating the shared model.Given a prompt batch, independently scored generated images provide the inputs for this harmonization process.

3 Method

MARBLE preserves reward-specific supervision by computing independent per-reward gradients and harmonizing their normalized directions through a convex quadratic program. An amortized implementation reduces backward-pass cost, while EMA smoothing stabilizes balancing coefficients against rollout fluctuations.

  • Motivation: Scalar reward aggregation obscures which reward dimensions each sample informs, producing poorly aligned multi-reward updates.Weighted-sum updates have negative worst-reward alignment in 80% of measured mini-batches, while MARBLE remains positive in all measured mini-batches.
  • Per-reward advantage decomposition: MARBLE maintains independent prompt-normalized advantage estimators and reward-specific NFT losses, so each sample receives dimension-specific credit.All K gradients use the same sampled batch and differ only in their advantage signals.
  • Gradient normalization and harmonization: MARBLE normalizes per-reward gradients and solves a convex quadratic program whose minimum-norm convex-hull direction balances reward objectives.The solver concentrates on shared directions when rewards align and adaptively reweights conflicting aspects according to the current batch.
  • Efficiency and stability: The full procedure requires K+1 backward passes and can produce fluctuating balancing weights, motivating amortization and coefficient smoothing.Single-batch coefficient fluctuations can cause undesirable visual artifacts at later training stages, even when average reward scores remain strong.
  • Amortized procedure: Amortized updates refresh balancing coefficients every N steps and use one reward backward pass between refreshes, reducing average cost from (K+1)× to (K + N)/N times a single-reward baseline.The cached coefficients are combined in advantage space, preserving the current NFT loss’s natural gradient scale.
  • Coefficient smoothing: EMA smoothing prevents rollout failures from eliminating reward signals across an amortization window while retaining adaptation to estimated gradient geometry.The experiments use EMA decay ρ = 0.7, and coefficient smoothing improves training efficiency and effectiveness.

4 Experiments

Experiments on Stable Diffusion 3.5 Medium evaluate MARBLE across five jointly optimized rewards and compare it with specialist, weighted-sum, and sequential baselines. MARBLE improves all five training rewards in one model while preserving near-baseline training speed through amortization.

  • Experimental setup: MARBLE jointly optimizes five rewards—PickScore, HPSv2, CLIPScore, OCR accuracy, and GenEval—using LoRA adapters on Stable Diffusion 3.5 Medium.The adapters use rank 32, alpha 64, the NFT loss, and AdamW with a constant learning rate of 3 × 10−4.
  • Performance: MARBLE improves all five training rewards within a single model, whereas single-reward FlowGRPO specialists transfer poorly across objectives.Qualitative comparisons also show MARBLE satisfying diverse reward dimensions while the weighted-sum baseline fails to do so.
  • Performance: DiffusionNFT† reaches comparable quality but requires a manually scheduled multi-stage curriculum whose reward ordering, stage lengths, and datasets are hyperparameter-sensitive.Adding a new reward can also require revisiting previously trained rewards to mitigate forgetting.
  • Performance: MARBLE matches or surpasses DiffusionNFT† on every reward except PickScore and CLIPScore, and ranks first on four held-out quality metrics.The held-out metrics are HPSv2.1, Aesthetic, ImageReward, and UniReward; MARBLE also attains the highest Composite score across eight metrics.
  • Training efficiency: 0.97× relative speed is achieved by MARBLE’s amortized variant, compared with 0.56× for full per-reward harmonization, with both normalized to the weighted-sum baseline on 8×H200.Full harmonization incurs overhead from multiple reward-specific backward passes, while amortization approaches baseline speed.
  • Ablations: Fixed uniform coefficients αk = 0.2, unnormalized gradients, and per-step coefficient solving respectively cause imbalanced reward convergence, scale-dominated or unstable coefficients, and higher cost with fluctuating updates.MARBLE instead dynamically adjusts coefficients, normalizes gradients before solving the QP, and amortizes coefficient estimation.

5 Conclusion · A Appendix Overview

MARBLE introduces multi-reward balancing for diffusion RL, preserving reward-specific supervision while keeping training cost near baseline, though current validation is primarily on image generation. The appendix documents qualitative comparisons, analyses, ablations, implementation details, and future directions.

  • 5 Conclusion: MARBLE is presented as the first multi-reward balancing method for diffusion model RL fine-tuning.
  • 5 Conclusion: Per-reward advantage decomposition and gradient harmonization preserve reward-specific supervision and avoid specialist-sample dilution from weighted-sum aggregation.
  • 5 Conclusion: MARBLE’s amortized formulation keeps training cost close to the baseline.
  • 5 Conclusion: The current study primarily validates MARBLE on image generation, while video diffusion and generative world models remain important extensions.
  • A Appendix Overview: The appendix provides qualitative examples, paper-level takeaways, extended ablations, implementation details, and future directions supporting MARBLE.
  • A Appendix Overview: Additional qualitative comparisons examine text rendering, attribute and spatial understanding, counting, and coherent visual quality in a single MARBLE model.
  • A Appendix Overview: Extended analyses cover scalar reward aggregation, learned coefficients, update-direction harmony, training dynamics, coefficient adaptation, amortization intervals, and EMA smoothing.
  • A Appendix Overview: The appendix also compares heuristic balancing strategies, studies human preferences and metric correlations, and details distributed-training extraction, synchronization, and harmonization.

B Additional Qualitative Comparisons

MARBLE’s qualitative comparisons show that one model can improve text rendering, attribute and spatial understanding, and counting while maintaining coherent visual quality. The comparisons also expose limitations of weighted-sum and DiffusionNFT baselines, including degraded visual quality and blur or distortion artifacts.

  • Qualitative capabilities: MARBLE produces legible, semantically consistent text while preserving fine-grained attribute-object bindings and spatial layouts across diverse contexts.These qualitative results are shown in Figure S1.
  • Baseline comparisons: The weighted-sum baseline often fails to improve all aspects simultaneously, and its overall visual quality can be visibly degraded.It can sometimes satisfy specific requirements such as object counts or attributes.
  • Baseline comparisons: Compared with DiffusionNFT, MARBLE generates sharper images with fewer blur and distortion artifacts.The comparison is reported in Figures S8–S11.
  • Qualitative capabilities: MARBLE follows counting constraints while maintaining coherent visual quality in a single model.Figure S1 summarizes simultaneous improvements in text rendering, attribute and position understanding, and counting.

C Additional Ablations and Analyses · C.1 Key Insights and Takeaways · C.2 Update-Direction Harmony Diagnostics

MARBLE attributes multi-reward diffusion RL failures to scalar reward aggregation, motivating reward-specific gradient preservation and harmonization. The analyses further motivate multi-dimensional evaluation and report defaults for coefficient refresh and EMA smoothing.

  • C.1 Key Insights and Takeaways: MARBLE identifies scalar reward aggregation, rather than scalar-weight selection alone, as the central multi-reward optimization problem.Weighted-sum optimization underperforms on several reward dimensions, while gradient-alignment analysis indicates weakened reward-specific signals.
  • C.1 Key Insights and Takeaways: N = 10 is the default coefficient-refresh setting because it matches N = 5 while slightly improving speed and avoids degradation at N = 20.The comparison is reported in Table S2.
  • C.1 Key Insights and Takeaways: ρ = 0.7 is the default EMA setting used in the reported results.This default is reported in Table S3.
  • C.1 Key Insights and Takeaways: Image-quality improvement is inherently multi-dimensional and cannot be fully captured by any single reward model.The passage lists PickScore, CLIPScore, HPSv2, Aesthetic Score, and ImageReward among the relevant metrics.
  • C.2 Update-Direction Harmony Diagnostics: Update-direction harmony evaluates per-batch worst, mean, and variance of cos(d, g_k) for weighted-sum and MARBLE directions.Figure S2 defines these diagnostics, while Table S1 reports aggregate statistics averaged over n = 5 mini-batches.
  • C.2 Update-Direction Harmony Diagnostics: 80% of measured mini-batches have negative worst-reward alignment under the weighted-sum direction, which is negative on average.A negative cos(d, g_k) means the shared update is anti-aligned with that reward’s own gradient on the same batch.

C.3 Training Dynamics and Coefficient Adaptation … C.6 Alternative Heuristic Strategies

MARBLE’s adaptive coefficients favor difficult specialist rewards while preserving improvements across reward dimensions, outperforming fixed uniform weighting on harder objectives. Sensitivity analyses characterize refresh-frequency and EMA-smoothing tradeoffs, while heuristic reweighting fails to improve all rewards simultaneously.

  • C.3 Training Dynamics and Coefficient Adaptation: Fixed α_k = 0.2 progresses faster initially on HPSv2, but MARBLE reaches a higher final value; on GenEval, MARBLE trains faster and ends higher.HPSv2 is described as an easier image-quality reward, whereas GenEval is a harder specialist reward.
  • C.3 Training Dynamics and Coefficient Adaptation: All five rewards continue improving during MARBLE training, including general image-quality and specialist rewards.The training curves in Figure S3 show continued improvement across the optimized reward dimensions.
  • C.3 Training Dynamics and Coefficient Adaptation: MARBLE’s smoothed coefficients reflect optimization difficulty rather than raw reward values, staying below 0.2 for HPSv2 and often rising to around 0.3 for GenEval.The uniform baseline for five rewards is 0.2, while GenEval requires precise compositional correctness.
  • C.4 Amortization Interval: Smaller amortization intervals refresh α* more often and respond better to current gradient geometry, but increase training overhead; larger intervals reduce overhead while risking stale coefficients.The interval N controls how frequently full harmonization recomputes α* before cached coefficients are reused.
  • C.5 EMA Decay for Coefficient Smoothing: EMA decay ρ trades adaptivity against stability: smaller values react faster but amplify mini-batch noise, whereas larger values smooth trajectories but adapt slowly.The analysis evaluates ρ ∈ {0.1, 0.3, 0.5, 0.7, 0.9}; the supplied passage does not report the quantitative results.
  • C.6 Alternative Heuristic Strategies: Fixed uniform coefficients and manually increasing weights for challenging rewards such as OCR and GenEval are tested as alternative heuristic strategies.These strategies modify reward weighting in the weighted-sum objective rather than using MARBLE’s gradient-space balancing.
  • C.6 Alternative Heuristic Strategies: Heuristic strategies fail to improve all reward dimensions simultaneously and can degrade several metrics, indicating that manual reweighting is unreliable for multi-reward diffusion RL.The result further supports MARBLE’s effectiveness in balancing multiple rewards.

C.7 Human Preference Evaluation and Metric Correlations

MARBLE achieves the highest human-rated text-image alignment and image quality in a blind user study, despite lower PickScore and CLIPScore than DiffusionNFT†. Correlation and qualitative analyses support evaluating generation quality with multiple complementary criteria rather than relying on a single automatic metric.

  • Human Preference Evaluation: A blind study had 20 participants rate 30 anonymized images per method on 1–5 text-image alignment and image quality scales.Images were randomly shuffled and shown with their corresponding prompts; higher scores indicate better perceived performance.
  • Human Preference Evaluation: MARBLE receives the highest average score on both text-image alignment and image quality.The study is presented as complementary human-centered evidence without a claim of statistical significance.
  • Human Preference Evaluation: MARBLE’s lower PickScore and CLIPScore do not correspond to lower human-rated quality, aligning better with broader automatic metrics including Composite score.DiffusionNFT† has higher PickScore and CLIPScore, whereas MARBLE achieves the best Composite score and performs better on several broader quality- and preference-oriented metrics.
  • Metric Correlations: HPSv2.1, Aesthetic Score, UniReward, and ImageReward positively correlate with both human-rated axes, with HPSv2.1 showing the strongest agreement.PickScore and CLIPScore are weaker predictors of human ratings, indicating that no single automatic metric sufficiently characterizes human-perceived quality.
  • Qualitative Comparison: Across diverse prompts, MARBLE better preserves text rendering, attribute binding, spatial layout, and counting while maintaining sharper and more coherent visual quality.The weighted-sum baseline often fails to improve all aspects simultaneously, while DiffusionNFT† sometimes produces less sharp or detailed images despite strong proxy scores.

D Additional Implementation Details · E Future Work

MARBLE’s distributed implementation extracts and synchronizes per-reward gradients before local harmonization, requiring explicit control of DDP synchronization and graph retention. Future work targets scaling beyond five rewards and extending MARBLE to video generation and world models.

  • D Additional Implementation Details: The implementation details are provided for reproducing MARBLE on the DiffusionNFT codebase with distributed training.This establishes the stated implementation setting for the synchronization procedure.
  • D Additional Implementation Details: MARBLE’s DDP implementation requires per-reward gradients to be extracted before gradient harmonization, beyond standard cross-rank gradient averaging.Each GPU processes a different data shard, while standard DDP synchronization averages gradients across ranks before the optimizer step.
  • D Additional Implementation Details: Each reward loss is backpropagated under no_sync(), its flattened gradient is stored and cleared, then stored gradients are averaged across ranks with explicit all_reduce.This prevents premature averaging and avoids K unnecessary collective operations before harmonization.
  • D Additional Implementation Details: The first K−1 reward backward passes retain the shared computation graph, while the final pass releases it to free memory.All per-reward losses share the same forward computation graph.
  • D Additional Implementation Details: After cross-rank synchronization, all ranks solve harmonization locally with identical α∗, restore the combined gradient to model parameters, and perform the optimizer step.The procedure stores K flattened gradient vectors and synchronizes each with dist.all_reduce using AVG.
  • E Future Work: MARBLE currently studies five reward dimensions, while future work will investigate scalability to larger and more diverse, potentially conflicting reward sets.The authors also aim to develop more effective balancing strategies for increasingly diverse reward signals.
  • E Future Work: Extending MARBLE to video generation and generative world models is proposed as a future direction involving visual fidelity, temporal consistency, motion realism, and physical plausibility.World models additionally require coherent dynamics and plausible long-horizon evolution.
Loading 2605.06507v1…