[논문리뷰] Generative Modeling by Estimating Gradients of the Data Distribution(Score-based Model) NCSN 리뷰

0. 들어가며

DDPM 논문을 읽다보면 다음과 같은 말이 등장합니다.

Our best results are obtained by training on a weighted variational bound designed according to a novel connection between diffusion probabilistic models and denoising score matching with Langevin dynamics

 

처음 DDPM을 읽었을 때 무슨 뜻인지 모르고 그냥 넘어갔던 기억이 있습니다. 본 논문에서는 Denoising score matching with Langevin dynamics 방법론을 다룹니다. 해당 방법을 이해하고 이 포스팅을 끝까지 보시면 결국 DDPM이 Score-based model과 weight만 다를 뿐 같은 동작을 한다는 것을 알 수 있게 됩니다(Chapter 6에서 증명).

 

1. Score Matching

일반적인 Likelihood based model은 Dataset에 대한 분포를 maximum likelihood 방법을 통해 근사하게 됩니다. 하지만 해당 방법에는 여러 intractable problem이 존재합니다. VAE 같은 경우에는 Variational approximation이라는 방법을 통해서 intractable한 $p(z\x)$를 근사하는 방법을 택하는 것 처럼 likelihood based model은 intractable problem을 우회하려고 노력합니다.

 

https://rla020.tistory.com/43

[논문리뷰] Auto-Encoding Variational Bayes(VAE) 모든 수식 알아보기

 

어떤 분포의 PDF는 다음과 같이 쓸 수 있습니다. $$p_\theta(x) = \frac{e^{-f_\theta(x)}}{Z_\theta}$$

위 식을 살펴보면 분자에는 어떤 input에 대한 값이 들어가고, 분모에는 분자 값을 확률 분포(더해서 1)이 될 수 있는 값이 들어갑니다.(이를 nomalization constant라고 합니다.) $Z_\theta$는 모든 x에 대해 $e^{-f_\theta(x)}$를 다 합친 값이라고 생각하면 좋을 것 같습니다. 하지만 당연히 $Z_\theta$를 구하는 것은 intractable합니다. continuous하다면 모든 x에 대한 적분, discrete하다면 모든 x에 대한 summation이 필요하기 때문이죠. 한번 nomalization constant 값을 구한다고 해도 $\theta$를 업데이트 하면 다시 모든 x에 대한 constant를 구해야 합니다. 따라서 무조건 intractable할 수 밖에 없습니다.

 

이러한 문제를 본 논문에서는 Score를 통해 해결하려고 합니다. Score 라는 것은 통계학에서 많이 사용되는 것으로 어떤 함수의 log 미분이라고 생각하면 쉽습니다. 본 논문에서는 어떤 함수를 input x로 미분한 함수를 score function이라고 정의합니다.

• $p_\theta(x) = \frac{e^{-f_\theta(x)}}{Z_\theta}$
• Score function = $\nabla_xlogp(x)$

$p_\theta(x) = \frac{e^{-f_\theta(x)}}{Z_\theta}$에 대한 score를 계산하면 다음과 같습니다.

재미있게도 log를 씌워 score를 구하게 되면 $Z_\theta$항이 사라지게 됩니다. 그러면 결국 score는 tractable한 연산이 됩니다.

 

Score를 통해서 PDF를 구할 수 있을 까요? 다음 그림을 보면 어느정도 감이 잡힙니다.

위의 그림은 2차원 gaussian의 score를 구한 vector field입니다. 결국 score란 함수의 log 미분이기 때문에 Mode에서는 작은 값을 갖고 멀어질 수록 값이 커지게 됩니다. 신기하게도 log를 적용해서 미분을 하게되면 mode와 멀 수록 큰 값을 갖게 됩니다. 그러면 위의 그림에서 vector field의 화살표만 잘 따라가면 어떠한 공간에 도착을 할 텐데 그 공간이 바로 PDF의 likelihood가 큰 공간이 됩니다. 이런식으로 score function을 이용하여 PDF를 구할 수 있게 됩니다.

 

그리고 우리는 Score function을 근사하는 모델을 만들 수 있습니다. 이 모델은 Fisher divergence를 통해 학습시킬 수 있습니다. Fisher divergence란 두 분포의 유사도를 측정하기 위해 derivative 즉 도함수를 사용하는 것을 말합니다.

$$ Fisher Divergence = E_{p(x)}[||\nabla_x logp(x) - \nabla_x log q(x)||_2^2] $$

 

위의 식을 우리의 모델을 포함시켜서 바꾸면

$$ Fisher Divergence = E_{p(x)}[||\nabla_x logp(x) - s_\theta(x)||_2^2] $$

가 됩니다.

위 식의 장점은 $S_\theta(x)$에 대한 가정이 필요 없어지고 단순히 GT와 model output의 L2 distance만 계산하면 된다는 점 입니다.

 

1.1 Detail in Score matching

하지만 VAE와 비슷한 문제가 발생합니다. 우리는 Score를 구해야 하는데 $p(x)$를 알지 못하기 때문에 당연히 $\nabla \log p(x)$도알지 못하고, 그렇다면 GT가 없기 때문에 Fisher divergence를 통해 모델을 학습 시킬 수 없게 됩니다. 즉

$$ Fisher Divergence = E_{p(x)}[||\nabla_x logp(x) - s_\theta(x)||_2^2] $$

를 구할수 없다는 것입니다.

 

그래서 이를 우회할 수 있는 방법이 필요한데, Hyvärinen이라는 분이 쓴 Estimation of non-normalized statistical models by score matching(2005)논문에서 그 방법을 찾습니다.

$$ E_{p(x)}[\frac{1}{2}||s_\theta(x)||_2^2 + tr(\nabla_xs_\theta(x))] $$

수식을 이용하면 $\nabla_x \log p(x)$ 없이도 학습이 가능하게 됩니다. 어떻게 저 수식이 나오는지는 제가 손으로 증명해봤습니다.

$ E_{p(x)}[\frac{1}{2}||s_\theta(x)||_2^2 + tr(\nabla_xs_\theta(x))] $ 식을 보면 두가지 term으로 나뉘는데,

• $\frac{1}{2}||s_\theta(x)||_2^2$는 거리이기 때문에 minimize하려면 0이 되어야 하고
• $tr(\nabla_xs_\theta(x))$는 gradient이기 때문에 음수로 쭉 가게 됩니다.

재미있게도 어떤 함수 $p(x)$의 도함수(score)인 $s_\theta(x)$가 0으로 간다는 것은 $p(x)$의 극대 혹은 극소를 말하게 되는데, 이때 도함수의 미분인 이계도함수 $\nabla_x s_\theta(x)$가 음수로 간다는 것은 극대값이 되게 됩니다. 즉 위의 식은 $p(x)$의 maximum likelihood를 뜻하게 됩니다.

 

하지만 위의 식은 trace를 구해야 한다는 문제가 있었습니다. trace를 구하는것은 computationally high cost였고 trace를 구하지 않고 우회하는 두가지 방법을 제시합니다.

 

1. Denoising score matching

해당 방법은 data에 perturbation noise를 더해서 score matching을 진행하는 방식입니다.

• Data point x -> $q_\sigma(\tilde{x}|x)$ ($\sigma$를 갖는 gaussian noise로 puterbation)

이렇게 한다면 식이 다음과 같이 변하게 됩니다.

2. Sliced score matching

Sliced score matching이라는 방법도 있는데, 저도 해당 논문을 읽지 않아서 간단하게만 설명하겠습니다. 이 방법은 Random projection을 통하여 근사하는 방법으로 v는 simple distribution of random vector입니다.

 

위의 두가지 방법 중 논문에서는 1의 방법을 통해 실험을 진행합니다.

 

2. Langevin Dynamics

이제 Score function을 근사하는 모델은 만들었으니 이미지를 Sampling하는 방법도 필요하게 됩니다. 논문에서는 Langevin Dynamics라는 방법을 통해서 sampling을 진행하게 됩니다.

 

Langevin dynamics는 MCMC방법을 통해 Score function만 가지고 $p(x)$를 sampling 할 수 있는 방법론 입니다.

• $x_0 \sim \pi (x)$ : $\pi ()$는 우리가 구할 수 있는 simple distribution(gaussian ...)
• $z_i \sim N(0,I)$

위 식이 Langevin dynamics방법론 입니다. 어떤 gaussian noise $x_0$를 시작으로 $\epsilon$만큼의 step을 정해서 Score와 곱해 더해주고 normal distribution noise를 더해 $x_{i+1}$을 만들게 됩니다. 위 식이 다음을 만족하면 sampling error는 무시할 만한 수준이 됩니다.

• $\epsilon \to 0$, $K \to \infty$

우리는 $\nabla_x \log p(x)$는 모르지만 score matching을 통해 근사한 $S_\theta(x)$는 알기 때문에 Langevin dynamics를 사용할 수 있습니다.

 

간단히 생각하면 noise를 Score vector field에 뿌려서 gradient를 타고 mode로 수렴하는 방식이라고 생각하면 편할 것 같습니다.

 

전체적인 과정은 다음과 같습니다.

 

3. Problem in Score Matching

이렇게 간단하게 generative model이 만들어지면 좋겠지만, 단순히 위에서 제시한 방법론을 이용하면 문제가 발생합니다.

 

3.1 Manifold 가정

Manifold가정으로 인해 고차원의 space 대부분은 sparse합니다. 그런데 Score라는 것은 gradient이기 때문에 주변 공간이 정의가 되어 있지 않으면 제대로 구해지지 않습니다. 본문에서는 Data distribution의 support가 whole space 일 경우에만 score estimator가 작동한다고 하고 있습니다.

 

3.2 Inaccurate score estimation

3.1과 비슷하게 sample이 부족한 공간의 score는 부정확 하다고 주장하고 있습니다. 직관적으로 생각해도 sample이 없기 때문에 모델의 학습이 제대로 이루어지지 않는다는 것을 알 수 있습니다.

 

3.3 답은 perturbation!

위의 두가지 문제를 해결하기 위해 Image perturbation을 진행합니다.

위 그림의 왼쪽 이미지는 perturbation을 진행하지 않았을 때의 Sliced score matching loss이고, 오른쪽은 진행 했을 때의 loss 입니다. Perturbation을 통해 이미지가 blurry해지면서 data distribution의 support를 whole space로 펼쳤다고 생각 할 수 있습니다.

 

sample이 부족한 공간의 score 또한 perturbation을 통해 매꿔주면서 score를 정확하게 구할 수 있었다고 주장합니다.

 

4. Problem in Langevin Dynamics

sampling 과정인 Langevin Dynamics에서도 문제가 발생했습니다. 먼저 어떤 data 분포가 Mixture of gaussian이라고 가정하면 다음과 같은 수식이 됩니다.

위 수식에서 $\pi p_1(x)$의 score만 구해본다면

가 되고, weight가 사라지게 됩니다. 

(a)그림은 직접 Mixture of gaussian에서 iid sampling한 것이고 (b)는 Langevin dynamics를 통해 sampling한 분포인데, (b)에서는 확실히 weight가 무시되는 것을 볼 수 있습니다. 그래서 이러한 문제를 Annealed Langevin dynamics를 통해 해결합니다.

 

4.1 Annealed Langevin Dynamics

Annealed Langevin Dynamics 방법은 먼저 perturbation이 많이된 상태의 data 분포를 sampling하고 그 값을 초기 값으로 점점 더 작은 perturbed data 분포를 sampling하는 방법입니다. 그림을 보면 이해가 쉽습니다.

위 그림에서 $\sigma_3$는 perturbation이 많이된 데이터를 sampling 한 것 이고 그 마지막 값을 $\sigma_2$의 초기값으로 시작하여 sampling을 합니다. 이런식으로 sampling을 진행하면, 원래의 weight를 잘 보존할 수 있다고 합니다.

 

5. NCSN

논문에서 제시한 모델을 Noise Conditional Score Network로 NCSN이라고 합니다.

5.1 Training NCSN

이제 모든 문제를 해결했고, 모델을 만들게 됩니다. 먼저 perturbation을 위한 $\sigma$를 다음과 같이 정합니다.

$$\{\sigma_i \}_{i=1}^L, \frac{\sigma_1}{\sigma_2} = ... = \frac{\sigma_{L-1}}{\sigma_L} > 1 $$

• sigma level을 1부터 L level을 준비한다.
• 이 때 $\sigma_1$이 $\sigma_2$보다 커야하고 $\sigma_2$가 $\sigma_3$보다 커야한다. (그래야 위의 수식이 1보다 커진다)

4.1의 Annealed Langevin Dynamics와 헷갈릴 수 있지만 단순히 notation이 반대로 됐을 뿐입니다. 4.1에서는 $\sigma_3$이 더 컸지만, 논문의 수식에서는 $\sigma_1$이 가장 크게 설정합니다. 단순한 notation의 오류로 위의 그림에서의 $\sigma$정의가 반대로 됐다고 생각하면 좋을 것 같습니다.

 

위의 $\sigma$를 통해서 perturbation은 다음과 같이 진행할 수 있습니다.

$$p_{\sigma_i}(x) = x + \sigma_i z , z \sim N(0,I)$$

 

논문에서는 Denoising score matching을 사용해서 학습을 진행합니다.

1. perturb data x는 gaussian의 log 미분으로 간단하게 나타내 진다.
2. Loss는 모든 level로 perturbation을 진행하고, weight $\lambda$를 곱해서 다 더한다.

 

5.2 Sampling method

$S_\theta(x,\sigma)$가 학습되면 Langevin dynamics를 통해 sampling을 진행하게 됩니다. 이때 noise annealing방식을 사용한다고 했는데 알고리즘은 다음과 같습니다.

1. $x_0$는 noise이다.
2. noise부터 시작해서 T번 만큼 Langevin dynamics를 돌린다.
3. 이 때 perturbation이 $\sigma_1$만큼(가장 큰 sigma) 들어간 이미지로 생각하고 sampling한다.
4. T번이 끝나면 해당 값을 초기값으로 $\sigma_2$만큼 들어간 이미지로 생각하고 sampling한다.
5. 반복

6. DDPM과의 관계

Score Matching중 Denoising Score Matching의 수식을 보면 다음과 같습니다.

DDPM의 수식을 보면 다음과 같습니다.

 

Denoising Score Matching의 수식을 풀어보면 다음과 같습니다.

1. input x에 gaussian noise를 추가한 Denoising score matching
2. noise perturbation을 --2의 수식과 같이 쓸 수 있음
3. 이 수식을 Gaussian 공식에 대입하면 --3처럼 쓸 수 있음
4. 해당 식의 score를 구하면 $-\frac{\tilde{x}-x}{\sigma^2}$가 되고 reparameterization trick($\tilde{x} = x + \sigma \epsilon$)을 이용해서 다시 쓰면 빨간색 글씨 처럼 쓸 수 있음
5. 결국 우리의 모델 $S_\theta(\tilde{x},\sigma)$는 다음과 같이 해석 될 수 있음
   
   1. $D_{\theta}(\tilde{x},\sigma)$ : perturbed data $\tilde{x}$를 받아서 denoising하여 x를 output으로 주는 모델
   2. $\tilde{x}$를 받아서 perturbation $\epsilon$을 예측하는 모델
6. 5번의 결과를 Denoising Score matching 수식에 넣으면 --6과 같음
7. --7처럼 쓸 수 있고
8. 결국 실제 $\epsilon$을 예측하는 $\epsilon_\theta$를 모델링 하는 것과 같음

이 됩니다. 결국 따져보면 DDPM에서 noise $\epsilon$을 예측하는 것은 weight만 다를 뿐 score matching을 하는 것과 같다는 것입니다. 처음 증명을 해보고 정말 신기했습니다. noise를 예측하는 것이 score matching과 같다니.. 참 generative model은 신기한 것 같습니다. 그렇기 때문에 DDPM에 Score matching with Langevin Dynamics라는 말이 많이 등장합니다.

 

7. 마치며

 

위 이미지는 NCSN으로 sampling한 이미지 입니다. 저자들은 NCSN을 통해 GAN에 비교할만한 quality의 image를 생성할 수 있었다고 주장합니다.

 

사실 NCSN자체는 지금 쓰이지 않는다고 생각합니다. 하지만 결국 DDPM과 Score based model은 같고, 두 모델이 SDE라는 Framework안에서 존재한다는 논문이 나오는데, 그 논문을 통해 generative model을 한단계 깊숙히 이해하기 위한 발판이 되는 논문이라고 생각합니다.