Gradient Guidance for Diffusion Models: An Optimization Perspective

Authors: yingqing  guo
hui  yuan
yukang  yang
minshuo  chen
mengdi  wang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 06:15:49

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.

Main paper: 
 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 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 Õ(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. 
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 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.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. 
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., dX t = -1 2 q(t)X t dt + q(t) dW t for q(t) > 0,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) ,EQUATIONHere 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:EQUATIONwhere S is a given function class, E D denotes the empirical expectation over training data D and E xt|x0 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 / (2h(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 :dX ← t = 1 2 X ← t + s θ (X ← t , T -t) + G (X ← t , T -t) dt + dW 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). Σ := 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 = ∂fi(x) ∂xj .Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?• 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. 
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):dX ← t = 1 2 X ← t + s θ (X ← t , T -t)+G(X ← t , t) dt + dW t .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] .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.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.,EQUATIONHowever, the naive gradient guidance (5) would jeopardize the latent structure of data, which is demonstrated by the following proposition: 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.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 .[63] also observed this, explaining that ∇f is not computed for t = T i.e., it is not aligned with the clean data space. = " > A A A B / H i c b V B L S g N B F H w T f z H + o i 7 d N A b B V Z g R i S 6 D b l w m Y D 6 Q D K G n p 5 M 0 6 f n Q / U Y I Q 7 y A W 7 2 B O 3 H r X b y A 5 7 A n m Y V J L G g o q t 7 j V Z c X S 6 H R t r + t w s b m 1 v Z O c b e 0 t 3 9 w e F Q + P m n r K F G M t 1 g k I 9 X 1 q O Z S h L y F A i X v x o r T w J O 8 4 0 3 u M 7 / z x J U W U f i I 0 5 i 7 A R 2 F Y i g Y R S M 1 c V C u 2 F V 7 D r J O n J x U I E d j U P 7 p + x F L A h 4 i k 1 T r n m P H 6 K Z U o W C S z 0 r 9 R P O Y s g k d 8 Z 6 h I Q 2 4 d t N 5 0 B m 5 M I p P h p E y L 0 Q y V / 9 u p D T Q e h p 4 Z j K g O N a r X i b + 5 / U S H N 6 6 q Q j j B H n I F o e G i S Q Y k e z X x B e K M 5 R T Q y h T w m Q l b E w V Z W i 6 W b r i 6 y z a r G S K c V Z r W C f t q 6 p T q 9 a a 1 5 X 6 X V 5 R E c 7 g H C 7 B g R u o w w M 0 o A U M O L z A K 7 x Z z 9 a 7 9 W F 9 L k Y L V r 5 z C k u w v n 4 B B Q q V g A = = < / l a tFailure 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. EQUATIONwhere 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 asEQUATIONwhere β(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., dX t = -1 2 q(t)X t dt + q(t) dW t .The formula in (7) generalizes the intuition of a conditional score for any data distribution and objective function. It scales with the residual term yg ⊤ 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 (yg ⊤ E[x 0 |x t ]) 2 resembles the proximal term commonly used in first-order proximal optimization methods.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 thatEQUATIONProof Sketch. We have∇ xt y -g ⊤ E[x 0 | x t ] 2 ∝ ∇ xt E[x 0 | x t ] ⊤ g.We will show that the Jacobian ∇ xt 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 haveEQUATIONhere m(u) = ∇ log p LD t (u) + h -1 (t)u, p LD t (u) latent density (Appendix D.1). We see∇ xt 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.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] :EQUATIONand 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 byEQUATION< l a t e x i t s h a 1 _ b a s e 6 4 = " 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 lossM C f C y l l m U a v / 6 O d N R y R N f C V S f v 8 = " > A A A B / 3 i c b V D L S s N A F J 3 4 r P V V d e l m s A h 1 U x L x t S y 6 c V n B P q A N Z T K Z N E M n k z B z I 5 T Q h T / g V v / A n b j 1 U / w B v 8 N J m 4 V t P T B w O O d e 7 p n j J Y J r s O 1 v a 2 V 1 b X 1 j s 7 R V 3 t 7 Z 3 d u v H B y 2 d Z w q y l o 0 F r H q e k Q z w S V r A Q f B u o l i J P I E 6 3 i j u 9 z v P D G l e S w f Y Z w w N y J D y Q N O C e R S W I O z Q a V q 1 + 0 p 8 D J x C l J F B Z q D y k / f j 2 k a M Q l U E K 1 7 j p 2 A m x E F n A o 2 K f d T z R J C R 2 T I e o Z K E j H t Z t O s E 3 x q F B 8 H s T J P A p 6 q f z c y E m k 9 j j w z G R E I 9 a K X i / 9 5 v R S C G z f j M k m B S T o 7 F K Q C Q 4 z z j 2 O f K 0 Z B j A 0 h V H G T F d O Q K E L B 1 D N 3 x d d 5 t E n Z F O M s 1 r B M 2 u d 1 5 6 p + + X B R b d w W F Z X Q M T p B N e S g a 9 R A 9 6 i J W o i i E L 2 g V / R m P V v v 1 o f 1 O R t d s Y q d I z Q H 6 + s X n I m W V g = = < / l a t e x i t > h(t) < l a t e x i t s h a 1 _ b a s e 6 4 = " O g x o w 3 Z r p h f i c 4 L 4 L Q Y h m B z 4 Q a U = " > A A A C C X i c b V D L S s N A F J 3 U V 6 2 v q k s 3 g 0 W o C 0 s i v p Z F N y 4 r 2 A e 0 a Z l M J u 3 Q y Y O Z G 6 G E f I E / 4 F b / w J 2 4 9 S v 8 A b / D S Z u F b T 1 w 4 X D O v d z D c S L B F Z j m t 1 F Y W V 1 b 3 y h u l r a 2 d 3 b 3 y v s H L R X G k r I m D U U o O w 5 R T P C A N Y G D Y J 1 I M u I 7 g r W d 8 V 3 m t 5 + Y V D w M H m E S M d s n w 4 B 7 n B L Q U r 9 H R D Q i / e T M S q t w O i h X z J o 5 B V 4 m V k 4 q K E d j U P 7 p u S G N f R Y A F U S p r m V G Y C d E A q e C p a V e r F h E 6 J g M W V f T g P h M 2 c k 0 d Y p P t O J i L 5 R 6 A s B T 9 e 9 F Q n y l J r 6 j N 3 0 C I 7 X o Z e J / X j c G 7 8 Z O e B D F w A I 6 e + T F A k O I s w q w y y W j I C a a E C q 5 z o r p i E h C Q R c 1 9 8 V V W b S 0 p I u x F m t Y J q 3 z m n V V u 3 y 4 q N R v 8 4 q K 6 A g d o y q y 0 D W q o 3 v U Q E 1 E k U Q v 6 B W 9 G c / G u / F h f M 5 W C 0 Z + c 4 j m Y H z 9 A g z b m m g = < / l a t e x i t > ↵ 1 (t) < l a t e x i t s h a 1 _ b a s e 6 4 = " M X p + D K i S l Y D t 6 9 E s 0 R l l C 1 6 j z L c = " > A A A B / n i c b V D L S s N A F L 3 x W e u r 6 t J N s A i u S i K + l k U 3 L i v a B 7 S h T C a T d u h k E m Z u x F I K / o B b / Q N 3 4 t Z f 8 Q f 8 D i d t F r b 1 w M D h n H u 5 Z 4 6 f C K 7 R c b 6 t p e W V 1 b X 1 w k Z x c 2 t 7 Z 7 e 0 t 9 / Q c a o o q 9 N Y x K r l E 8 0 E l 6 y O H A V r J Y q R y B e s 6 Q 9 u M r / 5 y J T m s X z A Y c K 8 i P Q k D z k l a K T 7 p y 5 2 S 2 W n 4 k x g L x I 3 J 2 X I U e u W f j p B T N O I S a S C a N 1 2 n Q S 9 E V H I q W D j Y i f V L C F 0 Q H q s b a g k E d P e a B J 1 b B 8 b J b D D W J k n 0 Z 6 o f z d G J N J 6 G P l m M i L Y 1 / N e J v 7 n t V M M r 7 w R l 0 m K T N L p o T A V N s Z 2 9 m 8 7 4 I p R F E N D C F X c Z L V p n y h C 0 b Q z c y X Q W b R x 0 R T j z t e w S B q n F f e i c n 5 3 V q 5 e 5 x U V 4 B C O 4 A R c u I Q q 3 E I N 6 k C h B y / w C m / W s / V u f V i f 0 9 E l K 9 8 5 g B l Y X 7 + m s J Z q < / l a t e x i t > xt Gradient Compute < l a t e x i t s h a 1 _ b a s e 6 4 = " q E X a E S V r 1 f T B b n N b r K v E P v C + E 8 A = " > A A A C G H i c b V D L S s N A F J 3 4 r P V V d d n N Y B F c l U S k u i y K 4 L K C f U A a w m Q y b Y d O H s z c S E v M w t / w B 9 z q H 7 g T t + 7 8 A b / D S d u F b T 0 w c D j n X u 6 Z 4 8 W C K z D N b 2 N l d W 1 9 Y 7 O w V d z e 2 d 3 b L x 0 c t l S U S M q a N B K R 7 H h E M c F D 1 g Q O g n V i y U j g C d b 2 h t e 5 3 3 5 g U v E o v I d x z J y A 9 E P e 4 5 S A l t x S u T s g k H Y D A g P P S 2 + y z B 6 5 J n 7 E I x c c t 1 Q x q + Y E e J l Y M 1 J B M z T c 0 k / X j 2 g S s B C o I E r Z l h m D k x I J n A q W F b u J Y j G h Q 9 J n t q Y h C Z h y 0 s k n M n y i F R / 3 I q l f C H i i / t 1 I S a D U O P D 0 Z J 5 W L X q 5 + J 9 n J 9 C 7 d F I e x g m w k E 4 P 9 R K B I c J 5 I 9 j n k l E Q Y 0 0 I l V x n x X R A J K G g e 5 u 7 4 q s 8 W l b U x V i L N S y T 1 l n V q l V r d + e V + t W s o g I q o 2 N 0 i i x 0 g e r o F j V Q E 1 H 0 h F 7 Q K 3 o z n o 1 3 4 8 P 4 n I 6 u G L O d I z Q H 4 + s X n g a g n w = = < / l a t e x i t > Ê[x0|xt] Gradient w.r.t < l a t e x i t s h a 1 _ b a s e 6 4 = " M X p + D K i S l Y D t 6 9 E s 0 R l l C 1 6 j z L c = " > A A A B / n i c b V D L S s N A F L 3 x W e u r 6 t J N s A i u S i K + l k U 3 L i v a B 7 S h T C a T d u h k E m Z u x F I K / o B b / Q N 3 4 t Z f 8 Q f 8 D i d t F r b 1 w M D h n H u 5 Z 4 6 f C K 7 R c b 6 t p e W V 1 b X 1 w k Z x c 2 t 7 Z 7 e 0 t 9 / Q c a o o q 9 N Y x K r l E 8 0 E l 6 y O H A V r J Y q R y B e s 6 Q 9 u M r / 5 y J T m s X z A Y c K 8 i P Q k D z k l a K T 7 p y 5 2 S 2 W n 4 k x g L x I 3 J 2 X I U e u W f j p B T N O I S a S C a N 1 2 n Q S 9 E V H I q W D j Y i f V L C F 0 Q H q s b a g k E d P e a B J 1 b B 8 b J b D D W J k n 0 Z 6 o f z d G J N J 6 G P l m M i L Y 1 / N e J v 7 n t V M M r 7 w R l 0 m K T N L p o T A V N s Z 2 9 m 8 7 4 I p R F E N D C F X c Z L V p n y h C 0 b Q z c y X Q W b R x 0 R T j z t e w S B q n F f e i c n 5 3 V q 5 e 5 x U V 4 B C O 4 A R c u I Q q 3 E I N 6 k C h B y / w C m / W s / V u f V i f 0 9 E l K 9 8 5 g B l Y X 7 + m s J Z q < / l a t e x i t > xt + < l a t e x i t s h a 1 _ b a s e 6 4 = " J j N j o y y 0 8 E 2 S T a h Q + 0 G X Q t F m l q w = " > A A A C D n i c b V D L S s N A F J 3 4 r P W V 6 t J N s A g V p C T i a 1 l 0 4 7 K C f U A b w m Q y a Y d O H s z c q C X k H / w B t / o H 7 s S t v + A P + B 1 O 2 i x s 6 4 E L h 3 P u 5 R 6 O G 3 M m w T S / t a X l l d W 1 9 d J G e X N r e 2 d X r + y 1 Z Z Q I Q l s k 4 p H o u l h S z k L a A g a c d m N B c e B y 2 n F H N 7 n f e a B C s i i 8 h 3 F M 7 Q A P Q u Y z g k F J j l 6 R T t q H I Q W c 1 Z 4 c O I F j R 6 + a d X M C Y 5 F Y B a m i A k 1 H / + l 7 E U k C G g L h W M q e Z c Z g p 1 g A I 5 x m 5 X 4 i a Y z J C A 9 o T 9 E Q B 1 T a 6 S R 6 Z h w p x T P 8 S K g J w Z i o f y 9 S H E g 5 D l y 1 G W A Y y n k v F / / z e g n 4 V 3 b K w j g B G p L p I z / h B k R G 3 o P h M U E J 8 L E i m A i m s h p k i A U m o N q a + e L J P F p W V s V Y 8 z U s k v Z p 3 b q o n 9 + d V R v X R U U l d I A O U Q 1 Z 6 B I 1 0 C 1 q o h Y i 6 B G 9 o F f 0 p j 1 r 7 9 q H 9 j l d X d K K m 3 0 0 A + 3 r F 5 4 P n F Q = < / l a t e x i t > s✓(xt, t) < l a t e x i t s h a 1 _ b a s e 6 4 = " 8 R b S 9 9 2 U i r R 9 7 / / c n a 4 G H 7 i d u + Q = " > A A A B / H i c b V B L S g N B F H w T f z H + o i 7 d N A b B V Z g R i S 6 D b l w m Y D 6 Q D K G n 5 0 3 S p O d D d 4 8 Q Q r y A W 7 2 B O 3 H r X b y A 5 7 A n m Y V J L G g o q t 7 j V Z e X C K 6 0 b X 9 b h Y 3 N r e 2 d 4 m 5 p b / / g 8 K h 8 f N J W c S o Z t l g s Y t n 1 q E L B I 2 x p r g V 2 E 4 k 0 9 A R 2 v P F 9 5 n e e U C o e R 4 9 6 k q A b 0 m H E A 8 6 o N l J z O C h X 7 K o 9 B 1 k n T k 4 q k K M x K P / 0 / Z i l I U a a C a p U z 7 E T 7 U 6 p 1 J w J n J X 6 q c K E s j E d Y s / Q i I a o 3 O k 8 6 I x c G M U n Q S z N i z S Z q 3 8 3 p j R U a h J 6 Z j K k e q R W v U z 8 z + u l O r h 1 p z x K U o 0 R W x w K U k F 0 T L J f E 5 9 L Z F p M D K F M c p O V s B G V l G n T z d I V X 2 X R Z i V T j L N a w z p p X 1 W d W r X W v K 7 U 7 / K K i n A G 5 3 A J D t x A H R 6 g A S 1 g g P A C r / B m P V v v 1 o f 1 u R g t W P n O K S z B + v o F 8 E O V c w = = < / l a t e x i t > g < l a t e x i t s h a 1 _ b a s e 6 4 = " O 0 N x y O D e 2 a a Q y e B Y z g s n x e d h T h E = " > A A A C J 3 i c b V D L S s N A F J 3 U V 6 2 v q E s 3 g 0 V Q w Z K I r 2 V R B J c K V g t N D J P p t B 0 6 e T B z I 4 a Y f / A 3 / A G 3 + g f u R J d u / A 4 n t Q t b P T B w O O d e 7 p n j x 4 I r s K w P o z Q x O T U 9 U 5 6 t z M 0 v L C 6 Z y y t X K k o k Z Q 0 a i U g 2 f a K Y 4 C F r A A f B m r F k J P A F u / b 7 J 4 V / f c u k 4 l F 4 C W n M 3 I B 0 Q 9 7 h l I C W P H N 7 M 8 U 7 u H v j Q B R j p 0 c g c w I C P d / P T v O 8 d e d Z + B 7 f e e B u 3 e x 6 Z t W q W Q P g v 8 Q e k i o a 4 t w z v 5 x 2 R J O A h U A F U a p l W z G 4 G Z H A q W B 5 x U k U i w n t k y 5 r a R q S g C k 3 G / w p x x t a a e N O J P U L A Q / U 3 x s Z C Z R K A 1 9 P F o H V u F e I / 3 m t B D p H b s b D O A E W 0 p 9 D n U R g i H B R E G 5 z y S i I V B N C J d d Z M e 0 R S S j o G k e u t F U R L a / o Y u z x G v 6 S q 9 2 a f V D b v 9 i r 1 o + H F Z X R G l p H m 8 h G h 6 i O z t A 5 a i C K H t A T e k Y v x q P x a r w Z 7 z + j J W O 4 s 4 p G Y H x + A 1 6 D p Y 8 = < / l a t e x i t > (y g > Ê[x0|xt]) 2y -g ⊤ Ê[x 0 |x t ] 2via a forward pass of s θ and a backward pass using the auto-gradient feature of deep-leaning frameworks such as PyTorch and TensorFlow. See Figure 3 for illustration. 
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.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 1: Input: Pre-trained score network s θ (•, •), differentiable objective function f . 2: Tuning Parameter: Strength parameters β(t), {y k } K-1 k=0 , number of iterations K, batch sizes {B k }.Initialization: G0 = NULL. 4: for k = 0, . . . , K -1 do 5: Generate: Sample z k,i ∼ Guided_BackwardSample(s θ , G k ) using Module 1, for i ∈ [B k ]. 6:Compute Guidance: 7), using s θ , gradient vector g k , and parameters y k and β(t). 7: end for 8: Output: (s θ , GK ).(i) Compute the sample mean zk := (1/B k ) B k i=1 z k,i . (ii) Query gradient g k = ∇f (z k ). (iii) Update gradient guidance G k+1 (•, •) = G loss (•, •) via (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.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 .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 S = s(x, t)EQUATION)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. f x * A,λ -f (µ K ) = λ L λ K O d log K δ ,where x * A,λ is an optimal solution of the regularized objective:EQUATIONwhere μ, Σ are empirical mean and covariance of pre-training data D.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.(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. 
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] .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:EQUATIONwhere 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 } 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 θ , G0 = 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 zk = (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) . (7) , using s θ k+1 and g k , with parameter y k , β(t). 7: end for 8: Output: (s θ K , GK ).(iii) Compute G k+1 (•, •) = G loss (•, •) inFinally, 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 classEQUATION)where Ĉ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δ,EQUATIONwheref * 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. 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) yg ⊤ 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.  
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.Figure 6 : Reward increase and effect on images across iterations. Left: Reward increases and converges across iterations. Larger guidance strength y (smaller regularizer strength) results in higher convergent reward value. Right: Images become more abstract, shifting from realistic to virtual backgrounds as reward increases.. 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. 
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.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). 
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? 
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). 
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]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. 
 
 Appendix: 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.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.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.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) isS = s(x, t) = C t x + b t : C t ∈ R D×D , b t ∈ R D , (recall (12))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 s θ (x t , t) = C t (x t -α(t)x) , where x = k i=0 w k,i E x∈Di [x] k i=0 w k,i .Moreover, minimizing the score matching objective (14) over the class (12) yieldss θ (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.Proof. Using the linear score network class S with freezing C t , we cast the score matching loss ( 14) intoT 0 k i=0 w k,i E x0∈Di E xt∼N (α(t)x0,h(t)I D ) ∥∇ xt log ϕ t (x t |x 0 ) -s(x t , t)∥ 2 dt = T 0 k i=0 w k,i E x0∈Di E xt∼N (α(t)x0,h(t)I D ) - 1 h(t) (x t -α(t)x 0 ) -C t x t -b t 2 dt = T 0 k i=0 w k,i E x0∈Di E xt∼N (α(t)x0,h(t)I D ) C t + 1 h(t) I D (x t -α(t)x 0 ) + (α(t)C t x 0 + b t ) 2 dt (i) = T 0 k i=0 w k,i E x0∈Di ∥b t + α(t)C t x 0 ∥ 2 + w T 0 trace h(t) C t + 1 h(t) I D ⊤ C t + 1 h(t) I D dt,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 minimizek i=0 w k,i E x0∈Di ∥b t + α(t)C t x 0 ∥ 2 for any t, which leads to bt = -α(t)C t x.Now, we solve C t for the second result. Substituting bt into the optimization objective ( 14) yields:T 0 k i=0 w k,i E x0∈Di ∥α(t)C t x -α(t)C t x 0 ∥ 2 + w T 0 trace h(t) C t + 1 h(t) I D ⊤ C t + 1 h(t) I D dt = T 0 α 2 (t) k i=0 w k,i E x0∈Di ∥C t (x 0 -x)∥ 2 + w T 0 trace h(t) C ⊤ t C t + 1 h(t) C t + 1 h(t) C ⊤ t + 1 h 2 (t) I D dt.Taking the gradient for C t , we get2α 2 (t)C t k i=0 w k,i E x0∈Di (x 0 -x)(x 0 -x) ⊤ + 2wh(t)C t + 2wI D .Setting the gradient above to 0, we get the solution for C t asĈt = -α 2 (t)w -1 k i=0 w k,i E x0∈Di (x 0 -x)(x 0 -x) ⊤ + h(t)I D -1 .Therefore, the proof is completed.When w k,0 = 1, w k,i = 0, i ∈ [k],EQUATIONThe 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 → ∞,X ← T d = N μ + y -g ⊤ µ σ 2 + g ⊤ Σg Σg, Σ - Σgg ⊤ Σ σ 2 + g ⊤ Σg .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 getX t Y d = N α(t)μ g ⊤ μ , α 2 (t) Σ + h(t)I D α(t) Σg α(t)g ⊤ Σ σ 2 + g ⊤ Σg .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 beEQUATION)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): TV (p 0 , p ← T ) ≤ TV (p T , φ) , 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ū and Σ = A Σu A ⊤ , whereū = E x∈D,u=A ⊤ x [u], Σu = E x∈D,u=A ⊤ x [(u -ū) (u -ū) ⊤ ]. Thus, the covariance matrix of X ← T is Σ - Σgg ⊤ Σ σ 2 + g ⊤ Σg = A Σu - Σu A ⊤ gg ⊤ A Σu σ 2 + g ⊤ Σg A ⊤ ,and due to X ← T follows Gaussian distribution, we get X ← T ∈ Span(A). Thus, the proof is completed. EQUATION)where p LD t (u ′ ) = ϕ t (u ′ |u)p u (u) du 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).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:dX ← t,⊥ = 1 2 - 1 h(T -t) X ← t,⊥ dt + b(t)gdt + I D -AA ⊤ dW t .Solving this SDE, we get the expectation of the final state followingE[X ← T,⊥ ] = T 0 exp - t 0 h -1 (s)ds e t/2 b(T -t)gdt.For the schedule h(t) = 1exp(-√ t), we have the coefficient of direction g is larger thanT 0 exp(-t/2 -2 √ t)b(T -t)dt > 1 0 exp(-5/2)b 0 dt > 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.To further clarify the derivation of our gradient guidance, we present a variant of Lemma 1.Proposition 3. Under Assumption 2, we haveEQUATION)where E[x 0 |x t ] denotes the conditional expectation of x 0 given x t in the forward process (1), andβ(t) = α(t)/(σ 2 + g ⊤ Σ -1 I D + α 2 (t)/h(t) • Σ -1 g).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:G(x t , t) ∝ y -g ⊤ E[x 0 |x t ] • g.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 yg ⊤ 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 theoretical 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.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 ∇ xt log p t (y | x t ) can be derived as∇ xt log p t (y | x t ) = -∇ xt 1 2 y -m y (x t ) σ y (x t ) 2 -∇ xt log σ y (x t ).Due to the linearity of Y with regard to X 0 , m y (x t ) can be computed asEQUATIONTo 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 haveX t Y = α(t)I D 0 I D g ⊤ 1 0   x 0 ϵ Z t   .Observing that the joint distribution of (x 0 , ϵ, z t ) is Gaussian, we deduceX t Y d = N α(t)µ g ⊤ µ , α 2 (t)Σ + h(t)I D α(t)Σg α(t)g ⊤ Σ σ 2 + g ⊤ Σg .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∇ xt log p t (y | x t ) = - 1 2σ 2 y (x t ) ∇ xt y -g ⊤ E[x 0 | x t ] 2 = 1 σ 2 y (x t ) y -g ⊤ E[x 0 | x t ] ∇ xt E[x 0 | x t ]g.To get E[x 0 | x t ], we derive the joint distribution of (X 0 , X t ):X 0 X t d = N µ α(t)µ , Σ α(t)Σ α(t)Σ α 2 (t)Σ + h(t)I D . Thus, we get E[x 0 |x t ] = µ + α(t) α 2 (t)Σ + h(t)I D -1 Σ (x t -α(t)µ).As a consequence, we have∇ xt log p t (y | x t ) = 1 σ 2 y (x t ) y -g ⊤ E[x 0 | x t ] α(t) α 2 (t)Σ + h(t)I D -1 Σg.Together with the following equality by Woodbury identity, we get the result.σ -2 y (x t ) = σ 2 + g ⊤ Σg -α 2 (t)g ⊤ Σ α 2 (t)Σ + h(t)I D -1 Σg -1 = σ 2 + g ⊤ Σ -1 I D + α 2 (t) h(t) Σ -1 g -1.Proof. We expand the derivative in G loss asG loss (x t , t) = 2β(t)(y -g ⊤ E[x 0 |x t ]) (∇ xt E[x 0 |x t ]) ⊤ g. 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∇ log p t (x t ) = A∇ log p LD t (A ⊤ x t ) - 1 h(t) I D -AA ⊤ x t ,where p LD t denotes the diffused latent distribution, i.e., p LD t (u ′ ) = ϕ t (u ′ |u 0 )p u (u 0 ) du 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 thatE[x 0 |x t ] = 1 α(t) x t + h(t) Am(A ⊤ x t ) - 1 h(t) x t = h(t) α(t) Am(A ⊤ x t ).As a consequence, we can verify that∇ xt E[x 0 |x t ] = h(t) α(t) A ∇m(A ⊤ x t ) A ⊤ ,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 thatG loss (x t , t) = 2β(t)(y -g ⊤ E[x 0 |x t ]) (∇ xt E[x 0 |x t ]) ⊤ g = ℓ t • g for ℓ t = 2β(t)(y -g ⊤ E[x 0 |x t ]) and g ′ = (∇ xt E[x 0 |x t ]) ⊤ g = h(t) α(t) A ∇m(A ⊤ x t ) ⊤ A ⊤ g ∈ Span(A).The proof is complete. . 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 * λ , andEQUATION)where D is the ambient dimension of data, and x * λ is a regularized maximizer of f given byEQUATIONwhere μ, Σ 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 .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 followss θ (x t , t) = -α 2 (t) Σ + h(t)I D -1 (x t -α(t)μ) .With proper choices of β(t), gradient guidance G loss leads to the following output distribution at the end of round k:N μ + y k -g ⊤ k μ σ 2 + g ⊤ k Σg k Σg k , Σ - Σg k g ⊤ k Σ σ 2 + g ⊤ k Σg k .Thus, we obtain the mean of the above distribution, i.e., µ k+1 = μ + η k Σ∇f (z k ), where zk is the empirical mean of previous samples, η k is a stepsize determined by y k . By a rearrangement, we obtain a recursive formulaEQUATIONWe observe that (24) resembles a gradient ascent update from µ k ≈ zk 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-1k=0 with H k be the information accumulated after k rounds of Alg.1.H 0 := σ(μ),H k := σ H k-1 , σ z k-1,1 , . . . , z k-1,B k-1 , k ∈ [K].Define the expectation of samples generated at k-th round asµ k := E[z k,i | H k-1 ], k ∈ [K -1].Applying Corollary 1, we get the pre-trained score ass θ (x t , t) = -α 2 (t) Σ + h(t)I D -1 (x t -α(t)μ) ,If we set y k as followsy k = η • σ 2 + g ⊤ k Σg k + g ⊤ k μ, 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:EQUATIONDefine the empirical covariance matrix of the latent variableU as Σu = E x∈D,u=A ⊤ x [(u -ū) (u -ū) ⊤ ] where ū = 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µ k = E[z k,i | H k-1 ] = AA ⊤ μ + η • A Σu A ⊤ g k-1 ,where g k-1 = ∇f (z k-1 ) and zk-1 = (1/B) B i z k-1,i . We rearrange the update rule to show a gradient ascent formula as followsµ k = AA ⊤ µ k-1 -AA ⊤ (µ k-1 -μ) + η • A Σu A ⊤ ∇f (µ k-1 ) + η • A Σu A ⊤ (g k-1 -∇f (µ k-1 )) = AA ⊤ µ k-1 -A Σu A ⊤ A Σ-1 u A ⊤ (µ k-1 -μ) + η • A Σu A ⊤ ∇f (µ k-1 ) + η • A Σu A ⊤ (g k-1 -∇f (µ k-1 )) = AA ⊤ µ k-1 + η • A Σu A ⊤ ∇f (µ k-1 ) -λA Σ-1 u A ⊤ (µ k-1 -μ) + η • A Σu A ⊤ (g k-1 -∇f (µ k-1 )) . where λ = 1/η. Define h(x) := f (x) -λ/2 ∥x -μ∥ 2 Σ-1 , we have µ k = AA ⊤ µ k-1 + η • Σ∇h(µ k-1 ) + η • Σ (g k-1 -∇f (µ k-1 )) .Recall the notation for the optimum:EQUATIONWe bound the second term I 2 first. According to f is L-smooth with respect to ∥•∥ Σ-1 , we haveI 2 = η ∥g k-1 -∇f (µ k-1 )∥ Σ ≤ ηL ∥z k-1 -µ k-1 ∥ Σ-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∥z k-1 -µ k-1 ∥ 2 Σ-1 ≤ 2 log 2K δ • trace V(z k-1,i ) • Σ-1 B k-1 .We have traceV(z k-1,i ) • Σ-1 ≤ trace Σ • Σ-1 = d.Therefore, I 2 is bounded byI 2 ≤ M 0 / B k-1 , 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 ≥ ζ -4k (1 -ζ) -2 , together with z 0 = μ, we get µ K -x ⋆ A,λ Σ-1 ≲ ζ K μ -x ⋆ A,λ Σ-1 + M 1 • √ d , w.p. 1 -δ, with η = 2 L + 2λ , ζ = |1 -η(L + λ)| and M 1 = 2L (1 + η 2 ) log 2K δ. Also, we can get (28) as in proof in Appendix E.2. We restate it here:EQUATIONΣ-1 . Due to (32) , we have, it holds w.p. 1δ,f x * A,λ -f (µ K ) ≲ λ ∥x * A -μ∥ 2 Σ-1 + (L + λ)ζ K μ -x * A,λ 2 Σ-1 + M 1 d . We choose λ = L log K/(4K) and get f x * A,λ -f (µ K ) ≲ L log K K • ∥x * A -μ∥ 2 Σ-1 + μ -x * A,λ 2 Σ-1 + M 1 d , w.p. 1 -δ.With assuming x ⋆ A,λ is bounded, we derivef x * A,λ -f (µ K ) = O dL 2 log K K • log K δ , w.p. 1 -δ.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 holdsEQUATIONProof. See Bubeck et al. [9, Lemma 3.11] for a proof.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.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 ⊤ )θ 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.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 Ê[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 . 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. 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.(a) θ = Aβ * (b) ∥θ ⊥ ∥ ∥θ ∥ ∥ = 9 (c) b = 4 • 1D (d) b ∼ N (4 • 1D, 9 • ID)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.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] .We summarize the time cost of our experiments on one NVIDIA A100 GPU in Table 1 . 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.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). 


Decision: The paper provides new theoretical and algorithmic insights into gradient guidance of diffusion models, which is the practice of adding an offset vector to the pre-trained score function estimate in the reverse SDE of a diffusion model with the goal of controlling the output toward a user-specified objective. The contributions are primarily theoretical and deal with guidance based on the “look-ahead” loss for estimating x0 given xt. The theory connects this type of gradient guidance to optimizing a user-specified loss function over the data subspace under the assumption that data lies on a linear subspace. The main convergence results are developed under the additional assumption that the pre-trained score-function estimator is a linear function of xt. 

Reviewers unanimously appreciated the theory, calling it “interesting and novel”, “solid”, “very impressive”, etc.; Reviewer 2yVt appreciated that it provided theoretical evidence favoring one type of guidance over another. The primary concerns were about: (1) strong assumptions underlying the theoretical results, (2) efficiency in practice, (3) the tone of some claims in the introduction being too strong given the assumptions and/or prior evidence. The efficiency concerns were adequately addressed during the rebuttal. The assumptions and their implications were also discussed and some clarifications made. The meta-reviewer notes that assumptions were all clearly stated in the methods section, and all reviewers found the theory interesting and relevant despite not being completely general. I believe the issue is the matching of the claims in the into to the evidence provided by the theorems. I urge the authors to follow the suggestion of Reviewer 2yVt to temper the optimism surrounding the convergence results in the introduction by stating clearly in the introduction the assumption of a linear score network. I also echo the advice to more clearly reference prior evidence surrounding the use of “look-ahead” vs “naive” gradient guidance. With that said, it is clear that even with slightly toned down claims the paper provides valuable theoretical understanding for a fast-moving area.

Reviews: review: paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: - 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. 
- 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.
questions: questions: - 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?

## **Reference**
- [A] Towards Practical Plug-and-Play Diffusion Models, CVPR 2023.

- [B] ELUCIDATING THE DESIGN SPACE OF CLASSIFIER GUIDED DIFFUSION GENERATION, ICLR 2024

- [C] ADJOINTDPM: ADJOINT SENSITIVITY METHOD FOR GRADIENT BACKPROPAGATION OF DIFFUSION PROBABILISTIC MODELS, ICLR 2024.

- [D] Towards Accurate Guided Diffusion Sampling through Symplectic Adjoint Method, Arxiv 2023.

- [E] Loss-guided Diffusion Models for Plug-and-Play Controllable Generation, ICML 2023

- [F] Manifold preserving guided diffusion, Arxiv 2023

- [G] Freedom: Training-free energy guided conditional diffusion model, ICCV 2023.
limitations: limitations: Linear approximation of the score function in their proof is mentioned as the limitation.
paper_summary: paper_summary: Under 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).  
For concave reward, and under the additional assumption of linear score function, they also establish convergence results for the mean of the guided process.
strength_weakness: strength_weakness: The assumptions made in the paper are strong (as discussed below).

A 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.

There 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.

More comments are stated below.
questions: questions: 1.  

The tight connection between gradient-guided diffusion models and proximal optimization algorithms is discussed in detail in:  
Tomer Garber and Tom Tirer, "Image Restoration by Denoising Diffusion Models with Iteratively Preconditioned Guidance," CVPR, 2024.  
This relevant literature should be mentioned.

2.  

The statement in line 48:  
"Why does naively guiding diffusion models using gradient never work in practice?"  
Is 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.   
Only 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.  
This 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.  

3.

Based on your Eq. 1, it seems that you consider the variance preserving (VP) SDE formulation.  
Why do you state in line 85 that without loss of generality all your analysis is done for q=1?  
This 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).

4.  

There 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.

5.  

The "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.

6.  

You 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.

7.  

The assumptions of the theoretical analyses in the paper are quite strong.  
Especially Assumption 1 on the signal belonging to low dimensional linear subspace and the assumption on linear score model (Eq. 12).  
Already in the contribution list you should state the assumption on the score model, which essentially leads to convex optimization.

8.  

Any idea how can Theorem 1 be generalized to low dimensional manifold rather than low dimensional linear subspace?  
Under the assumptions, in what aspects Theorem 2 differs from standard convex analysis? (maximization of an L-smooth concave function?). 
Aren'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?

Valentin De Bortoli, "Convergence of denoising diffusion models under the manifold hypothesis," TMLR, 2022. 

9.  

Under the subspace assumption 1, your data covariance matrix is a DxD matrix of rank d<D.  
Therefore 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.  
Note that there are some peculiar differences between inverse and pseudoinverse, e.g., generally, pinv(AB) is not equal to pinv(B)pinv(A).

10.  

Regarding 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").

11.

What 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).  
The 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.
limitations: limitations: n/a
paper_summary: paper_summary: This 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
strength_weakness: strength_weakness: 1. 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.

2. 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.

3. 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 
$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.

[1] Song, Jiaming, et al. "Loss-guided diffusion models for plug-and-play controllable generation." International Conference on Machine Learning. PMLR, 2023.

[2] Chung, Hyungjin, et al. "Diffusion posterior sampling for general noisy inverse problems." arXiv preprint arXiv:2209.14687 (2022).

[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
Computer Vision, pages 7280–7290, 2023b.

[4] 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.

[5] Korrawe Karunratanakul, Konpat Preechakul, Emre Aksan, Thabo Beeler, Supasorn Suwajanakorn,
and Siyu Tang. Optimizing diffusion noise can serve as universal motion priors. arXiv preprint
arXiv:2312.11994, 2023.

[6] Tang, Zhiwei, et al. "Tuning-Free Alignment of Diffusion Models with Direct Noise Optimization." arXiv preprint arXiv:2405.18881 (2024).
questions: questions: 1. I wonder how the algorithm proposed in this work compares to the direct optimization approaches in [3][4][5][6] (references mentioned above)?

2. 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.
limitations: limitations: See my comments above. In summary, the main limitations of this work are:

1. Incomplete literature review and comparison with existing works.

2. The performance of the proposed algorithm appears to be impractical.

3. The presentation can be improved.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: No questions
limitations: limitations: N/A

Comments: title: Thanks for your rebuttal
Author Reply to Reviewer 2yVt
comment: Hello thanks for the response! We are glad that **your primary concern about backpropagation is now resolved**. In regards to further comments: 

> "the main contribution of this work relies on a simplified model (linear or concave). "

Sorry there could still be a major misunderstanding. Our contribution does not **rely** on a simplified model. 

1. The simplified model is only used to **motivate** the design of gradient guidance. We choose to use a simple linear model to intuitively explain why certain forms of guidance are better than others. 
But our results are not limited to the simplified model, as follows.

2. **Theorem 1** holds for  **arbitrary** distributions and does **not** require a linear score model.

3. The guidance term constructed in our paper, G_loss (Eqn 7 and Eqn 11 in submission), applies to any pre-trained score network. It is not limited to linear models. 

4. Our first experiment (Sec 6.1 of submission) uses a 15M-parameter, nonlinear, U-Net score network.

5. Our second experiment (Sec 6.2) generated human-interpretable images, validating that our method preserves the latent manifold structure underlying image data. 

6. The image experiment (Sec 6.2) also validated the method on nonconcave objectives. 

7. Per your request for additional nonlinear experiments, we have provided a new set of experiments for synthetic data on nonlinear, spherical manifolds (**Sec 1 of [pdf][pdf-link]**).

While our methodology extends beyond "simplified models", we intentionally begin with simple mathematical models as the foundation for exploring complex theories. Simple models, such as linear models, have been fundamental to the developments of deep learning and diffusion model theories, e.g., [38] Marion et al, 2024. Philosophically, mastering these basic models often provides the deepest insights, most practical utility, as well as generalizability. Therefore, we firmly believe that robust theoretical research must start with simplicity.



> “several aspects of algorithm design also appear in the works on gradient guidance”

Our paper focuses on theory. As we already explained, gradient guidance via backpropagation is a common practice (universal guidance does it; [1],[2] also does it). Our contribution is to provide theoretical motivation for this design and establish the first optimization convergence theory.  

> “Direct optimization method also targets the scenario "adapting diffusion models with self-play and online data collection to find the global optimum of a user-provided objective function.", and hence this work is not "the first methodology".

By “the first methodology”, we mean the first methodology to provably solve optimization problems with convergence/optimality guarantees, with both theory and experiments. We are not aware of any other method that probably adapts a pretrained diffusion model to solve optimization problems to global optimum. 

We appreciate that you now agree that guidance methods and direct optimization “have different computation requirements”. They are different approaches in nature, and fully understanding their limits would be future work. While our focus is guidance, we appreciate that you highlight this alternative technological route. We’d be happy to incorporate more discussions about it in our related work section. 

**Could you please point us to the specific literature on direct optimization with self-play and online data collection to find global optimum of any objective function?** It’s possible that we have missed something. We are happy to go over them and include them in our literature review. Thank you!


- Please also note that the discussion period will end soon today. So we might not be able to respond again, but we appreciate any constructive comment!



[1] Song, Jiaming, et al. "Loss-guided diffusion models for plug-and-play controllable generation." International Conference on Machine Learning. PMLR, 2023.

[2] Chung, Hyungjin, et al. "Diffusion posterior sampling for general noisy inverse problems." arXiv preprint arXiv:2209.14687 (2022).


[38] Marion et al. Implicit diffusion: Efficient optimization through stochastic sampling, 2024


[pdf-link]: https://openreview.net/attachment?id=WsELdf2puF&name=pdf
Dear Reviewer 8pJj,

We sincerely appreciate your insightful review and are glad we have successfully addressed your concerns. If you feel that all your questions have been satisfactorily answered, we would be grateful if you would consider raising the rating.

Thanks again for your time and efforts in this work!

Best regards,

Authors,
I would like to thank the author for addressing some of my concerns and correcting my mistake. Thus I decided to increase my score by 1. However, I still do not lean towards acceptance for two reasons.

1. The design of the algorithm and theory, which are claimed to be the main contribution of this work relies on a simplified model (linear or concave). This makes it unclear what the fundamental factor to drive the improvement over other baselines like LGD, and also dims the meaningfulness of the theoretical results. Besides, I think several aspects of algorithm design also appear in the works on gradient guidance.

2. I disagree with the author's comment "We respectfully disagree that our method and direct optimization target the same tasks. Indeed they are relevant. However, we consider general optimization over data's latent manifold. Our Section 5 provides the first methodology and theory for adapting diffusion models with self-play and online data collection to find the global optimum of a user-provided objective function." Direct optimization method also targets the scenario "adapting diffusion models with self-play and online data collection to find the global optimum of a user-provided objective function.", and hence this work is not "the first methodology". As far as I can tell, all the experiments in the current manuscript can also be tackled by the Direct optimization method. While saying so, I do agree that the gradient guidance method has lower time complexity. I know it could hard to run the direct numerical comparison during the short rebuttal period. But I do wish the author to acknowledge the importance of this comparison because it is still possible that: The direct optimization method, while having high time-complexity for each gradient step, can still have a better reward-time trade-off compared to the gradient-guidance method.
I sincerely appreciate your thorough response. I carefully read all the reviewer's comments and discussions, and my concerns are addressed.
> I thank the authors for their response, which addresses most of my concerns.

We are glad that our response addressed most of your concern and you increased the rating. Thank you again for raising some insightful questions, which help us improve the paper. With that being said, we will include in the revised text version: a discussion on similar guidance designs in existing literature, mentioning assumptions made in contribution list, as well as a discussion on the previously made connection between gradient-guided diffusion and proximal optimization, and its differece to the connection drawn in our paper.
I thank the authors for their response, which addresses most of my concerns.  
As mentioned above, I still believe that the claims on the algorithmic novelty should be toned done (e.g., the core guidance idea is common), and that already in the contribution list the strong assumption on the linear score model (Eq. 12) should be stated.  
I decided to increase the rating, provided that the next version will include these changes, as well as discussion on the connection of gradient-guided diffusion schemes and proximal optimization methods that is stated in exiting work (where "pretrained model has an effect of regularization" as well).
>**Q2:** The tight connection between gradient-guided diffusion models and proximal optimization algorithms is discussed in detail in Garber and Tirer (2024). This relevant literature should be mentioned.

**A2:** This work is certainly related and has been included in our revised literature review (Sec 4 of [pdf][pdf-link]). However the connection between the proximal-gradient and gradient-based guidance discussed there is different from the connection we made. The connection in their paper is an optimization trick used to propose a two-step alternating algorithm for guided generation, which is classic in the research line of PnP (Plug in and Play). In contrast, our connection is an intuitive understanding distilled from our optimization theorem: pretrained model has an effect of regularization, guidance serves to optimize the objective.

> **Q3:** Is the algorithm slow?

**A3:** There might be misunderstanding on the reported "76 min" running time. We newly added a computation efficiency analysis with detailed running time records, please refer to **Table 1** and **computational efficiency** section in [our rebuttal][re-link] for more details.

>**Q4:** What's the difference between Alg.2 and other fine-tuning diffusion models works?

**A4:** Alg.2 simultaneously fine-tunes the pretrained model and adds a guidance term during inference. In contrast, existing works on fine-tuning diffusion models usually involves only fine-tuning the model weights.
>**Q5:** What is the motivation for considering convergence to the global optimum of the objective if one pays in ignoring the prior information of the score model?

**A5:** In our optimization framework, our goal is to find the global optimum **within the subspace**, which is also prior information or the underlying constraint. Thm 3 proves the exact convergence to the global optimum within the subspace. For example, in protein design, the generated samples should closely mimic natural proteins and comply with biological principles. Failure to do so can result in unstable structures, significantly degrading performance metrics. We agree that to which extent we balance reward optimization with regularization depends on the specific use case.

>**Q6:**  Why do you state without loss of generality that $q=1$ for the noise schedule in line 85?


**A6:** Our theoretical work actually holds for all general values of $q$, not just $q=1$. We appreciate your observation and we've clarified this point in our revision.


>**Q7:** Why are Eqn 2 and Eqn 6 in [5] Song et al. (2020) inconsistent?


**A7:** In our Eqn 2, $X_0$ represents the standard Gaussian noise distribution, whereas in Song et al.'s Eqn 6, $X_T$ represents the noise distribution. This difference in notation leads to a discrepancy in the $\mathrm{d}t$ term, specifically in its sign. 



>**Q8:** Suggests using the dagger symbol to distinguish between inverse and pseudoinverse.

**A8:** Thank you for the suggestion. We will add more clarification in the revision. Regarding the data covariance matrix, it has the form $A\Sigma_u A^T$, where $\Sigma_u$ is full rank and the columns of $A$ are orthonormal, thus the pseudoinverse is $A\Sigma_u^{-1} A^T$.

[re-link]: https://openreview.net/forum?id=X1QeUYBXke&noteId=WsELdf2puF
[pdf-link]: https://openreview.net/attachment?id=WsELdf2puF&name=pdf