Score-Based / Diffusion Model[4] - NCSNv2

 

Introduction

이번에는 Score-based generative model의 두번째 논문인 NCSN v2에 대해서 리뷰를 해보고자 합니다. 본 논문은 Score-based generative model에 대한 이론적인 접근으로 하이퍼파라미터를 설정할 때 어떻게 접근해야 하는지에 대해서 정리합니다.

 

우선 Noise Conditional Score Network를 통해 score-based generative model에서의 가능성을 보여주었습니다. 또한, CIFAR-10 데이터셋에서도 좋은 성능을 보여주었습니다. 하지만, 이런 문제점에도 불구하고 고해상도를 생성했을 때 (from 64 $\times$ 64 to 256 $\times$ 256) poor한 성능을 보여준다고 합니다. 그래서 저자는 이것에 대해서 설명하는데,

 

첫번째로, score function은 denoising score matching을 통해서 학습한다고 합니다. 직관적으로 score network는 Gaussian noise로 블러처리된 이미지들을 denoising하기 위해 학습됩니다. 여기서 주요한 insight는 multiple noise scales를 사용해서 score network는 coarse-grained와 fine-grained한 이미지를 둘 다 학습하게 됩니다. 그러나, 이 noise scales를 어떻게 고르냐에 대한 문제점이 있습니다. NCSN에서 추천했던 파라미터 셋팅은 32 $\times$ 32 이미지에서는 잘 작동을 하지만, 그 이상에서는 잘 안된다고 합니다.

두번째로, Langevin dynamic을 사용해 생성한다고 합니다. white noise로 부터 시작되고 score network를 사용해서 progressive하게 이미지를 denoising 을 합니다. 그러나 이 procedure는 high-dimension일 때 필연적으로 imperfect score network에서 사용 될 때 수렴이 굉장히 오래걸리거나 실패할 가능성이 있다.

 

그래서, score-based generative model에서 high-resolution image에서 사용할 수 있는 techniques들을 제안을 합니다. 본 논문에서는 크게 3가지 category (noise scales, Langevin dynamics, stability) 로 총 5가지의 테크닉을 제안합니다.

 

Choosing noise scales

NCSN에서는, pixel value의 range가 [0,1] 일 때는, $L=10, \sigma_{1}=1$, and $\sigma_{L}=0.01$ 로 설정하는 것이 기본이라고 합니다. 마지막에서는 작은 noise를 더하고 싶기 때문에 $\sigma_{min}=0.01 << 1$는 충분히 이유가 있다고 합니다. 하지만, 이것만으로는 high-resolution에서 안되는 이유가 설명되지 않는다고 합니다.

 

1. $\sigma_{min}=1$이 가장 적절한가? 아니면 다른 데이터셋에서는 조절을 해야하는가?
2. geometric progression은 어떻게 설정해야 하는가? (noise scheduling)
3. $L=10$은 과연 다른 데이터셋에서도 적절한가?

 

1. Initial Noise Scales

먼저, Annealed Langevin Dynamic (ALD)은 Large noise로 부터 Coarse한 sample을 만들어내고, Small noise로 부터 Fine한 sample을 반복적으로 refinement해서 만들어낸다고 합니다. 그렇기 때문에 $\sigma_{max}$가 sample의 diversity를 결정하는 중요한 요소가 된다고 합니다. 그렇기 때문에 diversity를 높이기 위해서 $\sigma_{max}$가 가능한 크게 키워야 합니다. 그러나, 무작정 키우게 된다면 ALD에서 많은 리소스가 필요하다고 합니다. 따라서 이것을 guide하는 이론적인 아이디어를 제공한다고 합니다.

 

Proposition 1. $\hat{p}_{\sigma_1}(x)=\frac{1}{N}\sum_{i=1}^{N}p^{(i)}(x)$이고 $p^{(i)}(x)=\mathcal{N}(x | x^{(i)}, \sigma_{1}^2I)$ 일 때, $r^{(i)}(x)=\frac{p^{(i)}}{\sum_{k=1}^{N}p^{(k)}(x)}$라면, score function은 $\nabla_x \log{\hat{p}_{\sigma_1}(x)}=\sum_{i=1}^{N}r^{(i)}(x)\nabla_x\log{p^{(i)}(x)}$이다. 또한, 

$$ \mathbb{E}_{p^{(i)}(x)}[r^{(j)}(x)] \leq \frac{1}{2} \exp{(-\frac{\Vert x^{(i)} - x^{(j)} \Vert_2^2}{8\sigma_1^2}) } $$

을 만족한다.

 

proof)$$\begin{align} \nabla_x\log{\hat{p}_{\sigma_1}}(x)&=\nabla_x\log{(\frac{1}{N}\sum_{i=1}^Np^{(i)}(x))}\\ &=\sum_{i=1}^N\frac{p^{(i)}(x)\nabla_x\log{p^{(i)}(x)}}{\sum_{j=1}^Np^{(j)}(x)} \\ &=\sum_{i=1}^Nr^{(i)}\nabla_x\log{p^{(i)}(x)} \end{align}$$

 

으로 정리됩니다. 또한, $x\in\mathbb{R}^D$라고 한다면, 

 

$$\begin{align} \mathbb{E}_{p^{(i)}(x)}[r^{(j)}(x)]&= \int \frac{p^{(i)}(x)p^{(j)}(x)}{\sum_{k=1}^N p^{(k)}(x)}dx \leq \int \frac{p^{(i)}(x)p^{(j)}(x) }{p^{(i)}(x) +p^{(j)}(x)}dx \\ &= \frac{1}{2} \int \frac{1}{ \frac{1}{p^{(i)}(x)} + \frac{1}{p^{(j)}(x)}}dx \leq \frac{1}{2}\int \sqrt{p^{(i)}(x)p^{(j)}(x)}dx \\ &= \cdots \\& = \frac{1}{2} \exp{(-\frac{\Vert x^{(i)} - x^{(j)}\Vert_2^2}{8\sigma_1^2})} \end{align} $$ 

 

결국, Langevin dynamic으로 인해 $p^{(i)}$에서 부터 $p^{(j)}$ 까지 전이를 시키려고 한다면, $\mathbb{E}_{p^{(i)}(x)}[r^{(j)}(x)]$값이 커야 합니다.  즉, 이 값을 키우기 위해서 Hyperparameter를 조절해야 한다는 의미이며, $\frac{1}{2} \exp{(-\frac{\Vert x^{(i)} - x^{(j)}\Vert_2^2}{8\sigma_1^2})}$ 을 이용해 $\sigma_1$을 키워야한다고 말하고 있습니다. 

 

Technique 1. Training set의 모든 pairs 중 Euclidean Distance가 가장 큰 값으로 $\sigma_{max}$ 으로 설정한다.

 

$\sigma_{1}$을 키웠을 때, 생성되는 이미지가 눈에 띄게 더 다양한 이미지를 생성해내는 것을 볼 수 있습니다. 정량적으로도, 데이터들의 pairwise distance도 18.65로 실제 데이터인 17.78과 비슷하게 나오고 $\sigma_1=1$일 때 10.12인 것과 확연히 차이가 나는 것을 볼 수 있습니다.

 

2. Initial Noise Scales

$\sigma_{1}$과 $\sigma_{L}$을 설정했으면, L과 그 사이의 $\{\sigma_{i}\}$ 에 대해서도 특정할 필요성이 있다고 합니다. 이전 논문에서, Score-based generative model의 성공은 모든 $i$들에 대해 $p_{\sigma_{i-1}}$의 High density region을 학습할 수 있는 충분한 training data 숫자가 보장되어야 한다고 합니다. 직관적으로, Langevin dynamic을 통해서 $p_{\sigma_{i-1}}$으로부터 $p_{\sigma_i}$으로 갈 때 신뢰성 있는 gradient signal이 필요하다는 뜻입니다.

 

하지만, 모든 $\{\sigma_{i}\}$의 grid search는 매우 비싼 작업입니다. 그래서 본 논문에서는 좋은 noise guidance를 찾기 위해서 이론적인 지침을 제공한다고 합니다. 특히, 높은 차원의 $x$에서 $p_{\sigma_i}$의 분포를 이해하는 것이 먼저입니다. $p_{\sigma_i}(x)$를 구면좌표계로 바꾸면 $p(\phi)p_{\sigma_i}(r)$로 분해했을 때, $p_{\sigma_i}(r)$이 isotropic Gaussian이므로 angular 성분인 $p(\phi)$는 uniform하고 모든 noise scaling을 공유한다고 합니다. ( $p_{\sigma_i}(r)$이 isotropic [방향에 상관없이]하기 때문에 그에 따른 각도는 고르게 분포되어 있고 (uniform) noise에도 영향을 받지 않으므로 공유한다는 뜻으로 해석됩니다.) $p_{\sigma_i}(r)$을 위해서 아이디어를 제공하면,

 

 

Proposition 2. $x \in \mathbb{R}^{D} \sim \mathcal{N}(0, \sigma^2I)$와 $r=||x||_2$라고 하자. 그렇다면, $$p(r)=\frac{1}{2^{D/2-1}\Gamma({D/2})}\frac{r^{D-1}}{\sigma^{D}}exp(-\frac{r^2}{2\sigma^2}) \text{ 그리고}, r-\sqrt{D}\sigma \xrightarrow{d} \mathcal{N}(0, \sigma^2/2), \text{when }D\rightarrow\infty.$$

proof) $x\sim\mathcal{N}(0,\sigma^2I)$이므로, $s=||x||_2^2/\sigma^2\sim\chi^2_{D}$를 얻을 수 있고, 즉

$$p_s(s)=\frac{1}{2^{D/2}\Gamma(D/2)}s^{D/2-1}e^{-s/2}$$

는 것을 의미한다.

$r=||x||_2=\sigma\sqrt{s}$이기 때문에 변수변환을 사용하면,

$$p(r)=\frac{2r}{\sigma^2}p_s(s)=\frac{1}{2^{D/2-1}\Gamma(D/2)}\frac{r^{D-1}}{\sigma^D}\exp(-\frac{r^2}{2\sigma^2})$$

$x\sim\mathcal{N}(0,\sigma^2)$일 때, $x^2/\sigma^2\sim\chi^2_1$이고 $\mathbb{E}[x]=\sigma^2, \text{Var}[x]=2\sigma^4$를 얻을 수 있다. 결과적으로, $x_1, x_2, ..., x_D $${i.i.d.} \atop \sim$ $\mathcal{N}(0, \sigma^2)$라면, 큰 수의 법칙과 $D\rightarrow\infty$ 로 중심극한 정리를 적용하면, 

$$\frac{x^2_1+x^2_2+...x^2_D}{D}\xrightarrow{p}\sigma^2$$

로 확률수렴하고

$$\sqrt{D}(\frac{r^2}{D}-\sigma^2)\xrightarrow{d}\mathcal{N}(0,2\sigma^4)$$

에서, delta method를 적용하면,

$$\sqrt{D}(\frac{r}{\sqrt{D}}-\sigma)\xrightarrow{d}\mathcal{N}(0,\sigma^2/2)$$

임을 알 수 있고, 따라서 $r-\sqrt{D}\sigma\xrightarrow{d}\mathcal{N}(0,\sigma^2/2)$가 성립합니다.

 

실제로, 이미지 데이터의 차원의 범위가 다양하고, 무시해도 될 만큼의 오차를 갖고 $p(r)\approx\mathcal{N}(r|\sqrt{D}\sigma, \sigma^2/2)$를 보증할 정도의 차원이 크다고 합니다. 그러므로, $p_{\sigma_i}(r)=\mathcal{N}(r|m_i, s^2_i)$로 간단하게 정의하기로 합니다. ($m_i=\sqrt{D}\sigma, s_i^2=\frac{\sigma^2}{2}$)

 

목표를 다시 리마인드하자면, $p_{\sigma_i}(x)$에서의 sample이 $p_{\sigma_{i-1}}(x)$의 high denstiy region을 커버하도록 하는 것입니다. $p(\phi)$가 모든 noise scale에 공유되므로 (각 $i$에 영향을 미치지 않는다는 뜻입니다.), $p_{\sigma_i}$는 $p_{\sigma_{i-1}}$의 angular 성분을 커버하고 있습니다. 그러므로 우리의 목표는 $p_{\sigma_i}$와 $p_{\sigma_{i-1}}$의 radial 성분이 ovelap 되기를 원합니다. $p_{\sigma_{i-1}}(r)$은 three-sigma 즉, $\mathcal{I}_{i-1}=[m_{i-1}-3s_{i-1}, m_{i-1}+3s_{i-1}]$에서 high density를 갖기 때문에, 일반적인 선택은 $p_{\sigma_i}(r\in\mathcal{I}_{i-1})=\Phi(\sqrt{2D}(\gamma_i-1)+3\gamma_i)-\Phi(\sqrt{2D}(\gamma_i-1)-3\gamma_i)=C>0$ ($\gamma_i=\frac{\sigma_{i-1}}{\sigma_i}$)으로 고정시키는 것입니다. 이 선택은 $\gamma_i, \forall i$를 만족하고, $\{\sigma_i\}_{i=1}^L$은 기하급수를 의미합니다.

 

이론적으로는, $C\approx1$이 되는 것이어야하지만, Langevin dynamic에서 너무 많은 noise scaleing이 필요하게 됩니다. 반면에, $L=10$으로 설정하면 (이전 논문의 original setting) 너무 작고, $C=0$에 가깝습니다. 균형을 맞추기 위해서, $C \approx 0.5$를 권장하고, 요약하자면,

Technique 2. Ratio $\gamma$로 기하 급수 $\{\sigma_i\}_{i=1}^L$을 $\Phi(\sqrt{2D}(\gamma-1)+3\gamma)-\Phi(\sqrt{2D}(\gamma-1)-3\gamma)\approx0 .5$가 만족하도록  설정하는 것이 좋다.

 

3. Incorporating the noise information

High resolution의 이미지를 생성해내기 위해서는, 필연적으로 Technique1, 2에 의해서 $\sigma_1$이 커질 수 밖에 없습니다. NCSN은 노이즈 척도를 input으로 받고, 그에 대응하는 score를 출력하는 single network입니다. NCSN에서는 normalization layer에서 별도의 scale과 bias를 이용해 각각의 noise scale에 대한 정보를 통합합니다. 하지만 이는 $L$에 의해 메모리가 linear하게 증가하게 되고, normalization layer가 없는 경우에는 해당되지 않습니다.

 

그래서 저자들은 좀 더 쉽게 구현하고 넓게 사용될 수 있도록 효과적인 대안을 제시합니다. Section 3.2에서 언급된 내용으로, $p_{\sigma}(x)=N(x \vert 0, \sigma^2I)$은 $\mathbb{E}[\Vert \nabla_x \log p_{\sigma}(x) \Vert_2]\approx \frac{\sqrt{D}}{\sigma}$ 입니다. 또한, NCSN에서 경험적으로 $\Vert s_{\theta}(x, \sigma)\Vert_2 \propto \frac{1}{\sigma}$ 이라고 합니다. 이는 score network의 norm이 $\sigma$에 반비례하기 때문에, score network의 output을 $\frac{1}{\sigma}$로 조정해서 노이즈 정보를 통합할 수 있다고 합니다. 따라서 저자들은 다음과 같이 추천하고 있습니다.

 

Technique 3. (Noise Conditioning) $S_{\theta}(x, \sigma)=\frac{S_{\theta}(x)}{\sigma}$이고 $S_{\theta}(x)$는 unconditional score network입니다.

 

4. Configuring annealed Langevin dynamics

NCSN에서는 Langevin dynanimc의 noise scheduling에서 $\epsilon=2 \times 10^{-5}$ 와 $T=100$을 제안하고 있습니다. 하지만, CIFAR-10 dataset을 제외하고 어떻게 이 configuration을 변경해야 하는지는 아직 불분명합니다.

 

조금만 더 이론적인 접근을 하자면, Annealed Langevin dynamic는 인접하는 두 분포 $p_{\sigma_i}$와 $p_{\sigma_{i-1}}$을 연결해주는 역할을 합니다. 우리는 $x_{t+1} = x_t + \alpha \nabla_{x} \log p_{\sigma_i}(x_{t}) + \sqrt{2\alpha}z_t$ 를 통해 $x_{t+1}$을 sampling할 수 있습니다. 따라서, $x_T$ (data)의 분포를 계산하는 것이 가능합니다.

 

Proposition 3. $\gamma=\frac{\sigma_{i-1}}{\sigma_i}$, $\alpha=\epsilon \frac{\sigma_i^2}{\sigma_L^2}$ 이고 $x_T \sim \mathcal{N}(0, s_T^2)$라고 한다면,

$$\frac{s_T^2}{\sigma_i^2}=(1-\frac{\epsilon}{\sigma_L^2})^{2T}(\gamma^2- \frac{2\epsilon}{\sigma_L^2-\sigma_L^2(1 -  \frac{\epsilon}{\sigma_L^2})^2})+\frac{2\epsilon}{\sigma_L^2-\sigma_L^2(1-{\frac{\epsilon}{\sigma_L^2}})^2}$$

 

$\{\sigma_i\}_i^L$이 Technique 2에 의해 기하 수열로 정의될 때, $\gamma$가 같기 때문에, 모든 $i$에 대해서 $\frac{s_T^2}{\sigma_i^2}$가 같다는 것을 알 수 있고, 또한 이는 데이터의 dimension에 영향을 받지 않는다는 것을 알 수 있습니다. 따라서 저자들은 $\frac{s_T^2}{\sigma_i^2}$가 1에 가까워지기를 바라면서 noise scale을 정한다고 합니다. 그러나, 이는 보통 T가 커짐에 따라 많은 computation budget이 든다고 합니다. 따라서, 저자들은 합리적으로 T를 설정한 후에 $\frac{s_T^2}{\sigma_i^2}$가 1에 가까워 지도록 $\epsilon$를 설정하는 것을 추천한다고 합니다.

 

Technique 4. ($T$와 $\epsilon$에 대한 selection) computation budget에 맞게 T를 선택한 후에 $\frac{s_T^2}{\sigma_i^2}$가 1에 가까워 지도록 $\epsilon$을 선택하는 것을 추천합니다.

5. Improving stability with moving average

GAN과 다르게, score-based generative model은 하나의 unified된 objective로 학습을 진행하고, adversarial training이 필요 없습니다. 그러나, NCSN의 loss는 학습 중에 꾸준하게 감소하지만 unstable한 이미지가 생성이 되는 것을 관찰했고, 특히 큰 resolution의 이미지에서 관측되었다고 합니다. Appendix C.3에서 보이는 것처럼 같은 checkpoint에서 공통적인 color shift가 일어나는 경향이 굉장히 강했다고 합니다. 따라서 이러한 에러를 고치기 위해서 Exponential Moving Average (EMA)를 사용했다고 합니다. $\theta_{i}$를 NCSN의 $i$-th training iteration후라고 정의하고, $\theta'$를 parameter의 copy라고 한다면, $\theta'\leftarrow m\theta'+(1-m)\theta_{i}$으로 $\theta'$을 업데이트한다고 합니다. 또한, sampling을 진행할 때, $s_{\theta_i}(x, \sigma)$를 사용하기 보단 $s_{\theta'}(x, \sigma)$를 사용했을 때, 안정적인 FID와 artifact가 지워지는 것을 볼 수 있었습니다.

Technique 5. (EMA) Sampling할 때, EMA를 적용하자.

 

Conclusion

이 논문은 효과적이고 안정적인 score-based generative model의 학습 방법과 sampling 방법에 대해서 서술한 논문입니다. 이전 논문은 굉장히 휴리스틱한 설계와 Hyperparameter들을 이용해 학습과 sampling이 진행됐지만, 이 논문은 휴리스틱한 방법들을 수식적으로 이루어진 이론들을 통해 좀 더 robust한 method로 발전시켰다는 의미가 있습니다. 또한, image의 resolution을 올려도 무리 없이 생성이 가능해지고, 이론적으로 좀 더 탄탄해지기 때문에 이는 후에 다루게 될 SDE논문의 초석이 됩니다. 

 

개인적으로는, score-based generative model을 사용하기 위해서는 저 수식이나 이론을 사용해서 modeling하지는 않을 것 같지만, score-based generative model의 작동 원리와 저자의 시선을 알 수 있는 것 같아서 생각보다 꽤 흥미있는 논문이었습니다. 단순히 Technique을 보는 것이 끝이 아니라 이러한 Technique이 어떻게 나오게 되었는지 배경을 알면 좀 더 흥미로운 논문일 듯 합니다.

 

다음으로는 Denoising Diffusion Implicit Model로 찾아오도록 하겠습니다.

 

이전글 : https://ivdevlog.tistory.com/4

Score-Based / Diffusion Model[1] - Score Network

이전글 : https://ivdevlog.tistory.com/8

Score-Based / Diffusion Model[2] - NCSN

이전글 : https://ivdevlog.tistory.com/14

Score-Based / Diffusion Model[3] - DDPM