위 그림이 본 논문의 전부를 잘 표현하는 그림입니다. Diffusion model의 가장 큰 단점은 sampling이 느리다는 데에 있습니다. 그 이유 중 NFE 가 많이 요구된다는 점도 있지만, high dimension (data dimension) 에서 sampling이 진행된다는 점이 큰 bottleneck으로 작용하죠. 따라서 본 논문에서는 latent space $z$ 에서의 diffusion을 제안합니다.
또한, discrete data를 modeling하는 VAE를 이용하여 non-continuous data에도 이용이 가능합니다.
Forward SDE는 아래와 같이 표현됩니다.
$$ \mathrm{d} \mathbf{z}=f(t) \mathbf{z} \mathrm{d} t+g(t) \mathrm{d} \mathbf{w} $$
Diffusion 이 latent space $\mathbf{z}$ 에서 진행되기 때문에, $\mathbf{z}$가 variable입니다. 여타 diffusion 모델들과 마찬가지로, $\mathbf{z}_1 \sim \mathcal{N}(0, \mathbf{I})$ 입니다. Reverse SDE는 반대로
$$ \mathrm{d} \mathbf{z}=\left[f(t) \mathbf{z}-g(t)^{2} \nabla_{\mathbf{z}} \log q_{t}(\mathbf{z})\right] \mathrm{d} t+g(t) \mathrm{d} \overline{\mathbf{w}} $$
위와 같고, 이 때 score function인 $\nabla_{\mathbf{z}} \log q_{t}(\mathbf{z})$를 구하기 위해 DSM 을 이용해 아래와 같이 optimize 하게 됩니다.
$$ \min {\boldsymbol{\theta}} \mathbb{E}{t \sim \mathcal{U}[0,1]}\left[\lambda(t) \mathbb{E}{q\left(\mathbf{z}{0}\right)} \mathbb{E}{q\left(\mathbf{z}{t} \mid \mathbf{z}{0}\right)}\left[\left\|\nabla{\mathbf{z}{t}} \log q\left(\mathbf{z}{t} \mid \mathbf{z}{0}\right)-\nabla{\mathbf{z}{t}} \log p{\boldsymbol{\theta}}\left(\mathbf{z}{t}\right)\right\|{2}^{2}\right]\right]+C $$
위 식에서 $C =\mathbb{E}{t \sim \mathcal{U}[0,1]}\left[\lambda(t) \mathbb{E}{q\left(\mathbf{z}{0}\right)} \mathbb{E}{q\left(\mathbf{z}{t} \mid \mathbf{z}{0}\right)}\left[\left\|\nabla_{\mathbf{z}{t}} \log q\left(\mathbf{z}{t}\right)\right\|{2}^{2}-\left\|\nabla{\mathbf{z}{t}} \log q\left(\mathbf{z}{t} \mid \mathbf{z}{0}\right)\right\|{2}^{2}\right]\right]$이 constant로 표현된 이유는, $\theta$ 에 대해 independent하기 때문이며, 이 부분이 constant로 표현될 수 있다는 점 때문에 SM이 아닌 DSM 을 이용해서 score function을 training할 수 있는 것입니다. 여기서 $\lambda(t)$ 의 choice에 대한 논의가 많이 이루어졌는데, Song et al. 에서 보였듯이 $\lambda(t) = g(t)^2 / 2$ 인 경우 maximum likelihood training을 근사할 수 있습니다.
$$ \mathrm{KL}\left(q\left(\mathbf{z}{0}\right) \| p{\boldsymbol{\theta}}\left(\mathbf{z}{0}\right)\right) \leq \mathbb{E}{t \sim \mathcal{U}[0,1]}\left[\frac{g(t)^{2}}{2} \mathbb{E}{q\left(\mathbf{z}{0}\right)} \mathbb{E}{q\left(\mathbf{z}{t} \mid \mathbf{z}{0}\right)}\left[\left\|\nabla{\mathbf{z}{t}} \log q\left(\mathbf{z}{t}\right)-\nabla_{\mathbf{z}{t}} \log p{\boldsymbol{\theta}}\left(\mathbf{z}{t}\right)\right\|{2}^{2}\right]\right] $$