The optimization of the latents and parameters of diffusion models with respect to some differentiable metric defined on the output of the model is a challenging and complex problem. The sampling for diffusion models is done by solving either the probability flow ODE or diffusion SDE wherein a neural network approximates the score function or related quantity, allowing a numerical ODE/SDE solver to be used. However, naïve backpropagation techniques are memory intensive, requiring the storage of all intermediate states, and face additional complexity in handling the injected noise from the diffusion term of the diffusion SDE. We propose a novel method based on the stochastic adjoint sensitivity method to calculate the gradientwith respect to the initial noise, conditional information, and model parameters by solving an additional SDE whose solution is the gradient of the diffusion SDE. We exploit the unique construction of diffusion SDEs to further simplify the formulation of the adjoint diffusion SDE and use a change-of-variables to simplify the solution to an exponentially weighted integral. Using this formulation we derive a custom solver for the adjoint SDE as well as the simpler adjoint ODE. The proposed adjoint diffusion solvers can efficiently compute the gradients for both the probability flow ODE and diffusion SDE for latents and parameters of the model. Lastly, we demonstrate the effectiveness of the adjoint diffusion solvers onthe face morphing problem.
Following in the research on ODE solvers for the Probability Flow ODE we explicit the semi-linear structure of the empirical adjoint Probability Flow ODE to eliminate the linear discretization error and simplify the computation of the ODE. We propose a bespoke family of ODE solvers which minimize the discretization error in calculating the ODE.
AdjointDEIS-1 First-order solver of the empirical adjoint Probability Flow ODE.
AdjointDEIS-2M Second-order multi-step solver of the empirical adjoint Probability Flow ODE.
SDEs introduce a lot of complexity over ODEs; however, we show that the adjoint diffusion SDE actually reduces to an ODE! This enables the usage of the same solvers for the adjoint Probability Flow with the addition of a factor of 2 on the non-linear term.
SDE-AdjointDEIS-1 First-order solver of the empirical adjoint diffusion SDE.
Our method is interoperable with a variety of solvers for the Probability Flow ODE and supports calculating the adjoint for diffusion SDEs.
@inproceedings{blasingame2024adjointdeis,
title = {Adjoint{DEIS}: Efficient Gradients for Diffusion Models},
author = {Blasingame, Zander W. and Liu, Chen},
booktitle = {The Thirty-eighth Annual Conference on Neural Information Processing Systems},
year = {2024},
url = {https://openreview.net/forum?id=fAlcxvrOEX},
}