This document proposes a method called Diffusion Posterior Sampling (DPS) to perform posterior sampling for inverse problems using diffusion models. DPS approximates the intractable posterior gradient using a denoising score model and Tweedie's formula. Specifically, it replaces the likelihood gradient with the gradient of the denoised data, allowing efficient posterior sampling via Langevin dynamics. The method is demonstrated on various inverse problems like super-resolution, inpainting, deblurring, achieving high-quality reconstructions from noisy measurements.

1. Diffusion models for inverse problems
ICLR 2023 (spotlight)

2. Background: Diffusion models and Stein score
http://yang-song.github.io/blog/2021/score/
𝑝𝑑𝑎𝑡𝑎(𝐱)
∇𝐱log 𝑝𝑑𝑎𝑡𝑎(𝐱)
Training data
𝑠𝜃 𝐱 : ℝ𝑑
↦ ℝ𝑑
≃
2 / 63

3. http://yang-song.github.io/blog/2021/score/
Background: Diffusion models and Stein score
Explicit score matching
Denoising score matching
Equivalent (Vincent et al. 2010)
𝜽∗
= argmin
𝜽
𝔼[‖∇𝐱log 𝑝 𝐱 − 𝐬𝜽(𝐱)‖2
2
]
𝜽∗
= argmin
𝜽
𝔼[‖∇𝐱log 𝑝 ෤
𝐱 | 𝐱 − 𝐬𝜽(෤
𝐱)‖2
2
]
3 / 63

4. • Once the score model is trained to optimality,
• i.e. 𝑠𝜃 𝐱 ≃ ∇𝐱𝑝(𝐱)
• Use Langevin dynamics to draw samples
𝐱𝑖+1 ← 𝐱𝑖 + 𝜖∇𝐱log 𝑝 𝐱 + 2𝜖𝐳𝑖
𝑖 = 0, 1, … , 𝐾
http://yang-song.github.io/blog/2021/score/
Background: Diffusion models and Stein score 4 / 63

5. http://yang-song.github.io/blog/2021/score/
Background: Diffusion models and Stein score 5 / 63

6. Background: Inverse problems
Imaging system
𝒜
𝐱
Ground truth image
𝜼
noise
𝐲
Measurement
• Problem: recover 𝐱 from noisy measurement 𝐲
• Ill-posed: Infinitely many solutions may exist
• We need to know the prior of the data distribution: how should the image look like?
6 / 63

7. Background: Examples of inverse problems
𝐲
𝐱
• Inpainting
⨀
• Deblurring (Deconvolution)
𝐱 𝐲
∗
• CS-MRI
𝓕
⨀
7 / 63

8. Posterior sampling for inverse problems
𝐱 ∼ 𝑝(𝐱|𝐲)
𝐲 = 𝒜(𝐱) + 𝜼
8 / 63

9. Posterior sampling for inverse problems
𝑝 𝐱 𝐲 =
𝑝 𝐲 𝐱 𝑝(𝐱)
𝑝(𝐲)
log 𝑝 𝐱 𝐲 = log 𝑝 𝐲 𝐱 + log 𝑝 𝐱 − log 𝑝(𝐲)
∇𝐱log 𝑝 𝐱 𝐲 = ∇𝐱log 𝑝 𝐲 𝐱 + ∇𝐱log 𝑝 𝐱
What we know
Measurement
process
9 / 63

10. Posterior sampling for inverse problems
Prior sampling
10 / 63
𝐱𝑖 ← 𝐱𝑖+1 + 𝑓 𝐱𝑖+1, 𝑖 + 1 − 𝑔 𝑖 + 1 2
∇𝐱𝑖
log 𝑝 𝐱𝑖 + 𝑔 𝑖 + 1 𝒛, 𝒛 ∼ 𝒩(𝟎, 𝐈)

11. Posterior sampling for inverse problems
Prior sampling
𝐱𝑖 ← 𝐱𝑖+1 + 𝑓 𝐱𝑖+1, 𝑖 + 1 − 𝑔 𝑖 + 1 2
∇𝐱𝑖
log 𝑝 𝐱𝑖 + 𝑔 𝑖 + 1 𝒛, 𝒛 ∼ 𝒩(𝟎, 𝐈)
Posterior sampling
𝐱𝑖 ← 𝐱𝑖+1 + 𝑓 𝐱𝑖+1, 𝑖 + 1 − 𝑔 𝑖 + 1 2
[∇𝐱𝑖
log 𝑝 𝐱𝑖 + ∇𝐱𝑖
log 𝑝(𝐲|𝐱𝑖)] + 𝑔 𝑖 + 1 𝒛
11 / 63

12. Posterior sampling for inverse problems
Prior sampling
Posterior sampling
Intractable!
12 / 65
𝐱𝑖 ← 𝐱𝑖+1 + 𝑓 𝐱𝑖+1, 𝑖 + 1 − 𝑔 𝑖 + 1 2
[∇𝐱𝑖
log 𝑝 𝐱𝑖 + ∇𝐱𝑖
log 𝑝(𝐲|𝐱𝑖)] + 𝑔 𝑖 + 1 𝒛
𝐱𝑖 ← 𝐱𝑖+1 + 𝑓 𝐱𝑖+1, 𝑖 + 1 − 𝑔 𝑖 + 1 2
∇𝐱𝑖
log 𝑝 𝐱𝑖 + 𝑔 𝑖 + 1 𝒛, 𝒛 ∼ 𝒩(𝟎, 𝐈)

13. The devil is in the likelihood
∇𝐱𝑡
log 𝑝(𝐲|𝐱𝑡)
𝑝 𝐲 𝐱0 = 𝒩(𝐲|𝒜𝐱𝟎, 𝜎2
𝑰)
∇𝐱0
log 𝑝 𝐲 𝐱0 = − 𝒚 − 𝒜𝐱𝟎 𝟐
𝟐
/𝜎2
Intractable!
13 / 63

14. Key 1. Factorization of the probabilistic graph
𝑝 𝐲 𝐱𝑡 = න 𝑝 𝐲 𝐱0, 𝐱t 𝑝 𝐱0 𝐱t 𝑑𝐱0
14 / 63

15. 𝑝 𝐲 𝐱t = න 𝑝 𝐲 𝐱0, 𝐱t 𝑝 𝐱0 𝐱t 𝑑𝐱0
= න 𝑝 𝐲 𝐱0 𝑝 𝐱0 𝐱t 𝑑𝐱0
Key 1. Factorization of the probabilistic graph 15 / 63

16. 𝑝 𝐲 𝐱0 = න 𝑝 𝐲 𝐱0, 𝐱t 𝑝 𝐱0 𝐱t 𝑑𝐱0
= න 𝑝 𝐲 𝐱0 𝑝 𝐱0 𝐱t 𝑑𝐱0
known partially
known
Key 1. Factorization of the probabilistic graph 16 / 63

17. 𝑝 𝐲 𝐱0 = න 𝑝 𝐲 𝐱0, 𝐱t 𝑝 𝐱0 𝐱t 𝑑𝐱0
= න 𝑝 𝐲 𝐱0 𝑝 𝐱0 𝐱t 𝑑𝐱0
Key 2. Tweedie’s formula 17 / 63
= 𝔼𝑝(𝐱𝟎|𝐱𝒕)[𝑝(𝐲|𝐱𝟎)]
Key 1. Tweedie denoising

18. 𝑝 𝐲 𝐱0 = න 𝑝 𝐲 𝐱0, 𝐱t 𝑝 𝐱0 𝐱t 𝑑𝐱0
= න 𝑝 𝐲 𝐱0 𝑝 𝐱0 𝐱t 𝑑𝐱0
Key 2. Tweedie’s formula 18 / 63
= 𝔼𝑝(𝐱𝟎|𝐱𝒕)[𝑝(𝐲|𝐱𝟎)]
≃ 𝑝(𝐲|𝔼[𝐱0|𝐱𝑡])]
Push
expectation
inside

19. 𝑝 𝐲 𝐱0 = න 𝑝 𝐲 𝐱0, 𝐱t 𝑝 𝐱0 𝐱t 𝑑𝐱0
= න 𝑝 𝐲 𝐱0 𝑝 𝐱0 𝐱t 𝑑𝐱0
Key 2. Tweedie’s formula 19 / 63
= 𝔼𝑝(𝐱𝟎|𝐱𝒕)[𝑝(𝐲|𝐱𝟎)]
≃ 𝑝(𝐲|𝔼[𝐱0|𝐱𝑡])]
= 𝑝(𝐲|ො
𝐱0)

20. Key 3. Jensen Bound 20 / 63

21. ∇𝐱𝑡
log 𝑝(𝐱𝑡|𝐲) = ∇𝐱𝑡
log 𝑝(𝐱𝑡) + ∇𝐱𝑡
𝑝 𝐲 𝐱𝑡
Diffusion Posterior Sampling (DPS) 21 / 63

22. ∇𝐱𝑡
log 𝑝(𝐱𝑡|𝐲) = ∇𝐱𝑡
log 𝑝(𝐱𝑡) + ∇𝐱𝑡
𝑝 𝐲 𝐱𝑡
Diffusion Posterior Sampling (DPS)
≃ ∇𝐱𝑡
log 𝑝(𝐱𝑡) + ∇𝐱𝑡
𝑝 𝐲 ො
𝐱0
Theorem 1.
22 / 63

23. ∇𝐱𝑡
log 𝑝(𝐱𝑡|𝐲) = ∇𝐱𝑡
log 𝑝(𝐱𝑡) + ∇𝐱𝑡
log 𝑝 𝐲 𝐱𝑡
Diffusion Posterior Sampling (DPS)
≃ ∇𝐱𝑡
log 𝑝 𝐱𝑡 + ∇𝐱𝑡
log 𝑝 𝐲 ො
𝐱0
≃ 𝑠𝜃∗ 𝐱𝑡, 𝑡 + ∇𝐱𝑡
log 𝑝 𝐲 ො
𝐱0
∇𝐱t
log 𝑝 𝐲 ො
𝐱0 = −𝜌∇𝐱t
‖𝐲 − 𝓐(ො
𝐱0)‖2
2
1. Gaussian
2. Poisson ∇𝐱t
log 𝑝 𝐲 ො
𝐱0 ≃ −𝜌∇𝐱t
‖𝐲 − 𝓐(ො
𝐱0)‖Λ
2
23 / 63

24. Diffusion Posterior Sampling (DPS) - Algorithm 24 / 63

25. Results: Super-resolution x16
Measurement (Gauss, 𝜎 = 0.05) DPS
25 / 63

26. Results: Inpainting (92%)
Measurement (Gauss, 𝜎 = 0.05) DPS
26 / 63

27. Results: Gaussian deblurring
Measurement (Gauss, 𝜎 = 0.05) DPS
27 / 63

28. Results: Motion deblurring
Measurement (Gauss, 𝜎 = 0.05) DPS
28 / 63

29. Results: Phase retrieval
Measurement (Gauss, 𝜎 = 0.05) DPS
29 / 63

30. Results: Non-uniform deblurring
Measurement (Gauss, 𝜎 = 0.05) DPS
30 / 63