A First-order Generative Bilevel Optimization Framework for Diffusion Models
Authors: Quan Xiao
Hui Yuan
A F M Saif
Gaowen Liu
Ramana Rao Kompella
Mengdi Wang
Tianyi Chen
Venue: ICML 2025 poster (2025)
Keywords: Bilevel Optimization; diffusion model; hyperparameter optimization; fine-tuning; noise scheduling
Primary area: deep_learning->algorithms

Abstract: Diffusion models, which iteratively denoise data samples to synthesize high-quality outputs, have achieved empirical success across domains. However, optimizing these models for downstream tasks often involves nested bilevel structures, such as tuning hyperparameters for fine-tuning tasks or noise schedules in training dynamics, where traditional bilevel methods fail due to the infinite-dimensional probability space and prohibitive sampling costs. We formalize this challenge as a generative bilevel optimization problem and address two key scenarios: (1) fine-tuning pre-trained models via an inference-only lower-level solver paired with a sample-efficient gradient estimator for the upper level, and (2) training diffusion model from scratch with noise schedule optimization by reparameterizing the lower-level problem and designing a computationally tractable gradient estimator. Our first-order bilevel framework overcomes the incompatibility of conventional bilevel methods with diffusion processes, offering theoretical grounding and computational practicality. Experiments demonstrate that our method outperforms existing fine-tuning and hyperparameter search baselines.

A First-order Generative Bilevel Optimization Framework for Diffusion Models
Quan Xiao
Hui Yuan? 3
A F M Saif'
Gaowen Liu
Ramana Kompella" Mengdi Wang" Tianyi Chen
Abstract
upper-level objective, has evolved from its theoretical origins in the 1970s (Bracken & McGill, 1973) into cor- nerstone of modern machine learning. Its ability to model hierarchical dependencies makes it ideal for complex leamn- ing tasks, such as hyperparameter tuning (Pedregosa, 2016), meta-learning (Finn et al., 2017), reinforcement learning (Stadie et al., 2020; Shen et al., 2024), adversarial training (Zhang et al., 2022), neural architecture search (Liu et al., 2019) and LLM alignment (Zakarias al., 2024; Shen et al., 2025a; Gong al., 2022; Qin et al., 2024).
Diffusion models, which iteratively denoise data samples to synthesize high-quality outputs, have achieved empirical success across domains. How- ever, optimizing these models for downstream tasks often involves nested bilevel structures, such as tuning hyperparameters for fine-tuning tasks or noise schedules in training dynamics, where traditional bilevel methods fail due to the infinite- dimens--dimeniionll probability space and pro- hibitive sampling costs. We formalize this chal- lenge generative bilevel optimization prob- lem and address two key scenarios: (1) fine- tuning pre-trained models via an inference-only lower-level solver paired with sample-efficient gradient estimator for the upper level, and (2) training diffusion model from scratch with noise schedule optimization by reparameterizing the lower-level problem and designing a computation- ally tractable gradient estimator. Our first-order bilevel framework overcomes the incompatibil- ity of conventional bilevel methods with diffu- sion processes, offering theoretical grounding and computational practicality. Experiments demon- strate that our method outperforms existing fine- tuning and hyperparameter search baselines. Our code has been released at https: /ttp:/gittttt com/ afmaaif wwwffssaf//illiii /bwlebiilllleeeeeeeeeeeeeeeeeeee diffusion.
Mcanwhile, the diffusion model has achieved remark- able success across various domains, particularly in image (Croitoru et al., 2023; Ho et al., 2020; Song et al., 202 audio (Yang et al., 2023; Liu al., 2023a), and biological sequence generation (Jing al., 2022; Wu et al., 2022). Yet. optimizing these models for downstream tasks often intro- duces nested objectives: for instance, fine-tuning pre-trained models to maximize task-specific rewards (c.g., aesthetic quality (Yang et al., 2024)) risks reward over-optimization, where generated samples become unrealistic despite high scores. To mitigate this, incorporating an auxiliary objective can balance the optimization process and prevent excessive focus single metric; see Section 3.1. Similarly, design- ing noise schedules for the forward/reverse processes (Chen, 2023) requires balancing sample quality against computa- tional cost a problem inherently requiring coordination between training dynamics (lower-level) and scheduler pa- rameters (upper-level); see Section 3.2. An overview of the bilevel generative problem is shown in Figure 1. All of the above illustrate the need for developing bilevel algorithms that are friendly to diffusson models. However, applying existing bilevel methods to diffussion models which oper- ate in dimento-dimensional probability spaces and require costly sampling steps-poses unique challenges. Tradi- tional gradient-based bilevel approaches (Franceschi et al., 2017; Maclaurin et al., 2015) rely on exact lower-level solu- tions and dense gradient backpropagation, both infeasible for diffusion processes where generating single sample can involve hundreds of neural network evaluations.
1. Introduction
Bilevel optimization, which involves nested problems where a lower-level optimization is constrained by the solution of
*This work was done when the authors were at Renssclear Polytcchnic Institute. 'Department of Electrical, Computer, and Systems Enginecring, Rensselacr Polytechnic Institute, Troy, NY "Department of Electrical and Computer Engineering. Cornell Tech, Cornell University. New York, NY Department of Electrical and Computer Engineering, Princeton University. NJ 'Cisco Re- scarch. Correspondence to: Quan Xiao <quanxI800 1808 @ com,, Tianyi Chen <chentianyil9 @eatriinnyl@@gmmlilcommm .omm.
In this paper, we focus on designing computationally eifi- cient and diffusion- -riendll approaches for the following problem that we call the generative bilevel problem:
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
min xEX
f(x,)), s.t. y E S(x) = arg min g(x, 3) (1)
PEP

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
Pre-trained model
Reward fine-tuning
f
loss
score
Net
Prompt
so(sk,
lower level
HH
'Hor standing the ground
max ())) AKL (Paut))
lower level
H+
I-H
S(g) org arg (s ((,)))
Image quality
Avoid cour-ppimmiattion
min
Laq (klo)) (T)
upper level
02
Forward (noising) process
We
parameterize
FID
&
0g
0-
CLIP score
level
Dostne/Sigmoiddt.ce (c-G.)
Backward (sampling) process
max Fj
Gaussian naise
Figure 1. An overview of bilevel generative optimization problems. (Left) Fine-tuning diffusion model with entropy regularization strength parameter X. (Right) Noise parameter q. scheduling problem in the diffusion model.
where x is some hyperparameter in the diffusion model and y represents a distribution we aim to learn, which can be either an image distribution or a noise distribution. Both the upper-level f Rd x P > IR and lower-level objective functions g : dR P R are continuously differentiable, X' C IRd is closed set, and P is the probability space. We study the optimistic setting where we select the best response distribution y S(x) to minimize the upper-level loss. Let us denote the minimal lower-level objective value as the value function g (z) = minyes(s) g(z, y) and the nested objective as F(x) minges(s) (1x,y))
ing): see an overview in Figure 1. Two applications we considered are essentially bilevel hyperparameter optimiza- tion. To the best our knowledge, this is the first study to explore bilevel hyperparameter optimization in the con- text of diffusion models. We develop a first-order bilevel framework in both categories to solve (1). The primary differences with non-generative bilevel methods are:
D1) for fine-tuning diffusion models that include a pre- trained model, we adopt guidance based approaches rather than gradient-bused methods for the lower-level and penalty problems concerning distribution y, ensuring the process training-free and inference-only;
The key challenges of solving the generative bilevel prob- lem (1) are threefold. First, the lower-level variable y is typically a distribution that operates in infinite dimensional probability space. However, direct access to or optimization over distribution is not feasible; we only have access to sam- ples, and estimating distributions from samples is highly sample-inefficient. Therefore, the gradient over the distri- bution is inaccessible, making traditional gradient-based bilevel approaches (Shen et al., 2025b; Kwon al., 2023; Ji et al., 2021; Chen et al., 2021; Hong et al., 2023; Ghadimi & Wang, 2018) generally inapplicable. Second, the objec- tives [(x,,) y) and ((x, y) are usually some measures of the sample quality and might not have an explicit form in terms of the hyperparameter x, so that V.,( y) and V.g(a, y) are either not explicitly given requires sample efficient approximation. Third, for fine-tuning the diffusion model, existing literature related to bilevel fine-tuning on diffusion model (Clark et al., 2024; Marion et al., 2024) requires back- propagations over the pre-trained model, which often suffers from high computational and memory costs. In contrast, we design an inference-only bilevel method for this task.
D2) for the noise scheduling problem without a pre-trained model, we optimize noise proxy parameterized by score neural network, rather than performing noise distribution matching directly in the lower-level problem of (1); and,
D3) we design scalable methods to estimate V (x. y) and V Vzg(z, yy for two specific applications, i.e. leveraging the closed form of them and proposing sample-efficient estima- tion for the fine-tuned diffusion model, and using zeroth- order method to estimate Vzf y) in noise scheduling problem without a pre-trained diffusion model.
2. Bilevel Optimization and Diffusion Models
In this section, we will review some preliminaries on bilevel optimization, diffusion models, and guided generation, as well as two motivating applications of studying diffusion models in bilevel optimization.
2.1. Bilevel optimization
To address these challenges, we classify generative bilevel problems into two categories: (1) those with pre-trained model, where the target lower-level distribution is the im- age distribution (c.g., fine-tuning diffusion models), and (2) those without a pre-trained model, where the target lower- level distribution is the noise distribution in the forward process (e.g., noise scheduling during diffusion model train-
An efficient approach to solving bilevel optimization in (1) is to reformulate (1) to its single-level penalty problem (Shen et al., 2025b; Kwon et al., 2024), given by
min WEP
(x. = f(x, v) + ((((,,))))))))) v) gx()) (2)
and then optimize the upper-level and lower-level variables jointly. By setting the penalty constant y = (c--0.))) in-
2

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
versely proportional to the target accuracy e, the single-level problem (2) is (c)) approximate problem to the origi- nal bilevel problem (I). This method builds upon equilib- rium backpropagation (Scellier& & Bengio, 2017; Zucchet & Sacramento, 2022), which introduced a similar framework for strongly convex lower-level problems. However, recent works extend the study to accommodate some nonconvex lower-level problems and lay the foundation for broader applications (Shen et al., 2025b; Kwon al., 2024).
where the initial Up is a random variable drawn from the data distribution Pdata (Wer denotes the standard Wiener process, g(t) is a nondecreasing noise scheduling function, and U, represents the noise- noise-corrupted data distribution at time t. Under mild conditions and for a sufficiently large timestep T' (5) transforms the original distribution Pdatu into a distribution close to Gaussian random noise. (0, In).
Backward process. Given the forward process in (5), the reverse-tine time SDE is defined by
Gradient-based bilevel approaches. In the context of diffusion models, the upper-level and lower-level variables, x and y, usually have different meanings. To facilitate al- gorithm design tailored to diffusion models, we decompose (2) into separate y- and z-optimization problems, that
1 dū, - 2 200 ((1) log (T:) dt
+ qt)ddW for te (0,7),
(6)
where dW, is the reverse-time Wiener process, pt(-) (-) is the marginal distribution of U2 in the forward process. Let un denote the realization of a random variable U2 and s(u, t) == V log Ps (u) is the score function that has to be estimated in practice or given by the pre-trained model.
mEX min (x),, with (x) = min (x(
(ES,(x)
and (z) arg min C. (z) yEP
Under some mild conditions specified in (Kwon et al., 2024), (x) is differentiable with the gradient given by
Score matching loss. To estimate the score function s(u, we train a parameterized score network sæ (u, t) by tracking the gradient of the log-likelihood of probability from the forward process, which eliminates the need for distribution normalization. Specifically, we minimize the following score- matching loss (Song et al., 2021b):
V- V=g(z,v")) (4) VC = (<(()>>V_[[aa) (2,f2222) +> 22 --)====----==a
where 2* E (z) in (3) and y" E S(x) in (1) are any solutions. Moreover, (x)(() is a proxy of the original gradient, with the error bounded by [z) (=(x)) z O(1/7); see details in Section B. Therefore, when the lower- level problem ((x, y) and the penalty problem (z(2z, y) are solvable with solutions ,²,=", we can approximate (x)) and update the upper-level variable x.
min LsM
:= E My
[ log (r, - so su,, t)]?3] (7)
where t is uniformly sampled over the interval [o, 7]. and pt (u (uo) is the conditional distribution of un over the ini- tial image ug. Typically, the score network so(u, t) is parametrized by U-Net model (Ronneherger et al., 2015).
2.2. Diffusion models
The goal of diffusion models (Song et al., 2b; Ho et al., 2020; Song et al., 2021a) is to generate samples that match the some target distribution. This is achieved through two complementary stochastic differential equations (SDEs): a forward process, which gradually transforms input data into random noise, and backward process, which reconstructs the data by denoising the noise. Central to this framework is the learning of a score function, which captures the gradient of the log-likeliin bigeliiel of the data distribution and is invariant to the input. This score function, learned during the forward process, is then used to guide the reverse sampling process. Below, we provide brief overview of these processes.
Guided generation for optimization. By leveraging the estimated score function sø(u, T - t) from (7) to re- place V pr- samples can be generated through the backward process (6). Furthermore, we can introduce guidance terms into the backward process to steer the gen- eration toward the desired reward V v. Using the con- ditional score function, the goal is to estimate the condi- tional distribution PUIVV = t). By Bayes' rule, we have V., "e log v) V (u.,) + V., log log | w:). Given the pre-trained score so (u,,t) t) so V log Pe(w)) for the unconditional forward process, and the guidance term G V., log ut), the backward SDE defined by
Forward process. The forward process of diffusion model gradually transforms the original data Uo E R² into pure noise by incrementalyy adding noise dW,- This transformation simplifies complex data distributions into a tractable form, enabling efficient modeling and sampling. The following SDE formally describes the process:
aū, [ 1 -1(0) - (110)).-a(000oll ((()so(Ur,))) a(T.)) dt
+ ValldWW, LE (0,7).
(8)
aU, = 1 J00U di + Va(dd,, for q(1) > 0 (5)
With a proper guidance term and an increasing reward value e, the guided sampling process in (8) generates samples that
3

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
maximize the given reward function r(+) with an entropy regularization to the pre-trained distribution (Guo et al., 2024; Uehara et al., 2024). The design of the guidance term and the complete procedure of guided generation are outlined in Algorithm 5 and can be found in Appendix E.
ages with diffusion models. Instead of relying on cross- validation, bilevel optimization can automatically learn an effective noise scheduler, enabling the model to learn useful features while efficiently transforming imputs into noise. In this application, bilevel problem (10) optimizes the noise scheduler in the upper level to minimize a quality score, while the lower level learns the noise distribution from the forward process to match the true Gaussian noise. Instead of framing the lower-level problem distribution matching task, we optimize the distribution's 's parameter 0 as follows.
3. Applications of Generative Bilevel Problems
In this section, we introduce two bilevel optimization prob- lems in diffusion models: one in the fine-tuning stage with a pre-trained model, and another in the pre-training stage.
min flq. 0) == (E(g) [LL((((()))) (20.0ES(o)
3.1. Reward fine-tuning in diffusion models
s.t. S(q) = arg min g(q. 8) == Lsm((', u(1)) (10)
E'ER"
where pe is the probability distribution generated by the backward SDE (6) associated with parameter 0 in the score network, u(q) collects samples in the forward pass from [0, 7] defined by schedule q. L-M(-) is the score matching loss defined in (7), and (.) measures the scheduling qual- ity. For a given schedule q = {=(t)) at the lower-level, we generate samples {ue )== according to q and optimize the score network 0 to predict a good proxy of V log Pe Pe (ux). Then, at the upper level, we automatically tune the schedul- ing parameter by sampling from the reverse process with probability parameterized by 0 and evaluating the quality of the generated samples by L-s(-). Typical choice of Lss()) can be Fréchet Inception Distance (FID) score, commonly used metric to measure the quality of the generated image, and is differentiable; see details in Appendix E.3.
(a) A 0.01 (b) A 55.5 (c) 0.01 (d) A 44.3
Figure 2. Visualization of generated images: (a) Horse and (c) elephant generated with A = 0.01, leading to reward over- optimization and resulting in more abstract images misaligned with captions. In contrast, (b) horse and (d) elephant, generated using the bilevel method with À optimized via CLIP score. This suggests CLIP score is proper metrics for À selection. More vi- sualizations are shown in Figure and can be found Appendix.
In practice, fine-tuning pre-trained diffusion model is often necessary to generate samples that achieve high reward. However, if the model over-focuses on reward maximization, it may generate overly aggressive samples that diverge from realistic distributions (Gao et al., 2023). Therefore, a well- tuned model must carefully balance reward optimization with adherence to the pre-trained data distribution. This balance can be formulated as a bilevel optimization problem:
4. Diffusion-friendly Bilevel Methods
In this section, we first present the meta-algorithh for the generative bilevel problem (1), followed by tailored subrou- tines for two specific applications in Section 3.
min (ES(A))
f(A.p) = Euupplli (w))
A meta bilevel algorithm. To solve the generative bilevel problem (1), we can compute the gradient V VZz(x) accord- ing to (4), where the solutions 2' S2 and y² E S(x) can be approximated using numerical oracles (e.g., gradient descent or Adam (Kingma, 2015)). Specifically, at each iteration k. we first solve the y-problem over (x* v) and g(ah; yy to obtain near-optimal solutions y² y² and * = 2". We then perform gradient descent updates for the x- problem according to (4) with y² and 2², The procedure is detailed in Algorithm 1.
s.t. S(A) = arg min
(ruwwp(22(()) + (F IPaas))
y'EP
g(A.p)==
where the lower level adjusts data distribution by optimizing an entropy regularized reward learning problem (Uchara et al., 2024; Fan et al., 2024), and the upper level selects the best-response entropy strength X by another realistic- measured reward. As shown in Figure 2, the upper-level reward ri (-) can be caption alignment score (CLIP), further refining the generated samples to align with the provided captions. Since the pre-trained diffusion model is available, Pama is measured by the pre-trained score so (u, t).
4.1. Strategy with pre-trained diffussin models
We first focus on algorithms to guide the generated data distribution of the pre-trained diffusion model toward the optimal solution of (9) in the application of reward fine- tuning diffusion model, with a pre-trained score network.
3.2. Noise scheduling in diffusion models
Tuning the noise magnitude q()) in the forward and back- ward processes is essential for generating high-quality im-
Guided-sampling from the lower-level variable. For
4

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
given X, we want solve two optimization problems minp g\, p) and min, C, LA, pA. For single-level optimiza- tion, the distribution generated by the guided backward sam- pling in Algorithm 5 converges to the optimal distribution for maximizing the reward function r with an entropy regu- larization term to ensure the generated samples remain close to the pre-training data (Guo et al., 2024). This suggests Algorithm 5 can be used to solve the lower-level problem and penalty problem with respect to p. The guidance terms added in the backward sampling process for the lower-level and penalty problems, are defined as
Algorithm 1 A meta generative bilevel algorithm
1: Inputs: Initialization Zng Yo. Zp; target error cgi step- sizes Th: penalty constant " 2: for k = 0, 1........ K - do
3: estimate yı: arg minyer g(8 v) with eg error.
4: estimate 2 = arg mingeP (a 2) with eg error
estimate V (#x) with y$ R5 y" and 2t a 2" in (4)
6: update I = ( me (2 (xx))
7: end for
8: outputs: (xK, (xK,2k)
Glower (1, A) = G (u,, (:122 A Gpenalty (t, u, A) = G (u,, t; TI / +r2
(11a)
Algorithm Bilevel approach with pre-trained model
(11b)
1: Input: Pre-trained score network so(-:)) .). differentiable reward (1))), r2 r2).,, stepsizes Tjs penalty constant Tk.
where (ue, r)) is defined in (32). Then Algorithm 5 is able to generate samples that are approximately from the optimal lower-level and penalty distribution.
sample {i/m 1 m= from reverse SDE (6) using se 3: for == 1,1, ...., K - 1 do
4: estimate V Cr (Ax) by (33)
Gradient update for the upper-level X. By substituting the definition of the objective function into (4), we obtain
update Ax++1 = Projæ, (>x 7/ VC (Ax)) end for
7: sample (u I.M mm from Algorithm 5 using (so, ri + r2, Gpeaaltyy (1 >xx)) in (11b)
(A) = (KL((p, (A) (Pama) - KL (A)llamm)))) (12)
7x
8: Output: (AK, fuk JM
where (A) E aarg min, (2, p) and p*()) E arg min, p). Following the gradient- based bilevel method (Kwon et al., 2023; Shen et al., 2025b), direct way to estimate (A) is to use guided backward sam- pling in Algorithm 5 with (11a) and (1Ib) to obtain samples from p* (A) and p, (A), and then compute the KL diver- gence from samples using kernel-based probability estima- tion. However, it has two drawbacks: 1) each À update requires guided backward sampling, which computation- ally expensive; and 2) when the number of samples is less than the dimensionality of the data, the covariance matrix becomes rank-deficient, which makes kernel-based density estimation impossible.
with samples ] Pdata is estimated by empirical mean with detailed form deferred to Appendix E.2.
The full procedure is summarized in Algorithm 2.
4.2. Strategy without pre-trained diffusion models
We next design bilevel algorithm solve (10) in the ap- plication of noise scheduling for training diffusion models.
Lower-level problem solver. When q is given, the sam- ples from the forward process u(q) are determined. Then optimizing the score matching objective LsM(.) on the score network gives the optimal lower-level solution, i.e. optimal weights of the score network arg ming L0M u(q)).
To address these issues, by leveraging the marginal density induced by the SDE, we can derive closed-fore expression for the upper-level gradient in terms of samples.
Proposition 1. The gradient in (12) takes the form
Penalty problem solver with respect to 6. Gradient- based approaches on the penalty function (. 0) are effective in solving the penalty problem over 0. The gradient of Cg, 0) with respect to / takes the form of
(A,(A) -Eurpsan
[w)] 7 log Eurpsnn
ra(u)
y log Burppmo (10)///220))
(13)
Vol,(1.0) =
+ VoLsM u(4))
This proposition enables us estimate V VL (A) directly from the pre-trained distribution, eliminating the need for guided backward sampling in Algorithm 5 to com- pute p" (X) and (A, for each X and thus, significantly reducing computational cost. The proof is in Appendix D.I.
Remark 1.
where the second term can be directly calculated by differ- entiating the score-matching loss over 0, and the first term can be estimated by the mean of gradients over batch of samples (ūggg MM m=1 ) generated by the backward process (6). Although it is possible to obtain the gradient of V (eLs) using PyTorch 's auto-differentiation, it requires differen- tiating through the backward sampling trajectory. Since backward sampling involves 50-100 steps, even for effi-
Based on Proposition 1, the leper-level gradient in (13) can be estimated via the Monte Carlo method. In other words,
5

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
cient methods like the Denoising Diffusion Implicit Model (DDIM), auto-differentiation is memory-intensive. Instead, we estimate ( by zeroth-order (ZO) approxima- tion (Nesterov & Spokoiny, 2017; Shamir, 2017)
5. Theoretical Guarantee
In this section, we quantify the theoretical benefits of the bilevel algorithms in both applications. We make the follow- ing assumption, which is standard in bilevel optimization literature (Kwon et al., 2023; Ji et al., 2021; Chen et al., 2021; Franceschi et al., 2018; Hong et al., 2023).
5 (Lso(u8tve,,) Lso(")-v(.g) (14) 2v
V VoLso (
Assumption 1. The objective g(z,-) is Pa py-strongly con- vex, f(z y) and g(z, y) are jointly smooth over (z. y) with constant f and kg,1 for all E X and y P. More- over, f(x,)) is E5,0 Lipschitz continuous and g(æ, y) has 9g Lipschitz Hessian jointly with respect to (x, v).
where v > 0 is the perturbation amount and [ ~ (0,, Ta) is randomly drawn from standard Gaussian distribution. At each round, (14) executes two backward processes to get the query of Lss(-) with two perturbation 0 + v and 0 v
Gradient for the upper-level scheduler q. According to (4), V C.yy (g) takes the following form
We will justify the validity of this assumption in two gen- erative bilevel applications in Section 3 after we prove the hyperparameter improvement theorem below.
v C (g) = V Pot [Tss(((()))
+ 7( (V V.LLsM(0 w(g) V VgLisM(0g u(q)))
Theorem 1. Under Assumption I, given any intial hyperpa- rameter zo, letting the inner loop accuracy Ek S S. then there exists stepsize mi S T.p and penalty constant "Ik such that the next updates 1x+1 generated by Algorithm 1 satisfy
where 8y E arg min, g(f, q) and 0- E arg ming L- (0,9) are given by the lower-level and penalty problem solver. For V, (0, u(a)), according to the chain rule, we have
F(Ix+1) - F(()) S
7k
4 [| (xx))] + 6BPTM
VgI.sM(. (0, u(1))= = Duly) V.I.sM u()))-
(15)
Vk
dq
where G (x) Prolx(t-n&C,(e))= is the projected gra- dient and Lp, B (()) are defined in Lemma 1.
Since the score matching loss for popular choices of diffu- sion models has explicit dependency on the generated noisy samples un in the forward process and ue has one-step closed form with respect to the initial sample. V LsIsM(0, (f, u(q)) has closed form for popular choices of diffusion model; see de- tails in Appendix D.2. For V -a()) [Lso (ū(4)) whose explicit expression is not available, we can use similar 70 approaches in (14) with perturbation on noise scheduler q to get the gradient estimator.
"
This theorem demonstrates that when Ganen (fxx) # O. and setting penalty constant " 2 [[G (sr))] the bilevel al- gorithm will have strict descent over the hyper-function F(x). Otherwise, since Vz,,()) - Vf(=)]I ((1/Tk),
4/3B
Gne.n (2z) = 0 is approximately the stationary points of F(x). The proof can be found in Appendix D.3.
Implications on generative bilevel applications. For fine- tuning diffusion models, the KL divergence with respect to probability distribution is strongly convex if p is strongly log- concaee (Vempala & Wibisono, 2019). This condition is usually met in diffusion models (Song et al., 2021a; Ho et al., 2020), as they generate Gaussian distribution with pos- itive definite covariance matrices. Therefore, when reward function is concave (Guo et al., 2024), Assumption 1 holds. In this application, C = since the upper-level gradient estimator does not depend on the inner loop accuracy; see (13). For noise scheduling problem, score matching loss is a composite quadratic function with respect to the noise network (see (24)), which can be viewed as a strongly con- vex function over the parameterized probability (functional) space (Petrulionyte et al., 2024), so that Assumption I holds. In this application, the upper-level gradient estimation er- ror includes an additional term dependent on the error of the 70 estimator, which can be made sufficiently small by appropriately choosing perturbation amount P.
Parametrization for noise scheduler. To further reduce the memory cost and ensure the nondecreasing of q(t), we parameterize the noise scheduler by the commonly used cosine and sigmoid function (Nichol & Dhariwal, 2021: Chen, 2023; Kingma et al., 2021) and optimize the param- eters within these functions instead of directly optimizing {T( In both parametrization, we optimize just scalar parameters, significantly reducing the dimensionality of the optimization from T to 4. Specifically, we use
(t(qe- qa) + 4o)/T + ge
24
((t)= cos
x
(16)
1 ++0 qe
where (s: Ges Tes 9r represent the start, end, offset error, and power effect, respectively. Sigmoid parameterization is defined similarly in (34) and can be found in Appendix F.2. Since q(() should be nondecreasing, we assign 4(1)= = 1- I(t)/T(t- 1). After parameterization, 70 perturbation will be added on qa qe, le.9: instead of directly on q(t).
The full procedure is summarized in Algorithm 6 and can be found in Appendix E.
Therefore, Theorem I applies to both settings, indicating that, regardless of initialization whether from a random
6

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
Baselines
FID 4
CLIP t
Time L
We evaluate bilevel reward fine-tuning Algorithm 5 on the image generation task with complex prompts (Wang et al., 2023; 2024; Clark et al., 2023), comparing it to gradient guidance generation approach in (Guo et al., 2024) com- bined with conventional hyperparameter search methods for tuning A. Inspired by Figure 2, we use CLIP score as the upper-level loss to automatically tune in a bilevel algo- rithm. To rule out the impact of the additive effect on the upper-level loss and lower-level reward, we also compare our approach with the weighted sum method, which naively combines the CLIPP score and lower-level reward r2 (-), with weight selected by grid search. A detailed description of the baselines provided in the Appendix F.2.
Grid search
125.77 +1.3
31.72 #2.1
3.34
Random search
123.70 #4.3
35.70 #1.1
3.37
Bayesian search
140.35 42.4
33.29 12.9
12.7
Weighted sum
116.52 +3.9
36.50 11.1
4.31
Bilevel method
102.82
43.5
39.54
=2.2
3.35
Table 1. Best FID and CLIP score given by different baselines and our method with penalty constant ry = 10 for fine-tuning diffusion model application using synthetic lower-level reward (Yuan et al., 2024). Running time is measured in hours.
Pre-train
Step:1
Step:3
Step:7
Table presents the average FID, CLIP score, and execution time for each method over prompts. The bilevel method outperformed traditional hyperparameter tuning methods, achieving superior FID and CLIP scores. Its compara- ble time complexity arises from requiring backpropagation through the CLIP score, unlike standard methods. For the weighted sum method, which also involves backpropagation through the CLIP score, the bilevel method is faster. More- over, the bilevel method achieved an 11.76% improvement in the FID score and 8.32% improvement in the CLIP score over the best-performing weighted sum method.
(a)
(b)
Iteration
Figure 3. Visualization of images generated at different steps: (a) Images generated with 0.1 become progressively more ab- stract at each step, while (b) images generated with bilevel method (A = 55.5) are more colorful and vivid than the pre-trained images and achieve perfect balance of quality across steps.
Figure 3 shows the generated images over generation iter- ation in Algorithm using À = 0.1 and the optimized by the bilevel approach. The results demonstrate that the entropy strength selected by the bilevel approach achieves a better balance between image quality and realism. Figure shows the optimization process of X, revealing that the optimal value of X varies across different prompts.
point selected through valis-validation the distribution generated using the hyperparameter of bilevel algorithms is guaranteed to perform better. For fine-tuning diffusion models, the bilevel algorithm generates images with higher CLIP scores, while the bilevel noise scheduling algorithm produces images with lower FID scores.
70
Lambda Scale Values
60
50-
6. Numerical Experiments
40
30
In this section, we present the experimental results of the pro- posed -ileeel-iifusion algorithms in two applications: re- ward fine-tuning and noise scheduling for diffusion models, and compare them with baseline hyperparameter optimiza- tion methods: grid search, random search, and Bayesian search (Snock al., 2012).
20
Man
Elephant
10
Horse
100
200
300
400
500
Number of Iterations
Figure 4. Change over iteration given by bilevel approach for different prompts.
6.1. Reward fine-tuning in diffusion models
More visualizations are provided in Figures 9-11 and can be found in Appendix, which visually illustrate the impact of over-aggressive reward optimization, which tends to gen- erate more abstract images (e.g., as observed in the results from grid search and Bayesian search methods).
For this experiment, use the StableDiffusion V1.5 model pre-trat pre-trained model and employ a ResNet-18 archi- tecture (trained on the ImageNet dataset) as the synthetic (lower-level) reward model, following (Yuan et al., 2024), to enhance colorfulness and vibrancy. To enable scalar reward outputs, we replace the final prediction layer of R8 with randomly initialized linear layer.
Furthermore, to showcase the robustness of our approach with respect to different reward functions the lower level,
7

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
We compare our method DDIM combined with baseline hyperparameter optimization methods. We chose greedy grid search over standard grid search as the baseline, as the latter is computationally intensive for searching across multiple hyperparameters. In greedy grid search, we se- quentially optimize parameters based on sensitivity, fixing each optimized parameter before tuning the next. Additional experimental setup can be found in Appendix F.2.
Baselines
FID 4
CLIP t
Time L
Grid search
142.76 +2.4
36.3 #1.7
3.34
Random search
139.21 #3.1
35.20 #3.1
3.37
Bayesian search
153.25 11.2
34.98 12.9
12.7
Weighted sum
140.56 +1.3
35.40 #2.1
4.31
Bilevel method
137.23
+1.7
37,10 =3.2
3.35
Table 2. Best FID and CLIP score given by different baselines and our method with penalty constant ry = 10 fine-tuning diffusion model application using HPSv2 (Wu et al., 2023) the lower-level reward. Running time is measured in hours.
Bilevel noise scheduling algorithm in Algorithm 6 alternates between optimizing the weights 0 in U-Net and finding the best noise scheduler (t) online, so that is computationally efficient and outperforms fixed noise schedulers. Figure shows the learned hyperparameter 9s. 9e 9r. 9e by bilevel optimization versus iteration l: and the corresponding q(t) at four timesteps. We observe the parameters are nontriv- ial: 9r is the most influential factor, varying significantly throughout the training process; start ys and end le update inversely and are utilized more extensively at the beginning of training; and ge increases progressively over the course of training. These parameters determine the noise scheduler (t), which introduces more noise at the beginning and the latter middle stages of training step k.
we test the performance of each method for benchmarking reward function, HPSv2 (Wu et al., 2023), as the lower- level reward function. Comparisons of our method and different baselines are given in Table Similar to the results given by the synthetic reward function (Yuan al., 2024), bilevel approach also outperform other baselines in terms of image quality in comparable time complexity. We also provided the visualization of the generated images in Figure 5. While all HPO on fine-tuned models enhance the aesthetic, clarity and sharpness compared to the pre-trained image, the random, grid, Bayesian search, and weighted sum approaches messed up the legs and trunk, and fail to generate the right number of elephant's legs, In comparison, proposed bilevel approach not only generate colorful images, but also preserve correct elephant biological feature.
9.
1.0
300
Values
4
0
0.0
200
400
500
1000
k
1
Figure 6. Varying hyperparameters start 9.. end qe. power q- and offset in cosine parameterization learned by bilevel method along the training steps and corresponding noise scheduler ((t) at iteration & = 0, 100, 300, 469.
Pre-trained
Random
Grid
Table 3 presents the best FID, inception score (IS), and time complexity achieved by each method. The bilevel method achieves comparable performance with the hyperparameter optimization baselines in both FID and IS while being 6x time faster. Generated images by each method are shown in Figure 7. While the Bayesian approach achieves rela- tively better FID and IS, it tends to focus on generating simpler numbers, such as 1 and 7. In contrast, the images produced by the bilevel method show excellent diversity across numbers 9 while maintaining image fidelity.
Bayesian
Weighted sum
Bilevel
Figure 5. Visualization of images generated by different methods using the prompt "clephant" and HPSv2 reward (Wu et al., 2023).
6.2. Noise scheduling in diffusion models
7. Related works
We evaluated our bilevel noise scheduling method, detailed in Algorithm 6, paired with DDIM backward sampling for the image generation on the MNIST dataset. We trained a U-Net with 178 layers and 10 parameters following the github repository I. We considered both cosine and sigmoid parametrization and tuned 4 parameters qa: ger Gr. qe jointly.
Fine-tuning diffusion models. Fine-tuning diffusion mod- els aims adapt pre-trained models to boost the reward on downstream tasks. Methods in this domain include directly backpropagating the reward (Clark et al., 2024), RL-based fine- tuning (Fan et al., 2024; Black et al., 2023), direct la- tent optimization (Tang et al., 2024; Wallace et al., 2023; Hoogeboom et al., 2023), guidance-bar -based approach (Guo
8

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
Cosine
Sigmoid
6
0
9
7
7
3
6
7
7
2
06 2 A 0 1 ?
Methods
3
8
2
4
8
0
5
5
7
9
4
7
1 2 3 3 e
FID |
IS 1
Time
FID
IS
Time
523
3
3 0 2 3
3
0
5
3
0
x
3
4
9
6 6 5
Grid search
67.30
1.76
31.53
65.31
1.76
39.12
3
8
4 3
3
5
7
6
8 7
8
4
5
Random scarch
68.97
1.61
29.62
66.02
1.69
35.94
3
1
1
6
4
9
4
9
7
/
/
Bayesian seurch
67.16
1.69
26.85
65.17
1.65
29.13
7
0
3
8
0
3
5
1
0
0
DDIM (default)
103.27
1.43
1.59
85.79
1.54
1.78
Bilevel method
65.41
1.78
3.88
65.16
1.79
3.94
(a) Bilevel
(b) DDIM (default)
(c) Bayesian
Figure 7. Visualization the final generated images by different
Table 3. Comparison of FID, IS, and running time (in hours) for different baselines and our method for the noise scheduling appli- cation with cosine and sigmoid parameterization. Default DDIM parameters are from (Nichol & Dhariwal, 2021) for cosine and (Vidhya, 2024) for sigmoid parameterization.
methods using cosine parameterization.
noise scheduling for training diffusion model from scratch. For fine-tuning, we propose an inference-only bilevel ap- proach to guide the diffusion model toward the target dis- tribution and leverage the closed-form of KL divergence to update the entropy strength. For training from scratch, we optimize the parameters of the noise distribution match the true noise and use zeroth-order optimization to deter- mine the optimal noise scheduler for generating high-quality images in the backward process "on the fly." Experiments demonstrate the effectiveness of the proposed method.
et al., 2024; Chung et al., 2022; Bansal et al., 2023) and optimal control (Uehara et al., 2024). Although entropy reg- ularization is often incorporated into the reward prevent over-optimization, existing work has explored designing an efficient bilevel method to tune its strength.
Noise scheduling in diffusion models. Noise schedule is crucial to balance the computational efficiency with data fi- delity during image generation. Early works, such as DDPM (Ho et al., 2020), employed simple linear schedules for noise variance, while Nichol & Dhariwal (2021) and Kingma al. (2021) introduced cosine and sigmoid schedules enhance performance. Recent studies (Lin et al., 2024; Chen, 2023) have highlighted limitations in traditional noise schedules and proposed new parameterization to improve the image quality. Notably, Sahoo et al. (2024) learned the noise sched- uler by optimizing the log-likelihood, which yields tighter lower bound (ELBO) and thus improves the generation qual- ity of the diffusion model. However, none of the prior works considered using bilevel optimization to automatically learn the noise schedule for directly optimizing sample quality.
A Acknowledgment
The work of Q. Xiao and T. Chen was supported by National Science Foundation (NSF) McDL-SCALE project 2401297, NSF project 2412486, and the Cisco Research Award.
Impact Statement
This paper aims to advance diffusion models with bilevel generative optimization, providing a novel approach hy- perparameter tuning for improved image generation. By addressing key challenges in fine-tuning diffusion model and noise scheduling. our work contributes to the broader development of more efficient and adaptive generative mod- els. Potential societal impacts include applications in cre- ative content generation, data augmentation, and machine learning- simulations. While we acknowledge the pos- sibility of unintended uses, do not identify any specific societal risks that need to be highlighted in this context.
Bilevel hyperparameter optimization. Bilevel opti- mization has been explored as an efficient hyperparame- ter optimization framework, including hypernetwork search (Mackay al., 2019; Liu al., 2019), iyper-representtaiin (Franceschi et al., 2018), regularization learning (Shaban et al., 2019) and data reweighting (Shaban et al., 2019; Franceschi et al., 2017). Recently, it has been explored in federated learning (Tarzanagh et al., 2022) and LLM fine- tuning (Shen et al., 2025a: Zakarias et al., 2024). None of the existing works have explored hyperparameter optimiza- tion in diffusion models, and the methods proposed so far inapplicable due the difiise-dimensionall probability space and the high computational cost of sampling.
References
Arbel, M. and Mairal, J. Amortized implicit differentiation for stochastic bilevel optimization. In Proc. International Conference on Learning Representations, virtual, 2022.
Bansal, A., Chu, H.-M., Schwarzschild, A., Sengupta, S., Goldblum, M., Geiping, J., and Goldstein, T. Universal guidance for diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 843- 352, 2023.
8. Conclusions
In this paper, we analyze two types of generative bilevel hyperparameter optimization problems in diffusion models: fine-tuning a diffusion model (with a pre-trained model) and
Black, K., Janner, M., Du, Y., Kostrikov, I., and Levine, S.

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
Training diffusion models with reinforcement learning- arXiv preprint (.tI::2305./3301 2023.
Franceschi, L., Frasconi, P., Salzo, S., Grazzi, R., and Pontil, M. Bilevel programming for hyperparameter optimization and meta-learning. In Proc. International Conference on Machine Learning, Stockholm, Sweden, 2018.
Bracken, and McGill, J. T. Mathematical programs with optimization problems in the constraints. Operations Research, 21(1):37-44, 1973.
Gao, L., Schulman, J., and Hilton, J. Scaling laws for reward model overoptimization. In Proc. International Confer- ence on Machine Learning. pp. 10835-10866, Honolulu, HI, 2023.
Chen, L.., Xu, J., and Zhang, J. On finding small hyper- gradients in bilevel optimization: Hardness results and improved analysis. In The Thirty Seventh Annual Confer- ence on Learning Theory, pp. 947-980. PMLR, 2024.
Ghadimi, S. and Wang, M. Approximation methods for bilevel programming. arXiv preprint arXiv: 1802.0226, 2018.
Chen, T. On the importance of noise scheduling for diffusion models. arXiv preprint arXiv:2301.10977 2023.
Chen, T., Sun, Y., and Yin, W. Closing the gap: Tighter analysis of alternating stochastic gradient methods for bilevel problems. In Proc. Advances in Neural Informa- tion Processing Systems, virtual, 2021.
Gong, S., Zhang, S., Yang. J., Dai, D., and Schiele, B. Bi- level alignment for cross-domain crowd counting. In Proc. IEEE/CVF Conference on Computer Vision and Pattern Recognition, New Orleans, LA, 2022.
Chung, H., Kim, J., Mccann, M. T., Klasky, M. L., and Ye, C. Diffusion posterior sampling for general noisy in- verse problems. arXiv preprint arXiv:2209. 14687, 2022.
Grazzi, R., Franceschi, L., Pontil, M., and Salzo, S. On the iteration complexity of hypergradient computation. In Proc. International Conference on Machine Learning, virtual, 2020.
Clark, K., Vicol, P., Swersky, K., and Fleet, J. Directly fine-tuning diffusion models on differentiable rewards. arXiv preprint arXiv:2309. 17400, 2023.
Guo, Y., Yuan, H., Yang, Y., Chen, M., and Wang, M. Gra- dient guidance for diffusion models: An optimization perspective. arXiv preprint arXiv:2404.14745 2024.
Clark, K., Vicol, P., Swersky, K., and Fleet, J. Directly fine-tuning diffusion models on differentiable rewards. In Proc. International Conference on Learning Representa- tions, Vienna, Austria, 2024.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. In Proc. Advances in Neural Information Processing Systems, virtual, 2020.
Croitoru, F-A., Hondru, V., Ionescu, R. T., and Shah, M. Diffusion models in vision: A survey. IEEE Transac- lions Pattern Analysis and Machine Intelligence, 45 (9):10850-1086) 2023.
Ho, J., Saharia, C., Chan, W., Fleet, D. J., Norouzi, M., and Salimans, T. Cascaded diffusion models for high fidelity image generation. Journal of Machine Learning Research, 23(47): 2022.
Denker, A., Vargas, F,, Padhy, S., Didi, K., Mathis, S. V., Barbano, R., Dutordoir, V., Mathieu, E., Komorowska, J., and Lio, P. Deft: Efficient fine-tuning of diffusion models by learning the generalised h-transform. In Proc. Advances in Neural Information Processing Systems, Van- couver, BC, Canada, 2024.
Hong, M. Wai, H.-T., Wang, Z., and Yang. Z. A two- timescale stochastic algorithm framework for bilevel op- timization: Complexity analysis and application to actor- critic. SIAM Journal on Optimization, 33(1):147-180, 2023.
Fan, Y., Watkins, O., Du, Y., Liu, H., Ryu, M., Boutilier, C., Abbeel, P., Ghavamzadeh, M., Lee, K., and Lee, K. Reinforcement learning for fine-tuning text-to-image dif- fusion models. In Proc. Advances in Neural Information Processing Systems, Vancouver, BC, Canada, 2024.
Hoogeboom, E., Heek. J., and Salimans, T. simple diffu- sion: End-to-end diffusion for high resolution images. In Proc. International Conference on Machine Learning, Honolulu, HI, 2023.
Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta- learning for fast adaptation of deep networks. In Proc. International Conference Machine Learning, Sydney, Australia, 2017
Ji, K., Yang, J., and Liang, Y. Bilevel optimization: Con- vergence analysis and enhanced design. In Proc. Interna- tional Conference Machine Learning, virtual, 2021.
Jiang, I.., Xiao, Q., Tenorio, V. M., Real- F,, Marques, A. G., and Chen, T. A primal-dual-assien penalty ap- proach to bilevel optimization with coupled constraints. In Proc. Advances in Neural Information Processing Sys- tems, Vancouver, BC, Canada, 2024.
Franceschi, L., Donini, M., Frasconi, P., and Pontil, M. Forward and reverse gradient-baat based hyperparameter opti- mization. In Proc. International Conference on Machine Learning, Sydney, Australia, 2017.
10

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
Jing, B., Corso, G., Chang, J., Barzilay, R., and Jaakkola, T. Torsional diffusion for molecular conformer genera- tion. In Proc. Advances in Neural Information Processing Systems, New Orleans, LA, 2022.
Liu, R., Liu, Y., Yao, W., Zeng, S., and Zhang. J. Averaged method of multipliers for bi-level optimization without lower-level strong convexity. In Proc. International Con- ference on Machine Learning, Honolulu, HI, 2023b.
Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion-based generative models. In Proc. Advances in Neural Information Processing Sys- tems, New Orleans, LA, 2022.
Lu, Z. and Mei, S. First-order penalty methods for bilevel optimization. arXiv preprint arXiv:230.001716, 2023.
Mackay, M., Vicol, P., Lorraine, J., Duvenaud, D., and Grosse, R. Self-tuning networks: Bilevel optimization of hyperparameters using structured best-response func- tions. In Proc. International Conference on Learning Representations, 2019.
Khanduri, P., Zeng. S., Hong, M., Wai, H.-T., Wang, Z., and Yang, Z. A near-optimal algorithm for stochastic bilevel optimization via double- in Neural Information Processing Systems, virtual, 2021.
In Proc. Advances
Maclaurin, D., Duvenaud, D., and Adams, R. Gradient- based hyperparameter optimization through reversible learning. In Proc. International Conference on Machine Learning, Lille, France, 2015.
Kingma, D., Salimans, T., Poole, B., and Ho, J. Varia- tional diffusion models. In Proc. Advances in Neural Information Processing Systems, virtual, 2021.
Kingma, D. P. Adam: A method for stochastic optimization. In Proc. International Conference on Learning Represen- tations, 2015.
Marion, P., Korba, A., Bartleut P., Blondel, M., De Bor- toli, V., Doucet, A., Llinares-López, F,, Paquette, C., and Berthet, Q. Implicit diffusion: Efficient opti- mization through stochastic sampling. arXiv preprint arXi::2402.0542 054666 2024.
Kwon, J., Kwon, D., Wright, S., and Nowak, R. D. A fully first-order method for stochastic bilevel optimization. In Proc. International Conference on Machine Learning, Honolulu, HI, 2023.
Mathiasen, A. and Hvilshøj. R. through frechet inception distance. arXiv preprint 20xhv:2009.14075, 2020.
Backpropagating
Kwon, J.. Kwon, D., Wright, S., and Nowak, R. On penalty methods for nonconvex bilevel optimization and first- order stochastic approximation. In Proc. International Conference on Learning Representations, Vienna, Aus- tria, 2024.
Nesterov, Y. and Spokoiny, V. Random gradient- mini-
mization of convex functions. Foundations of Compula- tional Mathematics, 17(2)5227-566 2017.
Nesterov, Y. et al. Lectures on convex optimization, volume 137. Springer, 2018
Li, Gu, B., and Huang, H. A fully single loop algorithm for bilevel optimization without hessian inverse. In Proc. Association for the Advancement of Artificial Intelligence, virtual, 2022.
Nichol, A. Q. and Dhariwal, Improved denoising diffusion probabilistic models. In Proc. International Conference on Machine Learning, pp. 8162 8171, 2021.
Lin, S., Liu, B., Li, J., and Yang, X. Common diffusion noise schedules and sample steps are flawed. In Proceedings of the IEEE/CVF winter conference on applications of computer vision, pp. 5404-544 5,11, 2024.
Pedregosa, F. Hyperparameter optimization with approx- imate gradient. In Proc. International Conference on Machine Learning, New York City, NY, 2016.
Petrulionyte, I., Mairal, J., and Arbel, M. Functional bilevel optimization for machine learning. In Proc. Advances in Neural Information Processing Systems, Vancouver, BC, Canada, 2024.
Liu, B., Ye, M., Wright, S., Stone, P., et al. Bome! bilevel optimization made easy: A simple first-order approach. In Proc. Advances in Neural Information Processing Sys- tems, New Orleans, LA, 2022.
Qin, P., Zhang, R., and Xie, P. Bidora: Bi-level optimization- based weight-deccomposed low-rank adaptation. arXiv preprint arXiv:2410.09755, 2024.
Liu, H., Simonyan, K., and Yang, Y. DARTS: Differentiable architecture search. In Proc. International Conference on Learning Representations, New Orleans, LA,, 2019.
Ronneberger, O., Fischer, P., and Brox, T. U-net: Con- volutional networks for biomedical image segmenta- tion. In Medical image computing and computer-assisted Intervention-MICC 2015: 18th international confer- ence, Munich, Germany, October 5-9, 2015, proceedings. part III 18, pp. 234 241, 2015.
Liu, H., Chen, Z., Yuan, Y., Mei, X., Liu, X., Mandic, D., Wang, W., and Plumbley, M. D. Audioldm: Text-oo-audio generation with latent diffusion models. In Proc. Inter- national Conference on Machine Learning, pp. 21450- 21474, Honolulu, HI, 2023a.

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
Sahoo, S., Gokaslan, A., De Sa, C. M., and Kuleshov, V. Dif- fusion models with learned adaptive noise. In Proc. Ad- vances in Neural Information P Processing Systems, 2024.
International Conference on Learning Representations, virtual, 2021b.
Stadie, B., Zhang, L.., and Ba, J. Learning intrinsic rewards as a bi-level optimization problem. In Conference on Uncertainty Artificial Intelligence, virtual, 2020.
Sambharya, R., Hall, G., Amos, B., and Stellato, B. Learn- ing to warm-start exeet-point optimization algorithms. Journal of Machine Learning Research, 25(166):1-46, 2024.
Tang, W. Fine-tuning of diffusion models via stochastic con-
trol: entropy regularization and beyond. arXiv preprint arXiv:2403.06273 2024.
Scellier, B. A deep learning theory for neural networks grounded in physics. arXiv preprint arXiv:2103.0995. 2021.
Tang, Z., Peng. J., Tang, J., Hong, M., Wang. F., and Chang, T.-H. Tuning-free free alignment of diffusion models with di- rect noise optimization, arXiv preprint arXiv:2405.1888 2024.
Scellier, B. and Bengio, Y. Equilibrium propagation: Bridg- ing the gap between energy- based models and backprop- agation. Frontiers in computational neuroscience, 11:24, 2017.
Tarzanagh, D. A., Li, M., Thrampoulidis, C., and Oymak, S. FEDNEST: Federated bilevel, minimax, and compo- sitional optimization. In Proc. International Conference on Machine Learning, Baltimore, MD, 2022.
Seitzer, M. pytorch-fid: FID Score for PyTorch. https:// g thub. com/ / msse wwseetzere/ptttttoo t August 2020. Version 0.3.0.
Uchara, M., Zhao, Y., Black, K., Hajiramezanali, E., Scalia, G., Diamant, N. L.., Tseng, A. M., Biancalani, T., and Levine, S. Fine-tuning of continuous-time diffusion models as rngropy-regulariize control. arXiv preprint arXiv:2402.1 15194, 2024.
Shaban, A., Cheng, C.-A., Hatch, N., and Boots, Trun- cated paok-propagation for bilevel optimization. In Proc. International Conference Artificial Intelligence and Statistics, Naha, Japan, 2019.
Shamir, O. An optimal algorithm for bandit and zero-order convex optimization with two-point feedback. The Jour- nal of Machine Learning Research, 18(1-1):1703-1713, 2017.
Vempala, S. and Wibisono, A. Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices. In Proc. Advances in Neural Information Processing Sys- tems, Vancouver, Canada, 2019.
Shen, H., Yang, Z., and Chen, T. Principled penalty-based methods for bilevel reinforcement learning and RLHF. In Proc. International Conference on Machine Learning, Vienna, Austria, 2024.
VIcol, P., Lorraine, J. P., Pedregosa, F., Duvenaud, D., and Grosse, R. B. On implicit bias in overparameterized bilevel optimization. Proc. International Conference on Machine Learning, Baltimore, MD, 2022.
Shen, H., Chen, P.-Y., Das, P., and Chen, T. Seal: Safety- enhanced aligned L.LM fine-tuning via bilevel data se- lection. In Proc. International Conference on Learning Representations, 2025a.
Vidhya, A. Noise schedules in stable diffusion. https:// www. ana yt csvi dhya cwt 20tmm///l///2220777/77// noise scweduleeeeeesseee -nn-inn---- 2024.
Shen, H., Xiao, Q. and Chen, T. On penalty-based bilevel gradient descent method. Mathematical Programming, pp. 1-51, 2025b.
Wallace, B., Gokul, A., Ermon, S., and Naik, N. End- to-end diffusion latent optimization improves classifier guidance. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 7280-7290, 2023.
Snoek, J., Larochelle, H., and Adams,
Practical bayesian optimization of machine learning algorithms. In Proc. Advances in Neural Information PrciionmPrrrre Processing Sys- tems, 2012.
Wang, R., Liu, T., Hsieh, C.-.,, and Gong, B. On discrete prompt optimization for diffusion models. arXiv preprint 1.XI::2077.01066 2024.
Song, J., Meng, C., and Ermon, S. Denoising diffusion implicit models. In Proc. International Conference on Learning Representations, virtual, 2021a.
Wang, 7., Jiang, Y., Lu, Y., He, P., Chen, W., Wang, Z.,, and Zhou, M. In-context learning unlocked for diffusion models. Advances in Neural Information Processing Systems, pp. 8542-8562, 2023.
Song. Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. baore-based generative mod- eling through stochastic differential equations. In Proc.
Wu, L., Gong, C., Liu, X., Ye, M., and Liu, Q. Diffusion- based molecule generation with informative prior bridges.
12

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
In Advances in Neural Information Processing Systems, New Orleans, LA, 2022.
Wu, X., Hao, Y., Sun, K., Chen, Y.. Zhu, F., Zhao, R., and Li, H. Human preference score v2: solid benchmark for evaluating human preferences of text-to-image synthesis. arXiv preprini arXiv:2306.0934) 2023.
Xiao, Q. Lu, S., and Chen, A generalized alternat- ing method for bilevel optimization under the polyak- tojasiewicz condition. In Proc. Advances in Neural Infor- mation Processing Systems, New Orleans, LA, 2023.
Yang, D., Yu, J., Wang, H., Wang, W., Weng, C., Zou, Y., and Yu, D. Diffsound: Discrete diffusion model for text- to-sound generation. IEEE/ACM Transactions on Audio, Speech, and Language Processing, 31:1720-1733, 2023.
Yang, J., Ji, K., and Liang, Y. Provably faster algorithms for bilevel optimization. In Proc. Advances Neural Information Processing Systems, virtual, 2021.
Yang, K., Tao, J., Lyu, J., Ge, C., Chen, J., Shen, W., Zhu, X., and Li, X. Using human feedback to fine-tune diffusion models without any reward model. In Proc. IEEE/CVF Conference on Computer Vision and Pattern Recognition, Seattle. WA. 2024.
Yao, W., Yin, H., Zeng, S., and Zhang, J. Overcoming lower-level constraints in bilevel optimization: A novel approach with regularized gap functions. arXiv preprint arXiv:2406.00 01 09900099992010 2024.
Yuan, H., Huang, K., Ni, C., Chen, M., and Wang, M. conditional diffusion: Provable distribu- tion estimation and reward improvement. In Proc. Ad- vances in Neural Information Processing Systems, Van- couver, BC, Canada, 2024.
Zakarias, G. W., Hansen, K., and Tan, Z.-H. Bissl: Bilevel optimization for self-supervised pre-training and fine- tuning, arXiv preprint arXhv:2410.02356 2024.
Zhang, Y., Zhang, G., Khanduri, P., Hong, M., Chang, S., and Liu, S. Revisiting and advancing fast adversarial training through the lens of bi-level optimization. Proc. International Conference on Machine Learning, Baltimore, MD, 2022.
Zucchet, N. and Sacramento, Beyond backpropagation: bilevel optimization through implicit differentiation and cquilibrium propagation. Neural Computation, 34(12) (4(12): 2309-2346, 2022.
13

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
Appendix for "A First-order Generative Bilevel Optimization Framework for Diffusion Models"
A. Additional related works
Bilevel optimization methods. Bilevel optimization has a long history that dates back to (Bracken & McGill, 1973). Recent
efforts have focused on developing efficient gradient-bar -based bilevel optimization methods with aon-asymptotic convergence guarantees, inspired by works such as (Ghadimi & Wang, 2018; et al., 2021; Hong et al., 2023; Chen et al., 2021). Since the gradient of the bilevel nested objective depends on the Hessian of the lower-level objective, existing literature proposed different Hessian inversion approximation methods including unrolling differentiation (Franceschi al., 2017;
2018; Grazzi et al., 2020), implicit differentiation (Chen et al., 2021; Ghadimi & Wang, 2018; Hong al., 2023; Pedregosa, 2016; Khanduri et al., 2021). conjugate gradients (Ji et al., 2021; Yang et al., 2021) and its warm-started single-loop versions (Arbel & Mairal, 2022; Li et al., 2022; Liu et al., 2023b; Xiao et al., 2023), and equilibrium backpropagation (Scellier& & Bengio, 2017; Scellier, 2021); see (Zucchet & Sacramento, 2022) for a comparison. Among these methods, equilibrium backpropagation stands out as fully first-order approach, valued for its balance of efficiency, robustness, and simplicity. Building on this principle, recent works have extended its applicability from the strongly convex setting to convex, nonconvex and constrained settings by reformulating the bilevel optimization problem as single-level penalty problem and solving via first-order approaches (Shen et al., 2025b; Liu et al., 2022; Kwon al., 2023; 2024; Chen et al., 2024; Lu & Mei, 2023; Jiang et al., 2024; Yao et al., 2024).
B. Background on bilevel optimization
In this section, we review some background knowledge for first order bilevel optimization.
The differentiability of the penalty problem relies the differentiability of the value function (x)) which is established through the extended Danskin theorem (Shen et al., 2025b, Proposition 4). Specifically, the gradient of value function takes
V (x) v=g(z,y"), Vy" E S(z).
(17)
This enables to solve (2) by gradient-based approach. Similarly, by applying the extended Danskin theorem to the penalty function (K (Kwon et al., 2024), we know
with V2" E (xx
which can be further rewritten according to (17) as
VC,()) V (V veg(r,z))),
(18)
Moreover, the following lemma shows that the penalty objective is a proxy of original bilevel hyper-function F(x). Lemma ((Kwon et al., 2023, Lemma 3.1)). Under Assumption 1, let n
21 we have F(x) is LE- smooth and
B [VF(x) (x)( S ?
where LF (1+
(es Pig
(1,1
Pa Bg +
21 21.olgl
O(1/k²) and B
(11,1 + 21ole12 O(1/x3) and
"3
f4 is the condition number. Mg
This lemma indicates that (z) is an approximation of FF(z)) with error controlled by enlarging penalty constant ry.
C. Background on diffusion models
In this section, we connect continuous to discrete diffusion model to enable the derivation of the closed-form gradient of the score matching function with respect to the noise scheduler in the discrete diffusion model implementation.
Denoising diffusion probabilistic model (DDPM) (Ho et al., 2020) and Denoising diffusion implicit model (DDIM) (Song al., 2021a) provide standard ways to discretize the continuous SDE in (5) and (6). To be self-contained, we provide a
14

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
derivation of connection between them. Let us recall the continuous forward process in (5)
du, - ((( dt + Va(t)dW,
(13)
which gives the following transition probabilities
p(ugluo) = N (moe 10" eigl da If( a(t)e- -0" d(e)dat))) N (uoe- So Lgddd (1 c1 s(s)do) I.))
See also (Song et al., 2021b, Appendix B) and (Denker et al., 2024, Appendix A), Therefore, by defining q(t) = e Se g(o)ds we get the form in DDPM (Ho et al., 2020)
p(u/luo) N (N (1 (())
(19)
Since the approximation
holds well when x is small, we have discrete approximation of q(t) as
N-1
a(t) e (())da as II (1-)((z))1))
7=
By choosing At 1, we get the expression of discrete DDPM in (Ho et al., 2020) follows.
Forward process. Given a data point sampled from data distribution ug ~ Pdatas the forward process in DDPM generates a sequence of samples u1, u2,..., ur by gradually adding noise
p(11 un)
I(uelee--)) p(u,lue1) =
N(V =N (6
(20)
t=1
where {- corresponds to the noise scheduler in discrete DDPM. (20) can be further expressed as
p(ur|uo) (V (1 L.)
(21)
where 9. - II 9s) is the variance scheduler defined by the noise scheduler {q. IT= Backward process of DDPM. The backward process aims to recover uo from ur by iteratively denoising
po(uu))) p(ur) II po (u--1/we..
(22)
where p(ur) -O, Ia) and each po is modeled as Gaussian distribution parameterized by Po (1)) = (Ne t), (u L)
with the mean and variance learned by optimizing the score-matee moteemnaching objective.
Score matching. In the discrete DDPM, matching loss also takes a simpler form. To learn μi and dg, we first estimate
the backward probability given the initial state using the Gaussian kernel estimation as follows
Pe(u--1lu, uo) -N(u (ue, o
where
H4 (14, un)
99--1
20
VI
qi (1
q
U1
(a)
(14
14
qt VI
8
(23a)
Tt
1 1- qt
VI
qt
gt
and
2 oi
1
1--1 41
(23b)
where (a) is earned by reparameterizing (21) as u1 62) voruo (1 )o for 8, ~ NO0 Ia). As b4 is proportional to 6, we can fit a neural network to proxy by optimizing the score matching loss in (7) in the following simplified form with explicit dependence on noise scheduler
(0 () = Euo... ||0- 50 (Vieno + /1 11-000
q:o,
(24)
15

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
where 8g is neural network approximator (e.g. U-Net) intended to predict Gaussian noise ó from We--
By optimizing the score matching objective LsM(0, q) with respect to 0, we obtain the proxy of òø and using de instead of 6; in (23a), we can sample the backward process by
1
1
(u
50) + 0,U
(25)
V
vI
qr
with v ~ Ia). The full training and backward sampling process in DDPM is summarized in Algorithm 3 and 4.
Algorithm 3 Score network training
Algorithm Backward sampling
1: repeat
1: {up M
mm=1 (O)
2: draw Uu Jm=1 ~ Pdata
1M
2: for l = T.
do
3: {fmfm=l ~ (niform((T))
LM
3: {u"h M
Jm=1 ~ NO, Ia) if > 1, else am = 0
4: {fm]m=) Io)
5: Take gradient descent step on Vo 1 MM
M m=1 115
4: " vT-qe i
9 60 (1 :1) + 0.a
"ol/ "0" + vI qt.m dm m
end for
6: until converged
6: return
Σ
M "0
DDIM (Song et al., 2021a) uses the same forward process and score network training as DDPM, but employs deterministic backward sampling strategy and eliminates redundant sampling steps to further accelerate the backward process as follows.
Backward process of DDIM. Letting {t,} be some selected time steps from [), (25) generalized by
Ur, I qr, (1+)) (ue,)
U4
14:11
+ V q; of, 66 (ux,) + V
(26)
Vtn,
which recovers DDPM in (25) when a; v(1 9-11)-1))))) q// (1 Φt) in (23b) and without skipping. i.e. t2 t, and the resulting deterministic model when = 0 is called DDIM. In DDIM, time steps (44) are selected using either linear
(t, = [ci] for some e) or a quadratic (t, [ci²] for some e) strategy. With these designs, backward sampling steps of DDIM can reduced from 1000 DDPM to 50 10 (Song et al., 2021a).
D. Theoretical analysis
In this section, we present the closed-form of (X, for fine-tuning diffusion model application, the gradient of the score matching function with respect noise scheduler in discrete diffusion models, and the theoretical guarantee of Algorithm 1.
D.1. Upper-level gradient: Proof for Proposition 1
Proof: According to (Tang, 2024, Equation (3.7)), we have the closed form of KL divergence of fine-tuning distribution and pre-trained distribution follows.
ru(
+ log Emn-pen [era(2//A]
KL (p" (A)()a)p))ma))
ri (
KL (p. (A)llaama)
+ log Eumpen
(27)
A
Then the proof can be obtained by plugging the closed-form of KL. divergence in (27) into (12). That is,
v (A) (KL (A) KT.(p" (A)llpamm))
Mr (1(2//7 (2(1))
Tr((u)
+ log Eusppn
Eumppn
log Europan
[
A
[} (r))] log Europan
e
r2(2]
+ y log E
Eumpon
which completes the proof.
16

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
D.2. Explicit lower-level noise scheduler's gradient in discrete diffusion models
In this section, we derive the explicit gradient expression of the score matching function with respect to the noise scheduler.
Since score matching loss in both DDPM and DDIM takes the form in (15), the noise scheduler's gradient in score matching objective can be earned by the chain rule
1 M Σ
Juyy
V VgLLM((,u(g))
V V.LSM(O, ug)
(28)
8q
where u = wi with fm uniformly chosen from & E [T] = , generated using noise scheduler q [91..... 9r]. In this way, using the reparameterization u vquuo + (1 G6 and the relation of G. = II. (1 - -o)), we have for any S tm.
, T and subscript q means this forward sample is
Ouy
Ouf..
aqe.. Jul 8q.. 0qu...
qtm
uo
8
T
(29)
0q:
aq.
2 9t
On the other hand, the gradient of the score-matching function with respect to sample w takes the form of
06 !
@LSM (0, ugy) (u
VuLsM um") Ouf,.
(ōo(ui7))) - 6mm)
(30)
06e (uf
(8
is the derivative of the score network with respect to the input samples that is directly obtainable via auto differentiation library in Pytorch, By plugging the above closed forms of partial derivative in (29) and (30) into (28), get for any E (T).
where 4 = V 400 + (1 T. 6m and the first term
Oupa
Oug
Va L0, u(4)) x
Σ
V v. (O um)
[M
dq:.
Mi [m/
21)
2
Σ
Dtm
5m
0og ( (u(mm)
(oo(u? 8mn.
(31)
[Me
qt
2tm
Ouf
M,::(m]t..224
In practice, we do not need to manually implement the closed form in (31), as PyTorch's auto-differentiation handles it automatically. The derivation in this section highlights the low computational cost of auto-differentiation, as only depends the U-Net structure, and this differential is commonly used in gradient guidance diffusion models (Guo et al., 2024: Bansal et al., 2023).
06g 06.(12 (1 )
Bum
D.3. Descent theorem: Proof of Theorem 1
Proof: By Taylor expansion and the Lr smoothness of F(x), we have
LF 2 | 2211
F(xx41) S F(xx) [F(())) -xx)+
LF Projx 17x V.. (xx)) xx21
< F(xx) ( 7h V VL (xx)) -xx) +
2
LF Projx 7h V VC (Ix)) 2x
= = + ( (roxx - v C, (xx)) - xx) +
2
+ +VFF((t) - V.ž (roi)) 7h V VL; (xx)) xx)
1/6 V VCg (xx)) - zn) + LF Projj (xx Th VL (Tx)) xx
< F(xx) ( (w+.), Projx
2
+ 2a
1
( - (xx)
a Projx (x= - 7%6 V (xx)) - xxl1
(4)
1
F(zL)
|| Projx The V (xx)) xx 21 m> VF (xx) - VCT
4rk
17

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
1 47jk
2B
F(x))
P|oooxxx (2 77 ( =x11² 1/++
+
where (a) comes from the descent lemma of projected gradient (e.g. (Nesterov et al., 2018, Theorem 2.2.13)) and choosing 2Tk and (6) is because
2
[F(xz) (x)) ( -V 2 (2L) (2 (TW)) _
4
7
where the last inequality comes from Lemma and the estimation error of penalty and lower-level problem. By defining
the projected gradient as (n)( = (rolx((--V and letting S %
B we get the conclusion.
E. Complete Algorithms
In this section, we present the complete algorithms with additional details for gradient guided diffusion model for (single- level) generative optimization (Guo et al., 2024), and the bilevel diffusion algorithm for fine-tuning and noise scheduling problem proposed in this work.
E.1. Guided Diffusion algorithm for Generative Optimization
To generate samples that optimize a given reward function r, we can iteratively apply the backward SDE in (6) with the pre-trained score network 8o and the guidance defined as
G(ü,, t;r) = = - u ((ūx + (o(t)s((u))) t)))
(32)
Vg(()
where g a gradient vector associated with the reward function r()) evaluated the current sample ü,, is a given target reward value that increase along the optimization, ((() exp(- So ((s)ds) h(()) = 1 @(t) are the mean and variance of t-th sample and p(t) the tunning parameter.
We can iteratively update the gradient guidance to steer the sample generation process maximize the reward function. Specifically, at each iteration nı, the backward SDE (8) is stimulated using the current gradient guidance from (32), evaluated at the current samples, to generate new samples. Subsequently, the gradient guidance term is updated at the newly generated samples. After N steps of guidance updates, we are able to generate samples approximately that follow the target distribution with (log( effectively minimizing r(-) while incorporating regularization to align with the pre-trained model (Guo et al., 2024). The complete algorithm for guided diffusion model for generative optimization outlined in Algorithm 5.
E.2. Bilevel fine-tuning algorithm
Bilevel fine-tuning diffusion algorithm is summarized in Algorithm 2 with the following upper-level gradient estimation.
Monte Carlo estimation of upper-level gradient.
The upper-level gradient can be estimated from the following way
1 Mo
Mo
Ma
(C = AMo
ri (ūm)) - y log Mo
ma(mm)
+ y log
Σ
(8m to
(33)
Mo
m=1
m=1
where {ūmm
] Me are samples from pre-trained distribution.
E.3. Bilevel noise scheduling algorithm
With ZO estimation and the parametrization, the complete algorithm for bilevel noise scheduling problem is summarized in Algorithm 6. Besides, the upper-level schedule quality loss is differentiable according to (Mathiasen & Hvilshøj. 2020).
Differentiable L.)(.) loss. FID score is a commonly used metric in computer vision to measure the distance of the generated distribution and the true distribution. Given {um } m=1 ~ pe generated from diffusion model algorithm and Mo
(ūm)) } | ~ Pdata we first encode all samples Um, üm by the pre-trained Inception network (Seitzer, 2020) and then FID
18

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
Algorithm 5 Guided Diffusion for Generative Optimization
Input: Pre-trained score network sg(:,).. differentiable reward r(-), guidance G.
2: Parameter: Strength parameters p(t), (Un J number of iterations N, batch sizes (Bn).
3: Initialization: Gg NULL.
4: for iteration . N
do
5: Generate: Sample i for i E [B,] by backward SDE in (8) with (sa, Gx) until time T
6: Compute Guidance:
ii Sample mean := 1 B. EBn Un
(ii) Query gradient In = vr(ūn).
(iii) Update gradient guidance Gn+i(:5-)= = G(5)5)) via (32), using 80, gradient vector In. and reward target en. and (t)).
7: end for
8: Generate: Sample ün for i [Bx] by backward SDE in (8) with (so, GN) until time T
9: Output: {
BN
score is computed by the Wasserstein distance between the two multivariate normal distributions. Therefore, when the Inception network used for encoding is differentiable with respect its input, as the proposed by Seitzer (2020) does, Lsq(-) is differentiable with respect to the sample and then by the chain rule, it is differentiable with respect to and A
F. Experimental details
In this section, we introduce the details of experimental setup for two applications. All experiments were conducted on two servers: one with four NVIDIA A6000 GPUs, and 256 GB of RAM; one with an Intel i9-9960X CPU and two NVIDIA A5000 GPUs.
F.1 Application 1: fine-tuning diffusion model with bilevel entropy regularization learning
For the hyperparameter settings, we set the initial value noise scheduler q(t) and tune the entropy strength À using different methods. We use batch size of 3 for the fine-tuning step and set optimization step 7 and repeat the optimization for 4 times. The prompts we used for generating these figures are mentioned in the figure 9 - 11. We compare our results against various baseline methods including grid search, random search, bayesian method, and weighted sum.
Grid Search, For the grid search method, we selected the following À values and conducted simulations for each value:
X E (0.01, 0.1, 1.0, 10.0, 100}
Random search. We fine-tuned the diffusion model using À values generated by random value generator. Specifically, sampled 5 random values logarithm uniformly from the range [0.01, 100). The À values generated are as follows:
X (62.3, 74.0, 74.18, 79.52, 94.25)
Bayesian search. We utilized Bayesian optimization with the objective maximizing the reward function and the CLIP score. Specifically, the search space the hyperparameter Ascale was defined as a continuous range [0.01, 100), sampled on logarithmic scale using log-uniform distribution. The optimization process was conducted using the gp ap_minimiia function from the scikit-ope imize ze library, which employs Gaussian process-based Bayesian optimization. To balance computational efficiency and optimization quality, the number of function evaluations was limited to Mealls 15. Additionally, fixed random seed (random_state = 42) was set to ensure the reproducibility of results.
Since Bayesian optimization minimizes the objective function by default, we reformulated the problem by negating the combined reward and CLIP score, thereby transforming the maximization task into a minimization problem. This reformulation allowed to identify the optimal Awaate value that best balances reward maximization and adherence the
19

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
Algorithm 6 Bilevel Approach without Pre-trained Diffusion Model
1: Input: Differentiable loss functions (-s and [sM()), initial samples (ūm I m== iteration number K, Sz, Sy, initial noise scheduler parametrization parameter Oparam = (ga, ge Qr qe (cosine or sigmoid), feasible set for Qparam E Q stepsizes B, T7--
2: for k = 0, ..... K do
3: sample [
from the forward process (5) with noise scheduler 9k-
q0.
3param
4: for 0, 11.................... S do
5: update 0% 1
3 Σ MI
m==
VoLsu( K,st u.
D 6% = 0% 60 k+1 KN.5
6: end for
7: for 8 0, ,, 11 do
8: sample {ug.m "k m= from (6) with gi and O7 + v0 pertur and 07, - 10 perturb.
D. 0% = 0%
9: estimate (VoT.o (1. )
m== by ZO in (14)
update 0i. 0h. & D M
(VLLso(ūū, + V VoLIM (0 uk,m))
0 07
11: end for
12:
calculate parameterization perturbation Ix,param = param + 1 Uppertub Ok.param parau Upeettub
D 0 0f
calculate noise scheduler qf- 9x from o puram Tk param by cosine or sigmoid parameterization
14:
IM sample (#k+1,ms 1k+,, 1==1 from backward process (6) with 2 9x and 01+1- estimate Quean
9k.pomeh
2v [Iso(ui
Lso ( by ZO
16:
calculate {V Yronn SM k+1 (k,m), ); V grann LSM (0% k+11 2k.m ) by auto-differentiation Dhe M
17:
update 4k.1ppa = paraar
EME (V Vorn Lso( 1m+) ) + VV 4ramm LsM (f uk. Oponm LSM 6+1 2k
update 9i++1.param Proje (9kio((x++1ppaam)
19: end for
20: calculate noise scheduler qk from qa.param by cosine or sigmoid parameterization
21: sample [ux,m MM m==1 from the backward process (5) with qι and 0
22: Output: (IK: [4 1M 1)
original data distribution. The acquisition function used was the expected improvement (EI), defined as:
-EI(A) -
where f(x²) represents the best observed value at iteration 6.
Weighted sum. For the weighted sum method, we jointly optimized the reward function and the CLIP score during the fine-tuning of the diffusion model. The optimization was performed by taking the weighted sum the reward value and the CLIP score, with the weight for the CLIP score set to 0.5. The values used for the weighted sum method are selected by grid search with the search grid: X (0.01, 0.1, 1.0, 10.0, 100).
F.2. Application 2: bilevel noise scheduling learning
This cosine parametrization (16) covers both the cosine noise scheduler in (Nichol & Dhariwal, 2021) when 9. = 0, ge fr and (Chen, 2023) when de = 0. Sigmoid parameterization is defined similarly by
9a T - t(ge 9x)
I(t) sigmoid
Le
(34)
7T
which covers (Chen, 2023) when de 0. Since (t) should be nondecreasing, we assign q(t) 1 1(t)/((t 1). With the use of parameterization, ZO perturbation will be added on qs Ye, Ge, 4r instead of directly on g()).
For hyperparameter optimization for noise scheduler, we compare our method against greedy grid search, random search, Bayesian search and the default DDIM. Default parameter choices of DDIM with cosine noise scheduler in (Nichol &
Dhariwal, 2021) are qs = 0. qe = 1, gr = 1, 9. = 0.008. Default parameter choices of DDIM with sigmoid noise scheduler are 4s = -3, gr = 3, 4- 0.1,g. g 0.5 according to (Vidhya, 2024).
20

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
(Greedy) grid search. According the sensitivity analysis shown in Figure 6, the most sensitive parameter is dr, while the last three almost equally important. Therefore, in greedy grid search, we tuned the parameters in the order Ir; Ges Os Ge- We adopt the following search grid for cosine parametrization and best parameter given by greedy grid search is highlighted:
a {0, 0.1, 0.3, 0.4]
qe E {1, 0.9, 0.8}
4r E {1, 2,3,4)
qe {0.005, 0.008, 0.01, 0.02, 0.03, 0.04]
For sigmoid parametrization, the following search grid and the best parameter is highlighted in black:
9s E (-6, -5, -4 -, 2, -1,
le E (2,3,4)
4r E (0.1, 0.3,
q. (-2,-1,,0.5,,
Random search. We sample 16 random combinations of 9a. qe: 9r ge from 9, [0.005, 0.04) for cosine parameterization and E [-6, 0), E [2. 4/ 9r E tion. We report the best-performing results given by the random combination.
[0,0.4], E [0.8,1].9, E [1,4], E [-2, for sigmoid parameteriza-
Bayesian search. We use the same range as the random search for Bayesian search and employ the same implementation the first application.
Bilevel algorithm. We employ bilevel algorithm in Algorithm 6 and set the initialization of the noise scheduler parameter 9e, 4r 9e the default values in DDIM (Nichol & Dhariwal, 2021; Vidhya, 2024). We batch size of 128 and choose the number of inner loop S₂ for 0² updates as 1. Empirically, we found that, at the beginning of the training process (i.e. when & = 0), the number of inner loop Sy for updating & should be larger to obtain relatively reasonable U-Net, but later on, we do not need large inner loop, i.c. set S 10 for k 1. We formalize this stage initial epoch, where we traverse every batch and set Sy 20. We choose the ZO perturbation amount leverages the warm-start strategy.
0.01. Moreover, Algorithm 6
Warm-start strategy. To further accelerate convergence, we avoid fully optimizing the penalty and lower-level problem with respect to for every q. Instead, we employ warm-start strategy, initializing using value from the previous epoch (Arbel & Mairal, 2022; Vicol et al., 2022; Sambharya et al.. 2024). Empirically, this approach effectively reduces the inner loop for optimizing 0 in the lower-level and penalty problems to 10 and 1, respectively. Moreover, only 3 4 outer epochs needed for optimizing q. Compared to the 100 epochs required for single-level diffusion model training, this significantly enhances the computational efficiency of our method, as shown in Table 3. With only 2.5x the training time of a single-level diffusion model, the bilevel method achieves a 30% improvement over the default model while using just 15% of the time for Bayesian search.
Exponential moving average (EMA). We also incorporate EMA, which an indispensable strategy in all high-quality image generation methods stabilize training (Nichol Dhariwal, 2021; Song al., 2021b; Ho al., 2022: Karras et al., 2022). EMA maintains running average of model parameters time, where recent updates weighted more heavily than older ones that it smooths out fluctuations in the training process.
(a) A = 0.01
(b) X = 0.1
(c) A = 1.0
(d)
10.0
(e) A = 44.3
(f) 100.0
Figure 8. Balancing the realism and acsthetic in the image generation by controlling the entropy regularization strength parameter X. Prompt: "An African elephant foggy morning, with hot air balloons landing in the background."
21

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
(a) Grid Search (FID= 127.04, CLIP= 31.97)
(b) Bayesian Search (FID= 140.35, CLIP= 32.68)
(c) Random Search (FID= 117.46, CLIP= 34.67)
(d) Weighted Sum (FID= 108.95, CLIP= 35.87)
(e) Bilevel (FID= 104.35, CLIP=
(f A = 0.01 (FID= 403.66, CLIP= 17.58)
Figure Visualization of the final generated images (step-7) by different methods. Prompt: realistic photo horse standing on lush green grass in countryside meadow on sunny day, with clear blue sky in the background." (a) Grid Search: The generated images not fully adhere the prompt, as the clear blue sky is often missing. Some images appear more abstract. (b) Bayesian Search: Most images lack blue sky in the background, and some horses are deformed. (c) Random Search: In certain images, the mane is not well-defined. (d) Weighted Sum: Some images exhibit imperfections in the mane and facial features the horses. (e) Bilevel: Generates visually striking, highly realistic images that closely align with the given prompt.
22

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
(a) Grid Search (FID= 131.95, CLIP= 31.66)
(b) Bayesian Search (FID= 158.15, CLIP= 34.16)
(c) Random Search (FID= 127.91, CLIP= 36.34)
(d) Weighted Sum (FID= 117.95, CLIP= 37.85)
(e) Bilevel (FID= 105.72, CLIP= 39.05)
(f) À = 0.01 (FID= 440.40, CLIP= 18.78)
Figure Visualization the final generated images by different methods. Prompt: "An African elephant foggy morning, with hot air balloons landing in the background." (a) Grid Search: Some images exhibit deformed elephant figures, and the hot air balloons are missing, (b) Bayesian Search: The images appear more abstract, with deformed elephants and trees. Elephant even missing one figure, while another depicts elephants on top of tree. (c) Random Search: The elephant and hot air balloons are unclear in some images. (d) Weighted Sum: Struggles to generate recognizable elephant figure, producing a deformed trunk instead. (e) Bilevel: Generates relatively high-quality images with visible deformations.
23

---

A First-order Generative Bilevel Optimization Framework for Diffusion Models
(a) Grid Search (FID= 144.23, CLIP= 31.37)
(b) Bayesian Search (FID= 142.45, CLIP= 32.71)
(c) Random Search (FID= 121.23, CLIP= 33.97)
(d) Weighted Sum (FID= 118.91, CLIP= 34.53)
(c) Bilevel (FID= 112.78, (6.6536.65)
( X 0.01 (FID= 450.34, CLIP= 16.26)
Figure 11. Visualization of the final generated images by different methods. Prompt: "A gentleman wearing white clothes and a beard, posing seaside setting." (a) Grid Search: Struggles to generate human faces; images appear blurry. (b) Bayesian Search: Some faces are blurry, and hands are deformed. (c) Random Search: Some face images appear blurry. (d) Weighted Sum: Produces comparatively good-quality images. (e) Bilevel: Generates relatively high-quality images with no visible deformations.
24


Reviews:
[{"review_id": "57Na4Y4oSt", "reviewer": "Reviewer_8QSP", "summary": "The paper explores the problem of bilevel optimization with diffusion models - a hierarchical framework consisting of a higher and lower level objectives which are jointly optimized. The authors frame the following two problems as bilevel optimization:\n1. KL regularized reward maximization as the lower level objective for diffusion model finetuning, and the higher level objective optimizing the KL weight $\\lambda$ using CLIP as the reward.\n2. Tuning the noising schedule of a diffusion model in the upper level and score matching loss in the lower level.\nFor the reward guided finetuning task, the authors use a color/vibrancy reward for the lower level objective and CLIP score for the higher level objective, and demonstrate improvement over other search methods such as grid search and Bayes opt. For the second task, the authors finetune the noise schedule of an MNIST generative model using DDIM inference.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for appreciating our topic. We hope our response to your comments below can resolve your minor concerns. \n\n**Q1. Motivation of two hyperparameter optimization (HPO) problems.** \n\nIn 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. \n\n- **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. \n\n- **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. \n\n> [R1] Simple diffusion: End-to-end diffusion for high resolution images. E Hoogeboom, et. al. ICML 2023. \n\n**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. \n\n> [R2] Stochastic compositional gradient descent: algorithms for minimizing compositions of expected-value functions. M. Wang, et. al. Math. Programming, 2017. \n\n> [R3] Solving stochastic compositional optimization is nearly as easy as solving stochastic optimization. T. Chen, et. al. IEEE Trans. on Signal Processing, 2021. \n\n\n**Q3. Generating from optimal entropy guided distribution.** \n\nWe 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.  \n \n\n**Q4. Applicability of other guidance strategies.** \n\nYes, 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. \n\n**Q5. Applicability of the noise scheduling task in faster sampling with a diffusion model.** \n\nDue to the space limit, please see the response to **Q4** for Reviewer XKpE.  \n\n**Q6. Noise scheduling task on MNIST dataset.**\n\nDue to the space limit, please see the response to **Q1** for Reviewer XKpE. \n\n**Q7. Writing of this paper.**\n\nWe 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.", "claims_and_evidence": "The paper claims to propose an efficient framework to perform bilevel optimization with diffusion models. I think the claims for the most part are justified by the experiments, but I have not gone through the implementation details thoroughly enough to know if the methods being compared against are fairly tuned.", "methods_and_evaluation_criteria": "I am not fully convinced that the two tasks the paper explores are directly useful themselves. Usually in previous work, entropy regularized reward guided sampling of diffusion models (the lower level problem) is investigated in isolation. I can see the utility of tuning the KL weight, but the complexity added with bilevel optimization doesn't seem useful in practice. However, it is a reasonable experiment to demonstrate the method in the context of this paper. The noise scheduler optimization with score matching is a lower level loss is also a somewhat confusing experiment, and only demonstrating qualtitative results on MNIST is a weak result.", "theoretical_claims": "I did not thoroughly check the theoretical claims and math, however I did notice a couple of potential issues:\n1. The gradient in Proposition 1, is not correctly estimated using the MC estimator in Appendix E.2. The expectation is inside the log, so taking MC average inside the log would provide a biased estimator. However, section 4.1 seems to imply we can tractably estimate this gradient.\n2. In page 4, final paragraph (section 4.1) it is written that algorithm 5 sampled from the optimal entropy guided distribution. But this is also an intractable task, and I believe the sampling process would only provide approximate guidance.", "experimental_designs_or_analyses": "I did not notice any issues with the experimental design.", "relation_to_broader_scientific_literature": "I am not aware of previous works that are directly related to the proposed bilevel optimization framework. The noise scheduling task could be applicable to faster sampling with a diffusion model. The bilevel reward finetuning task is potentially more relevant in my opinion, and could we used alongside other strategies from the diffusion guidance [1] for solving the low level optimization task.\n\n[1] Inference-Time Alignment in Diffusion Models with Reward-Guided Generation: Tutorial and Review, https://arxiv.org/abs/2501.09685", "essential_references_not_discussed": "No essential references are missing to my knowledge.", "other_strengths_and_weaknesses": "### Weaknesses\n1. I found the paper quite confusing to read throughout. After reading the paper, I still do not fully understand the motivation.", "other_comments_or_suggestions": "NA", "questions_for_authors": "1. Could the authors clarify the two issues I brought up regarding theoretical claims?\n2. Is there any reason the noise scheduling task was only done for MNIST sampling? It would be more convincing if done on with a more difficult dataset.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "w2RKRVtsMO", "reviewer": "Reviewer_XKpE", "summary": "This paper introduces a practical, first-order bilevel framework for diffusion models, outperforming standard methods in fine-tuning and training scenarios. The proposed method eliminates the high dimensionality and sampling costs in traditional methods.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for appreciating our novelty. We hope our response to your comments below can resolve your minor concerns. \n\n**Q1. Insufficient numerical experiments.** \n\n- **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. \n- **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. \n- **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. \n\n**Q2. New challenges and novelty from the optimization perspective.** \n\nThank 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. \n\n**Q3. Hyperparameter tuning is a standard application of bilevel optimization.** \n\nThank 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. \n\n**Q4. Applicability of noise scheduling tuning in faster sampling with a diffusion model.** \n\nWe 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.", "claims_and_evidence": "The claims are supported by experimental results and theoretical convergence guarantees.", "methods_and_evaluation_criteria": "Yes, the proposed method and evaluation criteria make sense.", "theoretical_claims": "The main proof is provided in D.3 and it is correct.", "experimental_designs_or_analyses": "Since the main objective of this paper is to address practical challenges in diffusion models, the numerical experiments presented are insufficiently representative. \n\nSpecifically, Section 6.1 optimizes only one hyperparameter, resulting in limited improvement over naive search methods. \n\nAdditionally, the experiments in Section 6.2 rely solely on the MNIST dataset, which might be simple for large-scale cases for diffusion model applications and thus does not adequately reflect the performance of modern diffusion frameworks. More large-scale datasets could be added to make the results more convincing.", "relation_to_broader_scientific_literature": "The main contribution of this paper is framing diffusion hyperparameter tuning as a bilevel optimization problem. \n\nOn the positive side, this formulation appears novel in the diffusion context. However, the motivation is not that exciting because hyperparameter tuning is already a standard application of bilevel optimization. A more promising direction may involve using a bilevel optimization approach to address fundamental diffusion challenges, such as improving sampling efficiency.", "essential_references_not_discussed": "References are good to me.", "other_strengths_and_weaknesses": "The application of bilevel optimization to tuning hyperparameters in diffusion models is novel and interesting. From the bilevel optimization side, the novelty is not very significant. The paper could be strengthened if more comprehensive and large-scale experiments can be done.", "other_comments_or_suggestions": "Line 128 left, $S_\\gamma^*(x)$ is not defined. \n\nLine 201 right, $y^*$ and $z^*$ are mismatched. \n\nLine 437 right, the quotation marks need some adjustments.", "questions_for_authors": "Although I appreciate the novel application of gradient-based hyperparameter tuning for diffusion models. I still want to understand from the optimization perspective, what new challenges the diffusion models introduce to the algorithmic development in bilevel optimization? In other words, what is the main technical or algorithmic novelty on the bilevel optimization side?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LceZsDblIL", "reviewer": "Reviewer_FnLj", "summary": "This paper explores the application of bilevel optimization in diffusion models, focusing on two key applications. The first optimizes the trade-off parameter that balances the reward and proximity to the pre-trained distribution during fine-tuning. The second optimizes the noise schedule in diffusion models. To enhance scalability, the paper proposes a first-order optimization framework specifically designed for diffusion models. Experimental results demonstrate that the proposed method significantly reduces hyperparameter optimization time compared to grid search, random search, and Bayesian optimization.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for appreciating the theoretical guarantees for our algorithm and its empirical performance. Our response to your comments follows. \n\n**Q1. Optimization merely on $x$ in equation 1?** \n\nThank 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. \n\n**Q2. Applicability of method on non-differentiable objective.** \n\nThank 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]. \n\n> [R4] An Algorithm with Optimal Dimension-Dependence for Zero-Order Nonsmooth Nonconvex Stochastic Optimization. G. Kornowski, O. Shamir. JMLR, 2024.", "claims_and_evidence": "The proposed method is theoretically well-founded and demonstrates strong empirical performance.", "methods_and_evaluation_criteria": "The proposed method adapt the existing fully-first order method to diffusion model which makes sense.", "theoretical_claims": "I check the proofs of proposition 1 and theorem 1.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "The proposed method builds on the fully first-order approach from Kwon et al. (2023) and introduces improvements to enhance its scalability for diffusion models", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "The proposed method is theoretically grounded and specifically designed for diffusion models, demonstrating strong empirical performance.", "other_comments_or_suggestions": "In eq 1, should the optimization variable in upper level be only x?", "questions_for_authors": "Could the method be extended to the setting whether the upper-level objective is not differentiable?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8h1YYvmGyC", "reviewer": "Reviewer_ExjZ", "summary": "The authors have proposed a bilevel optimization framework tailored for diffusion models, specifically addressing two scenarios:-\nFine Tuning a pre-trained diffusion models (via Inference-only Solver):\n\nTo fine-tune pre-trained diffusion models, to maximize the task-specific rewards, while preserving the aesthetic realism, the authors propose a bilevel optimization problem as:\n- Upper-level problem: Select optimal hyperparameters such as entropy regularization strength λ, that balances reward maximization and realism.\n- Lower-level problem: Adjusts the generated data distribution to maximize a reward function with entropy regularization, ensuring closeness to the pre-trained distribution.\nThis method avoids expensive backpropagation through diffusion steps by using guided sampling and closed-form gradient estimation.\n\nNoise Schedule Optimization (Training from Scratch):\nThe authors optimize noise schedules used during training diffusion models, and the bilevel problem is defined as:\n- Upper-level problem: Optimizes parameters controlling the noise schedule to minimize metrics like FID score of the generated images.\n- Lower-level problem: Learn parameters of a score function that approximates gradients of log-likelihoods of noisy data distributions.\nTo efficiently solve this nested structure without differentiating through multiple sampling steps (which would be computationally intensive), the authors use:\n- Reparameterization of noise schedules (using cosine or sigmoid functions with only four parameters).\nZeroth-order gradient estimation, allowing gradient approximation without explicit backpropagation through sampling trajectories\n\n\nThe authors convert nested bilevel optimization problems into a single-level penalty problems solvable via first-order methods, providing theoretical guarantees under strong convexity assumptions", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for appreciating the theoretical innovation and our algorithm design. Our response to your comments follows. \n\n**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. \n\n**Q2. Reproducibility.** Thank you for your suggestion. We will release it with the final version. \n\n**Q3. Presentation of images.** Thank you for your suggestion. We will change accordingly. \n\n**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]. \n\n> [R5] Loss landscapes and optimization in over-parameterized non-linear systems and neural networks. C. Liu, et al. 2022. \n\n**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. \n\n**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. \n\n\n> [R6] Implicit Diffusion: Efficient optimization through stochastic sampling\n> [R7] SPARKLE: A Unified Single-Loop Primal-Dual Framework for Decentralized Bilevel Optimization\n> [R8] Bi-level Guided Diffusion Models for Zero-Shot Medical Imaging Inverse Problems\n> [R9] Aligning Text-to-Image Diffusion Models with Reward Backpropagation\n> [R10] Aligning Diffusion Models with Noise-Conditioned Perception\n> [R11] Diffusion Model Alignment Using Direct Preference Optimization", "claims_and_evidence": "Proposition 1 derives a closed-form gradient estimator for entropy regularization strength (λ) using pre-trained data samples, avoiding backpropagation through diffusion steps. Theorem 1 establishes convergence guarantees under strong convexity assumptions (Assumption 1) Reformulates bilevel problems into single-level penalty objectives (Eq. 2), enabling gradient updates via inference-only guided sampling (Alg. 5)", "methods_and_evaluation_criteria": "Methods: The authors have proposed a Bilevel Optimization for guidance of diffusion models to maximize the reward function, at the same time, to maintain the realism of the generated images. They have proposed this method for two scenarios: (i) inference based fine tuning of pretrained models, (ii) noise schedule optimization when training from scratch. \n\nEvaluation Criteria: The authors use FID score, CLIP score, IS score and time (is seconds) to discuss the effectiveness of their proposed method.", "theoretical_claims": "Theorem 1 establishes convergence guarantees for the proposed bilevel optimization framework under the following conditions:\n- Strong convexity: The lower-level objective is strongly convex.\n- Smoothness: f(x,y) and g(x,y) are jointly smooth over (x,y) with constants l_{f,1} and l_{g,1}.\n- Lipschitz continuity: f(x,⋅) is l_{f,0} 0-Lipschitz, and g(x,y) has l_{g,2} Lipschitz Hessian.\n\nUnder Assumption 1, the bilevel algorithm ensures a strict descent in the upper-level objective F(x). This theorem bridges theory and practice, ensuring that the framework’s hyperparameter updates (e.g., entropy strength λ or noise schedule parameters) provably guide diffusion models toward better performance under realistic Assumptions1.", "experimental_designs_or_analyses": "Experimental Setup: \n\nFor pre-trained models, reward finetuning, the authors use StableDiffusion V1.5 model as their pre-trained model and employ a ResNet-18 architecture (trained on the ImageNet dataset) as the synthetic (lower-level) reward model. The bilevel method achieved an 11.76% improvement in the FID score and an 8.32% improvement in the CLIP score over the best-performing weighted sum method.\n\nFor Noise Schedule Optimization, the authors train a U-Net model on MNIST from scratch, and use Cosine/sigmoid schedules with 4 parameters. The authors claim their bilevel method achieved 30% lower FID than default DDIM with only 2.5× training time.", "relation_to_broader_scientific_literature": "The authors claim, this is the first work related bilevel optimization to diffusion models. Their work fits into the broader space of AI Alignment, and Reward based optimization of generative models. Their work is foundational towards understanding alignment from a theoretical perspective.\n\nThe authors discuss reward alignment from bilevel optimization perspective. There are few recent papers on this topic, such as:-\n\n1. Implicit Diffusion: Efficient optimization through stochastic sampling\n2. SPARKLE: A Unified Single-Loop Primal-Dual Framework for Decentralized Bilevel Optimization\n3. Bi-level Guided Diffusion Models for Zero-Shot Medical Imaging Inverse Problems\n\nApart from that, in a broader sense, since the method is towards inference time alignment as well, the authors can comment on other test time alignment and reward based methods such as:-\n\n1. Aligning Text-to-Image Diffusion Models with Reward Backpropagation\n2. Aligning Diffusion Models with Noise-Conditioned Perception\n3. Diffusion Model Alignment Using Direct Preference Optimization", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "**Strengths**\n\nTheoretical Innovation:\n- Closed-form gradient estimation: Proposition 1 enables direct computation of gradients for entropy regularization strength (λ) using pre-trained samples, eliminating backpropagation through diffusion steps.\n- Convergence guarantees: Theorem 1 establishes convergence under strong convexity assumptions, matching standard bilevel optimization rates.\n\nAlgorithmic Design:\n- Inference-only fine-tuning: Avoids backpropagation through sampling trajectories by using guided sampling (Algorithm 5), reducing memory costs.\n- Parameterized noise schedules: Reduces noise schedule optimization to 4 parameters (e.g., cosine/sigmoid functions) instead of tuning per-step values.\n- Zeroth-order (ZO) gradients: Enables gradient estimation for noise schedules without differentiating through sampling steps (Equation 14).\n\n**Weaknesses**\n\n- Baseline Comparisons: The authors should compare their bi-level optimization approach with other reward model based approaches (mentioned above), to see how much their method performs. \n- Reproducibility: The authors should release the code, and their experimental setup, to support their work.", "other_comments_or_suggestions": "The images provided by the authors are a bit difficult to understand in terms of their aesthetics. If the authors can put 1 image per prompt rather than a collage, will help to differentiate their methods effectiveness.", "questions_for_authors": "- Is there a relaxation for strongly convex setup, which might be applicable here. If a better proof can be provided for the same, would strengthen their method.\n- Memory cost comparison can be further explained, as it has been mentioned in few of the places.\n- Can the authors comment further on other related works (as mentioned above), in terms of aesthetics, computation, methods, and comparison metrics?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a first-order bilevel optimization framework tailored to diffusion models, with two main applications: (1) fine-tuning pretrained diffusion models via reward-guided sampling by optimizing a KL regularization weight, and (2) optimizing noise schedules during training from scratch to improve generation quality. To ensure practical tractability, they avoid backpropagation through the generative process and use guidance for instance in (1). Experiments are conducted on Stable Diffusion (reward tuning) and MNIST (noise schedule tuning), showing improved sample quality and efficiency over baseline HPO methods like grid search and Bayesian optimization.
Strengths:
* Novel Application of Bilevel Optimization: While bilevel optimization is a known tool in hyperparameter tuning, applying it directly and efficiently to diffusion models is relatively novel and non-trivial. 
* Practical Efficiency: The authors convincingly argue that their approach reduces the computational cost of tuning reward weights and noise schedules in diffusion settings. Tables and comparisons with traditional methods support this.
Weaknesses:
* Limited Experimental Scope: A recurring concern among reviewers (especially XKpE and 8QSP) was that experiments, particularly for the noise scheduling task, were confined to MNIST.  The authors specify that noise-tuning/training on mnist is as costly as fine-tuning on imagenet.
* Writing Clarity and Motivation: Some reviewers found the motivation and technical exposition confusing or unclear (notably 8QSP). The complexity of bilevel optimization combined with diffusion model details makes parts of the paper hard to follow.
* Limited Differentiation on the Optimization Side: While the application to diffusion models is novel, reviewers noted that the optimization algorithm itself is an adaptation of existing first-order bilevel methods. More innovation on this front could strengthen the contribution. Also, the theoretical results (eg Theorem 1) appear to be standard in bilevel optimization.
All reviews were positive (3/5,3/5,4/5,4/5) and remained so after rebutal.
While the paper's core contribution lies more in its application of bilevel optimization than in fundamentally new algorithmic techniques and theoretical contributions, this application is meaningful, non-trivial, and supported by both theory and experiments. The authors have thoughtfully addressed reviewer feedback provided additional experiments.
With additional large-scale experiments and clearer exposition, the paper would be stronger, yet in its current state it could make a notable contribution to both the diffusion modeling and bilevel optimization communities. It is therefore recommended for acceptance.
Decision: Accept (poster)