--- title: Centered Reward Distillation in Diffusion RL url: https://www.emergentmind.com/topics/centered-reward-distillation-crd type: topic --- # Centered Reward Distillation in Diffusion RL 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 [2603.14128]. ## 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 \(\beta \log Z(c)\). Direct regression of \(R_\theta\) onto \(r\) 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 [2603.14128]. ## 2. Within-prompt centering and the CRD loss The key identity of CRD is obtained by centering rewards within a group of \(K\) samples \(\{x_i\}_{i=1}^K\) generated for the same prompt \(c\). For any weights \(\{w_j\}_{j=1}^K\) satisfying \(\sum_j w_j=1\), \[ r(c,x_i) - \sum_{j=1}^K w_j r(c,x_j) = R_{\theta^*}(c,x_i) - \sum_{j=1}^K w_j R_{\theta^*}(c,x_j). \] The prompt-specific constant cancels exactly because it appears in every term with total coefficient \(1\). This is the defining “within-prompt centering” step. CRD therefore replaces absolute reward matching by centered reward matching. For each sample, \[ \Delta_{r,w}^i \triangleq r(c,x_i)-\sum_{j=1}^K w_j\, r(c,x_j), \qquad \Delta_{\widehat{R},w}^i \triangleq \widehat{R}_\theta(c,x_i)-\sum_{j=1}^K w_j\,\widehat{R}_\theta(c,x_j). \] The training objective is \[ \mathcal{L}_{\mathrm{CRD}}^{(\tau)}(p_\theta;\rho) = \mathbb{E}_{\rho(c,\{x_i\})}\; \frac{1}{K}\sum_{i=1}^K \left(\Delta_{r,w}^i - \Delta_{\widehat{R},w}^i\right)^2. \] The centering weights are defined by a reward softmax, \[ w_i(c,\{x_j\}_{j=1}^K;\tau) = \frac{\exp\!\left(r(c,x_i)/\tau\right)}{\sum_{j=1}^K \exp\!\left(r(c,x_j)/\tau\right)}. \] As \(\tau\to\infty\), the weights become uniform; as \(\tau\to 0\), they concentrate on the best sample. The appendix reports that CRD is largely insensitive to \(\tau\) across \(1,10,100,\infty\), and the main experiments use uniform weighting. The same appendix identifies two notable special cases: \(K=2,\tau\to 0\) recovers pairwise reward distillation, while \(\tau\to\infty\) recovers a GVPO-style centered mean-matching limit [2603.14128]. ## 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: \[ \widehat{R}_\theta(c,x) \triangleq -\beta\, \mathbb{E}_{t,\epsilon} \left[ w(t)\Big( \|v_\theta(x_t,t\mid c)-v_{\mathrm{target}}\|_2^2 - \|v_{\mathrm{ref}}(x_t,t\mid c)-v_{\mathrm{target}}\|_2^2 \Big) \right], \] where \(x_t\) is the forward-noised latent, \(v_\theta\) is the predicted velocity, and \(v_{\mathrm{target}}=\epsilon-x\). 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 \(\theta\), a moving “old” reference \(\theta_{\mathrm{old}}\), and a sampling model \(\theta_{\mathrm{samp}}\). For each prompt \(c\), CRD samples \(K\) latent noises, generates \(K\) clean samples with an ODE sampler from the sampling model, and evaluates raw rewards \(r_{\mathrm{raw}}^i=r_\psi(c,x_0^i)\). These raw rewards are scaled to \([0,1]\) and then group-normalized as \[ r(c,x_0^i)= \frac{r_{\mathrm{raw}}^i-\mathrm{mean}(\{r_{\mathrm{raw}}\}^{1:K})}{\mathrm{std}(\{r_{\mathrm{raw}}\}^{1:K})}. \] The method then samples \(t\sim\mathcal U[0,1]\) and \(\epsilon\sim\mathcal N(0,1)\), constructs \[ x_t = t\epsilon + (1-t)x_0, \qquad v_{\mathrm{target}}=\epsilon-x_0, \] and estimates the implicit reward using the current model against the old model, \[ R_\theta(c,x_0^i) \leftarrow -\beta_{\mathrm{old}} \Big( \|v_\theta-v_{\mathrm{target}}\|^2 - \|v_{\theta_{\mathrm{old}}}-v_{\mathrm{target}}\|^2 \Big). \] After centering external and implicit rewards within the prompt group, CRD minimizes \(\mathcal L_{\mathrm{CRD}}\) and adds a KL anchor term. The total loss is \[ \mathcal{L}_{\mathrm{CRD\_KL}} = \mathcal{L}_{\mathrm{CRD}}+\mathcal{L}_{\mathrm{KL}}. \] The two EMA models are then updated separately: \[ \theta_{\mathrm{old}} \leftarrow \eta_{\mathrm{old}}\theta_{\mathrm{old}} + (1-\eta_{\mathrm{old}})\theta, \qquad \theta_{\mathrm{samp}} \leftarrow \eta_{\mathrm{samp}}\theta_{\mathrm{samp}} + (1-\eta_{\mathrm{samp}})\theta. \] In the reported implementation, the base model is Stable Diffusion 3.5 Medium, LoRA rank is \(r=32\) with \(\alpha=64\), resolution is \(512\times512\), and the main experiments run on up to 4 NVIDIA H100 GPUs with runtime under 36 hours [2603.14128]. ## 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: \[ p^{(k+1)}(x\mid c) = \frac{1}{Z^{(k)}(c)}\, p^{(k)}(x\mid c)\, \exp\!\left(\frac{1}{\beta} r(c,x)\right), \] which unrolls to \[ p^{(K)}(x\mid c) \propto p^{(0)}(x\mid c)\, \exp\!\left(\frac{K}{\beta} r(c,x)\right). \] Second, if the old/reference model tracks the current model too aggressively, the ratio signal collapses because \(\log(p_\theta/p_{\mathrm{old}})\approx 0\). 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: \[ \beta_{\mathrm{init}} \mathrm{KL}\!\left(p_\theta(\cdot\mid c)\;\|\;p_{\phi}^{\mathrm{CFG}}(\cdot\mid c)\right), \] implemented through a diffusion-space approximation \[ \mathcal{L}_{\mathrm{KL}} = \frac{1}{K}\sum_{i=1}^K \hat\beta_{\mathrm{init}}(c,x^i)\, \mathbb{E}_{t,\epsilon}\!\left[\lambda(t)\, \left\|v_\theta(x_t^i,t\mid c)-v_{\phi}^{\mathrm{CFG}}(x_t^i,t\mid c)\right\|_2^2 \right]. \] 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: \[ \hat{\beta}_{\mathrm{init}}(c,x) = r_{\mathrm{raw}}(c,x)\,\beta_{\mathrm{init}}. \] 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 [2603.14128]. ## 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 [2603.14128]. | 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 \(0.63\), Flow-GRPO (w/ KL) \(0.95\), AWM (w/ CFG) \(0.86\), DiffusionNFT \(0.92\), and CRD \(0.93\). 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 \(5.44\) against SD3.5-M’s \(5.39\). Its unseen preference metrics on DrawBench are ImageReward \(0.98\), PickScore \(22.48\), and HPSv2.1 \(0.284\). On OCR, CRD reaches \(0.92\), matching Flow-GRPO (w/ KL), while AWM (w/o CFG) and DiffusionNFT reach \(0.97\) but have negative ImageReward, \(-0.64\) and \(-0.81\) respectively. The paper presents this pattern as evidence of reduced reward hacking rather than maximal single-reward optimization [2603.14128]. The appendix reports that CRD is largely insensitive to centering temperature \(\tau\), that larger group size \(K\) improves convergence and final reward, and that even \(K=2\) yields meaningful learning. Moderate CFG at inference often improves preference metrics further: in the GenEval setting, ImageReward rises to \(1.10\) with CFG sampling \(3.0\); in the OCR setting, it rises to \(1.04\) with CFG sampling \(3.0\) [2603.14128]. ## 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 \(k=8\) teacher responses per prompt, computes verifier-based rewards, and normalizes rewards within each prompt group by \[ A_j = \frac{r_j - \mu}{\sigma + \epsilon}, \] 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 [2507.00054]. 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_t(\hat{y}_t) \;=\; r_t(\hat{y}_t) \;-\; \frac{\lambda}{C}\sum_{k=1}^{C} \log \pi_{\text{ref}}(\hat{y}_t \mid x'_k, y_{