AdjointDEIS: Efficient Gradients for Diffusion Models

TL;DR

The paper tackles memory-efficient gradient computation through guided diffusion models by formulating continuous adjoint equations and introducing AdjointDEIS, a family of bespoke exponential-integrator solvers for the adjoint dynamics. It shows that adjoint equations for diffusion SDEs conveniently reduce to an ODE, enabling efficient, high-order solutions with convergence guarantees, and extends the framework to time-varying conditioning. The approach is demonstrated via training-free guided generation on a face-morphing task, where AdjointDEIS outperforms prior methods in gradient-guided morphing while providing a principled numerical foundation. The authors also provide theoretical analyses, practical solvers (AdjointDEIS-1 and AdjointDEIS-2M), and open-source code, enabling reproducibility and broader applicability to diffusion-based guidance tasks. Overall, AdjointDEIS offers a scalable, accurate means to backpropagate through diffusion processes for guided generation, with potential impact across editing, inversion, and adversarial-context applications.

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.

Figures (6)

• Figure 1: A high-level overview of the AdjointDEIS solver to the continuous adjoint equations for diffusion models. The sampling schedule consists of $\{t_n\}_{n=0}^N$ timesteps for the diffusion model and $\{\tilde{t}_n\}_{n=0}^M$ timesteps for AdjointDEIS. The gradients $\mathbf{a}_\mathbf{x}(T)$ can be used to optimize $\mathbf{x}_T$ to find some optimal $\mathbf{x}_T^*$.
• Figure 2: Example of guided morphed face generation with AdjointDEIS on the FRLL dataset.
• Figure 3: Comparison of DiM morphs on the FRLL dataset. From left to right, identity $a$, DiM-A, Fast-DiM, Morph-PIPE, AdjointDEIS (ODE), AdjointDEIS (SDE), and identity $b$.
• Figure 4: Morphed faces created by guided generation with AdjointDEIS with differing number of discretization steps.
• Figure 5: Morphed faces created by guided generation with AdjointDEIS with different learning rates. All used $M = 20$ the ODE variant.

Theorems & Definitions (16)

• Proposition 2.1: Exact solution of adjoint diffusion ODEs
• Theorem 2.1: AdjointDEIS-$k$ as a $k$-th order solver
• Theorem 2.2
• Theorem 3.1
• Proposition 3.1: Exact solution of adjoint diffusion SDEs
• Remark 3.1
• Remark B.1
• Remark B.2
• Proposition B.1
• Lemma C.1: Vector-Jacobian Product is Lipschitz.