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

• URL: http://arxiv.org/abs/2608.06768v1
• Date: Fri, 07 Aug 2026 03:41:58 GMT
• Title: Explore or Converge? Stage-Guided Per-Step Optimization for Diffusion Models
• Abstract summary: Reinforcement learning for preference alignment in diffusion models is promising but limited by reward sparsity.<n>We propose Stage-Guided Per-Step Optimization (SGPO) for diffusion models.<n>Our method achieves 26.7% average gains in generative quality and 36.7% higher convergence speed.
• Score: 11.180362540642689
• License: http://creativecommons.org/publicdomain/zero/1.0/
• Abstract: 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.

Related papers

• PAST: Prompt-Adaptive Sampling Termination for Efficient Diffusion Model [11.180362540642689]
  PAST enhances computational efficiency of existing RL fine-tuning methods by up to 66.7%.<n>It also improves preference optimization quality by up to 29.5% through its dual adaptive regulation mechanism.
  arXiv  Detail & Related papers  (2026-08-07T04:35:26Z)
• Beyond Uniform Forgetting: A Study of Sequential Direct Preference Optimization Across Preference Settings [11.551054698858266]
  We study sequential Direct Preference optimisation across four preference settings.<n>We find that sequential DPO does not produce a single forgetting pattern.<n>Mechanistic diagnostics show that Stage2 gradients and adapter updates are near-orthogonal to the previous objective across all settings.
  arXiv  Detail & Related papers  (2026-06-18T03:20:41Z)
• Do Less, Achieve More: Do We Need Every-Step Optimization for RL Fine-tuning of Diffusion Models? [30.438505697106496]
  Most studies apply RL to the full denoising trajectory, making it computationally costly and weakening preference alignment.<n>We propose AdaScope, an RL-enhanced plug-in that improves generation quality while reducing computational cost.<n>Compared with state-of-the-art methods, AdaScope improves performance by 66% while cutting computational cost by 59%.
  arXiv  Detail & Related papers  (2026-05-15T11:14:13Z)
• TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment [52.570581883709345]
  We propose Trajectory Matching Policy Balance Optimization (TMPO), which replaces reward with human-level reward distribution matching.<n>TMPO improves generative diversity over state-of-the-art methods by 9.1%, and achieves competitive performance in all downstream and efficiency metrics.<n>To further reduce multi-trajectory training time on large-scale flow- prefixes, TMPO incorporates Dynamic Tree Sampling models, where trajectories share denoising and branch at dynamically scheduled steps.
  arXiv  Detail & Related papers  (2026-05-09T04:41:02Z)
• Anchoring Values in Temporal and Group Dimensions for Flow Matching Model Alignment [61.80228667422234]
  VGPO redefines value estimation across both temporal and group dimensions.<n>It transforms the sparse terminal reward into dense, process-aware value estimates.<n>It replaces standard group normalization with a novel process enhanced by absolute values to maintain a stable optimization signal.
  arXiv  Detail & Related papers  (2025-12-13T16:31:26Z)
• SynCast: Synergizing Contradictions in Precipitation Nowcasting via Diffusion Sequential Preference Optimization [62.958457694151384]
  We introduce preference optimization into precipitation nowcasting for the first time, motivated by the success of reinforcement learning from human feedback in large language models.<n>In the first stage, the framework focuses on reducing FAR, training the model to effectively suppress false alarms.
  arXiv  Detail & Related papers  (2025-10-22T16:11:22Z)
• Discrete Diffusion Trajectory Alignment via Stepwise Decomposition [80.66701883088935]
  We propose an offline preference optimization method to approach trajectory alignment for discrete diffusion models.<n>We decompose the problem into a set of stepwise alignment objectives by matching the per-step posterior.<n> Experiments across multiple domains including DNA sequence design, protein inverse folding, and language modeling consistently demonstrate the superiority of our approach.
  arXiv  Detail & Related papers  (2025-07-07T09:52:56Z)
• Reward-Guided Iterative Refinement in Diffusion Models at Test-Time with Applications to Protein and DNA Design [87.58981407469977]
  We propose a novel framework for inference-time reward optimization with diffusion models inspired by evolutionary algorithms.<n>Our approach employs an iterative refinement process consisting of two steps in each iteration: noising and reward-guided denoising.
  arXiv  Detail & Related papers  (2025-02-20T17:48:45Z)
• Aligning Few-Step Diffusion Models with Dense Reward Difference Learning [81.85515625591884]
  Stepwise Diffusion Policy Optimization (SDPO) is an alignment method tailored for few-step diffusion models. SDPO incorporates dense reward feedback at every intermediate step to ensure consistent alignment across all denoising steps. SDPO consistently outperforms prior methods in reward-based alignment across diverse step configurations.
  arXiv  Detail & Related papers  (2024-11-18T16:57:41Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.