Recast assigns reward importance by timestep in diffusion training
ReCAST: Reward Credit Assignment across Timesteps for Online Diffusion Reinforcement
Machine LearningArtificial Intelligence
Summary
Training AI models to create images or data involves fine-tuning them with rewards that measure quality, but these rewards matter differently at each step of the process. The authors propose ReCAST, a method that figures out when and how much to trust each reward at each step, based on how informative that reward is. This helps the model learn better by giving more attention to useful rewards at the right moments. Their tests show it improves outcomes and generalizes well beyond just the rewards used during training.
What this means in practice
- •For machine learning engineers: Improve training of generative diffusion models by assigning adaptive rewards at optimal denoising timesteps.
- •For ai service developers: Enhance quality control in AI-generated content by weighting multiple reward signals according to their informativeness over time.
Authors
Yihang Chen, Yuanhao Ban, Kuei-Chun Kao, Cho-Jui Hsieh
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.