Gradient Guidance for Diffusion Models: An Optimization Perspective

Summary

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

Strengths

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

Weaknesses

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

Questions

No questions

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

N/A

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

Strengths

1. The idea of deriving this form of gradient guidance is interesting and novel, particularly the use of a low-dimensional linear sample space to demonstrate the problem of out-of-distribution in gradient-guidance sampling. 2. The theoretical results are solid, and I appreciate the finding that the proposed method can indeed preserve the low-dimensional linear structure theoretically.

Weaknesses

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

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.

Rating

4

Confidence

4

Soundness

2

Presentation

2

Contribution

2

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.

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.

Strengths

Overall providing theory for guided diffusion is appreciated. The paper formally shows an advantage of a certain guidance over the other. The presentation is clear and the paper is convenient to read.

Weaknesses

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

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.

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

n/a

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.

Strengths

- Strong theories are presented in this paper. Rethinking the guidance formulation as the optimization procedure, that aims to optimize 1) guidance and 2) convergence of the data subspace, is interesting. - The paper is well-written and easy to follow.

Weaknesses

- 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

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

Rating

5

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

Linear approximation of the score function in their proof is mentioned as the limitation.

Thanks for your rebuttal

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.