Post-Training & Alignment

---

A pair of trajectories from one noise, every step bent toward the reward-weighted image difference — the whole sampling run as one action.

• Why per-step updates are mostly noise (§1, §3.2). In an MDP each stochastic step is an action reinforced by the trajectory's advantage: a small part of every update helps the reward, the rest drifts what no reward watches — styles cycle, grid artifacts appear past ~500 epochs.

• The whole run is one action (§4). Rollouts $x_{0:T}$, ${\hat{x}}_{0:T}$ from one noise $ϵ$ with mild stochasticity; $\Delta x={\hat{x}}_T-x_T$, $\Delta R=R({\hat{x}}_T)-R(x_T)$, and the velocities at all steps of both bend toward $\Delta R\Delta x$, which points from the worse image to the better — no step credited for its own noise.

• Why an endpoint direction works mid-trajectory (§4.1, App. C.4). The update needs $\nabla R(x_T{)}^{\top }{J}_i[\Delta R\Delta x]\ge 0$, $J_i$ the Jacobian of the flow from step $i$; it holds in expectation if $J_i$ is positive semi-definite — exact for optimal transport, which diffusion flows approximate, SDEdit's bet too.

• A stochastic sampler that stays on schedule (Alg. 1, B.1). Overshoot the Euler step to lower noise, then add fresh noise back to $t_{i+1}$ — EDM's trick, so the network's time stays right; ${\gamma }_i$, the noise fraction re-drawn, is 0.0025 everywhere. Flow-GRPO alone: 71.06 → 76.74 (Fig. 8 B):

$${\tilde{t}}_{i+1}=\frac{t_{i+1}}{1-{\gamma }_i{t}_{i+1}+{\gamma }_i},{\tilde{x}}_{i+1}=x_i+({\tilde{t}}_{i+1}-t_i)v_i,x_{i+1}=\frac{{\tilde{x}}_{i+1}+{\tilde{t}}_{i+1}\sqrt{{\gamma }_i^2+2{\gamma }_i}{ϵ}_i}{{\gamma }_i{\tilde{t}}_{i+1}+1}$$

• Normalize the difference (§4.3). $\Delta R$ is, to first order, proportional to $\Vert \Delta x\Vert$, so $\Delta R\Delta x$ counts the pair's distance twice; the signal used is $\overline{\Delta x}=\Delta x/(\Vert \Delta x{\Vert }_{\mathrm{R}\mathrm{M}\mathrm{S}}^2+{10}^{-6})$, the difference over its own mean square. Without it 82.13 → 79.74 and training turns unstable (Fig. 8, H).

• Clipping without a likelihood (Alg. 2). With no likelihood for PPO's ratio, each stored velocity gets the target $v_{\mathrm{t}\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{e}\mathrm{t}}=v_{\mathrm{r}\mathrm{e}\mathrm{f}}-\overline{\Delta x}$ and a proxy ratio $\rho$ from two regressions; SPO then penalizes straying far from the rollout's velocity, $\epsilon =0.03$. PPO ties (G, 81.99) only under the interval schedule:

$$\rho =\exp (\Vert v_{\mathrm{t}\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{e}\mathrm{t}}-v_{\mathrm{r}\mathrm{e}\mathrm{f}}{\Vert }^2-\Vert v_{\mathrm{t}\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{e}\mathrm{t}}-v_{\mathrm{c}\mathrm{u}\mathrm{r}}{\Vert }^2),L=-\Delta R\rho -\frac{|\Delta R|}{2\epsilon }(\rho -1{)}^2$$

• Batching, and what is switched off (§4.3, Tab. 1). 432 prompt pairs per epoch at $T=40$, 34,560 stored samples, 4 batches, Flow-GRPO's memory; LoRA rank 32 on SD3.5-M; EMA off (strictly harmful), KL and CFG off, so the algorithms and not the regularizers are compared.

• Uniform stochasticity wins (B.2, Fig. 20). Against constant ${\gamma }_i$ (U), an interval schedule perturbs one random window (I), a prior schedule only the seed (P), with optional gradient weighting: U 82.13 / 84.25 at 200 epochs / best, I 82.29 / 83.50 yet weighting-sensitive, P 70.79 / 80.51.

• Finite differences are a smoothed gradient (App. C). Perturb the step-$j$ image, $x^{\prime }~𝒩(x,{\sigma }^{2}I)$, finish both ODEs, form $\tilde{g}=(R(x_T^{\prime })-R(x_T))/{\sigma }^2·(x_T^{\prime }-x_T)$. With $f$ the flow from $j$, $J$ its Jacobian, Taylor and Stein's lemma make its expectation a smoothed gradient:

$${𝔼}_{ϵ}[(R(f(x+ϵ))-R(f(x)))(f(x+ϵ)-f(x))]\approx {𝔼}_{ϵ}[(R(f(x)+Jϵ)-R(f(x)))Jϵ]={𝔼}_{{ϵ}^{\prime }~𝒩(f(x),{\sigma }^{2}J{J}^{\top })}[R({ϵ}^{\prime })({ϵ}^{\prime }-f(x))]$$

$$={\sigma }^{2}J{J}^{\top }{𝔼}_{{ϵ}^{\prime }}[\nabla R({ϵ}^{\prime })]={\sigma }^{2}J{J}^{\top }\nabla \tilde{R}(f(x)),\tilde{R}=R*𝒩(0,{\sigma }^{2}J{J}^{\top })$$

• What the identity buys (App. C.3–C.5). $J^{\top }$ backpropagates the gradient of a reward smoothed by a Gaussian of covariance ${\sigma }^{2}J{J}^{\top }$ — discontinuity smoothed away — the extra $J$ distorts it, and §4.1's condition is a quadratic form in a near-positive-definite matrix, so it ascends:

$$𝔼[\nabla R(f(x){)}^{\top }J[\Delta R\Delta x]]={\sigma }^2\nabla R(f(x){)}^{\top }JJ{J}^{\top }\nabla \tilde{R}(f(x))\ge 0$$

• Experiments (§5). SD3.5-M on Pick-a-Pic prompts with PickScore, VLM alignment and their sum, against Flow-GRPO at equal cost by reward, OneIG and HPSv2: combined reward 30.3 to Flow-GRPO's 28.5, OneIG alignment 82.13 vs 71.06, one reward in 19× fewer GPU-hours.

• The exact gradient does not beat the difference (Fig. 8 I–J). The reward model's true gradient for $\Delta R\Delta x$ gives 81.47 against 82.13, the smoothing worth the gradient; pushed through the SDE steps it falls to 66.69 — one endpoint direction per step beats per-step credit.

• Four hours of a human as the reward (§5.6). One author rated 3,200 on-policy pairs in 4 hours, 50 epochs, HPDv2 prompts, no reward model or offline data — possible since one pairwise preference fixes an update direction, a budget the paper calls far too small for Flow-GRPO.

The release pins two things the paper leaves loose and breaks one of its own statements: the pair's shared seed is itself churned, so the pair differs from step one; every loss is a mean over pixels, not the appendix's sum; and the CFG path, KL term and weighting ship, all off in every run.

---

DDPO is denoising score matching with noisy targets; keep the pretraining loss instead and weight each sample by its advantage.

• The two objectives do not match (§3.2). In LLMs both stages maximize one sequence log-likelihood; in diffusion, pretraining is flow matching on the forward process (Eq. 5) while DDPO, which Flow-GRPO inherits, maximizes a per-step reverse-transition Gaussian likelihood (Eq. 3–4).

• DDPO is score matching in disguise (Thm 1, Eq. 6). The joint law of $(x_{t-\Delta t},x_t)$ is the same forward and reverse, so up to Euler–Maruyama error a DDPO step at $t$ targets $𝔼\Vert s_{\theta }(x_t,t)-\nabla \log p(x_t\mid x_{t-\Delta t}){\Vert }^2$ — denoising score matching conditioned on a noisy intermediate instead.

• The target is still the right one (Lemma 1). For any $s<t$, score matching conditioned on $x_s$ has the same population minimizer as the plain objective — the textbook clean-data result is the case $s=0$ — so what separates DDPO from pretraining is not bias but its noise.

• Noisier conditioning, wider target (Thm 2). Both conditional scores are unbiased, yet conditioning on $x_s$ inflates the target covariance by $\kappa (s,t)I$, where $\kappa (s,t)=(1-t{)}^2{s}^2/[t^2(t^2(1-s{)}^2-s^2(1-t{)}^2)]$ is zero at $s=0$, climbs strictly with $s$, and diverges as $s$ approaches $t$ from below.

• The cost shows without any RL (§3.3, Eq. 7–8). Replace EDM's clean target $\Vert D_{\theta }-x_0{\Vert }^2$ with the noisy proxy that mirrors DDPO and change nothing else: the run learns in the right direction but reaches every FID level later, so the penalty is the objective, not the RL machinery.

• A different MDP (§4). DDPO's state is $(c,t,x_t)$, its action $x_{t-1}$, its policy $p_{\theta }(x_{t-1}|x_t)$; AWM's state is the prompt $c$, its action the finished image $x_0$, its policy $p_{\theta }(x_0|c)$ — one decision per finished sample, not one per denoising step, which is what makes the trajectory disposable.

• The matching loss is the log-likelihood (Eq. 10). $\log {\pi }_{\theta }(x_0|c)$ is read through its ELBO surrogate $-{𝔼}_t[w(t)\Vert v_{\theta }(x_t,t,c)-(ϵ-x_0){\Vert }^2]$; the uniform $w(t)=1$ beats the ELBO weight $(1-t)/t$ of Eq. 2, visual quality and likelihood pulling apart here as they have ever since DDPM.

• GRPO on that surrogate (Eq. 9, 11–12). The ratio becomes $\exp (-{𝔼}_t[w(t)(\Vert v_{\theta }-v{\Vert }^2-\Vert v_{{\theta }_{\mathrm{o}\mathrm{l}\mathrm{d}}}-v{\Vert }^2)])$, with $t$ and $ϵ$ shared across the two policies to cut variance, and the KL a velocity-space $\Vert v_{\theta }-v_{{\theta }_{\mathrm{r}\mathrm{e}\mathrm{f}}}{\Vert }^2$ — $\beta =0.2$ collapses into instability, 2.0 crawls, and 0.4 to 1.0 holds both ends.

• What the gradient does, and what it frees. On-policy the update is $-A_i{\nabla }_{\theta }{𝔼}_t[w(t)\Vert v_{\theta }-(ϵ-x_0){\Vert }^2]$: a positive advantage pulls $v_{\theta }$ toward that target, a negative one away. Sampling and training decouple — 20 sampler steps against 4 training ones on its grid, any solver, no CFG.

• Experiments (§5). SD3.5-M and FLUX with LoRA on GenEval, OCR and PickScore, reward per GPU-hour vs Flow-GRPO: GenEval 0.95 matches it at 8.02× less compute, OCR 0.89 costs 17.6 GPU-hours against 415.9 (23.6×), PickScore 23.02 10.5× less; on FLUX 8.5× and 6.8×.

• Where $t$ is drawn from matters (§5.2, Fig. 7). Uniform and the sampler's grid tie; logit-normal, pretraining's default, lags, then degrades as RL runs on — mid-range mass is the wrong prior once the loss is advantage-weighted; reusing half the last policy's batches changes nothing.

The released configs disagree with the paper on the very knobs the paper reports: the loss is not a plain square, the sampler is not Euler–Maruyama, the training timesteps are six rather than four, and the trust region doing the work is a KL to an EMA copy the paper never names.

# scripts/train_sd3_awm.py with config/dgx_awm.py:geneval_sd3_no_cfg_1node.
# LoRA "default" is v_theta, LoRA "ema" is a slow copy of it, adapters off is v_ref.
x0 = sample["next_latents"][:, -1]  # the rollout, drawn by SA-Solver - not the paper's Euler-Maruyama
t = sample["timesteps"][0][idx] / 1000  # 6 of the sampler's 14 steps, the first dropped; paper says 4 of 20
xt = (1 - t) * x0 + t * noise

# The loss standing in for log pi(x0|c) is neither uniform nor squared: a generalized
# Huber at power 0.25, times t. The paper writes w(t) = 1 and a plain L2.
def log_p(v):  # v = v_theta(xt, t, c), and mse is meaned over every non-batch axis
    return -(((v - (noise - x0)) ** 2).mean(dims) ** 0.25) * t / 0.25

ratio = exp(log_p(v_theta) - log_p(v_old).detach())  # Eq. 11, one t and one noise for both
adv = clip(adv, -5, 5) / 5  # rescaled to +-1, and the std is global, not per prompt
policy_loss = maximum(-adv * ratio, -adv * clip(ratio, 0, 2)).mean()  # clip_range = 1, so barely a clip

# Two KL terms, and the one the paper names carries a thousandth of the weight.
loss = policy_loss + 0.001 * mse(v_theta, v_ref) + 1.0 * mse(v_theta, v_ema)
ema.data = decay * ema.data + (1 - decay) * theta.data  # decay = min(0.3, 0.001 * step)

---

Online RL on the forward process: reward splits the rollouts into implicit positive and negative policies, and one flow-matching loss contrasts them.

• Why the reverse process is a trap (§2.2). Diffusion admits no exact likelihood, so GRPO-style work discretizes the reverse SDE into tractable Gaussians — a policy gradient bought with first-order SDE samplers, stored trajectories and a model untied from its own forward process.

• Split the rollouts by reward (§3.1). Read $r\in [0,1]$ as the probability that a sample is good, $p(o=1\mid x_0,c)$. The $K$ images then weight two implicit policies, ${\pi }^{+}\propto r{\pi }^{\mathrm{o}\mathrm{l}\mathrm{d}}$ and ${\pi }^{-}\propto (1-r){\pi }^{\mathrm{o}\mathrm{l}\mathrm{d}}$: each image counts $r$ toward the first and $1-r$ toward the second, and ${\pi }^{+}\succ {\pi }^{\mathrm{o}\mathrm{l}\mathrm{d}}\succ {\pi }^{-}$.

• A direction, not a target (§3.2, Eq. 3). Improved velocity $v^{*}=v^{\mathrm{o}\mathrm{l}\mathrm{d}}+(1/\beta )\Delta$: the old velocity plus a correction $\Delta$ scaled by $1/\beta$ — CFG's shape, which the paper reads as the same move, done once and frozen: the conditional model as positive policy, unconditional as negative, $w$ as $1/\beta$.

• Why the velocities align (Thm 3.1, Eq. 4). Bayes' rule gives the positive-group probability $\alpha =({\pi }_t^{+}/{\pi }_t^{\mathrm{o}\mathrm{l}\mathrm{d}}){𝔼}_{{\pi }^{\mathrm{o}\mathrm{l}\mathrm{d}}}r$. The posterior mixture makes $v^{\mathrm{o}\mathrm{l}\mathrm{d}}=\alpha v^{+}+(1-\alpha )v^{-}$: the old velocity lies between the two at fixed noise level, position, and prompt. Thus $\Delta =\alpha (v^{+}-v^{\mathrm{o}\mathrm{l}\mathrm{d}})=(1-\alpha )(v^{\mathrm{o}\mathrm{l}\mathrm{d}}-v^{-})$ points toward the positive policy; substituting $\beta =\alpha >0$ into $v^{*}=v^{\mathrm{o}\mathrm{l}\mathrm{d}}+\Delta /\beta$ gives $v^{*}=v^{+}$, the reward-weighted positive policy.

• One loss, two implicit policies (Eq. 5–6). $ℒ=𝔼[r\Vert v_{\theta }^{+}-v{\Vert }^2+(1-r)\Vert v_{\theta }^{-}-v{\Vert }^2]$ over $v_{\theta }^{\pm }=(1\mp \beta )v^{\mathrm{o}\mathrm{l}\mathrm{d}}\pm \beta v_{\theta }$ has optimum $v^{\mathrm{o}\mathrm{l}\mathrm{d}}+(2/\beta )\Delta$ (Thm 3.2): flow matching that has merely seen a reward, the guidance already in it. Delete the negative term and the reward collapses at once (§4.4).

• Reward to probability (Alg. 1). $r=1/2+1/2\mathrm{c}\mathrm{l}\mathrm{i}\mathrm{p}[(r^{\mathrm{r}\mathrm{a}\mathrm{w}}-{𝔼}_{{\pi }^{\mathrm{o}\mathrm{l}\mathrm{d}}}{r}^{\mathrm{r}\mathrm{a}\mathrm{w}})/Z_c,-1,1]$ with $Z_c$ a reward std — group-relative like GRPO, but landing on a split probability rather than an advantage, which is what lets a single batch weight the two branches by $r$ and $1-r$ at every noise level.

• Off-policy by construction (§3.3). The sampler is an EMA copy, ${\theta }^{\mathrm{o}\mathrm{l}\mathrm{d}}\leftarrow {\eta }_i{\theta }^{\mathrm{o}\mathrm{l}\mathrm{d}}+(1-{\eta }_i)\theta$, so no importance ratio is needed; on-policy $\eta =0$ collapses and $\eta =0.9$ crawls, while the ramp ${\eta }_i=min(0.001i,0.5)$ holds both ends — only the OCR reward needs a 0.999 ceiling.

• Weighting and guidance strength (§3.3, §4.4). The distillation method DMD's self-normalized $x_0$ regression stands in for the weighting $w(t)$: $\Vert x_{\theta }-x_0{\Vert }^2/\mathrm{s}\mathrm{g}(\mathrm{m}\mathrm{e}\mathrm{a}\mathrm{n}|x_{\theta }-x_0|)$, beating every fixed schedule while $w(t)=1-t$ collapses; $\beta$ near 1 is stable, 0.1 buys a faster climb.

• Solver freedom, and no CFG (§3.2, §4.4). Data collection is a black-box call, so the 2nd-order ODE sampler it permits beats the first-order SDE GRPO is pinned to, clearest on noise-sensitive PickScore; CFG is dropped outright and the policy starts from the conditional model alone.

• Experiments (§4). CFG-free SD3.5-M (2.5B), five rewards in sequence, preferences first: at 1.7k iterations GenEval 0.94, OCR 0.91, PickScore 23.80, past guided SD3.5-L (8B) everywhere, FLUX.1-Dev (12B) on all but ClipScore; single rewards run 3–25× faster than FlowGRPO.

• Positives alone collapse (§4.4). Without $v^{-}$ the reward collapses at once — unlike language models, where rejection-sampling fine-tuning on positives is a strong baseline — so the negative half is what diffusion RL cannot do without — the finding against the LLM playbook.

The released code fills in what the paper leaves out: the three policies are three rank-32 LoRA states of one frozen backbone, the training timesteps are the sampler's own 10-step grid rather than a prior, and an unmentioned anchor on the pretrained model rides along with the objective.

# scripts/train_nft_sd3.py. LoRA "default" is v_theta, LoRA "old" is v_old,
# adapters disabled is the frozen v_ref - one 2.5B backbone, three rank-32 LoRA states.
x0 = batch["latents_clean"]  # rollouts store the clean image only, never the trajectory
t = batch["timesteps"][:, j] / 1000  # t walks the sampler's own 10-step grid, not a prior
xt = (1 - t) * x0 + t * randn_like(x0)  # Alg. 1 line 9, rectified-flow schedule

# Optimality reward, Alg. 1 lines 4-5. Z_c is 5x the per-prompt std, and that std is
# accumulated over the whole epoch's rollouts rather than over the group of 24 alone.
r = clip((adv / 5) / 2 + 0.5, 0, 1)  # adv = clip((r_raw - mean_c) / (std_c + 1e-4), -5, 5)

v_pos = beta * v_theta + (1 - beta) * v_old  # Eq. 5; beta = 1.0, or 0.1 for OCR and multi-reward
v_neg = (1 + beta) * v_old - beta * v_theta

# Adaptive weighting is applied per branch, each normalized by its own error - the paper
# writes one w(t). The velocity is scored through its x0 predictor, never directly.
def nmse(v_hat, dims=(1, 2, 3)):  # per sample, scaled by that sample's own mean error
    e = xt - t * v_hat - x0
    return (e ** 2 / sg(e.abs().mean(dims, keepdim=True).clip(min=1e-5))).mean(dims)

loss = (r * nmse(v_pos) + (1 - r) * nmse(v_neg)).mean() * 5 / beta  # 5/beta pins the step size
loss += 1e-4 * ((v_theta - v_ref) ** 2).mean()  # never in the paper: an anchor on the pretrained model

# Soft update of the sampling LoRA, once per iteration; eta_i = min(0.001 i, 0.5),
# and for OCR min(0.0075 (i - 75), 0.999) - flat for 75 iterations before it ramps.
old.data = eta * old.data + (1 - eta) * new.data

---

One GRPO recipe stretched across diffusion and rectified flow and across image, video and image-to-video, with the fixes each scale demanded.

• Diffusion and rectified flow are one update (§2.1). Both read as ${\tilde{z}}_s={\tilde{z}}_t+(\text{network output})({\eta }_s-{\eta }_t)$; only the meaning of $\tilde{z}$ and $\eta$ changes — $z/\alpha$ and $\sigma /\alpha$ for $ϵ$-prediction, $z$ and $t$ for flow. One MDP then covers both, which is why a single recipe here can span all four base models.

• Two SDEs, one purpose (Eq. 7–8). Diffusion gets $d{z}_t=[f_t{z}_t-\frac{1+{\epsilon }_t^2}{2}{g}_t^2\nabla \log p_t]dt+{\epsilon }_t{g}_{t}dw$, rectified flow $d{z}_t=(u_t-\frac{1}{2}{\epsilon }_t^2\nabla \log p_t)dt+{\epsilon }_{t}dw$; under a Gaussian $p_t$ the score is closed-form, so both give a Gaussian policy whose per-step likelihood is as cheap to read as DDPO's.

• The objective, minus the KL (Eq. 9–10). $A_i=(r_i-\mathrm{m}\mathrm{e}\mathrm{a}\mathrm{n}\{r\})/\mathrm{s}\mathrm{t}\mathrm{d}\{r\}$ applied at every timestep because the reward exists only at the last one, inside a clipped ratio. The KL is dropped outright — the paper reports no measurable difference from keeping it, so nothing anchors the policy.

• One noise per group. All $G$ samples for a prompt start from the same $z_T$. Giving each its own noise, as DDPO does, reliably produces reward hacking and unstable training in video — the group comparison has to be about the policy, not about which seed happened to be lucky.

• Timestep dropout. 40% of the denoising steps are masked at random each iteration. The first 30% carry the most signal, yet training only on them is worse than the full sequence, so a uniform random subset is the compromise that keeps late-stage refinement in the picture.

• Several rewards, added as advantages. Reward models live on different scales, so it is the normalized advantages that are summed, $A_i={\sum }_k(r_i^k-{\mu }^k)/{\sigma }^k$, not the rewards. HPS alone drives images toward an oily look; adding CLIP holds them back toward realism.

• Best-of-N as a data filter (§3.2). Train on the best 8 and worst 8 of a pool of 16, 64 or 256 per prompt: the larger the pool the faster the reward climbs, paid for in sampling time — which is why the paper leaves it an optional extension rather than a part of the default recipe.

• Binary rewards train (§3.2). Thresholding HPS at 0.28 and CLIP at 0.39 into $\{0,1\}$ still moves the model, which is what makes rule-based rewards plausible for images at all, and echoes what verifiable rewards have already done for reasoning in the language models.

• Experiments (§3). SD1.4, FLUX, HunyuanVideo and SkyReels-I2V on HPS-v2.1, CLIP and VideoAlign, 10K+ prompts: SD1.4 goes 0.239 → 0.365 on HPS, past DDPO's 0.297 and ReFL's 0.357; HunyuanVideo gains 56% visual and 181% motion quality, SkyReels-I2V 91% motion.

• What it costs, said plainly (§3.2–3.3). Optimizing one reward moves the others: FLUX tuned on HPS alone falls 0.659 → 0.561 on GenEval, and HunyuanVideo's text alignment drops 1.75 → 1.59 while its quality rises — VideoAlign's alignment head was too unstable to train against.

The release fixes the two numbers the paper leaves as symbols, and the noise schedule is not the one its closest neighbour uses.

# fastvideo/train_grpo_flux.py, with scripts/finetune/finetune_flux_grpo_8gpus.sh.
# eta = 0.3 for FLUX, 0.25 for HunyuanVideo, and the step noise is flat in t -
# Flow-GRPO scales its own as a * sqrt(t / (1 - t)).
std_dev_t = eta * sqrt(delta_t)
prev_mean = prev_sample_mean_from(pred, z, sigmas, index, log_term=-0.5 * eta**2 * score)
log_prob = -((z_next - prev_mean)**2) / (2 * std_dev_t**2) - log(std_dev_t) - log(sqrt(2*pi))

# Group advantage, or a global one when --use_group is off.
adv[g] = (rewards[g] - rewards[g].mean()) / (rewards[g].std() + 1e-8)

# The 40% dropout of Section 3.6 is a shuffled index list, truncated.
perms = stack([randperm(len(timesteps)) for _ in range(batch)])
train_timesteps = int(len(timesteps) * 0.6)  # timestep_fraction, one value for every model

ratio = exp(new_log_probs - sample["log_probs"][:, k])
loss = maximum(-adv * ratio, -adv * clip(ratio, 1 - 1e-4, 1 + 1e-4)).mean()  # DDPO's clip range

---

A flow model has no randomness to explore with — convert its ODE into an equivalent SDE, and GRPO then carries over almost unchanged.

• The MDP is inherited (§3). $s_t=(c,t,x_t)$, $a_t=x_{t-1}$, $\pi (a_t|s_t)=p_{\theta }(x_{t-1}|x_t,c)$, reward only at $t=0$ — DDPO's mapping taken whole. What is new is making a flow model fit into it at all, since a flow has no per-step conditional distribution to write down at all, let alone to differentiate.

• A flow model has no policy (§4). Rectified-flow sampling is the ODE $d{x}_t=v_{t}dt$: the step is a bijection, so $p(x_{t-1}|x_t)$ would need a divergence estimate, and beyond the initial seed there is nothing to explore — given the seed, the trajectory is fixed, every rollout the same.

• ODE to SDE, the target (App. A, Eq. 10–14). The sampler is the ODE $d{x}_t=v_{t}dt$; wanted is an SDE $d{x}_t=fdt+{\sigma }_{t}dw$ with the same marginals $p_t$. Each moves its density by its own law, Fokker–Planck or the continuity equation, so equal marginals means equal right-hand sides:

$$-\nabla ·[f{p}_t]+\frac{1}{2}{\nabla }^2[{\sigma }_t^2{p}_t]={\partial }_t{p}_t=-\nabla ·[v_t{p}_t]$$

• The diffusion term is a drift in disguise (Eq. 15–17). Since $\nabla p=p\nabla \log p$, ${\nabla }^2[{\sigma }_t^2{p}_t]={\sigma }_t^2\nabla ·(p_t\nabla \log p_t)$: what noise does to $p_t$ is a flow of $p_t$ along its own score. Fold it into the drift and the equation is solved by $f=v_t+({\sigma }_t^2/2)\nabla \log p_t$, an SDE that spreads and re-gathers equally:

$$d{x}_t=[v_t(x_t)+\frac{{\sigma }_t^2}{2}\nabla \log p_t(x_t)]dt+{\sigma }_{t}dw\text{(forward SDE, Eq. 17)}$$

• Run it backwards (Eq. 18–21). Generation runs $t=1\to 0$, noise to data, so the reverse-time SDE is what is sampled: Anderson's rule turns $dx=fdt+gdw$ into $dx=[f-g^2\nabla \log p_t]dt+gd\overline{w}$, and with $g={\sigma }_t$ the score's coefficient flips, ${\sigma }_t^2/2-{\sigma }_t^2=-{\sigma }_t^2/2$ — the main text's Eq. 7:

$$d{x}_t=[v_t(x_t)-\frac{{\sigma }_t^2}{2}\nabla \log p_t(x_t)]dt+{\sigma }_{t}dw\text{(backward SDE, Eq. 21)}$$

• First the score (Eq. 22–24). No network outputs $\nabla \log p_t$, but for $x_t={\alpha }_t{x}_0+{\beta }_t{x}_1$ the conditional $p_t(x_t|x_0)=𝒩({\alpha }_t{x}_0,{\beta }_t^{2}I)$ has score $-(x_t-{\alpha }_t{x}_0)/{\beta }_t^2=-x_1/{\beta }_t$, and the marginal score is its posterior average — $\nabla p_t=\int p(x_0)\nabla p_t(x_t|x_0)d{x}_0$, divided through by $p_t$:

$$\nabla \log p_t(x_t)=𝔼[\nabla \log p_t(x_t|x_0)\left|x_t\right]=-\frac{𝔼[x_1|x_t]}{{\beta }_t}$$

• Then the velocity (Eq. 25–27). $v_t(x)=𝔼[{\dot{\alpha }}_t{x}_0+{\dot{\beta }}_t{x}_1|x_t=x]$; write $x_0=(x_t-{\beta }_t{x}_1)/{\alpha }_t$ and every term is $x$ or $𝔼[x_1|x_t]$, which Eq. 24 turns into the score. With ${\alpha }_t=1-t$, ${\beta }_t=t$ the bracket collapses to $t+t^2/(1-t)=t/(1-t)$, and solving for the score then gives Eq. 27:

$$v_t(x)=\frac{{\dot{\alpha }}_t}{{\alpha }_t}x-({\dot{\beta }}_t{\beta }_t-\frac{{\dot{\alpha }}_t{\beta }_t^2}{{\alpha }_t})\nabla \log p_t(x)=-\frac{x}{1-t}-\frac{t}{1-t}\nabla \log p_t(x)$$

$$\nabla \log p_t(x)=-\frac{x}{t}-\frac{1-t}{t}{v}_t(x)$$

• The sampler (Eq. 28–29 = Eq. 8–9). Eq. 27 into the reverse SDE removes the score, and Euler–Maruyama makes every transition a Gaussian: mean the drift step, variance ${\sigma }_t^2|\Delta t|I$. Time runs $1\to 0$, so $\Delta t<0$ and the noise scale is $\sqrt{|\Delta t|}$, the release's sqrt(-dt) term:

$$d{x}_t=[v_t(x_t)+\frac{{\sigma }_t^2}{2t}(x_t+(1-t)v_t(x_t)\left)\right]dt+{\sigma }_{t}dw$$

$$x_{t+\Delta t}=x_t+[v_{\theta }(x_t,t)+\frac{{\sigma }_t^2}{2t}(x_t+(1-t)v_{\theta }(x_t,t)\left)\right]\Delta t+{\sigma }_t\sqrt{\Delta t}ϵ,ϵ~𝒩(0,I)$$

• What the Gaussian buys. ${\pi }_{\theta }(x_{t-1}|x_t,c)=𝒩({\mu }_{\theta },{\sigma }_t^2|\Delta t|I)$, ${\mu }_{\theta }$ the bracket above: the ratio $r_t(\theta )$ is a quotient of two Gaussian densities and the KL a difference of means, both closed-form, so DDPO's machinery carries over — and the marginals stay the pretrained model's.

• A KL that costs nothing (§4). Both policies are Gaussians of one variance, so the KL is a squared difference of predicted means — and, as the means differ only through the velocities, of predicted velocities: no second sampling pass, no estimator, no extra gradient variance:

$$D_{\mathrm{K}\mathrm{L}}({\pi }_{\theta }\Vert {\pi }_{\mathrm{r}\mathrm{e}\mathrm{f}})=\frac{\Vert x_{t+\Delta t,\theta }-x_{t+\Delta t,\mathrm{r}\mathrm{e}\mathrm{f}}{\Vert }^2}{2{\sigma }_t^2\Delta t}=\frac{\Delta t}{2}(\frac{{\sigma }_t(1-t)}{2t}+\frac{1}{{\sigma }_t}{)}^2\Vert v_{\theta }(x_t,t)-v_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_t,t){\Vert }^2$$

• The noise level is the exploration knob (§5.3). ${\sigma }_t=a\sqrt{t/(1-t)}$, which keeps ${\sigma }_t^2/2t$ finite at $t\to 0$ and the noise nil at the data end (my reading): $a=0.1$ explores too thinly, 0.7 is best, past 1.0 nothing more, and too much noise destroys the samples and the reward with them.

• The objective (Eq. 4–5). the advantage $A^i=(R^i-\mathrm{m}\mathrm{e}\mathrm{a}\mathrm{n}R)/\mathrm{s}\mathrm{t}\mathrm{d}R$ over a group of $G$ images for one prompt, the same scalar at every denoising step, against a per-step clipped ratio and the KL above. No value network appears anywhere, which is the reason GRPO and not PPO.

• Denoising Reduction. Roll out at $T=10$ and still evaluate at $T=40$: over 4× faster on all three tasks with no loss in final reward, since low-quality trajectories carry the same ranking signal. Dropping to 5 is not reliably faster and sometimes slower, so ten is where it stops.

• Group size decides stability (§5.3). $G=24$ trains through; at 12 and 6 the advantage estimate is noisy enough that both runs climb, peak and then collapse outright — a small group makes the denominator of Eq. 4 unreliable long before the numerator itself ever becomes so.

• Experiments (§5). SD3.5-M on GenEval composition, OCR text rendering and PickScore, DrawBench quality watching for hacking: GenEval 0.63 → 0.95, past GPT-4o's 0.84, OCR 0.59 → 0.92, PickScore 21.72 → 23.31, ahead of SFT, Flow-DPO and their online variants.

• KL is not early stopping (§5.2–5.3). Without it GenEval still reaches 0.95, but DrawBench aesthetics fall 5.39 → 4.93 and DeQA 4.07 → 2.77; on PickScore, quality holds and diversity collapses instead, different seeds landing on one style — which no early stop would have caught.

• What a template teaches (§5.3, Tab. 3–4). Trained on 60 GenEval classes and counts of 2–4, it scores 0.90 on 20 unseen classes from 0.64, counts 5–6 objects at 0.48 (0.13), 12 at 0.12 (0.02), and lifts T2I-CompBench 2D-spatial 0.2850 → 0.5447 — relations, not prompts.

The release differs from the paper in the parts a reader would reuse: a second SDE discretization that is not Euler–Maruyama, a filter that throws away prompt groups whose rewards all agree, and an EMA of the trained weights that no section of the paper ever thinks to mention at all.

# scripts/train_sd3.py + flow_grpo/diffusers_patch/sd3_sde_with_logprob.py.
std_dev_t = sqrt(sigma / (1 - sigma)) * 0.7  # the paper's a, pinned to 0.7 as the default
prev_mean = x * (1 + std_dev_t**2 / (2*sigma) * dt) \
          + v_theta * (1 + std_dev_t**2 * (1 - sigma) / (2*sigma)) * dt  # Eq. 9
log_prob = -((x_next - prev_mean)**2) / (2 * (std_dev_t * sqrt(-dt))**2) - log(std_dev_t * sqrt(-dt))

# A second branch nobody writes about: a DDIM-style step whose noise is
# std_dev_t = sigma_prev * sin(a * pi / 2), not the paper's Euler-Maruyama.

ratio = exp(log_prob - sample["log_probs"][:, j])
adv = clip(sample["advantages"][:, j], -5, 5)  # per-prompt std, or a global one by config
policy_loss = maximum(-adv * ratio, -adv * clip(ratio, 1 - 1e-3, 1 + 1e-3)).mean()
kl = ((prev_mean - prev_mean_ref)**2).mean(dims) / (2 * std_dev_t**2)  # the closed form above
loss = policy_loss + 0.04 * kl  # beta = 0.04 for GenEval and OCR, 0.004 in the base config

# Groups where every image scored alike give a zero advantage and are dropped
# before the gradient step; the fraction is logged as zero_std_ratio.
keep = samples["advantages"].abs().sum(dim=1) != 0

---

DPO for diffusion — a pairwise preference loss asking only that the model improve on the winner more than it does on the loser beside it.

(see notes in jupyter)

• Why not a reward model (§1). RL routes so far tune on small prompt sets and, run against a learned scorer, drift into mode collapse and reward hacking; DPO drops both reward network and rollouts, learning from pairs $(c,x_0^w\succ x_0^l)$ a reference model drew and people ranked.

• Bradley–Terry (Eq. 3–4). Preference is a sigmoid of a reward gap, $p_{\mathrm{B}\mathrm{T}}(x_0^w\succ x_0^l\mid c)=\sigma (r(c,x_0^w)-r(c,x_0^l))$, and a reward net $r_{\varphi }$ is fit by $L_{\mathrm{B}\mathrm{T}}(\varphi )=-𝔼[\log \sigma (r_{\varphi }(c,x_0^w)-r_{\varphi }(c,x_0^l))]$ — plain binary classification on a static, human-labeled dataset, which is all a reward model ever learns.

• RLHF, and its closed form (Eq. 5–7). Maximize ${𝔼}_{p_{\theta }}[r(c,x_0)]-\beta D_{\mathrm{K}\mathrm{L}}[p_{\theta }(x_0|c)\Vert p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_0|c)]$; the unique optimum is $p_{\theta }^{*}=p_{\mathrm{r}\mathrm{e}\mathrm{f}}\exp (r/\beta )/Z(c)$ with $Z(c)={\sum }_{x_0}{p}_{\mathrm{r}\mathrm{e}\mathrm{f}}\exp (r/\beta )$, so $r(c,x_0)=\beta \log (p_{\theta }^{*}(x_0|c)/p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_0|c))+\beta \log Z(c)$ — the reward is a function of the policy and nothing else.

• DPO (Eq. 8). Put that reward into Eq. 4: $Z(c)$ cancels across the pair, no reward model is ever fit and no sample is ever drawn, and the policy is trained directly by $L_{\mathrm{D}\mathrm{P}\mathrm{O}}=-𝔼[\log \sigma (\beta \log (p_{\theta }(x_0^w|c)/p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_0^w|c))-\beta \log (p_{\theta }(x_0^l|c)/p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_0^l|c)))]$ — the policy is its own reward model.

• What diffusion breaks (§4, Eq. 9–10). Eq. 8 needs $p_{\theta }(x_0|c)$, which marginalizes each path $x_{1:T}$ ending at $x_0$ — the one thing a diffusion model never gives. So define a chain reward $R$, $r(c,x_0)={𝔼}_{p_{\theta }(x_{1:T}|x_0,c)}[R(c,x_{0:T})]$, and bound the KL by the joint $D_{\mathrm{K}\mathrm{L}}[p_{\theta }(x_{0:T}|c)\Vert p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{0:T}|c)]$.

• The lift, line by line (S2, Eq. 17). Eq. 5 over $\beta$, sign flipped; each line below is one move — the ELBO bound, the reward swap of Eq. 9, completing the KL with $Z(c)$ — and the last minimizer is read off: $p_{\theta }^{*}=p_{\mathrm{r}\mathrm{e}\mathrm{f}}\exp (R/\beta )/Z(c)$, so $R=\beta \log (p_{\theta }^{*}/p_{\mathrm{r}\mathrm{e}\mathrm{f}})+\beta \log Z(c)$, Eq. 7 on paths.

$${min}_{p_{\theta }}{𝔼}_{p_{\theta }(x_0|c)}[-r(c,x_0)/\beta ]+D_{\mathrm{K}\mathrm{L}}(p_{\theta }(x_0|c)\Vert p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_0|c))$$

$$\le {min}_{p_{\theta }}{𝔼}_{p_{\theta }(x_0|c)}[-r(c,x_0)/\beta ]+D_{\mathrm{K}\mathrm{L}}(p_{\theta }(x_{0:T}|c)\Vert p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{0:T}|c))$$

$$={min}_{p_{\theta }}{𝔼}_{p_{\theta }(x_{0:T}|c)}[-R(c,x_{0:T})/\beta ]+D_{\mathrm{K}\mathrm{L}}(p_{\theta }(x_{0:T}|c)\Vert p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{0:T}|c))$$

$$={min}_{p_{\theta }}{𝔼}_{p_{\theta }(x_{0:T}|c)}[\log \frac{p_{\theta }(x_{0:T}|c)}{p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{0:T}|c)\exp (R(c,x_{0:T})/\beta )/Z(c)}-\log Z(c)]$$

$$={min}_{p_{\theta }}{D}_{\mathrm{K}\mathrm{L}}(p_{\theta }(x_{0:T}|c)\Vert p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{0:T}|c)\exp (R(c,x_{0:T})/\beta )/Z(c)).$$

• Path DPO (Eq. 11). $R$ into Eq. 4, and $Z(c)$ cancels again: the first line below is Eq. 8 with path ratios inside an expectation over the reverse process $p_{\theta }(x_{1:T}|x_0)$, which is both intractable to sample and $T=1000$ steps long. The chain below ($c$ dropped) removes both in turn.

$$L(\theta )=-\log \sigma \left(\beta {𝔼}_{x_{1:T}^w~p_{\theta }(x_{1:T}|x_0^w),x_{1:T}^l~p_{\theta }(x_{1:T}|x_0^l)}\right[\log \frac{p_{\theta }(x_{0:T}^w)}{p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{0:T}^w)}-\log \frac{p_{\theta }(x_{0:T}^l)}{p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{0:T}^l)}\left]\right)$$

$$L_1(\theta )=-\log \sigma \left(\beta {𝔼}_{x_{1:T}^w~q(x_{1:T}|x_0^w),x_{1:T}^l~q(x_{1:T}|x_0^l)}\right[{\sum }_{t=1}^T\log \frac{p_{\theta }(x_{t-1}^w|x_t^w)}{p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{t-1}^w|x_t^w)}-\log \frac{p_{\theta }(x_{t-1}^l|x_t^l)}{p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{t-1}^l|x_t^l)}\left]\right)$$

$$=-\log \sigma \left(\beta T{𝔼}_{t,x_{t-1,t}^w~q(·|x_0^w),x_{t-1,t}^l~q(·|x_0^l)}\right[\log \frac{p_{\theta }(x_{t-1}^w|x_t^w)}{p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{t-1}^w|x_t^w)}-\log \frac{p_{\theta }(x_{t-1}^l|x_t^l)}{p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{t-1}^l|x_t^l)}\left]\right)$$

$$\le -{𝔼}_{t,x_t^w~q(x_t|x_0^w),x_t^l~q(x_t|x_0^l)}\log \sigma \left(\beta T{𝔼}_{x_{t-1}^w~q(·|x_t^w,x_0^w),x_{t-1}^l~q(·|x_t^l,x_0^l)}\right[\log \frac{p_{\theta }(x_{t-1}^w|x_t^w)}{p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{t-1}^w|x_t^w)}-\log \frac{p_{\theta }(x_{t-1}^l|x_t^l)}{p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{t-1}^l|x_t^l)}\left]\right)$$

$$=-{𝔼}_{t,x_t^w,x_t^l}\log \sigma (-\beta T[D_{\mathrm{K}\mathrm{L}}(q(x_{t-1}^w|x_{0,t}^w)\Vert p_{\theta }(x_{t-1}^w|x_t^w))-D_{\mathrm{K}\mathrm{L}}(q(x_{t-1}^w|x_{0,t}^w)\Vert p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{t-1}^w|x_t^w))$$

$$-D_{\mathrm{K}\mathrm{L}}(q(x_{t-1}^l|x_{0,t}^l)\Vert p_{\theta }(x_{t-1}^l|x_t^l))+D_{\mathrm{K}\mathrm{L}}(q(x_{t-1}^l|x_{0,t}^l)\Vert p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{t-1}^l|x_t^l))\left]\right)$$

• Off-policy by substitution (S2, Eq. 18). Line two swaps the unsampleable reverse process for the forward $q(x_{1:T}|x_0)$ — what lets training run off a fixed dataset — and splits the path ratio into per-step ratios; line three trades the sum over $T$ steps for $T$ times one random step $t$.

• Jensen, then KLs (Eq. 12–13). $-\log \sigma$ is convex, so the expectation over $x_{t-1}$ moves outside it (line four, an upper bound); against the posterior $q(x_{t-1}|x_0,x_t)$, each ratio's expectation is a difference of two KLs (line five), and every one of them is a Gaussian, by Eq. 1's form.

• The objective (Eq. 14). Gaussian KLs are squared errors: $L(\theta )=-𝔼\log \sigma (-\beta T\omega ({\lambda }_t)[(\Vert {ϵ}^w-{ϵ}_{\theta }{\Vert }^2-\Vert {ϵ}^w-{ϵ}_{\mathrm{r}\mathrm{e}\mathrm{f}}{\Vert }^2)-(\Vert {ϵ}^l-{ϵ}_{\theta }{\Vert }^2-\Vert {ϵ}^l-{ϵ}_{\mathrm{r}\mathrm{e}\mathrm{f}}{\Vert }^2)])$ — four denoising losses at one $t$: denoise the winner better than the reference, the loser worse; $T$ folds into $\beta$, hence a $\beta$ in the thousands.

• The other approximation (S2, $L_2$). Per step, $L_1$ stands $q(x_{t-1,t}|x_0)$ in for $p_{\theta }(x_{t-1,t}|x_0)$; using $q(x_t|x_0)p_{\theta }(x_{t-1}|x_t)$ is closer, since $D_{\mathrm{K}\mathrm{L}}(q(x_t|x_0)p_{\theta }(x_{t-1}|x_t)\Vert p_{\theta }(x_{t-1,t}|x_0))=D_{\mathrm{K}\mathrm{L}}(q(x_t|x_0)\Vert p_{\theta }(x_t|x_0))<D_{\mathrm{K}\mathrm{L}}(q(x_{t-1,t}|x_0)\Vert p_{\theta }(x_{t-1,t}|x_0))$ — one factor fewer to guess.

$$L_2(\theta )=-\log \sigma \left(\beta T{𝔼}_{t,x_t^w~q(x_t|x_0^w),x_t^l~q(x_t|x_0^l)}{𝔼}_{x_{t-1}^w~p_{\theta }(·|x_t^w),x_{t-1}^l~p_{\theta }(·|x_t^l)}\right[\log \frac{p_{\theta }(x_{t-1}^w|x_t^w)}{p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{t-1}^w|x_t^w)}-\log \frac{p_{\theta }(x_{t-1}^l|x_t^l)}{p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{t-1}^l|x_t^l)}\left]\right)$$

$$\le -{𝔼}_{t,x_t^w,x_t^l}\log \sigma \left(\beta T\right[D_{\mathrm{K}\mathrm{L}}(p_{\theta }(x_{t-1}^w|x_t^w)\Vert p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{t-1}^w|x_t^w))-D_{\mathrm{K}\mathrm{L}}(p_{\theta }(x_{t-1}^l|x_t^l)\Vert p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{t-1}^l|x_t^l))\left]\right)$$

$$=-{𝔼}_{t,{ϵ}^w,{ϵ}^l}\log \sigma (-\beta T\omega ({\lambda }_t)[\Vert {ϵ}_{\theta }(x_t^w,t)-{ϵ}_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_t^w,t){\Vert }^2-\Vert {ϵ}_{\theta }(x_t^l,t)-{ϵ}_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_t^l,t){\Vert }^2\left]\right).$$

• Why $L_2$ is not the one used. It asks the model to stay near the reference on winners and leave it on losers — but $\theta$ starts at ${\theta }_{\mathrm{r}\mathrm{e}\mathrm{f}}$, so both squared differences are zero at initialization and so is the gradient, whereas $L_1$ has the data's $ϵ$ to pull against from the very first step.

• A noisy classifier gives it too (S4, Eq. 44–46). Condition the policy on $(t,x_t)$; with the DDIM likelihood $p(x_0|x_t)=𝒩(x_0^{\mathrm{p}\mathrm{r}\mathrm{e}\mathrm{d}},{\sigma }_t^{2}I)$, each log-ratio in Eq. 8 is a gap in squared noise error times $b_t^2/(a_t^2{\sigma }_t^2)\approx 1$, so $-\log \sigma (-(\beta /2)[({\mathrm{S}\mathrm{E}}_{\theta }^w-{\mathrm{S}\mathrm{E}}_{\mathrm{r}\mathrm{e}\mathrm{f}}^w)-({\mathrm{S}\mathrm{E}}_{\theta }^l-{\mathrm{S}\mathrm{E}}_{\mathrm{r}\mathrm{e}\mathrm{f}}^l)])$ — Eq. 14, and the code's $-\beta /2$.

• The same MDP as DDPO, solved off-policy (S3, Eq. 19–23). State $(c,x_t,t)$, action $x_{t-1}$, reward $r(c,x_0)$ at $t=0$; soft optimal control gives $p^{*}(x_{t-1}|x_t,c)=p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{t-1}|x_t,c)e^{(Q^{*}-V^{*})/\beta }$, exact since the right side integrates to one, so $\log (p^{*}/p_{\mathrm{r}\mathrm{e}\mathrm{f}})=(Q^{*}-V^{*})/\beta$ exactly, Eq. 25.

• Telescoping to Eq. 11 (Eq. 24–27). The inverse Bellman equation $r(c,x_t)=V^{*}(x_{t-1},c)-Q^{*}$ (Eq. 24), summed down the chain, leaves $r(c,x_0)={\sum }_t\log (p^{*}(x_{t-1}|x_t,c)/p_{\mathrm{r}\mathrm{e}\mathrm{f}}(x_{t-1}|x_t,c))-V^{*}(x_T,c)$; $V^{*}(x_T)$ cancels across a pair — Eq. 11, reached with no rollout at all.

• Data and the scaling rule (§5.1). Pick-a-Pic v2: 851,293 pairs over 58,960 prompts once the ~12% ties are dropped; batches of 2048 pairs; $\beta =2000$ for SD1.5 and $5000$ for SDXL, learning rate $(2000/\beta )·2.048·{10}^{-8}$, 25% warmup — $1/\beta$ scaling, since the gradient's norm grows with $\beta$.

• Experiments (§5). SD1.5 and SDXL-base tuned on Pick-a-Pic v2, five labelers on PartiPrompts and HPSv2: DPO-SDXL beats SDXL-base 70.0% / 64.7%, appeal and alignment moving in step, and the 6.6B base+refiner 69% / 64% from 3.5B; and wins TEdBench edits 65% to 24%.

• A reward model falls out for free (§5.5, Tab. 2). The inner term of Eq. 14, averaged over 10 random $t$, estimates a reward gap, so the tuned model is also a preference classifier: DPO-SDXL scores 72.0% on Pick-a-Pic v2 validation, above PickScore's 64.2, HPS's 59.3 and CLIP's 57.1.

• Pseudo-labels beat labels; SFT hurts SDXL (§5.4–5.5). Reranking pairs by PickScore, not people, lifts the PartiPrompts win rate 59.8% → 63.3%; fine-tuning on winners alone helps SD1.5 (55.5%) yet hurts SDXL at any rate: SDXL-1.0 already outdraws the models that made the data.

The released code tracks the paper closely, but three of its own lines carry decisions that the derivation in Section 4 never mentions at all.

# train.py from the official Salesforce release. The pair arrives concatenated
# along the channel axis and is split apart; the winner is the first half.
timesteps = timesteps.chunk(2)[0].repeat(2)  # one t for the pair, not two - never stated in Eq. 14
noise = noise.chunk(2)[0].repeat(2, 1, 1, 1)  # and one epsilon, which halves the gradient noise

model_losses = (model_pred - noise).pow(2).mean(dim=[1, 2, 3])
model_diff = model_losses_w - model_losses_l  # the model's margin
ref_diff = ref_losses_w - ref_losses_l  # the frozen reference's margin, no grad

# omega(lambda_t) is dropped outright; the -beta/2 is Eq. 46's coefficient, b_t^2/(a_t^2 sigma_t^2) ~ 1.
inside_term = -0.5 * beta_dpo * (model_diff - ref_diff)  # beta_dpo = 5000 by default
loss = -F.logsigmoid(inside_term).mean()
implicit_acc = (inside_term > 0).float().mean()  # the implicit reward model, scored live

# AI feedback (5.4) is this one line: a scorer reorders the pair before training.
if choice_model_says_flip(batch): feed_pixel_values = feed_pixel_values.flip(0)

---

The same multi-step MDP as DDPO, arrived at independently, and a KL to the pretrained model carried the whole way as an implicit reward.

• The MDP, indexed backwards (Eq. 4). $s_t=(z,x_{T-t})$, $a_t=x_{T-t-1}$, ${\pi }_{\theta }(a_t|s_t)=p_{\theta }(x_{T-t-1}|x_{T-t},z)$, transitions $({\delta }_z,{\delta }_{a_t})$, $P_0=(p(z),𝒩(0,I))$, reward only at $t=T-1$ — the same mapping DDPO found, written so that MDP time runs forward while the denoising runs backward.

• The objective and its gradient (Eq. 5–6). Minimize ${𝔼}_{p(z)}{𝔼}_{p_{\theta }(x_0|z)}[-r(x_0,z)]$; Lemma 4.1 turns that into $𝔼[-r(x_0,z){\sum }_t{\nabla }_{\theta }\log p_{\theta }(x_{t-1}|x_t,z)]$, REINFORCE over the $T$-step chain. The DDPM covariance stays frozen at $\Sigma$ and only ${\mu }_{\theta }$ is trained, which is what leaves the policy stochastic.

• A KL you can actually compute (Lemma 4.2, Eq. 7). The divergence that matters is on the image, and that marginal is intractable — so bound it over all $T$ steps by ${\sum }_t𝔼\mathrm{K}\mathrm{L}(p_{\theta }(x_{t-1}|x_t,z)\Vert p_{\mathrm{p}\mathrm{r}\mathrm{e}}(x_{t-1}|x_t,z))$, which makes every term a Gaussian KL with a closed form of its own.

• The regularized objective (Eq. 8). ${𝔼}_{p(z)}[\alpha 𝔼[-r(x_0,z)]+\beta {\sum }_t𝔼\mathrm{K}\mathrm{L}(p_{\theta }\Vert p_{\mathrm{p}\mathrm{r}\mathrm{e}})]$ — $\alpha$ and $\beta$ trade reward against drift, and the KL arrives as a second reward evaluated online, per update, on fresh samples, which is precisely the thing a supervised fine-tuning pass cannot do.

• The gradient they actually use (Eq. 9) is not that objective's. $𝔼[-\alpha r{\sum }_t{\nabla }_{\theta }\log p_{\theta }+\beta {\sum }_t{\nabla }_{\theta }\mathrm{K}\mathrm{L}(p_{\theta }\Vert p_{\mathrm{p}\mathrm{r}\mathrm{e}})]$ differentiates the KL's value but not the state distribution that the KL is measured under — a gap the main text flags plainly in the line right after the equation, not in the code.

• The term they drop (App. A.3, Eq. 16). The KL sits under $p_{\theta }(x_t)$, the policy's own state distribution, so the product rule adds a second term, ${\sum }_t𝔼[\mathrm{K}\mathrm{L}(p_{\theta }\Vert p_{\mathrm{p}\mathrm{r}\mathrm{e}}){\sum }_{s>t}{\nabla }_{\theta }\log p_{\theta }(x_{s-1}|x_s,z)]$: each step then carries the KL of every step still ahead of it, dropped purely for its cost.

• Reusing a round of samples (App. A.6, Eq. 25–26). Writing $p_{\mathrm{o}\mathrm{l}\mathrm{d}}$ for the sampling policy, a ratio against it lets five gradient steps run off one round, and a PPO clip bounds the drift: $𝔼[-{\sum }_t(r-V){\nabla }_{\theta }\mathrm{c}\mathrm{l}\mathrm{i}\mathrm{p}(p_{\theta }/p_{\mathrm{o}\mathrm{l}\mathrm{d}},1-ϵ,1+ϵ)]$, with $V(x_t,z)$ a learned value baseline that cuts variance.

• KL for supervised fine-tuning too (Eq. 10–13). The baseline is $𝔼[-r\log p_{\theta }(x_0|z)]$; Lemma 4.3 bounds its KL-regularized form two ways — KL-D shifts the weight to $r+\gamma$, smoothing it toward uniform, and KL-O adds $\gamma \Vert f_{\mathrm{p}\mathrm{r}\mathrm{e}}-f_{\theta }{\Vert }^2$, pinning the denoising direction to the pretrained one.

• Three ways online differs (§4.3). RL may search outside the pretrained support while SFT only imitates inside it; RL's KL is measured on new samples each update, SFT's on a frozen set; and RL queries the reward model off-distribution, so it inherits whatever that model generalizes.

• Experiments (§5). SD v1.5, LoRA, ImageReward reward, one prompt per color, count, composition and location, judged by ImageReward, aesthetics, raters: RL beats SFT on all four, aesthetics intact; on 104 COCO / 183 DrawBench prompts it goes 0.22 → 0.55 / 0.13 → 0.58.

• Whiskey to flowers (§5.4). "Four roses" draws the whiskey from the base model; RL against ImageReward turns it into the flower, the score rising −0.52 → 1.12 — human feedback undoing a bias the web corpus put in, the clearest case for a preference model over a caption.

The release fixes what the paper leaves as symbols, and one of those numbers openly disagrees with the value its own appendix reports.

# dpok/train_online_pg.py from google-research. alpha and beta of Eq. 8 are
# flags: reward_weight defaults to 100, though App. B reports alpha = 10.
log_prob, kl_regularizer = pipe.forward_calculate_logprob(..., unet_copy=unet_copy)

# v_flag off by default: the "advantage" is the raw reward, uncentred - no
# per-prompt baseline of the kind DDPO normalizes with.
adv = reward - value_function(x_t, txt_emb, t) if v_flag else reward

# Not PPO. The ratio is clamped outright, so a sample outside the trust region
# contributes a constant instead of being discarded by a max() over two branches.
ratio = clamp(exp(log_prob - old_log_prob), 1 - 1e-4, 1 + 1e-4)
loss = (-100.0 * adv.detach() * ratio).mean()
if step > kl_warmup: loss += 0.01 * kl_regularizer.mean()

# log_prob is meaned over the latent axes, not summed, so its scale depends on
# resolution - part of why the reward weight has to be so large.