AdjointDEIS: Efficient Gradients for Diffusion Models

Summary

This paper presents a way to optimize the initial seed or control parameters to a diffusion network to optimize a differentiable loss applied to the final samples drawn from a discretization of the diffusion SDE / probability flow ODE. Put differently the idea is to apply the white box integration trick from DPM / DPM++ paper, and the method from the neural ODE paper to differentiate through an ODE solver. In addition, the paper presents methods for backpropagating through an SDE. In terms of experiments the method is demonstrated on a face morphing application.

Strengths

The paper is well written and the derivation is mostly clear. I did not go through section 4 for the derivation of the SDE Adjoint Diffusion SDE too closely since I am not very familiar with the Stratonovich stochastic integral. The paper is also fairly comprehensive since it presents the derivations for both the probability flow formulation as well as the SDE formulation. Overall I think the problem tackled in the paper is interesting and the solutions are elegant.

Weaknesses

The experimental results seem to be weak. This paper presents a way to control the final outcome from a diffusion sampler using a differentiable loss, but the only application demonstrated in the paper is on face-morphing where the NFE are not controlled across the baselines. Some of the baselines have fewer NFE and it's not possible to say conclusively whether they underperform because they are worse methods or just because they had less computational budget. AdjointDEIS-2M is not even implemented and compared. Also, a lot more compelling demonstrations could be shown of these new algorithms. For example, a number of "controlnet" type applications could have been tried to create an image that matches a sketch of an output.

Questions

Mainly I would like to understand a little bit more about the rationale for selecting the face morph application selection, and whether the authors think the empirical comparisons in Table 1 are apples-to-apples given the different NFE for the methods. Some more minor comments - eq (3.1) has a factor of 2 missing. - around eq(2.5) it might be good to explicitly clarify that while VP typically uses alpha^2 + sigma^2 = 1 ibut you are not making that assumption here.

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

Yes limitations have been addressed.

Summary

In this paper, the authors proposed an accelerated method for differentials of pretrained diffusion models with respect to its latent valuables or parameters by making use of 1. the Taylor expansion of the log-SNR parameter, and 2. the exact integral formula of the derivatives related to the probability flow ODE, which they call Adjoint Diffusion ODE. They applied this method to morphing of two generated images. As they summarized in Table 1, their results achieved good performance. For more details: - The setting - Data diffusion SDE $d\mathbf{x}_t = f(t) \mathbf{x}_t + g(t) d\mathbf{w}_t$ - Its integral $q(\mathbf{x}_t|\mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t|\alpha_t \mathbf{x}_0, \sigma_t^2)$ where $f(t), g(t)$ are related to $\alpha_t, \sigma_t$ by Eq.(2.5). - A pretrained diffusion model $\epsilon_\theta(\mathbf{x},\mathbf{z}, t)$ - Its generating process: the probability flow ODE $\frac{d\mathbf{x}\_t}{dt} = f(t)\mathbf{x}\_t + \frac{g(t)^2}{2 \sigma_t}\epsilon_\theta(\mathbf{x}, \mathbf{z}, t)$ - They name the velocity (the RHS of the ODE) as $\mathbf{h}_\theta(\mathbf{x}, \mathbf{z}, t)$ - The purpose - Calculation of gradients of $\mathcal{L}(\mathbf{x}_0(\mathbf{x}_T, \mathbf{z}, \theta))$ with respect to $\mathbf{z}$ or $\theta$ - which are defined by ODEs in Eq.(3.3) and (3.4). - The method - One can recover the derivatives based on a vector $\hat{\mathbf{a}}_t = \frac{\partial \mathcal{L}}{\partial \mathbf{y}_t}$ where $\mathbf{y}_t = \alpha_t^{-1}\mathbf{x}_t$. - For $t<s$, one can get $\hat{\mathbf{a}}_s = \hat{\mathbf{a}}_t + \int\_{\lambda_t}^{\lambda_s} ... d\lambda$ (Proposition 3.1). - Applying the Taylor expansion on ... with $\lambda$ around $\lambda_t$, and drop $O((\lambda_s - \lambda_t)^{k+1})$ term, we get a concrete relation between $\hat{\mathbf{a}}_s$ and $\hat{\mathbf{a}}_t$ (Eq.(3.11)), - which provides explicit (approximated) formulas for the derivatives. - The empirical result - They apply the proposed method to morphing (Section 5 Experiments).

Strengths

- **quality(+)** The explanations are well written. The authors provide concrete proofs on their theorems and experiments with real data.

Weaknesses

- **significance(-)** It is difficult to understand in which respect the proposed method is superior. - **clarity(-)** I am almost convinced by the explanations, however there is a concern as explained in Questions.

Questions

- for significance improvements: - **Q1(significance)** In the theoretical sections, the authors apply omitting the higher order terms of the Taylor expansion, so the resultant method should be just an approximation. Why is it good performance in Table 1 compared to other methods nonetheless the method is just an approximation? - for clarity improvements: - **Q1(clarity)** In the Adjoint Empirical Diffusion SDE cases, the integral results in Eq.(4.12) and (4.13) look very similar to ODE cases in Eq.(3.13) and (3.14). The only difference seem to be (doubled) 2nd term. Why such similarity happens? - **Q2(clarity)** In the beginning of section 4, the authors wrote their motivation to analyze SDE cases as it possibly improve the result, however, it turns to be that the SDE version works worse than the ODE cases. Why does it happens?

Rating

5

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

They leave comments on 1. Lack of conditional generation in numerical experiments 2. No application of parameter differentiation in numerical experiments in the final page.

Summary

The paper proposes an adjoint sensitivity method -- AdjointDEIS -- for efficiently calculating gradients of diffusion SDE models. Current methods for naive backpropagation rely on discrete adjoints which are memory intensive. The authors introduce an approach based on the stochastic adjoint sensitivity method to solve the gradients with respect to initial noise, conditional information, and model parameters. The authors develop custom solvers for the adjoint problems along with the proposed sensitivity method. The methods are validated on a face morphing problem.

Strengths

* **Great Flow/Presentation**: The general presentation of the paper is well done! Especially starting off with a more straightforward setting of an ODE and extending them to SDEs * **Specialized Solver**: Diffusion models are quite widespread and developing specialized solvers leveraging the structure of the problem demonstrates a practical application.

Weaknesses

* **Limited Experimental Scope**: While the overall analysis and formulations are sound, the overall experimental validation for the method is limited to just a single experiment.

Questions

1. Point 2 in contributions: Aren't the continuous adjoint sensitivity methods used in Neural SDEs already general-purpose enough to handle diffusion problems?

Rating

7

Confidence

2

Soundness

3

Presentation

4

Contribution

3

Limitations

N/A

Summary

AdjointDEIS uses the method of adjoint sensitivity to compute gradients of diffusion models, which is more efficient and less memory intensive, and robust to the injected noise. This work proposes efficient solvers for both the adjoint probability flow ODE and the adjoint diffusion SDE. Experiments demonstrate the efficacy of the solvers on guided generation for a face morphing problem.

Strengths

- AdjointDEIS reduces the adjoint of the PF ODE to the problem of exponential integrators, which enables use of the vast amount of literature in the areas. - Obtaining accurate gradients wrt initial noise x_T and the latent z can open the door to possibly novel applications of diffusion models.

Weaknesses

- The number of experimental applications of the method seem lacking (just one application to face morphing attacks is included). I’d encourage the authors to think about applications to data assimilation or inverse problems for instance. - Since the primary contribution of the paper is solver, having a proof of convergence rates or empirical studies of convergence rate would be valuable additions to the paper.

Questions

- It would be interesting to visually plot example forward and backward trajectories of the diffusion processes. What is the typical magnitude of the numerical errors that accumulate during the backward ODE solve?

Rating

6

Confidence

2

Soundness

3

Presentation

2

Contribution

2

Limitations

N/A

Response to authors

Thank you for addressing my comment. I have updated my score. I still encourage the authors to do more extensive evaluation for the method in the future.