SOAR: Self-Correction for Optimal Alignment and Refinement in Diffusion Models

You Qin; Linqing Wang; Hao Fei; Roger Zimmermann; Liefeng Bo; Qinglin Lu; Chunyu Wang

cs.LGApr 14, 2026

SOAR: Self-Correction for Optimal Alignment and Refinement in Diffusion Models

Authors: You Qin, Linqing Wang, Hao Fei, Roger Zimmermann, Liefeng Bo, Qinglin Lu, Chunyu Wang

Abstract

The post-training pipeline for diffusion models currently has two stages: supervised fine-tuning (SFT) on curated data and reinforcement learning (RL) with reward models. A fundamental gap separates them. SFT optimizes the denoiser only on ground-truth states sampled from the forward noising process; once inference deviates from these ideal states, subsequent denoising relies on out-of-distribution generalization rather than learned correction, exhibiting the same exposure bias that afflicts autoregressive models, but accumulated along the denoising trajectory instead of the token sequence. RL can in principle address this mismatch, yet its terminal reward signal is sparse, suffers from credit-assignment difficulty, and risks reward hacking. We propose SOAR (Self-Correction for Optimal Alignment and Refinement), a bias-correction post-training method that fills this gap. Starting from a real sample, SOAR performs a single stop-gradient rollout with the current model, re-noises the resulting off-trajectory state, and supervises the model to steer back toward the original clean target. The method is on-policy, reward-free, and provides dense per-timestep supervision with no credit-assignment problem. On SD3.5-Medium, SOAR improves GenEval from 0.70 to 0.78 and OCR from 0.64 to 0.67 over SFT, while simultaneously raising all model-based preference scores. In controlled reward-specific experiments, SOAR surpasses Flow-GRPO in final metric value on both aesthetic and text-image alignment tasks, despite having no access to a reward model. Since SOAR's base loss subsumes the standard SFT objective, it can directly replace SFT as a stronger first post-training stage after pretraining, while remaining fully compatible with subsequent RL alignment.

Figures & tables

Figure 1: Conceptual comparison of post-training regimes from a trajectory perspective. Left: SFT trains only on ground-truth denoising states. Once an inference rollout exits the narrow corridor of well-optimized states, subsequent steps sharpen the wrong structure. Middle: RL separates good and bad final outcomes with a terminal reward, but the sparse signal arrives after the erroneous rollout has already formed. Right: SOAR performs a one-step rollout to generate off-trajectory states and supervises the model to correct back toward the clean target, acting before semantic errors fully materialize.

	Training Signal	On-policy	Signal Density	Reward Model	Credit Assign.	Reward Hacking	Resolves Model Bias	Training Cost
SFT	Supervised (GT traj.)	✗	Dense	✗	✗	✗	✗	Low
DPO	Pref. pairs (implicit reward)	✗	Sparse	✗	✗	✗	Partial	Medium
RL	Reward function	✓	Sparse	✓	✓	✓	Partial	High
SOAR	Supervised (traj. correction)	✓ (single-step)	Dense	✗	✗	✗	✓	Low

Table 1: Qualitative comparison of post-training paradigms for diffusion models. SOAR combines on-policy training with dense, reward-free supervision while maintaining low training cost.

Figure 2: Noise-level analysis illustrating the ideal trajectory, the biased rollout, and the local correction direction. The off-trajectory state is constructed by one ODE step with the current model and then re-noised to auxiliary levels.

Model	#Iter	Rule-Based (Flow-GRPO test sets)		Model-Based (OOD DrawBench)
Model	#Iter	GenEval	OCR	PickScore	ClipScore	HPSv2.1	Aesthetic	ImgRwd
SD-XL ‡	–	0.55	0.14	22.42	0.287	0.280	5.60	0.76
SD3.5-L ‡	–	0.71	0.68	22.91	0.289	0.288	5.50	0.96
FLUX.1-Dev †	–	0.66	0.59	22.84	0.295	0.274	5.71	0.96
SD3.5-M	–	0.63	0.59	22.34	0.285	0.279	5.36	0.85
+ SFT	10k	0.70	0.64	22.71	0.295	0.284	5.35	1.04

Table 2: Main quantitative comparison at

512{\times}512

(

\mathrm{cfg}{=}4.5

). Rule-based metrics are evaluated on the Flow-GRPO test sets; model-based metrics are evaluated on OOD DrawBench. † Official checkpoint. ‡ At

1024{\times}1024

. Bold : best overall; underline : second best.

Figure 3: Reward-specific training dynamics at

512{\times}512

. (a) OOD DrawBench Aesthetic score after training on the 3,725-pair High-Aesthetic subset (

\geq 6.8

). (b) OOD DrawBench ClipScore after training on the 6,857-pair High-ClipScore subset. All methods start from the same SD3.5-M pretrained checkpoint (step 0). SOAR achieves the highest final values in both settings with monotonic improvement.

Variant	Rule-Based (Flow-GRPO test sets)		Model-Based (OOD DrawBench)
Variant	GenEval	OCR	PickScore	ClipScore	HPSv2.1	Aesthetic	ImgRwd
SOAR (ODE-only, shared $z_{1}$ )	0.78	0.67	22.86	0.295	0.289	5.46	1.09
$+$ SDE branch	0.78	0.65	22.87	0.296	0.287	5.47	1.07
$+$ random $z_{1}$ per path	0.76	0.64	22.89	0.292	0.287	5.45	1.08

Table 3: Ablation of SOAR design choices on SD3.5-M (

512\times 512

, cfg

=

4.5, 10k steps). Rule-based metrics use the Flow-GRPO test sets; model-based metrics use OOD DrawBench. Best per column in bold .

Figure 4: Training visualization on the high-aesthetic subset, covering portraits, landscapes, and anime styles. SOAR rapidly acquires coherent composition and structural fidelity as training progresses.

Figure 5: Training visualization on poster data. SOAR learns layout structure, typographic placement, and visual hierarchy with high fidelity.

Figure 6: Training visualization on WebUI data. SOAR accurately captures interface elements, shadow effects, and fine-grained stylistic details.

Explore similar work

On-Policy Self-Distillation in Diffusion Models
Aug 25, 2026Wei Zhou, Xiongwei Zhu, Lingdong Kong +14Diffusion Model Alignment On-Policy Self-Distillation
Correction Training for Diffusion Models
Oct 9, 2026Donghoon Ahn, Runqian Wang, Alexei A. Efros +2Diffusion-Based Image Generation Diffusion Models
University of California, Berkeley · Reve
Explicit Critic Guidance for Aligning Diffusion Models
May 26, 2026Zhengyang Liang, Qihang Zhang, Ceyuan YangDiffusion Model Alignment Diffusion Model Guidance
University of Toronto · Vector Institute · The Chinese University of Hong Kong

Aug 25, 2026J/K move · Enter open · S save

On-Policy Self-Distillation in Diffusion Models

Wei Zhou, Xiongwei Zhu, Lingdong Kong, Bo Chen+13

Diffusion Model Alignment On-Policy Self-Distillation Diffusion Model Fine-Tuning On-Policy Distillation RL for Diffusion Models

Open paper AskPDF arXiv

Reinforcement learning can align diffusion models with human preferences and task-specific objectives, but endpoint rewards do not specify how an intermediate denoising prediction should change. We introduce DiffusionOPSD as an on-policy self-distillation framework that converts image-level reward guidance into explicit targets for clean-output predictions at sampled queries. At each outer iteration, a frozen behavior policy generates trajectories and supplies query states and anchors. Reward gradients construct bounded positive and negative targets around each anchor. The trainable policy fits these targets as detached supervision through finite fitting before an exponential moving average update refreshes the behavior policy. This setup lets us measure target construction and finite realization separately. Controlled same-query experiments show that larger target-construction gains do not necessarily translate into larger realized gains after a single fitting update. Across SD 3.5-M and the step-distilled Z-Image-Turbo, our approach achieves the best final held-out scores in 19 of 20 reward-matched settings across two backbones and ten evaluators. It outperforms the strongest competing method by up to 44.0% and reduces training GPU-hours relative to DiffusionNFT by 40% on SD 3.5-M and 63% on Z-Image-Turbo. These results support on-policy self-distillation as an efficient and analyzable approach to diffusion post-training by converting image-level reward guidance into explicit and continually refreshed intermediate supervision, thereby opening a path toward more efficient and diagnosable alignment.

Abstract

Explore similar work

On-Policy Self-Distillation in Diffusion Models

Correction Training for Diffusion Models

Explicit Critic Guidance for Aligning Diffusion Models