Score-Based / Diffusion Model[2] - NCSN

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

Score-Based Generative Model[1] - Score Network

 

-------------------------------- 2022.04.05 자로 많은 내용이 수정 및 추가되었습니다. --------------------------------

 

 

이전 포스팅에서 Score-Based Model이라는 것을 알아봤습니다.

 

결국, Score based 모델에서 $ \text{tr}(\nabla_x s_\theta (x)) $의 연산량을 해결하는 것이 Score-Based 모델의 핵심이라고 생각할 수 있습니다.

 

이것을 해결한다면, Score Matching을 이용해 Sampling을 해서 data를 generation할 수 있게 됩니다.

 

그래서 나온 논문이 총 2가지가 있는데, 여기서는 그 중 하나인 NCSN (Noise Conditional Score Network)를 설명하고자 합니다. 

 

Denoising Score Matching

score function을 학습하기 위해서, 결국 $ \text{tr}(\nabla_x s_\theta (x)) $의 연산량을 해결하기 위한 방법 중에 하나로, denoising score matching 입니다. $q_{\sigma}(\tilde{x})=\int{p_{data}(x)}q_{\sigma}(\tilde{x}\vert x)dx$로 정의할 때, $ \mathbb{E}_{q_{\sigma}(\tilde{x}\vert x)} \mathbb{E}_{p_{data}}[s_{\theta}(\tilde{x})-\nabla_{\tilde{x}}\log q_{\sigma}(\tilde{x} \vert x)] $ 으로 모델을 학습할 수 있습니다. 하지만 $ s_{\theta}(x)=\nabla_{x}{\log}q_{\sigma}(x)\approx\nabla_{x}{\log}_{data}(x) $ 이기 위해서는 $\sigma$가 굉장히 작아야 한다고 합니다.

 

Sliced Score Matching

다른 방법으로는 sliced score matching 입니다. $ \mathbb{E}_{p_v} \mathbb{E}_{p_{data}}[\mathbf{v}^{T} \nabla_{x}s_{\theta}(x)\mathbf{v}+\frac{1}{2}\Vert s_{\theta}(x) \Vert_2^2] $의 objective function을 갖는데, 식을 보시면 알겠지만, data의 distribution term이 들어가있지 않고, x의 score function을 학습할 수 있지만, auto-differentiation보다 4배의 연산량을 갖는다고 합니다. 아무래도 한번에 backward하는 것이 아니라서 연산이 번거롭다고 할 수 있습니다.

 

Langevin Dynamic

Langevin dynamic은 Bownian motion을 stochastic process로 나타내는 미분방정식입니다. 

 

$$ \mathbf{\tilde{x}}_t = \mathbf{\tilde{x}}_{t-1} + \frac{\epsilon_{t}}{2}\nabla_{x}\log p(\mathbf{\tilde{x}}_{t-1})+\sqrt{\epsilon_{t} }\mathbf{z}_t $$

 

다음과 같은 식으로 정의할 수 있는데, $p(x)$ 에서 오직 $\nabla_{x} \log p(x)$ 만 이용해서 sampling이 가능합니다. 그래서 이 stochastic process의 각 step마다 score를 학습시켜 process를 이용해 sampling을 하는 방식이고, 따라서, score function을 학습시켜 Langevin Dynamic을 이용해 sampling을 하는 것이 score-based generative model이라고 정의했습니다.

 

하지만 score-based generative model을 기반으로 Langevin dynamic을 이용해 sampling을 할 경우에 문제점이 생기는데, 첫 번째는 Manifold Hypothesis 이고, 두 번째는 Inaccurate low data density data 라고 논문에서 서술하고 있습니다.

 

Manifold hypothesis

첫번째로, Manifold 가설은 data가 어떤 manifold상에서의 point로 표시된다라는 가설이지만, data의 score는 정의할 수 없습니다. 그림을 보면, 

와 같이 Data Manifold가 존재하는 가정이지만, score는 vector field이기 때문에 manifold상에 존재하는지에 대한 의문이 있습니다.

즉, 주변 space에서의 gradient이기 때문에, low demensional manifold라고 한다면 존재되지 않습니다.

 

두번째로는, data distribution의 support가 whole space일 경우에만 정상적으로 score estimator로써 작동할 수 있다고 합니다. 그렇기 때문에, 저자들은 간단한 실험을 제안합니다. Unperturbated data와 perturbated data를 이용해 Sliced Score Matching을 진행했을 때, 노이즈를 추가하지 않은 unperturbated data의 loss는 Figure 1과 같이 불규칙하게 나오지만, perturbated data의 loss는 일정하게 수렴하는 것으로 나오게 됩니다. 그리고 이것은 Gaussian noise의 distribution이 support가 $R^D$로 정의되기 때문이라고 이야기 하고 있습니다.

 

Inaccurate low data density data

두번째로는, Low Data-Density에서는 score function을 학습을 시켜도, 정상적으로 score가 estimate되지 않는다는 문제점을 제시합니다. [2021.08.04 Low Dimension -> Low Data-Density 수정]

 

다음과 같이 2개의 정규분포의 가중치를 다르게 했을 때 Data Score 이지만, Sliced Score Matching을 이용해 학습해 예측한 score는 다르게 나옵니다. 이것은 실제로 수렴하는 것에 대한 이슈이며, sampling을 했을 때 정상적으로 sampling되는지 의문을 던지게 됩니다.

 

Slow mixing of Langevin dynamic

또한, Langevin Dynamic에 대해서도 문제점을 제시합니다. Langevin Dynamic을 이용해 Sampling을 했을 때, low density를 충분히 반영하지 못한다는 문제점이 있습니다. $p_{data}(x)=\pi p_{1}(x) + (1-\pi)p_{2}(x)$ 라고 정의했을 때, $p_{data}$의 score는 $\pi$에 영향을 끼치지 않기 때문에 sampling을 할 때 low density를 반영하지 못합니다. 이론적으로는 수렴하게 하기 위해선 굉장히 작은 step size ($\epsilon$)와 굉장히 큰 수의 step을 사용해야 한다고 합니다.

 

결과적으로, real data의 distribution보다는 perturbated data의 distribution을 이용하는 것이 score-based modeling에 더 적합하다는 것을 알 수 있었습니다. 첫번째로, Gaussian distribution은 whole space에서 정의되기 때문에, low demension manifold에 국한되지 않으며, 이는 manifold가설에 위배되지 않습니다. 두번째로, Large noise가 unperturbated data distribution에서 low density region을 채워줄 수 있다라고 합니다. 또한 multiple noise level을 사용하여 실제 데이터 분포로 수렴하는 sequence를 얻을 수 있다고 합니다.

 

또한, simulated annealing과 annealed importance sampling의 정신으로(라고 적혀있는데...) intermediate distribution을 활용해서 Langevin Dynamic을 향상시킬 수 있다고 합니다.

 

이러한 관점으로, 저자들은 총 두 가지를 제안합니다.

첫 번째로 다양한 레벨의 노이즈를 이용해 data를 perturbing하는 것이고,

(perturbing the data using various levels of noise)

이것은 위에 언급했듯이, unperturbation data distribution의 support부분이 잘 정의되지 않고, Gaussian noise를 섞어서 만든 perturbation data distribution은 whole space에서 정의되기 때문에 학습하기 유용하기 때문입니다. 또한, large noise의 score를 사용하고 점차적으로 noise의 level을 낮추는 것이 원래 데이터와 구분할 수 없는 낮은 noise로 전달할 때 이점이 있습니다. 

 

두 번째로 single conditional score network로 모든 noise level을 학습시키는 것입니다.
(simultaneously estimating scores corresponding to all noise levels by training a single conditional score network)

 

이 두 가지의 목적을 이루기 위해서 다음과 같은 score network, the training objective, annealing schedule for Langevin dynamic을 설명합니다.

Noise Conditional Score Network

 

Perturbation schedule을 L step으로 정의하고 gaussian noise를 이용해 만든 perturbation distribution의 score function을 학습합니다. perturbation distribution의 std는 data방향으로 갈수록($\sigma_{L}$) 작아지고 prior방향으로 갈수록 ($\sigma_{1}$) 크게 perturbation schedule을 정의합니다.

 

즉, conditional perturbation distribution $q_{\sigma_{i}}(\tilde{x} \vert x)=\mathcal{N}(\tilde{x}; x, \sigma_{i}^2\mathbb{I}) $ 으로 정의할 수 있습니다. 따라서 conditional perturbation distribution의 score matching을 하기 위해서,

$$ \theta^* = \arg\min_{\theta}\sum_{i\in[1,2,..,N]}\lambda(\sigma_{i})\mathbb{E}_{p_{data}(x)}\mathbb{E}_{\rm p_{\sigma_i}(\tilde{x}\vert x)}{{1}\over{2}}[\vert\vert \nabla_x \log p_{\sigma_i}(\tilde{x}\vert x) - s_{\theta}(x, \sigma_i) \vert\vert^{2}_{2}] $$

 

의 objective를 갖습니다. 그리고 $p_{\sigma_i}(\tilde{x}\vert x) = N(\tilde{x};x,\sigma_i^2)$ 이므로, $\nabla_x \log p_{\sigma_i}(\tilde{x}\vert x)=-\frac{z}{\sigma_i}$ 로 간단하게 계산이 가능합니다. 

 

따라서 최종적으로 전체 loss는 $ \mathcal{L}_{\sigma_{i}} (\theta) = \lambda ( \sigma_{1} ) l_{\sigma_{1}} (\theta) + \cdots + \lambda ( \sigma_{L} ) l_{\sigma_{L}} (\theta)$ 으로 정의됩니다. 그리고 $\lambda(\sigma_{i})=\sigma_{i}^2$로 정의했을 때, 경험적으로 더 잘 나온다고 하고 $\nabla_x \log p_{\sigma_i}(\tilde{x}\vert x)=-\frac{z}{\sigma_i}$ 에서 분모의 $\sigma_i$ term을 지우기 위해서 사용한다 라고합니다. 

Annealed Langevin Dynamic

다음으로는, sampling을 하는 과정에서 Low data density를 반영하지 못하는 issue에서 Langevin Dynamic의 문제점을 지적했었습니다. 그래서 Langevin Dynamic에서 step마다 annealed를 해서 sampling을 하는 Annealed Langevin dynamic을 제시합니다. 즉, Langevin Dynamic의 각 step 을 T번 반복으로 실행하는 알고리즘으로써, 그렇게 된다면 실제 data distribution을 조금 더 반영해서 sampling한다고 합니다.

 

식으로 표현하자면, 다음과 같습니다.

$$ \mathbf{x}_i^{m} = \mathbf{x}_i^{m-1} + \epsilon_{i}\mathbf{s}_{\theta^{*}}(\mathbf{x}_i^{m-1}, \sigma_i)+\sqrt{2 \epsilon_{i} }\mathbf{z}_i^{m}$$

 

실제로 i.i.d sampling을 했을 때와 Langevin dynamic, Annealed Langevin dynamic을 했을 때 실제 low data density를 반영하는 결과의 차이가 큰 것을 알 수 있습니다.

 

 

그러므로, 각 step의 perturbation distribution을 정의하고 하나의 score network로 각각 perturbation distribution의 score를 학습한 후에 Annealed Langevin dynamic을 통해 prior distribution (본 논문에서는 uniform distribution으로 정의했습니다.)에서 부터 sampling하는, 즉 Noise를 Condition으로 줘서 만들어지는 perturbated data를 이용해 score network를 학습시키는 알고리즘을 제안합니다. 그것을 Noise Conditional Score Network (NCSN) 라고 부릅니다. 또한, score모델을 score matcing으로 학습하고 Langevin dynamic 기반으로 sampling하는 framework를 Score Matching Langevin Dynamic (SMLD) 라고 합니다.

 

Sampling을 하기 위해선 Langevin dynamic을 사용하기 때문에 각 step마다의 결과를 보게되면 다음과 같습니다. 이는, prior distribution부터 점점 denoising되면서 sampling되는 것을 볼 수 있습니다.

 

GAN과 비교될 만큼 퀄리티가 좋게 나오지만, 아직은 FID Score를 비교했을 때는 살짝 떨어지긴 합니다. 하지만 NCSN 즉, Noise를 denoising하면서 만들어내는 결과라고 생각하기엔 화질이 좋게 나오고 GAN 만큼의 가능성을 충분히 보여주지 않았나 라는 생각이 듭니다. 

 

마지막으로, NCSN은 image inpainting이 굉장히 잘 작동합니다. 이는 inpainting하고자 하는 mask를 이용하는 Algorithm 2를 통해 sampling하는 기법을 제안합니다. Algorithm 1에 score network의 input인 perturbated data를 만들 때, 원본의 정보에 perturbing을 시켜 mask를 이용해 원하는 부분만 generation합니다.

 

이상 튜토리얼에 대한 코드를 공유하고 NCSN에 대한 설명을 마무리 하겠습니다. [Link]

다음 논문으로는 Denoising Diffusion Probabilistic Mode (DDPM) 에 대해 설명하도록 하겠습니다.

 

 

 

Reference

[1] Generative Modeling by Estimating Gradients of the Data Distribution

[2] Estimation of non-normalized statistical models by score matching