Diffusion Post Training Tutorial En

§0 TL;DR

1. Why it's hard: diffusion generation runs $T$ steps (typically 20–50), with reward given only once at the terminal state $x_0$ — so sparse terminal reward + long denoising trajectory + credit assignment combine, giving one extra "trajectory integration" dimension beyond LLM RLHF.
2. Three main lines: (i) RL on denoising MDP (DDPO / DPOK — treat $T$-step denoising as an MDP); (ii) Direct reward backprop (DRaFT / AlignProp / ReFL — treat reward as a differentiable loss and backprop through the $T$-step sampler); (iii) Preference optimization (Diffusion-DPO / D3PO / SPO / Diffusion-KTO / MaPO — port the LLM DPO family to diffusion).
3. DDPO (Black et al. 2024 ICLR, arXiv 2305.13301): view denoising as a $T$-step MDP, state $= (x_t, t, c)$, action $= x_{t-1}$, per-trajectory reward $R(x_0, c)$; use REINFORCE or PPO-clip to update $\log p_\theta(x_{t-1} \mid x_t, c)$.
4. AlignProp (Prabhudesai et al. 2024 ICLR, arXiv 2310.03739) & DRaFT (Clark et al. 2024 ICLR, arXiv 2309.17400): when reward $R$ is differentiable in $x_0$, directly backprop $R(x_0)$ through the $T$-step sampler into $\theta$. The key engineering problem: memory is $\mathcal{O}(T)$; DRaFT-K / AlignProp only backprop the final $K$ steps (typically $K \in \{1, 5\}$), and combined with gradient checkpointing the memory drops to $\mathcal{O}(K)$.
5. Diffusion-DPO (Wallace et al. 2024 CVPR, arXiv 2311.12908): replace the LLM DPO's $\log\pi/\pi_\text{ref}$ with diffusion's per-step ELBO surrogate — specifically, use $-\|\epsilon - \epsilon_\theta(x_t, t)\|^2$ as one lower-bound term of $\log p_\theta(x_0)$, and assemble it into a DPO contrastive loss over the pair $(y_w, y_l)$.
6. D3PO (Yang et al. 2024 CVPR, arXiv 2311.13231): completely RM-free — directly plug human thumbs up/down feedback on generated images into the KL-regularized optimal solution's implicit reward; the derivation parallels DPO but operates on the diffusion per-step Markov chain.
7. SPO (Liang et al. 2024, arXiv 2406.04314): observes that different denoising steps have different preferences (high-noise steps learn composition, low-noise steps learn detail), and extends DPO to be step-aware — for each $t$, independently sample an in-step pair $(x_{t-1}^w, x_{t-1}^l)$, and weight the loss along the step dimension.
8. Flow-GRPO (Liu et al. 2025, arXiv 2505.05470): the first work bringing GRPO to Flow Matching. Two key tricks: an ODE→SDE equivalent conversion that turns the deterministic flow into an explorable stochastic process; and denoising reduction — fewer steps for training, full steps for inference. RL-tuned SD3.5-M raised GenEval from 63% to 95%.
9. Reward hacking is the real boss: over-saturated colors, monotonous composition, style convergence, high PickScore but ugly to humans — mitigations include reward ensembles (HPSv2 + PickScore + ImageReward + CLIP-Score), KL anchor (Diffusion-DPO's $\beta$), and early stopping on reward plateau. SD3 / FLUX barely disclose post-training details, but the community consensus is that SD3.5 Turbo and FLUX.1 dev use a DPO + distillation hybrid.

§1 Intuition: why diffusion post-training is hard

1.1 Single-step vs multi-step generation: the essential difference

LLM rewards are generally also sequence-level, but tokens are discrete, trajectory length $L \sim 10^3$, vocabulary moderate. The diffusion "trajectory" is $T$ steps of denoising, each operating on a continuous high-dimensional tensor $x_t \in \mathbb{R}^{C \times H \times W}$ (SDXL latent is $4\times128\times128 = 65536$ dimensions), with $T$ typically 20–50.

1.2 The three main lines

• Line A (RL on denoising MDP): treat diffusion as an RL environment, not requiring reward to be differentiable. Representatives: DDPO, DPOK, Flow-GRPO.
• Line B (Direct reward backprop): when reward is differentiable in $x_0$, directly gradient-descend, akin to treating the reward as a new loss. Representatives: DRaFT, AlignProp, ReFL.
• Line C (Preference optimization, DPO-style): port the LLM DPO family to diffusion, no longer requiring on-policy sampling. Representatives: Diffusion-DPO, D3PO, SPO, Diffusion-KTO, MaPO.

                              preference / reward signal
                                    │
            ┌───────────────────────┼───────────────────────┐
            │                       │                       │
       reward non-diff           reward diff             preference pairs (offline)
            │                       │                       │
        Line A: RL              Line B: backprop         Line C: DPO family
        DDPO, DPOK,             DRaFT, AlignProp,        Diffusion-DPO,
        Flow-GRPO               ReFL                     D3PO, SPO, KTO, MaPO
            │                       │                       │
        most general            memory-friendly (K-step)   off-policy fast
        but sampling expensive  but requires diff reward   needs preference data

1.3 One-line intuitions

• Line A treats $T$ denoising steps as an RL trajectory: each step is a stochastic policy output.
• Line B treats reward as a differentiable loss: backprop through $T$ steps, but $O(T)$ memory crushes you.
• Line C bypasses reward entirely: uses the implicit reward of preference pairs $= \beta \log(p_\theta/p_\text{ref})$.

1.4 Convention (used throughout)

§2 RL for Diffusion: DDPO and DPOK

2.1 Treating denoising as an MDP (DDPO viewpoint)

Black et al. 2024 ICLR Training Diffusion Models with Reinforcement Learning (arXiv 2305.13301) key observation: the DDPM/DDIM reverse process is itself a finite-horizon MDP.

Definitions:

• State: $s_t = (x_t, t, c)$, with time reversed $t = T, T-1, \dots, 1$
• Action: $a_t = x_{t-1}$ (sampled from $p_\theta(\cdot \mid x_t, c)$)
• Transition: deterministic — $s_{t-1} = (x_{t-1}, t-1, c)$
• Reward: $r_t = 0$ for $t > 1$, $r_1 = R(x_0, c)$ (terminal-only)
• Policy: $\pi_\theta(a_t \mid s_t) = p_\theta(x_{t-1} \mid x_t, c)$

Policy gradient theorem:

$$\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim p_\theta}\!\left[\sum_{t=1}^{T} \nabla_\theta \log p_\theta(x_{t-1} \mid x_t, c)\, R(x_0, c)\right]$$

Key: $\log p_\theta(x_{t-1} \mid x_t, c)$ is Gaussian in DDPM, so the log-prob is analytical and the gradient can be computed directly.

2.2 DDPO-SF (Score Function) algorithm

The simplest variant (DDPO-SF, score function estimator):

1. Sampling phase — starting from prompt $c$, run $T$ DDPM reverse steps to get a trajectory $\tau = (x_T, x_{T-1}, \dots, x_0)$; compute $R(x_0, c)$.
2. Update phase — REINFORCE-style gradient estimator:

$$\hat{g} = \frac{1}{N}\sum_{n=1}^{N} \sum_{t=1}^{T} \nabla_\theta \log p_\theta(x_{t-1}^{(n)} \mid x_t^{(n)}, c)\, (R^{(n)} - b)$$

where $b$ is a baseline (typically the batch-mean reward).

2.3 DDPO-IS (Importance Sampling, PPO-style)

The practical recommended DDPO variant uses PPO-clip with per-step importance ratios:

$$\rho_t = \frac{p_\theta(x_{t-1} \mid x_t, c)}{p_{\theta_\text{old}}(x_{t-1} \mid x_t, c)}, \qquad L^\text{CLIP}_t = \min\!\big(\rho_t R, \text{clip}(\rho_t, 1-\epsilon, 1+\epsilon) R\big)$$

2.4 DDPO's two reward experiments

Black et al. run DDPO + SD-1.5 on four rewards:

Across all four rewards, DDPO yields large gains over reward-weighted regression (RWR baseline), and can shift styles (e.g. emoji → oil painting) — proof that RL is genuinely exploring, not just mode-seeking.

2.5 DPOK: KL-regularized RL for diffusion

Fan et al. 2023 NeurIPS DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models (arXiv 2305.16381) is contemporaneous with DDPO; the difference is an explicit KL anchor:

$$\boxed{\;\max_\theta\; \mathbb{E}_{\tau \sim p_\theta}\!\left[R(x_0, c)\right] - \beta\, \mathbb{E}_c\!\left[\text{KL}\!\big(p_\theta(\cdot \mid c) \,\big\Vert\, p_\text{ref}(\cdot \mid c)\big)\right]\;}$$

The KL term expanded to per-step:

$$\text{KL}(p_\theta \Vert p_\text{ref}) = \sum_{t=1}^{T} \mathbb{E}\!\left[\text{KL}\!\big(p_\theta(x_{t-1} \mid x_t, c) \,\Vert\, p_\text{ref}(x_{t-1} \mid x_t, c)\big)\right]$$

Since the DDPM $p_\theta(x_{t-1} \mid x_t, c)$ is Gaussian, the KL between two Gaussians has a closed form, so per-step KL is directly computable. DPOK uses policy gradient + this KL penalty, the diffusion analog of RLHF's "$\beta \log(\pi/\pi_\text{ref})$ added to reward".

• DDPO: pure RL (REINFORCE or PPO-clip), KL is implicit (via ratio clipping).
• DPOK: explicit KL term, matching LLM RLHF's "KL added to reward".
• Empirically, DPOK is more stable on reward-prompt alignment (ImageReward); DDPO is more aggressive on rule-based rewards like compressibility.

2.6 DDPO failure modes and mitigations

§3 Direct Reward Fine-Tuning: DRaFT / AlignProp / ReFL

3.1 Core idea: treat reward as a differentiable loss

If the reward $R(x_0, c)$ is differentiable in $x_0$ (true for typical CNN/ViT RMs) and the diffusion sampler is differentiable, then we can directly backprop to $\theta$:

$$\theta \leftarrow \theta + \eta\, \nabla_\theta R\!\big(x_0(\theta), c\big), \quad x_0(\theta) = \text{Sample}_\theta^T(c)$$

where $\text{Sample}_\theta^T(c)$ denotes running $T$ reverse steps from $x_T \sim \mathcal{N}(0,I)$ to obtain $x_0$. This treats the entire reverse trajectory as one giant differentiable computation graph and optimizes reward end-to-end.

3.2 DRaFT (Clark et al. 2024 ICLR, arXiv 2309.17400)

Directly Fine-Tuning Diffusion Models on Differentiable Rewards. Two core tricks:

Trick 1: DRaFT-K, only backprop the last $K$ steps.

The full chain rule (denoise step $\epsilon_\theta(x_t,t)$ both influences the next step $x_{t-1}$ and directly depends on $\theta$):

$$\nabla_\theta R(x_0) = \frac{\partial R}{\partial x_0} \cdot \sum_{t=1}^{K}\left(\prod_{s=1}^{t-1} \frac{\partial x_{s-1}}{\partial x_s}\right) \cdot \frac{\partial x_{t-1}}{\partial \theta}\bigg|_{\text{direct}}$$

where $\partial x_{t-1}/\partial\theta|_\text{direct}$ is the partial derivative of step-$t$ on $\theta$ via $\epsilon_\theta(x_t,t)$ directly (not through the indirect $x_t \to x_t$ path), and $\prod_s$ is the Jacobian product propagation during backward. The first $T-K$ steps run under torch.no_grad(); only the last $K$ steps retain the graph. K=1 (DRaFT-1) already gives a very strong signal — the last step has the largest direct effect on $x_0$. Autograd automatically sums the $\partial/\partial\theta|_\text{direct}$ across all $K$ steps, so the code only needs loss.backward().

Trick 2: LoRA fine-tune + high learning rate.

Only train a LoRA adapter (~1% parameters), with the base UNet frozen. Combined with gradient checkpointing, memory fits within 24 GB on a single GPU.

Pseudo-code:

# DRaFT-K one-step training
x_t = torch.randn(B, C, H, W).to(device)  # x_T
with torch.no_grad():
    for t in range(T-1, K, -1):           # first T-K steps without grad
        x_t = ddim_step(unet_lora, x_t, t, cond)
for t in range(K, 0, -1):                  # last K steps with grad
    x_t = ddim_step(unet_lora, x_t, t, cond)
x_0 = x_t
reward = image_rm(x_0, prompt)             # ImageReward / HPSv2
loss = -reward.mean()                       # note negative sign — we maximize reward
loss.backward()                             # memory O(K)
optimizer.step()

3.3 AlignProp (Prabhudesai et al. arXiv 2310.03739, 2023-10; ICLR 2024 venue; the arXiv version was later superseded/withdrawn)

Aligning Text-to-Image Diffusion Models with Reward Backpropagation — almost parallel to DRaFT (both appeared on arXiv in late 2023), with the same core idea: reward backprop through denoising. Differences:

AlignProp's key contribution is theorizing "why reward backprop works" — proving that under fixed-point assumptions, the truncated BPTT gradient is a biased but low-variance estimator of the true gradient.

3.4 ReFL (Xu et al. 2023 NeurIPS, arXiv 2304.05977)

ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation. This is the original ImageReward paper, which also proposes the ReFL (Reward Feedback Learning) algorithm.

ReFL is close to DRaFT in spirit but predates it. It treats reward directly as a loss and supervises only at one randomly selected intermediate step, an early implementation of the DRaFT idea:

$$\mathcal{L}_\text{ReFL} = \mathcal{L}_\text{simple} - \lambda \cdot \mathbb{E}_{t' \sim [t_\text{min}, t_\text{max}]}\!\big[R\big(\hat{x}_0(x_{t'}, t')\big)\big]$$

where $\hat{x}_0(x_{t'}, t') = (x_{t'} - \sqrt{1-\bar\alpha_{t'}}\epsilon_\theta(x_{t'}, t'))/\sqrt{\bar\alpha_{t'}}$ is the $x_0$ estimate from a single-step $\epsilon$-prediction (one-step Tweedie unfolding).

Key difference: ReFL is "single-step backprop + $L_\text{simple}$ trained jointly", while DRaFT/AlignProp is "multi-step backprop + pure reward loss". ReFL is more stable but reward gains are smaller, because it only sees the one-step $x_0$ estimate rather than the true sampling trajectory.

3.5 Three-way comparison

3.6 Reward hacking in direct reward backprop

Direct backprop is more prone to hacking than RL because the gradient signal is "too precise":

Mitigations: reward ensemble (mean or min of HPSv2 + PickScore + ImageReward), KL anchor, early stop, small LoRA scale.

§4 Preference Optimization: the Diffusion-DPO family

4.1 Diffusion-DPO (Wallace et al. 2024 CVPR, arXiv 2311.12908)

Diffusion Model Alignment Using Direct Preference Optimization. The key challenge in porting LLM DPO to diffusion: diffusion's $\log p_\theta(x_0 \mid c)$ has no closed form — must use an ELBO substitute.

4.1.1 Derivation (key steps)

The LLM DPO core is the KL-regularized RL optimal solution:

$$\pi^*(y \mid x) \propto \pi_\text{ref}(y \mid x) \exp\!\big(r(x, y) / \beta\big)$$

Solving for $r = \beta \log(\pi^*/\pi_\text{ref}) + \beta \log Z$ and substituting into Bradley-Terry, the $\log Z$ cancels.

For diffusion, replace "sample $y$" with "sample trajectory $(x_T, \dots, x_0)$"; the optimal form is the same, but $\log p$ is the full-trajectory likelihood:

$$\log p_\theta(x_{0:T} \mid c) = \log p(x_T) + \sum_{t=1}^{T} \log p_\theta(x_{t-1} \mid x_t, c)$$

This trajectory log-likelihood is analytical (each term is a Gaussian log-prob). However: if every update step requires a full trajectory run, the compute cost explodes.

Wallace et al.'s trick: use an ELBO surrogate.

DDPM's $L_\text{simple}$ is one of the (negative) ELBO terms of $\log p_\theta(x_0)$; specifically:

$$-\log p_\theta(x_0 \mid c) \le L_\text{simple}(x_0, c, \theta) = \mathbb{E}_{t, \epsilon}\!\left[\|\epsilon - \epsilon_\theta(x_t, t, c)\|^2\right] + \text{const}$$

Use $-L_\text{simple}$ as a single-sample estimate of $\log p_\theta(x_0 \mid c)$ (Jensen's inequality strictly gives a lower bound, but numerically it is usable as a DPO implicit-reward proxy) and substitute into the DPO framework:

$$\boxed{\;\mathcal{L}_\text{Diff-DPO}(\theta) = -\mathbb{E}_{(x_0^w, x_0^l, c, t, \epsilon)}\log\sigma\!\left(-\beta T\!\left[\|\epsilon^w - \epsilon_\theta(x_t^w, t, c)\|^2 - \|\epsilon^w - \epsilon_\text{ref}(x_t^w, t, c)\|^2 - \|\epsilon^l - \epsilon_\theta(x_t^l, t, c)\|^2 + \|\epsilon^l - \epsilon_\text{ref}(x_t^l, t, c)\|^2\right]\right)\;}$$

4.1.2 Implementation details

• $(x_0^w, x_0^l)$ come from a human-preference pair conditioned on a single prompt $c$ (Pick-a-Pic is the main dataset).
• Each training step randomly samples $t \in \{1, \dots, T\}$ and $\epsilon \sim \mathcal{N}(0,I)$, constructing $x_t^w = \sqrt{\bar\alpha_t} x_0^w + \sqrt{1-\bar\alpha_t}\epsilon$ (the same $\epsilon$ is used for both $y_w$ and $y_l$ — paired noise).
• $\pi_\text{ref}$ is a frozen base UNet (usually the original SDXL checkpoint).

4.1.3 Results (on SDXL)

• Consistently better than SDXL base on PickScore / HPSv2.
• Training cost is ~1.5–2× SFT (must run UNet twice: policy + ref).
• Simpler than DDPO: fully offline, no sampling needed.

4.2 D3PO (Yang et al. 2024 CVPR, arXiv 2311.13231)

Using Human Feedback to Fine-tune Diffusion Models without Any Reward Model. Posted to arXiv around the same time as Diffusion-DPO (2023-11); the difference is in derivation path:

• Diffusion-DPO: KL-regularized RL → ELBO surrogate → DPO loss.
• D3PO: directly port the LLM DPO derivation step-by-step onto the diffusion Markov chain — each denoising step is treated as an MDP step, using the same "solve for implicit reward + plug into BT" framework.

The final D3PO loss form is nearly the same as Diffusion-DPO's:

$$\mathcal{L}_\text{D3PO}(\theta) = -\mathbb{E}_{(\tau^w, \tau^l)} \log\sigma\!\left(\beta \sum_{t=1}^{T}\!\left[\log\frac{p_\theta(x_{t-1}^w \mid x_t^w, c)}{p_\text{ref}(x_{t-1}^w \mid x_t^w, c)} - \log\frac{p_\theta(x_{t-1}^l \mid x_t^l, c)}{p_\text{ref}(x_{t-1}^l \mid x_t^l, c)}\right]\right)$$

What's needed is a full trajectory pair $(\tau^w, \tau^l)$; if the preference pair only has the final image $(x_0^w, x_0^l)$, reconstruct the trajectory using $q(x_{1:T} \mid x_0)$ (with DDPM's forward q-sample).

• Diffusion-DPO uses single-$t$ estimation (cheaper), D3PO sums over the full trajectory (more accurate but expensive).
• Diffusion-DPO is stable on Pick-a-Pic; D3PO is stable on collected thumbs up/down data.
• Industrial deployments mostly use Diffusion-DPO (simpler compute).

4.3 SPO (Liang et al. 2024, arXiv 2406.04314)

Step-aware Preference Optimization: Aligning Preference with Denoising Performance at Each Step. Key observation: different denoising steps are responsible for different aspects of the image:

• High-noise steps ($t \approx T$): determine global structure (composition, object placement).
• Low-noise steps ($t \approx 0$): determine local detail (texture, edges).

Diffusion-DPO's "single-$t$ sampling" treats all steps equally — but human preferences have different "importance" across steps.

4.3.1 The two SPO modifications

Modification 1: In-step preference — for the same $x_t$, independently sample two $x_{t-1}^w, x_{t-1}^l$, and use a step-wise reward model to judge "which $x_{t-1}$ is better at step $t$".

Modification 2: Step-aware weighting — the SPO loss is weighted along the step dimension:

$$\mathcal{L}_\text{SPO}(\theta) = -\mathbb{E}_{t \sim w(t),\; x_t}\!\left[\log\sigma\!\left(\beta\!\log\frac{p_\theta(x_{t-1}^w \mid x_t, c)}{p_\text{ref}(x_{t-1}^w \mid x_t, c)} - \beta\!\log\frac{p_\theta(x_{t-1}^l \mid x_t, c)}{p_\text{ref}(x_{t-1}^l \mid x_t, c)}\right)\right]$$

where $w(t)$ is the step sampling distribution (typically uniform or biased toward middle $t$).

4.3.2 The in-step reward model

To obtain the in-step preference $(x_{t-1}^w, x_{t-1}^l)$, SPO trains a step-wise reward model $R(x_{t-1}, x_t, c, t)$ that judges "given $x_t$, is $x_{t-1}$ a good transition at step $t$?" It does not directly score pixel quality of $x_{t-1}$ but rather estimates whether this transition at step $t$ leads to a high-quality $x_0$.

4.4 Diffusion-KTO (Li et al. 2024 NeurIPS, arXiv 2404.04465)

Aligning Diffusion Models by Optimizing Human Utility. The diffusion version of LLM KTO (Ethayarajh 2024, arXiv 2402.01306).

LLM KTO core idea: replace the BT preference model with Kahneman-Tversky prospect theory, requiring only per-sample binary feedback (thumbs up/down), no pairs needed:

$$L_\text{KTO} = \mathbb{E}_{x, y}\!\left[\lambda_y v\!\big(\beta \log\frac{\pi_\theta(y|x)}{\pi_\text{ref}(y|x)} - z_0(x)\big)\right]$$

where $v(\cdot)$ is the prospect-theoretic value function (thumbs up uses $1 - \sigma(\cdot)$, thumbs down uses $\sigma(\cdot)$), and $z_0$ is the reference utility.

Diffusion-KTO replaces $\log(\pi_\theta/\pi_\text{ref})$ with the Diffusion-DPO $\epsilon$-distance ELBO surrogate:

$$L_\text{Diff-KTO} = \mathbb{E}_{x_0, c, \text{label}}\!\left[\lambda_\text{label}\, v\!\left(\beta T \left[\|\epsilon - \epsilon_\text{ref}\|^2 - \|\epsilon - \epsilon_\theta\|^2\right] - z_0(c)\right)\right]$$

4.5 MaPO (Hong et al. 2024, arXiv 2406.06424)

Margin-aware Preference Optimization for Aligning Diffusion Models without Reference. Core idea: drop the reference model entirely — similar in spirit to LLM SimPO.

MaPO loss simultaneously optimizes two things:

1. Likelihood margin: $\log p_\theta(x_0^w) - \log p_\theta(x_0^l)$ (estimated via the ELBO surrogate $\|\epsilon - \epsilon_\theta\|^2$).
2. Likelihood of preferred: $\log p_\theta(x_0^w)$ itself must be high (to prevent "both sides dropping").

$$\mathcal{L}_\text{MaPO}(\theta) = -\mathbb{E}\!\left[\log\sigma\!\big(\beta(\hat{\ell}_w - \hat{\ell}_l) - \gamma\big) + \alpha \hat{\ell}_w\right]$$

where $\hat{\ell} = -\|\epsilon - \epsilon_\theta(x_t, t, c)\|^2$ is the likelihood surrogate, $\gamma$ is the margin, and $\alpha$ is the likelihood-term weight.

Advantages:

• No reference UNet needed: saves half the memory (from $2 \times M$ to $M$).
• Solves reference mismatch: when fine-tuning to a new style (where the reference is far from the target distribution), Diffusion-DPO training collapses, but MaPO is stable.
• 15% faster training (paper-reported, validated across 5 domains).

4.6 DPO family overview

§5 Flow-GRPO: RL for Flow Matching

5.1 Why Flow Matching also needs post-training

SD3 / FLUX / Lumina all moved to Flow Matching (Rectified Flow); post-training needs are the same:

• Improve composition benchmarks like GenEval / DPG (color, count, spatial relations).
• Improve OCR / text rendering accuracy.
• Improve prompt-image alignment (VLM judge).

But Flow Matching is a deterministic ODE ($\dot x_t = v_\theta(t, x, c)$) — the stochastic transition assumed by DDPO/DPOK does not exist. Directly plugging it into the RL framework fails.

5.2 Flow-GRPO's two core tricks

Liu et al. 2025 Flow-GRPO: Training Flow Matching Models via Online RL (arXiv 2505.05470) solves two fundamental Flow + RL problems:

Trick 1: ODE → SDE equivalent conversion

For the Rectified Flow ODE $\dot x_t = v_\theta(t, x_t, c)$, construct an equivalent SDE:

$$dx_t = \big[v_\theta(t, x_t, c) + \tfrac{1}{2}\sigma(t)^2 \nabla_x \log p_t(x_t)\big]\,dt + \sigma(t)\,dW_t$$

Key property (Song et al. 2021 score SDE framework): this SDE has exactly the same marginal $p_t$ as the original ODE. The difference is that the SDE provides stochastic exploration (the $dW_t$ noise), allowing RL to sample different trajectories.

For Flow Matching, $\nabla_x \log p_t = -\epsilon/\sigma_t$ (under a Gaussian path), and the score can be derived from $v_\theta$. Setting $\sigma(t)$ as a schedule (typically $\sigma(t) = \sqrt{1-t}$) gives the SDE sampler used by Flow-GRPO during training.

Trick 2: Denoising reduction

GRPO needs a group of $G$ trajectories per sample, typically $G = 16$–$32$. Flow Matching inference normally uses 25–50 steps, so a single training sample needs $\approx 25G$ forwards — too expensive.

Flow-GRPO uses fewer steps during training (e.g. 10 steps); inference still uses 25–50 steps. The SDE is more uniform on the schedule, so the "exploration quality" of few-step training is enough:

$$\text{Training}: T_\text{train} = 10, \quad \text{Inference}: T_\text{infer} = 28$$

Empirically, no quality drop on GenEval / OCR / Aesthetic.

5.3 Flow-GRPO advantage computation

Completely parallel to GRPO for LLM — for the same prompt $c$, sample $G$ final images $\{x_0^{(1)}, \dots, x_0^{(G)}\}$, score each with reward $r_i$, and normalize within the group:

$$\hat{A}_i = \frac{r_i - \text{mean}_{j}(r_j)}{\text{std}_j(r_j) + \epsilon}$$

All steps within a trajectory share the same $\hat{A}_i$ (same as per-token sharing in LLM GRPO).

5.4 Flow-GRPO loss

Let the SDE Euler step transition log-prob be $\log p_\theta(x_{t-1} \mid x_t, c)$ (Gaussian), with importance ratio $\rho_{i,t} = p_\theta / p_{\theta_\text{old}}$. PPO-clip:

$$L^\text{Flow-GRPO} = \mathbb{E}\!\left[\frac{1}{G}\sum_i\!\frac{1}{T_\text{train}}\!\sum_t \min\!\big(\rho_{i,t}\hat{A}_i, \text{clip}(\rho_{i,t}, 1-\epsilon, 1+\epsilon)\hat{A}_i\big) - \beta\, \text{KL}_{i,t}(p_\theta \Vert p_\text{ref})\right]$$

KL still uses the K3 estimator (same as GRPO for LLM).

5.5 Geometric meaning of advantage on the vector field

• LLM GRPO's advantage reweights in the token logit space;
• Flow-GRPO's advantage directly reweights the direction correction of $v_\theta$ — specifically, when $\hat A_i > 0$, push $v_\theta(t, x_t, c)$ toward the direction $(x_{t-1}^{(i)} - x_t^{(i)})/dt$ that trajectory $\tau_i$ actually took.
• This is a "directional gradient" on the $v_\theta$ space, equivalent to importance-weighted $\epsilon$-prediction reweighting under a Gaussian path.

5.6 Flow-GRPO empirical results

The paper reports for SD3.5-M:

§6 Code Patterns (readable pseudo-code)

6.1 DDPO REINFORCE-style update

import torch
import torch.nn.functional as F

def ddpo_step(unet, ref_unet, scheduler, prompts, reward_fn,
              T=20, B=4, lr=1e-5, beta=0.0):
    """
    DDPO-SF one-step training (REINFORCE + optional KL anchor).
    prompts: list of B text prompts
    reward_fn: callable (x0_batch, prompts) -> [B] scalar
    """
    # ── 1. Rollout: sample G=B trajectories ──
    x = torch.randn(B, 4, 64, 64, device=device)
    traj_log_probs = []
    with torch.set_grad_enabled(False):                   # rollout needs no grad
        x_t = x
        for t in reversed(range(T)):
            # predict noise + sample x_{t-1}
            eps_pred = unet(x_t, t, prompts)
            mean, std = scheduler.step_mean_std(x_t, eps_pred, t)
            x_tm1 = mean + std * torch.randn_like(mean)   # stochastic transition
            traj_log_probs.append((mean.detach(), std.detach(), x_tm1.detach()))
            x_t = x_tm1
        x_0 = x_t

    # ── 2. Reward ──
    R = reward_fn(x_0, prompts)                            # [B]
    A = (R - R.mean()) / (R.std() + 1e-8)                  # batch baseline

    # ── 3. Policy gradient: re-forward to get log p_θ ──
    x_t = x.detach()
    loss_pg = 0.0
    for t, (mean_old, std_old, x_tm1) in zip(reversed(range(T)), traj_log_probs):
        eps_pred = unet(x_t, t, prompts)                   # grad needed
        mean, std = scheduler.step_mean_std(x_t, eps_pred, t)
        # Gaussian log-prob
        log_p = -0.5 * (((x_tm1 - mean) / std) ** 2).sum([1, 2, 3])
        log_p -= std.log().sum([1, 2, 3])
        loss_pg = loss_pg - (log_p * A).mean()             # REINFORCE
        if beta > 0:
            ref_eps = ref_unet(x_t, t, prompts).detach()
            mean_ref, std_ref = scheduler.step_mean_std(x_t, ref_eps, t)
            # Gaussian-Gaussian KL closed form
            kl = ((mean - mean_ref) ** 2 / (2 * std_ref ** 2)
                  + (std / std_ref) ** 2 / 2
                  - 0.5 - (std / std_ref).log()).sum([1, 2, 3])
            loss_pg = loss_pg + beta * kl.mean()
        x_t = x_tm1.detach()
    return loss_pg

• Use set_grad_enabled(False) during rollout, then enable grad in the policy-gradient pass — avoids $O(T)$ memory.
• DDPM stochastic transition is critical: DDIM is deterministic with no $dW$ dimension to optimize, so DDPO must use DDPM or DDIM-eta=1.
• The batch baseline $A = (R - \bar R)/\sigma_R$ is much more stable than no baseline.
• Train LoRA rather than full fine-tune, otherwise the base drifts quickly.

6.2 Diffusion-DPO loss

def diffusion_dpo_loss(unet, ref_unet, scheduler,
                       x0_w, x0_l, prompt_embeds, beta=2000.0):
    """
    Diffusion-DPO (Wallace 2024) one-step training.
    x0_w, x0_l: [B, 4, H, W]  preferred / dispreferred latents
    beta: paper uses 2000~5000 (note: this is the combined β·T coefficient, larger than LLM DPO)
    """
    B = x0_w.shape[0]
    t = torch.randint(0, scheduler.num_train_timesteps, (B,), device=x0_w.device)
    noise = torch.randn_like(x0_w)                          # paired noise!

    xt_w = scheduler.add_noise(x0_w, noise, t)
    xt_l = scheduler.add_noise(x0_l, noise, t)

    # ── policy ε-prediction ──
    eps_w = unet(xt_w, t, prompt_embeds)
    eps_l = unet(xt_l, t, prompt_embeds)

    # ── reference ε-prediction (frozen) ──
    with torch.no_grad():
        ref_eps_w = ref_unet(xt_w, t, prompt_embeds)
        ref_eps_l = ref_unet(xt_l, t, prompt_embeds)

    # ── ELBO surrogate: -‖ε - ε_θ‖² is a proxy for log p_θ ──
    err_w_pol = ((noise - eps_w) ** 2).mean([1, 2, 3])      # [B]
    err_w_ref = ((noise - ref_eps_w) ** 2).mean([1, 2, 3])
    err_l_pol = ((noise - eps_l) ** 2).mean([1, 2, 3])
    err_l_ref = ((noise - ref_eps_l) ** 2).mean([1, 2, 3])

    # DPO log-ratio: smaller err = better likelihood
    #   log(π_θ/π_ref)(y_w) ≈ -(err_w_pol - err_w_ref)
    diff_w = -(err_w_pol - err_w_ref)
    diff_l = -(err_l_pol - err_l_ref)

    inner = beta * (diff_w - diff_l)
    loss = -F.logsigmoid(inner).mean()

    with torch.no_grad():
        margin = inner.mean()
        accuracy = (inner > 0).float().mean()
    return loss, {"margin": margin.item(), "acc": accuracy.item()}

6.3 AlignProp / DRaFT backprop with checkpointing

def alignprop_step(unet_lora, scheduler, prompts, reward_fn,
                   T=50, K=1, B=4, lr=1e-5):
    """
    DRaFT-K / AlignProp one-step training.
    K: last K steps retain gradient
    Memory O(K); K=1 is single-forward magnitude on SDXL
    """
    x = torch.randn(B, 4, 64, 64, device=device)

    # ── First T - K steps without grad ──
    with torch.no_grad():
        for t in reversed(range(K, T)):
            eps = unet_lora(x, t, prompts)
            x = scheduler.step_ddim(x, eps, t)             # deterministic DDIM

    # ── Last K steps with grad ──
    for t in reversed(range(K)):
        eps = unet_lora(x, t, prompts)                     # gradient ON
        x = scheduler.step_ddim(x, eps, t)

    x_0 = x
    # ── Backprop ──
    reward = reward_fn(x_0, prompts)                       # [B]
    loss = -reward.mean()                                  # max reward = min -reward
    return loss

# Memory analysis:
#   K=1:  ~24 GB on SDXL (single forward + grad)
#   K=5:  ~60 GB
#   K=10: ~120 GB (needs multi-GPU)
#   K=T=50: ~600 GB (completely infeasible)

6.4 SPO step-aware preference loss

def spo_loss(unet, ref_unet, scheduler, step_rm,
             x_t, t, prompt_embeds, beta=500.0):
    """
    SPO (Liang 2024) in-step preference.
    Given x_t and t, independently sample two x_{t-1}, let step_rm judge winner.
    """
    # ── 1. Use policy to sample two x_{t-1} candidates (must no_grad,
    #      otherwise DPO log-prob backprop would flow back into the sampler) ──
    with torch.no_grad():
        eps_sample = unet(x_t, t, prompt_embeds)
        mean_s, std_s = scheduler.step_mean_std(x_t, eps_sample, t)
        noise_a, noise_b = torch.randn_like(mean_s), torch.randn_like(mean_s)
        x_a = mean_s + std_s * noise_a
        x_b = mean_s + std_s * noise_b

    # ── 2. step-wise reward model judges winner ──
    with torch.no_grad():
        r_a = step_rm(x_a, x_t, t, prompt_embeds)          # [B]
        r_b = step_rm(x_b, x_t, t, prompt_embeds)
        winner = (r_a > r_b).long()                        # [B], 1 if a wins
    x_w = torch.where(winner.bool()[:, None, None, None], x_a, x_b).detach()
    x_l = torch.where(winner.bool()[:, None, None, None], x_b, x_a).detach()

    # ── 3. compute log p_θ / log p_ref for x_w, x_l (grad-aware forward) ──
    eps = unet(x_t, t, prompt_embeds)
    mean, std = scheduler.step_mean_std(x_t, eps, t)
    log_p_w = -0.5 * ((x_w - mean) / std).pow(2).sum([1, 2, 3])
    log_p_l = -0.5 * ((x_l - mean) / std).pow(2).sum([1, 2, 3])
    with torch.no_grad():
        ref_eps = ref_unet(x_t, t, prompt_embeds)
        ref_mean, ref_std = scheduler.step_mean_std(x_t, ref_eps, t)
        log_pref_w = -0.5 * ((x_w - ref_mean) / ref_std).pow(2).sum([1, 2, 3])
        log_pref_l = -0.5 * ((x_l - ref_mean) / ref_std).pow(2).sum([1, 2, 3])

    inner = beta * ((log_p_w - log_pref_w) - (log_p_l - log_pref_l))
    return -F.logsigmoid(inner).mean()

6.5 Flow-GRPO group-relative advantage

def flow_grpo_step(flow_net, ref_flow, prompts, reward_fn,
                   G=16, T_train=10, sigma_fn=lambda t: (1 - t) ** 0.5,
                   eps_clip=0.2, beta=0.04):
    """
    Flow-GRPO one-step training.
    G: sample G trajectories per prompt.
    T_train: SDE step count for training (inference uses 28-50).
    """
    P = len(prompts)
    # Repeat each prompt G times
    prompts_rep = sum([[p] * G for p in prompts], [])      # [P*G]

    # ── 1. SDE rollout: ODE→SDE equivalent conversion ──
    x_t = torch.randn(P * G, 4, 64, 64, device=device)
    log_probs_old = []                                     # for PPO importance ratio
    trajectory = [x_t.clone()]
    with torch.no_grad():
        for i in range(T_train):
            t_now = 1.0 - i / T_train
            t_next = 1.0 - (i + 1) / T_train
            dt = t_next - t_now
            sigma = sigma_fn(t_now)
            v = flow_net(x_t, t_now, prompts_rep)
            # SDE Euler: drift = v + 0.5 σ² ∇log p (PF-ODE → SDE conversion, Song 2021)
            # !!! IMPORTANT: the following drift is a simplified pedagogical placeholder.
            #     Production implementations follow Flow-GRPO paper Eq.(6) and derive score
            #     correctly from score = (data_pred - x_t)/σ_t² with the specific
            #     Rectified Flow / EDM schedule. Refer to paper + official repo for deployment;
            #     the -v/σ here is only illustrative.
            drift = v + 0.5 * sigma ** 2 * (-v / (sigma + 1e-6))  # placeholder, see paper Eq.(6)
            noise = torch.randn_like(x_t)
            x_next = x_t + drift * dt + sigma * noise * abs(dt) ** 0.5
            # Gaussian log-prob (transition)
            mean = x_t + drift * dt
            std = sigma * abs(dt) ** 0.5
            log_p = -0.5 * ((x_next - mean) / std).pow(2).sum([1, 2, 3])
            log_probs_old.append(log_p)
            x_t = x_next
            trajectory.append(x_t.clone())
        x_0 = x_t

    # ── 2. Group-relative advantage ──
    R = reward_fn(x_0, prompts_rep)                        # [P*G]
    R = R.view(P, G)
    mean_R = R.mean(dim=1, keepdim=True)
    std_R = R.std(dim=1, keepdim=True) + 1e-8
    A = ((R - mean_R) / std_R).view(P * G)                 # [P*G]

    # ── 3. PPO-clip loss with KL ──
    loss = 0.0
    x_t = trajectory[0]
    for i in range(T_train):
        t_now = 1.0 - i / T_train
        v = flow_net(x_t, t_now, prompts_rep)              # grad ON
        sigma = sigma_fn(t_now)
        drift = v + 0.5 * sigma ** 2 * (-v / (sigma + 1e-6))
        dt = -1.0 / T_train
        mean = x_t + drift * dt
        std = sigma * abs(dt) ** 0.5
        log_p_new = -0.5 * ((trajectory[i+1] - mean) / std).pow(2).sum([1, 2, 3])
        ratio = (log_p_new - log_probs_old[i]).exp()
        surr1 = ratio * A
        surr2 = ratio.clamp(1 - eps_clip, 1 + eps_clip) * A
        loss = loss - torch.min(surr1, surr2).mean()

        # K3 KL estimator
        with torch.no_grad():
            v_ref = ref_flow(x_t, t_now, prompts_rep)
            drift_ref = v_ref + 0.5 * sigma ** 2 * (-v_ref / (sigma + 1e-6))
            mean_ref = x_t + drift_ref * dt
            log_p_ref = -0.5 * ((trajectory[i+1] - mean_ref) / std).pow(2).sum([1,2,3])
        delta = log_p_ref - log_p_new
        kl_k3 = (delta.exp() - delta - 1)
        loss = loss + beta * kl_k3.mean()

        x_t = trajectory[i + 1].detach()
    return loss

6.6 Combined reward signal

def combined_reward(images, prompts, weights=None):
    """
    Weighted multi-reward combination — mitigates single-RM hacking.
    """
    weights = weights or {"image_reward": 0.4, "hps_v2": 0.3,
                          "pickscore": 0.2, "clip_score": 0.1}
    rewards = {}
    rewards["image_reward"] = image_reward_model(images, prompts)        # [-1, 4]
    rewards["hps_v2"] = hps_v2(images, prompts)                          # [0, 1]
    rewards["pickscore"] = pickscore(images, prompts)                    # logits
    rewards["clip_score"] = clip_cosine(images, prompts)                 # [-1, 1]

    # ── z-score each individually (different rewards have very different scales) ──
    normed = {k: (v - v.mean()) / (v.std() + 1e-8) for k, v in rewards.items()}

    # ── weighted + length / safety penalty ──
    R = sum(weights[k] * normed[k] for k in weights)

    # NSFW penalty (rule-based)
    nsfw_score = nsfw_detector(images)                                   # [0, 1]
    R = R - 5.0 * nsfw_score

    return R

• Z-score each reward separately: the scales differ enormously (HPSv2 ~0.25, ImageReward ~1.5, CLIP-Score ~0.3); not normalizing means ImageReward dominates.
• Min is more stable than mean: R = min(normed.values()) forces all RMs to agree on "good", drastically lowering hacking risk (classic reward-ensemble strategy).
• Keep rule-based safety overrides: NSFW / political / copyright rewards must not be optimized away by RL.

§7 Reward Design & failure modes

7.1 Reward model choices

7.2 Reward hacking gallery (diffusion-specific)

7.3 Step-level vs trajectory-level reward

7.4 Core mechanisms to mitigate reward hacking

1. Reward ensemble: take min or mean across multiple RMs (HPSv2 + PickScore + ImageReward is the mainstream combo).
2. KL anchor: DPO's $\beta$, DPOK's explicit KL, Flow-GRPO's K3 KL term.
3. LoRA scale: full fine-tune drifts fast; LoRA scale caps reward-hacking ceiling.
4. Early stop on reward plateau: reward up + FID up = hacking signal.
5. Composite reward: rule-based (object count, OCR) + neural RM (aesthetic, alignment) weighted.
6. Adversarial RM: include hacking samples as negatives when training RM.

§8 Production landscape: what SD3 / FLUX actually use

8.1 What the public papers / reports actually say

8.2 Did SD3 / FLUX use post-training?

• SD3 paper (arXiv 2403.03206) "Improving Rectified Flow Transformers" section discusses sampling + reflow, with no mention of reward fine-tuning.
• FLUX has no paper at all; the community infers distillation from the model card (FLUX schnell is a 4-step distilled version).
• Stability AI mentioned in the SD3.5-Large release that it was "fine-tuned with improved aesthetics", which could be SFT rather than RL.
• DALL-E 3 paper (OpenAI 2023) explicitly says caption-faithful RLHF was used.

Industry consensus (from HuggingFace community + Reddit r/StableDiffusion): closed-source large models (FLUX pro, DALL-E 3, Midjourney v6+) have reward-based fine-tuning but the specifics are undisclosed; open-source bases (SD3.5 base, FLUX dev base) have public training pipelines without RL, but Stability AI's internal dev branches may have it.

8.3 An industrial-grade pipeline hypothesis

                    ┌─────────────────┐
                    │  LDM Pretrain   │  several 100M / B images, $L_simple$
                    └────────┬────────┘
                             │
                    ┌────────▼────────┐
                    │  SFT on Curated │  high-quality prompt-image pairs
                    │     dataset     │  (Aesthetic > 6.0, no watermark)
                    └────────┬────────┘
                             │
            ┌────────────────┼────────────────┐
            │                                  │
   ┌────────▼────────┐              ┌─────────▼────────┐
   │ Diffusion-DPO   │              │  DRaFT / AlignProp│
   │ on Pick-a-Pic    │              │   on multi-RM    │
   └────────┬────────┘              └─────────┬────────┘
            │                                  │
            └────────────────┬─────────────────┘
                             │
                    ┌────────▼────────┐
                    │   Distillation  │  4-step / 1-step turbo
                    │   (ADD / LCM)   │
                    └────────┬────────┘
                             │
                       Production

§9 vs LLM RLHF comparison

9.1 The full table

9.2 Shared lessons

1. KL anchor is mandatory: pure RL without ref policy always reward-hacks (LLMs produce vacuous fluff, diffusion produces over-saturated monotonous compositions).
2. DPO family >> on-policy RL (engineering-wise): no sampling, no value model, offline — true for both LLM and diffusion.
3. Reward ensemble counters hacking: min-of-K RMs is the common mitigation across domains.
4. Group-based advantage: LLM's GRPO/RLOO and diffusion's Flow-GRPO both bypass the value model via group statistics.

9.3 Diffusion-specific differences

• Diffusion has a "backprop" option: reward differentiability makes DRaFT/AlignProp viable — LLMs cannot do this because sampling is discrete.
• Diffusion has "step-aware" preference: denoising steps have explicit semantics (high-noise: composition, low-noise: detail); SPO exploits this — LLM tokens lack such natural stratification.
• Diffusion's $\beta$ scale is 1000× larger: because the ELBO surrogate absorbs the $T$-fold trajectory term.

§10 25 frequently-asked interview questions

Three difficulty tiers: L1 = multimodal/diffusion-role basics; L2 = research/alignment-focused; L3 = top-lab hardcore.

L1 must-know (10 questions)

L2 advanced (10 questions)

L3 top-lab questions (5 questions)

§A Appendix

A.1 Key papers (with arXiv IDs)

A.2 Common reward model resources

• ImageReward: https://github.com/THUDM/ImageReward
• HPSv2: https://github.com/tgxs002/HPSv2
• PickScore: https://github.com/yuvalkirstain/PickScore
• CLIP: OpenAI / OpenCLIP, multiple backbones available

A.3 Open-source training code

• TRL (HuggingFace): diffusers + DPO Trainer for Diffusion-DPO (most mature)
• DDPO original repo: https://github.com/kvablack/ddpo-pytorch
• AlignProp: https://github.com/mihirp1998/AlignProp
• DRaFT (Google research): https://github.com/clarkjkr/draft (Clark et al. 2024 ICLR)
• MaPO: https://github.com/mapo-t2i/mapo
• Flow-GRPO: see paper arXiv 2505.05470 for the official implementation

A.4 Engineering footgun checklist

A.5 Mapping back to §0 TL;DR

• Can verbally describe the Diffusion-DPO loss form + paired-noise detail
• Can explain why AlignProp's $K=1$ is enough + memory $\mathcal{O}(K)$
• Can write DDPO's state/action/reward + per-step ratio
• Can explain why Flow-GRPO's ODE→SDE conversion is necessary + denoising reduction
• Know the honest answer about whether SD3/FLUX used RL (publicly undisclosed)