Gradient Guidance for Diffusion Models: An Optimization Perspective

We appreciate your recognition of our theoretical contributions and your valuable suggestions! We've added more detailed explanations of the derivations in the proof.

Q1: The proposed method is prohibitively slow due to gradient backpropagation over the neural network. Other approaches like [1][2] do not require gradient backpropagation over the neural network.

A1: There seems to be some confusion. Your understanding is incorrect. Both [1] and [2] require backpropagation through the score network. Please refer to:

• In [1], Eqn.2 includes score network, and Eqn.8 takes gradient over it.

• In [2], Eqn.10 includes score network, and Eqn.16 takes gradient over it.

Our algorithm is as efficient as other gradient-based guidance methods while enjoying additional theoretical guarantees. Please see our new Table 1 in pdf for a detailed analysis of our runtime efficiencies.

Q2: The method can't rival direct optimization. These works target the same task and have the same computational requirements.

A2: Thank you for pointing us to these papers. They are relevant and we have added discussions about these works in our expanded literature review in Sec 4 of pdf.

• We target more general tasks and focus on theory. 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.

• They don't have the same computation requirement: Direct optimization has a high memory burden and ours does not. While our gradient guidance method enables sampling using $O(1)$ memory, direct optimization methods [3-6] suffer from O(T) memory, a substantially higher burden, because they need to backpropagate through the ODE solver, which requires storing all intermediate gradients when utilizing the chain rule. For example, [4] states: "backpropagating through the solver can be expensive in memory", similarly, [5] notes: "maintaining the intermediate activations for solving the ODE during backpropagation can be memory-intensive. This issue can be addressed with gradient checkpointing or an invertible ODE, at the cost of more computation or model complexity." Thus, our method is much more light-weighted.

Q3: A concurrent work DNO [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.

A3: Thank you for pointing out this paper. After carefully checking their paper, precise running time was only reported for SD1.5 (Table 1 in DNO[6]), the same model tested in our experiments. For SDXL, there's only a figure with a vague runtime axis, prevents direct comparision. We added a new Table 1 our rebuttal for analysis of computation efficiency. For adapting the SD 1.5 model, our method takes 15.8 seconds per optimization round, and <5 iterations to converge (as shown in Figure 6). Our total runtime is < 2min, which is comparable to the 10min runtime reported by Table 1 in DNO .

Q4: Could you conduct some simple experiments on synthetic data with a nonlinear sample space?

A4: Thanks for the suggestion. We have conducted additional experiment (Sec 1 in pdf ) on data on nonlinear space. Our new result demonstrates that our guidance better preserves the structure of nonlinear manifolds than naive gradient guidance.

Q5: When moving from Section 3 to Section 4, why does the target objective function switch from a linear function to a general function?

A5: Yes our task is to generate solutions for general nonlinear, concave optimization problems. Our Section 3 uses the linear model as a motivating example to derive the gradient guidance. Our Sections 4, 5 provide our main results for general nonlinear optimization.

References:

[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." 2022.

[3] End-to-end diffusion latent optimization improves classifier guidance. ICCV, 2023.

[4] D-flow: Differentiating through flows for controlled generation, 2024.

[5] Optimizing diffusion noise can serve as universal motion priors, 2023.

[6] Tuning-Free Alignment of Diffusion Models with Direct Noise Optimization. arXiv preprint, 2024.

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

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

Thank you for appreciating our effort in providing theory for guidance. We first respond to your main concerns.

Weakness 1 Assumptions made are strong. Any idea how can Theorem 1 be generalized to low dimensional manifold rather than low dimensional linear subspace?

A: Given the challenge in developing theory for guided diffusion, adopting linear assuptions on data subspace and score funtion should be considered reasonable, more detailed discussion can be found in Linear Assumption part of our rebuttal.

Theorem 1 may have the potential to generalize to low-dimensional manifolds. For low-dimensional manifolds, the local geometry around any point can be approximated by its linear tangent space, and current Thm1 shows the guidance lies in this tangent space. By controlling the strength of guidance, the guidance vector on the tangent space will remain close to the manifold.

Weakness 1 (Cont.) In what aspects Thm 2 differ from standard convex analysis? There are works on convergence of diffusion models, e.g., (De Bortoli, 2022). What prevents extending them to the guided case?

A: Diffusion theory cannot be simply extended to guided diffusion. The addition of a guidance term disrupts the dynamics of the reverse SDE, thus the output distribution cannot be characterized. Previous theoretical work, such as [17](De Bortoli, 2022), heavily relies on the fact that distribution errors stem from approximation errors in the score function and the initial distribution of SDE. However, the training-free approach doesn't provide an approximation of $\nabla_{x_t} \log p_t(y|x_t)$ but instead introduces a plugin, implementation-convenient gradient guidance term. This discrepancy creates a barrier to directly extending existing theories to guided diffusion. [38](Marion et al, 2024) is the only work we know involving the theory of guided diffusion dynamics, and it also requires the linear score assumption. Therefore, this assumption should be considered reasonable. Our key insight is characterizing the output sample distribution of guided diffusion as taking a proximal gradient step. We derive the reverse process as a proximal step. This established relationship forms the premise for optimization and distinguishes our analysis from conventional convex optimization.

Weakness 2: A certain weakness of Theorems 2 and 3 is that they do not say anything about the variance.

A: This is a misunderstanding, you might have missed the formula we provided for variance in line 235, which points to Eqn 25 in Appx E. This formula demonstrates that the variance of the output distribution is smaller than that of the training data (empirical covariance). Similarly, for Theorem 3, we present the variance in Eqn 31 of Appx E.3.

Weakness 3: There are no new algorithmic ideas in the paper, as the method is common, and also there exist works on fine-tuned/adaptive diffusion models.

A: Algorithmically, we introduce an iterative optimization algorithm that apply the prompted gradient guidance to the local linearization of objective function(Alg 1). This approach is new within the optimization context. In addition, our main focus is more on theoretical understaning for guided diffusion than proposing new method.

Response to your other questions:

Q1: The statement in line 48 on "naive gradient guidance" is not clear. The explanation in Sec 3.2 that naive gradient doesn't work is not convincing because it ignores that fact that the step-size decreases and that the noise injection can mitigate error propagation.

A1: We appreciate the reviewer's suggestion, we added a clear definition for "Naive gradient guidance" around line 48. To give a more convincing explanation for the failure of naive gradient-based guidance, we add a new lemma showing naive gradient suffers at least a constant error:

Lemma(Failure of naive guidance) For naive guidance $`G`(X_t^{\leftarrow}, t) = b(t) \nabla f(X_t^{\leftarrow})$, suppose $b(t) >b_0>0$ for $t >t_0.$ For data in subspace under Assumption 1 and reward $f(x)=g^\top x$, $g \perp Span(A)$ with $h(t)=1-\exp(-\sqrt{t})$, then the orthogonal component of the generated sample is consistently large:

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

Solving this SDE, we get the expectation of the final state following $\mathbb{E}[X_{T, \perp}^{\leftarrow}] = \int_0^T \exp \left(- \int_0^{t}h^{-1}(s)\mathrm{d}s \right) e^{t/2}b(T-t) g \mathrm{d}t$. For $h(t) = 1 - \exp(-\sqrt{t})$, we have the coefficient of direction $g$ is larger than $\int_0^T \exp(-t/2-2\sqrt{t})b(T-t)\mathrm{d}t > \int_{0}^1 \exp(-5/2)b_0 \mathrm{d}t >0$ where we can assume $T>1$. Thus, $\mathbb{E}[ X_{T, \perp}^{\leftarrow}] \neq 0$. This means the generated sample is going out of the subspace, i.e., naive gradient guidance will violate the latent structure.

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

Thank you for acknowledging our theoretical value and other comments! We first respond to your main concerns.

Weakness 1: While the proposed algorithm demonstrates considerable theoretical value, its practical implementation appears to be quite slow.

A1: During rebuttal, we add a new running time analysis for our experiments, please refer to Table 1 and the computational efficiency section in our rebuttal for more details.

Weakness 2: The paper discusses Universal Guidance, yet it omits mention of other similar training-free and training-required guidance mechanisms. Additionally, it would be beneficial to address the practical and theoretical advantages compared to these mechanisms.

A2: We appreciate your mention of these related works that were omitted, and we've already included them in revision (Sec 4 in pdf) and discussed our practical and theoretical advantages compared to them. Here is a summary of the comparison.

• First of all, to adapt a pre-trained diffusion model for guided generation, there are training-required methods [A] and training-free methods, which can be further classified into gradient guidance methods [E,F,G,B] and latent optimization methods [C,D]. Our method falls into gradient guidance methods.

• Generally, classifier guidance [A] requires training/fine-tuning the classifier on noisy input even based on an off-the-shelf classifier, introducing extra training cost than training-free methods. Latent optimization methods are also training free, but are memory-intensive due to the need to store all intermediate gradients when utilizing the chain rule.

• Within gradient-based methods, our work focuses on the theoretical aspect, providing understanding of the guidance method in [E,F,G,B]. Additionally, in our work, we are the first to iteratively apply the gradient guidance as a module for solving optimization problems, and provide theoretical guarantees.

Response to your other questions:

Question 1: How reasonable is the assumption? Aren't guidance classifiers and score functions typically more complex than linear functions?

A1: Please note that we do not assume the guidance classifier to be linear in our theorems. The form of our guidance is motivated by Gaussian data + linear classifier but its guarantees do not rely on the classifier to be a linear model. Thm 1 also holds for non-linear score, i.e., arbitrary distribution.

We do assume linear score network in Thm 2 and 3, though being restrictive in practice, it is a reasonable assumption for establishing the first optimization theorem for guidance and yielding mathematical insights. The same assumption is also adopted by [38] (Marion et al, 2024), which is the only work we know involving the theory of guided diffusion dynamics. In addition, the effectiveness of our method is beyond linear assumption, justified by our experiments and new simulations (Sec 1 in pdf ) added during the rebuttal. More detailed discussion please refer to Linear Assumption section in our rebuttal.

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,