A First-order Generative Bilevel Optimization Framework for Diffusion Models

We thank the reviewer for appreciating the theoretical innovation and our algorithm design. Our response to your comments follows.

Q1. Baseline comparison. Thank you for your question. We have added numerical comparisons using different reward functions during the rebuttal period; please see Table R1 and Figure R1 on anonymous link: https://anonymous.4open.science/r/bilevel-diffusion-11A1/bilevel_diffusion_rebuttal.pdf. We test the performance of our method for another widely used lower-level reward function, HPSv2. Bilevel approach also outperform other baselines in terms of image quality in comparable time complexity, which showcase the robustness of our approach with respect to different reward functions. Due to time constraints, we could not include all the baselines you mentioned, as they all correspond to the lower-level reward fine-tuning task and require additional HPO on top of that. We believe these methods, though differing in their lower-level fine-tuning strategies, could similarly benefit from an upper-level HPO. We will discuss these points further in our revised manuscript.

Q2. Reproducibility. Thank you for your suggestion. We will release it with the final version.

Q3. Presentation of images. Thank you for your suggestion. We will change accordingly.

Q4. Relaxation of strongly convexity assumption. Thank you for your question. Yes, it is possible to relax strongly convexity assumption to so-called Polyak-Lojasiwcz condition following (Kwon et. al 2024, Shen et. al 2024). This condition covers nonconvex objectives and are satisfied by loss of overparameterized neural network [R5].

[R5] Loss landscapes and optimization in over-parameterized non-linear systems and neural networks. C. Liu, et al. 2022.

Q5. Memory cost comparison. For the first application of reward fine-tuning, since Algorithm 2 seprates the hyperparameter optimization stage and sampling stages, it does not introduce additional memory cost. For the second application of noise scheduling, we add a Table R2 in attached PDF to compare the memory usage of bilevel HPO and other approaches. By utilizing ZO and noise scheduler parameterization, the memory overhead of bilevel method is not severe.

Q6. Comparison with related works. Thank you for providing these literatures. Most of related works are related to the lower-level task for the first application on the fine-tuning diffusion model stage, including [R6-R11]. They belong to related works in "fine-tuning diffusion models" section and we adopt the guidance based fine-tuning method for the lower-level task. All of the methods in [R6-R11] can be further benefits from tuning the KL regularization using our framework. We will add a discussion in the paper.

[R6] Implicit Diffusion: Efficient optimization through stochastic sampling [R7] SPARKLE: A Unified Single-Loop Primal-Dual Framework for Decentralized Bilevel Optimization [R8] Bi-level Guided Diffusion Models for Zero-Shot Medical Imaging Inverse Problems [R9] Aligning Text-to-Image Diffusion Models with Reward Backpropagation [R10] Aligning Diffusion Models with Noise-Conditioned Perception [R11] Diffusion Model Alignment Using Direct Preference Optimization

We thank the reviewer for appreciating the theoretical guarantees for our algorithm and its empirical performance. Our response to your comments follows.

Q1. Optimization merely on $x$ in equation 1?

Thank you for your question. In equation 1, we state the general bilevel HPO formulation where the lower-level objective is not necessarily strongly convex so that $\mathcal{S}(x)$ may contain multiple solutions. Therefore, we also need to optimize on $y$ to select one solution. To eliminate any confusion, we will revise our formulation throughout to assume a strongly convex lower-level objective so that optimizing $x$ merely is enough.

Q2. Applicability of method on non-differentiable objective.

Thank you for your question. Our approach in the second application can extend to non-differentiable upper-level metrics, as ZO estimation is compatible with non-differentiable objectives. For the first application, if the reward function is non-differentiable, we would similarly employ a ZO estimator for its gradient; see [R4].

[R4] An Algorithm with Optimal Dimension-Dependence for Zero-Order Nonsmooth Nonconvex Stochastic Optimization. G. Kornowski, O. Shamir. JMLR, 2024.

We thank the reviewer for appreciating our novelty. We hope our response to your comments below can resolve your minor concerns.

Q1. Insufficient numerical experiments.

• One hyperparameter in the fine-tuning diffusion model experiment. Although KL regularization is just one hyperparameter, careful tuning it is essential. The appropriate KL strength $\lambda$ prevents reward over-optimization on downstream tasks by keeping the model close to the pre-trained distribution, while still allowing for the necessary variability to improve the reward (see Uehara et al., 2024; Fan et al., 2024) and Figure 2. We also consider multiple hyperparameter tuning on the second experiment for noise scheduling in training diffusion model.
• Noise scheduling tasks only on MNIST dataset. The noise scheduling problem we consider arises during the (pre-)training stage of diffusion models, which is particularly expensive compared to fine-tuning. Notably, the cost of training a diffusion model on MNIST is already comparable to the cost of fine-tuning on ImageNet (see Tables 1 and 2). Moreover, training requires tuning multiple hyperparameters in noise scheduler parameterization, further increasing the computational burden (see grid search method in Table 2). Our goal for the second application is to propose a new automatic HPO framework for training diffusion models via bilevel optimization, so we demonstrate our method using the MNIST dataset in this version. More datasets will be investigated in future work.
• More comprehensive experiments. During the rebuttal period, we add more comparisons on HPO for the fine-tuning diffusion model task. We test the performance of our method for another widely used lower-level reward function, HPSv2. See Table R1 and Figure R1 on anonymous link: https://anonymous.4open.science/r/bilevel-diffusion-11A1/bilevel_diffusion_rebuttal.pdf. Bilevel approaches also outperform other baselines in terms of image quality in comparable time complexity, which showcase the robustness of our approach with respect to different reward functions.

Q2. New challenges and novelty from the optimization perspective.

Thank you for your question. As we highlighted at the end of Section 1, there are sufficient novelty in the context of bilevel optimization as well. The first challenge is that we cannot directly optimize the distribution itself but can only work with samples. In HPO for fine-tuning a diffusion model, we use guided backward sampling to generate samples approximately from the desired distribution. Thanks to (Guo et al, 2024), we bridge the gap between guided sample-generation guarantees and optimization guarantees for the underlying probability distribution in this settings. In HPO for training a diffusion model, we parameterize the noise scheduler and noise distribution via cosine/sigmoid functions and a score network, respectively, and optimize their parameters instead. The second challenge is the numerical feasibility and computational overhead. For Application 1, we derived a backward-process-free approach to estimate the upper-level gradient (Proposition 1), reducing complexity. For Application 2, we employ ZO estimation to avoid both computational and memory costly backpropagation. Finally, we validate the assumption of strong convexity in probability space for diffusion models which are discussed in "implications for generative bilevel applications" in Section 5.

Q3. Hyperparameter tuning is a standard application of bilevel optimization.

Thank you for your question. Note that although prior work has explored HPO via bilevel optimization, most existing methods rely on implicit gradient or unrolling differentiation approaches that entail costly second-order computations. In contrast, we leveraged a fully first order bilevel method, which is new in HPO. Moreover, the techniques in existing works do not readily extend to the infinite-dimensional probability space of diffusion models. The new challenges we highlight in Q2 are fundamental when applying bilevel HPO to diffusion models.

Q4. Applicability of noise scheduling tuning in faster sampling with a diffusion model.

We are happy to see our work may stimulate promising direction! This is indeed an interesting future direction for us. We note the emerging line of work on noise optimization for faster fine-tuning (e.g., Tang et al., 2024), but the stages they target differ from ours. Those methods focus on reward fine-tuning using a pre-trained model so that it is essentially a single-level reward maximization task, similar to non-HPO version of our first application on reward fine-tuning, whereas we target automatic HPO for the fundamental diffusion model training stage. Moreover, as our method explored the better choice of noise scheduler at each iteration, our method can potentially reduce the overall sample complexity needed to achieve the target image quality in diffusion model training.

We thank the reviewer for appreciating our topic. We hope our response to your comments below can resolve your minor concerns.

Q1. Motivation of two hyperparameter optimization (HPO) problems.

In the first application, the primary concern is whether adding additional computational cost for HPO is worthwhile, whereas in the second application, the concern is on its rationale.

• First application - Reward fine-tuning. Our research is complementary to the entropy-regularized reward-guided diffusion models since no matter how to choose the guidance terms, the tuning of KL regularization is inevitable. Conventionally, tuning $\lambda$ via cross-validation (e.g. grid search) also incurs computational overhead due to the lack of universal guidelines across datasets and prompts (Uehara et al., 2024; Fan et al., 2024). In contrast, reward-guided diffusion model associated with our bilevel-based HPO approach yields better FID and CLIP scores while preserving time complexity (see Table 1 and the 3rd paragraph in Section 6.1). Furthermore, without bilevel, a common strategy to use CLIP for enhancing realistic is to include it in a weighted sum alongside the reward guidance, which requires the costly CLIP gradient. In contrast, our bilevel method, thanks to Prop. 1, is CLIP gradient-free, which cuts down overall time costs; see Table 1.

• Second application - Noise scheduling. The noise scheduler, the variance of the added noise, controls the generation quality in diffusion models. If the variance is too large, the forward process quickly becomes noise and fails to learn meaningful representation for backward generation; if too small, it never fully reaches Gaussian noise. Prior work has highlighted the need to tune this noise scheduler to balance this trade-off (Nichol & Dhariwal, 2021; Lin et al., 2024; Chen, 2023; R1). In our bilevel HPO framework, we first fix the noise scheduler $q(t)$ and train the diffusion model using score-matching function as in (Song et al., 2021a,b; Ho et al., 2020; Nichol et al., 2021), then measure backward image quality via FID on the upper-level. The proposed method not only outperforms and significantly accelerates runtime of conventional HPO, but also enhances DDIM with empirically chosen schedules in comparable time; see Table 2.

[R1] Simple diffusion: End-to-end diffusion for high resolution images. E Hoogeboom, et. al. ICML 2023.

Q2. Biased Montel Carlo (MC) estimation. Yes, the vanilla MC estimator is biased, but Theorem 1 can accommodate $\epsilon_k$ error. Moreover, since the reward is always positive, we have $e^{r_2(\cdot)/\lambda}, e^{(r_1(\cdot) / \gamma+r_2(\cdot))/\lambda}\geq 1$. Then since $\log(a)$ is Lipschitz continuous when $a\geq 1$, the MC estimation error can be controlled by large sampling batch size (see (Ji et al. 2021; Arbel & Mairal, 2022; Ghadimi & Wang, 2018)) or momentum updates (see [R2,R3]). However, empirical evidence suggests that we do not need a very large batch size and gradient updates without momentum also work well. We will add a comment in the revision and leave the theoretical improvement based on momentum updates for future work.

[R2] Stochastic compositional gradient descent: algorithms for minimizing compositions of expected-value functions. M. Wang, et. al. Math. Programming, 2017.

[R3] Solving stochastic compositional optimization is nearly as easy as solving stochastic optimization. T. Chen, et. al. IEEE Trans. on Signal Processing, 2021.

Q3. Generating from optimal entropy guided distribution.

We only need generation from the $\epsilon$-optimal entropy guided distribution and the error in Algorithm 5 can be accommodated by $\epsilon_k$ in Theorem 1. Moreover, Algorithm 5 in (Guo et al, 2024) converges linearly under a concave reward function and a strongly convex regularization term, implying that this accuracy can be achieved in $\mathcal{O}\bigl(\log(\epsilon_k^{-1})\bigr)$ iterations of Algorithm 5, which is not a huge computational burden. We will add a comment in the revision.

Q4. Applicability of other guidance strategies.

Yes, other guidance terms can be applied as long as they ensure the backward process converges on the optimal entropy-guided distribution, since our entropy-weight tuning approach is guidance-agnostic. We specifically choose the guidance from (Guo et al., 2024) because of its finite-time optimization guarantees. A clarifying paragraph will be provided.

Q5. Applicability of the noise scheduling task in faster sampling with a diffusion model.

Due to the space limit, please see the response to Q4 for Reviewer XKpE.

Q6. Noise scheduling task on MNIST dataset.

Due to the space limit, please see the response to Q1 for Reviewer XKpE.

Q7. Writing of this paper.

We follow the bilevel HPO framework, so some diffusion model details are moved to the Appendix due to space limit. We will clarify and improve the writing in the revision.