Interpreting and improving diffusion models
from an optimization perspective
Chenyang Yuan (Joint work with Frank Permenter)
Blogpost/Code: chenyang.co/diffusion.html
Paper: arxiv.org/abs/2306.04848
Toyota Research Institute

Friday 14th June, 2024

Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

1 / 25

Introduction
Diffusion models achieve state-of-the-art results in multiple domains such as:

Image generation

Trajectory planning

A powerful generative framework for sampling from multimodal distributions

Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

2 / 25

Motivation
Diffusion models are motivated by probabilistic models

Reversal of stochastic process that adds noise to data (Ho et.al. 2020)

Sampling of data distribution using score function (Song et.al. 2021)

However, commonly used sampling procedures (e.g. DDIM) are deterministic
Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

3 / 25

Motivation

Given commonly used diffusion training and sampling algorithms,
Is there a deterministic model that motivates the same algorithms?
Can we make reasonable assumptions on learned NN model to analyze
performance of sampling algorithm?
Our optimization-based interpretation:
Denoising approximates projection under manifold hypothesis
Diffusion sampling finds projection to data manifold by minimizing distance
via gradient descent

Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

4 / 25

Training diffusion models
Denoising diffusion models estimate a noise vector ϵ ∈ Rn from a given noise level
σ > 0 and noisy input xσ ∈ Rn such that for some x0 in the data manifold K,xσ ≈ x0 + σϵ
A denoiser ϵθ : Rn × R+ → Rn is learned by minimizing
L(θ) := Ex0 ,σ,ϵ ∥ϵθ (x0 + σϵ, σ) − ϵ∥2*Note: to get expressions commonlyused in literature, change of coordinates√zt = αt xt , where σt2 = (1 − αt )/αt .

Visualization of training process
Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

5 / 25

Distance and projection
The distance function distK : Rn → R to a set K ⊆ Rn , is defined via
distK (x) := inf{∥x − x0 ∥ : x0 ∈ K}
The projection of x ∈ Rn , is the set of points that attain this distance
projK (x) := {x0 ∈ K : distK (x) = ∥x − x0 ∥}Intuitively, projK (x) − x is the direction of steepest
descent (i.e. neg. gradient) of distK (x).
Proposition
Suppose K ⊆ Rn is closed and x ∈/ K. If projK (x) isa singleton, then∇ 12 distK (x)2 = x − projK (x)Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

6 / 25

Smoothed distance and approximate projection
Distance function is not differentiable everywhere, ∇distK (x) is not continuous,thus hard to learn!
Solution: σ-smoothed distance function
2σ 2 ∥x − x∥ = −σ 2 log
distK (x, σ) := softmin
0x0 ∈KP∥x0 −x∥2exp−2x0 ∈K
2σ

Smoothed distance function has continuous gradients
Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

7 / 25

Ideal denoisers
The ideal denoiser is the minimizer of training loss, a function of data distribution
K and noise level σϵ∗ := arg min Ex0 ,σ,ϵ ∥ϵθ (x0 + σϵ, σ) − ϵ∥2ϵθFor finite K, there is a closed-form solution:
P2(xσ − x0 ) exp(− ∥xσ − x0 ∥ /2σ 2 )ϵ∗ (xσ , σ) = x0 ∈K
P2σ x0 ∈K exp(− ∥xσ − x0 ∥ /2σ 2 )Plot of direction of ϵ∗ (x, σ) for different x and σ
Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

8 / 25

Ideal denoiser equivalent to gradient of smoothed distance
Theorem
For all σ > 0 and x ∈ Rn , we have1
∇x dist2K (x, σ) = σϵ∗ (x, σ).
2Does ideal denoiser approximate
projection?
We can compute relative error oflearned denoiser v.s. ideal denoiser for
CIFAR-10 dataset
Plot error distribution for 10k different
DDIM sampling trajectories

Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

9 / 25

Denoising approximates projection: Low noise
Manifold hypothesis: “real-world” datasets are (approximately) contained in
low-dimensional manifolds K of of Rn .
tanK (x0 )⊥xσtanK (x0 )x0projK (xσ ) − xσ
KGiven xσ = x0 + σϵ, most of the added noise lies in NK (x0 ) with high probability,
thus denoising approximates projection

Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

10 / 25

Denoising approximates projection: Low noise
The reach of K is the largest τ so that projK (x) is unique when distK (x) ≤ τtanK (x0 )⊥xσtanK (x0 )x0projK (xσ ) − xσ
KTheorem√Fix σ > 0 and suppose that reach(K) ≳ σ n. Given x0 ∈ K and ϵ ∼ N (0, I ), letxσ = x0 + σϵ. With high probability, we have:√∥projK (xσ ) − x0 ∥ ≲ σ d.
Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

11 / 25

Denoising approximates projection: High noise
Diffusion models often add large levels of noise to x0 in training, in order to startsampling from a Gaussian distribution
xσ
projK (xσ ) − xσKx0

When σ is large, both denoising and projection point in the same direction
towards K
Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

12 / 25

Denoising approximates projection
We claim that the denoiser learned from diffusion objective approximates
projection with small relative error
When σ small, manifold hypothesis tells us that most of noise added isorthogonal to data manifold
When σ large, any weighted mean of data has small relative error
Denoising with ideal denoiser is a σ-smoothing of projK (x) with small
relative error
Next:
Introduce relative error model
Prove that diffusion sampling minimizes distance to data manifold under thiserror model

Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

13 / 25

Sampling from diffusion models (Deterministic)
Given noisy xσ and noise level σ, the learned denoiser ϵθ (xσ , σ) estimatesx0 ≈ x̂0 (xσ , σ) := xσ − σϵθ (xσ , σ).
xt(σt−1 − σt )ϵθ (xt )xt−1x0x1

K x̂0tSampling algorithms (e.g. DDIM) construct a sequence x̂0t := x̂0 (xt , σt ) of
estimates from a sequence of points xt using the update:
xt−1 = xt + (σt−1 − σt )ϵθ (xt , σt )Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

14 / 25

Sampling from diffusion models (Probabilistic)
Deterministic (DDIM) update:
xt−1 = xt + (σt−1 − σt )ϵθ (xt , σt )Probabilistic (DDPM) update:
xt−1 = xt + (σt ′ − σt )ϵθ (xt , σt ) + ηwt
2
Where wt ∼ N (0, I ), σt ′ = σt−1/σt and η =q2
σt−1− σt2′22(Matches norm of update in expectation if E ∥wt ∥ = ∥ϵθ (xt , σt )∥ )√Note: σt−1 = σt σt ′ , thus σt ′ < σt−1 < σt .
These iterations look like gradient descent! But on which function?

Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

15 / 25

Our error model√
Let f (x) := 21 distK (x)2 . Intuitively, ∇f (x) = x − projK (x) ≈ distK (x)ϵθ (x)/ nAssumption (Projection with relative error)
There exists ν ≥ 1 and η ≥ 0 such that if ν1 distK (x) ≤∇f (x) exists, then ∥σt ϵθ (x, t) − ∇f (x)∥ ≤ ηdistK (x).
If√√nσt ≤ νdistK (x) andnσt closely tracks distK (x), then σt ϵθ (x, t) is approximately ∇f (xt )Relative error model where error depends on distance to K
Implications can be empirically tested on real datasets

DDIM is approximate gradient descent on f with stepsize 1 − σσt−1, witht∇f (xt ) estimated by ϵθ (xt , σt )Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

16 / 25

Intuitions on error modelIf√nσt closely tracks distK (x), then σt ϵθ (x, t) is approximately ∇f (xt )xt(σt−1 − σt )ϵθ (xt )xt−1x0x1

K x̂0tProjection at large noise levels is relatively inaccurate (tends to return data mean)
Projection at smaller noise levels is more accurate, but denoiser requires xt tohave noise level σtRelative error assumption captures this intuition
Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

17 / 25

Analysis under the error modelA schedule is {σt }Nt=0 is (η, ν)-admissible when σt is decreased slow enough tomaintain relative error assumption
Log-linear (geometrically decreasing) schedules are (η, ν)-admissible
σt−1 = (1 − β)σt , with β < C (η, ν)
We show convergence under the relative error assumption
Theorem (DDIM with relative error)
Let xt denote the sequence generated by DDIM and suppose that the gradient off (x) := 12 distK (x)2 exists for all xt . Then for all t:√1nσt ≤ νdistK (xt ),ν distK (xt ) ≤QNQNdistK (xN ) i=t (1−βi (η +1)) ≤ distK (xt−1 ) ≤ distK (xN ) i=t (1+βi (η −1)).
Admissible schedule =⇒ Control of relative error =⇒ distK decreases
Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

18 / 25

Improving sampling by gradient estimation
Our error model asserts that ϵθ (x, σ) ≈√n∇distK (x) when distK (x) ≈√nσ.

xt+1Since ∇distK (x) is invariant between x andproj√ K (x), we aim to minimize estimation errorn∇distK (x) − ϵθ (xt , σt ), with the updateϵ̄t = ϵθ (xt+1 ) + γ(ϵθ (xt ) − ϵθ (xt+1 ))Replaces ϵθ (xt , σt ) in sampling algorithm

xtϵθ (xt+1 )ϵ̄ϵθ (xt )Corrects for error made in previous step using
current estimate

K
Illustration of our choice of ϵ̄t

Empirically, γ = 2 achieves best results across many datasets and numberof sampling steps
Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

19 / 25

Improved sampling algorithm
Given (σN , . . . , σ0 ), xN ∼ N (0, I ) and ϵθ , to compute x0 with N evaluations of ϵθ :Algorithm 1 DDIM sampler
Algorithm 2 Our sampler
for t = N, . . . , 1 do
xN−1 ← xN + (σN−1 − σN )ϵθ (xN , σN )xt−1 ← xt + (σt−1 − σt )ϵθ (xt , σt )for t = N − 1, . . . , 1 doϵ̄t ← 2ϵθ (xt , σt ) − ϵθ (xt+1 , σt+1 )return x0
xt−1 ← xt + (σt−1 − σt )ϵ̄t
return x0

Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

20 / 25Experiments on noise schedule
How should we choose σt ? Relative noise model suggests log-linear schedule

Plot of different choices of log(σt ) for
N = 10.

Chenyang Yuan (Toyota Research Institute)

Schedule

CIFAR-10

CelebADDIM
DDIM Offset
EDM
Ours

16.86
14.18
20.85
13.25

18.08
15.38
16.72
13.55

FID scores of the DDIM sampler withdifferent σt schedules on the CIFAR-10
model for N = 10 steps.

Diffusion approximates projection

Friday 14th June, 2024

21 / 25

Sampler comparison experiments (Visual)

Visualizing x̂0t throughout the denoising process:A comparison of our sampler with DDIM on the CelebA dataset with N = 5 steps.

Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

22 / 25

Sampler comparison experiments (FID)

CIFAR-10 FID
CelebA FID
Sampler N = 5 N = 10 N = 20 N = 50 N = 5 N = 10 N = 20 N = 50
Ours
DDIM

12.57 3.79
47.20 16.86

PNDM 13.9
DPM
DEIS
18.43
UniPC 23.22
A-DDIM

7.03
6.37
7.12
3.87
14.00

3.32
8.28

3.41
4.81

10.76 4.41
32.21 18.08

3.19
11.81

3.04
7.39

5.00
3.72
4.53

3.95
3.48
3.78

11.3

7.71
5.83
25.07 6.95

5.51
2.82
3.41

3.34
2.71
2.95

5.81*

4.04

15.62

9.22*

6.13

FID scores of our sampler compared to that of other samplers for pretrained CIFAR-10
and CelebA models with a discrete linear schedule. *Results for N = 25

Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

23 / 25

Comparison on latent diffusion models
Ours
FID 13.77

UniPC
15.59

DPM++
15.43

PNDM
19.43

DDIM
14.06

Example outputs on text-to-image Stable Diffusion when limited to N = 10 function
evaluations. FID scores for text-to-image generation on MS-COCO 30K.
Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

24 / 25

Conclusion

Elementary deterministic framework for analyzing and generalizing diffusion
modelsSimplified exposition of existing algorithms and methods
New fast and simple-to-implement sampler designed with our interpretation
Framework for incorporating ideas from optimization into diffusion models
Constraining diffusion models ↔ constrained optimization
Use diffusion models in optimization problems (e.g. as a regularizer forcompressed sensing)

Chenyang Yuan (Toyota Research Institute)Diffusion approximates projection

Friday 14th June, 2024

25 / 25