SDEdit: Guided image synthesis and editing with SDEs

Google slides

• https://docs.google.com/presentation/d/1eV-FlFZl1avKjVGgCsQ9tJJ7yCp5AYcW2ScocUuPi2E/edit?usp=sharing
  네이버 부스트캠프 내 피어세션에서 사용한 발표 자료입니다!

SDEdit 개요

SDEdit은 SDE(Stochastic Differential Equation)로 input image(guided)를 denoising하여 이미지를 합성 or 편집하는 방법입니다.
GAN 기반의 Guided image synthesis는 추가적인 data나 loss를 요구하는 반면에, SDEdit에서는 pre-trained 모델로 task 수행이 가능합니다.
또한, 특정 task에만 specific하게 작동하지 않는다(general)는 점도 주목할 만합니다.

about SDE

$dX_t=\mu(X_t,t)dt+ \sigma(X_t,t)d\bar{w}_t$

• $\mu(X_t,t)dt$: drift term
  - 시간에 따라 식이 변하는 정도 표현
• $\sigma(X_t,t)d\bar{w}_t$: wiener process
  - random process (with noise)를 표현
  SDE에 대해 살펴보겠습니다. SDE는 확률적 미분 방정식으로, 시간 t에 따라 변하는 state를 표현합니다. 일반적인 ODE와 다른 점은, noise를 통하여 시간에 따른 randomness가 부여된다는 점입니다. 더 자세하게 forward 과정과 backward 과정을 따로 알아보겠습니다.

Forward

$\mathbf{x}(t) = \alpha(t)\mathbf{x}(0) + \sigma(t)\mathbf{z}$

• $\mathbf{x}(t)$: timestep $t$에 대한 diffused image
• $\mathbf{x}(0)$: data 분포로부터 얻은 샘플
• $\sigma(t)$: noise $\mathbf{z}$를 조정하는 scalar function
  - 사전에 정의됨! (i.e. linear funciton 등)
  - $\mathbf{z}$: noise <- $\mathbf{z} \sim \mathcal{N}(0,\mathbf{I})$
• $\alpha(t)$: 샘플 이미지 반영 정도
  - VE-SDE에서는 1로 고정됨.
  해당 과정을 거쳐 샘플 이미지에 noise를 주입하여 diffusing합니다.

Forward in SDEdit

SDEdit에서 forward가 적용되는 과정은 다음과 같습니다. (stroke-based image synthesis라는 가정에서) image 분포와 storke 분포가 존재할 때, forward를 통해서 prior 분포와 image 분포를 섞어주게 됩니다. 논문에서는 다음의 과정을 Perturb로 표현하였습니다.

Reverse process

$d\mathbf{x} = [\mathbf{f}(\mathbf{x},t) - g^{2}(t)\nabla_\mathbf{x}\log p_t(\mathbf{x})]dt + g(t)d\mathbf{w}$

• $\nabla_\mathbf{x}\log p(x)$: score function
  - log probability의 gradient입니다.
  - 실제 분포 p(x)를 직접 구하는 것이 불가능하므로, score function로 이를 추정하는 방식입니다.
  - gradient를 사용하는 이유는, log-likelihood를 구할 때 계산이 더 편리해지는 경우가 생기기 때문입니다. (i.e. energy-based model)
• loss func: $\frac {1} {2} \mathbb{E}_{p_{data}}[||\nabla_\mathbf{x}\log {p_{data}}(x) - \nabla_\mathbf{x}\log p_\theta(x)||^2_2]$
  - data 분포의 score function과 모델의 score function을 L2 loss로 최소화!
  score function을 모델로 학습하여 위의 loss function을 최적화하는 것으로, reverse SDE를 구할 수 있습니다. 해당 과정을 통해 원하는 데이터를 생성할 수 있게 됩니다.

SDEs (VE, VP)

$d \mathbf{x}(t) = [ - \frac{d[\sigma^{2}(t)]}{dt}\nabla_\mathbf{x}\log p_t(\mathbf{x}) ] dt + \sqrt{\frac{d[\sigma^{2}(t)]}{dt}}d\bar{\mathbf{w}}$

논문에서는 주로 사용되는 SDE인 VE-SDE(Variance Exploding SDE), VP-SDE(Variance Preserving SDE) 총 2개의 SDE를 사용하였습니다. 그 중에서도 VE-SDE에 대해서만 이번 포스트에서는 다뤄보겠습니다. VE-SDE의 diffusing을 거치면, diffused image가 다음의 분포를 띄게 됩니다.

• diffused: $x(1) \sim \mathcal{N}(0,\sigma^2(1))$

SDE solving (Euler-maruyama)

SDE를 해결하기 위해서는, 각 변수에 대해 적분하는 과정이 필요합니다. 이에 대한 계산이 매우 복잡하기 때문에, 다음과 같이 timestep의 변화량을 적용하여 iteration 과정으로 SDE를 풀이하는 Euler-maruyama방법을 사용합니다.

Reverse process in SDEdit

Forward 과정에서 얻어진 perturbed 분포에서, reverse process를 거쳐 이미지(파란 점)가 image 분포로 이동하는 것을 확인할 수 있습니다.

Pesudo Code

VE-SDE로 풀이한 SDEdit의 대략적인 코드는 다음과 같습니다. forward 적용한 후에, N step에 대해서 reverse SDE를 적용하여 원하는 이미지를 생성합니다.

Results

SDEdit은 Realism과 Faithfulness, 두 가지의 지표로 평가되었습니다.

• Realism: 더 현실성 있는지에 대한 지표입니다. 정성 평가와 같은 뜻으로 해석됩니다.
  - 사람들: MTurk
  - Neural network: FID (GAN-based evaluation)
• Faithfulness: 생성된 이미지가 가이드(input) 이미지와 비슷한 정도를 측정합니다. 정량 평가와 같은 뜻으로 해석됩니다.
  - L2 distance 사용
  왼쪽 그림에서 알 수 있듯이, 두 지표 사이에는 trade-off가 존재합니다. 실험 결과 $t_0$를 [0.3, 0.6]으로 설정하였을 때 가장 좋은 성능이 나온다고 합니다. (guide 이미지가 흰 색 사진일 경우, 더 큰 $t_0$ 값을 설정해야 합니다.)

Faithfulness + Realism

LSUN (bedroom), CelebA-HQ 데이터셋에 대한 성능 지표입니다. SDEdit 모델이 다른 모델들에 비해 두 영역 모두에서 좋은 성능을 보이고 있는 것을 확인할 수 있습니다.

Realism

GAN baseline (StyleGAN2-ada)와 SDEdit 간의 realism을 비교한 결과입니다. SDEdit이 더 현실성 있는 이미지를 생성하는 것을 확인할 수 있습니다.
예시로 두 빨간 박스를 비교해보면, GAN baseline에서는 흰색 부분을 모호하게 처리한 반면에 SDEdit에서는 꽤 정확하게 창문을 생성한 것을 확인할 수 있습니다.

모델 사용해보기

• 공식 코드에서 제공된 dataset과 model이 모두 막혀있습니다...
• 모델의 경우는, 따로 학습된 모델이 아니기 때문에 hugging face에서 pre-trained ddpm 모델을 사용하였습니다.
• dataset은 이미지와 mask가 페어의 형태로 존재하는 것으로 확인하였고, 논문에 기재된 대로 마스크의 값을 1로 설정하여 stroke-based 이미지 합성을 시도하였습니다.
  

결과 (troubled)

• pinta로 그린 custom image를 input으로, mask는 1.로 설정하여 feed한 결과입니다.
• 모델이 이미지 합성을 하지 못하고, 그대로 복원하는 데에 그친 것을 확인할 수 있습니다.
• 여러 테스트 결과(with 부스트캠프 멘토님[박성현 멘토님 감사합니다ㅠㅠ]), 논문에 기재된 마스크 value가 반대임을 알게되었습니다..!
  - 논문에선 editable 영역을 1.로 설정하였다고 했지만, 코드에선 0.이 editable 영역이었습니다.

결과 (문제 해결 후)

모델 실행 결과입니다. pre-trained ddpm(LSUN church)으로 score function 예측 후, reverse SDE를 적용하여 다음의 이미지들을 합성할 수 있었습니다. Stable diffusion 공식 코드에서 img2img를 SDEdit으로 구현하였다고 작성되었는데, 해당 코드에서는 DDIM을 바로 적용한 것으로 보였습니다. (DDIM의 probability flow ODE가 지금까지 살펴봤던 VE-SDE의 probability flow ODE와 동일하기 때문입니다.)

Reference

[1] Chenlin Meng, Yutong He, Yang Song (stanford), SDEDIT: GUIDED IMAGE SYNTHESIS AND EDITING WITH STOCHASTIC DIFFERENTIAL EQUATIONS, https://arxiv.org/pdf/2108.01073
[2] Jiaming Song, Chenlin Meng & Stefano Ermon, DENOISING DIFFUSION IMPLICIT MODELS, https://arxiv.org/pdf/2010.02502
[3] hugging face, pre-trained ddpm model, https://huggingface.co/google/ddpm-church-256