ReCAST: Reward Credit Assignment across Timesteps for Online Diffusion Reinforcement

Yihang Chen; Yuanhao Ban; Kuei-Chun Kao; Cho-Jui Hsieh

cs.LGSep 11, 2026

ReCAST: Reward Credit Assignment across Timesteps for Online Diffusion Reinforcement

Authors: Yihang Chen, Yuanhao Ban, Kuei-Chun Kao, Cho-Jui Hsieh

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 $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.

Figures & tables

Figure 1: Gain curves and the weight matrix they project to. Left: per-step gain

\Delta D_{i,t}

. Center: mean-normalized gain

\overline{\Delta D}_{i,t}

(Eq. ( 13 )), the exponent of the affinity kernel (Eq. ( 14 )). Right: the Sinkhorn matrix

W_{i,t}^{\star}

(Eq. ( 18 )) at uniform

\lambda_{i}{=}1/m

, solved over the four rewards that are trained together in the OCR setting; both marginals hold exactly, every column summing to

1/T

and every row to

\lambda_{i}

. In all panels the horizontal axis is

t/T

, running from the cleanest step

t{=}1

(left) to the noisiest

t{=}T

(right). OCR and GenEval (dashed) are resampled onto the shared

T{=}10

grid from their native

T{=}25

curves, so the right panel is illustrative: the training-stage kernel is solved at

T{=}25

on OCR ’s own grid (Sec. 4.1 ).

	Training stage-OCR ( $5$ budgets $\times$ $3$ seeds)		Training stage-GenEval ( $5$ budgets $\times$ $3$ seeds)
Training reward	base / static / Ours	win	base / static / Ours	win
Aggregate $r$	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

Table 1: Training-reward scores at the final checkpoint, reported as base / static / Ours, where base is the untuned SD3.5-Medium, static is static reweighting by

\bm{\lambda}

, and Aggregate

r

is the sum of the individual rewards weighted by the integer

\bm{\lambda}

. Both settings were trained end-to-end (warmup stage + training-stage sweep, both methods) under

3

seeds; entries are the mean of the three seed-level means, standard deviation as a subscript, and win pools all

15

matched (budget, seed) pairs. Bold marks the better of static and Ours, and both when the two agree at the reported precision. Per-

\bm{\lambda}

results are in App. D .

	Training stage-OCR ( $15$ pairs)		Training stage-GenEval ( $15$ 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

Table 2: Held-out judge scores at the final checkpoint, reported as base / static / Ours. Both settings cover

5

budgets

\times

3

independent training seeds, and entries, subscripts, win , and the bolding rule are as in Tab. 1 ; the HPSv3 row is at two decimals because its scale is an order of magnitude larger. Per-

\bm{\lambda}

results are in App. D .

$\bm{\lambda}$	$(1,1,1,1)$	$(1,1,1,2)$	$(1,1,2,1)$	$(1,2,1,1)$	$(2,1,1,1)$	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

Table 3: MMRBv2 pairwise evaluation of ReCAST against static on the same subset as Tab. 2 , position-debiased over both presentation orders. Every entry is ReCAST’s win rate, overall and under the two rubric criteria, with

0.5

indicating parity;

\bm{\lambda}

is ordered ( ClipScore , HPSv2 , PickScore , OCR ).

$\bm{\lambda}$	$(1,1,1,1)$	$(1,1,1,2)$	$(1,1,2,1)$	$(1,2,1,1)$	$(2,1,1,1)$
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

Table 4: Aggregate training-reward score for each

\bm{\lambda}=(\text{clip},\text{hps},\text{pick},\text{ocr})

budget in the training stage-OCR experiment, comparing static weighting, the adaptive kernel alone (no-Sink Rényi), and the full method. Scores are on the held-out OCR dataset at the last checkpoint (step

180

), single matched training seed.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 2:

\alpha

-sensitivity of the

\pi^{+}

-sampling Rényi estimator on the SD3.5-Medium run of Sec. 3.5 . Rows:

\alpha\in\{0.25,\,0.5,\,0.75,\,1.5,\,2,\,4,\,8\}

. Columns: the same three quantities as Fig. 1 , per-step gain

\Delta D_{i,t}

(left), mean-normalized gain

\overline{\Delta D}_{i,t}

(Eq. ( 13 ), center), and the Sinkhorn matrix

W_{i,t}^{\star}

(Eq. ( 18 )) at uniform

\lambda_{i}{=}1/m

(right).

Figure 3: Cumulative discriminability

D_{i,t}

, per-step gain

\Delta D_{i,t}

, mean-normalized gain

\overline{\Delta D}_{i,t}

, and Sinkhorn matrix

W_{i,t}^{\star}

at uniform

\lambda_{i}{=}1/m

, left to right in the order they are derived, under the two sampling distributions. The leftmost column is the estimator’s own

D_{i,t}

, shifted by the single

t

-independent constant that sets

D_{i,T}{=}0

at pure noise, where

\pi_{i,T}^{+}{=}\pi_{T}^{\mathrm{old}}

; it is not a cumulative sum of the clipped gains beside it. Under

\pi^{\mathrm{old}}

almost all of the estimated discriminability sits on OCR , the other three rewards coming out nearly indiscriminable at every step, and the OCR curve is not even monotone in

t

: it dips below zero, which Theorem 3.1 rules out in population. Top:

\pi^{+}

samples from GPT Image 1.5 , at the

N{=}8

K{=}16

budget used throughout the paper. Middle:

\pi^{\mathrm{old}}

self-rollouts at that same budget. Bottom: the same

\pi^{\mathrm{old}}

run at the full

N{=}32

K{=}64

estimation budget,

16\times

the rollout cost. All three rows are matched estimator runs on the training stage-OCR curve: same four training rewards ( ClipScore , HPSv2 , OCR , PickScore ), same prompt set, same

\alpha{=}2

and

T{=}25

; the top two differ only in the sampling distribution, the bottom two only in

(N,K)

. Both

\pi^{\mathrm{old}}

rows are visibly noisier at every step than the

\pi^{+}

row, consistent with the higher-variance

\pi^{\mathrm{old}}

-side estimator, and the

16\times

budget does not smooth them.

Reward	Pearson $r$	Spearman $\rho$	Peak $t$	Total gain
ClipScore	$0.879$	$0.912$	$17\to 19$	$0.94\times$
HPSv2	$0.847$	$0.829$	$15\to 16$	$0.81\times$
OCR	$0.862$	$0.844$	$17\to 17$	$0.92\times$
PickScore	$0.919$	$0.915$	$15\to 16$	$0.84\times$

Appendix

Table 5: Surrogate sensitivity of the training stage-OCR gain curve: GPT Image 1.5 vs. Nano Banana Pro as the

\pi^{+}

generator, all other estimator settings identical.

r

and

\rho

are Pearson and Spearman correlations between the two per-step curves over the

T{=}25

steps. “Peak

t

” gives the argmax timestep of

\Delta D_{i,t}

in the paper’s convention (

t{=}1

clean,

t{=}T

pure noise); “total gain” is the ratio of

\sum_{t}\Delta D_{i,t}

under Nano Banana Pro to that under GPT Image 1.5 .

Figure 4: Per-step gain

\Delta D_{i,t}

, mean-normalized gain

\overline{\Delta D}_{i,t}

, and Sinkhorn matrix

W_{i,t}^{\star}

at uniform

\lambda_{i}{=}1/m

, for the training stage-OCR curve under the two

\pi^{+}

surrogate generators. Top: GPT Image 1.5 . Bottom: Nano Banana Pro . The two rows produce near-identical reward-specific shapes and peak orderings (quantified in Tab. 5 ).

Reward	Base policy	ckpt- $180$	Remaining
ClipScore	$0.252$	$0.036$	$0.14\times$
HPSv2	$0.501$	$0.024$	$0.05\times$
OCR	$2.278$	$0.028$	$0.01\times$
PickScore	$0.209$	$0.024$	$0.11\times$
All four	$3.241$	$0.113$	$0.03\times$

Appendix

Table 6: Total discriminability gain

\sum_{t}\Delta D_{i,t}

of the training stage-OCR curve with

\pi^{\mathrm{old}}

at the untrained base policy vs. at the trained checkpoint-

180

, with the GPT Image 1.5

\pi^{+}

samples and all other estimator settings identical. By Eq. ( 7 ) each entry is the data-level divergence remaining between that policy and the reward’s positive distribution.

Figure 5: Cumulative discriminability

D_{i,t}

of the training stage-OCR curve at the two policies, with

\pi^{+}

set to GPT Image 1.5 and all other estimator settings identical. Left:

\pi^{\mathrm{old}}

at the untrained base. Right:

\pi^{\mathrm{old}}

at checkpoint-

180

of the

\bm{\lambda}{=}(1,1,1,1)

ReCAST stage-OCR run, drawn on the same vertical axis; the inset repeats it at its own scale. Each curve telescopes the per-step gains up from

D_{i,T}=0

(Eq. ( 7 )), so its clean endpoint

D_{i,0}

is the full data-level divergence between that policy and reward

i

’s positive distribution. Training leaves

D_{i,t}

smaller at every timestep, for every reward.

$\bm{\lambda}$	Method	Aggregate $r$	ClipScore	HPSv2	PickScore	OCR	Pareto
$(1,1,1,1)$	static	2.387 ±0.058	0.315 ±0.003	0.288 ±0.007	0.890 ±0.002	0.894 ±0.051	Ours $\succ$ static
$(1,1,1,1)$	Ours	2.435 ±0.058	0.321 ±0.001	0.306 ±0.006	0.900 ±0.004	0.908 ±0.055	Ours $\succ$ static
$(1,1,1,2)$	static	3.251 ±0.030	0.308 ±0.004	0.255 ±0.022	0.869 ±0.008	0.910 ±0.001	Ours $\succ$ static
$(1,1,1,2)$	Ours	3.383 ±0.031	0.319 ±0.004	0.290 ±0.022	0.887 ±0.011	0.944 ±0.006	Ours $\succ$ static
$(1,1,2,1)$	static	3.259 ±0.039	0.313 ±0.003	0.296 ±0.004	0.898 ±0.002	0.854 ±0.035	Ours $\succ$ static
$(1,1,2,1)$	Ours	3.338 ±0.043	0.320 ±0.004	0.306 ±0.008	0.903 ±0.005	0.907 ±0.056	Ours $\succ$ static

Appendix

Table 7: Training stage-OCR, per-

\bm{\lambda}

training rewards on the held-out dataset at each run’s final checkpoint, static vs. ReCAST. Every entry is the mean over the

3

seed replicates with its standard deviation as a subscript, Aggregate

r

is the

\bm{\lambda}

-weighted sum as in Tab. 1 , and bold marks the better method per cell, neither when the two agree at the reported precision. Pareto is the verdict over the four training rewards.

$\bm{\lambda}$	Method	Aggregate $r$	ClipScore	HPSv2	PickScore	GenEval	Pareto
$(1,1,1,1)$	static	2.390 ±0.003	0.298 ±0.001	0.316 ±0.005	0.914 ±0.002	0.861 ±0.002	mixed
$(1,1,1,1)$	Ours	2.392 ±0.022	0.298 ±0.003	0.308 ±0.007	0.909 ±0.006	0.877 ±0.014	mixed
$(1,1,1,2)$	static	3.300 ±0.017	0.299 ±0.001	0.309 ±0.004	0.908 ±0.008	0.892 ±0.007	mixed
$(1,1,1,2)$	Ours	3.290 ±0.045	0.299 ±0.000	0.307 ±0.007	0.906 ±0.005	0.889 ±0.019	mixed
$(1,1,2,1)$	static	3.321 ±0.026	0.298 ±0.001	0.319 ±0.006	0.914 ±0.007	0.875 ±0.014	mixed
$(1,1,2,1)$	Ours	3.322 ±0.014	0.297 ±0.004	0.322 ±0.001	0.913 ±0.004	0.877 ±0.009	mixed

Appendix

Table 8: Training stage-GenEval, per-

\bm{\lambda}

training rewards on the held-out dataset at each run’s final checkpoint, static vs. ReCAST. Entries, subscripts, bolding, and the Pareto verdict are as in Tab. 7 .

$\bm{\lambda}$	Method	ImageReward	Aesthetic	HPSv3	UnifiedReward-2	Pareto
$(1,1,1,1)$	static	1.433 ±0.056	5.493 ±0.096	12.86 ±0.33	3.179 ±0.071	mixed
$(1,1,1,1)$	Ours	1.476 ±0.027	5.574 ±0.063	13.16 ±0.50	3.154 ±0.030	mixed
$(1,1,1,2)$	static	1.341 ±0.060	5.411 ±0.036	12.06 ±0.71	3.146 ±0.061	mixed
$(1,1,1,2)$	Ours	1.420 ±0.096	5.542 ±0.027	12.44 ±0.81	3.129 ±0.055	mixed
$(1,1,2,1)$	static	1.438 ±0.042	5.525 ±0.083	13.00 ±0.18	3.159 ±0.008	Ours $\succ$ static
$(1,1,2,1)$	Ours	1.499 ±0.029	5.571 ±0.045	13.25 ±0.61	3.190 ±0.065	Ours $\succ$ static

Appendix

Table 9: Training stage-OCR, per-

\bm{\lambda}

held-out judge scores at each run’s final checkpoint, static vs. ReCAST. Entries, subscripts, bolding, and the Pareto verdict are as in Tab. 7 , now over the four judges; the HPSv3 column is at two decimals because its scale is an order of magnitude larger.

$\bm{\lambda}$	Method	ImageReward	Aesthetic	HPSv3	UnifiedReward-2	Pareto
$(1,1,1,1)$	static	1.372 ±0.011	6.057 ±0.020	8.73 ±0.51	2.745 ±0.025	mixed
$(1,1,1,1)$	Ours	1.358 ±0.046	6.006 ±0.076	9.30 ±0.12	2.741 ±0.025	mixed
$(1,1,1,2)$	static	1.268 ±0.073	5.882 ±0.128	8.43 ±1.01	2.705 ±0.065	Ours $\succ$ static
$(1,1,1,2)$	Ours	1.317 ±0.106	5.938 ±0.129	8.61 ±0.72	2.709 ±0.042	Ours $\succ$ static
$(1,1,2,1)$	static	1.234 ±0.203	5.913 ±0.228	8.19 ±1.58	2.698 ±0.125	Ours $\succ$ static
$(1,1,2,1)$	Ours	1.386 ±0.061	6.040 ±0.090	8.70 ±0.62	2.736 ±0.030	Ours $\succ$ static

Appendix

Table 10: Training stage-GenEval, per-

\bm{\lambda}

held-out judge scores at each run’s final checkpoint, static vs. ReCAST. Entries, subscripts, bolding, and the Pareto verdict are as in Tab. 9 .

Reward	Architecture / checkpoint	Range	Role
ClipScore [ 12 ]	CLIP ViT-L/14 ( openai/clip-vit-large-patch14 ) image–text cosine similarity	$\approx[0,1]$	train
HPSv2 [ 39 ]	open_clip ViT-H/14 backbone fine-tuned on human preference pairs ( HPS_v2.1_compressed.pt checkpoint)	$\approx[0,1]$	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	$\approx[0,3.8]$	train
OCR	PaddleOCR (English, angle classification off) text extraction	$[0,1]$	train (training stage-OCR)
GenEval [ 9 ]	Mask2Former (Swin-S, COCO) detector + CLIP zero-shot color/counting classifiers	$[0,1]$	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)	$\approx[1,10]$	held-out

Appendix

Table 11: Reward models used in the paper. “Range” is the reward’s output scale; several are unbounded in principle but concentrate in the interval shown in practice.

Panel of Figure 6 — Figure 6: Qualitative examples from the OCR setting (training stage-OCR family, held-out $1000$ -prompt OCR dataset of Tab. 2 ): two prompts per $\bm{\lambda}$ budget, static baseline (left of each pair) vs. ReCAST (right of each pair). All $10$ pairs are a curated selection of decisive, judge-verified ReCAST wins, not a random or representative sample; see App. G for the selection methodology and Tab. 3 for this setting’s aggregate win rates.

Panel of Figure 7 — Figure 7: Qualitative examples from the GenEval setting (training stage-GenEval family, held-out GenEval prompts): two prompts per $\bm{\lambda}$ budget, static baseline (left of each pair) vs. ReCAST (right of each pair). All $10$ pairs are a curated selection of decisive, judge-verified ReCAST wins, not a random or representative sample; see App. G for the selection methodology and Tabs. 1 and 2 for this setting’s aggregate results; Sec. 4.4 reports the pairwise judge for the OCR setting only.

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

Explicit Critic Guidance for Aligning Diffusion Models
May 26, 2026Zhengyang Liang, Qihang Zhang, Ceyuan YangDiffusion Model Alignment Diffusion Model Guidance
University of Toronto · Vector Institute · The Chinese University of Hong Kong
CAST: Causal Advantage-Structured Training with Spatially Grounded Compositional Rewards for Diffusion Models
Sep 30, 2026Shu Yu, Chaochao LuDiffusion Model Fine-Tuning Compositional T2I Generation
Shanghai Artificial Intelligence Laboratory, Shanghai, China · Shanghai Innovation Institute, Shanghai, China · Fudan University, Shanghai, China
Latent Reward Registers for Diffusion Preference Alignment
Aug 4, 2026Yuanshen Guan, Zipeng Feng, Chengru Song +2Diffusion Model Alignment Diffusion Model Guidance
1Kling Team · University of Science and Technology of China

May 26, 2026J/K move · Enter open · S save

Explicit Critic Guidance for Aligning Diffusion Models

Zhengyang Liang, Qihang Zhang, Ceyuan Yang

University of Toronto · Vector Institute · The Chinese University of Hong Kong

Diffusion Model Alignment Diffusion Model Guidance RL for Diffusion Models

Open paper AskPDF arXiv

Online reinforcement learning is becoming increasingly important for aligning diffusion models with non-differentiable objectives. However, existing methods still face limitations in assigning fine-grained credit along denoising trajectories and in realizing stable value-based optimization. We propose a state-aligned latent actor-critic framework for diffusion post-training, in which the diffusion model serves as its own timestep-conditioned value function and predicts values directly on noisy latent states. This enables trajectory-level PPO training, supports stable actor-critic optimization with simple conditioning and value pretraining strategies, and naturally allows the learned critic to be reused for inference-time steering. We further extend the framework to multi-reward optimization, where joint training with complementary rewards helps alleviate reward hacking. Across both UNet- and DiT-based backbones, our method consistently outperforms prior group-relative RL and actor-critic baselines on single-reward and multi-reward benchmarks, while test-time steering provides additional gains in generation quality.

Abstract

Diffusion Model Fine-Tuning

Explore similar work

Explicit Critic Guidance for Aligning Diffusion Models

CAST: Causal Advantage-Structured Training with Spatially Grounded Compositional Rewards for Diffusion Models

Latent Reward Registers for Diffusion Preference Alignment