ICLR 2023 spotlight

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