[Paper Review] Gradient Guidance for Diffusion Models: An Optimization Perspective

Yingqing Guo, Hui Yuan+3 more|arXiv (Cornell University)|Apr 23, 2024

Spacecraft Dynamics and Control4 citations

TL;DR

This paper proposes a novel gradient guidance method for diffusion models that frames guided sampling as a regularized optimization problem, preserving latent data structure through a forward prediction loss. It proves convergence to global optima under adaptive fine-tuning, achieving $Ó(1/K)$ rate for concave objectives while maintaining sample quality.

ABSTRACT

Diffusion models have demonstrated empirical successes in various applications and can be adapted to task-specific needs via guidance. This paper studies a form of gradient guidance for adapting a pre-trained diffusion model towards optimizing user-specified objectives. We establish a mathematical framework for guided diffusion to systematically study its optimization theory and algorithmic design. Our theoretical analysis spots a strong link between guided diffusion models and optimization: gradient-guided diffusion models are essentially sampling solutions to a regularized optimization problem, where the regularization is imposed by the pre-training data. As for guidance design, directly bringing in the gradient of an external objective function as guidance would jeopardize the structure in generated samples. We investigate a modified form of gradient guidance based on a forward prediction loss, which leverages the information in pre-trained score functions and provably preserves the latent structure. We further consider an iteratively fine-tuned version of gradient-guided diffusion where guidance and score network are both updated with newly generated samples. This process mimics a first-order optimization iteration in expectation, for which we proved O(1/K) convergence rate to the global optimum when the objective function is concave. Our code will be released at https://github.com/yukang123/GGDMOptim.git.

Motivation & Objective

• To address the lack of theoretical understanding in gradient-guided diffusion models, especially regarding structure preservation and convergence guarantees.
• To resolve the issue of naive gradient guidance degrading sample quality by introducing a structured guidance signal based on a lookahead loss.
• To formalize guided diffusion as an optimization process that balances objective maximization and data distribution priors from pre-training.
• To develop an adaptive fine-tuning mechanism that iteratively improves the score network using self-generated samples.
• To establish convergence rates and theoretical guarantees for guided diffusion in optimizing concave objectives.

Proposed method

• Introduces a gradient guidance signal based on a forward prediction loss, defined as $\texttt{G}_{\text{loss}}(x_t,t) := \beta(t)\left(y - g^\top\mathbb{E}[x_0|x_t]\right)g$, which leverages the pre-trained score function to maintain latent structure.
• Uses a look-ahead estimator $\hat{\mathbb{E}}[x_0|x_t]$ based on Tweedie's formula to approximate the conditional expectation in practice.
• Proposes Algorithm 1: a non-adaptive guided sampling process that solves a regularized optimization problem, where the pre-trained score acts as a prior.
• Proposes Algorithm 2: an adaptive fine-tuning loop that queries gradients at generated samples and updates the score network using new data.
• Employs a score-based sampling process where the guidance signal is injected at each time step to steer the diffusion trajectory toward high-objective regions.
• Theoretical analysis shows that Algorithm 2 converges to the global optimum at rate $\mathcal{O}(1/K)$ when the objective is concave and differentiable.

Experimental results

Research questions

• RQ1 Why does naive gradient guidance fail in practice, and what causes degradation in sample quality?
• RQ2 How can gradient guidance be designed to preserve the latent low-dimensional structure of pre-trained data?
• RQ3 Can guided diffusion models be theoretically guaranteed to converge to optimal solutions under adaptive fine-tuning?
• RQ4 What is the relationship between gradient-guided diffusion and first-order optimization methods?
• RQ5 What is the convergence rate of the guided sampling process when the objective function is concave?

Key findings

• The proposed gradient guidance based on the forward prediction loss ($\texttt{G}_{\text{loss}}$) preserves latent subspace structure significantly better than naive gradient guidance, as shown by lower off-support errors in experiments.
• Algorithm 1, which uses non-adaptive guidance, converges to a sub-optimal solution due to the regularization effect of the pre-trained score function, consistent with theoretical analysis.
• Algorithm 2, which performs iterative fine-tuning using self-generated samples, achieves convergence to the global optimum for concave objectives, with a provable $\mathcal{O}(1/K)$ convergence rate.
• Experiments show that even under distribution shift from fine-tuning, $\texttt{G}_{\text{loss}}$ maintains better structural fidelity than naive gradient guidance.
• The distribution of samples generated by Algorithm 2 moves outside the original data distribution (as captured by the pre-training contour), indicating effective exploration of high-objective regions.
• The theoretical framework establishes that gradient-guided diffusion is equivalent to solving a regularized optimization problem, where the pre-trained score function acts as a structural prior.