Related papers: MARBLE: Multi-Aspect Reward Balance for Diffusion RL

MARBLE: Multi-Aspect Reward Balance for Diffusion RL

URL: http://arxiv.org/abs/2605.06507v1
Date: Thu, 07 May 2026 16:20:42 GMT
Title: MARBLE: Multi-Aspect Reward Balance for Diffusion RL
Abstract summary: Reinforcement learning has become the dominant approach for aligning diffusion models with human preferences.<n>Existing practice deal with multiple rewards by training one specialist model per reward.<n>We propose MARBLE, a gradient-space optimization framework that maintains independent advantage estimators for each reward.
Score: 71.6241143519038
License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
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.

Related papers

ReCAST: Reward Credit Assignment across Timesteps for Online Diffusion Reinforcement [49.49571101362551]
We propose ReCAST (Reward Credit ASsignment across Timesteps) as a method for per-reward, timestep-dependent credit assignment in diffusion reward fine-tuning.<n>We show that ReCAST yields improvements that generalize beyond the training rewards to support its core principle: assigning each reward greater weight at the denoising timesteps where its feedback is most informative.
arXiv Detail & Related papers (2026-09-11T18:41:09Z)
GRAS: Guided Reduced-Variance Proposals and Adaptive Selection for Training-Free Reward Alignment in Discrete Diffusion [8.46069844016289]
We study the combined regime and identify two weaknesses in how it is usually run.<n>We call the resulting method Guided Reduced-variance proposals and Adaptive Selection (GRAS)<n>GRAS is simple yet effective: across regulatory DNA and protein design it attains the best training-free reward, outperforming prior training-free methods and matching or surpassing a reward-fine-tuned model, and it remains effective even for non-differentiable rewards.
arXiv Detail & Related papers (2026-08-27T03:53:59Z)
Sample-Adaptive Latent Rewards for Uncertainty-Guided Diffusion Post-Training [51.18425726646089]
We propose textscSURE, a unified latent-space framework for image and video diffusion models.<n>It learns reward distributions and directly uses their reliability to guide dense post-training.<n>textscSURE-REFL achieves the highest VBench quality, semantic, and total scores among the evaluated methods.
arXiv Detail & Related papers (2026-08-06T14:55:42Z)
Optimizing Visual Generative Models via Distribution-wise Rewards [57.46109819595665]
We present a novel framework that finetunes generative models using distribution-wise rewards, ensuring better alignment with real-world data distributions.<n>Our approach significantly improves FID-50K across various base models, from 5.77 for SiT and from 3.74 to 3.52 for EDM2.
arXiv Detail & Related papers (2026-07-02T15:08:56Z)
Reinforce Adjoint Matching: Scaling RL Post-Training of Diffusion and Flow-Matching Models [44.082049677105964]
Diffusion and flow-matching models scale because pretraining is supervised regression.<n>We show that the structure extends to RL post-training.<n>Under KL-regularized reward, the optimal generative process tilts the clean-endpoint distribution towards samples with higher reward.<n>We derive Reinforce Adjoint Matching (RAM): a consistency loss that corrects the pretraining target with the reward.
arXiv Detail & Related papers (2026-05-11T15:56:29Z)
On Training in Imagination [69.97419830683606]
State-of-the-art model-based reinforcement learning methods train policies on imagined rollouts.<n>We study how errors in learned dynamics and reward models affect returns and policy optimization.
arXiv Detail & Related papers (2026-05-07T12:51:32Z)
Reward Score Matching: Unifying Reward-based Fine-tuning for Flow and Diffusion Models [54.597200388738656]
Reward-based fine-tuning aims to steer a pretrained diffusion or flow-based generative model toward higher-reward samples.<n>We show that many can be written under a common framework, which we call reward score matching (RSM)
arXiv Detail & Related papers (2026-04-19T12:47:52Z)
SOAR: Self-Correction for Optimal Alignment and Refinement in Diffusion Models [48.335262141752715]
Post-training pipeline for diffusion models has two stages: supervised fine-tuning (SFT) on curated data and reinforcement learning (RL) with reward models.<n>We propose SOAR (Self-Correction for Optimal Alignment and Refinement), a bias-correction post-training method that fills this gap.<n>It is on-policy, reward-free, and provides dense per-time supervision with no credit-assignment problem.
arXiv Detail & Related papers (2026-04-14T11:45:15Z)
FAIL: Flow Matching Adversarial Imitation Learning for Image Generation [52.643484089126844]
Post-training of flow matching models-aligning the output distribution with a high-quality target-is mathematically equivalent to Imitation learning.<n>We propose Flow Matching Adrial Learning (FAIL), which minimizes policy-expert divergence through adversarial training without explicit rewards or pairwise comparisons.
arXiv Detail & Related papers (2026-02-12T16:36:33Z)
Nested-ReFT: Efficient Reinforcement Learning for Large Language Model Fine-Tuning via Off-Policy Rollouts [25.205293698698867]
We introduce Nested-ReFT, where a subset of layers of the target model acts as the behavior model to generate off-policy completions during training.<n>Our theoretical analysis shows that Nested-ReFT yields unbiased gradient estimates with controlled variance.<n>Our empirical analysis demonstrates improved computational efficiency measured as tokens/sec across multiple math reasoning benchmarks and model sizes.
arXiv Detail & Related papers (2025-08-13T18:37:46Z)
RED: Unleashing Token-Level Rewards from Holistic Feedback via Reward Redistribution [50.171320156632866]
Reinforcement learning from human feedback offers a promising approach to aligning large language models with human preferences.<n>Current reward models operate as sequence-to-one models, allocating a single, sparse, and delayed reward to an entire output sequence.<n>We propose a more fine-grained, token-level guidance approach for RL training.
arXiv Detail & Related papers (2024-11-13T02:45:21Z)

This list is automatically generated from the titles and abstracts of the papers in this site.