독서정리 - 밑바닥부터 시작하는 딥러닝 5
Contents
Ch1. 정규분포와 MLE
확률 분포를 근사하는 가장 간단한 방법은 정규 분포에 근사시키는 것이다.
- 정규분포
- Central limit theorem : 평균이 $\mu$이고 분산이 $\sigma^2$인 임의 확률 분포에서 뽑은 $N$ 크기의 표본의 평균($\bar{x}$)의 분포는 평균이 $\mu$이고 분산이 $\frac{\sigma^2}{N}$인 정규분포에 근사한다.
- 표본합의 분포 : Central limit theorem을 변형해서 표본합의 분포도 정규분포를 따르며 그 분포는 평균이 $N\mu$이고 분산이 $N\sigma^2$임을 보일 수 있다.
- 생성 모델을 만드는 순서
- Modeling : 어떤 분포를 따를 것이다라는 것을 가정하는 과정
- Inference(parameter estimation) : 샘플 데이터에 맞도록 파라메터를 조정하는 과정
- Likelihood function : parameter 샘플 $\mathcal{D}=\lbrace e_1,e_2,\dots,e_k\rbrace$가 뽑힐 확률 밀도 함수에 log를 씌운 함수
-
Maximum Likelihood Estimation(MLE) : Likelihood가 최대가 될 때 샘플 데이터가 뽑힐 확률이 최대이다.
-
정규분포의 MLE : 정규분포는 MLE를 해석적(analytic, closed form solution)으로 구할 수가 있다
- 다변수 정규분포
단일 변수일때와 마찬가지로 MLE를 해석적으로 구할 수 있다.
\[L(D)=\log p(\mathcal{D})=\log\prod_{i=1}^{k}\mathcal{N}(e_i;\hat\Sigma,\hat\mu)=\sum_{i=1}^k{\log \mathcal{N}(e_i;\hat\Sigma,\hat\mu)}\] \[\frac{\partial}{\partial \hat\mu}L(\mathcal{D})=0,\hat\mu=\frac{1}{k}\sum_{i=1}^k{e_i}\] \[\frac{\partial}{\partial\hat\Sigma}L(\mathcal{D})=0,\hat\Sigma=\frac{1}{k}\sum_{i=1}^k{(e_i-\hat\mu)(e_i-\hat\mu)^{\top}}\]Ch2. GMM과 EM 알고리즘
GMM (Gaussian mixture model) : peak가 여러 개 있도록 개량한 확률 분포 함수
\[\phi=\lbrace\phi_1,\phi_2,\dots,\phi_k\rbrace,\ \sum_{i=1}^k{\phi_i}=1,\ p(z=n)=\phi_n\]이라고 할 때
\[p(x,z=n)=p(x\mid z=n)p(z=n)\]※ $p(a, b)$는 joint probability고 $p(a\mid b)$는 conditional probability
\[p(x)=\sum_{i=1}^k{p(x,z=i)}=\sum_{i=1}^k{p(x\mid z=i)p(z=i)}\]따라서
\[p(x;\phi,\Sigma,\mu)=\sum_{i=1}^k\phi_i\mathcal{N}(x;\Sigma_i,\mu_i)\]GMM에서 $\phi$에 의해 결정되는 확률 변수 $z$는 관측되는 변수가 아니기 때문에 latent variable(잠재 변수)이라고 한다.
MLE를 구하기 위해서 log를 씌운 다음에 편미분을 해보면 log-sum 꼴 이라서 해석적으로 해를 못 구하는 것을 알 수 있다. 그래서 GMM에서는 EM 알고리즘을 통해 parameter estimation을 한다.
- KL divergence : 두 확률 분포 사이의 차이(거리)를 구할 수 있는 방법
따라서
\[\begin{flalign}\operatorname*{arg\,min}_{\theta}\ D_{KL}(p_{*}(x)\Vert p_{\theta}(x))&=\operatorname*{arg\,min}_{\theta}\ -\frac{1}{k}\sum_{i=1}^k\log p_\theta (x_i)\\&=\operatorname*{arg\,max}_{\theta}\ \sum_{i=1}^k\log p_\theta (x_i)\\&=\operatorname*{arg\,max}_{\theta}\ p_{\theta}(\mathcal{D})\end{flalign}\]MLE 과정은 모집단 확률분포 $p_*(x)$와 모델 확률분포 $p_{\theta}(x)$의 KL divergence distance를 최소화하는 $\theta$를 구하는 것과 동일한 과정인 것을 알 수 있다.
GMM을 포함해 여러 생성 모델은 latent variable을 가진다. $p_\theta(z\mid x_i)$를 근사하는 임의 확률분포 $q_i(z)$를 가정하여 latent variable이 포함된 확률 분포의 log likelihood를 아래와 같이 표현할 수 있다. $\sum q(z)=1$이므로 우변에 곱할 수 있다.
\[\begin{flalign} p_\theta(\mathcal{D})&=\sum_{i=1}^k{\sum_{j=1}^nq_i(z_j)\log p_\theta (x_i)}\\&=\sum_{i=1}^k{\sum_{j=1}^nq_i(z_j)\log \frac{p_\theta(x_i,z_j)}{p_\theta(z_j\mid x_i)}\cdot\frac{q_i(z_j)}{q_i(z_j)}}\\&=\sum_{i=1}^k{\lbrace\sum_{j=1}^nq_i(z_j)\log \frac{p_\theta(x_i,z_j)}{q_i(z_j)}+\sum_{j=1}^nq_i(z_j)\log \frac{q_i(z_j)}{p_\theta(z_j\mid x_i)}\rbrace}\\&= \sum_{i=1}^k{\lbrace\sum_{j=1}^nq_i(z_j)\log \frac{p_\theta(x_i,z_j)}{q_i(z_j)}+D_{KL}(q_i(z)\Vert p_\theta(z\mid x_i))\rbrace}\\&= \sum_{i=1}^k{\lbrace ELBO(x_i;q_i,\theta)+D_{KL}(q_i(z)\Vert p_\theta(z\mid x_i))\rbrace}\end{flalign}\]KL divergence 항은 항상 음이 아닌 값이기 때문에 $p_\theta(\mathcal{D})$은 모든 $x$에 대한 ELBO 항의 합보다 항상 크다.
ELBO 항(Evidence Lower BOund)은 sum-log 꼴이라서 MLE를 구하기 쉽다. 다만 파라미터가 $q$, $\theta$로 2개가 되기 때문에 하나를 고정한 상태로 번갈아가며 각 파라미터에 대해서 최적화를 진행한다. $q(z)$를 최적화하는 과정은 Expectation value 단계(E-step)라고 하고 $\theta$에 대해서 최적화하는 과정은 Maximization 단계(M-step)라고 한다. E-step과 M-step을 반복하는 것을 EM 알고리즘이라고 한다.
- E-step : 거치면 ELBO가 log likelihood랑 같아진다
- M-step : log likelihood가 증가한다
최적화의 종료 시점을 판단하기 위해서 log likelihood의 평균을 사용한다. log likelihood의 평균은 점점 작아지며 어느 시점에 거의 변화가 없어지는데 임계값을 설정해서 종료 시점을 판단한다.
\[\frac{1}{k}\sum_{i=1}^k{\log p_\theta (x_i;\theta)}\]- GMM에 EM 알고리즘 적용하기
- E-step 적용
- M-step 적용
먼저 $\phi_s$에 대해서 최적화를 적용한다. $\phi$는 합이 1이 되어야 한다는 조건이 있어서 Lagrange multiplier method(라그랑주 승수법)로 구하면 된다.
\[\begin{align}f(\phi_s)={\sum_{i=1}^k\sum_{j=1}^n q_i(z_j)\log\frac{\phi_j \mathcal{N}(x_i;\mu_j,\Sigma_j)}{q_i(z_j)} }\\ g(\phi_s)=\sum_{j=1}^n{\phi_j}-1 \end{align}\] \[\begin{align} J(\phi_s)=f(\phi_s)+\beta g(\phi_s)\\ \frac{\partial J}{\partial \phi_s}=0,\ \frac{\partial J}{\partial \beta}=0\end{align}\]해를 풀면 다음과 같다.
\[\begin{align}\frac{\partial J}{\partial \beta}=\sum_{j=1}^n{\phi_j}-1=0\\ \therefore \sum_{j=1}^n{\phi_j}=1 \end{align}\] \[\begin{align}\frac{\partial J}{\partial \phi_s}=\sum_{i=1}^k{\frac{q_i(z_s)}{\phi_s}}+\beta=0\\\therefore \phi_s=-\frac{\sum_{i=1}^k{q_i(z_s)}}{\beta} \end{align}\] \[\begin{align}\sum_{i=1}^k\sum_{j=1}^n -\frac{q_i(z_j)}{\beta}=1\\ -\frac{k}{\beta}=1\\\beta=-k\\\therefore \phi_s=\frac{\sum_{i=1}^k{q_i(z_s)}}{k}\end{align}\]마찬가지로 $\mu_s$와 $\Sigma_s$에 대해서도 해를 구해줄 수 있다.
\[\begin{align}\mu_s=\frac{\sum_{i=1}^{k}{q_i(z_s)x_i}}{\sum_{i=1}^{k}q_i(z_s)} \\\\ \Sigma_s=\frac{\sum_{i=1}^{k} q_i(z_s)(x_i-\mu_s)(x_i-\mu_s)^{\top}}{\sum_{i=1}^{k} q_i(z_s)} \end{align}\]Ch3. VAE (Variational AutoEncoder)
VAE에서는 확률 변수 $z$가 연속 확률 변수이다.
\[\begin{flalign} \log p_\theta(\mathcal{D})&=\sum_{i=1}^n\lbrace\int q(z)\log \frac{p_\theta(x_i,z)}{q(z)}dz+\int q(z)\log \frac{q(z)}{p_\theta(z\mid x_i)}dz\rbrace \\&= \sum_{i=1}^n\lbrace\int q(z)\log \frac{p_\theta(x_i,z)}{q(z)}dz+D_{KL}(q(z)\Vert p_\theta(z\mid x_i))\rbrace \\&= \sum_{i=1}^n\lbrace ELBO(x_i)+D_{KL}(q(z)\Vert p_\theta(z\mid x_i))\rbrace \end{flalign}\]E-step에서
\[q(z)=p_\theta(z\mid x)=\frac{p_\theta(x,y)}{p_\theta(x)}=\frac{p_\theta(x,y)}{\int p_\theta(x, z)dz}\]$z$가 연속 확률 변수이기 때문에 $p_\theta(z\mid x)$를 계산하는 과정에 적분이 필요하다. 따라서 $p_\theta(z\mid x)$를 계산해서 $q(z)$로 사용하는 것이 어렵다. VAE에서는 $q(z)$를 평균이 $\vec{0}$이고 분산이 $I$인 정규 분포로 한정하여 $q_{\Phi}(z\mid x_i)$를 $p(z)$로 근사시킨다.
- Encoder : 관측 변수를 잠재 변수로 바꾸는 과정
* nn은 neural net을 말한다. $\Sigma$는 대각 행렬로 한정한다.
\[q_{\Phi}(z\mid x)=\mathcal{N}(z;\mu,\Sigma)\]$P(\mathcal{D})$ 값은 $q_\Phi(z\mid x)$의 형태와 무관하기 때문에 $\Phi$에 대해서 ELBO항을 최대화시켜주면 $q_\Phi(z\mid x)$와 $p_\theta(z\mid x_i)$의 거리가 최소화된다. 근사적으로 최적값을 구하는 것이기 때문에 variational approximation이라고 한다.
위 내용을 수식에 적용하면 최적화 대상이 다음과 같아진다.
\[\begin{flalign}\log p_\theta(\mathcal{D})&=\sum_{i=1}^n\int q_{\Phi}(z\mid x_i)\log \frac{p_\theta(x_i,z)}{q_{\Phi}(z\mid x_i)}dz\\&= \sum_{i=1}^n\int q_{\Phi}(z\mid x_i)\log \frac{p_\theta(x_i\mid z)p_\theta(z)}{q_{\Phi}(z\mid x_i)}dz\\&=\sum_{i=1}^n\lbrace\int q_{\Phi}(z\mid x_i)\log p_\theta(x_i\mid z)dz-\int q_{\Phi}(z\mid x_i)\log \frac{q_\Phi(z\mid x_i)}{p_\theta(z)}dz\rbrace \\&=\sum_{i=1}^n\lbrace \mathbb{E}_{q_{\Phi}(z\mid x_i)}[\log p_\theta(x_i\mid z)]-D_{DL}(q_{\Phi}(z\mid x_i)\Vert p_\theta(z))\rbrace \end{flalign}\]VAE에서는 decoder를 사용하여 $p(x)$를 모델링한다.
- Decoder : 잠재 변수를 관측 변수로 바꾸는 과정
기댓값을 구하는 과정에 Monte Carlo method를 이용하여 다음과 같이 식을 나타낼 수 있다.
\[\begin{flalign} \operatorname*{arg\,max}_{\theta,\Phi}\log p_\theta(D)&=\operatorname*{arg\,max}_{\theta,\Phi}\sum_{i=1}^n\lbrace{\log \mathcal{N}(x;\hat{x},I)-D_{DL}(q_{\Phi}(z\mid x_i)\Vert p_\theta(z))}\rbrace\\ &=\operatorname*{arg\,max}_{\theta,\Phi}\sum_{i=1}^n\lbrace\log \frac{1}{\sqrt{2\pi}}\exp(-\frac{1}{2}\vert x_i-\hat x_i\vert^2)-D_{DL}(q_{\Phi}(z\mid x_i)\Vert p_\theta(z))\rbrace\\&=\operatorname*{arg\,min}_{\theta,\Phi}\sum_{i=1}^n\lbrace\vert x_i-\hat x_i\vert^2+D_{DL}(q_{\Phi}(z\mid x_i)\Vert p_\theta(z))\rbrace \end{flalign}\]한편으로
\[p_\theta(z)=\mathcal{N}(z;\vec{0},I)\] \[q_\Phi(z\mid x_i)=\mathcal{N}(z;\mu,\Sigma)\]두 정규 분포의 KL divergence는 다음과 같이 구할 수 있다.
\[q(z)=\mathcal{N}(z;\mu_1,\sigma_1^2 I),\ p(z)=\mathcal{N}(z;\mu_2,\sigma_2^2 I)\] \[D_{KL}(q\Vert p)=-\frac{1}{2}\sum_{h=1}^H(1+\log\frac{\sigma_{1,h}^2}{\sigma_{2,h}^2}-\frac{(\mu_{1,h}-\mu_{2,h})^2}{\sigma_{2,h}^2}-\frac{\sigma_{1,h}^2}{\sigma_{2,h}^2})\]* $H$는 $z$의 차원수
좀 더 일반적으로는 다음이 성립한다.
\[q(z)=\mathcal{N}(z;\mu_1,\Sigma_1),\ p(z)=\mathcal{N}(z;\mu_2,\Sigma_2)\] \[D_{KL}(\mathcal{N}(z;\mu_1,\Sigma_1)\Vert \mathcal{N}(z;\mu_2,\Sigma_2))=-\frac{1}{2}(H-\log (\frac{\det \Sigma_2}{\det \Sigma_1})-(\mu_2-\mu_1)^{\top}\Sigma_2^{-1}(\mu_2-\mu_1)-tr(\Sigma_2^{-1}\Sigma_1))\]따라서
\[\begin{flalign}\operatorname*{arg\,max}_{\theta,\Phi}\log p_\theta(D)&=\operatorname*{arg\,min}_{\theta,\Phi}\sum_{i=1}^n\lbrace\vert x_i-\hat x_i\vert^2-\sum_{h=1}^H (\log \sigma^2_{i,h}-\mu_{i,h}^2-\sigma_{i,h}^2)\rbrace\\&=Loss(x;\theta,\Phi)\end{flalign}\]신경망에서 위 함수를 loss function으로 계산하면 된다.
encoder에서 $z$를 샘플링하는 과정은 미분이 어렵기 때문에 backpropagation할 때 기울기가 흘러갈 수가 없는 문제가 생긴다.
이 문제를 해결하기 위해서 reparameterization trick을 사용한다.
\[z\sim \mathcal{N}(z;\mu,I\sigma^2)\] \[\epsilon\sim \mathcal{N}(\epsilon;\vec{0},I)\] \[z=\mu+\epsilon\odot\sigma\]* $\odot$은 hadamard product
다음과 같이 샘플링을 진행하면 기울기가 흘러갈 수 있게 된다.
참조 : MNIST 학습 VAE 구현 코드 (pytorch)
- hierarchical VAE
latent variable의 변환을 여러 번 하면 hierarchical VAE가 된다. 다음은 2계층 VAE에 대한 설명이다.
\[\begin{flalign}z_2&\sim\mathcal{N}(z_2;\vec{0},I)\\\mu_1,\sigma_1&=\operatorname*{nn}_{\phi_1}(x) \\ q_{\phi_1}(z_1\mid x)&=\mathcal{N}(z_1;\mu_1,\sigma_1^2 I)\\ \mu_2,\sigma_2&=\operatorname*{nn}_{\phi_2}(z_1) \\q_{\phi_2}(z_2\mid z_1)&= \mathcal{N}(z_2;\mu_2,\sigma_2^2 I) \\ \hat x&=\operatorname*{nn}_{\theta_1}(z_1) \\ p_{\theta_1}(x\mid z_1)&=\mathcal{N}(x;\hat x,I)\\ \hat z_1&=\operatorname*{nn}_{\theta_2}(z_2) \\ p_{\theta_2}(z_1\mid z_2)&=\mathcal{N}(z_1;\hat z_1,I)\end{flalign}\] \[\begin{flalign} \log p(x)&=\int q(z_1,z_2\mid x)\log \frac{p(x,z_1,z_2)}{q(z_1,z_2\mid x)}dz_1 dz_2 +\int q(z_1,z_2\mid x)\log \frac{q(z_1,z_2\mid x)}{p(z_1,z_2\mid x)}dz_1 dz_2\\ &=\int q(z_2\mid z_1)q(z_1\mid x)\log\frac{p(x\mid z_1)p(z_1\mid z_2)p(z_2)}{q(z_2\mid z_1)q(z_1\mid x)}dz_1 dz_2+D_{KL}(q(z_1,z_2\mid x)\Vert p(z_1,z_2\mid x)) \end{flalign}\] \[ELBO(x)=\mathbb{E}_{q(z_1\mid x)}[p(x\mid z_1)-D_{KL}(q(z_2\mid z_1)\Vert p(z_2))]-\int q(z_2\mid z_1)q(z_1\mid x)\log \frac{q(z_1\mid x)}{p(z_1\mid z_2)}dz_1 dz_2\]이때 $q(z_2\mid z_1)=q(z_2\mid z_1, x)\approx q(z_2\mid x)$ 라는 가정을 하면
\[ELBO(x)=\mathbb{E}_{q(z_1\mid x)}[\log p(x\mid z_1)-D_{KL}(q(z_2\mid z_1)\Vert p(z_2))] -\mathbb{E}_{q(z_2\mid x)}[D_{KL}(q(z_1\mid x)\Vert p(z_1\mid z_2))]\]참조 : MNIST 학습 conditional hierarchical VAE 구현 코드 (pytorch)
Ch4. Diffusion
Diffusion에서는 encoder 과정을 gaussian noise 추가로 대체하고 잠재 변수와 관측 변수의 차원을 동일하게 설정한다.
\[\begin{align}\beta=\{\beta_1,\beta_2,\cdots,\beta_T\}\\\alpha_t=1-\beta_t\end{align}\] \[q(x_{t}\mid x_{t-1})=\mathcal{N}(x_t;x_{t-1}\sqrt{\alpha_t},(1-\alpha_t)I)\]특히 $t=T$일 때 $q(x_T)$는 표준정규분포가 된다.
\[q(x_T)=\mathcal{N}(x_T;\vec{0},I)\]decoder는 VAE와 동일한 계산 과정을 거친다. decoder는 denoise하는 역할을 한다.
\[\hat{x}_{t-1}=\operatorname*{nn}_{\theta}(x_t)\] \[p(x_{t-1}\mid x_t)= \mathcal{N}(x_{t-1};\hat x_{t-1},I)\]Diffusion에서 ELBO는 다음과 같다. VAE의 ELBO와 차이점은 $z$가 $x_0$으로 바뀌었다는 점과 $x$가 $x_{1:T}=x_1,x_2,\cdots,x_T$로 바뀌었다는 점, $q_\Phi(z\mid x)$가 noise 추가로 대체되었기 때문에 $\Phi$가 제거됐다.
\[\begin{flalign}ELBO(x_0)&=\int q(x_{1:T}\mid x_0)\log\frac{p_\theta(x_{0:T})}{q(x_{1:T}\mid x_0)}dx_{1:T}\\&=\int q(x_{1:T}\mid x_0)\log\frac{p_\theta(x_{0:T-1}\mid x_T)p(x_T)}{q(x_{1:T}\mid x_0)}dx_{1:T}\\&=\int q(x_{1:T}\mid x_0)\log p_\theta(x_{0:T-1}\mid x_T)dx_{1:T}\end{flalign}\]$p(x_T)$와 $q(x_{1:T}\mid x_0)$은 $\theta$가 없기 때문에 제외시켰다.
\[p_\theta(x_t\mid x_{t-1},x_{t-2},\cdots,x_0)=p_\theta(x_t\mid x_{t-1})\] \[p_\theta(x_{0:T-1}\mid x_T)=\prod_{t=1}^T p_\theta(x_{t-1}\mid x_t)\]Markov property에 따라서
\[\begin{flalign} &= \int q(x_{1:T}\mid x_0) \log \prod_{t=1}^T p_\theta(x_{t-1}\mid x_t)dx_{1:T}\\ &=\int q(x_{1:T}\mid x_0) \sum_{t=1}^T \log p_\theta(x_{t-1}\mid x_t)dx_{1:T} \\ &= \mathbb{E}_{q(x_{1:T}\mid x_0)}[\sum_{t=1}^T \log p_\theta(x_{t-1}\mid x_t)]\end{flalign}\]단계 별로 T개, 2개, 1개의 샘플에서 loss를 유도하는 과정을 설명하겠다.
- T개 샘플에서 loss 계산
$x_0$에서 T개의 $x_t$를 샘플링한 후에 Monte Carlo method로 근사해서 loss를 구하면 된다.
\[\hat{x}_{t-1}=\operatorname*{nn}_\theta (x_{t})\] \[p_\theta(x_{t-1}\mid x_t)= \mathcal{N}(x_{t-1};\hat{x}_{t-1},I)\] \[\begin{flalign} Loss(x_0)&= -\sum_{t=1}^T \log p_\theta(x_{t-1}\mid x_t) \\ &=-\sum_{t=1}^T \log \mathcal{N}(x_{t-1};\hat{x}_{t-1},I) \\ &= \sum_{t=1}^T \Vert x_{t-1}-\hat{x}_{t-1}\Vert^2\end{flalign}\]computation cost가 커서 사용하기 어려운 방법이다.
- 2개 샘플에서 loss 계산
다음은 기댓값과 분산에 적용되는 몇 가지 정리들이다.
\[\begin{flalign} \mathbb E_{f(x)}[x]&=\int xf(x)dx\\ \mathbb E_{f(x)}[kx]&=\int kxf(x)dx\\ &=k\int xf(x)dx\\&=k\mathbb E_{f(x)}[x]\\ \mathbb E_{f(x)}[x+k]&= \int (x+k)f(x)dx\\&=\int xf(x)dx+\int kf(x)dx\\&=\mathbb E_{f(x)}[x]+k\\ \mathbb E_{f(x,y)}[x,y]&=\int (x+y)f(x,y)dxdy\\&=\int xf(x,y)dxdy+\int yf(x,y)dxdy\\&=\int xf(x)dx+\int yf(y)dy\\&=\mathbb E_{f(x)}[x]+\mathbb E_{f(y)}[y]\end{flalign}\]평균($\mu$)이 0이라면
\[\begin{flalign}Var_{f(x)}[x]&=\mathbb E_{f(x)}[x^2]=\sigma^2 \\ Var_{f(x)}[kx]&=\mathbb E_{f(x)}[k^2x^2]\\&=k^2\mathbb E_{f(x)}[x^2]\\&=k^2\sigma^2 \end{flalign}\]각 정규분포를 따르는 연속확률변수에 대해서
\[\begin{flalign}a\sim \mathcal{N}(a;\mu_a, \sigma_a^2I)\\b\sim \mathcal{N}(b;\mu_b, \sigma_b^2I)\\ c=a+b\\ \Longrightarrow c\sim \mathcal{N}(a;\mu_a+\mu_b, (\sigma_a^2+\sigma_b^2)I) \end{flalign}\]gaussian noise의 특징은 두 개의 gaussian noise를 합쳐도 하나의 gaussian noise가 된다는 것이다. 따라서 $q(x_t\mid x_0)$도 정규분포를 따르며 해석적으로 구할 수 있다.
\[\begin{flalign} q(x_1\mid x_0)&=\mathcal{N}(x_1;\sqrt{\alpha_1}x_0,(1-\alpha_1)I)\\ \epsilon_1&\sim \mathcal{N}(\epsilon_1;\vec{0},I)\\x_1&=\sqrt{\alpha_1}x_0+\sqrt{(1-\alpha_1)}\epsilon_1\\q(x_2\mid x_1)&=\mathcal{N}(x_2;\sqrt{\alpha_2}x_1,(1-\alpha_2)I)\\ \epsilon_2&\sim \mathcal{N}(\epsilon_2;\vec{0},I) \\ x_2&=\sqrt{\alpha_2}x_1+\sqrt{(1-\alpha_2)}\epsilon_2\\&=\sqrt{\alpha_2}(\sqrt{\alpha_1}x_0+\sqrt{(1-\alpha_1)}\epsilon_1)+\sqrt{(1-\alpha_2)}\epsilon_2 \\ &=\sqrt{\alpha_1 \alpha_2}x_0+\sqrt{\alpha_2(1-\alpha_1)}\epsilon_1+\sqrt{(1-\alpha_2)}\epsilon_2\\ q(x_2\mid x_0)&=\mathcal{N}(x_2;\sqrt{\alpha_1 \alpha_2}x_0,\lbrace\alpha_2(1-\alpha_1)+(1-\alpha_2)\rbrace I )\\ &=\mathcal{N}(x_2;\sqrt{\alpha_1 \alpha_2}x_0,(1-\alpha_1\alpha_2)I) \end{flalign}\]$q(x_2\mid x_0)$에 대해서 $q(x_3\mid x_2), q(x_4\mid x_3), \dots$를 연쇄적으로 적용할 수 있기 때문에 다음과 같이 일반화하여 식을 작성할 수 있다.
\[\begin{flalign} q(x_t\mid x_0)&=\mathcal{N}(x_t;\sqrt{\alpha_1\alpha_2\cdots\alpha_t}x_0,(1-\alpha_1\alpha_2\cdots\alpha_t)I) \\ &= \mathcal{N}(x_t;\sqrt{\bar{\alpha}_t}x_0,(1-\bar{\alpha}_t)I)\\ (\bar{\alpha}_t&= \alpha_1\alpha_2\cdots\alpha_t) \end{flalign}\]이를 바탕으로 유도한 ELBO식을 더 변형한다.
\[\begin{flalign}ELBO(x_0)&=\mathbb{E}_{q(x_{1:T}\mid x_0)}[\sum_{t=1}^T \log p_\theta(x_{t-1}\mid x_t)]\\&=\sum_{t=1}^T \mathbb{E}_{q(x_{1:T}\mid x_0)}[\log p_\theta(x_{t-1}\mid x_t)]\\&= \sum_{t=1}^T \mathbb{E}_{q(x_{t},x_{t-1}\mid x_0)}[\log p_\theta(x_{t-1}\mid x_t)]\end{flalign}\]discrete uniform distribution $\mathcal U\lbrace 1,T\rbrace$에 대해서
\[\begin{flalign} \mathbb E_{\mathcal U(t)}[f(x)]&=\sum_{t=1}^T \mathcal U(t)f(x)=\frac{1}{T}\sum_{t=1}^T f(x)\\ \sum_{t=1}^T f(x)&=T\cdot \mathbb E_{\mathcal U(t)}[f(x)] \end{flalign}\]$u(t)$를 이용하여 합을 변형하고 Monte Carlo method로 근사하면
\[\begin{flalign}ELBO(x_0)&=T\cdot \mathbb E_{\mathcal U(t)}[\mathbb{E}_{q(x_{t},x_{t-1}\mid x_0)}[\log p_\theta(x_{t-1}\mid x_t)]]\\ &= T\cdot \log p_\theta(x_{t-1}\mid x_t)\\&=-\frac{T}{2}\Vert x_{t-1}-\hat x_{t-1}\Vert^2\end{flalign}\]최종적인 Loss 계산은 다음과 같다.
\[\begin{flalign}t\sim \mathcal U\lbrace 1,T\rbrace\\ x_{t-1}\sim q(x_{t-1}\mid x_0)\\x_t\sim q(x_t\mid x_{t-1})\\ \hat x_{t-1}=\operatorname*{nn}_{\theta}(x_t, t) \\ Loss(x_0)=\frac{T}{2}\Vert x_{t-1}-\hat x_{t-1}\Vert^2 \end{flalign}\]T개 샘플링 할 필요없이 2번만 샘플링하면 되지만 더 최적화가 가능하다.
- 1개 샘플에서 loss 계산
1개 샘플로 loss를 구하기 위해서는 $q(x_{t-1}\mid x_t,x_0)$를 이용해야한다.
\[\begin{flalign} q(x_{t-1}\mid x_t,x_0)&=\frac{q(x_t\mid x_{t-1},x_0)q(x_{t-1},x_0)}{q(x_t,x_0)}\\ &=\frac{q(x_t\mid x_{t-1},x_0)q(x_{t-1}\mid x_0)}{q(x_t\mid x_0)}\\&=\frac{q(x_t\mid x_{t-1})q(x_{t-1}\mid x_0)}{q(x_t\mid x_0)} \end{flalign}\]3개의 정규분포가 곱 또는 나누기 연산에 의해 결합한 형태이기 때문에 $q(x_{t-1}\mid x_t,x_0)$ 역시 정규분포 형태를 따른다.
\[q(x_t\mid x_{t-1})=\mathcal{N}(x_t;\sqrt{\alpha_t}x_{t-1},(1-\alpha_t)I)\] \[q(x_{t-1}\mid x_0)=\mathcal{N}(x_{t-1};\sqrt{\bar \alpha_{t-1}}x_{0},(1-\bar \alpha_{t-1})I)\] \[q(x_{t}\mid x_0)=\mathcal{N}(x_{t};\sqrt{\bar \alpha_{t}}x_{0},(1-\bar \alpha_{t})I)\]exp 항만 계산하면 평균과 분산을 구할 수 있다.
\[\exp(-\frac{(x_t-\sqrt{\alpha_t}x_{t-1})^2}{2(1-\alpha_t)}-\frac{(x_{t-1}-\sqrt{\bar \alpha_{t-1}}x_{0})^2}{2(1-\bar \alpha_{t-1})}+\frac{(x_t-\sqrt{\bar \alpha_{t}}x_{0})^2}{2(1-\bar \alpha_t)})\] \[=\exp(-x_{t-1}^2\lbrace\frac{\alpha_t}{2(1-\alpha_t)}+\frac{1}{2(1-\bar\alpha_{t-1})}\rbrace + x_{t-1}\lbrace \frac{\sqrt{\alpha_t}x_t}{1-\alpha_t}+\frac{\sqrt{\bar \alpha_{t-1}}x_0}{1-\bar \alpha_{t-1}} \rbrace+C)\]* $C$는 $x_{t-1}$과 관련이 없는 항이다.
$\mathcal{N}(x;\mu,\sigma^2 I)$의 exp 항은 다음과 같다.
\[\exp(-\frac{(x-\mu)^2}{2\sigma^2}) =\exp(-\frac{x^2-2\mu x+\mu^2}{2\sigma^2})\]따라서
\[\begin{flalign} \sigma^2&=\frac{1}{2}/\lbrace\frac{\alpha_t}{2(1-\alpha_t)}+\frac{1}{2(1-\bar\alpha_{t-1})}\rbrace=\frac{(1-\alpha_t)(1-\bar\alpha_{t-1})}{1-\bar \alpha_{t}}\\ &=\sigma^2_q(t) \end{flalign}\] \[\begin{flalign}\mu&=\frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_t+\sqrt{\bar \alpha_{t-1}}(1-\alpha_t)x_0}{1-\bar \alpha_t}\\ &=\mu_q(x_t,x_0) \end{flalign}\]엄밀하게는 함수에 정규화 상수를 곱해야하지만 상수항은 무시해도 되므로 정규분포로 취급한다.
\[q(x_{t-1};x_t,x_0)=\mathcal{N}(x_{t-1};\mu_q(x_t,x_0),\sigma_q^2(t)I)\]다음은 loss를 유도할 차례이다. $J(x_0)$은 ELBO를 변형한 최적화 대상이다.
\[\begin{flalign} &t\sim \mathcal{U}\lbrace1,T\rbrace \\J(x_0)&=\mathbb{E}_{q(x_t,x_{t-1}\mid x_0)}[\log p_\theta(x_{t-1}\mid x_t)]-\mathbb{E}_{q(x_t,x_{t-1}\mid x_0)}[\log q(x_{t-1}\mid x_t,x_0)]\\&=T\cdot \mathbb{E}_{q(x_t,x_{t-1}\mid x_0)}[\log \frac{p_\theta (x_{t-1}\mid x_t)}{q(x_{t-1}\mid x_t, x_0)}] \\ &= T\cdot \int q(x_t,x_{t-1}\mid x_0)\log \frac{p_\theta (x_{t-1}\mid x_t)}{q(x_{t-1}\mid x_t, x_0)}dx_{t-1}dx_t \\ &= -T\cdot \int q(x_t\mid x_0)q(x_{t-1}\mid x_t,x_0)\log \frac{q(x_{t-1}\mid x_t, x_0)}{p_\theta (x_{t-1}\mid x_t)}dx_{t-1}dx_t\\&=-T\cdot \mathbb{E}_{q(x_t\mid x_0)}[D_{KL}(q(x_{t-1}\mid x_t, x_0)\Vert p_\theta (x_{t-1}\mid x_t))]\\&= -T\cdot \mathbb{E}_{q(x_t\mid x_0)}[D_{KL}(\mathcal{N}(x_{t-1};\mu_q(x_t,x_0),\sigma_q^2(t)I)\Vert \mathcal{N}(x_{t-1};\hat x_{t-1},\sigma_q^2(t)I))]\\&=-T\cdot D_{KL}(\mathcal{N}(x_{t-1};\mu_q(x_t,x_0),\sigma_q^2(t)I)\Vert \mathcal{N}(x_{t-1};\hat x_{t-1},\sigma_q^2(t)I))\\&=-\frac{T}{2}\cdot \sum_{h=1}^H(\frac{\Vert \mu_q(x_t,x_0)_h-\hat x_{t-1,h}\Vert^2}{\sigma_{q}^2(t)})\end{flalign}\]$\sigma^2_q(t)$는 스칼라라서 sum 밖으로 뺄 수 있다. 분자의 개별 원소에 대해서 제곱 후 합한 것은 L2 norm의 제곱과 동일하기 때문에 최종 loss는 다음과 같다.
\[\begin{flalign}t&\sim \mathcal{U}(1, T)\\x_t&\sim q(x_t\mid x_0)\\ \hat x_{t-1}&=\operatorname*{nn}_\theta(x_t, t) \\ Loss(x_0)&=\frac{\Vert \mu_q(x_t,x_0)-\hat x_{t-1}\Vert^2}{\sigma_{q}^2(t)} \end{flalign}\]neural net이 꼭 $\hat x_{t-1}$을 예측할 필요는 없다. 다른 2가지 방식을 소개한다.
- neural net이 $x_0$를 예측하게 하는 방식
$\mu_q(x_t,x_0)$의 구조와 유사하게 $\hat x_{t-1}$를 구성한다. neural net은 $x_0$를 예측한다.
\[\hat x_{t-1}=\frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_t+\sqrt{\bar \alpha_{t-1}}(1-\alpha_t)\cdot \operatorname*{nn}\limits_\theta(x_t,t)}{1-\bar \alpha_t}\] \[\begin{flalign} Loss(x_0)&=\frac{\bar \alpha_{t-1}(1-\alpha_t)^2}{\sigma^2_q(t)(1-\bar \alpha_t)^2}\Vert x_0-\operatorname*{nn}\limits_\theta(x_t,t)\Vert ^2 \\ &= \Vert x_0-\operatorname*{nn}\limits_\theta(x_t,t)\Vert ^2\end{flalign}\]- neural net이 $q(x_t\mid x_0)$의 noise를 예측하게 하는 방식
$x_0$를 예측하는 방식과 유사한 구조에서 $x_0$가 아닌 $q(x_t\mid x_0)$의 noise를 예측한다.
\[\begin{flalign}q(x_t\mid x_0)&=\mathcal{N}(x_t;\sqrt{\bar\alpha_{t}}x_0,(1-\bar\alpha_{t})I)\\ \epsilon&=\mathcal{N}(\epsilon;\vec{0},I)\\ x_t&=\sqrt{\bar\alpha_{t}}x_0+\sqrt{1-\bar\alpha_{t}}\epsilon \\x_0&=\frac{x_t-\sqrt{1-\bar\alpha_{t}}\epsilon}{\sqrt{\bar\alpha_{t}}} \end{flalign}\] \[\hat x_0=\frac{x_t-\sqrt{1-\bar\alpha_{t}}\operatorname*{nn}\limits_\theta(x_t,t)}{\sqrt{\bar\alpha_{t}}}\] \[\begin{flalign}\hat x_{t-1}&=\frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_t+\sqrt{\bar \alpha_{t-1}}(1-\alpha_t)\cdot \frac{x_t-\sqrt{1-\bar\alpha_{t}}\operatorname*{nn}\limits_\theta(x_t,t)}{\sqrt{\bar\alpha_{t}}}}{1-\bar \alpha_t}\\ &= \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}} \operatorname*{nn}_\theta(x_t,t) \right) \end{flalign}\] \[\begin{flalign}\mu_q(x_t,x_0)&=\frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_t+\sqrt{\bar \alpha_{t-1}}(1-\alpha_t)x_0}{1-\bar \alpha_t}\\ &=\frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_t+\sqrt{\bar \alpha_{t-1}}(1-\alpha_t)\cdot \frac{x_t-\sqrt{1-\bar\alpha_{t}}\epsilon}{\sqrt{\bar\alpha_{t}}}}{1-\bar \alpha_t}\\&=\frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}} \epsilon\right)\end{flalign}\]기존의 loss function에 대입하여 식을 정리하면 다음과 같이 loss가 나온다.
\[Loss(x_0)=\Vert \operatorname*{nn}\limits_\theta(x_t,t)-\epsilon\Vert^2\]논문에서는 noise를 예측하는 방식의 성능이 더 좋았다고 한다.
Diffusion에서 새로운 데이터를 샘플링하는 과정은 다음과 같다. noise를 예측하는 구조를 기준으로 설명하지만 다른 경우도 중간 과정을 예측하기 때문에 동일한 과정을 따른다. 아래 과정을 T부터 연쇄적으로 적용하면 된다.
\[\begin{flalign} \hat \epsilon &= \operatorname*{nn}\limits_\theta(x_t,t)\\ \hat x_{t-1}&=\frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_t+\sqrt{\bar \alpha_{t-1}}(1-\alpha_t)\cdot \frac{x_t-\sqrt{1-\bar\alpha_{t}}\operatorname*{nn}\limits_\theta(x_t,t)}{\sqrt{\bar\alpha_{t}}}}{1-\bar \alpha_t}\\&=\frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}} \operatorname*{nn}_\theta(x_t,t) \right)\\ \sigma_q^2(t)&=\frac{(1-\alpha_t)(1-\bar\alpha_{t-1})}{1-\bar \alpha_{t}}\\p_\theta(x_{t-1}\mid x_t)&=\mathcal{N}(x_{t-1};\hat x_{t-1},\sigma_q^2(t)I)\\\epsilon&\sim\mathcal{N}(\epsilon;\vec{0},I)\\x_{t-1}&=\hat x_{t-1}+\epsilon\cdot \sigma^2_q(t) \end{flalign}\]마지막 $x_0$를 구하는 과정에서는 noise를 추가하지 않는 것이 더 좋은 성능을 보인다고 한다.
참조 : MNIST 학습 Diffusion 구현 코드 (pytorch)
generative model을 사용해서 어떠한 조건에서 데이터를 생성하기 위해서는 조건부 확률 분포를 모델링해야 된다.
\[\log p(x_0\mid y)=\int q(x_{1:T}\mid x)\log \frac{p(x_{0:T}\mid y)}{q(x_{1:T}\mid x)}dx_{1:T}+\int q(x_{1:T}\mid x)\log \frac{q(x_{1:T}\mid x)}{p(x_{1:T}\mid x_0,y)}dx_{1:T}\] \[\begin{flalign}ELBO(x_0\mid y) &= \int q(x_{1:T}\mid x_0)\log \frac{p(x_{0:T}\mid y)}{q(x_{1:T}\mid x_0)}dx_{1:T}\\&=\int q(x_{1:T}\mid x_0)\log \frac{p(x_{0:T-1}\mid x_T,y)p(x_T\mid y)}{q(x_{1:T}\mid x_0)}dx_{1:T}\\&=\int q(x_{1:T}\mid x_0)\log p(x_{0:T-1}\mid x_T,y)dx_{1:T}\\&=\int q(x_{1:T}\mid x_0)\log\prod_{t=1}^T p(x_{t-1}\mid x_{t},y)dx_{1:T}\\&=T\cdot \mathbb{E}_{q(x_t,x_{t-1}\mid x_0)}[\log p(x_{t-1}\mid x_t,y)]\end{flalign}\]위 방식으로 Loss function을 유도해서 모델을 학습시킬 수 있다. 하지만 이렇게 신경망에 조건만 추가하여 학습시키면 경우에 따라 조건을 무시하게 될 수 있다. 따라서 score function을 기반으로 한 guidance의 도입이 필요하다.
score를 유도하기 위해서는 Tweedie’s formula가 사용된다.
- Tweedie’s formula
$x\sim \mathcal{N}(x;\mu,\Sigma)$일 때
\[\mathbb{E}(\mu\mid x)=x+\Sigma \nabla_{x} \log p(x)\]이를 이용하여
\[\begin{flalign} x_t&\sim \mathcal{N}(x_t;\sqrt{\bar \alpha_{t}}x_0,(1-\bar \alpha_t)I)\\ \epsilon&\sim \mathcal{N}(\epsilon;\vec{0},I) \\ x_t&=\sqrt{\bar \alpha_{t}}x_0+\sqrt{1-\bar \alpha_{t}}\epsilon\\ \mathbb{E}(x_t-\sqrt{1-\bar \alpha_{t}}\epsilon\mid x_t)&=x_t+(1-\bar \alpha_t)\nabla_{x_t}\log p(x_t)\\&=\mathbb{E}(x_t\mid x_t)-\sqrt{1-\bar \alpha_{t}}\cdot\mathbb{E}(\epsilon\mid x_t)\\ \mathbb{E}(\epsilon\mid x_t)&= -\sqrt{1-\bar \alpha_{t}}\cdot \nabla_{x_t}\log p(x_t)\\&\approx \epsilon\end{flalign}\]이때 $\nabla_{x_t}\log p(x_t)$를 score라고 한다. score는 $\epsilon$과 상수배만큼 차이나기 때문에 신경망이 score를 예측하도록 할 수 있다.
\[\begin{flalign}\nabla_{x_t}\log p(x_t\mid y)&=\nabla_{x_t}\log \frac{p(y\mid x_t)p(x_t)}{p(y)}\\ &=\nabla_{x_t}\log p(y\mid x_t)+\nabla_{x_t}\log p(x_t) \end{flalign}\]따라서 $\nabla_{x_t}\log p(y\mid x_t)$만 알면 조건부 score를 구할 수 있다. $\nabla_{x_t}\log p(y\mid x_t)$를 구하는 방식은 2가지가 있다.
- classifier guidance
classification을 할 수 있는 다른 모델을 이용하는 방법이다. 정답인 조건의 softmax에 대한 $x_t$의 gradient를 구하기만 하면 되지만 classifier 모델을 따로 학습시켜야 한다는 단점이 있다. $\gamma$는 classifier의 gradient를 얼마나 반영할 것인지에 대한 하이퍼파라미터이다.
- classifier-free guidance
classifier guidance의 식을 조금 변형하면 classifier-free guidance 식을 얻을 수 있다.
\[\begin{flalign}\nabla_{x_t}\log p(x_t\mid y)&=\gamma \nabla_{x_t}\log p(y\mid x_t)+\nabla_{x_t}\log p(x_t)\\ &=\gamma\nabla_{x_t}\log \frac{p(x_t\mid y)p(y)}{p(x_t)}+\nabla_{x_t}\log p(x_t)\\&=\gamma(\nabla_{x_t}\log p(x_t\mid y)-\nabla_{x_t}\log p(x_t))+\nabla_{x_t}\log p(x_t) \end{flalign}\]이때 $\nabla_{x_t}\log p(x_t)$를 조건이 0인 $\nabla_{x_t}\log p(x_t\mid \emptyset)$ 경우로 생각하면 다음과 같이 식이 얻어진다.
\[\nabla_{x_t}\log p(x_t\mid y)= \gamma(\nabla_{x_t}\log p(x_t\mid y)-\nabla_{x_t}\log p(x_t\mid \emptyset))+\nabla_{x_t}\log p(x_t\mid \emptyset)\]classifier-free guidance에서는 조건이 없을 때와 있을 때 각각의 score를 계산해 조건부 score를 구할 수 있다.
참조 : MNIST 학습 score based classifier free guidance diffusion 구현 코드 (pytorch)