💻 computer science

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

This paper proposes Stage-Guided Per-Step Optimization (SGPO), a novel method that adaptively assigns stage-specific objectives to diffusion models based on signal-to-noise ratios and semantic changes to overcome reward sparsity and temporal mismatch issues in reinforcement learning, thereby significantly improving generative quality and convergence speed.

Original authors: Renye Yan, Jikang Cheng, You Wu, Wei Peng, Zongwei Wang, Ling Liang, Yimao Cai

Published 2026-08-10

📖 6 min read🧠 Deep dive

CC0 1.0

Original authors: Renye Yan, Jikang Cheng, You Wu, Wei Peng, Zongwei Wang, Ling Liang, Yimao Cai

Original paper dedicated to the public domain under CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot artist to paint. You don't just want it to copy a photo perfectly; you want it to create something beautiful that matches a specific description, like "a happy dog wearing a hat." This is the world of Diffusion Models, a type of artificial intelligence that creates images by starting with a chaotic mess of static (like TV snow) and slowly cleaning it up, step-by-step, until a clear picture emerges.

To make these robots follow your specific wishes, scientists use a technique called Reinforcement Learning. Think of this as a game where the robot gets a "reward" (like a gold star) only when it finishes the whole painting and the result looks good. The problem is that the robot only gets that gold star at the very end. It doesn't know if its brushstrokes in the middle were good or bad. So, the robot tries to guess which moves led to the star and applies that same "good job" signal to every single step of the painting process, from the first messy scratch to the final detail.

But here's the catch: painting isn't the same at every moment. The beginning is pure chaos, the middle is where the shape forms, and the end is just adding tiny details. Treating the messy beginning exactly the same as the final polish confuses the robot. It starts taking shortcuts, painting weird patterns that trick the scoring system but look terrible to humans. This paper, titled "Explore or Converge? Stage-Guided Per-Step Optimization for Diffusion Models," by researchers from Peking University, Nanjing University, and Stanford, proposes a smarter way to teach the robot. Instead of giving the same reward signal for the whole journey, they suggest changing the rules depending on which "stage" of painting the robot is in.

The Problem: One Size Does Not Fit All

The authors argue that the old way of training these models is like trying to teach a child to ride a bike by yelling "Pedal harder!" at the exact same volume whether the child is wobbling on training wheels, cruising down a hill, or trying to balance on one foot.

In the early days of diffusion, the image is just noise. The robot is essentially guessing in the dark. If you tell it "You're doing great!" based on the final image, it gets confused because the current messy pixels have almost nothing to do with the final result. This leads to reward hacking, where the robot learns to game the system. It might start repeating the same safe, boring patterns or amplifying tiny details just to get a high score, even if the picture looks weird or broken. It's like a student who memorizes the answer key instead of learning the math.

The Solution: A Three-Act Play

The researchers, led by Renye Yan and Jikang Cheng, realized that the generation process naturally splits into three distinct acts, and each act needs a different teacher. They call their new method SGPO (Stage-Guided Per-step Optimization).

Act 1: The Chaotic Stage (Get Out of the Noise)
At the very beginning, the image is just static. The robot is in a "chaotic state." The authors found that trying to optimize for the final goal here is useless because the connection between the current mess and the final picture is too weak.

The SGPO Strategy: Instead of worrying about the final reward, the robot is told to simply "get out of the noise." The goal is to move away from the initial static as fast as possible. It's like telling a swimmer in a stormy ocean to just stop flailing and find solid ground before worrying about swimming style.

Act 2: The Stable Exploration Stage (Find the Shape and Try New Things)
Once the noise clears up, the main shapes and structures of the image start to appear. This is the "sweet spot." The robot can now see what it's making, and the final reward actually makes sense.

The SGPO Strategy: Here, the robot is encouraged to chase the final reward (make the picture look good) but is also told to explore. The authors add a special rule that encourages the robot to try different variations of the image. This prevents the robot from getting stuck in a rut or copying the same boring pattern over and over. It's like a chef who has the main ingredients ready and is now encouraged to experiment with spices to make the dish unique, rather than just sticking to one safe recipe.

Act 3: The Convergence Stage (Polish and Stop)
In the final moments, the image is almost done. The big shapes are fixed, and only tiny details are changing. If the robot keeps trying to "optimize" the reward here, it might start overthinking, tweaking pixels until the image looks weird or loses its natural feel. This is where overfitting happens.

The SGPO Strategy: The robot is told to converge. It should stop exploring and just settle down, making sure the final details are stable and true to the original plan. It's like a sculptor who has finished the statue and is now just smoothing the surface, not trying to change the shape of the nose.

How They Knew It Worked

The researchers didn't just guess; they built a system that watches the robot in real-time. They measure two things: how much "noise" is left (Signal-to-Noise Ratio) and how much the meaning of the image is changing (Semantic Change). When these numbers hit specific points, the system automatically switches the robot's training rules from "Escape Chaos" to "Explore" to "Converge."

To test if this actually helps, they ran 16 comparative experiments. They compared their method against other popular ways of training AI, like DDPO and DPOK.

The results were quite impressive. The SGPO method didn't just make better pictures; it made them faster. The authors report that their method achieved 26.7% average gains in generative quality and was 36.7% faster at reaching a good result.

They also tested for "reward hacking." In many other methods, when the AI tries to maximize a score for "aesthetic beauty," it often ends up creating images that look like a glitchy mess of repeating patterns. The authors showed that with SGPO, the AI maintained high scores for beauty while keeping the images diverse and natural. In fact, they measured a 41.93% improvement in detail fidelity, meaning the tiny textures and edges looked much more realistic and less "hacked."

The Takeaway

This paper suggests that the secret to teaching AI to create art isn't just giving it a better goal, but understanding the journey it takes to get there. By recognizing that the beginning, middle, and end of image creation are fundamentally different, and by changing the training rules to match, the robot learns to create better, more diverse, and more reliable images. It's a reminder that sometimes, to get the best result, you have to stop treating every step of the process as if it's the same.

Technical Summary: Stage-Guided Per-Step Optimization for Diffusion Models

Problem Statement

Diffusion models excel at generative tasks but struggle to align with specific user preferences because their standard training objective focuses on reconstructing data distributions rather than optimizing for specific goals. While Reinforcement Learning (RL) offers a pathway for preference alignment, it faces significant challenges in this domain:

Reward Sparsity: Preference rewards can only be accurately assessed upon the completion of the generation process (the final timestep), leaving intermediate latent states without direct feedback.
Temporal Objective Mismatch: Existing methods typically backpropagate the single final reward to all previous denoising steps to mitigate sparsity. However, this assumes all denoising stages are optimally equivalent, which contradicts the stage-wise nature of diffusion dynamics.
Reward Hacking: Applying a uniform final reward across all steps creates a temporal mismatch. Early stages are chaotic and noise-dominated, making the correlation between actions and the final reward weak. This encourages the model to find "shortcuts" (reward hacking) rather than learning meaningful behaviors, leading to mode collapse, repetitive structures, or over-amplified local patterns that degrade diversity and fidelity.

Methodology: Stage-Guided Per-Step Optimization (SGPO)

The authors propose SGPO, an adaptive RL framework that recognizes the intrinsic stage-wise dynamics of diffusion generation. Instead of uniform optimization, SGPO dynamically switches optimization objectives based on the current generation stage, identified by real-time monitoring of semantic changes and signal-to-noise ratio (SNR).

1. Adaptive Stage Detection

The method divides the generation process into three distinct stages using two metrics:

Semantic Evolution Rate ( $\Delta E(t)$ ): The change in semantic alignment (e.g., CLIP score) between adjacent latents and the prompt.
Signal-to-Noise Ratio ( $\gamma(t)$ ): The ratio of signal to noise in the latent state.

The transition points are defined by the second-order derivatives of these metrics approaching zero:

Stage I (Chaotic/Noisy): $0 \le t < t_c^{(1)}$ . Characterized by high noise and unstable semantics.
Stage II (Stable Preference): $t_c^{(1)} \le t < t_c^{(2)}$ . Semantics stabilize, and the correlation between actions and the final reward strengthens.
Stage III (Steady Convergence): $t \ge t_c^{(2)}$ . Global structure is fixed; only local details remain sensitive.

2. Stage-Specific Optimization Objectives

SGPO assigns distinct reward functions to each stage to align with the generative dynamics:

Stage I (Chaotic): The goal is to exit the chaotic noise regime quickly. Since the final reward is unreliable here, the objective focuses on maximizing the distance from the initial noise origin.
- Reward Function: $r_t^{(I)} = \lambda_t \|x_t - x_T\|^2$ , where $\lambda_t$ scales with SNR.
Stage II (Stable Preference): This is the primary phase for preference optimization. The model optimizes the final preference reward ( $r(x_0, z)$ ) while actively maintaining diversity to prevent early convergence to a single mode.
- Reward Function: $r_t^{(II)} = r(x_0, z) + \Delta E(t) \| \hat{x}_0(x_t) - \hat{x}_0(x_{t-1}) \|^2$ .
- The second term acts as an exploration bonus, modulated by the semantic evolution rate $\Delta E(t)$ , encouraging controlled deviations that enhance diversity without breaking semantic alignment.
Stage III (Convergence): The global structure is largely fixed. Further preference optimization risks overfitting and destabilizing the distribution. The goal shifts to stable convergence toward the true distribution.
- Reward Function: $r_t^{(III)} = -\frac{1}{2} \| \hat{x}_0(x_t) - x_0^{pre} \|^2$ , where $x_0^{pre}$ is the output of the pre-trained model. This penalizes unnecessary drift from the stable pre-trained distribution.

Key Contributions

Stage-Aware RL Framework: The paper introduces a paradigm that explicitly models the stage-wise generative dynamics of diffusion models, departing from the assumption of uniform optimization across all timesteps.
Mitigation of Reward Hacking: By aligning training signals with phase-dependent characteristics (chaos exit, diversity exploration, and stable convergence), SGPO effectively mitigates reward hacking and improves generation quality.
Adaptive Mechanism: The method utilizes real-time semantic and SNR monitoring to dynamically switch optimization strategies, ensuring the training signal matches the requirements of different generation phases.

Experimental Results

The authors conducted 16 comparative experiments across various datasets (HPSv2, Pick-a-Pic, GenEval, Simple Animals) and model architectures (Stable Diffusion v1.4–v2.1, SDXL, SD 3.5, FLUX).

Performance Gains: SGPO achieved an average 26.7% gain in generative quality (measured by reward scores) and a 36.7% improvement in convergence speed compared to baselines.
Diversity and Fidelity: The method improved diversity by 8.51% and detail fidelity by 41.93%, significantly outperforming baselines in metrics like LPIPS, IS, and FID.
Reward Hacking Resistance: In evaluations where models were trained to match a fixed Aesthetic Score (AES), SGPO consistently maintained top-two performance across fidelity, diversity, and richness metrics, whereas other methods exhibited signs of over-optimization and mode collapse.
Controlled Generation: SGPO demonstrated superior instruction-following capabilities in complex prompts involving object counts, colors, and spatial relationships compared to baselines like DDPO, DPOK, and Flow-GRPO.
Efficiency: Despite the overhead of stage detection, the overall wall-clock time efficiency was significantly superior to baselines due to faster convergence.

Significance and Claims

The paper claims that the indiscriminate application of final rewards across all denoising steps is a fundamental cause of reward hacking and optimization inefficiency in diffusion RL. By respecting the stage-wise structure of diffusion processes—specifically the transition from chaotic noise to stable semantics and finally to convergence—SGPO enables more robust and efficient fine-tuning.

The authors assert that their method provides a necessary correction to current RL practices for diffusion models, allowing for:

Better alignment with specific preferences without sacrificing diversity.
Prevention of the "reward shortcut" phenomenon where models exploit spurious correlations.
Enhanced controllability in complex generation tasks.

The work suggests that future RL fine-tuning for generative models must account for the temporal dynamics of the generation process rather than treating the trajectory as a homogeneous sequence of states.