Implicit Diffusion: Efficient Optimization through Stochastic Sampling

Pierre Marion, Anna Korba, Peter Bartlett, Mathieu Blondel, Valentin De Bortoli, Arnaud Doucet, Felipe Llinares-López, Courtney Paquette, Quentin Berthet

Keywords: Diffusion, Differentiation, Sampling, Bilevel

TL;DR: We develop an algorithm to optimize through a sampling operation, borrowing from concepts in implicit differentiation and in bilevel optimization. We use it to train energy-based models and finetune diffusion models.

Abstract: We present a new algorithm to optimize distributions defined implicitly by parameterized stochastic diffusions. Doing so allows us to modify the outcome distribution of sampling processes by optimizing over their parameters. We introduce a general framework for first-order optimization of these processes, that performs jointly, in a single loop, optimization and sampling steps. We showcase it in training and finetuning applications.

Submission Number: 26