Centered Reward Distillation in Diffusion RL

• The paper introduces the CRD method, a diffusion RL framework that uses within-prompt centering to cancel prompt-specific normalization constants.
• CRD replaces absolute reward matching with centered reward matching by applying softmax-based weighting to stabilize the reward estimation process.
• Empirical results indicate that CRD improves aesthetics and reduces reward hacking compared to existing methods in tasks like image generation.

Centered Reward Distillation (CRD) is a forward-process diffusion reinforcement learning method for post-training diffusion or flow models with black-box rewards. It is derived from KL-regularized reward maximization and is built on forward-process-based fine-tuning. Its defining claim is that the intractable prompt-dependent normalizing constant in the reward-tilted optimum cancels exactly under within-prompt centering, which turns an otherwise ill-posed absolute reward-matching problem into a well-posed centered reward-matching objective (Zhu et al., 14 Mar 2026).

1. Formal definition and objective

CRD starts from the KL-regularized objective

$$\max_{p_\theta}\; \mathbb{E}_{c}\Big[ \mathbb{E}_{x\sim p_\theta(\cdot\mid c)}[r(c,x)] -\beta\,\mathrm{KL}\!\left(p_\theta(\cdot\mid c)\,\|\,p_{\mathrm{ref}}(\cdot\mid c)\right) \Big].$$

Here $r(c,x)$ is an external reward, $p_{\mathrm{ref}}$ is a reference model, and $\beta>0$ is the KL regularization strength. The corresponding optimum has exponential-tilt form,

$$p_{\theta^*}(x\mid c)\;\propto\;p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right),$$

so the reward can be written as

$$r(c,x) \;=\; \beta \log \frac{p_{\theta^*}(x\mid c)}{p_{\mathrm{ref}}(x\mid c)} +\beta \log Z(c),$$

with prompt-dependent normalizer

$$Z(c)=\int p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right)dx.$$

The paper defines the model-implied reward

$$R_\theta(c,x)\triangleq\beta \log \frac{p_\theta(x\mid c)}{p_{\mathrm{ref}}(x\mid c)}.$$

At the optimum, $R_{\theta^*}(c,x)$ matches $r(c,x)$ only up to the additive prompt-dependent constant $r(c,x)$0. Direct regression of $r(c,x)$1 onto $r(c,x)$2 is therefore ill-posed. CRD addresses precisely this identifiability problem: only relative reward structure within a prompt is recoverable, and the method is built around that restriction rather than treating it as a nuisance term (Zhu et al., 14 Mar 2026).

2. Within-prompt centering and the CRD loss

The key identity of CRD is obtained by centering rewards within a group of $r(c,x)$3 samples $r(c,x)$4 generated for the same prompt $r(c,x)$5. For any weights $r(c,x)$6 satisfying $r(c,x)$7,

$r(c,x)$8

The prompt-specific constant cancels exactly because it appears in every term with total coefficient $r(c,x)$9. This is the defining “within-prompt centering” step.

CRD therefore replaces absolute reward matching by centered reward matching. For each sample,

$p_{\mathrm{ref}}$0

The training objective is

$p_{\mathrm{ref}}$1

The centering weights are defined by a reward softmax,

$p_{\mathrm{ref}}$2

As $p_{\mathrm{ref}}$3, the weights become uniform; as $p_{\mathrm{ref}}$4, they concentrate on the best sample. The appendix reports that CRD is largely insensitive to $p_{\mathrm{ref}}$5 across $p_{\mathrm{ref}}$6, and the main experiments use uniform weighting. The same appendix identifies two notable special cases: $p_{\mathrm{ref}}$7 recovers pairwise reward distillation, while $p_{\mathrm{ref}}$8 recovers a GVPO-style centered mean-matching limit (Zhu et al., 14 Mar 2026).

3. Diffusion-specific surrogate and algorithmic realization

Because exact log densities for diffusion models are intractable, CRD uses an ELBO-style surrogate for the implicit reward: $p_{\mathrm{ref}}$9 where $\beta>0$0 is the forward-noised latent, $\beta>0$1 is the predicted velocity, and $\beta>0$2. The appendix states that the practical implementation uses an adaptive-weighting variant with stop-gradient normalization.

The training loop uses three model instances initialized from the pretrained model: a trainable model $\beta>0$3, a moving “old” reference $\beta>0$4, and a sampling model $\beta>0$5. For each prompt $\beta>0$6, CRD samples $\beta>0$7 latent noises, generates $\beta>0$8 clean samples with an ODE sampler from the sampling model, and evaluates raw rewards $\beta>0$9. These raw rewards are scaled to $$p_{\theta^*}(x\mid c)\;\propto\;p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right),$$0 and then group-normalized as

$$p_{\theta^*}(x\mid c)\;\propto\;p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right),$$1

The method then samples $$p_{\theta^*}(x\mid c)\;\propto\;p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right),$$2 and $$p_{\theta^*}(x\mid c)\;\propto\;p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right),$$3, constructs

$$p_{\theta^*}(x\mid c)\;\propto\;p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right),$$4

and estimates the implicit reward using the current model against the old model,

$$p_{\theta^*}(x\mid c)\;\propto\;p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right),$$5

After centering external and implicit rewards within the prompt group, CRD minimizes $$p_{\theta^*}(x\mid c)\;\propto\;p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right),$$6 and adds a KL anchor term. The total loss is

$$p_{\theta^*}(x\mid c)\;\propto\;p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right),$$7

The two EMA models are then updated separately: $$p_{\theta^*}(x\mid c)\;\propto\;p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right),$$8

In the reported implementation, the base model is Stable Diffusion 3.5 Medium, LoRA rank is $$p_{\theta^*}(x\mid c)\;\propto\;p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right),$$9 with $$r(c,x) \;=\; \beta \log \frac{p_{\theta^*}(x\mid c)}{p_{\mathrm{ref}}(x\mid c)} +\beta \log Z(c),$$0, resolution is $$r(c,x) \;=\; \beta \log \frac{p_{\theta^*}(x\mid c)}{p_{\mathrm{ref}}(x\mid c)} +\beta \log Z(c),$$1, and the main experiments run on up to 4 NVIDIA H100 GPUs with runtime under 36 hours (Zhu et al., 14 Mar 2026).

4. Drift control, anchoring, and stability mechanisms

CRD is presented not only as a centering trick but as a drift-controlled forward-process diffusion RL framework. The paper isolates three failure modes. First, repeated online reward tilting can compound distribution drift: $$r(c,x) \;=\; \beta \log \frac{p_{\theta^*}(x\mid c)}{p_{\mathrm{ref}}(x\mid c)} +\beta \log Z(c),$$2 which unrolls to

$$r(c,x) \;=\; \beta \log \frac{p_{\theta^*}(x\mid c)}{p_{\mathrm{ref}}(x\mid c)} +\beta \log Z(c),$$3

Second, if the old/reference model tracks the current model too aggressively, the ratio signal collapses because $$r(c,x) \;=\; \beta \log \frac{p_{\theta^*}(x\mid c)}{p_{\mathrm{ref}}(x\mid c)} +\beta \log Z(c),$$4. Third, anchoring to the non-CFG conditional pretrained model can pull training toward a weaker conditional baseline than the inference-time CFG semantics.

The first stabilizer is decoupling the sampler from the moving reference. The sampling model is updated faster to remain near-on-policy, while the old/reference model is updated more slowly to preserve a meaningful nonzero ratio signal. The second stabilizer is KL anchoring to a CFG-guided pretrained model: $$r(c,x) \;=\; \beta \log \frac{p_{\theta^*}(x\mid c)}{p_{\mathrm{ref}}(x\mid c)} +\beta \log Z(c),$$5 implemented through a diffusion-space approximation

$$r(c,x) \;=\; \beta \log \frac{p_{\theta^*}(x\mid c)}{p_{\mathrm{ref}}(x\mid c)} +\beta \log Z(c),$$6

The paper explicitly notes that if the RL signal disappeared, this term alone would reduce to CFG distillation.

The third stabilizer is reward-adaptive anchor strength: $$r(c,x) \;=\; \beta \log \frac{p_{\theta^*}(x\mid c)}{p_{\mathrm{ref}}(x\mid c)} +\beta \log Z(c),$$7 The stated interpretation is that low-reward samples receive weaker anchoring, allowing larger corrective updates, while high-reward samples receive stronger anchoring, which reduces late-stage exploitation of reward-model loopholes. The paper reports that this schedule improves optimization speed while only minimally harming generation quality (Zhu et al., 14 Mar 2026).

5. Empirical performance and observed behavior

CRD is evaluated on two black-box reward settings: GenEval reward for compositional image generation and OCR reward for visual text rendering. Training and evaluation use the same prompt splits as Flow-GRPO, with additional evaluation on DrawBench prompts. Task-specific metrics are GenEval and OCR, and unseen preference or quality metrics are PickScore, CLIPScore, HPSv2.1, Aesthetics, and ImageReward (Zhu et al., 14 Mar 2026).

Setting	CRD	Selected comparison
GenEval	0.93	Flow-GRPO (w/ KL): 0.95
OCR	0.92	Flow-GRPO (w/ KL): 0.92
Aesthetics	5.44	SD3.5-M: 5.39

On GenEval, the reported scores are SD3.5-M $$r(c,x) \;=\; \beta \log \frac{p_{\theta^*}(x\mid c)}{p_{\mathrm{ref}}(x\mid c)} +\beta \log Z(c),$$8, Flow-GRPO (w/ KL) $$r(c,x) \;=\; \beta \log \frac{p_{\theta^*}(x\mid c)}{p_{\mathrm{ref}}(x\mid c)} +\beta \log Z(c),$$9, AWM (w/ CFG) $$Z(c)=\int p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right)dx.$$0, DiffusionNFT $$Z(c)=\int p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right)dx.$$1, and CRD $$Z(c)=\int p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right)dx.$$2. Although CRD is slightly below Flow-GRPO on pure GenEval reward, it is the only method in that comparison reported to improve Aesthetics over the base model, reaching $$Z(c)=\int p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right)dx.$$3 against SD3.5-M’s $$Z(c)=\int p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right)dx.$$4. Its unseen preference metrics on DrawBench are ImageReward $$Z(c)=\int p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right)dx.$$5, PickScore $$Z(c)=\int p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right)dx.$$6, and HPSv2.1 $$Z(c)=\int p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right)dx.$$7. On OCR, CRD reaches $$Z(c)=\int p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right)dx.$$8, matching Flow-GRPO (w/ KL), while AWM (w/o CFG) and DiffusionNFT reach $$Z(c)=\int p_{\mathrm{ref}}(x\mid c)\exp\!\left(\frac{r(c,x)}{\beta}\right)dx.$$9 but have negative ImageReward, $$R_\theta(c,x)\triangleq\beta \log \frac{p_\theta(x\mid c)}{p_{\mathrm{ref}}(x\mid c)}.$$0 and $$R_\theta(c,x)\triangleq\beta \log \frac{p_\theta(x\mid c)}{p_{\mathrm{ref}}(x\mid c)}.$$1 respectively. The paper presents this pattern as evidence of reduced reward hacking rather than maximal single-reward optimization (Zhu et al., 14 Mar 2026).

The appendix reports that CRD is largely insensitive to centering temperature $$R_\theta(c,x)\triangleq\beta \log \frac{p_\theta(x\mid c)}{p_{\mathrm{ref}}(x\mid c)}.$$2, that larger group size $$R_\theta(c,x)\triangleq\beta \log \frac{p_\theta(x\mid c)}{p_{\mathrm{ref}}(x\mid c)}.$$3 improves convergence and final reward, and that even $$R_\theta(c,x)\triangleq\beta \log \frac{p_\theta(x\mid c)}{p_{\mathrm{ref}}(x\mid c)}.$$4 yields meaningful learning. Moderate CFG at inference often improves preference metrics further: in the GenEval setting, ImageReward rises to $$R_\theta(c,x)\triangleq\beta \log \frac{p_\theta(x\mid c)}{p_{\mathrm{ref}}(x\mid c)}.$$5 with CFG sampling $$R_\theta(c,x)\triangleq\beta \log \frac{p_\theta(x\mid c)}{p_{\mathrm{ref}}(x\mid c)}.$$6; in the OCR setting, it rises to $$R_\theta(c,x)\triangleq\beta \log \frac{p_\theta(x\mid c)}{p_{\mathrm{ref}}(x\mid c)}.$$7 with CFG sampling $$R_\theta(c,x)\triangleq\beta \log \frac{p_\theta(x\mid c)}{p_{\mathrm{ref}}(x\mid c)}.$$8 (Zhu et al., 14 Mar 2026).

6. Relation to adjacent centered or reward-relative distillation methods

Several post-training methods outside diffusion RL implement closely related centered or relative reward transformations without using the name CRD. In small-language-model distillation, AdvDistill samples $$R_\theta(c,x)\triangleq\beta \log \frac{p_\theta(x\mid c)}{p_{\mathrm{ref}}(x\mid c)}.$$9 teacher responses per prompt, computes verifier-based rewards, and normalizes rewards within each prompt group by

$R_{\theta^*}(c,x)$0

then converts the resulting advantages into softmax weights for supervised fine-tuning. The paper describes these values as showing “how much better or worse each response is relative to the average response for the same prompt,” which makes it a direct instance of within-prompt reward centering in dataset distillation (Padarha, 25 Jun 2025).

In on-policy self-distillation, CREDIT decomposes the standard token reward into an input-specific teacher component and an input-generic teacher component, then subtracts a batch-contrastive estimate of the generic term: $R_{\theta^*}(c,x)$1 The paper argues that this subtraction is not merely a control variate, because it changes the reward semantics from generic feedback-predictiveness to input-specific feedback-predictiveness (Shen et al., 12 May 2026).

In preference optimization, TVKD introduces teacher value-based knowledge distillation through the potential-difference shaping term

$R_{\theta^*}(c,x)$2

with $R_{\theta^*}(c,x)$3. This is a difference-based teacher reward proxy rather than an absolute reward target, and the paper emphasizes its potential-based reward shaping interpretation and policy invariance (Kwon et al., 21 Sep 2025).

A closely related diffusion-distillation line appears in $R_{\theta^*}(c,x)$4, which recasts distribution matching as a reward and then applies group normalization: $R_{\theta^*}(c,x)$5 The paper explicitly characterizes GNDM and GNDMR as centered and variance-normalized internal reward estimators, and presents them as highly relevant analogues or precursors to CRD-style diffusion distillation (Fan et al., 30 Mar 2026).

By contrast, rubric-conditioned self-distillation is explicitly described as not implementing “a centered-reward objective in the usual sense.” It replaces scalar reward supervision with rubric-conditioned teacher distributions and therefore addresses the same credit-assignment problem through structured criterion-conditioned token supervision rather than reward centering or advantage normalization (Gu et al., 17 Jun 2026).

7. Limitations, caveats, and open questions

The CRD paper itself notes several open issues. It does not provide full convergence or stability theory under the practical approximations used in training. It also depends on base-model coverage: if the pretrained model rarely generates desirable behaviors, self-generated RL data may not bootstrap them well. Training remains materially more expensive than ordinary finetuning because it requires online generation, reward evaluation, and forward-process ELBO-style estimation, and the paper explicitly identifies video diffusion as a particularly expensive future setting (Zhu et al., 14 Mar 2026).

Adjacent work suggests that centered or relative reward transformations remain sensitive to how the baseline is estimated and how groups are formed. CREDIT reports that performance is best around $R_{\theta^*}(c,x)$6, with monotonic degradation as $R_{\theta^*}(c,x)$7 increases, and attributes this to over-debiasing that weakens the policy gradient and makes the policy more generic (Shen et al., 12 May 2026). AdvDistill shows that multi-sample centered reward distillation carries substantial compute overhead: AdvDistill with $R_{\theta^*}(c,x)$8 responses costs $R_{\theta^*}(c,x)$9 and $r(c,x)$0 GPU-hours for distilled SFT (Padarha, 25 Jun 2025). GNDM further shows that group construction matters materially: sharing both $r(c,x)$1 and $r(c,x)$2 within a group gives the best result, and shared noise within a group outperforms random noise (Fan et al., 30 Mar 2026).

Taken together, these caveats indicate that CRD is best understood not as a generic baseline subtraction trick, but as a family of reward-relative objectives whose behavior depends on identifiability, grouping, drift control, and the semantics of the centered signal. In the direct arXiv usage of the term, CRD denotes the diffusion RL formulation in which within-prompt centering removes the intractable prompt-wise normalizer and enables reward matching of implicit density-ratio rewards under explicit drift control (Zhu et al., 14 Mar 2026).