A First-order Generative Bilevel Optimization Framework for Diffusion Models
Quan Xiao, Hui Yuan, A F M Saif, Gaowen Liu, Ramana Rao Kompella, Mengdi Wang, Tianyi Chen
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) finetuning 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 finetuning and hyperparameter search baselines. Our code has been released at https://github. com/afmsaif/bilevel_diffusion .
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 9e6ef983-0bb3-457f-804b-5b1133b4424cCited by top-tier papers2
- Beyond Value Functions: Single-Loop Bilevel Optimization under Flatness ConditionsLiuyuan Jiang, Quan Xiao, Lisha Chen, Tianyi ChenNeurIPS 2025 · 11 citations
- Constrained Flow Optimization via Sequential Fine-Tuning for Molecular DesignSven Gutjahr, Riccardo De Santi, Luca Schaufelberger, Kjell Jorner et al.ICML 2026 · 3 citations
Builds on43
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Denoising Diffusion Implicit ModelsJiaming Song, Chenlin Meng, Stefano ErmonICLR 2021 · 11,743 citations
- Improved Denoising Diffusion Probabilistic ModelsAlexander Quinn Nichol, Prafulla DhariwalICML 2021 · 5,234 citations
- Elucidating the Design Space of Diffusion-Based Generative ModelsTero Karras, Miika Aittala, Timo Aila, Samuli LaineNeurIPS 2022 · 3,959 citations
- Score-Based Generative Modeling through Stochastic Differential EquationsYang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar et al.ICLR 2021 · 1,270 citations
Related papers
- Efficient Diffusion Models via Time Step Optimization with Consistent Training and Inference ConstraintsBinrui Wu, Zihao Cheng, Yuesen Liao, Weizhong ZhangICML 2026
- Score-Optimal Diffusion SchedulesChristopher Williams, Andrew Campbell, Arnaud Doucet, Saifuddin SyedNeurIPS 2024 · 19 citations
- Gradient Guidance for Diffusion Models: An Optimization PerspectiveYingqing Guo, Hui Yuan, Yukang Yang, Minshuo Chen et al.NeurIPS 2024 · 79 citations
- Learning Fast Samplers for Diffusion Models by Differentiating Through Sample QualityDaniel Watson, William Chan, Jonathan Ho, Mohammad NorouziICLR 2022 · 224 citations
- Self-diffusion for Solving Inverse ProblemsGuanxiong Luo, Shoujin HuangNeurIPS 2025 · 5 citations
