ReCAST: Reward Credit Assignment across Timesteps for Online Diffusion Reinforcement
- URL: http://arxiv.org/abs/2609.13425v1
- Date: Fri, 11 Sep 2026 18:41:09 GMT
- Title: ReCAST: Reward Credit Assignment across Timesteps for Online Diffusion Reinforcement
- Abstract summary: 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.
- Score: 49.49571101362551
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: Training diffusion models with multiple rewards requires distinguishing user preference from reward informativeness. User preference determines how much each reward should contribute to the overall objective; reward informativeness determines when its feedback is useful during denoising. Some rewards can meaningfully evaluate a sample as soon as global structure emerges, but others become informative only when the sample is nearly clean. To address both questions jointly, we propose ReCAST (Reward Credit ASsignment across T}imesteps), the first method, to our knowledge, for per-reward, timestep-dependent credit assignment in diffusion reward fine-tuning. ReCAST separates user preferences from temporal allocation through a reward-by-timestep weight matrix $W$, whose row sums match the user-specified reward budgets $λ$, while its column sums are equal, assigning the same total weight to each denoising step. Under these marginal constraints, ReCAST allocates weight according to each reward's informativeness, quantified by its Rényi discriminability gain at each step. These gains telescope to the total discriminability between the reward-induced positive policy and the current policy, providing a basis for temporal credit assignment. We evaluate ReCAST by training SD3.5-Medium under two distinct four-reward settings, each across five reward budgets $λ$. ReCAST improves the training rewards in one setting and matches them in the other, improves every held-out judge in both, and is preferred by an independent LLM-as-a-Judge. Together, these results show that ReCAST yields improvements that generalize beyond the training rewards and support its core principle: assigning each reward greater weight at the denoising timesteps where its feedback is most informative.
Related papers
- Learn What's Left, Not What's Mastered: Saturation Aware Advantage Reweighting for Multi-Reward Policy Optimization [60.89431018071735]
Reinforcement learning with group-relative advantages has become the de facto standard for post-training language model reasoners.<n>We show that rollouts with distinct reward profiles can receive identical advantages, and all objectives are optimized with fixed relative weights regardless of their current level of saturation.<n>We introduce saturation-aware reweighting, which standardizes each reward objective independently and adaptively discounts its contribution according to a batch-level estimate of objective saturation.
arXiv Detail & Related papers (2026-08-17T04:07:50Z) - TRACE: A Unified Rollout Budget Allocation Framework for Efficient Agentic Reinforcement Learning [52.67756371231985]
Reinforcement learning with verifiable rewards (RLVR) is a promising approach for enhancing reasoning and agentic behavior in large language models.<n>We introduce Tree Rollout Allocation for Contrastive Exploration (TRACE), a unified rollout allocation framework.<n>Technically, TRACE allocates rollout budget to both prompt roots and intermediate prefixes that are most likely to yield mixed terminal rewards.
arXiv Detail & Related papers (2026-06-09T17:16:03Z) - MARBLE: Multi-Aspect Reward Balance for Diffusion RL [71.6241143519038]
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.
arXiv Detail & Related papers (2026-05-07T16:20:42Z) - Learning Explainable Dense Reward Shapes via Bayesian Optimization [45.34810347865996]
We frame reward shaping as an optimization problem focused on token-level credit assignment.<n>We use explainability methods such as SHAP and LIME to estimate per-token rewards from the reward model.<n>Our experiments show that achieving a better balance of token-level reward attribution leads to performance improvements over baselines.
arXiv Detail & Related papers (2025-04-22T21:09:33Z) - Information-Theoretic Reward Decomposition for Generalizable RLHF [51.550547285296794]
We decompose the reward value into two independent components: prompt-free reward and prompt-related reward.<n>We propose a new reward learning algorithm by prioritizing data samples based on their prompt-free reward values.
arXiv Detail & Related papers (2025-04-08T13:26:07Z) - 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.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.