앞 글에서 score-based model에 대한 계산들을 살펴보았다. 본 글에서는 improved techniques for training score-based generative models 논문에 대한 계산과 결과들을 살펴본다.
Score-based model은 단순히 noise를 추가하는 것이 아니라 다음과 같은 iterative Langevin dynamics의 inference 과정을 따른다.
$$ x_t\leftarrow x_{t-1} + \alpha\nabla_x\log p(x_{t-1}) + \sqrt{2\alpha}z_t $$
for $1\leq t\leq T$. 이 때 $z_t$는 표준 다변량 정규분포를 따르는 noise term이다. 하지만 앞선 글에서 살펴보았듯 $\nabla_x\log p(x_{t-1})$을 알기는 불가능하므로 우리는 학습된 score model $s_{\theta}(x)$를 사용하여 이를 근사한다. 즉 실제로 사용하게 되는 Langevin dynamics는 다음과 같은 형태이다:
$$ x_t\leftarrow x_{t-1} + \alpha s_{\theta}(x) + \sqrt{2\alpha}z_t $$
이 과정은 $\log p(x)$라는 log-density를 찾아가는 noisy gradient ascent로 해석될 수도 있다.
데이터셋 $\{x^{(1)}, x^{(2)}, \cdots,x^{(N)}\}$이 있다고 하고 이들은 $p_{data}(x)$라는 분포에서 iid로 샘플링되었다고 가정하자. N이 충분히 크고, 우리가 다음의 근사를 할 수 있다고 하자:
$$ p_{data}(x)\approx\hat{p}{data}(x):=\frac{1}{N}\sum{i=1}^N\delta(x=x^{(i)}) $$
이 때 $\delta$는 Dirac measure (다른 말로는 point mass distribution)이다. 이 데이터가 $\mathcal{N}(0,\sigma_1^2I)$로 pertubation되었다면 이 경험적 분포인 $\hat{p}_{data}$를 다음의 정의처럼 수정하는 것이 가능하다:
$$ \hat{p}{\sigma_1}:=\frac{1}{N}\sum{i=1}^Np^{(i)}(x) $$
이 때 $p^{(i)}(x):=\mathcal{N}(x|x^{(i)},\sigma_1^2I)$로 정의한다. 즉, 각 데이터들에 perturbation을 거는 것을 의미한다. 결국 우리가 사용할 것은 $\hat{p}{\sigma_1}$이므로 inference시의 Langevin dynamics는 score function $\nabla_x\log\hat{p}{\sigma_1}(x)$에 의존한다.
위의 정의들에 더불어, $r^{(i)}(x):=\frac{p^{(i)}(x)}{\sum_{k=1}^Np^{(k)}(x)}$로 정의하자. 그러면 자연스레 score function은
$$ \nabla_x\log\hat{p}{\sigma_1}(x) = \sum{i=1}^Nr^{(i)}(x)\nabla_x\log p^{(i)}(x) $$
이 된다. 더 나아가서,
$$ \begin{align}\mathbb{E}{p^{(i)}(x)}[r^{(j)}(x)]&=\int\frac{p^{(i)}(x)p^{(j)}(x)}{\sum{k=1}^Np^{(k)}(x)}dx\nonumber\\ &\leq\int\frac{p^{(i)}(x)p^{(j)}(x)}{p^{(i)}(x)+p^{(j)}(x)}dx\nonumber\\ &=\frac{1}{2}\int\frac{2}{\frac{1}{p^{(i)}(x)}+\frac{1}{p^{(j)}(x)}}\\ &\leq\frac{1}{2}\int\sqrt{p^{(i)}(x)p^{(j)}(x)}\\ &=\frac{1}{2}\frac{1}{(2\pi\sigma_1^2)^{D/2}}\int\exp\bigg(-\frac{1}{4\sigma_1^2}\bigg(\bigg\|x-x^{(i)}\bigg\|^2+\bigg\|x-x^{(j)}\bigg\|^2\bigg)\bigg)dx\nonumber\\ &=\frac{1}{2}\frac{1}{(2\pi\sigma_1^2)^{D/2}}\int\exp\bigg(-\frac{1}{4\sigma_1^2}\bigg(\bigg\|x-x^{(i)}\bigg\|^2+\bigg\|x-x^{(i)}+x^{(i)}-x^{(j)}\bigg\|^2\bigg)\bigg)dx\nonumber\\ &=\frac{1}{2}\frac{1}{(2\pi\sigma_1^2)^{D/2}}\int\exp\bigg\{-\frac{1}{2\sigma_1^2}\bigg(\bigg\|x-x^{(i)}\bigg\|^2+(x-x^{(i)})^T(x^{(i)}-x^{(j)})+\frac{\big\|x^{(i)}-x^{(j)}\big\|^2}{2}\bigg)\bigg\}\nonumber\\ &=\frac{1}{2}\frac{1}{(2\pi\sigma_1^2)^{D/2}}\int\exp\bigg(-\frac{1}{2\sigma_1^2}\bigg(\bigg\|x-x^{(i)}+\frac{x^{(i)}-x^{(j)}}{2}\bigg\|^2+\frac{\big\|x^{(i)}-x^{(j)}\big\|^2}{4}\bigg)\bigg\}dx\nonumber\\&=\frac{1}{2}\exp\bigg(-\frac{\big\|x^{(i)}-x^{(j)}\big\|}{8\sigma_1^2}\bigg)\frac{1}{(2\pi\sigma_1^2)^{D/2}}\int\exp\bigg\{-\frac{1}{2\sigma_1^2}\bigg\|x-x^{(i)}+\frac{x^{(i)}-x^{(j)}}{2}\bigg\|^2\bigg)\bigg\}dx\nonumber\\&=\frac{1}{2}\exp\bigg(-\frac{\|x^{(i)}-x^{(j)}\|^2}{8\sigma_1^2}\bigg)\nonumber\end{align} $$
이 때 $(1)$에서 $(2)$로 넘어가는 과정에서 기하-조화평균(geometric mean-harmonic mean) 부등식이 쓰였다. 따라서 이를 정리하면