Divergence Minimization Preference Optimization for Diffusion Model Alignment

TL;DR

DMPO (Divergence Minimization Preference Optimization) is a preference-alignment method for diffusion models that replaces Diffusion-DPO’s implicit forward KL objective with an explicit reverse KL between the learned policy and the RLHF-optimal Boltzmann policy. The paper shows that Diffusion-DPO’s gradient is equivalent (in the multi-sample generalization) to minimizing forward KL on the sampling-path distribution, which is mean-seeking — it covers diverse modes but blurs them. Reverse KL is mode-seeking, so the post-trained model concentrates probability mass on the highest-preference modes; the authors derive a tractable per-step objective for diffusion via Jensen’s inequality on a convexity-preserving interval, and prove the resulting gradient direction agrees with vanilla RLHF under mild assumptions. On Stable Diffusion v1.5 and SDXL fine-tuned on Pick-a-Pic V2, DMPO beats Diffusion-DPO, Diffusion-KTO, and DSPO across PickScore / HPS / CLIP / Aesthetics / ImageReward on Pick-a-Pic V2, HPS V2, and PartiPrompt — with at-least-64.6% PickScore win-rate over every preference-pair-trained baseline. This is the cleanest “swap forward KL for reverse KL” recipe in diffusion alignment to date.

Key claims

Diffusion-DPO’s gradient (generalized from the pairwise to the multi-sample setting with preferences sampled from the reference policy) is proportional to the gradient of KL(π* ‖ π_θ) over the sampling-path distribution — i.e. forward KL, which is mean-seeking [Theorem 1, §3.1, Eq. 8].
Reverse KL KL(π_θ ‖ π*) is mode-seeking — it only penalizes the model when it places probability in regions unsupported by π*, driving the post-trained policy toward the high-reward modes of preferences rather than covering their full support [§3.2, “the reverse KL is given by …”].
The reverse-KL objective can be rewritten in terms of the relative logit λ_θ(x|c) = log[π_θ(x|c)/π_ref(x|c)] and the (unnormalized) preference-optimal distribution π_y* — giving a divergence objective that does not require evaluating the partition function Z(c) [§3.2, Eq. 9–10].
For pairwise preference data with smoothing coefficient α (probability assigned to the loser), DMPO reduces to a per-prompt scalar function f(z) of the log-likelihood-ratio difference z = λ_θ(x_w) − λ_θ(x_l); the paper shows f''(z) > 0 over the practical optimization interval, making Jensen’s inequality usable to push the expectation inside the convex function [§3.2, Eq. 11–12, convexity argument].
The intractable trajectory likelihood is handled by approximating the reverse posterior with the forward q(x_{t−1} | x_t, x_0), yielding a per-step DMPO loss that adds no computational overhead over Diffusion-DPO [§3.2, Eq. 14].
Theorem 2: as the smoothing α → 0, the DMPO gradient direction equals the RLHF objective’s gradient direction — i.e. DMPO and reward-model RLHF push the policy the same way under the same KL regularization [§3.3, Theorem 2, Eq. 16].
On SD1.5 fine-tuned on Pick-a-Pic V2, DMPO is the only method that consistently improves every reward metric (PickScore, HPS, CLIP, Aesthetics, ImageReward) over the base on all three test sets (Pick-a-Pic V2, HPS V2, PartiPrompt) [Table 1].
DMPO posts ≥64.6% PickScore win-rate over every preference-pair-trained baseline (Diffusion-DPO, Diffusion-KTO, DSPO, MAPO) across the three test sets, and a ≥65% win-rate vs SD1.5 itself on HPS, ImageReward, and Aesthetics [Table 2].
Diffusion-DPO underperforms an SFT baseline on Pick-a-Pic V2 ImageReward (0.0750 vs 0.4953) — the paper attributes this to forward-KL’s mean-seeking diluting the preference signal, not to insufficient data [Table 1].
Image-editing transfer: applying SDEdit with the DMPO-fine-tuned model improves text-guided editing on TEd-bench and InstructPix2Pix bench over the base and over Diffusion-DPO; the authors attribute this to better prompt-following from the sharper alignment [§5.3, Figure 4].
Sensitivity ablations: too small a β collapses the KL regularizer and degrades performance; too large makes the objective overly restrictive. The smoothing coefficient α must be small but non-zero — at very small α the bound from Jensen becomes looser, and the model can occasionally favor the loser sample [§5.4, Figure 3a/b].
The reverse-KL choice is taken from autoregressive-LM alignment work (EXO, f-PO) — DMPO is the first to port it to diffusion’s Markov-chain factorization, requiring the convexity-interval argument absent in the LM derivation [§4, Related Work].

Method

The setup is standard RLHF for diffusion: a reward r(x|c) learned from human preference pairs, regularized by KL to a reference policy π_ref, with the analytical optimal π*(x|c) ∝ π_ref(x|c) exp(r(x|c)/β). Diffusion-DPO substitutes the closed-form r(x|c) = β log[π*(x|c)/π_ref(x|c)] + β log Z(c) into the Bradley–Terry preference likelihood, giving a tractable pairwise loss — and, in the multi-sample limit, minimizing forward KL(π* ‖ π_θ) over sampling paths.

DMPO swaps the objective to reverse KL(π_θ ‖ π_y*) between the learned policy and an unnormalized preference-optimal distribution π_y* over {x_w, x_l}. The derivation has three moves:

Express reverse KL in λ-space. Let λ_θ(x|c) = log[π_θ(x|c)/π_ref(x|c)]. The reverse-KL objective becomes an expectation, over winner / loser pairs, of a scalar function f(λ_θ(x_w) − λ_θ(x_l), α) where α is a smoothing coefficient (probability of the loser).
Apply Jensen across the diffusion trajectory. Because the marginal likelihoods are intractable, the per-step DMPO loss is an upper bound on the trajectory-level objective via Jensen’s inequality — which requires f''(z) > 0. The paper shows the convexity interval includes the practical optimization trajectory: optimization starts near z = 0 (model near reference) and increases monotonically toward the preference signal.
Approximate the reverse posterior with the forward. Standard diffusion trick: replace p_θ(x_{t−1} | x_t) with q(x_{t−1} | x_t, x_0) to get a per-step regression loss in λ_θ. The resulting objective is structurally identical to Diffusion-DPO — same forward passes, same KL regularizer, same hyperparameter β — but the loss function around the log-ratio is the reverse-KL surrogate rather than the BT likelihood.

Training uses Pick-a-Pic V2 (851,293 pairs after dropping ties, 58,960 prompts), 4× A100 with global batch size 2048 via gradient accumulation, learning rate 1e-8, 500-step linear warmup. SD1.5 uses β = 5000, α = 0.005; SDXL uses β = 2000, α = 0.005. Evaluation is reward-model scores + win-rates + a 100-prompt user study covering general preference, visual appeal, and prompt alignment.

Results

Reward-model scores on Pick-a-Pic V2 (SD1.5 base): DMPO beats every baseline on every metric — PickScore 0.2165 vs 0.2120 (DSPO) / 0.2120 (KTO) / 0.2106 (DPO) / 0.2066 (SD1.5); ImageReward 0.5415 vs 0.5341 (KTO) / 0.4683 (DSPO) / 0.0750 (DPO) / -0.1478 (SD1.5) [Table 1].
HPS V2 test set: DMPO ImageReward 0.6350 vs SFT 0.5607 / KTO 0.5652 / DSPO 0.5385 / DPO 0.1271 [Table 1].
PartiPrompt: DMPO ImageReward 0.6614 vs SFT 0.5293 / KTO 0.5551 / DSPO 0.5021 / DPO 0.2243 [Table 1].
Win rate vs SD1.5 (PickScore): DMPO 82.40% (Pick-a-Pic V2) / 85.50% (HPS V2) / 76.72% (PartiPrompt) — uniformly higher than DPO, KTO, DSPO, SFT [Table 2].
Pairwise win rate vs DPO: DMPO wins on every metric on every test set, e.g. PickScore 73.40% on Pick-a-Pic V2, 74.00% on HPS V2, 69.30% on PartiPrompt [Table 2].
Pairwise win rate vs KTO / DSPO (the closest competitors): DMPO wins on PickScore (≈65–69%) on all three test sets but is statistically a coin flip on HPS / Aesthetics / ImageReward (47–55%) — i.e. DMPO is robustly better on the training-set-aligned preference metric and roughly matches on the others [Table 2].
User study: human evaluators consistently prefer DMPO over SD1.5 and over all baselines on general preference / visual appeal / prompt alignment (100 prompts, details in Appendix C) [§5.2, Figure 2].
Image editing (SDEdit pipeline): DMPO improves over base and DPO on TEd-bench and InstructPix2Pix bench — qualitative figures show DMPO correctly renders “pet entrance” and “half-eaten pizza” where baselines fail [§5.3, Figure 4].
Ablations: PickScore is unimodal in β (peaks at the chosen 5000 for SD1.5); too-small α makes the optimization occasionally invert the preference pair [§5.4, Figure 3].
No additional cost: per-step compute and memory are identical to Diffusion-DPO; the headline gains come from the loss-function change alone [§3.2 closing remark, §5.1 training details].

Why it’s interesting

DMPO is the diffusion analog of the “swap forward KL for reverse KL” move that’s been quietly happening on the LM side (EXO, f-PO), and it provides the missing piece — the convexity-interval argument that justifies Jensen on the diffusion trajectory — without which the LM derivation does not transfer. The empirical headline is interesting on its own (a free win over Diffusion-DPO at identical cost), but the more useful framing for Luma post-training is that this paper diagnoses what Diffusion-DPO is actually optimizing (Theorem 1) — the existing default is mean-seeking, which explains the “diffuse generations” failure mode every team using Diffusion-DPO on a video model has observed. The result lands directly in the design space DanceGRPO: Unleashing GRPO on Visual Generation and Finite Difference Flow Optimization (FDFO): RL Post-Training for Flow-Based Image Generators are exploring from different angles: DanceGRPO uses GRPO with shaped rewards and no KL term, FDFO avoids the denoising-MDP framing entirely with a flow-native finite-difference update, and DMPO stays within the DPO formulation but fixes the divergence direction. All three are post-training recipes for visual generators that diverge from vanilla Diffusion-DPO; head-to-head on the same base model is now a real open question. Contrast with Stabilizing Reinforcement Learning for Diffusion Language Models which addresses RL stability for diffusion language models — DMPO addresses preference alignment for diffusion image models with a different objective family, but both are pushing the same “DPO is not enough” thread.