DiffusionOPSD: On-Policy Self-Distillation for Reward Post-Training of Diffusion Models

The supervision gap DiffusionOPSD is built to close

Reward optimization for diffusion models has a structural timing problem. The reward model scores a finished image, but the denoising network makes dozens of intermediate predictions before that image exists. The README frames this as a supervision gap at intermediate predictions: the image-level score arrives only after a multi-step trajectory, so every intermediate state is trained without knowing how the final sample will be judged. DiffusionOPSD's stated goal is to close that gap by constructing explicit intermediate targets rather than waiting for the endpoint reward to propagate backwards through the whole trajectory. The intended user is someone doing post-training research on text-to-image models who already has a reward model and a working sampling pipeline. The README supports two backbones, SD3.5-M at 512x512 and Z-Image-Turbo at 1024x1024 in its native few-step regime, which tells you the target is a lab with enough GPU memory to fine-tune a mid-size diffusion transformer. This is not a tool for prompt engineering or inference-time guidance.

Anchors, bounded targets and the EMA refresh loop

The mechanism has four moving parts and the README gives the equations for each. At a query state s, defined by a caption c, a noisy latent z_q and a noise level sigma_q, the frozen behaviour policy produces a clean-output anchor y_0 = z_q minus sigma_q times the old velocity prediction. That anchor is detached. Around it, the code runs a normalized reward-gradient ascent and descent inside what the README calls a relative trust region, producing a positive target y_+ and a negative target y_-. The normalization divides by the gradient norm plus epsilon, and the step size h bounds how far each target can move from the anchor. The trainable policy then fits those detached targets through two branches, a positive branch and a negative branch, combined by a weight omega in the loss. The implementation detail that matters here is that the reward and decoder graphs are discarded before the policy update, so the gradient the policy sees comes from a mean squared residual against a fixed target, not from backpropagating through the reward model. An EMA of the behaviour policy is refreshed after each outer update, which regenerates trajectories, anchors and targets for the next round. The README states this separation makes target construction and finite realization independently observable, which is the actual research claim: you can measure whether a bad result came from bad targets or from the policy failing to fit good ones.

What the training loop actually optimizes

The loss combines a positive residual and a negative residual, each normalized by a detached mean-absolute residual normalizer, with elementwise-mean squared residuals inside. The README specifies c_adv = 5 as the multiplier on the whole objective. The pseudocode in the README is short enough to read as a contract: roll out the behaviour policy on a prompt, select a low-noise state from the trajectory, detach the clean-output anchor, compute a group-normalized endpoint weight from the trajectory reward, build bounded positive and negative targets, run the trainable policy's clean-output prediction on the detached query, compute the loss, backpropagate, step, then update the behaviour policy by EMA. Two things are worth noting. First, the query selection is restricted to low-noise states, which is consistent with the anchor formula: at high noise the one-step clean estimate is unreliable, so the method concentrates supervision where the estimate is meaningful. Second, the endpoint weight is group-normalized, meaning rewards are compared within a group of trajectories rather than in absolute terms. That is a design choice with a consequence: if your reward model produces near-identical scores across a batch, the normalized weight carries little signal.

Getting it running and the configuration surface

The README does not include an installation section, a requirements file listing, or a quickstart command block in the material provided, so the exact invocation cannot be quoted here. What the README does state is that the release supports single-reward and mixed-reward training, that public presets cover all seven open-weight evaluators, and that arbitrary positive weighted sums are allowed. The paper example is given as PickScore weighted at 26 plus CLIPScore plus HPSv2.1, which is the configuration to look for in the preset directory if you want to reproduce the reported setting. The backbone choices are SD3.5-M at 512x512 and Z-Image-Turbo at 1024x1024. The pseudocode names the functional units you should expect to find in the codebase: rollout, select_low_noise_state, clean_output, group_normalized_endpoint_weight, bounded_reward_ascent, bounded_reward_descent, detached_target_loss, update_behavior_policy_ema. If those names do not appear in the source tree, the README and the code have drifted apart. Treat that as the first thing to check after cloning, because the repository has no releases and the only version you can get is whatever is on main at the time you pull.

The cost of holding two policies at once

The frozen behaviour policy and the trainable policy both have to be resident during a training step, and the behaviour policy is refreshed by EMA rather than re-loaded, so you are paying for two copies of the model plus optimizer state plus activations for both the rollout and the fitting pass. The README's efficiency claim is relative: it reports that DiffusionOPSD reduces training GPU-hours relative to DiffusionNFT by 40% on SD3.5-M and 63% on Z-Image-Turbo. That is a comparison against one specific baseline, not a statement about absolute cost, and the README does not give the absolute GPU-hour figures in the material provided. The method also depends on a differentiable reward model, because the positive and negative targets are built from reward gradients with respect to the clean output y. A reward model that is only available as a black-box scoring endpoint cannot be used for target construction without changing the approach. That is the sharpest constraint in the design: the reward has to be differentiable, not just callable.

Where it sits next to reward fine-tuning and distillation

The obvious alternative is direct reward fine-tuning, where you backpropagate the reward gradient through the sampling trajectory into the policy. DiffusionNFT appears in the README as the GPU-hour baseline for that family, and the README's comparison is that DiffusionOPSD is faster on the two backbones tested. The difference in approach is where the gradient flows. Direct reward fine-tuning carries the reward gradient through the denoising steps, so the reward model sits inside the training graph and memory scales with trajectory length. DiffusionOPSD computes the reward gradient only at the clean-output anchor, converts it into a bounded target, detaches it, and then trains the policy against that fixed target with a squared-error loss. The reward model is never inside the policy's backward pass. The trade-off is that you lose the exact gradient of the endpoint reward with respect to the policy parameters and replace it with a supervised fit to a constructed target. If the target is wrong, the policy will fit it faithfully and still be wrong. That is precisely why the README emphasizes that the separation makes target quality and realization independently measurable; the method's value depends on that measurement being something you actually intend to do.

Maintenance, licensing and what the repository does not ship

The repository is Apache-2.0 and is not archived, with a last push dated 2026-08-27. There are no releases retrieved, so there is no tagged version to pin, no changelog to read, and no upgrade path other than tracking main. That matters for anyone planning to build on top of it: pin a commit hash rather than a branch name, because the only versioning signal available is the commit history. Apache-2.0 permits commercial use and modification and includes a patent grant, but it also requires that you preserve copyright and licence notices in redistributed copies and state significant changes. None of this is legal advice; if you plan to ship a model trained with this code, read the licence text and check the terms of the reward models you use, since those carry their own licences. The README also states in its opening line that this is an external implementation of the algorithm in the paper, not the authors' own code. That single sentence should shape how you read everything else: the equations are the specification, and the code is one group's reading of them.

Editorial conclusion

DiffusionOPSD is for research groups that already have a reward-model evaluation harness and want to study how target construction and finite realization behave separately; it is not for teams that need a stable editing API next week, since the repository is an algorithm implementation with no released checkpoints and no packaging. Before committing, verify three things in the repository itself: whether the training entry point exposes the reward-weight preset you intend to use, whether your hardware can hold a frozen behaviour policy and a trainable policy at the same time for your chosen backbone, and whether the code you pull matches the paper's equations, because the README states plainly that this is an external implementation rather than the authors' release.

Official sources

1. Issues
2. License: Apache-2.0
3. Project website
4. README
5. worldbench/DiffusionOPSD on GitHub