AdjointDEIS: Efficient Gradients for Diffusion Models

TL;DR

This paper introduces AdjointDEIS, a novel family of efficient ODE solvers for computing gradients in diffusion models, simplifying the process and reducing memory usage, with proven convergence guarantees and practical effectiveness demonstrated in face morphing tasks.

Contribution

The paper presents AdjointDEIS, a new class of bespoke ODE solvers tailored for diffusion models, leveraging exponential integrators to simplify adjoint equations and improve gradient computation efficiency.

Findings

Effective guided generation demonstrated in face morphing.

Significant reduction in memory usage for backpropagation.

Proven convergence order guarantees for the proposed solvers.

Abstract

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 allowing a numerical ODE/SDE solver to be used. However, naive 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 family of bespoke ODE solvers to the continuous adjoint equations for diffusion models, which we call AdjointDEIS. We exploit the unique construction of diffusion SDEs to further simplify the formulation of the continuous adjoint equations using exponential integrators. Moreover, we provide convergence order guarantees for our bespoke solvers. Significantly, we show that continuous adjoint equations for diffusion SDEs actually simplify to a simple ODE. Lastly, we demonstrate the effectiveness of AdjointDEIS for guided generation with an adversarial attack in the form of the face morphing problem. Our code will be released at https: //github.com/zblasingame/AdjointDEIS.