# Explore or Converge? Stage-Guided Per-Step Optimization for Diffusion Models

arXiv: https://arxiv.org/abs/2608.06768 · HTML: https://mlsanity.com/paper/2608.06768/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for robotics, one of 14 directions in Diffusion models.
- Builds on 28 of its 50 references that are in the map.
- First posted on arXiv 7 Aug 2026.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)

## Builds on (most cited first, 28 of 28)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [LAION-5B: An open large-scale dataset for training next generation image-text models](https://mlsanity.com/paper/2210.08402/) (2022)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Pick-a-Pic: An Open Dataset of User Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2305.01569/) (2023)
- [GenEval: An Object-Focused Framework for Evaluating Text-to-Image Alignment](https://mlsanity.com/paper/2310.11513/) (2023)
- [Improved Precision and Recall Metric for Assessing Generative Models](https://mlsanity.com/paper/1904.06991/) (2019)
- [Perception Prioritized Training of Diffusion Models](https://mlsanity.com/paper/2204.00227/) (2022)
- [Using Human Feedback to Fine-tune Diffusion Models without Any Reward Model](https://mlsanity.com/paper/2311.13231/) (2023)
- [Inverse Reward Design](https://mlsanity.com/paper/1711.02827/) (2017)
- [Helping or Herding? Reward Model Ensembles Mitigate but do not Eliminate Reward Hacking](https://mlsanity.com/paper/2312.09244/) (2023)
- [Steering Your Diffusion Policy with Latent Space Reinforcement Learning](https://mlsanity.com/paper/2506.15799/) (2025)
- [Scaling Laws for Reward Model Overoptimization in Direct Alignment Algorithms](https://mlsanity.com/paper/2406.02900/) (2024)
- [RLHF Deciphered: A Critical Analysis of Reinforcement Learning from Human Feedback for LLMs](https://mlsanity.com/paper/2404.08555/) (2024)
- [Dense Reward for Free in Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2402.00782/) (2024)
- [AutoDiffusion: Training-Free Optimization of Time Steps and Architectures for Automated Diffusion Model Acceleration](https://mlsanity.com/paper/2309.10438/) (2023)
- [Correlated Proxies: A New Definition and Improved Mitigation for Reward Hacking](https://mlsanity.com/paper/2403.03185/) (2024)
- [Towards Understanding the Working Mechanism of Text-to-Image Diffusion Model](https://mlsanity.com/paper/2405.15330/) (2024)
- [DyMO: Training-Free Diffusion Model Alignment with Dynamic Multi-Objective Scheduling](https://mlsanity.com/paper/2412.00759/) (2024)
- [SliderSpace: Decomposing the Visual Capabilities of Diffusion Models](https://mlsanity.com/paper/2502.01639/) (2025)
- [Measuring Diversity in Co-creative Image Generation](https://mlsanity.com/paper/2403.13826/) (2024)
- [The Evolution of Reinforcement Learning in Quantitative Finance: A Survey](https://mlsanity.com/paper/2408.10932/) (2024)
- [MIRA: Towards Mitigating Reward Hacking in Inference-Time Alignment of T2I Diffusion Models](https://mlsanity.com/paper/2510.01549/) (2025)
- [The Exploration-Exploitation Dilemma Revisited: An Entropy Perspective](https://mlsanity.com/paper/2408.09974/) (2024)
- [DRAGON: Distributional Rewards Optimize Diffusion Generative Models](https://mlsanity.com/paper/2504.15217/) (2025)
- [Alignment and Safety of Diffusion Models via Reinforcement Learning and Reward Modeling: A Survey](https://mlsanity.com/paper/2505.17352/) (2025)

## Built on by (most cited first, 1 of 1)

- [ExploreNet: Learning Where to Explore in Diffusion GRPO](https://mlsanity.com/paper/2609.38329/) (2026)

## Abstract (arXiv)

Diffusion models have strong generative capabilities. However, their maximum likelihood training objective only focuses on reconstructing the data distribution, making it difficult to align with specific preferences. Reinforcement learning (RL) for preference alignment in diffusion models is promising but limited by reward sparsity. Since a single reward cannot support optimization, existing RL methods usually backpropagate the final reward to all previous steps. However, denoising is stage-wise, with distinct semantics and controllability. Repeating the final reward across all steps creates a temporal objective mismatch, encouraging reward shortcuts that lead to reward hacking. At the same time, due to reward backfilling, each time step receives the same reward, making it impossible to distinguish between actions, thereby weakening the optimization process. To resolve this issue, we propose Stage-Guided Per-Step Optimization (SGPO) for diffusion models, which jointly leverages signal-to-noise ratio and semantic changes to identify generation stages and adaptively assign stage-specific objectives. Early denoising is chaotic and far from the final reward, resulting in weak reward-behavior correlation. This stage should prioritize exiting the chaotic state. In the mid stage, the latent transitions to a stable structure, where the final reward better corresponds to generative behavior. Therefore, this stage optimizes the final reward while exploring diversity to avoid early convergence to a single mode. In the late stage, the latent's core structure is largely fixed, and preference optimization mainly amplifies local details, risking overfitting. Therefore, stable convergence is preferred to avoid quality degradation. Results from 16 comparative experiments validate SGPO. Our method achieves 26.7% average gains in generative quality and 36.7% higher convergence speed.