ReCAST: Reward Credit Assignment across Timesteps for Online Diffusion Reinforcement
Organizations: University of California, Los Angeles · Arena AI
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 , 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.
Figures & tables
| Training stage-OCR ( budgets seeds) | Training stage-GenEval ( budgets seeds) | |||
|---|---|---|---|---|
| Training reward | base / static / Ours | win | base / static / Ours | win |
| Aggregate | 1.642 / 2.867 ±0.026 / 2.920 ±0.028 | 11/15 | 1.784 / 2.881 ±0.021 / 2.878 ±0.015 | 7/15 |
| ClipScore | 0.271 / 0.313 ±0.002 / 0.320 ±0.002 | 13/15 | 0.240 / 0.299 ±0.001 / 0.299 ±0.001 | 6/15 |
| HPSv2 | 0.195 / 0.288 ±0.008 / 0.303 ±0.004 | 12/15 | 0.213 / 0.313 ±0.005 / 0.312 ±0.002 | 5/15 |
| PickScore | 0.773 / 0.888 ±0.002 / 0.899 ±0.003 | 15/15 | 0.790 / 0.911 ±0.004 / 0.908 ±0.003 | 6/15 |
| Target reward | 0.129 / 0.893 ±0.014 / 0.900 ±0.032 | 9/15 | 0.244 / 0.873 ±0.011 / 0.874 ±0.007 | 9/15 |
| Training stage-OCR ( pairs) | Training stage-GenEval ( pairs) | |||
|---|---|---|---|---|
| Held-out model | base / static / Ours | win | base / static / Ours | win |
| ImageReward | 1.140 / 1.422 ±0.037 / 1.482 ±0.030 | 13/15 | 1.024 / 1.198 ±0.165 / 1.289 ±0.122 | 9/15 |
| Aesthetic | 5.388 / 5.486 ±0.063 / 5.578 ±0.009 | 12/15 | 5.949 / 5.864 ±0.120 / 5.968 ±0.102 | 10/15 |
| HPSv3 | 12.69 / 12.74 ±0.18 / 13.14 ±0.49 | 11/15 | 9.19 / 8.02 ±0.95 / 8.58 ±0.56 | 10/15 |
| UnifiedReward-2 | 3.015 / 3.154 ±0.025 / 3.176 ±0.040 | 11/15 | 2.682 / 2.671 ±0.085 / 2.705 ±0.046 | 9/15 |
| mean | ||||||
|---|---|---|---|---|---|---|
| Overall | 0.6538 | 0.5580 | 0.4780 | 0.5365 | 0.6325 | 0.5718 |
| Faithfulness | 0.6550 | 0.5570 | 0.5130 | 0.5655 | 0.6094 | 0.5800 |
| Aesthetics | 0.6438 | 0.5397 | 0.4736 | 0.4674 | 0.6242 | 0.5497 |
| static | 2.3848 | 3.2476 | 3.2948 | 2.5597 | 2.6980 |
|---|---|---|---|---|---|
| no-Sink Rényi | 2.4242 | 3.3514 | 3.3781 | 2.6721 | 2.6969 |
| ReCAST (Sinkhorn) | 2.4650 | 3.3564 | 3.3564 | 2.7960 | 2.7038 |
Appendix figures & tables13 assets
Supplementary material from the paper’s appendix.
Appendix
| Reward | Pearson | Spearman | Peak | Total gain |
|---|---|---|---|---|
| ClipScore | ||||
| HPSv2 | ||||
| OCR | ||||
| PickScore |
| Reward | Base policy | ckpt- | Remaining |
|---|---|---|---|
| ClipScore | |||
| HPSv2 | |||
| OCR | |||
| PickScore | |||
| All four |
| Method | Aggregate | ClipScore | HPSv2 | PickScore | OCR | Pareto | |
|---|---|---|---|---|---|---|---|
| static | 2.387 ±0.058 | 0.315 ±0.003 | 0.288 ±0.007 | 0.890 ±0.002 | 0.894 ±0.051 | Ours static | |
| Ours | 2.435 ±0.058 | 0.321 ±0.001 | 0.306 ±0.006 | 0.900 ±0.004 | 0.908 ±0.055 | ||
| static | 3.251 ±0.030 | 0.308 ±0.004 | 0.255 ±0.022 | 0.869 ±0.008 | 0.910 ±0.001 | Ours static | |
| Ours | 3.383 ±0.031 | 0.319 ±0.004 | 0.290 ±0.022 | 0.887 ±0.011 | 0.944 ±0.006 | ||
| static | 3.259 ±0.039 | 0.313 ±0.003 | 0.296 ±0.004 | 0.898 ±0.002 | 0.854 ±0.035 | Ours static | |
| Ours | 3.338 ±0.043 | 0.320 ±0.004 | 0.306 ±0.008 | 0.903 ±0.005 | 0.907 ±0.056 |
| Method | Aggregate | ClipScore | HPSv2 | PickScore | GenEval | Pareto | |
|---|---|---|---|---|---|---|---|
| static | 2.390 ±0.003 | 0.298 ±0.001 | 0.316 ±0.005 | 0.914 ±0.002 | 0.861 ±0.002 | mixed | |
| Ours | 2.392 ±0.022 | 0.298 ±0.003 | 0.308 ±0.007 | 0.909 ±0.006 | 0.877 ±0.014 | ||
| static | 3.300 ±0.017 | 0.299 ±0.001 | 0.309 ±0.004 | 0.908 ±0.008 | 0.892 ±0.007 | mixed | |
| Ours | 3.290 ±0.045 | 0.299 ±0.000 | 0.307 ±0.007 | 0.906 ±0.005 | 0.889 ±0.019 | ||
| static | 3.321 ±0.026 | 0.298 ±0.001 | 0.319 ±0.006 | 0.914 ±0.007 | 0.875 ±0.014 | mixed | |
| Ours | 3.322 ±0.014 | 0.297 ±0.004 | 0.322 ±0.001 | 0.913 ±0.004 | 0.877 ±0.009 |
| Method | ImageReward | Aesthetic | HPSv3 | UnifiedReward-2 | Pareto | |
|---|---|---|---|---|---|---|
| static | 1.433 ±0.056 | 5.493 ±0.096 | 12.86 ±0.33 | 3.179 ±0.071 | mixed | |
| Ours | 1.476 ±0.027 | 5.574 ±0.063 | 13.16 ±0.50 | 3.154 ±0.030 | ||
| static | 1.341 ±0.060 | 5.411 ±0.036 | 12.06 ±0.71 | 3.146 ±0.061 | mixed | |
| Ours | 1.420 ±0.096 | 5.542 ±0.027 | 12.44 ±0.81 | 3.129 ±0.055 | ||
| static | 1.438 ±0.042 | 5.525 ±0.083 | 13.00 ±0.18 | 3.159 ±0.008 | Ours static | |
| Ours | 1.499 ±0.029 | 5.571 ±0.045 | 13.25 ±0.61 | 3.190 ±0.065 |
| Method | ImageReward | Aesthetic | HPSv3 | UnifiedReward-2 | Pareto | |
|---|---|---|---|---|---|---|
| static | 1.372 ±0.011 | 6.057 ±0.020 | 8.73 ±0.51 | 2.745 ±0.025 | mixed | |
| Ours | 1.358 ±0.046 | 6.006 ±0.076 | 9.30 ±0.12 | 2.741 ±0.025 | ||
| static | 1.268 ±0.073 | 5.882 ±0.128 | 8.43 ±1.01 | 2.705 ±0.065 | Ours static | |
| Ours | 1.317 ±0.106 | 5.938 ±0.129 | 8.61 ±0.72 | 2.709 ±0.042 | ||
| static | 1.234 ±0.203 | 5.913 ±0.228 | 8.19 ±1.58 | 2.698 ±0.125 | Ours static | |
| Ours | 1.386 ±0.061 | 6.040 ±0.090 | 8.70 ±0.62 | 2.736 ±0.030 |
| Reward | Architecture / checkpoint | Range | Role |
|---|---|---|---|
| ClipScore [ 12 ] | CLIP ViT-L/14 ( openai/clip-vit-large-patch14 ) image–text cosine similarity | train | |
| HPSv2 [ 39 ] | open_clip ViT-H/14 backbone fine-tuned on human preference pairs ( HPS_v2.1_compressed.pt checkpoint) | train | |
| PickScore [ 18 ] | yuvalkirstain/PickScore_v1 , a CLIP ViT-H/14 fine-tuned end-to-end on preference pairs, with the laion/CLIP-ViT-H-14-laion2B-s32B-b79K processor | train | |
| OCR | PaddleOCR (English, angle classification off) text extraction | train (training stage-OCR) | |
| GenEval [ 9 ] | Mask2Former (Swin-S, COCO) detector + CLIP zero-shot color/counting classifiers | train (training stage-GenEval) | |
| Aesthetic [ 32 ] | Linear-probe MLP over frozen CLIP ViT-L/14 features, trained on LAION/AVA aesthetic ratings ( sac+logos+ava1-l14-linearMSE checkpoint) | held-out |
Diffusion Model Fine-Tuning
Diffusion model fine-tuning adapts pretrained generative models to new concepts, tasks, domains, and preferences through targeted updates.
See all 181 papers