Title: Gradient Guidance for Diffusion Models: An Optimization Perspective
Venue: NeurIPS 2024
Accepted: yes
Keywords: diffusion model; gradient guidance; reward optimization; convergence rate
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 $\tilde{\mathcal{O}}(1/K)$ convergence rate to the global optimum when the objective function is concave. Our code is released at https://github.com/yukang123/GGDMOptim.git.

Gradient Guidance for Diffusion Models: An Optimization Perspective

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 is released at https://github.com/yukang123/GGDMOptim.git.

1 Introduction

Diffusion models have emerged as a significant advancement in the field of generative artificial intelligence, offering state-of-the-art performance in image generation [58, 56, 22]. These models operate by gradually transforming a random noise into a structured output, utilizing the score function learned from data. One of the key advantages of diffusion models is their flexibility which allows controlled sample generation for task-specific interest, excelling diffusion models in a wide range of applications, such as content creation, sequential decision making, protein engineering [38, 1, 29, 11].

Controlling the generation of large generative models stands at the forefront of AI. Guidance and fine-tuning are two most prevalent approaches for controlling the generation of diffusion models. Unlike fine-tuning which changes the weights of pre-trained models, guidance mechanism enables a more directed and flexible control. Adding gradient-based guidance during inference was pioneered by classifier guidance [60, 22], which involves training a time-dependent classifier. Diffusion Posterior Sampling (DPS) [16] introduced a fully training-free form of gradient-based guidance, which removes the dependence on time. This method has since been explored in numerous empirical studies [17, 37, 43, 65, 57, 70, 4, 35]. However, despite these empirical successes, significant gaps

remain in the theoretical understanding and guarantees of gradient-based guidance in diffusion models.

Problem and Challenges Suppose we have a pre-trained diffusion model that can generate new samples faithfully, maintaining the latent structure of data. We study the problem of adapting this diffusion model to generate new samples that optimize task-specific objectives, while maintaining the learned structure in new samples. This problem has a strong connection to classic optimization, guided diffusion offers new possibilities to optimize complex design variables such as images, videos, proteins, and genomes [7, 66, 42] in a generative fashion. More comprehensive exposure to this middle ground can be found in recent surveys [69, 12, 30].

Given the optimization nature of this problem, it's critical to answer the following theoretical questions from an optimization perspective: (i) Why doesn't simply applying the gradient of the objective function w.r.t. the noised sample work? (ii) How to add a guidance signal to improve the target objective without compromising the sample quality? (iii) Can one guarantee the optimization properties of new samples generated by guided diffusion? (iv) What are the limits of adaptability in these guided models?

Noise

Pre-trained Diffusion

D

k

Noise

Weighted Fine-tuning

via all past generated data

s

θ

+

G

D

0

∪ · · · ∪ D

k

-

1

Gradient-Guided Diffusion for Generative Optimization

w

0

.

.

.

w

k

-

1

(optional) update

G

s

θ

(

θ

K

)

+

G

K

s

θ

Adapted Diffusion

Noise

s

θ

k

Compute Gradient Guidance

G

Weights

ˆ

E

[

x

0

|

x

t

]

∇

x

t

(

y

-

g

/latticetop

ˆ

E

[

x

0

|

x

t

])

2

Pre-trained Diffusion

Gradient Query

g

d

X

t

=

[

1

2

X

t

+

s

θ

(

X

t

, T

-

t

) +

G

(

X

t

, t

)

]

d

t

+d

W

t

Scope and Contribution. In this paper, we investigate guided diffusion from an optimization perspective. To answer the four questions above, we propose an algorithmic framework, see Figure 1 for an illustration. Our main contributions are summarized as follows:

Study structure-preserving guidance. We study the design of guidance under structural data distribution belonging to a latent low-dimensional subspace (Assumption 1). We diagnose the failure of naive gradient guidance and study the theoretical aspects of guidance based on forward prediction loss (Definition 1), which provably preserves any low-dimensional subspace structure (Theorem 1).

Establish a mathematical framework of guided diffusion. We build a mathematical framework for guided diffusion, which facilitates algorithm analysis and theory establishment. We propose and analyze an iterative guided diffusion using gradient queries on new samples (Algorithm 1; Figure 1 with fine-tuning block off). We give the first convergence theory showing generated samples converge to a regularized optimal solution w.r.t the objective (Theorem 4) with linear score class (12). The regularization is imposed by the pre-trained diffusion model, revealing a fundamental limit of adapting pre-trained diffusion models with guidance.

Provide rate-matching optimization theory. Furthermore, we propose an adaptive gradientguided diffusion, where both pre-trained score network and guidance are iteratively updated using self-generated samples (Algorithm 2; Figure 1 with the fine-tuning block turned on). We show in expectation its iteration converges to a global optima within the latent subspace, at a rate of ˜ O (1 /K ) (Theorems 3, K is # of iterations), matching the classical convergence in convex optimization.

Provide experimental justification. Simulation and image experiments are provided in Section 7 to support theoretical findings on latent structure-preserving and optimization convergence.

2 Related Works

To summarize the related work, we first give an overview of empirical studies relevant to our objective. We then discuss the theory of diffusion models, to which our main contribution is focused. Other

<latexi sh

1_b

64="AP2G

HoYK

MVT0

DS

kmNO

QvZU

>

C

3

c

L

FJ

r

d

u

XR

g

9

pn

E

B

f

+

7

/

z

w

5

8

q

j

I

W

y

<latexi sh

1_b

64="HM

LY

TUXy3E

oJ

8

q9

D

>A

C

c

V

S

N

F

rP

d

m

u

I

n

K

g

pB0

Z

f

+

/

v

7

2

W

w

G

z

R

5

Q

O

k

j

<latexi sh

1_b

64="

p

yF+KX2

39GP0TB

RnS

>A

CN

c

VDL

J

vq

g

I

fO

MZ

o

5

Y

z

EQ

/

r

8

w

d

7

m

H

<latexi sh

1_b

64="EZ

TYmX9nRyL

Nk

g

+u cW

>A

C

VD

S

FJ3U

2vq

0

w

K

r

o

d

M

j

B

f

/

7

p

H

G

z

O

P

I

5

Q

8

<latexi sh

1_b

64="D0

2uGXg

HU9

/RBQWcv

7Mq

>A

C

dVFL

N

E

wm

Ky ko

p

O

Y

j

P

5

I

8

Z

3

+

fn

J

T

Sr

z

<latexi sh

1_b

64="BoqC

zOu

VX

n

DFT3G

j

>A

H

c

LS

N

J

U

2v gk

w

ZK

y

m

p

5

Yd

W

I

/

9

7

M0

E

r

8

P

f

R

<latexi sh

1_b

64="j/K

U2

Q

Ov

X

oqD

Iz7NBJ

>A

C

H

c

V

LS

F

3

r

0M

u

R

g

W

8

m

w

f

Tn

p

5

d

Z

G

y

Y

P

9

+

<latexi sh

1_b

64="y5+B

UdWOZ

Pfr

S

9ug

/X

I

>A

CJ

c

VDL

N

F

3

o

ER8

q

Km

2

GT

zM

Y

n

H

w

vk

j

p

0

Q

7

<latexi sh

1_b

64="qEX

SVr fTB

nN

Kv

P

C+

8A

>

GH

c

DL

FJ3

d

Y

U

ku

y

wmQ

O

z

M

/

9

7g

0

j

Z

W

2

R

5

o

I

<latexi sh

1_b

64="8R

S9

2U

r

7/

cn

GH

du+Q

>A

B

V

L

gN

F

wTfz

o

Z

D

mY

K

503

pO

y

W

X

J

q

j

C

E

I

k

M

<latexi sh

1_b

64="BoqC

zOu

VX

n

DFT3G

j

>A

H

c

LS

N

J

U

2v gk

w

ZK

y

m

p

5

Yd

W

I

/

9

7

M0

E

r

8

P

f

R

<latexi sh

1_b

64="O

VjuL2U

k

RGn mE

w

3BC

>A

H

c

D

S

N

F

W

r

Y

8d

Q

JP

p

0

Z

f

/

7

+

M

X

v

y

I

T

o

5

z

q

<latexi sh

1_b

64="+

XvKO9

/Tk7

PL

V

r

AY

>

CI

c

DJSgNBE

2

G

p

j

d

Q

H

RM

Z

w

qU

8

z

0

5o

F

m

n

W

f

u

3

y

related topics, such as direct latent optimization in diffusion models and a detailed review of sampling and statistical theory of diffusion models, are deferred to Appendix A.

Classifier Guidance and Training-free Guidance. [22] introduced classifier-based guidance, steering pre-trained diffusion models towards a particular class during inference. This method offers flexibility by avoiding task-specific fine-tuning, but still requires training a time-dependent classifier. Training-free guidance methods [16, 57, 70, 4, 32, 50, 28] eliminate the need for a timedependent classifier, using only off-the-shelf loss guidance during inference. [16, 57, 32, 28] is a line of works solving inverse problems on image and [70, 50] aims for guided/conditional image generation. Though not originally developed for solving optimization problems, [16, 70] both propose a similar guidance to ours: taking gradient on the predicted clean data x 0 with respect to corrupted x t . Differently, our paper presents the first rigorous theoretical analysis of this gradient-based guidance approach. Furthermore, we propose an algorithm that iteratively applies the guidance as a module to the local linearization of the optimization objective, demonstrating provable convergence guarantees.

Fine-tuning of Diffusion Models. Several methods for fine-tuning diffusion models to optimize downstream reward functions include RL-based fine-tuning [7, 26] and direct backpropagation to rewards [18, 52, 68, 63]. However, these approaches often suffer from high computational costs and catastrophic forgetting in pre-trained models. Our guidance method is training-free and applied during the inference phase, eliminating the need to fine-tune diffusion models.

Theory of Diffusion Models. Current theory works primarily focus on unconditional diffusion models. Several studies demonstrate that the distributions generated by diffusion models closely approximate the true data distribution, provided the score function is accurately estimated [20, 2, 8, 39, 13, 40, 15, 14, 6]. For conditional diffusion models, [71, 27] establish sample complexity bounds for learning generic conditional distributions. Our novel analysis establishes a connection between the sampling process in gradient-based guided diffusion and a proximal gradient step, providing convergence guarantees.

3 Preliminaries: Diffusion Models and Guidance

Score-based diffusion models capture the distribution of pre-training data by learning a sequence of transformations to generate new samples from noise [60]. A diffusion model comprises a forward and a backward process, for which we give a review as follows.

Forward Process. The forward process progressively adds noise to data, and then the sample trajectories are used to train the score function. The forward process initializes with X 0 ∈ R D , a random variable drawn from the pre-training data D . It introduces noise via an Ornstein-Uhlenbeck process, i.e., where ( W t ) t ≥ 0 is Wiener process, and q ( t ) is non-decreasing. X t represents the noise-corrupted data distribution at time t . The conditional distribution X t | X 0 = x 0 is Gaussian, i.e., N ( α ( t ) x 0 , h ( t ) I D ) with α ( t ) = exp( -∫ t 0 1 2 q ( s ) ds ) and h ( t ) = 1 -α 2 ( t ) . In practice, the forward process will terminate at a large time T so that the marginal distribution of X T is close to N (0 , I D ) .

Backward Process. If reversing the time of the forward process, we can reconstruct the original distribution of the data from pure noise. With ( W t ) t ≥ 0 being another independent Wiener process, the backward SDE below [3] reverses the time in the forward SDE (1),

Here p t ( · ) denotes the marginal density of X t in the forward process. In the forward SDE (2), the score function ∇ log p t ( · ) plays a crucial role, but it has to be estimated from data.

Score Matching. To learn the unknown score function ∇ log p t ( · ) , we train a score network s θ ( x, t ) using samples from forward process. Let D denote the data for training. Then the score network is

learned by minimizing the following loss:

where S is a given function class, E D denotes the empirical expectation over training data D and E x t | x 0 denotes condition expectation over the forward process, ϕ t ( x t | x 0 ) is the Gaussian transition kernel, i.e., (2 πh ( t )) -D/ 2 exp( -∥ x t -α ( t ) x 0 ∥ 2 / (2 h ( t ))) .

Generation and Guided Generation. Given a pre-trained score function s θ , one generates samples by the backward process (2) with the true score replaced by s θ . Further, one can add additional guidance to steer its output distribution towards specific properties, as formulated in Module 1.

Module 1 Guided_BackwardSample ( s θ , G )

1: Input : Score s θ , guidance G default to be zero for unguided generation.

2: Hyper-parameter : T .

3: Initialized at X ← t ∼ N (0 , I ) , simulate the following SDE till time T :

4: Output : Sample X ← T .

A common goal of guided generation (Module 1) is to generate X with a desired property Y = y from the distribution P ( X | Y = y ) . To this end, it essentially needs to learn the conditional score function ∇ x t log p t ( x t | y ) . The Bayes rule gives

Whenapre-trained score network s θ ( x t , t ) ≈ ∇ log p t ( x t ) , what remains is to estimate ∇ x t log p t ( y | x t ) and add it as the guidance term G to the backward process in Module 1.

Classifier and Classifier-Free Guidance. Classifier guidance [60, 22] samples from P ( X | Y = y ) when Y is a discrete label. This method estimates ∇ x t log p t ( y | x t ) by training auxiliary classifiers, denoted as ˆ p ( y | x t , t ) , and then computing the gradient of the classifier logits as the guidance, i.e., G ( x t , t ) = ∇ x t log ˆ p ( y | x t , t ) . Alternatively, classifier-free guidance [33] jointly trains a conditional and an unconditional diffusion model, combining their score estimates to generate samples.

Notations. For a random variable X , P x represents its distribution, and p ( x ) denotes its density. For X , Y jointly distributed, P ( X | Y = y ) denotes the conditional distribution, and p ( x | y ) its density. The conditional expectation is denoted as E [ x | y ] . Let D be the pre-training data, and let E D be the empirical expectation over D . The empirical mean and covariance matrix are ¯ µ := E x ∈D [ x ] and ¯ Σ := E x ∈D [( x -¯ µ )( x -¯ µ ) ⊤ ] . For a matrix A , Span( A ) denotes the subspace spanned by its column vectors, and for a square matrix A , A -1 denotes its inverse or Moore-Penrose inverse. For any differentiable function f : R n → R m , ∇ f ∈ R m × n denotes Jacobian matrix, i.e., ( ∇ f ) ij = ∂f i ( x ) ∂x j .

4 A Primer on Gradient Guidance

Let us start with stating the problem we want to study: suppose given a pre-trained diffusion model where the score network s θ ( x t , t ) provides a good approximation to the true score function ∇ log p ( x t ) , the goal is to generate novel samples with desired properties that can be measured by a user-specified differentiable function f . We will refer to f as a reward or objective function later on. To achieve this goal, from Section 3, we know that guided generation with diffusion model is a good candidate, which deploys the following guided backward process (Module 1):

Here the guidance term G is what we focus on and wish to design. Specifically, we want to construct this guidance term G based on the gradient ∇ f of a general objective f . This is motivated by the

gradient methodology in optimization, a natural, intuitive way for adding guidance is to steer the generated samples towards the steepest ascent direction of f [16, 4, 18].

4.1 Structural Data Distribution with Subspace

When incorporating property optimization in the generation process, it's crucial to consider intrinsic low-dimensional structures of real-world data, such as local regularities, global symmetries, and repetitive patterns [62, 55, 51]. Blindly improving f at the cost of losing these structures degrades sample quality dramatically. This quality degradation, also known as 'reward over-optimization', is a common challenge for adapting diffusion models towards an external reward [71, 63].

To study the design of guidance that mitigates the risk of over-optimization, we focus on data that admits a low-dimensional latent subspace, formulated in the following assumption.

Assumption 1 (Subspace Data) . Data X ∈ R D can be represented as X = AU , where A ∈ R D × d is an unknown matrix with orthonormal columns, and the latent variable U ∈ R d follows some distribution P u with a density p u . Here d ≪ D . The empirical covariance of U is assumed full rank.

In the rest of this section, we investigate the principles for designing a guidance based on the gradient of f that ensures (i) improving the value of f , and at the same time, (ii) being adhere to the subspace structure, i.e. generated samples being close to the subspace spanned by A .

4.2 Naive Gradient Does't Work as Guidance

A tempting simple choice of the guidance G is by taking the steepest ascent direction ∇ f , which we refer to as naive gradient guidance i.e.,

However, the naive gradient guidance (5) would jeopardize the latent structure of data, which is demonstrated by the following proposition:

Proposition 1 (Failure of Naive Guidance) . For naive guidance G ( X ← t , t ) = b ( t ) ∇ f ( X ← t ) , suppose b ( t ) > b 0 > 0 for t > t 0 . For data in subspace under Assumption 1 and reward f ( x ) = g ⊤ x , g ⊥ Span( A ) with h ( t ) = 1 -exp( - √ t ) , then the off-subspace component of the generated sample is consistently large:

The intuition provided by Proposition 1 is, while the pre-trained score network effectively steers the distribution toward the latent subspace [10], the gradient vector ∇ f may point outside the subspace, causing the generated output to deviate from it (Figure 2). This is why naive gradient guidance fails. [63] also observed this, explaining that ∇ f is not computed for t = T i.e., it is not aligned with the clean data space.

Pre-trained

Random

Noise

Large

Reward

Region

Gradient Direction

Subspace

Naive Gradient

Guidance

O

ff

Subspace

naive gradient

Gloss

on support space

1.0

off support direction

0.8

0.6

0.4

0.2

0.0

t

4.3 Motivating Latent Subspace Preserving Guidance from Conditional Score Function

Failure of the naive gradient in maintaining data structure motivates us to seek alternatives. To get some inspiration, we start with the most elementary Gaussian probabilistic model and linear f . Later we will drop this assumption and consider general data distributions and general f .

Assumption 2 (Gaussian Linear model) . Let data follow a Gaussian distribution, i.e., X ∼ N ( µ, Σ) , and let f ( x ) = g ⊤ x be a linear function for some g ∈ R D . Let Y = f ( X ) + ϵ with independent, identically distributed noise ϵ ∼ N (0 , σ 2 ) for some σ > 0 .

<latexi sh

1_b

64="

cQ/

LzPoTK

0

>A

B

H

V

F

w

f

+

Z

R

m

Gn

5M

U

I

y

W

2

3

rX

J

q

k

8

u

C

E

By the Bayes' rule, the conditional score ∇ x t log p t ( x t | y ) takes the form of a sum given by

Under the Gaussian assumption, we derive the following closed-form formula of the guidance term log p t ( y | x t ) that we want to estimate. The proof is provided in Appendix D.4.

Lemma 1 (Conditional Score gives a Gradient-like Guidance) . Under Assumption 2, we have where E [ x 0 | x t ] denotes the conditional expectation of x 0 given x t in the forward process (1) , and σ 2 y ( x t ) is the variance of the conditional distribution Y | X t = x t .

The form of conditional score shown in Lemma 1 motivates our proposed gradient guidance:

Definition 1 (Gradient Guidance of Look-Ahead Loss) . Given a gradient vector g , define the gradient guidance of look-ahead loss as

The formula in (7) generalizes the intuition of a conditional score for any data distribution and objective function. It scales with the residual term y -g ⊤ E [ x 0 | x t ] , tuning the strength of guidance . Here, E [ x 0 | x t ] represents the expected clean data x 0 given x t in the forward process, which coincides with the expected sample in the backward view. This residual measures the look-ahead gap between the expected reward of generated samples and the target value. The look-ahead loss ( y -g ⊤ E [ x 0 | x t ]) 2 resembles the proximal term commonly used in first-order proximal optimization methods.

G loss ( x t , t ) := -β ( t ) · ∇ x t ( y -g ⊤ E [ x 0 | x t ] ) 2 , (7) where β ( t ) > 0 , y ∈ R are tuning parameters, and E [ x 0 | x t ] is the conditional expectation of x 0 given x t in the forward process (1) , i.e., d X t = -1 2 q ( t ) X t d t + √ q ( t ) d W t .

Remark. The gradient guidance (7) aligns with the groundtruth conditional score in (6) under the assumptions of Gaussian data and linear reward (Assumption 2). This theoretical motivation, rooted in a fundamental framework, distinguishes our work from the empirical practice, such as DPS [16] and universal guidance [4].

A key advantage of G loss is that it enables preserving the subspace structure, for any data distribution under Assumption 1. This is formalized in the following theorem, the full proof in Appendix D.5.

Theorem 1 (Faithfulness of G loss to the Low-Dimensional Subspace of Data) . Under Assumption 1, it holds for any data distribution and g ∈ R D that

We will show that the Jacobian ∇ x t E [ x 0 | x t ] maps any vector g ∈ R D to Span( A ) . To see this, we utilize the score decomposition result in Appendix D.1 and plug it into the equality E [ x 0 | x t ] = α -1 ( t ) ( x t + h ( t ) ∇ log p t ( x t )) (Tweedie's formula [24]), we have

E [ x 0 | x t ] = α -1 ( t ) ( x t + h ( t ) [ Am ( A ⊤ x t ) -h -1 ( t ) x t ]) = h ( t ) /α ( t ) · Am ( A ⊤ x t ) , (9) here m ( u ) = ∇ log p LD t ( u ) + h -1 ( t ) u , p LD t ( u ) latent density (Appendix D.1). We see ∇ x t E [ x 0 | x t ] ⊤ maps any vector to Span( A ) because m ( · ) takes A ⊤ x t as input in the expression of E [ x 0 | x t ] . ■

We highlight that the faithfulness of G loss holds for arbitrary data distribution supported on the latent subspace. It takes advantage of the score function's decomposition (19), having the effect of automatically adapting g onto the latent low-dimensional subspace of data.

Remark. We provide a rigorous guarantee for manifold preservation of gradient guidance, a property previously discussed by [16, 17, 32]. However, while [16, 17] claim manifold preservation, they do not present a formal mathematical proof. [32] relies heavily on pre-trained autoencoders for manifold projections, which are often unavailable in practical scenarios.

4.4 Estimation and Implementation of G loss

In this section, we discuss the estimation and computation of G loss based on a pre-trained score function s θ in practice. G loss involves the unknown quantity E [ x 0 | x t ] . One can construct estimate E [ x 0 | x t ] by considering the Tweedie's formula [24]: ∇ log p t ( x t ) = -h -1 ( t ) E [ x t -α ( t ) x 0 ∣ ∣ x t ] , which gives rise to ˆ E [ x 0 | x t ] := α -1 ( t ) ( x t + h ( t ) s θ ( x t , t )) , (10)

and we refer to it as the look-ahead estimator . The estimator (10) is widely adopted in practice [56, 4]. Here α ( t ) and h ( t ) are the noise scheduling used in the forward process (1).

Thus, we have obtained an implementable version of the gradient guidance G loss , given by

With a slight abuse of notation, we use G loss to refer to this implementable formula (11) in the remainder of this paper. Here, y is a target reward value from conditional score analysis under a Gaussian model and is treated as a tuning parameter in practice. The gradient guidance (11) is lightweight to implement. Given a pre-trained score function s θ in the form of a neural network, computing (11) involves calculating the squared loss

……

h

(

t

)

α

-

1

(

t

)

x

t

Gradient

Compute

ˆ

E

[

x

0

|

x

t

]

Gradient w.r.t

x

t

+

s

θ

(

x

t

, t

)

g

(

y

-

g

/latticetop

ˆ

E

[

x

0

|

x

t

])

2

using the auto-gradient feature of deep-leaning frameworks such as PyTorch and TensorFlow. See Figure 3 for illustration.

5 Gradient-Guided Diffusion Model as Regularized Optimizer

In this section, we study if gradient guidance steers pre-trained diffusion models to generate nearoptimal samples. Our results show that: 1) Iterative gradient guidance improves the objective values; 2) The pre-trained diffusion model acts as a regularizer from an optimization perspective.

5.1 Gradient-Guided Generation with A Pre-trained Score

Assuming access to a pre-trained score network s θ and the gradient of the objective function f , we present Algorithm 1 to adapt the diffusion model and iteratively update the gradient guidance (11). See Figure 1 for illustration.

Alg. 1 takes any pre-trained score function s θ as input. Each iteration evaluates ∇ f ( · ) at samples from the previous iteration (Line 5(i)), computes the gradient guidance G loss with the new gradient (Line 5(ii)), and generates new samples using the updated guidance (Module 1). The algorithm outputs an adapted diffusion model specified by ( s θ , G K ) .

Algorithm 1 Gradient-Guided Diffusion for Generative Optimization

5.2 Gradient-Guided Diffusion Converges to Regularized Optima in Latent Space

We analyze the convergence of Alg. 1 and show that in final iterations, generated samples center around a regularized solution of the optimization objective f within the subspace Span( A ) . Our theorems allow the pre-training data to have arbitrary distribution .

Assumption 3 (Concave smooth objective) . The objective f : R D → R is concave and L -smooth w.r.t. the (semi-)norm ∥·∥ ¯ Σ -1 , i.e., ∥∇ f ( x 1 ) -∇ f ( x 2 ) ∥ ¯ Σ ≤ L ∥ x 1 -x 2 ∥ ¯ Σ -1 for any x 1 , x 2 .

<latexi sh

1_b

64="MXp+DK

S

Y

9E

0R

C

jzLc

>A

B/n

V

N

F

3

W

ur

J

kU

v

7

T

d

mZ

I

o

Q

f8

w

H

5

2

q

yO

P

g

G

<latexi sh

1_b

64="JjN

oy

08E2ST

GX

Fm qw

>A

CDn

c

V

L

3

rPW

g

p

7K

fU

YdOH

z

k

/

B

v

u5R

M

9

Z

I

<latexi sh

1_b

64="qEX

SVr fTB

nN

Kv

P

C+

8A

>

GH

c

DL

FJ3

d

Y

U

ku

y

wmQ

O

z

M

/

9

7g

0

j

Z

W

2

R

5

o

I

<latexi sh

1_b

64="MCf

y

mU

v/

OdNR

VS

8

>A

B

3

c

DL

FJ

rP

n

q

ZTK

E

kz

I5

Q

g

j

w

7p

Y

2

X

u

H

o0

9

G

W

+

<latexi sh

1_b

64="Og ow3Zrp

f

c

L

QY

mBz

U

>A

C

X

VD

S

N

FJ

2vqk

0W

y

M

u

G

E

I

/

9

8

T

d

j

H

R

5

P

7

+

n

K

<latexi sh

1_b

64="MXp+DK

S

Y

9E

0R

C

jzLc

>A

B/n

V

N

F

3

W

ur

J

kU

v

7

T

d

mZ

I

o

Q

f8

w

H

5

2

q

yO

P

g

G

<latexi sh

1_b

64="O0N

y

D

2

Q

BYzg

n

d

T

E

>A

CJ3

c

V

LS

F

U

R

Pp f/

G

+

u

7

j

o

9

5

M

X

k

m

W

H

8

<latexi sh

1_b

64="8R

S9

2U

r

7/

cn

GH

du+Q

>A

B

V

L

gN

F

wTfz

o

Z

D

mY

K

503

pO

y

W

X

J

q

j

C

E

I

k

M

While Alg. 1 works with any pre-trained score network, we study its optimization properties focusing on the class of linear score functions given by

Remark on the linear parametrization of score network (12) : Analyzing the output distribution of guided diffusion is challenging because the additional guidance term destroys the dynamics of reverse SDE. A linear score is a natural and reasonable choice for characterizing the output distribution, as it was also adopted by [44].

With a linear score function (12), pre-training a diffusion model is equivalent to using a Gaussian model to estimate and sample from the estimated distribution. Thus, the guidance G loss is also linear in x t , and the final output follows a Gaussian distribution; see (25) in Appendix E. We focus on the mean, µ K , of the generated distribution from the backward sampling of ( s θ , G K ) (as T →∞ ), and establish its optimization guarantee.

Theorem 2 (Convergence to Regularized Maxima in Latent Subspace in Mean) . Let Assumptions 1 and 3 hold. Suppose we use the score function class (12) for pre-training and computing guidance. Then Alg.1 gives an adapted diffusion model that generates new samples that belong to Span( A ) . Further, for any λ > L, there exists β ( t ) , { y k } and batch size B k , such that with high probability 1 -δ , the mean of the output distribution µ K converges to be near x ∗ A,λ , and it holds where x ∗ A,λ is an optimal solution of the regularized objective:

Remarks. (1) The regularization term λ 2 ∥ x -¯ µ ∥ 2 ¯ Σ -1 (23) centers the data's mean ¯ µ and is stronger in directions where the original data has low variance. Thus, the pre-trained score acts as a "prior" in the guided generation, favoring samples near the pre-training data, even with guidance.

where ¯ µ, ¯ Σ are empirical mean and covariance of pre-training data D .

(2) The regularization term cannot be arbitrarily small, as our theorem requires λ ≥ L . Thus, only adding gradient guidance cannot achieve the global maxima. If the goal is global optima, the pre-trained score must be updated and refined with new data, as explored in Section 6.

(3) The convergence rate is linear in the latent dimension d , rather than data dimension D . Since G loss is faithful to the latent subspace (Theorem 2), the generated samples and optimization iterates of Alg. 1 remain within Span( A ) . This leverage of the latent structure results in faster convergence.

6 Gradient-Guided Diffusion with Adaptive Fine-Tuning for Global Optimization

In the previous section, we have seen that adding guidance to a pre-trained diffusion model can't improve the objective function unlimitedly due to the pre-trained score function acting as a regularizer. We consider adaptively fine-tuning pre-trained diffusion models to attain global optima. Empirically, fine-tuning diffusion models using self-generated samples has been explored by [7, 18].

6.1 Adaptive Fine-Tuning Algorithm with Gradient Guidance

We propose an adaptive version of the gradient-guided diffusion, where both the guidance and the score are iteratively updated using self-generated samples. The full algorithm is given in Algorithm 2.

We introduce a weighting scheme to fine-tune the score network using a mixture of pre-training data and newly generated samples. In Round k , let D 1 , . . . , D k be sample batches generated from the previous rounds. Let { w k,i } k i =0 be a set of weights. Conceptually, at Round k , we update the model by minimizing the weighted score matching loss:

where D 0 := D is the pre-training data. For illustration, please see also Figure 1, and the practical implementation of Alg. 2 is in Appendix F.

Algorithm 2 Gradient-Guided Diffusion with Adaptive Fine-tuning

1: Input : Pre-trained score s θ ( · , · ) , differentiable objective function f . 2: Tuning Parameter : strength parameter β ( t ) , { y k } K -1 k =0 , weights {{ w k,i } k i =0 } K -1 k =0 , number of iterations K , batch sizes { B k } . 3: Initialize : s θ 0 = s θ , G 0 = NULL . 4: for k = 0 , · · · , K -1 do 5: Generate : Sample a batch D k = { z k,i } B k i =1 from Guided_BackwardSample ( s θ k , G k ) (Module 1). 6: Compute Guidance: (i) Compute sample mean ¯ z k = (1 /B k ) ∑ B k i =1 z k,i , and query gradient g k = ∇ f (¯ z k ) . (ii) Update s θ k to s θ k +1 by minimizing the re-weighted objective (14). (iii) Compute G k +1 ( · , · ) = G loss ( · , · ) in (7), using s θ k +1 and g k , with parameter y k , β ( t ) . 7: end for 8: Output: ( s θ K , G K ) .

6.2 Guided Generation Finds Unregularized Global Optima

Finally, we analyze the optimization properties for gradient-guided diffusion model with iterative finetuning. We establish that the process of Algorithm 2 yields a final output distribution whose mean, denoted by µ K , converges to the global optimum of f .

For simplicity of analysis, we study the following function class where ˆ C t is set to stay the same in pre-trained scores, with only b t updated during iterative fine-tuning.

Theorem 3 (Convergence to Unregularized Maxima in Latent Subspace in Mean) . Let Assumptions 1 and 3 hold, and assume there exists M > 0 such that ∥ ∥ ∥ x ∗ A,λ ∥ ∥ ∥ < M for all λ ≥ 0 . Suppose we use the score function class (12) for pre-training s θ and the class (15) for finetuning it. Then Algorithm 2 gives an adapted diffusion model that generates new samples belonging to Span( A ) . Further, there exists { β ( t ) } , { y k } , { B k } and { w k,i } , such that with probability 1 -δ , where f ∗ A = max { f ( x ) | x ∈ Span( A ) } .

Theorem 3 illustrates that fine-tuning a diffusion model with self-generated data can reach global optima while preserving the latent subspace structure. The convergence rate matches standard convex optimization in terms of gradient evaluations, K . Compared to standard gradient solvers, guided diffusion models leverage pre-training data to solve optimization problems in a low-dimensional space, preserving desired structures and enabling more efficient exploration and faster convergence.

7 Experiments

7.1 Simulation

We conduct numerical simulations of Algorithms 1 and 2, following the subspace setup described in Assumption 1. Specifically, we set d = 16 , D = 64 , The latent variable u is drawn from N (0 , I d ) and used to construct x = Au , where A is a randomly generated orthonormal matrix. We define the objective function f ( x ) = 10 -( θ ⊤ x -3) 2 . To approximate the score function, we employ a version of the U-Net [54] with 14.8M trainable parameters. More details including how to set up θ are provided in Appendix F.1.

Preserving Subspace Structure. We first demonstrate that G loss preserves the subspace structure learned from the pre-trained model. For comparison, we also tested the naive guidance G ( x t , t ) := β ( t ) ( y -g ⊤ E [ x 0 | x t ] ) g (more details in Appendix F.1.). Figure 4 (left) shows that G loss performs much better than the naive gradient G in preserving the linear subspace. Figure 4 (right)

demonstrates that off-support errors increase with adaptive score fine-tuning (Alg. 2) due to distribution shift, with G resulting in more severe errors than G loss .

Convergence Results. Figure 5 (a) and (b) show that Alg. 1 converges to a sub-optimal objective value, leaving a gap to the maximal value. This aligns with our theory that the pre-trained model acts as a regularizer in addition to the objective function. Figure 5 (c) shows that Alg 2 converges to the maximal value of the objective function. As illustrated by Figure 5 (d), samples from Alg. 1 mostly stay close to the pre-training data distribution (dotted contour area), whereas samples from Alg. 2 move outside the contour as the diffusion model is fine-tuned with self-generated data.

(a) Alg. 1:

θ

=

Aβ

∗

(b) Alg. 1:

∥

θ

⊥

∥

∥

θ

∥

∥

= 9

(c) Convergence of Alg. 2 (d) Visualizing sample gener-

9

8

10

十

Gloss

fA

Gloss

fA

W.2

f(x)

(x)

2

huox Aig. 1

Sxkp. 1

1

Gloss

-1

0

fA

2

100

200

300

400

roundk

500

600

700

800

900

1000

100

200

300

400

500

600

700

800

900

1000

100200300400500600700

roundk

800900100011001200

attion

7.2 Image Generation

We validate our theory in the image domain for Algorithm 1. We employ the StableDiffusion v1.5 model [53] as the pre-trained model. For the reward model, we follow the approach outlined by [71] to construct a synthetic model. This model is based on a ResNet-18 [31] architecture pre-trained on ImageNet [21], with the final prediction layer replaced by a randomly initialized linear layer that produces scalar outputs. For more experiment details, refer to Appendix F.2.

'rabbit'

0.2

F00

Reward

'horse'

0.2

0.4

F 9°0-

y=2

"fox'

y=4

y=10

0

3

4

5

7

Iteration

Iteration

Results. By Algorithm 1, the reward increases and converges. Figure 6 (left) shows the reward changes with optimization iterations. The hyperparameter y tunes the strength of guidance and is inversely related to the strength of the regularizer (theoretical implications in Appendix E). A larger guidance strength (smaller regularizer strength) leads to a higher convergent reward value. Figure 6 (right) illustrates the changes in generated images across iterations. As the reward increases, the images become increasingly abstract, shifting from photo-realistic with detailed backgrounds to more virtual, stylized ones.

8 Conclusion

In this paper, we focus on gradient guidance for adapting or fine-tuning pre-trained diffusion models from an optimization perspective. We investigate the look-ahead loss based gradient guidance and two variants of diffusion-based generative optimization algorithms utilizing it. We provide guarantees for adapting/fine-tuning diffusion models to maximize any target concave differentiable reward function. Our analysis extends to linear subspace data, where our gradient guidance and adaptive algorithms preserve and leverage the latent subspace, achieving faster convergence to near-optimal solutions.

Acknowledgments

Mengdi Wang acknowledges the support by NSF IIS-2107304, NSF CPS-2312093, ONR 1006977 and Genmab.

References

Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making? arXiv preprint arXiv:2211.15657 , 2022.

Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797 , 2023.

Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications , 12(3):313-326, 1982.

Arpit Bansal, Hong-Min Chu, Avi Schwarzschild, Soumyadip Sengupta, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Universal guidance for diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 843-852, 2023.

Heli Ben-Hamu, Omri Puny, Itai Gat, Brian Karrer, Uriel Singer, and Yaron Lipman. D-flow: Differentiating through flows for controlled generation. arXiv preprint arXiv:2402.14017 , 2024.

Joe Benton, Valentin De Bortoli, Arnaud Doucet, and George Deligiannidis. Linear convergence bounds for diffusion models via stochastic localization. arXiv preprint arXiv:2308.03686 , 2023.

Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305.13301 , 2023.

Adam Block, Youssef Mroueh, and Alexander Rakhlin. Generative modeling with denoising auto-encoders and langevin sampling. arXiv preprint arXiv:2002.00107 , 2020.

Sébastien Bubeck et al. Convex optimization: Algorithms and complexity. Foundations and Trends® in Machine Learning , 8(3-4):231-357, 2015.

Minshuo Chen, Kaixuan Huang, Tuo Zhao, and Mengdi Wang. Score approximation, estimation and distribution recovery of diffusion models on low-dimensional data. arXiv preprint arXiv:2302.07194 , 2023.

Minshuo Chen, Song Mei, Jianqing Fan, and Mengdi Wang. Opportunities and challenges of diffusion models for generative ai. National Science Review , page nwae348, 2024.

Minshuo Chen, Song Mei, Jianqing Fan, and Mengdi Wang. An overview of diffusion models: Applications, guided generation, statistical rates and optimization. arXiv preprint arXiv:2404.07771 , 2024.

Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru R Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. arXiv preprint arXiv:2209.11215 , 2022.

Sitan Chen, Sinho Chewi, Holden Lee, Yuanzhi Li, Jianfeng Lu, and Adil Salim. The probability flow ode is provably fast. arXiv preprint arXiv:2305.11798 , 2023.

Sitan Chen, Giannis Daras, and Alex Dimakis. Restoration-degradation beyond linear diffusions: A non-asymptotic analysis for ddim-type samplers. In International Conference on Machine Learning , pages 4462-4484. PMLR, 2023.

Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209.14687 , 2022.

Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. Advances in Neural Information Processing Systems , 35:25683-25696, 2022.

Kevin Clark, Paul Vicol, Kevin Swersky, and David J Fleet. Directly fine-tuning diffusion models on differentiable rewards. arXiv preprint arXiv:2309.17400 , 2023.

Valentin De Bortoli. Convergence of denoising diffusion models under the manifold hypothesis. arXiv preprint arXiv:2208.05314 , 2022.

Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems , 34:17695-17709, 2021.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A largescale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee, 2009.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems , 34:8780-8794, 2021.

Zehao Dou, Minshuo Chen, Mengdi Wang, and Zhuoran Yang. Theory of consistency diffusion models: Distribution estimation meets fast sampling. In Forty-first International Conference on Machine Learning .

Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association , 106(496):1602-1614, 2011.

Ahmed El Alaoui, Andrea Montanari, and Mark Sellke. Sampling from mean-field gibbs measures via diffusion processes. arXiv preprint arXiv:2310.08912 , 2023.

Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Dpok: Reinforcement learning for fine-tuning text-to-image diffusion models. arXiv preprint arXiv:2305.16381 , 2023.

Hengyu Fu, Zhuoran Yang, Mengdi Wang, and Minshuo Chen. Unveil conditional diffusion models with classifier-free guidance: A sharp statistical theory. arXiv preprint arXiv:2403.11968 , 2024.

Tomer Garber and Tom Tirer. Image restoration by denoising diffusion models with iteratively preconditioned guidance. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 25245-25254, 2024.

Nate Gruver, Samuel Stanton, Nathan C Frey, Tim GJ Rudner, Isidro Hotzel, Julien LafranceVanasse, Arvind Rajpal, Kyunghyun Cho, and Andrew Gordon Wilson. Protein design with guided discrete diffusion. arXiv preprint arXiv:2305.20009 , 2023.

Zhiye Guo, Jian Liu, Yanli Wang, Mengrui Chen, Duolin Wang, Dong Xu, and Jianlin Cheng. Diffusion models in bioinformatics: A new wave of deep learning revolution in action. arXiv preprint arXiv:2302.10907 , 2023.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770-778, 2016.

Yutong He, Naoki Murata, Chieh-Hsin Lai, Yuhta Takida, Toshimitsu Uesaka, Dongjun Kim, Wei-Hsiang Liao, Yuki Mitsufuji, J Zico Kolter, Ruslan Salakhutdinov, et al. Manifold preserving guided diffusion. arXiv preprint arXiv:2311.16424 , 2023.

Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598 , 2022.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems , 33:6840-6851, 2020.

Kaixuan Huang, Yukang Yang, Kaidi Fu, Yanyi Chu, Le Cong, and Mengdi Wang. Latent diffusion models for controllable rna sequence generation. arXiv preprint arXiv:2409.09828 , 2024.

Korrawe Karunratanakul, Konpat Preechakul, Emre Aksan, Thabo Beeler, Supasorn Suwajanakorn, and Siyu Tang. Optimizing diffusion noise can serve as universal motion priors. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 1334-1345, 2024.

Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. Advances in Neural Information Processing Systems , 35:23593-23606, 2022.

Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. arXiv preprint arXiv:2009.09761 , 2020.

Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence for score-based generative modeling with polynomial complexity. arXiv preprint arXiv:2206.06227 , 2022.

Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence of score-based generative modeling for general data distributions. arXiv preprint arXiv:2209.12381 , 2022.

Xingchao Liu, Lemeng Wu, Mao Ye, and Qiang Liu. Let us build bridges: Understanding and extending diffusion generative models. arXiv preprint arXiv:2208.14699 , 2022.

Yixin Liu, Kai Zhang, Yuan Li, Zhiling Yan, Chujie Gao, Ruoxi Chen, Zhengqing Yuan, Yue Huang, Hanchi Sun, Jianfeng Gao, et al. Sora: A review on background, technology, limitations, and opportunities of large vision models. arXiv preprint arXiv:2402.17177 , 2024.

Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11461-11471, 2022.

Pierre Marion, Anna Korba, Peter Bartlett, Mathieu Blondel, Valentin De Bortoli, Arnaud Doucet, Felipe Llinares-López, Courtney Paquette, and Quentin Berthet. Implicit diffusion: Efficient optimization through stochastic sampling. arXiv preprint arXiv:2402.05468 , 2024.

Song Mei and Yuchen Wu. Deep networks as denoising algorithms: Sample-efficient learning of diffusion models in high-dimensional graphical models. arXiv preprint arXiv:2309.11420 , 2023.

Andrea Montanari and Yuchen Wu. Posterior sampling from the spiked models via diffusion processes. arXiv preprint arXiv:2304.11449 , 2023.

Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In Proceedings of the International Conference on Machine Learning , pages 8162-8171. PMLR, 2021.

Kazusato Oko, Shunta Akiyama, and Taiji Suzuki. Diffusion models are minimax optimal distribution estimators. arXiv preprint arXiv:2303.01861 , 2023.

Jiachun Pan, Jun Hao Liew, Vincent YF Tan, Jiashi Feng, and Hanshu Yan. Adjointdpm: Adjoint sensitivity method for gradient backpropagation of diffusion probabilistic models. arXiv preprint arXiv:2307.10711 , 2023.

Jiachun Pan, Hanshu Yan, Jun Hao Liew, Jiashi Feng, and Vincent YF Tan. Towards accurate guided diffusion sampling through symplectic adjoint method. arXiv preprint arXiv:2312.12030 , 2023.

Phillip Pope, Chen Zhu, Ahmed Abdelkader, Micah Goldblum, and Tom Goldstein. The intrinsic dimension of images and its impact on learning. arXiv preprint arXiv:2104.08894 , 2021.

Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, and Katerina Fragkiadaki. Aligning textto-image diffusion models with reward backpropagation. arXiv preprint arXiv:2310.03739 , 2023.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10684-10695, 2022.

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18 , pages 234-241. Springer, 2015.

Sam T Roweis and Lawrence K Saul. Nonlinear dimensionality reduction by locally linear embedding. science , 290(5500):2323-2326, 2000.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502 , 2020.

Jiaming Song, Qinsheng Zhang, Hongxu Yin, Morteza Mardani, Ming-Yu Liu, Jan Kautz, Yongxin Chen, and Arash Vahdat. Loss-guided diffusion models for plug-and-play controllable generation. In International Conference on Machine Learning , pages 32483-32498. PMLR, 2023.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems , 32, 2019.

Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. In Uncertainty in Artificial Intelligence , pages 574-584. PMLR, 2020.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Zhiwei Tang, Jiangweizhi Peng, Jiasheng Tang, Mingyi Hong, Fan Wang, and Tsung-Hui Chang. Tuning-free alignment of diffusion models with direct noise optimization. arXiv preprint arXiv:2405.18881 , 2024.

Joshua B Tenenbaum, Vin de Silva, and John C Langford. A global geometric framework for nonlinear dimensionality reduction. science , 290(5500):2319-2323, 2000.

Masatoshi Uehara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Finetuning of continuous-time diffusion models as entropy-regularized control. arXiv preprint arXiv:2402.15194 , 2024.

Bram Wallace, Akash Gokul, Stefano Ermon, and Nikhil Naik. End-to-end diffusion latent optimization improves classifier guidance. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 7280-7290, 2023.

Yinhuai Wang, Jiwen Yu, and Jian Zhang. Zero-shot image restoration using denoising diffusion null-space model. arXiv preprint arXiv:2212.00490 , 2022.

Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature , 620(7976):1089-1100, 2023.

Andre Wibisono, Yihong Wu, and Kaylee Yingxi Yang. Optimal score estimation via empirical bayes smoothing. arXiv preprint arXiv:2402.07747 , 2024.

Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation. arXiv preprint arXiv:2304.05977 , 2023.

Ling Yang, Zhilong Zhang, Yang Song, Shenda Hong, Runsheng Xu, Yue Zhao, Wentao Zhang, Bin Cui, and Ming-Hsuan Yang. Diffusion models: A comprehensive survey of methods and applications. ACM Computing Surveys , 56(4):1-39, 2023.

Jiwen Yu, Yinhuai Wang, Chen Zhao, Bernard Ghanem, and Jian Zhang. Freedom: Trainingfree energy-guided conditional diffusion model. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 23174-23184, 2023.

Hui Yuan, Kaixuan Huang, Chengzhuo Ni, Minshuo Chen, and Mengdi Wang. Reward-directed conditional diffusion: Provable distribution estimation and reward improvement. arXiv preprint arXiv:2307.07055 , 2023.

A Additional Related Works

Direct Latent Optimization in Diffusion Models. Besides guidance methods, an alternative training-free route by optimizing the initial value of reverse process [64, 5, 36, 61, 49]. These methods typically backpropagate the gradient of reward directly to the initial latent vector through an ODE solver, utilizing the chain rule. Thus at inference time, the reverse process is unchanged except for being fed with an optimized initialization, different from the guidance method we studied.

Sampling and Statistical Theory of Diffusion Model In contrast to the fruitful empirical advances, the theory of diffusion models is still limited. To the best of our knowledge, a theoretical understanding of fine-tuning diffusion models is absent. Existing results mainly focus on the sampling ability and statistical properties of unconditional diffusion models. In particular, for sampling ability, a line of works shows that the distribution generated by a diffusion model is close to the data distribution, as long as the score function is accurately estimated [20, 2, 8, 39, 13, 40]. The accuracy of the estimated score function is measured in terms of an L ∞ or L 2 -norm distance. More recently, [15, 14, 6] develop refined and tighter analyses using Taylor expansions of the discretized backward process and localization method. It is worth mentioning that the analysis in [15, 14, 6] extends to broader sample generation processes such as deterministic ones based on probabilistic ODEs. Going beyond distributions in Euclidean spaces, [19] analyzes diffusion models for sampling distribution supported on a low-dimensional manifold. Moreover, [46] consider sampling from symmetric spiked models, and [25] study sampling from Gibbs distributions using diffusion processes.

Turning towards the statistical theory of diffusion models, [59] and [41] provide asymptotic analyses, assuming a parametric form of the score function. Unfortunately, asymptotic analysis does not lead to concrete sample complexities. Later, concurrent works, [48] and [10], establish sample complexity bounds of diffusion models for estimating nonparametric data distributions. In high dimensions, their results highlight a curse of dimensionality issue without further assumptions, which also appears in [67] considering kernel methods. More interestingly, these works demonstrate that diffusion models can circumvent the curse of dimensionality issue if the data has low-dimensional structures. In the same spirit, [45] investigate learning high-dimensional graphical models using diffusion models, without the curse of dimensionality. [23] establishes statistical theory for consistency models. For conditional diffusion models, [71, 27] establish sample complexity bounds for learning generic conditional distributions. We refer readers to [12] for an overview of contemporary theoretical progress.

B Discussion

Limitations We do not see significant limitations in our analysis. The linear score class initiates the study of gradient-like guidance. We have already demonstrated that our proposed gradient guidance preserves low-dimensional structures for any score and data distribution (Theorem 1). Moving forward, we aim to analyze the convergence properties for general score functions.

Broader Impacts This paper's exploration of gradient-guided diffusion models offers theoretical implications that could inspire further algorithmic advancements. By establishing a connection between gradient-guided diffusion and optimization, the research opens new avenues for efficient and effective optimization through a generative approach. This work has the potential to enhance the precision and personalization of AI systems across various fields, including image synthesis and drug design. However, these techniques could also be misused to generate harmful content.

C Characterization for Output Distribution of Backward Process

In this section, we provide analytical characterizations for the output distribution of the backward process guided by G loss when the pre-trained score is linear. We first give the result of score matching as follows.

Lemma 2 (Pre-training with Linear Score Functions) . Suppose for pre-training the score network, the class in (14) is

x

If we freeze C t in (12) , that is, minimizing the score matching objective (14) over the class { s ( x, t ) = C t x + b t : b t ∈ R D } gives

α

(

t

)¯

x

)

,

s θ ( x t , t ) = -( α 2 ( t ) ¯ Σ+ h ( t ) I D ) -1 ( x t -α ( t )¯ x ) , where ¯ Σ = ( ∑ k i =0 w k,i E x ∈D k [ ( x -¯ x )( x -¯ x ) ⊤ ] ) /( ∑ k i =0 w k,i ) are weighted data covariance.

s θ ( x t , t ) = C t ( t -where ¯ x = ( ∑ k i =0 w k,i E x ∈D i [ x ] )/( ∑ k i =0 w k,i ) . Moreover, minimizing the score matching objective (14) over the class (12) yields

Proof. Using the linear score network class S with freezing C t , we cast the score matching loss (14) into where w = ∑ k i =0 w k,i , equality ( i ) follows from computing the expectation over the conditional Gaussian distribution of x t | x 0 . We note that b t should minimize ∑ k i =0 w k,i E x 0 ∈D i [ ∥ b t + α ( t ) C t x 0 ∥ 2 ] for any t , which leads to ˆ b t = -α ( t ) C t ¯ x.

Now, we solve C t for the second result. Substituting ˆ b t into the optimization objective (14) yields:

Taking the gradient for C t , we get

Setting the gradient above to 0, we get the solution for C t as

Therefore, the proof is completed.

When w k, 0 = 1 , w k,i = 0 , i ∈ [ k ] , Lemma 2 reduces to the pre-traning score matching.

Corollary 1. Let D be the pre-training data. Minimizing the score matching objective (3) over the function class (12) gives

The following lemma characterizes the output distribution of the backward process guided by G loss when the pre-trained score is linear.

Lemma 3. If the pre-trained score s θ ( x t , t ) is (17) , substituting the score function with s θ ( x t , t ) + G loss ( x t , t ) in the backward SDE (2) yields, when T →∞ , with β ( t ) assigned as β ( t ) = 1 2 ( σ 2 + g ⊤ ¯ Σ -1 ( I D + α 2 ( t ) ¯ Σ /h ( t ) ) -1 g ) -1 . Moreover, if pretraining data reside in Span( A ) following Assumption 1, it holds X ← T ∈ Span( A ) .

Proof. Consider X 0 d = N ( ¯ µ, ¯ Σ ) , Y = g ⊤ X 0 + ϵ where ϵ ∼ N ( 0 , σ 2 ) . Let X 0 be the initialization of the forward process. Similar to the proof in Appendix D.4, we get

Thus, we get s θ ( x t , t ) is exactly the score of marginal distribution of X t , i.e., ∇ log p t ( x t ) = s θ ( x t , t ) . According to the proof in Appendix D.4, we get s θ ( x t , t ) + G loss ( x t , t ) = ∇ log p t ( x t | y ) . Thus, the backward SDE turns out to be

The initial distribution p 0 ( x 0 | y ) of the forward process can also be obtained by (18) where we replace the initial distribution as p T ( x T | y ) . According to the data processing inequality, we get the bound of the total variation distance between the terminal distribution p ← T of (18) and p 0 ( x 0 | y ) :

where p 0 , p T are short hands for p 0 ( x 0 | y ) and p T ( x T | y ) , and φ ( · ) is the density for the standard normal distribution N (0 , I D ) . Since in the forward process, p T → φ when T → ∞ , we have TV ( p 0 , p ← T ) → 0 when T →∞ . We complete the first part of the lemma. As for the second part, if data reside in Span( A ) following Assumption 1, we have ¯ µ = A ¯ u and ¯ Σ = A ¯ Σ u A ⊤ , where and due to X ← T follows Gaussian distribution, we get X ← T ∈ Span( A ) . Thus, the proof is completed.

D Additional Materials for Section 4

Contents

D.1 Score decomposition for subspace data

Under Assumption 1, the score function ∇ log p t ( x ) decomposes to two orthogonal parts: an onsupport component belonging to the subspace; and an orthogonal component. We recall this key result in Proposition 2, which later plays a key role in deriving subspace preserving guidance.

Proposition 2 (Score Decomposition for Subspace Data ([10] Lem. 1, Thm. 3)) . Under Assumption 1, the score function ∇ log p t ( x ) decomposes as where p LD t ( u ′ ) = ∫ ϕ t ( u ′ | u ) p u ( u ) d u with ϕ t ( ·| u ) being the density of N ( α ( t ) u, h ( t ) I d ) for the same α ( t ) and h ( t ) in the forward process (1) .

D.2 Proof of Proposition 1

We give the proof of Proposition 1, which shows the failure of naive gradient guidance.

Proof. Under Assumption 1, the score can be decomposed to terms parallel and orthogonal to Span( A ) (Proposition 2). Applying naive guidance, we examine the orthogonal reverse process:

̸

Solving this SDE, we get the expectation of the final state following E [ X ← T, ⊥ ] = ∫ T 0 exp ( -∫ t 0 h -1 ( s )d s ) e t/ 2 b ( T -t ) g d t . For the schedule h ( t ) = 1 -exp( - √ t ) , we have the coefficient of direction g is larger than ∫ T 0 exp( -t/ 2 -2 √ t ) b ( T -t )d t > ∫ 1 0 exp( -5 / 2) b 0 d t > 0 where we can assume T > 1 . Thus, E [ X ← T, ⊥ ] = 0 . This means the generated sample is leaving the subspace, i.e., naive gradient guidance will violate the latent structure.

D.3 Discussion for gradient-like guidance

To further clarify the derivation of our gradient guidance, we present a variant of Lemma 1.

Proposition 3. Under Assumption 2, we have

Remarks. Observe that, when Σ = I , (20) suggests the following form of guidance that is aligned with the naive gradient, i.e., the steepest ascent direction:

However, even for Gaussian distributions, as long as Σ = I , the term of (6) is no longer proportional to g but becomes a pre-conditioned version of the gradient. We show the guidance above can maintain the subspace structure of data in the experiments Appendix F.

̸

Another observation is that this guidance scales with a residual term y -g ⊤ E [ x 0 | x t ] . In particular, the residual term y -g ⊤ E [ x 0 | x t ] tunes the strength of guidance . Recall E [ x 0 | x t ] denotes the posterior expectation of clean data x 0 given x t in the forward process. Thus, in a backward view, E [ x 0 | x t ] coincides with the expected sample to be generated conditioned on x t . In this sense, the quantity y -g ⊤ E [ x 0 | x t ] measures a look-ahead gap between the expected reward of generated samples and the target value. A larger absolute value of the residual means stronger guidance in the backward generation process.

Under q ( t ) ≡ 1 in the forward process (1), we plot the theoreti-

cal choice of β ( t ) and α ( t ) , h ( t ) to t in Figure 7. In practice, the choice of α ( t ) , h ( t ) can vary and they are determined by the forward process used for pre-training; and β ( t ) can be treated as a tuning parameter to adjust the strength of guidance.

The proof of Proposition 3 is included in the proof of Lemma 1.

D.4 Proof of Lemma 1

Proof. Recall { X t } t ≥ 0 is the stochastic process from the forward process. X 0 d = N ( µ, Σ) , Y = g ⊤ X 0 + ϵ where ϵ ∼ N (0 , σ 2 ) is independent with X 0 . Since X 0 , X t and Y are joint Gaussian, we have Y | X t also follows Gaussian distribution, denoted as N ( m y ( x t ) , σ 2 y ( x t ) ) .Then, the closed form of ∇ x t log p t ( y | x t ) can be derived as

Due to the linearity of Y with regard to X 0 , m y ( x t ) can be computed as

To get the variance σ 2 y ( x t ) , we compute the joint distribution ( X t , Y ) . In the forward process, given X 0 = x 0 , X t can be written as α ( t ) x 0 + Z t for Z t d = N (0 , h ( t ) I D ) independent of x 0 . Due to the linear function assumption, we have

Observing that the joint distribution of ( x 0 , ϵ, z t ) is Gaussian, we deduce

Thus, we get σ 2 y ( x t ) = σ 2 + g ⊤ Σ g -α 2 ( t ) g ⊤ Σ ( α 2 ( t )Σ + h ( t ) I D ) -1 Σ g . Together with the derivation of the mean m y ( x t ) (21), we get

To get E [ x 0 | x t ] , we derive the joint distribution of ( X 0 , X t ) :

Thus, we get E [ x 0 | x t ] = µ + α ( t ) ( α 2 ( t )Σ + h ( t ) I D ) -1 Σ( x t -α ( t ) µ ) . As a consequence, we have

Together with the following equality by Woodbury identity, we get the result.

D.5 Proof of Theorem 1

Proof. We expand the derivative in G loss as

It holds E [ x 0 | x t ] = 1 α ( t ) ( x t + h ( t ) ∇ log p t ( x t )) . Via the score decomposition under linear subspace data in Chen et al. [10, Lemma 1], we have

where p LD t denotes the diffused latent distribution, i.e., p LD t ( u ′ ) = ∫ ϕ t ( u ′ | u 0 ) p u ( u 0 ) d u 0 . Recall that ϕ t is the Gaussian transition kernel of the forward process and p u is the density of latent variable u 0 in Assumption 1.

To ease the derivation, we denote m ( u ) = ∇ log p LD t ( u ) + 1 h ( t ) u . It then holds that

As a consequence, we can verify that

where ∇ m ( A ⊤ x t ) ∈ R d × d is the Jacobian matrix of m at A ⊤ x t . Plugging the last display into G loss , we conclude that

E Additional Materials for in Sections 5 and 6 Contents

E.1 Convergence to Regularized Maxima in Mean

Here we present a more fundamental version of Theorem 2.

Theorem 4 (Convergence to Regularized Maxima in Mean) . Let Assumption 3 hold, and let the pre-training data D have arbitrary distribution with covariance matrix ¯ Σ ≻ 0 . Suppose the score function s θ is pre-trained via minimizing the score matching loss (3) over the linear function class (12) . Let Alg. 1 take s θ ( · , · ) and f as the input. For any λ > L, there exists { β ( t ) } , { y k } , { B k } such that, with probability ≥ 1 -δ , the mean of the output distribution µ K converges to be near x ∗ λ , and where D is the ambient dimension of data, and x ∗ λ is a regularized maximizer of f given by

where ¯ µ, ¯ Σ are empirical mean and covariance of pre-training data D .

Proof. The proof is a special case of Theorem 2 in Appendix E.2, via setting the representation matrix A = I D .

E.2 Proof of Theorem 2

We first provide a proof sketch.

Proof Sketch Solving the score matching problem (3) with a linear function class (12) yields a pre-trained score as follows

With proper choices of β ( t ) , gradient guidance G loss leads to the following output distribution at the end of round k :

Thus, we obtain the mean of the above distribution, i.e., µ k +1 = ¯ µ + η k ¯ Σ ∇ f (¯ z k ) , where ¯ z k is the empirical mean of previous samples, η k is a stepsize determined by y k . By a rearrangement, we obtain a recursive formula

We observe that (24) resembles a gradient ascent update from µ k ≈ ¯ z k to µ k +1 corresponding to a regularzed optimization problem (23). In this regularized objective, the original objective f ( x ) incorporates an additional proximal term with λ := 1 /η k . Therefore we can analyze the convergence of µ k by following the classical argument for gradient optimization.

Proof of Theorem 2. Define a filtration {H k } K -1 k =0 with H k be the information accumulated after k rounds of Alg.1.

Define the expectation of samples generated at k -th round as

Applying Corollary 1, we get the pre-trained score as

If we set y k as follows

where η = 1 /λ . And we choose β ( t ) at k -round as β ( t ) = 1 2 ( σ 2 + g ⊤ k -1 ¯ Σ -1 ( I D + α 2 ( t ) ¯ Σ /h ( t ) ) -1 g k -1 ) -1 . Then, Lemma 3 provides the generated distribution in k -th round:

Define the empirical covariance matrix of the latent variable U as ¯ Σ u = E x ∈D ,u = A ⊤ x [( u -¯ u ) ( u -¯ u ) ⊤ ] where ¯ u = E D [ u ] . Then in the subspace setting, the empirical mean and covariance of data X can be written as ¯ µ = AA ⊤ ¯ µ and ¯ Σ = A ¯ Σ u A ⊤ respectively. The mean of the sample z k,i follows

where g k -1 = ∇ f (¯ z k -1 ) and ¯ z k -1 = (1 /B ) ∑ B i z k -1 ,i . We rearrange the update rule to show a gradient ascent formula as follows where λ = 1 /η . Define h ( x ) := f ( x ) -λ/ 2 ∥ x -¯ µ ∥ 2 ¯ Σ -1 , we have

Recall the notation for the optimum: x ⋆ A,λ = argmax x = Au h ( x ) . We consider the distance of µ k to x ⋆ A,λ under the semi-norm ∥·∥ ¯ Σ -1 .

Lemma 3 shows the distribution of z k -1 ,i . Therefore, according to concentration inequality for Gaussian distribution, with the probability at least 1 -δ/K , it holds

We have trace ( V ( z k -1 ,i ) · ¯ Σ -1 ) ≤ trace ( ¯ Σ · ¯ Σ -1 ) = d . Therefore, I 2 is bounded by I 2 ≤ M 0 / √ B k -1 ,

where M 0 := ηL √ 2 log ( 2 K δ ) · d . Next, we consider the first term in (26). Since x ∗ A,λ is the optimum of h within Span( A ) , the gradient ∇ h ( x ∗ A,λ ) is in the orthogonal subspace, i.e., A ⊤ ∇ h ( x ∗ A,λ ) = 0 , thus ¯ Σ ∇ h ( x ∗ A,λ ) = 0 . The first term in (26) can be written as

Recall h is f adding a ∥·∥ ¯ Σ -1 regularized term. We get h is ( L + λ ) -smooth with respect to semi norm ∥·∥ ¯ Σ -1 which is derived from f L -smooth. Also, h is λ -strongly concave with respect to semi norm ∥·∥ ¯ Σ -1 since f is concave. According to Lemma 4, we derive

〈 µ k -1 -x ⋆ A,λ , ∇ h ( µ k -1 ) -∇ h ( x ∗ A,λ ) 〉 ≤ -λ ( L + λ ) L +2 λ ∥ ∥ µ k -1 -x ∗ A,λ ∥ ∥ 2 ¯ Σ -1 -1 L +2 λ ∥ ∥ ∇ h ( µ k -1 ) -∇ h ( x ∗ A,λ ) ∥ ∥ 2 ¯ Σ . Plugin the formula of I 1 , we get

Since η = 1 /λ , it holds η 2 -2 η L +2 λ > 0 . Due to h ( L + λ ) -smoothness, we get

I 1 I 1 ≤ ζ ∥ ∥ µ k -1 -x ⋆ A,λ ∥ ∥ ¯ Σ -1 , where ζ := | 1 -η ( L + λ ) | . Combing the upper bound of I 1 and I 2 , we get with probability at least 1 -δ/K , for 1 < k ≤ K ,

As for k = 1 , by similar derivation, we can obtain ∥ ∥ ∥ µ 1 -x ⋆ A,λ ∥ ∥ ∥ ¯ Σ -1 ≤ ζ ∥ ∥ ∥ z 0 -x ⋆ A,λ ∥ ∥ ∥ ¯ Σ -1 . By induction, we get with probability at least 1 -(( K -1) /K ) δ ,

Choose B k ≥ ζ -4 k (1 -ζ ) -2 for all k ∈ [ K -1] , then we can get

where M 1 := ηL √ 2 log ( 2 K δ ) . Since h is ( L + λ ) -smooth with respect to ∥·∥ ¯ Σ -1 , it holds

Considering that ∇ h ( x ∗ A,λ ) ⊥ Span( A ) yields 〈 ∇ h ( x ∗ A,λ ) , µ K -x ∗ A,λ 〉 = 0 , we obtain the following by rearranging the equation above

=

Since ∥ ∥ ∥ ¯ µ -x ∗ A,λ ∥ ∥ ∥ 2 ¯ Σ -1 = ∥ ∥ ∥ A ⊤ ( ¯ µ -x ∗ A,λ )∥ ∥ ∥ 2 ¯ Σ -1 u is the distance within Span( A ) , i.e., O ( d ) . Recall η = 1 /λ , ζ

final result:

|

1

-

η

(

L

+

λ

)

|

=

L/λ

,

λ > L

, and

M

1

=

ηL

2 log

2

K

δ

. Therefore, we get the

E.3 Proof of Theorem 3

We first provide a proof sketch.

Proof Sketch The proof idea is similar to the proof of Theorem 2. For simplicity, we analyze the case where only the most recent sample batch D k is merged with D 0 for finetuning the score function. More specifically, we let w k,i = 0 for 0 < i < k and w k, 0 = 1 -w k,k . Similar to the proof of Theorem 4, we obtain a recursive update rule given by

µ k +1 = ¯ z k + η k ¯ Σ [ ∇ f (¯ z k ) -(1 -w k,k ) η -1 k · ¯ Σ -1 (¯ z k -¯ µ ) ] , (29) where ¯ z k ≈ µ k is the empirical mean of previous samples. This update rule also closely resembles the gradient ascent iteration for maximizing a regularized objective. A key difference here is that we can control the weights w k,i to reduce the impact of D 0 and make the regularization term vanish to zero. Thus the mean µ k eventually converges to the global maxima.

Proof. Define

According to Lemma 2, with freezing C t in class (12), the pre-trained score in Round k is s θ k +1 ( x t , t ) = -( α 2 ( t ) ¯ Σ+ h ( t ) I D ) ( x t -α ( t )¯ x k ) where ¯ x k = ∑ k j =0 w k,j ¯ z j and ¯ z j = E x ∈D j [ x ] . By choosing y k , and weights w k,j as

where η k > 0 , 0 < w k < 1 will be specified later. And we choose β ( t ) at Round k as β ( t ) = 1 2 ( σ 2 + g ⊤ k -1 ¯ Σ -1 ( I D + α 2 ( t ) ¯ Σ /h ( t ) ) -1 g k -1 ) -1 . Lemma 3 gives the mean of distribution of z k +1 ,i as µ k +1 = ¯ x k + η k ¯ Σ g k , (30)

and the output distribution

Applying Lemma 3 yields z k,i ∈ Span(A) , thus, ¯ x k = AA ⊤ ¯ x k and ¯ Σ = A ¯ Σ u A ⊤ , we get the update rule reduced to

We set w k = 1 -η k λ and set η k = η , where λ, η > 0 will be specified later. Therefore, we have µ k +1 = AA ⊤ ¯ z k + ηA ¯ Σ u A ⊤ ∇ h λ (¯ z k ) , (32)

√

(

)

where h λ ( x ) := f ( x ) -( λ/ 2) ∥ x -¯ µ ∥ 2 ¯ Σ -1 . Define x ⋆ A,λ = argmax x = Au h λ ( x ) . With some similar steps in proof in Appendix E.2, by choosing B k ≥ ζ -4 k (1 -ζ ) -2 , together with z 0 = ¯ µ , we get with

η

=

L

2

+2

λ

,

ζ

=

|

1

-

η

(

L

+

λ

)

|

and

M

1

as in proof in Appendix E.2. We restate it here:

) log

(

2

K

δ

)

. Also, we can get (28)

= 2

L

√

(1 +

η

Adding (33), it holds

)

f ( x ∗ A,λ ) -f ( µ K ≤ ( ∥ ∥ -∥ ∥ Due to (32), we have, it holds w.p. 1 -δ ,

λ

2

x

A,λ

∗

µ

¯

2

-

1

-

∥

∥

We choose λ = L log K/ (4 K ) and get

1

-∥

µ

K

-

µ

¯

∥

-

1

+

+

2

x

A,λ

∗

-

x

A

∗

2

-

With assuming ∥ ∥ ∥ x ⋆ A,λ ∥ ∥ is bounded, we derive

)

∥

∥

∥

∥

E.4 Auxiliary Lemma

The following is a standard result in convex optimization utilized in previous proofs.

Lemma 4. Let f be α -strongly concave and β -smooth with respect to the (semi) norm ∥·∥ Σ -1 , for all x and y , it holds

Proof. See Bubeck et al. [9, Lemma 3.11] for a proof.

F Additional Materials for Experiments

F.1 Simulation

We experiment with our design of the gradient guidance as well as Algorithm 1 and Algorithm 2. Going beyond our theoretical assumptions, we adopt a 15M-parameter U-Net as the score function class for training and fine-tuning our diffusion model.

F.1.1 Experiment Setup

For linear data structure, we set the data's ambient dimension as D = 64 and the linear subspace dimension as d = 16 . The linear subspace is represented by an orthogonal matrix A ∈ R D × d . We randomly generate a matrix A and fix it once generated. After that, we sample a data point X by first randomly sampling a latent variable U ∼ N (0 , I d ) and computing X = AU . We independently sample a total of 65536 data points as our pre-training data set. For nonlinear data structure, data are uniformly sampled from a unit ball in R 64 .

The objective functions considered in our experiments are f 1 ( x ) = 10 -( θ ⊤ x -3) 2 and f 2 ( x ) = 5 -0 . 5 ∥ x -b ∥ . Here, θ and b are randomly generated and fixed afterward. Since our data assumes a low-dimensional subspace representation, it is convenient to decompose θ into θ ⊥ = ( I -AA ⊤ ) θ

2

2

L

λ

µ

K

-

x

∗

A,λ

∥

∥

2

-

1

.

and θ ∥ = AA ⊤ θ , representing the off-support and on-support components. We refer to ∥ θ ⊥ ∥ ∥ θ ∥ ∥ as the off/on-support ratio. Analogously, for a generated sample, we can also define its off/on-support ratio. Clearly, a small off/on-support ratio indicates close vicinity to the subspace.

Score Network Pre-training. We utilize a version of the U-Net [54], with 14.8M trainable parameters. Note that this is a complicated network going beyond the linear score function class considered in our theories. Following the implementation of Denoising Diffusion Probabilistic Models (DDPM, Ho et al. [34]), we train the U-Net o estimate the score function ∇ log p t , via minimizing the score matching loss introduced in Eqn. (3). We discretize the backward process to have 200 time steps as in Nichol and Dhariwal [47], and the U-Net is trained using our generated data set for 20 epochs. We use Adam as the optimizer, set the batch size as 32 , and set the learning rate to be 10 -4 . After the pre-training phase, we confirmed that the data subspace structure is well learned, as the generated samples using the pre-trained diffusion model have an average off/on-support ratio of 0 . 039 .

Implementation of Algorithm 1. In each iteration of Algorithm 1, we need to compute the gradient guidance G loss . We set the targeted y value at the k -th iteration as y k = δ + g ⊤ k z k , where δ specifies the increment per iteration. The choice on δ is instance-dependent and we set it via tuning for near-optimal in different experiments. For comparing naive gradient with gradient guidance in Figure 8, we set δ = 0 . 2 and 0 . 9 , respectively for using naive gradient G and gradient guidance G loss . In Figure 9, we choose δ to be (a) 0 . 05 , (b) 0 . 2 , (c) 1 , and (d) 1 , corresponding to each panel. We initialize Algorithm 1 with a batch of 32 samples generated by the pre-trained model. Each sample determines an optimization trajectory. We repeat Algorithm 1 for 5 times with different random seeds and report the error bars.

Implementation of Algorithm 2. Algorithm 2 differs from Algorithm 1 in that it allows additional fine-tuning of the pre-trained score network. In practice, to update the score network incorporating newly generated data, one does not have to exactly solve (14) by re-training the full model from scratch. Instead, (14) can be viewed as a guideline that motivates more computationally efficient ways for updating the pre-trained score. It is a common practice to only fine-tune the weights of the old model by performing gradient descent over a few batches of newly generated data, which is similar to the spirit of (14). To be more specific, we adopt a computationally lightweight fine-tuning strategy: We only perform one Adam optimization step using the re-weighted loss given by Eqn. (14) with a batch of 32 generated samples. We set the learning rate as 10 -6 . This simple strategy already demonstrates good performances as shown in Figure 10. Other implementation details are kept the same as those of Algorithm 1.

F.1.2 Results

We first demonstrate our gradient guidance G loss preserves the subspace structure learned from the pre-trained model. For comparison, we also tested the naive guidance G defined following Proposition 3 (with Σ = I ). For a quick reference, we repeat the definition here:

G ( x t , t ) := β ( t ) ( y -g ⊤ E [ x 0 | x t ] ) g, where β ( t ) > 0 and y ∈ R are tuning parameters, and E [ x 0 | x t ] is the conditional expectation of x 0 given noise corrupted data x t . For implementation, we replace E [ x 0 | x t ] by its look-ahead estimator ˆ E [ x 0 | x t ] based on the Tweedie's formular.

Comparing G and G loss on Preserving Subspace Structure. Figure 8 (a), (c) verify that the naive gradient G performs much worse than G loss in preserving the linear subspace structure. It is consistent with our theoretical finding that the gradient guidance G loss keeps the generated sample close to the latent subspace, with substantially smaller off-support errors. When allowing adaptive score fine-tuning in Algorithm 2, Figure 8 (b), (d) show that the off-support error increases as the model gets fine-tuned using self-generated data, due to increasing distribution shift. Even in this case, the naive gradient G leads to much more severe off-support errors as compared to G loss .

Algorithm 1 Converges to Regularized Optima. We plot the convergence of Algorithm 1 in terms of the objective value in Figure 9. Figure 9 (a),(b) are for the objective function f 1 = 10 -( θ ⊤ x -3) 2 as the objective function, while Figure 9(c),(d) are for the objective f 2 = 5 -0 . 5 ∥ x -b ∥ . We observe

(a) Algorithm 1

(b) Algorithm 2

(c) 300-350 round of (a)

(d) 1000-1200 round of (b)

1.1

1.1

1.0

0.9

Gloss

G

1.0

0.9

Gloss

G

on support space

Giass

Gies

on support space

G

0.8

0.7

Pretrained

0.8

0.7

Pretrained

3.0

0.6

0.5

off

0.5

0.6

support direction

0.4

0.4

2.0

0.3

0.3

1.5

0.2

0.2 -

0.1

0.1

1.3

0.0

001

200

300

400

500

600

700

800

900

1000

0.0

OOT

200

300

400

500

600

7D0

800

900 1000 1100 1200

that the algorithm converges to reach some sub-optimal objective value, but there remains a gap to the maximal value. This is consistent with our theory that the pre-trained model essentially acts as a regularization in addition to the objective function. Adding gradient guidance alone cannot reach global maxima. This coincides with our theoretical findings in Theorem 2.

(a)

θ

=

Aβ

∗

(b)

∥

θ

⊥

∥

∥

θ

∥

= 9

(c)

b

= 4

·

1

D

(d)

b

∼ N

(4

·

1

D

,

9

·

I

D

)

10

9

8

7

6

10

6

Gloss

fA

7.0

7.5

0'8=

S'8-

7.0

0°9=

7.5

8.5

Gloss

f

f(x)

5

4

f(x)

S’6=

0°6-

S°6-

0°6

10.0

10.0

=1

1

fA

Gloss

1

-10.5

-11.0

11.5

Gloss

fA

10.5

11.5

0

100

200

300

400

500600

700

800

006

1000

0

100

200

300

400

500

600

700

B00

900

1000

-12.0

10

15

20

25

30

35

40

45

50

12.0

5

10

15

20

25

30

35

40

45

50

∥

Algorithm 2 Converges to Global Optima. Algorithm 2 converges to the maximal value of the objective function f 1 = 10 -( θ ⊤ x -3) 2 as shown in Figure 10(a). In Figure 10(b), we visualize the distribution of generated samples of Algorithm 1 (blue) and 2 (red), respectively, as the iteration evolves. We see that samples from Algorithm 1 mostly stay close to the pre-training data distribution (area described by the dotted contour). In constrast, samples of Algorithm 2 move outside the contour, as the diffusion model gets fine-tuned using self-generated data.

Results for Nonlinear Data Structure. We apply Algorithm 1 to data uniformly sampled from a unit ball in R 64 . The objective reward function is defined as f ( x ) = θ ⊤ x , where ∥ θ ∥ = 1 . The left panel of Figure 11 demonstrates that rewards increase and converge when using Algorithm 1. Higher guidance strength ∆ (corresponding to lower regularization) results in a higher convergent reward. The right panel of Figure 11 shows that, for the same reward level, gradient guidance achieves a smaller deviation from the unit ball compared to the naive gradient approach. This suggests that gradient guidance can better preserve data structure for nonlinear manifolds.

F.2 Image Generation

Hyperparameters. Since tuning parameters y and β ( t ) are both for the strength of guidance, we can fix one of them. We choose β ( t ) = 100 as suggested by [71], and set a series of y ∈ { 2 , 4 , 10 } . We run batch size = 20 samples parallelly across iterations to evaluate the rewards. The prompt is uniformly sampled from the 1000 classes of ImageNet [21].

F.3 Time Efficiency

We summarize the time cost of our experiments on one NVIDIA A100 GPU in Table 1.

(a) Convergence of Algorithm 2

(b) Distribution of generated samples

5o Alg. 2

510e Alg. 2

S2se Alg. 2

So Alg. 1

51oe Alg. 1

S2oe Alg. 1

10

9

Gloss

8

fA

5400 Alg. 2

5400 Alg. 1

Ssoo Alg. 1

100

200

300400500600

700

800900100011001200

roundk

0.12

1.0

naive gradient

0.10

pretrain

gradient guidance

0.8

0.08

0.6

0.06

oddns

0.4

0.04

告

0.02

0.2

f*

= 0.5

0.00

△ = 0.3

0.0

= 0.2

0.02

0

10

15

20

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

puno

reward

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The introduction (Section 1) includes a paragraph 'Scope and Contribution", together with the abstract, clearly reflecting the paper's contributions and scope.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitations in Appendix B.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All the theoretical results Lemma 1, Theorem 1, Theorem 2 and Theorem 3 provide the full set of assumptions, and the correspond proofs are in Appendix D.4, Appendix D.5, Appendix E.2 and Appendix E.3, respectively.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The main information of the experiments are described in Section 7, and the details are provided in Appendix F.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in

some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The codebase is released at https://github.com/yukang123/GGDMOptim.git.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The experimental details are provided in Appendix F.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: The experiments that support the main claims of the paper provide error bars, as shown in Figure 4, Figure 5 and Figure 6.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: All information on the computer resources is included in Appendix F.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer:[Yes]

Justification: The research adheres to all guidelines and principles outlined in the NeurIPS Code of Ethics, including considerations for fairness, accountability, and transparency. Ethical practices were maintained throughout the study, ensuring compliance with all applicable standards.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss the social impacts of this paper in Appendix B.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: All creators and original owners of assets used in this paper are properly credited. The licenses and terms of use are explicitly mentioned and respected according to their respective guidelines.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We release our code and models along with documentation at https://github.com/yukang123/GGDMOptim.git.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing or research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
This paper investigate gradient guidance for adapting or fine-tuning pre-trained diffusion models from an optimization perspective.
The author proposed a look-ahead loss based gradient guidance and two variants of diffusion-based generative optimization algorithms utilizing it.
The author provided theoretical guarantees for adapting/fine-tuning diffusion models.

STRENGTHS:
Overall, although this is a theory paper, it is very impressive. It used theoretical approaches to explain some problems of diffusion with guidance:
1. naive guidance doesn't work on latent structure of data
2.  algorithm on guided diffusion using gradient queries on new samples
3. proposed an adaptive gradient guided diffusion, where both pre-trained score network and guidance are iteratively updated using self-generated samples.

WEAKNESSES:
No obvious weakness.
If possible, please include the reason for each step in your proof, i.e. the first step follows from???, the second step follows from???. Then it will be a perfect paper. I will increase the score if you can fix this issue.

QUESTIONS:
No questions

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-08

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new approach to the problem of gradient-guided generation for diffusion models. The main challenge of gradient-guided generation is to maintain the generated sample within the support of the sample distribution. To address this issue, this work starts with a simplified model, using a linear objective function and a low-dimensional linear sample space, and provides a closed form for the gradient guidance. Under this simplified model, the paper offers a theoretical analysis to show the converged distribution of the generated samples for concave and smooth objective functions. Lastly, it also proposes another variant that allows fine-tuning over pretrained models, further improving optimization performance\n\nSTRENGTHS:\n1. The idea of deriving this form of gradient guidance is interesting and novel, particularly the use of a low-dimensional linear sample space to demonstrate the problem of out-of-distribution in gradient-guidance sampling.\n\n2. The theoretical results are solid, and I appreciate the finding that the proposed method can indeed preserve the low-dimensional linear structure theoretically.\n\nWEAKNESSES:\n1. Although I appreciate the motivation and analysis in this work, I do not think this method is practical in reality. Firstly, the core advantage of gradient-guidance approaches like [1][2] is that they do not require gradient backpropagation over the neural network, thus adding only minor computational cost to direct sampling. In contrast, the method proposed in this work requires gradient backpropagation over the neural network, making it prohibitively slow. As mentioned in the paragraph above Algorithm 1, guiding stable diffusion requires 76 minutes for optimization. Secondly, since the method requires gradient backpropagation over the neural network, I think the author should compare it to the serial works on direct optimization in [3][4][5], which are completely ignored in this work. These works target the same task and have the same computational requirements. From my understanding, the algorithm proposed in this work cannot rival direct optimization methods, as the latter do not rely on simplified models to derive the algorithm. In fact, a concurrent work [6] shows that one can optimize SDXL, a 3B parameter model, within 10 minutes, which is in sharp contrast to the performance of Algorithm 1 proposed in this work.\n\n2. The theoretical analysis is certainly novel, but it heavily relies on the linear sample space assumptions. It is unclear whether this method can be applied to cases where the sample space is highly nonlinear. The author conducted a simple experiment with image diffusion, which somewhat investigates this point, but I think the author should also conduct experiments on a highly nonlinear synthetic data distribution to validate if their method works in this scenario.\n\n3. The writing and presentation of the paper can be improved. For example, when moving from Section 3 to Section 4, the target objective function switches from a linear function to a general function without much explanation. I can only guess from Algorithm 1 that the gradient \n$g$ is directly replaced by some stochastic gradient estimator of a general function $f$? I think the author should provide some explanation on this point.\n\n[1] Song, Jiaming, et al. \"Loss-guided diffusion models for plug-and-play controllable generation.\" International Conference on Machine Learning. PMLR, 2023.\n\n[2] Chung, Hyungjin, et al. \"Diffusion posterior sampling for general noisy inverse problems.\" arXiv preprint arXiv:2209.14687 (2022).\n\n[3] Bram Wallace, Akash Gokul, Stefano Ermon, and Nikhil Naik. End-to-end diffusion latent optimization improves classifier guidance. In Proceedings of the IEEE/CVF International Conference on\nComputer Vision, pages 7280–7290, 2023b.\n\n[4] Heli Ben-Hamu, Omri Puny, Itai Gat, Brian Karrer, Uriel Singer, and Yaron Lipman. D-flow:\nDifferentiating through flows for controlled generation. arXiv preprint arXiv:2402.14017, 2024.\n\n[5] Korrawe Karunratanakul, Konpat Preechakul, Emre Aksan, Thabo Beeler, Supasorn Suwajanakorn,\nand Siyu Tang. Optimizing diffusion noise can serve as universal motion priors. arXiv preprint\narXiv:2312.11994, 2023.\n\n[6] Tang, Zhiwei, et al. \"Tuning-Free Alignment of Diffusion Models with Direct Noise Optimization.\" arXiv preprint arXiv:2405.18881 (2024).\n\nQUESTIONS:\n1. I wonder how the algorithm proposed in this work compares to the direct optimization approaches in [3][4][5][6] (references mentioned above)?\n\n2. Could the author conduct some simple experiments on synthetic data with a nonlinear sample space? I am curious how well the proposed method will perform when the linear assumption is violated.\n\nLIMITATIONS:\nSee my comments above. In summary, the main limitations of this work are:\n\n1. Incomplete literature review and comparison with existing works.\n\n2. The performance of the proposed algorithm appears to be impractical.\n\n3. The presentation can be improved.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nUnder the assumption of the data belonging to a low-dimensional linear subspace, the authors investigate two common gradient-based guidance techniques of diffusion models, encouraging the use of one of them (computing the gradient at the estimate of x0 given xt, as done in many works).  \nFor concave reward, and under the additional assumption of linear score function, they also establish convergence results for the mean of the guided process.\n\nSTRENGTHS:\nOverall providing theory for guided diffusion is appreciated.\n\nThe paper formally shows an advantage of a certain guidance over the other.\n\nThe presentation is clear and the paper is convenient to read.\n\nWEAKNESSES:\nThe assumptions made in the paper are strong (as discussed below).\n\nA certain weakness of Theorems 2 and 3 is that they consider only convergence of the mean of the procedure do not say anything about the variance.\n\nThere are no new algorithmic ideas in the paper, as the promoted gradient-based guidance is known and common, and also there exist works on fine-tuned/adaptive diffusion models. Thus, the authors needs to tone down some claims in the contribution list and in the abstract.\n\nMore comments are stated below.\n\nQUESTIONS:\n1.  \n\nThe tight connection between gradient-guided diffusion models and proximal optimization algorithms is discussed in detail in:  \nTomer Garber and Tom Tirer, \"Image Restoration by Denoising Diffusion Models with Iteratively Preconditioned Guidance,\" CVPR, 2024.  \nThis relevant literature should be mentioned.\n\n2.  \n\nThe statement in line 48:  \n\"Why does naively guiding diffusion models using gradient never work in practice?\"  \nIs not clear, as there are leading approaches with gradient-based guidance that work, such as (Garber and Tirer, 2024) and the back-projections [53] / least-squares [14] methods that it generalizes.   \nOnly much later in the paper, the reader can see that the authors do not refer to **the common practice** of computing the gradient at the estimate of x0 given xt as a plain gradient-based guidance. This should be stated clearly already at early stage.  \nThis also means that you need to tone down the claim in the contribution list that you: \"introduce guidance based on forward prediction loss\", because this is already a well-known practice.  \n\n3.\n\nBased on your Eq. 1, it seems that you consider the variance preserving (VP) SDE formulation.  \nWhy do you state in line 85 that without loss of generality all your analysis is done for q=1?  \nThis needs to be explained, at least in the appendix, as q=1 is not the common setting (but rather some linearly increase from near 0 to 1).\n\n4.  \n\nThere seems to be some discrepancy between your Eq. 2, associated with Eq. 1, and the SDE formulation of diffusion models [50]. Recall that dt in Eq. 2 is negative.\n\n5.  \n\nThe \"general\" explanation in Section 3.2 that naive gradient-based guidance doesn't work is not convincing because it ignores that fact that the step-size decreases as the t get closer to t=0 and that the noise injection can mitigate error propagation.\n\n6.  \n\nYou state: \"Alg. 1 is light-weighted ... takes 76min overall.\" Recall that the reverse process is not performed in the offline pretraining stage. So obviously Alg. 1 is quite slow. Indeed, it is known that guidance methods that include the DNN's Jacobian computation are slow.\n\n7.  \n\nThe assumptions of the theoretical analyses in the paper are quite strong.  \nEspecially Assumption 1 on the signal belonging to low dimensional linear subspace and the assumption on linear score model (Eq. 12).  \nAlready in the contribution list you should state the assumption on the score model, which essentially leads to convex optimization.\n\n8.  \n\nAny idea how can Theorem 1 be generalized to low dimensional manifold rather than low dimensional linear subspace?  \nUnder the assumptions, in what aspects Theorem 2 differs from standard convex analysis? (maximization of an L-smooth concave function?). \nAren't there any other works that study convergence of guided diffusion models except Marion et al. [38]? There are several works on convergence of diffusion models, e.g., (De Bortoli, 2022). What prevents extending them to the guided case?\n\nValentin De Bortoli, \"Convergence of denoising diffusion models under the manifold hypothesis,\" TMLR, 2022. \n\n9.  \n\nUnder the subspace assumption 1, your data covariance matrix is a DxD matrix of rank d<D.  \nTherefore it is not invertible. I suggest using the dagger symbol to avoid confusion between inverse and pseudoinverse. Provide more details in the proof to clarify that this issue is taken into account.  \nNote that there are some peculiar differences between inverse and pseudoinverse, e.g., generally, pinv(AB) is not equal to pinv(B)pinv(A).\n\n10.  \n\nRegarding Section 5, you need to clearly state what are the differences of Alg 2 from many other works that fine-tune pretrained diffusion models (just google \"adaptive diffusion\" and \"diffusion personalization\").\n\n11.\n\nWhat is the motivation for convergence to the **global** optimum of the guidance objective if one pays in ignoring/omitting the prior information / regularization of the score model? (As seems to be promoted in Sections 5 and 6).  \nThe tradeoff between reward/loss and regularization is trivial and clearly in most case there is some delicate balance between them that a user should look for.\n\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper rethinks the gradient-based guidance methods through the optimization perspective.  Similar to the manifold assumption, the authors first rely on the assumption that the observed data is from the lower dimensional space. Then, they claim that the naive gradient guidance does not maintain the data subspace, and show how to preserve the data subspace and incorporate guidance by using the look-ahead formulation. For the proof, they rely on the Gaussian linear model. From such motivations from theoretical results, they propose two versions of gradient guidance methods: 1) Only update the gradient guidance components and 2) Simultaneously finetuning the score model. These formulations enable the convergence to the data subspace, unlike naive guidance formulation. Experimental results are conducted on toy examples and image generation settings.\n\nSTRENGTHS:\n- Strong theories are presented in this paper. Rethinking the guidance formulation as the optimization procedure, that aims to optimize 1) guidance and 2) convergence of the data subspace, is interesting.\n- The paper is well-written and easy to follow.\n\nWEAKNESSES:\n- While the proposed algorithm demonstrates considerable theoretical value, its practical implementation appears to be quite slow. This could be perceived as a weakness of the paper. \n- The paper discusses Universal Guidance, yet it omits mention of similar training-free [B, C, D, E, F, G] and training-required guidance [A] mechanisms for utilizing off-the-shelf models. Including a discussion on these mechanisms would enhance the comprehensiveness of the paper. Additionally, it would be beneficial to address the practical and theoretical advantages compared to these mechanisms. As it stands, the paper does not sufficiently clarify these comparative advantages, which leaves some ambiguity about the relative strengths compared to these methods.\n\nQUESTIONS:\n- The paper relies heavily on the linear assumption. How reasonable is this assumption in the context of the proposed algorithm? Aren't guidance classifiers and score functions typically more complex than linear functions?\n\n## **Reference**\n- [A] Towards Practical Plug-and-Play Diffusion Models, CVPR 2023.\n\n- [B] ELUCIDATING THE DESIGN SPACE OF CLASSIFIER GUIDED DIFFUSION GENERATION, ICLR 2024\n\n- [C] ADJOINTDPM: ADJOINT SENSITIVITY METHOD FOR GRADIENT BACKPROPAGATION OF DIFFUSION PROBABILISTIC MODELS, ICLR 2024.\n\n- [D] Towards Accurate Guided Diffusion Sampling through Symplectic Adjoint Method, Arxiv 2023.\n\n- [E] Loss-guided Diffusion Models for Plug-and-Play Controllable Generation, ICML 2023\n\n- [F] Manifold preserving guided diffusion, Arxiv 2023\n\n- [G] Freedom: Training-free energy guided conditional diffusion model, ICCV 2023.\n\nLIMITATIONS:\nLinear approximation of the score function in their proof is mentioned as the limitation.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}