A First-order Generative Bilevel Optimization Framework for Diffusion Models
This work addresses a specific problem for researchers and practitioners using diffusion models, offering a novel method for optimizing hyperparameters and fine-tuning tasks, though it is incremental in improving computational efficiency.
The paper tackles the challenge of optimizing diffusion models for downstream tasks, which involves nested bilevel structures that traditional methods fail to handle due to infinite-dimensional probability spaces and high sampling costs, and it demonstrates that their first-order bilevel framework outperforms existing fine-tuning and hyperparameter search baselines.
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.