Source-linked AI summary
ReCAST: Reward Credit Assignment across Timesteps for Online Diffusion Reinforcement
Yihang Chen, Yuanhao Ban, Kuei-Chun Kao, Cho-Jui Hsieh
TL;DR
Multi-reward diffusion training must allocate each reward’s total budget and decide when its feedback is informative. ReCAST distributes fixed reward budgets across denoising steps using per-reward discriminability gains, improving held-out judges in both evaluated settings even where training rewards tie.
Problem
Static scalarization fixes each reward’s relative influence across timesteps, although different rewards become informative at different stages of denoising.
Source: §1 · pp. 1–2Method
Results
On SD3.5-Medium, ReCAST improves training rewards in the OCR setting and matches them in GENEVAL, while every held-out judge improves and an independent LLM judge prefers it at four of five budgets.
Source: §5 · p. 9Takeaways & Limitations
When a reward is applied can matter as much as how much weight it receives.
Source: §5 · p. 9The estimated per-step gain uses a surrogate log-moment rather than an exact importance-sampling estimator.
Source: §3.3 · p. 5
Abstract
from arXiv · showhide
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 $\lambda$, 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\'enyi 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 $\lambda$. 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.
1 Introduction
Multi-reward diffusion training must decide both how much each reward contributes and when its feedback is informative. ReCAST assigns each reward a user-set budget, then allocates it across denoising steps using a reward-by-timestep weight matrix.
1 Introduction
- Static scalarization assumes a reward’s relative influence should remain constant across timesteps, even though its feedback may be informative at different stages.The paper frames multi-reward alignment as deciding both each reward’s total budget and when to spend it.Source: §1 · p. 1
- Global composition becomes visible before fine detail, but existing timestep-dependent schedules use one reward-agnostic curve rather than reward-specific timing.Prompt-alignment or rule-based rewards can score coarse structure earlier than aesthetic or fine-detail rewards can judge a nearly clean sample.Source: §1 · p. 2
- ReCAST introduces per-reward, timestep-dependent credit assignment for diffusion reward fine-tuning through a weight matrix whose entries specify each reward’s weight at each step.The authors describe it as the first such method to their knowledge.Source: §1 · p. 2
- The matrix preserves user preference budgets by row while enforcing equal total weight across denoising steps by column.Its row sums match the user-specified reward budgets; its column sums are each 1/T.Source: §1 · p. 2
2 Related work
Related work studies reward fine-tuning, combining multiple rewards, and timestep weighting in diffusion training. ReCAST’s distinction is to vary each reward’s weight across denoising steps using its own discriminability curve.
2 Related work
- ReCAST builds on DiffusionNFT [42], which regresses implicit positive and negative velocity fields using a forward-only loss.The authors select it for its performance and efficiency.Source: §2 · p. 2
- Prior multi-reward approaches use fixed scalarization, separate policies or post-hoc merging, hierarchies, or per-reward normalization; none varies reward weights across denoising timesteps.Examples include SafeRLHF [5], Personalized Soups [14], Rewarded Soups, ALaRM [19], and GDPO [25].Source: §2 · p. 2
- Existing timestep-weighting methods use a shared heuristic curve across rewards, whereas ReCAST derives timing from each reward’s own discriminability curve.The surveyed approaches include step-by-step preference optimization [21], MixGRPO [20], and TempFlow-GRPO [11].Source: §2 · pp. 2–3
3 Method
ReCAST estimates each reward’s per-step discriminability gains, turns them into weights satisfying reward-budget and timestep-total constraints, and illustrates reward-specific timing on SD3.5-Medium.
3 Method
- The method builds per-step reward-discriminability gains and then projects reward-specific shapes into a weight matrix satisfying both marginals.The method assumes nonnegative rewards and uses a rectified-flow setup with the DiffusionNFT objective.Source: §3 · p. 3
3.1 Setup and the density ratio
- A reward’s timestep density ratio is its posterior expected reward under the current policy divided by its marginal expected reward, so it is computable without sampling the positive policy.The conditional and marginal expected rewards are denoted µ_i,t and Z_i, respectively.Source: §3.1 · p. 3
3.2 Reward discriminability
- R´enyi discriminability between the positive and current policy marginals is nondecreasing through denoising, and its per-step gains telescope to total data-level discriminability.The gain isolates the new discriminability contributed at each transition, without repeatedly counting earlier gains.Source: §3.2 · pp. 3–4
3.3 Estimation details
- For α > 1, a zero conditional-reward estimate contributes zero before the logarithm, keeping the estimate finite when the empirical moment is positive; training uses α = 2.R´enyi divergence also upper-bounds KL for α > 1.Source: §3.3 · p. 4
- Because the positive-policy marginals are unknown and self-rollouts can miss high-reward regions, the estimator uses a stronger common proposal; its log-moment is a surrogate, not an exact importance-sampling estimator.The paper reports that πold self-rollouts yield near-zero discriminability estimates for several rewards.Source: §3.3 · p. 5
- Estimated curves remain essentially unchanged when the external proposal π+ is switched between GPT IMAGE 1.5 and NANO BANANA PRO, as reported in Appendix B.4.Source: §3.3 · p. 5
3.4 Sinkhorn projection
- Mean-normalized gains form an affinity kernel, which is projected onto the feasible set with row sums matching reward budgets and equal column totals.The method uses an entropic/KL projection, with a unique minimizer under the stated positive-kernel condition.Source: §3.4 · pp. 5–6
3.5 Empirical gain curves on SD3.5-Medium
- OCR, GENEVAL, and IMAGEREWARD peak toward noisy steps while other gains stay near-uniform; projected weights shift toward each reward’s high-gain steps with exact marginals.The figure plots gains, mean-normalized gains, and the Sinkhorn weight matrix; its horizontal axis runs from cleanest to noisiest step.Source: §3.5 · pp. 6–7, Figure 1 · p. 7
4 Experiments
Experiments compare ReCAST with static weighting across two four-reward settings, five budgets, and three seeds. ReCAST improves OCR training rewards, wins across held-out judges, and is preferred by an LLM judge; the Sinkhorn ablation separates adaptive timing from the shared step budget.
4.1 Training setup
- The study fine-tunes SD3.5-Medium LoRA adapters in OCR and GenEval settings, sweeping five reward budgets with three independent seeds per setting.The settings use separate OCR and compositional GenEval prompt sets; each budget is compared against static weighting at matched seeds.Source: §4.1 · p. 7
- The evaluation scores final checkpoints on held-out training rewards, held-out judges, and pairwise LLM judgments.Scores use matched seeds and report means and win rates over matched pairs.Source: §4.1 · pp. 7–8
4.2 Training rewards
- OCR aggregate score rose from 2.867 to 2.920 (+0.053±0.049), while GenEval aggregate shifted −0.003±0.031 and tied static.The comparison pools five budgets and three seeds; the table reports final-checkpoint scores.Source: Table 1 · p. 8, §4.2 · p. 8
4.3 Held-out judges
- All four held-out judges favored ReCAST in both settings; on GenEval, HPSV3 gained +0.56±0.41 despite tied training rewards.On OCR, the judges won 13, 12, 11, and 11 of 15 matched pairs; on GenEval, each won 9 or 10.Source: Table 2 · p. 8, §4.3 · p. 8
4.4 LLM-as-a-Judge
- For MMRBv2 pairwise evaluation against static, ReCAST’s mean overall win rate was 0.572 and it was preferred at four of five budgets, most strongly at (1, 1, 1, 1) and (2, 1, 1, 1).The evaluation was position-debiased over both presentation orders, and both faithfulness and aesthetics improve at those two budgets.Source: Table 3 · p. 9, §4.4 · p. 9
4.5 Ablation study: the Sinkhorn projection
- The adaptive kernel contributed +0.068 of ReCAST’s +0.099 gain over static, and Sinkhorn contributed +0.031 in the single-seed ablation.No-Sink R´enyi beat static at four of five budgets, while the full method beat no-Sink at four of five.Source: Table 4 · p. 9, §4.5 · p. 9
5 Conclusion
ReCAST fixes each reward’s total budget while distributing it over timesteps according to its R´enyi discriminability curve. On SD3.5-Medium, it improves OCR training rewards, transfers to held-out judges, and is preferred by an LLM judge.
5 Conclusion
- The reward budget λ fixes each reward’s total contribution, while its R´enyi discriminability curve distributes that contribution across timesteps.The method separates how much each reward counts from when its weight is applied.Source: §5 · p. 9
- On SD3.5-Medium, ReCAST improves OCR training rewards, every held-out judge in both settings, and LLM preference at four of five budgets.The GenEval training rewards tie between methods, while all held-out judges still favor ReCAST.Source: §5 · p. 9
A.2 Proof of Theorem 3.1
- Theorem 3.1’s proof differentiates the Rényi moment, cancels drift contributions, and rewrites the remaining diffusion term using ∇log ρ.The supplementary section provides the proof of Theorem 3.1.Source: A Supplementary theory · p. 13, §A.2 · pp. 14–15
A.1 Background: rectified flow matching and DiffusionNFT
- Rectified-flow sampling follows a straight noise-to-data interpolation and integrates the learned velocity backward from noise to clean data.A per-timestep regression coefficient can be attached without changing the estimator, and reward tilting changes the target while leaving the forward path untouched.Source: §A.1 · p. 13
- DiffusionNFT conditions the current policy on optimality to define a positive policy, then optimizes toward its velocity direction without estimating likelihood.The stated improvement direction follows from moving toward the positive model and away from the negative model.Source: §A.1 · pp. 13–14
A.3 Error induced by the surrogate
The appendix analyzes the surrogate used to estimate reward-specific discriminability gains, describes its three-phase estimator, and examines sensitivity to the Rényi order α.
A.3 Error induced by the surrogate
- A shared positive policy yields similar temporal weights when it is close to each reward-specific positive policy in KL divergence.Larger total discriminability makes the approximation less sensitive, while lower Sinkhorn temperature makes the final allocation more sensitive to profile error.Source: §A.3 · pp. 15–16
- With bounded rewards, the total-gain surrogate error is at most 2b_i; the per-timestep profile error is controlled when total discriminability exceeds 2b_i.The analysis assumes a sufficiently small mismatch between the shared and reward-specific positive policies.Source: §A.3 · pp. 15–16
B.1 Algorithm
- The estimator forward-noises clean positive-policy samples, rolls out the current policy to estimate conditional rewards, then differences log-moments to obtain per-step gains.The marginal expected reward cancels and need not be formed.Source: §B.1 · p. 16
B.2 α-sensitivity of the R´enyi estimator
- For α>0, α≠1, per-step gains remain nonnegative, but only α>1 gives the Rényi-divergence upper bound on KL.For α<1, the log-moment emphasizes the lower tail rather than the upper tail.Source: §B.2 · p. 16
- At α=2, projected weights span 0.29–1.60 times 1/(mT), milder than α=0.25 (0.24–1.89) and α=8 (0.41–2.24).The authors use α=2 as the default; projected weights track the changes in per-step gains.Source: §B.2 · p. 16, Figure 2 · p. 17
B.3 π+ vs. πold self-rollout gain curves
The self-rollout estimator produces noisy, often near-degenerate gain curves, whereas π+ sampling yields usable curves for constructing timestep weights. The estimated curves and resulting weights are also similar across the two tested π+ surrogate generators.
B.3 π+ vs. πold self-rollout gain curves
- Under πold, CLIPSCORE, HPSV2, and PICKSCORE have clean-end discriminability estimates of 0.016, 0.020, and 0.001, versus 0.245, 0.501, and 0.209 under π+; πold gains are jagged and non-monotonic.OCR retains appreciable discriminability under πold, but its per-step gain curve is unstable; the resulting Sinkhorn matrices also oscillate rather than smoothly redistributing weight.Source: §B.3 · pp. 18–19, Figure 3 · p. 19
- Increasing πold estimation to 16× the rollout cost leaves its gain curves just as jagged, while π+ produces usable curves at 1/16 of that cost.The comparison increases both samples per prompt and rollouts per noisy state by 4× for πold.Source: §B.3 · p. 19
B.4 Sensitivity to the π+ surrogate generator
- With all other estimator settings identical, GPT IMAGE 1.5 and NANO BANANA PRO as π+ generators yield near-identical reward-specific training stage-OCR gain-curve shapes and peak orderings.Source: §B.4 · p. 20, Table 5 · p. 20
B.5 Closing the divergence to π+ over a training stage
ReCAST training substantially reduces the measured discriminability between the policy and each reward’s positive distribution. The reduction occurs across all four rewards, rather than closing one reward’s gap at another’s expense.
B.5 Closing the divergence to π+ over a training stage
- By checkpoint-180, total gain falls to 0.14× its base-policy value for CLIPSCORE and 0.01–0.11× for the other three rewards, with the four-reward sum at 0.03×.The trained policy’s cumulative discriminability curves sit near zero across the trajectory on a shared vertical axis.Source: §B.5 · p. 20
- The reweighted objective closes nearly all of the measured gap for all four rewards simultaneously, without closing one reward’s gap at another’s expense.Because the measured Rényi divergence upper-bounds KL for α=2, the KL gap is closed at least as far.Source: §B.5 · p. 20
C.1 The log-domain iteration
The log-domain Sinkhorn solver converts normalized per-reward gains into a weight matrix satisfying the budget marginals. Boundary cases clarify how the construction relates to static weighting, hard assignments, and rewards without an estimated gain curve.
C.1 The log-domain iteration
- Algorithm 3 normalizes each reward’s gains, exponentiates them into a kernel, and alternates marginal updates in the log domain to stabilize scaling when kernel entries are tiny.The solver converges linearly, takes fewer than 100 iterations for m ≤ 4 and T ≤ 25, and is computed once per λ before training.Source: §C.1 · pp. 20–21
C.2 Boundary cases
- With flat gain curves, the kernel is rank one and the entropic optimum reduces to the static baseline.A reward with no estimated gain curve instead receives its full row budget spread uniformly, subject to the shared column scaling.Source: §C.2 · p. 22
- Removing entropy yields a transport solution that assigns each timestep essentially to the reward with the largest normalized gain, while entropy preserves a smooth, strictly positive matrix.The authors characterize hard assignment based on finite-sample gains as brittle.Source: §C.2 · p. 22
D Per-λ results
- Per-λ tables report final-checkpoint static-versus-ReCAST scores as means over three seeds, with standard deviations and a Pareto verdict for each budget.The verdict marks ReCAST or static as dominant only when it is no worse on every reward and strictly better on at least one.Source: D Per-λ results · p. 22
D.1 Training rewards
Across five reward budgets, ReCAST improves or matches training rewards depending on the setting, while held-out judges favor ReCAST in both settings. Per-budget results show the strength and variability of those comparisons.
D.1 Training rewards
- In training stage-OCR, ReCAST Pareto-dominates static at 3 of 5 budgets, is dominated at none, and has a higher aggregate at 4 of 5 budgets.At the two mixed budgets, ReCAST concedes OCR by 0.039 and 0.026, respectively.Source: §D.1 · p. 22, Table 7 · p. 23
- In training stage-GenEval, all five budgets are mixed, with no method dominating; aggregate shifts are smaller than either method’s seed standard deviation at every budget.The target GENEVAL reward rises at 3 budgets and falls at the other 2.Source: §D.1 · p. 22, Table 8 · p. 23
D.2 Held-out judges
- For held-out judges in stage-OCR, ReCAST Pareto-dominates at 3 of 5 budgets and is dominated at none; IMAGEREWARD, AESTHETIC, and HPSV3 improve at every budget.The two mixed budgets concede only UNIFIEDREWARD-2, by 0.025 and 0.017.Source: §D.2 · p. 22, Table 9 · p. 23
- For held-out judges in stage-GenEval, ReCAST dominates at 3 budgets, static dominates at none, and pooled scores favor ReCAST on all four judges.Per-budget seed standard deviations are wide enough that three replicates do not resolve individual budgets where means separate.Source: §D.2 · p. 22, Table 10 · p. 24
E.1 Reward models
- The appendix specifies each training reward’s computation and notes that the cosine-based CLIPSCORE, HPSV2, and PICKSCORE are clamped at zero, though their measured values remain positive across evaluated checkpoints.Their minimum observed values are 0.26, 0.23, and 0.77, respectively.Source: §E.1 · p. 24
E.2 Training hyperparameters and compute
The appendix specifies the SD3.5-Medium training configuration and compute costs, then details the MMRBv2 pairwise-judging protocol and its position debiasing.
E.2 Training hyperparameters and compute
- The static and ReCAST runs cost the same; a warmup run costs ≈68 GPU-hours, an OCR run ≈37, and a GenEval run ≈59.Training used single-node 8×NVIDIA H200 allocations.Source: §E.2 · p. 25
- Each gain curve is estimated once per stage and reused; a T=10 curve takes ≈15 GPU-hours, while a T=25 training-stage curve takes about 6× that.The P·N queries to GPT IMAGE 1.5 cost around $35 in API fees.Source: §E.2 · p. 25
E.3 MMRBv2 judge: protocol, position debiasing, and prompt
- Every matched image pair is judged in forward and reverse order, with each ordering crediting a win after mapping the judge’s choice back to method identity.Pairs use a fixed 1,000-prompt held-out OCR subset; calls that still fail after 8 retries are dropped.Source: §E.3 · pp. 25–26
- Faithfulness and aesthetics sub-scores come from the judge’s reasoning fields, with undecided or not-applicable orderings excluded from each criterion’s denominator.The fields are faithfulness to prompt and overall quality, respectively.Source: §E.3 · p. 26
F Limitations
The limitations concern estimated gain curves and dependence on a proprietary image-generation API.
F Limitations
- Using GPT IMAGE 1.5 as a surrogate for reward-specific positive policies can bias gain curves relative to theoretical ΔD_i,t when no importance correction is applied.The paper notes that swapping to an independent surrogate leaves the curve essentially unchanged, but does not remove this bias.Source: F Limitations · pp. 27–28
- Reproducing a curve from scratch depends on a proprietary, versioned, priced API and costs around $35, though the one-time per-stage expense can be amortized.Releasing gain-curve weight files would let downstream users avoid querying GPT IMAGE 1.5.Source: F Limitations · p. 28
G Qualitative examples
The galleries show paired static-baseline and ReCAST generations for OCR and GENEVAL, but the displayed examples are selected wins rather than representative samples.
G Qualitative examples
- All 20 displayed pairs are decisive, judge-verified ReCAST wins selected by a human from qualifying pairs, so they illustrate decisive wins rather than average-case behavior.The selection draws two pairs per λ budget in each setting.Source: G Qualitative examples · p. 28
- Figure 6 pairs static baseline images on the left with ReCAST images on the right for held-out OCR prompts, using curated wins rather than random examples.The figure shows two prompts per λ budget.Source: Figure 6 · p. 29
- Figure 7 uses the same left-right pairing for held-out GENEVAL prompts, with aggregate results reported separately from the gallery.The figure shows two prompts per λ budget; the examples are curated decisive wins.Source: Figure 7 · p. 30