[21’ ICLR] Score-Based Generative Modeling through Stochastic Differential Equations
카테고리: Generation
🔍 Abstract
- SDE(Stochastic Differential Equation)를 사용하여 지금까지의 score-based generative modeling과 DDPM을 통합하는 방법을 제안한다.
- Predictor-Corrector framework, Probability Flow ODE를 제안하였으며, 이를 통해 오차를 줄이고 likelihood를 계산할 수 있게 되었으며, 샘플링 속도를 높일 수 있었다.
- 방법을 확장하여 class-conditional generation, inpainting, colorization 등의 controllable generation을 수행할 수 있었다.
1. Introduction
Diffusion 계열에서 너무 많이 인용되는 논문이라 선정하여 읽게 되었다. 본 논문은 diffusion model을 설명하는 하나의 중요한 방식으로 이해하는 것이 좋은 것 같다. 저자 중에는 CS236n Deep Generative Models 강의를 진행했던 Stefano Ermon 교수님의 이름이 있다. 사실 대부분의 내용이 CS236n을 통해 이미 배웠던 내용이라 크게 어렵지 않게 읽을 수 있었다. 관련된 내용은 아래 링크를 참고하기 바란다.
위 그림도 아주 익숙하다. Reverse-time SDE를 푸는 것은 score-based generative model과 같다. 자세한 내용은 아래에서 다루도록 하겠다.
2. Score-based Generative Modeling with SDEs
2.1. Score-Based Generative Modeling
Diffusion process, 즉 forward SDE는 Ito SDE로 모델링된다. 이때 $f$는 drift coefficient, $g$는 diffusion coefficient로 알려져 있다. $w$는 standard Brownian motion이다.
\[dx = f(x, t)dt + g(t)dw\]한편 reverse-time SDE는 다음과 같이 얻을 수 있음이 알려져 있다. 따라서 score $\log p_t(x)$를 정확히 알 수 있다면 reverse process를 통해 $p_0$의 샘플을 얻을 수 있다.
\[dx = \left[f(x,t) - g(t) ^ 2 \nabla_ x \log p_t (x) \right] dt + g(t) d \bar{w}\]그러나 score를 정확하게 아는 것은 불가능하고 따라서 이를 예측할 수 있는 score function $s_\theta$를 학습시킨다. 다양한 방법이 있으나 논문에서는 denoisng score matching을 사용한다.
\[L = \mathbb{E}_ t \left\lbrace \lambda (t) \mathbb{E}_ {x(0)} \mathbb{E}_ {x(t) \vert x(0)} \left[ \left\Vert s_\theta (x(t), t) - \nabla_ {x(t)} \log p_ {0t} (x(t) \vert x(0)) \right\Vert ^ 2 \right] \right\rbrace\]이때 $p_ {0t} (x(t) \vert x(0))$는 $x(0)$에서 $x(t)$로의 transition kernel이다. 일반적으로 $\lambda$는 다음과 같이 설정한다. 이 뜻은, $t$가 커질수록, 즉 transition kernel이 smooth해질수록 $\lambda$가 커진다는 뜻이다. 즉, step이 커 예측하기 어려울수록 더 많은 가중치로 학습한다고 이해하면 된다.
\[\lambda(t) \propto \frac{1}{\mathbb{E} \left[ \left\Vert \nabla_ {x(t)} \log p_ {0t} (x(t) \vert x(0)) \right\Vert ^ 2 \right]}\]2.2. SMLD (VE-SDE)
Background
SMLD(Denoising Score Matching with Langevin Dynamics)는 score-based model을 대표하는 방법이다. 논문의 표현을 빌리자면 다음과 같이 설명할 수 있다.
Perturbation kernel $p_ \sigma (\tilde{x} \vert x) := \mathcal{N} (x, \sigma ^ 2 \mathbf{I})$이라 하고, $p_ \sigma (\tilde{x}) := \int p_ {\text{data}} (x) p_ \sigma (\tilde{x} \vert x) dx$라 하자. 그리고 noise scale $\sigma_ {\min} = \sigma_ 1 \lt \cdots \lt \sigma_ N = \sigma_ {\max}$를 정의하자. 이때 $\sigma_ {\min}$은 충분히 작아서 $p_ {\sigma_ {\min}} (\tilde{x}) \approx p_ {\text{data}} (\tilde{x})$이다. 그리고 $\sigma_ {\max}$는 충분히 커서 $p_ {\sigma_ {\max}} (\tilde{x}) \approx \mathcal{N} (0, \sigma_ {\max} ^ 2 \mathbf{I})$이다.
이때 NCSN(Noise Conditional Score Network) $s_ \theta (x, \sigma)$는 weighted sum of denoising score matching objective에 의해 훈련된다. 여기서 $\lambda(t) = \sigma_ i ^ 2$로 설정한다.
\[L = \frac{1}{N} \sum_ {i=1} ^ N \sigma_ i ^ 2 \mathbb{E} _ {p_{\text{data}}(x)} \mathbb{E} _ {p_ {\sigma_ i} (\tilde{x} \vert x)} \left[ \left\Vert s_ \theta (\tilde{x}, \sigma_ i) - \nabla_ {\tilde{x}} \log p_ {\sigma_ i} (\tilde{x} \vert x) \right\Vert ^ 2 \right]\]Sampling은 다음과 같은 Langevin MCMC를 통해 이루어진다. $p_ {\sigma_ i} (x)$의 sample을 얻기 위해 아래 과정을 $M$번 반복하고, 이렇게 얻은 $x_ i ^ M = x_ {i-1} ^ 0$이다.
\[x_ i ^ m = x_ i ^ {m-1} + \epsilon_ i s_ \theta (x_ i ^ {m-1}, \sigma_ i) + \sqrt{2 \epsilon_ i} z_ i ^ m\]VE SDE
SMLD의 forward process 과정을 SDE로 나타낼 수 있다. 먼저 $p_ {\sigma_ i} (\tilde{x} \vert x) := \mathcal{N} (x, \sigma _ i ^ 2 \mathbf{I})$로부터 다음과 같은 식을 얻는다.
\[x_ i = x_ {i-1} + \sqrt{\sigma_ i ^ 2 - \sigma_ {i-1} ^ 2} z_ {i-1}\]이때 noise scale의 수 $N \to \infty$라 하고, $x_i = x(i/N)$, $\sigma_i = \sigma(i/N)$, $z_i = z(i/N)$이라 하자. 또한, $\Delta t = 1/N$이라 하자. 이 가정 하에 continuous stochastic process $\lbrace x(t) \rbrace _ {t = 0} ^ 1$을 다음과 같이 나타낼 수 있다.
\[\begin{aligned} x(t + \Delta t) &= x(t) + \sqrt{\sigma(t + \Delta t) ^ 2 - \sigma(t) ^ 2} \cdot z(t) \\ &\approx x(t) + \sqrt{ \frac{d \left[\sigma ^ 2 (t) \right]}{dt} \Delta t} \cdot z(t) \\ \end{aligned}\]따라서 $\Delta t \to 0$이라 하면, 다음과 같은 forward SDE를 얻을 수 있다. 이를 점점 시간이 지날수록 variance $\sigma(t)$가 커진다고 하여 VE SDE(Variance Exploding SDE)라고 부른다.
\[dx = \sqrt{ \frac{d \left[\sigma ^ 2 (t) \right]}{dt}} dw\]2.3. DDPM (VP-SDE)
Background
DDPM(Denoising Diffusion Probabilistic Model)은 이미 다루었기 때문에 간단히 score-based model 방식의 표기만 정리하겠다. 먼저 noise scale $0 \lt \beta_1 \lt \cdots \lt \beta_ N \lt 1$, $\alpha_i := \prod_ {j=1} ^ i (1 - \beta_ j)$로 정의하자. 이때 forward process는 다음과 같이 나타낼 수 있다.
\[\begin{aligned} p(x_i \vert x_{i-1}) &= \mathcal{N} (x_i; \sqrt{1- \beta_i} x_{i-1}, \beta_i \mathbf{I}) \\ p_ {\alpha_ i} (x_i \vert x_0) &= \mathcal{N} (x_i; \sqrt{\alpha_i} x_0, (1- \alpha_i) \mathbf{I}) \\ \end{aligned}\]SMLD와 유사하게 이를 perturbed data distribution $p_ {\alpha_ i} (\tilde{x}) := \int p_ {\text{data}} (x) p_ {\alpha_ i} (\tilde{x} \vert x) dx$으로 나타낼 수 있다. 그리고 score function을 이용하여 objective를 다음과 같이 나타낼 수 있으며, 이는 DDPM의 $L_ {\text{simple}}$과 동일하다. 이때 $\lambda(t) = 1 - \alpha_ i$로 설정한다. $\alpha_i$는 $i$가 커질수록 작아지므로, 동일하게 $i$가 커지면 가중치가 커진다.
\[L = \frac{1}{N} \sum_ {i=1} ^ N (1 - \alpha_ i) \mathbb{E} _ {p_{\text{data}}(x)} \mathbb{E} _ {p_ {\alpha_ i} (\tilde{x} \vert x)} \left[ \left\Vert s_ \theta (\tilde{x}, i) - \nabla_ {\tilde{x}} \log p_ {\alpha_ i} (\tilde{x} \vert x) \right\Vert ^ 2 \right]\]Sampling은 다음과 같이 reparametrization trick의 결과로 이루어진다. 이를 ancestral sampling이라고 부른다.
\[x_ {i-1} = \frac{1}{\sqrt{1- \beta_ i}} (x_i + \beta _i s_ \theta (x_i, i)) + \sqrt{\beta_ i} z_i\]Derivation
DDPM 논문의 수식을 이용하여 위 두 식을 유도해보자. DDPM 원 논문에서 사용했던 $L_ {\text{simple}}$을 다시 생각해보면 다음과 같다.
\[L_ {\text{simple}} = \frac{1}{N} \sum_ {i=1} ^ N \mathbb{E} _ {p_ {\text{data}} (x)} \mathbb{E} _ {p_ {\alpha_ i} (\tilde{x} \vert x)} \left[ \left\Vert \epsilon - \epsilon_ \theta (\tilde{x}, i) \right\Vert ^ 2 \right]\]이때 $p_ {\alpha_ i} (\tilde{x} \vert x) = \mathcal{N} (\tilde{x}; \sqrt{\alpha_i} x, (1- \alpha_i) \mathbf{I})$이다. 따라서 $\nabla_ {\tilde{x}} \log p_ {\alpha_ i} (\tilde{x} \vert x)$는 다음과 같이 나타낼 수 있다. 이는 $\nabla_ {\tilde{x}} \log \mathcal{N}(x, \sigma ^ 2 \mathbf{I}) = - (\tilde{x}- x) / \sigma ^ 2$임을 활용한 것이다.
\[\begin{aligned} \nabla_ {\tilde{x}} \log p_ {\alpha_ i} (\tilde{x} \vert x) &= - \frac{1}{1 - \alpha_ i} (\tilde{x} - \sqrt{\alpha_i} x) \\ &= - \frac{1}{1 - \alpha_ i} ( \sqrt{1 - \alpha_i} \epsilon ) \\ \therefore \epsilon &= - \sqrt{1 - \alpha_i} \nabla_ {\tilde{x}} \log p_ {\alpha_ i} (\tilde{x} \vert x) \\ \end{aligned}\]이에 따라 $\epsilon_ \theta (\tilde{x}, i) = - \sqrt{1 - \alpha_i} s_ \theta (\tilde{x}, i)$로 parametrize하고, 다시 $L_ {\text{simple}}$에 대입하면 바로 본 논문의 objective $L$을 얻는다.
\[L = \frac{1}{N} \sum_ {i=1} ^ N (1 - \alpha_ i) \mathbb{E} _ {p_{\text{data}}(x)} \mathbb{E} _ {p_ {\alpha_ i} (\tilde{x} \vert x)} \left[ \left\Vert s_ \theta (\tilde{x}, i) - \nabla_ {\tilde{x}} \log p_ {\alpha_ i} (\tilde{x} \vert x) \right\Vert ^ 2 \right]\]한편 DDPM 원 논문의 sampling 과정은 다음과 같이 나타낼 수 있다.
\[x_ {i-1} = \frac{1}{\sqrt{1- \beta_ i}} \left(x_i - \frac{\beta_ i}{\sqrt{1 - \alpha_ i}} \epsilon_ \theta (x_i, i) \right) + \sigma_ t z_i\]DDPM에서는 $\sigma_ t ^ 2 = \beta_ i$로 설정하였으므로 이를 다시 정리하면 본 논문의 sampling 식을 얻는다.
\[x_ {i-1} = \frac{1}{\sqrt{1- \beta_ i}} (x_i + \beta _i s_ \theta (x_i, i)) + \sqrt{\beta_ i} z_i\]VP SDE
$p(x_i \vert x_{i-1}) = \mathcal{N} (x_i; \sqrt{1- \beta_i} x_{i-1}, \beta_i \mathbf{I})$로부터 다음 식을 얻을 수 있다.
\[x_i = \sqrt{1- \beta_i} x_{i-1} + \sqrt{\beta_i} z_{i-1}\]SMLD와 마찬가지로 noise scale의 수 $N \to \infty$라 하고, noise scale $\beta_i = \beta(i/N)$, $x_i = x(i/N)$, $z_i = z(i/N)$라 하자. 동일하게 다음과 같이 나타낼 수 있다.
\[\begin{aligned} x(t + \Delta t) &= \sqrt{1- \beta(t + \Delta t) \Delta t} x(t) + \sqrt{\beta(t + \Delta t) \Delta t} z(t) \\ &\approx x(t) - \frac{1}{2} \beta(t + \Delta t) \Delta t x(t) + \sqrt{\beta(t + \Delta t) \Delta t} z(t) \\ &\approx x(t) - \frac{1}{2} \beta(t) \Delta t x(t) + \sqrt{\beta(t) \Delta t} z(t) \\ \end{aligned}\]따라서 $\Delta t \to 0$이라 하면 다음과 같은 forward SDE를 얻을 수 있다. 이 경우에는 variance가 bounded되어 있다고 하여 VP SDE(Variance Preserving SDE)라고 부른다.
\[dx = - \frac{1}{2} \beta(t) x dt + \sqrt{\beta(t)} dw\]2.4. sub-VP SDE
저자들은 VP SDE에 영감을 얻어 variance가 어떤 time step에서도 VP SDE보다 작거나 같은 sub-VP SDE를 제안한다. 이는 다음과 같이 나타낼 수 있다.
\[dx = - \frac{1}{2} \beta(t) x dt + \sqrt{\beta(t) \left(1 - \exp \left(-2 \int _ 0 ^ t \beta(s) ds \right) \right)} dw\]Variance가 작기 때문에 더 deterministic하고, 정확한 샘플링이 가능할 것으로 기대할 수 있다.
3. Solving the Reverse SDE
3.1. Numerical SDE Solver
Forward SDE로부터 reverse SDE를 다음과 같이 나타낼 수 있었다.
\[\begin{aligned} \text{Forward} \quad & dx = f(x, t)dt + g(t)dw \\ \text{Reverse} \quad & dx = \left[f(x,t) - g(t) ^ 2 \nabla_ x \log p_t (x) \right] dt + g(t) d \bar{w} \\ \end{aligned}\]이를 discretized form으로 나타내면 다음과 같다. 참고로 DDPM의 ancestral sampling도 아래의 reverse sampling으로 근사할 수 있으며, 이 numerical SDE solver를 저자들은 reverse diffusion sampler라고 부른다. SDE 기반의 기본적인 sampler라고 생각하면 된다.
\[\begin{aligned} \text{Forward} \quad & x_{i+1} = x_i + f_i (x_i) + g_i z_i \\ \text{Reverse} \quad & x_i = x_ {i+1} - f_{i+1} (x_ {i+1}) + g_{i+1} ^ 2 s_ \theta (x_ {i+1}, i+1) + g_{i+1} z_ {i+1} \\ \end{aligned}\]3.2. Predictor-Corrector Samplers
저자들은 여기에서 더 나아가서 numerical SDE solver가 $x_{i+1}$로부터 예측값 $x_i ^ \prime$을 계산하고, 이 샘플의 marginal distribution을 교정하는 score-based MCMC를 사용하여 결과적으로 $x_i$를 얻는 predictor-corrector sampler를 제안한다. 여기에서 reverse diffusion sampler는 predictor, Langevin MCMC는 corrector로 사용된다. 이를 통해 더 정확한 샘플링을 기대할 수 있다.
3.3. Probability Flow ODE
Probability Flow ODE는 SDE와 동일한 결과를 얻을 수 있는 ODE를 대신 풀어 계산하는 것으로 SDE와는 달리 deterministic하다는 장점이 있다. SDE와 동일한 marginal probability density $\lbrace p_t(x) \rbrace _ {t=0} ^ T$를 얻기 위해서는 다음과 같은 ODE를 풀어야 한다. 유도는 본 논문의 Appendix에 수록되어 있다.
\[dx = \left[ f(x, t) - \frac{1}{2} g(t) ^ 2 \nabla_ x \log p_t (x) \right] dt\]이러한 Probability Flow ODE를 사용하는 경우, likelihood를 아래와 같은 식으로 계산할 수 있다. ELBO만 계산할 수 있었던 기존의 방법과 달리 likelihood를 직접 계산할 수 있게 되었다.
\[\begin{aligned} \log p_0(x(0)) &= \log p_T(x(T)) + \int_ 0 ^ T \nabla \cdot \tilde{f} _ \theta (x(t), t) dt \\ &= \log p_T(x(T)) + \int_ 0 ^ T \mathbb{E}_ {p(\epsilon)} \left[ \epsilon^ T \nabla \tilde{f} _ \theta (x(t), t) \epsilon \right] dt \\ \end{aligned}\]그리고 latent representation의 manipulation이 가능하다. 저자들은 이를 이용하여 interpolation, temperature scaling 등을 수행하였다.
아래는 두 이미지의 latent representation을 이용하여 spherical interpolation을 수행한 결과이다.
아래는 embedding의 norm을 조절하여 temperature scaling을 수행한 결과이다.
이러한 ODE 방식의 또 다른 장점은 이미지에 대해 unique encoding을 부여한다는 것이다. 만약 충분히 많은 데이터셋에 대해 충분히 잘 학습되었다면 항상 재현 가능한 encoding을 얻을 수 있다. 예를 들어 개별적으로 학습된 두 모델에 대하여 동일한 이미지에 대한 latent representation은 거의 동일하다.
또한 효율적인 샘플링이 가능하다. 샘플링 단계를 뛰어넘으면서 step을 진행할 수 있고, error tolerance를 높일수록 evaluation 수가 줄어들어 efficiency-accuracy trade-off를 조절할 수 있다. 그리고 아래 그림을 통해 어느 정도까지는 이미지의 특징을 유지하면서도 샘플링 속도를 높일 수 있다는 것을 확인할 수 있다.
3.4. Results
먼저 sampling 방법에 따른 결과를 보면 predictor-corrector sampler를 사용한 경우가 가장 좋았으며, score-based model의 Langevin MCMC만 사용한 경우 다른 방법에 비해 떨어지는 것을 확인할 수 있다. 앞으로 score-based model에서 diffusion 방식의 sampling이 필연적임을 보여준다. Probability Flow ODE는 PC sampler와 같이 사용할 때만 성능 개선이 있음을 알 수 있다.
Table 2를 보면, ODE를 사용하여 exact likelihood computation이 가능했기 때문에 기존의 DDPM에 비해 NLL이 낮아진 것을 볼 수 있고, 기존의 $L_ {\text{simple}}$이 아닌 본 논문에서 제안한 continuous objective $L$를 사용한 경우 더 NLL이 낮아진 것을 볼 수 있다. 여기에 기존의 VP SDE가 아닌 sub-VP SDE를 사용한 경우, 그리고 모델 구조를 개선한 DDPM++을 사용한 경우, network depth가 2배인 deep DDPM을 사용한 경우 성능이 개선되었다.
한편 Table 3에서 NCSN과 DDPM을 비교한 경우 NCSN이 sample quality 면에서는 좋았지만 DDPM의 경우 likelihood가 더 높았다. 따라서 저자들은 필요에 따라 VE SDE와 VP SDE를 사용하여 sample quality와 likelihood를 trade-off할 수 있다고 제안한다.
4. Controllable Generation
Condition $y$를 가지고 $x$를 생성하는 controllable generation을 수행하기 위해서는 $\nabla_ x \log p_t (x(t) \vert y)$를 계산할 수 있어야 한다. 이는 확률의 성질에 의해 간단히 다음과 같은 방식으로 계산될 수 있다.
\[\nabla_ x \log p_t (x(t) \vert y) = \nabla_ x \log p_t (x(t)) + \nabla_ x \log p (y \vert x(t))\]만약 $y$가 class condition인 경우 classifier $p(y \vert x(t))$를 훈련시켜 class-conditional generation을 수행할 수 있다. 만약 $y$가 low-resolution image, masked image, grayscale image 등인 경우 이를 예측하는 새로운 diffusion model을 만들어 inpainting, colorization 등을 수행할 수 있다.
💡 Summary
- SDE를 사용하여 score-based model과 diffusion model을 통합하였고, 각각 VE SDE, VP SDE로 나타냈다. 더 나아가 sub-VP SDE를 제안하였다.
- Predictor-Corrector sampler로 샘플링의 정확도를 향상시킬 수 있었다.
- Probability Flow ODE를 사용하여 likelihood를 직접 계산할 수 있게 되었고, deterministic latent representation을 활용한 manipulation이 가능해졌으며 efficient sampling이 가능해졌다.
- 이를 응용하여 controllable generation을 수행할 수 있었다.
- Opinion: 중요한 generative model을 하나의 framework로 통합시킨 인상적인 논문이다. 수학적으로는 가장 본질에 근접한 논문이기에, 이 논문을 이해한다면 diffusion model에서 문제가 발생했을 때 모델에 대한 본질적인 고민을 할 수 있을 것이다.
댓글 남기기