PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion Reinforcement Learning via Centered Reward Distillation

Yuanzhi Zhu, Xi Wang, Stéphane Lathuilière, Vicky Kalogeiton

arXiv:2603.14128v1cs.CVcs.AIcs.LG

TL;DR

Diffusion RL must address behaviors that pretraining objectives weakly specify, while existing approaches face limitations including rollout costs and distribution drift. CRD uses within-prompt centered reward matching with techniques to control drift; on GenEval and OCR it achieves competitive, fast, stable optimization with reduced reward hacking.

  • Problem

    Pretraining objectives weakly specify practical behaviors such as fine-grained prompt fidelity and legible text, while some diffusion RL methods have limited applicability to non-differentiable rewards.

    Source: §1 · pp. 1–2

  • Method

    CRD uses within-prompt centering to cancel the unknown normalizer in reward matching, with stabilization techniques that control distribution drift.

    Source: §1 · p. 3

  • Results

    CRD achieves competitive reward optimization with faster, more stable training and reduced reward hacking on GenEval and OCR tasks.

    Source: §1 · p. 4

  • Takeaways & Limitations

    CRD’s centered objective unifies prior reward-distillation methods and GVPO-style objectives as special cases.

    Source: §1 · p. 3

    Backward-process MDP methods face timestep-dependent variance, rollout overhead, and credit-assignment challenges that limit training efficiency.

    Source: §2.1 · p. 5

Abstract

from arXiv · show

Diffusion and flow models achieve State-Of-The-Art (SOTA) generative performance, yet many practically important behaviors such as fine-grained prompt fidelity, compositional correctness, and text rendering are weakly specified by score or flow matching pretraining objectives. Reinforcement Learning (RL) fine-tuning with external, black-box rewards is a natural remedy, but diffusion RL is often brittle. Trajectory-based methods incur high memory cost and high-variance gradient estimates; forward-process approaches converge faster but can suffer from distribution drift, and hence reward hacking. In this work, we present \textbf{Centered Reward Distillation (CRD)}, a diffusion RL framework derived from KL-regularized reward maximization built on forward-process-based fine-tuning. The key insight is that the intractable normalizing constant cancels under \emph{within-prompt centering}, yielding a well-posed reward-matching objective. To enable reliable text-to-image fine-tuning, we introduce techniques that explicitly control distribution drift: (\textit{i}) decoupling the sampler from the moving reference to prevent ratio-signal collapse, (\textit{ii}) KL anchoring to a CFG-guided pretrained model to control long-run drift and align with the inference-time semantics of the pre-trained model, and (\textit{iii}) reward-adaptive KL strength to accelerate early learning under large KL regularization while reducing late-stage exploitation of reward-model loopholes. Experiments on text-to-image post-training with \texttt{GenEval} and \texttt{OCR} rewards show that CRD achieves competitive SOTA reward optimization results with fast convergence and reduced reward hacking, as validated on unseen preference metrics.

1 Introduction

The introduction presents CRD as a forward-process diffusion RL framework that makes reward matching well-posed through within-prompt centering. It pairs this objective with drift-control techniques and reports competitive, faster, more stable optimization with reduced reward hacking.

1 Introduction

  • Diffusion RL seeks to optimize pretrained models using external rewards, including non-differentiable signals such as OCR accuracy that are difficult to differentiate through.The introduction notes that earlier approaches could require expensive data collection or differentiable rewards, limiting applicability to some signals.Source: §1 · p. 2
  • Trajectory-based RL methods can require storing full SDE trajectories and incur high-variance signals, memory and compute overhead, while forward-process methods can drift and amplify reward hacking.Forward-process objectives are described as simpler and more efficient, with lower-variance gradients that often converge faster, but moving references can enable iterative drift.Source: §1 · p. 3
  • Within-prompt centering cancels the unknown prompt-dependent normalizer, making reward matching well-posed and unifying prior reward-distillation and GVPO-style objectives as special cases.CRD matches an implicit log-density ratio, approximated by a diffusion ELBO, to centered external rewards within each prompt group.Source: §1 · pp. 3–4
  • CRD controls distribution drift by decoupling sampling from the moving reference, anchoring to a fixed CFG-guided pretrained model, and adapting KL strength to rewards.The stated goals include avoiding training instability, aligning regularization with inference-time semantics, accelerating early learning, and limiting late-stage reward-loophole exploitation.Source: §1 · pp. 3–4
  • CRD achieves competitive reward optimization with faster, more stable training and reduced reward hacking on GenEval and OCR tasks.The introduction also describes its performance as competitive against recent diffusion RL baselines.Source: §1 · p. 4

2 Related Works

The related work contrasts trajectory-based GRPO with forward-process diffusion RL, which separates sampling from training and can train faster. The paper places its method in the forward-process family while targeting reward hacking without sacrificing efficiency.

2 Related Works

  • Diffusion and flow models have motivated RL post-training to align generations with human preferences or task-specific rewards.Source: §2 · p. 4

2.1 GRPO in Visual Generation

  • GRPO variants broaden exploration, stabilize timestep-wise weighting, improve sampling efficiency, or enrich rewards, but add design sensitivities and retain backward-MDP efficiency limits.Reported trade-offs include sensitivity to branching or chunk choices, reward-proxy quality, calibration, and reward or model misspecification.Source: §2.1 · pp. 4–5

2.2 Diffusion RL Based on the Forward Process

  • Forward-process methods train on synthetically noised versions of generated samples rather than storing full sampling trajectories, and prior work reports faster training than backward-process methods.GPO connects ELBO-based likelihood estimation to advantage-weighted log-likelihood; AWM addresses estimator variance, with GPO's first-order approximation becoming more accurate near the reference model.Source: §2.2 · pp. 5–6
  • Other forward-process work includes ODE-sampler efficiency and reward-directed velocity updates, group-level preferences, probability-based cross-entropy objectives, and systematic estimator studies.DiffusionNFT targets an optimal velocity direction; DGPO uses within-group preferences, while GDRO derives a cross-entropy objective related to InfoNCA.Source: §2.2 · p. 6
  • The authors' forward-process method combines computational efficiency with new objectives and practical techniques intended to prevent reward hacking.Within-prompt centering cancels the prompt-dependent normalizer; the method also uses a fixed CFG-guided pretrained-model KL penalty.Source: §2.2 · p. 6

3 Diffusion RL using Reward Distillation

The section frames diffusion RL as reward maximization constrained by closeness to a reference, then develops CRD to match centered rewards without needing an intractable prompt-dependent normalizer. It also introduces drift-control techniques for stable fine-tuning.

3.1 Background

  • Under KL-regularized reward maximization, rewards encode a scaled log-density ratio plus a prompt-dependent normalizer, making direct absolute-reward regression ill-posed when that normalizer is unknown.The paper estimates the log-density ratio with a diffusion ELBO surrogate and uses group-normalized rewards.Source: §3.1 · p. 7

3.2 Centered Reward Distillation (CRD)

  • Within-prompt centering cancels the unknown prompt-dependent normalizer, making reward matching well-posed and motivating CRD’s objective.CRD matches centered external rewards to centered implicit model rewards.Source: §3.2 · p. 8
  • CRD recovers two-sample reward distillation and GVPO as special cases, and its ratio-based variant recovers InfoNCA under appropriate parameterization.The ratio-based variant is introduced in Appendix G.Source: §3.2 · p. 8
  • Decoupling the sampler from the moving log-ratio reference keeps data collection near on-policy while preserving a meaningful ratio signal.An aggressively updated reference can drift too close to the current model and collapse the log-ratio signal toward zero.Source: §3.2 · p. 9

3.3 Reward-Adaptive CFG based KL Regularization

  • CFG-guided pretrained-model anchoring aligns regularization with inference semantics and is intended to reduce reward hacking.A non-CFG reference can be weaker and pull training toward low-quality regions; without the RL gradient, the CFG anchor is functionally equivalent to CFG distillation.Source: §3.3 · p. 10
  • Reward-adaptive KL strength weakens anchoring on low-reward samples and strengthens it on high-reward samples, empirically accelerating training without sacrificing final performance.The raw reward is normalized to [0, 1] to scale the initial-reference KL coefficient.Source: §3.3 · p. 10

4 Experiments

The experiments test CRD on compositional image generation and visual text rendering, then examine how its KL and old-model choices affect training and generation quality. CRD achieves strong task-reward results alongside preference scores, while the ablations expose trade-offs between optimization speed and quality.

4.1 Experiment Setup

  • Experiments fine-tune Stable Diffusion 3.5-Medium on GenEval and OCR rewards, with image-quality and preference metrics evaluated on DrawBench prompts.The main results use group size K = 24; ablations use K = 6.Source: §4.1 · p. 11, Table 1 · p. 12

4.2 Main Results

  • CRD scores 0.93 GenEval versus SD3.5-M’s 0.63 and reaches 0.92 OCR accuracy, matching Flow-GRPO (w/ KL), with positive preference scores on unseen prompts.Flow-GRPO reaches the highest GenEval score (0.95); CRD attains the best Aesthetics (5.44) and improves Aesthetics over SD3.5-M (5.39).Source: Table 1 · p. 12, §4.2 · p. 12
  • CRD preserves SD3.5-M photorealism while improving prompt fidelity, especially in multi-object composition and integrated text, with more legible text and coherent object relations.The comparisons also report less natural backgrounds for Flow-GRPO, oversaturation for AWM (w/ CFG), and uniform backgrounds or blurred text regions for DiffusionNFT.Source: §4.2 · pp. 12–13
  • CRD converges significantly faster than Flow-GRPO variants, though it underperforms on some metrics; moderate test-time CFG further improves preference scores.The discussion attributes the CFG result possibly to stronger text conditioning and insufficient CFG distillation during RL fine-tuning.Source: §4.2 · p. 13

4.3 Ablation Studies

  • Fast old-model decay yields faster convergence but unstable KL divergence with model ϕ, and its generations collapse to flat-background text that satisfies OCR while losing scene coherence.The passage associates this instability and large KL with severe reward hacking.Source: §4.3 · pp. 13–14
  • Smaller initial KL strength accelerates optimization, while larger initial KL typically produces better perceptual quality.The ablation varies initial KL coefficients and CFG values.Source: §4.3 · p. 14, Figure 4 · p. 14
  • Adaptive initial-KL weighting improves training progress with minimal degradation in generation quality.The result is supported by reward curves and qualitative comparisons.Source: §4.3 · p. 14, Figure 4 · p. 14, Figure 5 · p. 14

5 Conclusion

The conclusion presents CRD as a forward-process diffusion RL framework whose within-prompt centering yields well-posed reward matching. It reports competitive optimization on GenEval and OCR with fast, stable training, and the appendices document related theory, methods, experiments, and further results.

5 Conclusion

  • CRD uses within-prompt reward centering to cancel the intractable prompt-dependent normalizer, yielding well-posed reward-matching objectives that unify prior methods as special cases.The conclusion also credits practical techniques targeting distribution drift and reward hacking with supporting competitive, fast, stable optimization on GenEval and OCR and competitive unseen-metric results.Source: §5 · p. 14

Appendix for CRD

  • Appendices A–H cover limitations and future work, related work, experimental details, theoretical discussion, accumulated tilting, special cases, InfoNCA connections, and additional results.The listed topics include GVPO and Reward Distill as special cases of CRD, and ratio-based reward distillation.Source: Appendix for CRD · p. 22

A Limitations and Future Work

The paper identifies open theoretical guarantees, dependence on pretrained-model coverage, and sampling costs as limitations of CRD. Its related-work discussion positions CRD as a unifying text-to-image framework while distinguishing its empirical focus from VMPO.

A Limitations and Future Work

  • CRD lacks rigorous convergence and stability guarantees under practical approximations, and self-generated sampling limits leverage when the pretrained model rarely produces a desired behavior.Continued self-training can further narrow diversity; proposed mitigations include mixing in SFT updates or adding data anchoring.Source: A Limitations and Future Work · p. 22

B More Discussion with Related Works

  • CRD is tailored to text-to-image diffusion, and the paper presents it as a unified framework subsuming GVPO and prior two-sample reward-distillation methods.GVPO is described as focused on gradient-variance analysis and primarily designed for LLM post-training.Source: B More Discussion with Related Works · pp. 22–23
  • VMPO minimizes variance of log importance weights rather than directly optimizing a KL-based objective; at submission, its experimental scope was narrower than CRD’s.The paper describes forward-process methods as complementary components of a broader recipe for faster training while mitigating reward hacking.Source: B More Discussion with Related Works · p. 23

C Experimental Details

The appendix connects DiffusionNFT to KL-regularized reward maximization and analyzes how repeated online reference updates can compound reward tilting. It also describes special-case connections and experimental implementation details.

C Experimental Details

  • CRD is implemented on the DiffusionNFT codebase; AWM and DiffusionNFT are compared under the same training budget at early checkpoints because both baselines suffer reward hacking.The experiments use the adaptive weighting variant of the implicit model reward.Source: C Experimental Details · p. 23

D.1 DiffusionNFT Score Update Implies an Exponential Tilt

  • DiffusionNFT’s population-optimal density is an exponential tilt of the old model by its optimality score.The appendix derives this density-level form from the score update; at the data level, the optimality score is identified with the training reward.Source: §D.1 · pp. 24–25, §D.2 · p. 25

D.2 Reduction to KL-regularized RL at t = 0

  • With a fixed reference, DiffusionNFT is equivalent at the population optimum to KL-regularized reward maximization targeting the same exponentially tilted distribution.The effective inverse-temperature is λ = 2/β; repeated online reference updates can accumulate tilt.Source: §D.2 · p. 25

E Online RL Leads to Reward Hacking via Accumulated Tilting

  • Online training with a moving reference applies reward tilting epoch by epoch, a process that can produce increasingly peaked policies and potential reward hacking.The section idealizes each epoch as reaching the KL-regularized population optimum under a fixed reward.Source: E Online RL Leads to Reward Hacking… · pp. 25–26, §E.1 · p. 26

E.1 Idealized Recursion under Exact Per-epoch Optima

  • Under exact per-epoch optimization and a fixed reward, online training compounds reward tilt, with effective inverse-temperature scaling linearly with epochs K.The analysis assumes a fixed reward and population-optimal updates at each epoch.Source: §E.1 · p. 26

E.2 Concentration and Reward Hacking

  • As K grows, the distribution concentrates on reward maximizers; with an imperfect reward model, this sharpening can amplify reward exploitation.The concentration claim assumes mild regularity, including initial mass near the reward-maximizer set.Source: §E.2 · p. 26

E.3 Remarks on EMA References

  • EMA references do not follow the exact multiplicative recursion, but without anchoring to the initial reference, repeated updates can still accumulate tilt and increase reward-hacking risk.EMA typically reduces the rate at which the tilt coefficient grows compared with fully replacing the reference.Source: §E.3 · pp. 26–27

F GVPO and Reward Distill as Special Cases

G Ratio-based Reward Distillation and Connection to InfoNCA

This section formulates reward distillation as matching normalized distributions over finite candidate sets, avoiding the unknown reward normalizer. It also connects the objective to divergence minimization and InfoNCA.

G Ratio-based Reward Distillation and Connection to InfoNCA

  • When model reward matches the true reward up to an x-independent additive constant, the two candidate distributions coincide.The section proposes fitting the model-induced distribution to the reward-induced distribution by minimizing a divergence over sets.Source: G Ratio-based Reward Distillation and… · p. 28
  • Normalizing the reward-induced and model-induced distributions over a candidate set removes dependence on the unknown Z(c), turning reward matching into distribution matching.The formulation compares candidate probabilities induced by the true reward and the model reward.Source: G Ratio-based Reward Distillation and… · p. 28
  • With uniform weights and forward KL, the distribution-matching objective recovers InfoNCA as cross-entropy from the reward-induced teacher to the model-induced student.The formulation also permits other divergences, including reverse KL and Jensen–Shannon.Source: G Ratio-based Reward Distillation and… · p. 28

H Additional Results

The section uses Best-of-N (BoN) as an inference-time reference for text-to-image reward performance. It reports that RL fine-tuning can match or exceed this reference without multiple test-time samples.

H.1 Best-of-N (BoN) Performance of SD3.5-M and SD1.5

  • RL fine-tuning can match or exceed BoN performance on GenEval without requiring multiple samples at test time.Figure 6 compares BoN performance as N increases with the average reward score of the current N samples for SD3.5-M and SD1.5.Source: §H.1 · p. 29
  • BoN generates N independent samples and selects the sample with the highest reward score.The section reports BoN performance for SD3.5-M and SD1.5 on GenEval as an inference-time reference.Source: §H.1 · p. 29
  • In Figure 6, the orange curve shows BoN performance as N increases, while the blue curve shows average performance across the current N samples.The reported BoN curves are loglinear for both models.Source: §H.1 · p. 29

H.2 Multi-reward Training

CRD is reported to perform competitively across metrics in multi-reward training, while CFG sampling trades a slight reduction in aesthetic quality at higher guidance values for improved CLIPScore. Adaptive KL regularization is also reported to mitigate reward hacking on DiffusionNFT while preserving visual quality.

H.2 Multi-reward Training

  • CRD achieves competitive performance across task metrics and preference scores in multi-reward training on DrawBench prompts.The reported metrics are PickScore, CLIPScore, HPSv2.1, Aesthetics, and ImageReward; Table 4 compares CRD with Flow-GRPO and Choi et al.Source: §H.2 · p. 29, Table 4 · p. 30
  • CFG sampling further improves CLIPScore, with a slight cost to aesthetic quality at higher CFG values.Source: §H.2 · p. 29

H.3 Adaptive KL Regularization on DiffusionNFT

  • On DiffusionNFT, adaptive KL regularization mitigates reward hacking while preserving visual quality, including fewer white backgrounds and high reward scores.Figure 7 presents this result.Source: §H.3 · p. 29, Figure 7 · p. 30

H.4 Additional Ablations

The ablations find that CRD is largely insensitive to temperature, while larger group sizes improve convergence and final rewards. The v-prediction adaptive weighting achieves the highest evaluation and sampling-model rewards while keeping KL divergence stable.

H.4 Additional Ablations

  • Reward curves remain remarkably similar across τ values from 1.0 to 100.0 and ∞, so the other experiments use uniform weighting (τ=∞).The reported insensitivity holds across both small and large initial KL coefficients βinit.Source: §H.4 · p. 30
  • Larger group sizes lead to faster convergence and higher final rewards; K=24 reaches near-perfect OCR scores, while K=2 still enables meaningful learning.The passage attributes the richer gradient signal per update to larger groups.Source: §H.4 · p. 31
  • The v-prediction adaptive weighting achieves the highest evaluation and sampling-model rewards while maintaining stable KL divergence throughout training.The original adaptive weighting uses an x0 representation, introducing an extra coefficient t compared with the estimator in Eq. (13).Source: §H.4 · p. 31, Figure 10 · p. 31

H.5 Additional Qualitative Results

Additional qualitative results show improving text rendering over training without collapse, alongside limits in high-CFG generations and colored-zebra outputs. The visualizations also compare CFG settings and show generations across checkpoints, prompts, and random seeds.

H.5 Additional Qualitative Results

  • Training checkpoints show text rendered more accurately over time without model collapse.Samples are shown from step 0 at 200-step intervals.Source: §H.5 · p. 31, Figure 11 · p. 32
  • Large CFG consistently degrades visual fidelity; text prompts can develop broken glyphs, while compositional prompts may follow instructions less reliably.Higher CFG can produce typographic artifacts even when the image appears more aligned with the prompt.Source: §H.5 · p. 31, Figure 12 · p. 33
  • At large CFG values, a photorealistic cow shifts toward an illustration-like style with simplified textures and exaggerated features.The cow appears in the final column of the visualization.Source: §H.5 · p. 32, Figure 12 · p. 33
  • Despite a roughly uniform training color distribution, the model struggles to produce convincing blue and green zebras, and usually replaces their black stripes with the target color.The zebra generations vary across colors and are shown for different random seeds under the same prompt template.Source: §H.5 · pp. 32–33, Figure 13 · p. 33

Related papers

Loading 2603.14128v1…