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

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.

• The paper presents SOAR, a post-training strategy that corrects trajectory bias in diffusion models to enhance compositional accuracy.
• It utilizes a single ODE rollout step to generate and correct off-trajectory states with dense analytic supervision, outperforming standard SFT and RL methods.
• Empirical results on SD3.5-Medium show improved semantic consistency, higher aesthetic metrics, and overall data efficiency compared to previous approaches.

SOAR: Self-Correction for Trajectory Bias Mitigation in Diffusion Model Post-Training

Motivation and Post-Training Paradigm Analysis

Diffusion and flow-matching models have delivered state-of-the-art performance in high-fidelity visual synthesis, yet persistent generation failures—semantic errors, compositional breakdowns, and malformed text rendering—remain endemic across sampled outputs. These failures are rooted in exposure bias, reflecting a distributional mismatch between training (ground-truth on-trajectory states) and inference (model-induced off-trajectory states), analogous to the classic problem afflicting autoregressive models but accumulated along the denoising trajectory.

Supervised fine-tuning (SFT) is inherently off-policy: the model is trained exclusively on ideal states generated from the true data distribution, never learning to recover from its own prediction errors, leading to compounded drift during inference. Reinforcement learning (RL) alignment methods, e.g., Flow-GRPO, mitigate this bias at the outcome level via sparse terminal reward signals but suffer from credit assignment difficulties, reward hacking, sample inefficiency, and lack of dense trajectory-level supervision.

Figure 1: Conceptual comparison of SFT, RL, and SOAR post-training regimes, emphasizing SOAR's early intervention at the trajectory step before irreversible semantic errors emerge.

SOAR Algorithmic Formulation

SOAR (Self-Correction for Optimal Alignment and Refinement) introduces a post-training bias-correction paradigm that synthesizes the strengths of SFT and RL while abrogating their respective limitations. The key mechanism is trajectory bias correction: SOAR directly supervises the model to recover from its own rollout-induced off-trajectory states towards the canonical clean data anchor.

Given a training pair $(z_0, c)$, SOAR applies a single ODE rollout step (stop-gradient) at noise level $\sigma_{t_0}$, generating an off-trajectory state representative of actual inference errors. This state is re-noised via interpolation with the same noise $z_1$ to auxiliary levels, maintaining proximity to the original transport ray. The analytically derived correction target, $v_{\text{corr}} = (z'_\sigma - z_0)/\sigma$, imposes endpoint consistency, steering off-trajectory states back toward $z_0$ with dense per-timestep supervision.

Figure 2: Noise-level analysis, illustrating ideal trajectories, biased rollouts, and the local correction direction from SOAR's auxiliary supervision.

This paradigm is reward-free, on-policy, and eliminates credit-assignment ambiguity, as the analytic supervision is available immediately at the error step. SOAR’s loss subsumes the flow-matching objective and directly replaces SFT as a first-stage post-training step, with compatibility for subsequent RL alignment.

Empirical Results and Quantitative Analysis

SOAR's efficacy was rigorously evaluated on SD3.5-Medium across broad and reward-curated corpora. SOAR outperforms SFT in all reported metrics, including rule-based compositional accuracy (GenEval: 0.78 vs. 0.70), OCR text fidelity (0.67 vs. 0.64), and model-based preference scores (PickScore, ClipScore, HPSv2.1, Aesthetic, ImageReward), all without recourse to a reward model.

Figure 3: Reward-specific training dynamics at $512{\times}512$, showing SOAR's monotonic improvement in target metrics across curated high-aesthetic and high-ClipScore subsets.

On reward-specific benchmarks, SOAR achieves the highest aesthetic and ClipScore values under fixed compute budgets, surpassing both SFT and explicit RL (Flow-GRPO). Notably, SOAR delivers monotonic metric improvement, whereas SFT saturates early, and Flow-GRPO risks non-target metric degradation—a direct manifestation of reward hacking. These results substantiate SOAR’s superior data efficiency and trajectory-level corrective supervision.

Domain-Generalization and Qualitative Insights

Qualitative visualization reveals that SOAR rapidly acquires structural fidelity and semantic consistency across complex domains, including high-aesthetic portraits, structured poster layouts, and precise WebUI renderings.

Figure 4: Training visualization on high-aesthetic data, demonstrating SOAR's acquisition of global structure and semantic consistency.

Figure 5: SOAR's learning trajectory on poster layouts, optimizing typographic placement and spatial hierarchy.

Figure 6: SOAR captures fine-grained stylistic and interface design details on WebUI data.

These outcomes demonstrate that dense correction at the trajectory level yields improved compositional alignment without sacrificing visual diversity or quality.

Ablation and Design Justifications

Ablation studies confirm design choices: ODE-only auxiliary rollout branches outperform SDE-augmented variants in compositional accuracy, and re-noising with shared noise preserves anchor validity and geometric correction. The small single-step rollout deviation keeps the correction target tightly coupled to the original data sample, safeguarding semantic fidelity.

Practical and Theoretical Implications

SOAR represents a formal advancement in post-training for diffusion models by tightly integrating trajectory-level supervision without relying on external preference annotation or reward models. Practically, SOAR is a data-efficient, low-compute paradigm that improves generalization and semantic consistency across broad and curated domains. Theoretically, it demonstrates that trajectory bias correction directly closes the inference-training distribution gap, with dense analytic supervision mitigating compounded error and exposure bias.

SOAR’s methodology can be generalized to other multi-step generative models where distributional drift emerges, e.g., video, 3D synthesis, or sequential distillation. Further research directions include formal diversity quantification, optimized noise-weighting schedules, and integration with sample-efficient RL pipelines.

Conclusion

SOAR introduces trajectory bias correction as a principled post-training strategy for diffusion models, supplanting SFT and outperforming RL alignment in both model-based and rule-based metrics. By delivering dense, analytic corrections at the timestep where deviation occurs, SOAR achieves superior compositional accuracy, semantic fidelity, and generalization without reward hacking or credit assignment challenges. Its implications extend across generative pipelines, highlighting the value of on-policy, reward-free correction for aligning inference distributions with training anchors (2604.12617).