Most reward-based post-training of diffusion models leaves intermediate denoising steps unspecified: an image-level reward tells you which final samples are better, but not how each denoising prediction should change. The core insight of this work is that converting image-level reward signals into explicit, bounded per-query targets—and refreshing those targets on the student’s own rollouts—makes alignment both more effective and more analyzable.
Key Findings
- Converting image-level reward gradients into bounded positive and negative targets around anchors yields dense, per-query supervision that the model can directly fit.
- Separating target construction from finite realization exposes when larger constructed gains fail to translate after a single fitting update, enabling diagnostic insight into training dynamics.
- Across two backbones (SD 3.5-M and a step-distilled Z-Image-Turbo) and ten evaluators, the method achieves the top held-out scores in 19 of 20 reward-matched settings, outperforms the strongest competing method by up to 44.0%, and reduces GPU-hours versus DiffusionNFT by ~40% (SD 3.5-M) and ~63% (Z-Image-Turbo).
Method overview
- Use a frozen behavior policy to generate on-policy denoising trajectories and collect query states and anchors.
- Build bounded positive/negative targets at each query from reward gradients around the anchor prediction.
- Train the student to fit these detached targets via a finite number of fitting updates; then refresh the behavior policy with an exponential moving average (EMA) of the student.
- Repeat: construction of targets and finite fitting are alternated so that target quality and realizable gains can be measured independently.
Who it's for and trade-offs
Great fit if you want diagnosable, compute-efficient post-training alignment of image-generation diffusion models and can run on-policy rollouts and short fitting cycles. Look elsewhere if you cannot afford on-policy sampling at scale, have no reward signal to shape targets, or require purely off-policy supervised fine-tuning without changing the rollout distribution. The approach introduces hyperparameters for target bounds and fitting steps that must be tuned; gains depend on the interplay between target construction and the student’s capacity to realize those targets.