A First-order Generative Bilevel Optimization Framework for Diffusion Models

TL;DR

This paper introduces a first-order bilevel optimization framework tailored for diffusion models, enabling efficient fine-tuning and noise schedule optimization, and demonstrating superior performance over existing methods.

Contribution

The paper develops a novel first-order bilevel optimization approach specifically designed for diffusion models, addressing the challenges of infinite-dimensional spaces and sampling costs.

Findings

Outperforms existing fine-tuning baselines

Efficient noise schedule optimization from scratch

Provides theoretical and computational advantages

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.