https://arxiv.org/abs/1907.05600
Generative Modeling by Estimating Gradients of the Data Distribution
We introduce a new generative model where samples are produced via Langevin dynamics using gradients of the data distribution estimated with score matching. Because gradients can be ill-defined and hard to estimate when the data resides on low-dimensional
arxiv.org
본 논문에서는 score 기반의 새로운 방식의 생성형 모델을 제시한다. 신경망 $s_{\theta}(\mathbf{x},\sigma)$는 $\sigma$만큼의 노이즈가 추가된 데이터 $\tilde{\mathbf{x}}$의 확률 분포의 미분값 $\nabla \mathrm{log} \, p(\tilde{\mathbf{x}})$을 추정한다. 이후 학습된 $ s_{\theta}(\mathbf{x},\sigma)$와 annealed Langevin dynamics 알고리즘을 통해 데이터를 샘플링한다.
1. Introduction
확률 기반 모델과 GAN은 생성형 모델의 대표적인 두 형태이지만 두 모델은 각각의 한계점을 내재하고 있다. 확률 기반 모델은 특별한 구조의 모델을 사용해야 하고 GAN은 적대적인 방식을 사용한 모델의 훈련이 어렵다.
본 논문에서는 데이터의 확률 분포 $p(\mathbf{x})$에 대해 $\nabla \mathrm{log} \, p(\mathbf{x})$의 값을 추정하는 score function $s_\theta(\mathbf{x})$를 도입한다. 또한 Langevin dynamics에서 사용되는 방식을 통해 확률 분포 $p(\mathbf{x})$를 따르는 데이터들을 생성한다.
이러한 생성형 모델은 특별한 매개변수나 구조가 필요없어 계산이 편리하다. 특별한 훈련 방식도 필요없어 훈련 시키기도 용이하다.
2. Score-based generative modeling
모델의 구조를 수학적으로 설명하기 위해 아래의 수식들을 사용한다.
$p_{data}(\mathbf{x})$ : 초기 데이터 확률 분포
$\{\mathbf{x}_i \in \mathbb{R}^{D}\}_{i=1}^{N}$ : $p_{data}(\mathbf{x})$에서 추출한 i.i.d. 샘플
$s_{\theta} : \mathbb{R}^{D} \rightarrow \mathbb{R}^{D}$ : $\nabla \mathrm{log} \, p(\mathbf{x})$의 값을 추정하도록 훈련되는 인공 신경망
2.1 Score matching for score estimation
모델은 아래의 objective function을 최소화하도록 훈련된다.
$$ \frac{1}{2} \mathbb{E}_{p_{data}(\mathbf{x})} \left[ \Vert s_{\theta}(\mathbf{x}) - \nabla_{\mathbf{x}} \mathrm{log} p_{data}(\mathbf{x}) \Vert \right]_{2}^{2} $$
위 수식을 아래와 같이 변형한다.
$ = \frac{1}{2} \int \left\Vert s_{\theta}(\mathbf{x}) - \frac{\nabla_{\mathbf{x}} p(\mathbf{x})}{p(\mathbf{x})} \right\Vert _{2}^{2} \, p(\mathbf{x}) d\mathbf{x} $
$ = \frac{1}{2} \int \Vert s_{\theta}(\mathbf{x}) \Vert _2^2 \, p(\mathbf{x})d\mathbf{x} - \int s_{\theta}(\mathbf{x}) \nabla_{\mathbf{x}} p(\mathbf{x}) d\mathbf{x} + \frac{1}{2} \int \frac{\Vert \nabla_{\mathbf{x}} p(\mathbf{x}) \Vert _2^2}{p(\mathbf{x})} d\mathbf{x} $
$ = \mathbb{E}_{p_{data}(\mathbf{x})} \left[ \frac{1}{2} \Vert s_{\theta}(\mathbf{x}) \Vert_2^2 + \mathrm{tr}(\nabla_{\mathbf{x}} s_{\theta} (\mathbf{x})) \right] + C $
이 때 C는 인공 신경망의 매개변수 $\theta$에 영향을 받지 않는 상수이다. 위 수식에 따라 objective function을 최소화하는 문제를 아래의 수식을 최소화하는 문제로 바꿀 수 있다.
$$ \mathbb{E}_{p_{data}(\mathbf{x})} \left[ \mathrm{tr}(\nabla_{\mathbf{x}} s_\theta(\mathbf{x})) + \frac{1}{2}\Vert s_{\theta}(\mathbf{x}) \Vert_{2}^{2} \right] $$
하지만 위 수식의 부분 중 $\mathrm{tr}(\nabla _{\mathbf{x}} s_{\theta}(\mathbf{x}))$는 고차원의 신경망에서 계산하기 어렵다. 이를 해결하기 위한 대표적인 두가지 방법 DSM, SSM을 아래에서 소개한다(이후 논문에서는 DSM 방법만 채택된다).
2.1.1 Denoising score matching
DSM은 초기 데이터 $\mathbf{x}$에 노이즈를 추가함으로써 $\mathrm{tr}(\nabla _{\mathbf{x}} s_{\theta}(\mathbf{x}))$의 계산을 우회한다. 인공신경망이 $\nabla \mathrm{log} p(\mathbf{x})$를 추정하는 대신 $\nabla \mathrm{log} q_{\sigma} (\tilde{\mathbf{x}} | \mathbf{x})$를 추정하도록 한다.
$q_{\sigma} (\tilde{\mathbf{x}} | \mathbf{x})$ : 노이즈가 추가된 데이터의 조건부 확률분포
$ q_{\sigma}(\tilde{\mathbf{x}}) \triangleq \int q_{\sigma}(\tilde{\mathbf{x}} | \mathbf{x}) p_{data}(\mathbf{x}) d\mathbf{x} $
이 때 objective function은 아래와 같이 정의된다.
$$ \frac{1}{2} \mathbb{E}_{p_{data}(\mathbf{x})} \, \mathbb{E}_{q_{\sigma}(\tilde{\mathbf{x}} | \mathbf{x})} \left[ \Vert s_{\theta}(\tilde{\mathbf{x}}) - \nabla _{\tilde{\mathbf{x}}} \mathrm{log} q_{\sigma} (\tilde{\mathbf{x}} | \mathbf{x}) \Vert_2^2 \right] $$
실제 알고리즘에서 $q_{\sigma} (\tilde{\mathbf{x}} | \mathbf{x})$는 정규분포로 설정하기 때문에 위 objective function은 쉽게 계산 가능하다.
2.1.2 Sliced score matching
$ \mathrm{tr} ( \nabla _{\mathbf{x}} s_{\theta}(\mathbf{x}) ) = \mathbb{E}_{p_{\mathbf{v}}} \left[ \mathbf{v}^{\mathbf{T}} \nabla_{\mathbf{x}} s_{\theta}(\mathbf{x}) \mathbf{v} \right] $임을 이용하여 forward mode auto-differentiation으로 효율적으로 이를 계산한다. DSM과 다르게 노이즈 없는 실제 데이터를 추정하지만 DSM 대비 4배 많은 계산이 필요하다.
2.2 Sampling with Langevin dynamics
Langevin dynamics에서 영감을 받아, 아래와 같은 방식으로 샘플링한다.
$$ \tilde{\mathbf{x}}_t = \tilde{\mathbf{x}}_{t-1} + \frac{\epsilon}{2} \nabla_{\mathbf{x}} \mathrm{log} \, p(\tilde{\mathbf{x}}_{t-1}) + \sqrt{\epsilon} \, \mathbf{z}_t $$
고정된 step size $\epsilon$ 값을 사용하며 $\mathbf{z}_t$는 표준정규분포이다. 자명한 분포 $\pi (\mathbf{x})$로부터 초기 값 $\tilde{\mathbf{x}}_0$을 샘플링한 후 데이터는 확률 밀도가 높아지는 방향으로 step size만큼 이동한다.
위 수식에서 볼 수 있듯이 초기 데이터 확률 분포 $p_{data}(\mathbf{x})$에 대한 추정 없이 $\nabla _{\mathbf{x}} \mathrm{log} p_{data}(\mathbf{x})$를 추정하는 인공신경망 $s_{\theta}(\mathbf{x})$만 있으면 데이터를 샘플링할 수 있다. 이것이 score-based generative modeling의 핵심 내용이다.
3. Challenges of score-based generative modeling
해당 섹션에서는 위에서 제시한 방법을 그대로 적용할 때 생기는 두 가지 주요한 문제에 대해 다룬다.
미리 말하자면, 두 문제는 모두 데이터에 노이즈를 더해줌으로써 해결할 수 있다.
3.1 The manifold hypothesis
manifold hypothesis는 실제 상황에서 고차원의 데이터가 특정 저차원 다양체에 속하는 경향이 있다는 가설이다.
쉽게 말해, 고차원의 데이터가 저차원의 데이터들로 표현가능한 경우를 의미한다. 3차원 공간의 원을 생각해보면 된다.
위의 수식에서 데이터가 고차원(ambient space)에 존재한다고 가정했지만 만약 데이터가 저차원 다양체에 기반한다면 문제가 생길 수 있다.
위 그래프에서 manifold hypothesis가 학습에 부정적인 영향을 끼침을 확인할 수 있다.
3.2 Low data density regions
데이터의 분포가 희미한 구간에서는 함수 추정이 정확하지 않은 문제와 샘플링이 정확하지 않은 문제가 발생할 수 있다.
3.2.1 Inaccurate score estimation with score matching
데이터 밀도가 낮은 지역에서는 샘플링하는 데이터 개수의 부족으로 값이 부정확하게 추론되는 문제가 발생한다. 목적함수 $ \frac{1}{2} \mathbb{E}_{p_{data}} \left[ \Vert s_{\theta}(\mathbf{x}) - \nabla_{\mathbf{x}} \mathrm{log} p_{data}(\mathbf{x}) \Vert_2^2 \right] $는 확률 분포 $p_{data}(\mathbf{x})$로부터 i.i.d. 샘플들을 샘플링하여 계산된다. 확률 분포 $p_{data}(\mathbf{x})$에서 밀도가 낮은 구간에서는 데이터가 거의 샘플링되지 않을 것이고, 따라서 인공신경망이 학습하는 과정에서 해당 구간의 데이터는 거의 영향을 미치지 못한다.
위 예시에서는 $p_{data} = \frac{1}{5}\mathcal{N}((-5,-5),I) + \frac{4}{5}\mathcal{N}((5,5),I) $에 대해 학습을 진행했다. 데이터의 밀도가 높은 빨간 사각형 내부에서는 값이 올바르게 추정되었지만, 이로부터 멀리 떨어진 중앙 부근에서는 학습이 제대로 이루어지지 않은 것을 확인할 수 있다.
3.2.2 Slow mixing of Langevin dynamics
위 예시와 같이 데이터가 집중적으로 분포되어 있는 서로 다른 두개의 구간이 있는 경우를 생각해보자. 이러한 경우에 Langevin dynamics를 이용한 데이터 샘플링은 제대로 작동하지 않는다.
두 정규분포를 합친 $p_{data}(\mathbf{x}) = \pi p_1(\mathbf{x}) + (1-\pi)p_2(\mathbf{x}) (0<\pi<1)$를 생각해보자. 해당 확률 분포는 아래와 같은 gradient 값을 가진다.
$$\nabla_{\mathbf{x}}\mathrm{log}p_{data}(\mathbf{x})=\begin{cases} \nabla_{\mathbf{x}}\mathrm{log}p_1(\mathbf{x}) &\text{if } \mathbf{x} \in supp(p_1(\mathbf{x})) \\ \nabla_{\mathbf{x}}\mathrm{log}p_2(\mathbf{x}) & \text{if } \mathbf{x} \in supp(p_2(\mathbf{x})) \end{cases} $$
여기서 $supp$은 지원집합을 의미한다. $supp(p(\mathbf{x}))$는 $\mathbf{x}$의 확률밀도가 너무 작지 않은 구간들의 집합이다.
위 수식에서 확인할 수 있듯이 $ \nabla_{\mathbf{x}}\mathrm{log}p_{data}(\mathbf{x}) $는 $\pi$ 값과 무관하다. section 2.2에서 소개한 샘플링 방식에서는 학습된 $\nabla_{\mathbf{x}}\mathrm{log}p_{data}(\mathbf{x}) $ 값을 사용하기 때문에 $\pi$ 값과 관련없이 샘플링 하게 된다.
4. Noise Conditional Score Networks: learning and inference
section 3에서 언급한 문제들은 데이터에 정규분포 노이즈를 더해줌으로써 해결할 수 있다. 저차원 다양체에 속하는 고차원 데이터에 노이즈를 더함으로써 manifold hypothesis를 해결할 수 있으며 이는 확률 분포의 밀도가 낮은 구간을 채우는 역할 또한 한다.
최종적으로 모델에 아래와 같은 특성을 추가한다.
- 데이터에 다양한 정도의 노이즈를 더해준다.
- 모든 노이즈 레벨에 상응하여 점수를 반환하는 하나의 신경망을 학습시킨다.
- 학습된 점수 모델을 이용해 annealed langevin dynamics 방식으로 데이터를 샘플링한다. section 2.2에서 설명한 방식에 더불어 초기에 큰 스케일의 노이즈를 사용하고 점차 노이즈의 크기는 줄어들어 최종적으로 노이즈를 더한 데이터와 그렇지 않은 데이터가 거의 같아지도록 설계한다.
4.1 Noise Conditional Score Networks
아래에서 정의된 $s_{\theta}(\mathbf{x},\sigma)$를 Noise Conditional Score Network(NCSN)라 부른다.
$$ q_{\sigma}(\mathbf{x}) \triangleq \int p_{data}(\mathbf{t})\mathcal{N}(\mathbf{x} | \mathbf{t}, \sigma^2I)d\mathbf{t} $$
$$ \forall \sigma \in \{ \sigma_i \}_{i=1}^{L} : s_{\theta}(\mathbf{x},\sigma) \approx \nabla_{\mathbf{x}} \mathrm{log}q_{\sigma}(\mathbf{x}) $$
노이즈 $\{ \sigma_i \}_{i=1}^{L}$은 $\frac{\sigma_1}{\sigma_2} = \ldots = \frac{\sigma_{L-1}}{\sigma_{L}} > 1 $를 만족시키는 양의 수열로 구성된다. $\sigma_1$은 section 3에서 언급된 문제들을 해결할만큼 적당히 큰 값으로 설정하고 $\sigma_L$은 원본 데이터에 끼치는 영향을 최소화하도록 최대한 작은 값으로 설정한다.
4.2 Learning NCSNs via score matching
DSM을 이용하여 loss 값을 계산한다.
$$ \ell(\theta;\sigma) \triangleq \frac{1}{2}\mathbb{E}_{p_{data}(\mathbf{x})}\mathbb{E}_{\tilde{\mathbf{x}} \sim \mathcal{N}(\mathbf{x},\sigma^2I)} \left[ \Vert s_{\theta}(\tilde{\mathbf{x}},\sigma) - \nabla_{\tilde{\mathbf{x}}} \mathrm{log}q_{\sigma} (\tilde{\mathbf{x}} | \mathbf{x}) \Vert_2^2 \right] $$
$$ = \frac{1}{2}\mathbb{E}_{p_{data}(\mathbf{x})}\mathbb{E}_{\tilde{\mathbf{x}} \sim \mathcal{N}(\mathbf{x},\sigma^2I)} \left[ \Vert s_{\theta}(\tilde{\mathbf{x}},\sigma) + \frac{\tilde{\mathbf{x}} - \mathbf{x}}{\sigma^2} \Vert_2^2 \right] $$
모든 노이즈에 대해서 식을 확장한다.
$$ \mathcal{L}(\theta;\{ \sigma_i \}_{i=1}^{L}) \triangleq \frac{1}{L} \sum\limits_{i=1}^{L} \lambda(\sigma_i) \ell(\theta;\sigma_i) $$
$\lambda(\sigma_i)$는 $\sigma_i$ 값에 따라 달라지는 함수값의 스케일을 조정하는 양의 계수이다. $\Vert s_{theta}(\mathbf{x},\sigma) \Vert_2 \propto \frac{1}{\sigma}$임을 실험적으로 확인했고, 이에 따라 $\lambda(\sigma)=\sigma^2$로 정한다. 이로써 $\lambda(\sigma)\ell(\theta;\sigma)$의 값의 스케일이 $\sigma$에 따라 달라지지 않게 만들었다.
4.3 NCSN inference via annealed Langevin dynamics
NCSN $s_{\theta}(\mathbf{x},\sigma)$의 훈련이 끝난 후 아래의 annealed Langevin dynamics 알고리즘에 따라 샘플링을 수행한다.
먼저 고정된 초기 확률 분포로부터 $\tilde{\mathbf{x}}_0$를 생성한다. 그 다음 step size $\alpha_1$을 이용하여 점화식을 $T$회 반복하여 확률 분포 $q_{\sigma_1}(\mathbf{x})$로부터 $\tilde{\mathbf{x}}_T$을 생성한다. 추출한 샘플 $\tilde{\mathbf{x}}_T$은 다음 step size $\alpha_2$를 이용한 계산에서 초기 데이터 $\tilde{\mathbf{x}}_0$으로 대입된다. 이후 이 과정을 반복하여 $q_{\sigma_L}(\mathbf{x})$로부터 $\tilde{\mathbf{x}}_T$을 생성한다. 이 때 충분히 작은 노이즈 $\sigma_L$에 대하여 확률 분포 $q_{\sigma_L}(\mathbf{x})$은 확률 분포 $p_{data}(\mathbf{x})$에 근사한다.
step size $\alpha_i$를 설정하는 방법에는 여러가지가 있을 수 있는데 본 논문에서는 알고리즘에서 signal-to-noise 비율인 $\frac{\alpha_i s_{\theta}(\mathbf{x},\sigma_i)}{2\sqrt{\alpha_i}\mathbf{z}}$가 일정한 스케일을 가지도록 설정한다.
$$\left\Vert \frac{\alpha_i s_{\theta}(\mathbf{x},\sigma_i)}{2\sqrt{\alpha_i}\mathbf{z}} \right\Vert_2^2 \approx \frac{\alpha_i \Vert s_{\theta}(\mathbf{x},\sigma_i) \Vert_2^2}{4} \propto \alpha_i \Vert s_{\theta}(\mathbf{x},\sigma_i) \Vert_2^2 \propto \frac{\alpha_i}{\sigma_i^2} $$
위 수식에서 $\Vert s_\theta(\mathbf{x},\sigma_i) \Vert_2^2 \propto \frac{1}{\sigma_i}$임을 활용하여 signal-to-noise의 스케일을 계산했다. 이에 따라 위 알고리즘에서와 같이 $\alpha_i \propto \sigma_i^2 $가 되도록 step size를 설정했다.
5. Experiments
본 섹션에서 NCSN 모델이 다양한 고품질 이미지 샘플을 만들어낼 수 있음을 실험적으로 보인다.
5.1 Setup
데이터셋
- MNIST, CelebA, CIFAR-10
특이사항
- CelebA 이미지들은 $140 \times 140$으로 center-crop한 후 $32 \times 32$로 resize했다.
- 모든 이미지들은 $[0,1]$의 픽셀 값을 가지도록 설정되었다.
Hyperparameter
- $L=10$, $\sigma_1=1$, $\sigma_10=0.01$
- $T=100$, $\epsilon=2\times10^{-5}$
- 초기 샘플 $\tilde{\mathbf{x}_0}$는 uniform noise에서 샘플링한다.
5.2 Image generation
위 이미지 샘플에서 볼 수 있듯이, 확률 기반 모델이나 GAN과 같은 생성형 모델과 비교했을 때 좋은 성능을 보인다.
위 그림에서 단순한 노이즈가 어떻게 정교한 이미지 데이터로 변화했는지 확인할 수 있다.
위 표에서 볼 수 있듯이 NCSN은 준수한 Inception score와 FID score를 얻었다. 특히 Inception score는 당시 최신 conditional model들과 비교해도 뒤지지 않는 성능을 보여주었다.
conditional model은 데이터의 라벨(MNIST의 경우 숫자, CIFAR-10의 경우 고양이, 개, 배 등의 대상)이 주어지고 해당 라벨의 이미지를 생성해내는 모델로 무엇을 생성해야 하는지가 명확하게 제시되기 때문에 일반적으로 성능이 더 높게 측정된다.
5.3 Image inpainting
위 그림에서 볼 수 있듯이 NCSN은 좋은 inpainting 성능을 보여준다. 이를 위해선 약간의 수정된 알고리즘만 적용하면 된다. 알고리즘은 아래와 같다.
기존의 annealed Langevin dynamics 알고리즘과 비교하면 두가지 차이점이 있다.
1. 마스크 $m$을 통해 노이즈로부터 이미지가 생성되는 부분과 이미 이미지가 있는 부분이 구분되어 알고리즘이 진행된다. 훼손되지 않은 구간은 마스크를 통해 기존의 데이터에 약간의 노이즈만 추가되어 표현된다.
2. $i$에 대해 반복적으로 알고리즘을 실행하는 동안 데이터의 훼손되지 않은 부분에도 step size만큼의 노이즈가 추가되어 NCSN 모델에 들어간다.
이렇게 간단한 알고리즘 수정을 통해 준수한 성능의 inpainting을 쉽게 구현할 수 있다.
6. Related work
- NCSN은 기존의 Markov chain 기반 학습 방식(GSN, NET, Infusion Training)과 달리, 학습 시 Markov chain 기반 샘플링이 필요 없어 효율적이고, 학습과 샘플링 과정이 완전히 분리되어 있다.
- 기존의 log-likelihhod을 사용하는 방식은 이미지가 흐릿해지는 문제가 있었는데 NCSN은 점수 기반 모델로 안정적으로 고품질 이미지를 생성한다.
- Annealing과 Denoising Score Matching을 결합하여 대규모 학습을 진행할 수 있고, 계산 부담이 덜하다.
7. Conclusion
본 논문에서는 데이터 확률 분포의 미분값을 추정하는 점수 기반의 생성형 모델의 구조를 제시하고 Langevin dynamics를 통해 샘플을 생성하는 방식을 제시하였다. 또한 이 방식을 그대로 사용할 때 나타날 수 있는 문제들을 지적하며 이를 해결하기 위해 Noise Conditional Score Networks와 annealed Langevin dynamics를 제시하였다. 해당 모델은 학습에 어려움이 없으며 학습 과정 중 샘플링이 필요없다. 뿐만 아니라 당시 최신 생성형 모델들과 비교하여도 뒤지지 않는 성능을 보여주었다.
