[논문링크] https://arxiv.org/abs/2011.13456
Score-Based Generative Modeling through Stochastic Differential Equations
Creating noise from data is easy; creating data from noise is generative modeling. We present a stochastic differential equation (SDE) that smoothly transforms a complex data distribution to a known prior distribution by slowly injecting noise, and a corre
arxiv.org
SMLD와 DDPM 논문을 읽지 않았다면 본 논문을 이해하는데 어려움이 있을 수 있다.
[SMLD 논문리뷰] https://outta.tistory.com/141
[2025-1] 최민서 - Generative Modeling by Estimating Gradients of the Data Distribution
https://arxiv.org/abs/1907.05600 Generative Modeling by Estimating Gradients of the Data DistributionWe introduce a new generative model where samples are produced via Langevin dynamics using gradients of the data distribution estimated with score matchin
blog.outta.ai
[DDPM 논문리뷰] https://outta.tistory.com/181
[2025-1] 최민서 - Denoising Diffusion Probabilistic Models
[DDPM] https://arxiv.org/abs/2006.11239 Denoising Diffusion Probabilistic ModelsWe present high quality image synthesis results using diffusion probabilistic models, a class of latent variable models inspired by considerations from nonequilibrium thermody
blog.outta.ai
1. Introduction
Diffusion Model은 여러 방식의 생성형 모델 중 하나로, 데이터에 노이즈를 더하여 확산시킨 후, 모델이 확산의 역과정을 학습하여 노이즈를 원래 데이터로 복원하는 방식의 모델이다. Score matching with Langevin dynmaics(SMLD)에서는 노이즈화된 데이터 확률분포의 로그의 미분값을 추정하는 신경망을 학습시키는데, 이 때 신경망은 노이즈의 레벨과 노이즈화된 데이터를 입력으로 받게 된다. Denoising Diffusion Probabilstic models(DDPM)에서는 특정 타임스텝 $t$에서 노이즈화된 데이터에 있는 노이즈의 값을 추정하는 신경망을 학습시키는데, 이 때 신경망은 타임스텝 $t$와 노이즈화된 데이터를 입력으로 받는다.
SMLD와 DDPM 모두 이산화된 노이즈나 타임스텝에 대해 해당 데이터의 점수를 추정하는 모델이라 볼 수 있다. 본 논문에서는 일반화한 확률미분방정식(SDE)으로 표현한다. 논문의 SDE는 이미지에 노이즈를 더해주는 forward process를 연속적인 시간 $t$에 대해 표현한다. 이를 $t$에 대하여 역방향으로 SDE를 풀어주면 reverse process를 구현할 수 있다. 아래 그림에서 논문의 주요 아이디어를 이해할 수 있다.
2. Background
2.1 Denosing Score Matching with Langevin dynamics (SMLD)
SMLD에서는 $\theta$를 아래와 같이 최적화시킨다.
$$ \theta^\ast = \arg \min_{\theta} \sum_{i=1}^{N} \sigma_i^2 \mathbb{E}_{p_{\text{data}}(\mathbf{x})} \mathbb{E}_{p_{\sigma_i}(\tilde{\mathbf{x}}|\mathbf{x})} \left[ \Vert s_{\theta}(\tilde{\mathbf{x}}, \sigma_i) - \nabla_{\tilde{\mathbf{x}}} \log p_{\sigma_i}(\tilde{\mathbf{x}} | \mathbf{x}) \Vert_2^2 \right] $$
이 때 $p_{\sigma_i}(\tilde{\mathbf{x}} | \mathbf{x})$는 $\mathcal{N}(\tilde{x};x,\sigma^2\mathbf{I})$와 같이 조건부 정규분포 노이즈로 정의된다.
샘플링은 아래와 같이 진행된다.
$$x^{m}_{i} = x^{m-1}_{i} + \epsilon_{i} s_{\theta^*} (x^{m-1}_{i}, \sigma_{i}) + \sqrt{2\epsilon_{i}} z^{m}_{i}, \quad m = 1, 2, \dots, M $$
위의 과정이 $i=N,N-1,\dots,1$을 따라 반복되면서 최종적으로 $p_{\sigma_{min}}(\mathbf{x})$가 $p_{data}(\mathbf{x})$와 거의 같아지게 된다.
2.2 Denoising Diffusion Probabilistic Models (DDPM)
DDPM에서는 $\theta$를 아래와 같이 최적화시킨다.
$$\theta^\ast = \arg \min_{\theta} \sum_{i=1}^{N} (1-\alpha_i) \mathbb{E}_{p_{\text{data}}(\mathbf{x})} \mathbb{E}_{p_{\sigma_i}(\tilde{\mathbf{x}}|\mathbf{x})}\left[\Vert s_{\theta}(\tilde{\mathbf{x}}, i) - \nabla_{\tilde{\mathbf{x}}} \log p_{\sigma_i}(\tilde{\mathbf{x}} | \mathbf{x}) \Vert_2^2\right] $$
이는 DDPM 논문의 $L_{simple}$을 이용한 최적화 방식을 SMLD와 비슷해 보이도록 변형한 것이다.
샘플링은 아래와 같이 진행된다.
$$ x_{i-1} = \frac{1}{\sqrt{1 - \beta_i}}
\left( \mathbf{x}_i + \beta_i s_{\theta^*} (\mathbf{x}_i, i) \right) + \sqrt{\beta_i} \mathbf{z}_i, \quad i = N, N-1, \dots, 1 $$
3. Score-based Generative Modeling with SDEs
3.1 Perturbing Data with SDEs
Section 2에서 볼 수 있듯이 기존의 score-based 모델들은 데이터에 노이즈를 추가한 뒤 score 값을 추정하도록 모델을 학습시켰다. Section 3.1에서는 데이터에 노이즈를 더하는 과정을 연속적인 시간 $t$에 대한 Ito SDE로 표현한다.
$$ d\mathbf{x} = \mathbf{f}(\mathbf{x},t)dt + g(t)d\mathbf{w} $$
위 SDE는 연속적인 시간 변수 $t\in[0,T]$에 대해 정의되었다. $\mathbf{f}\,:\,\mathbb{R}^d \rightarrow \mathbb{R}^d$는 drift coefficient라 불리고 $g\,:\, \mathbb{R} \rightarrow \mathbb{R}$는 diffusion coefficient라 불린다. $\mathbf{w}$는 standard Wiener process를 의미한다.
Ito SDE는 확률적인 현상을 반영하여 시스템의 동역학을 모델링하는 미분방정식이다. $\mathbf{f}$는 결정론적 부분으로, 시간에 따른 상태의 변화를 나타낸다. $g$는 확률적인 부분으로 브라운 운동을 나타내며 평균 $0$, 분산이 $dt$인 정규분포를 따른다. 시점 $t$에서의 데이터는 SDE를 적분하여 얻을 수 있다.
$\mathbf{f}$와 $g$를 적절히 설정하여 위 SDE가 SMLD, DDPM의 forward process와 같아지도록 할 수 있다.
3.2 Generating Samples by Reversing the SDE
Section 3.1의 SDE에서 시간 순서를 반대로 하면 reverse-time SDE를 아래와 같이 구할 수 있다.
$$ d\mathbf{x} = [\mathbf{f}(\mathbf{x},t)\,-\,g(t)^2 \nabla_\mathbf{x}\mathrm{log}p_t(\mathbf{x})]dt\,+\,g(t)d\bar{\mathbf{w}} $$
위 SDE를 통해 $\mathbf{x}(T)$로부터 $\mathbf{x}(0)$을 얻을 수 있다.
3.3 Estimating Scores for the SDE
학습은 아래와 같이 진행된다.
$$ \theta^\star = \arg\min_{\theta} \mathbb{E}_{t} \Biggl[ \lambda(t) \, \mathbb{E}_{\mathbf{x}(0)} \, \mathbb{E}_{\mathbf{x}(t)|\mathbf{x}(0)} \left\| s_\theta(\mathbf{x}(t), t) - \nabla_{\mathbf{x}(t)} \log p_{0t}(\mathbf{x}(t) \mid \mathbf{x}(0)) \right\|_2^2 \Biggr] $$
기존의 score-based 모델들과 마찬가지로 $\nabla_{\mathbf{x}(t)}\mathrm{log}\,p_{0t}(\mathbf{x}(t) | \mathbf{x}(0))$ 값을 신경망이 추정하도록 학습시켜준다. 또한 계수 $\lambda(t)$는 그 값이 $1/\mathbb{E}[\Vert \nabla_{\mathbf{x}(t)}\mathrm{log}\,p_{0t}(\mathbf{x}(t) | \mathbf{x}(0)) \Vert_2^2]$에 비례하도록 설정하고 이는 SMLD, DDPM의 방식과 동일하다.
이 때 시간 $t$는 기존과 다르게 $[0,T]$ 사이의 연속적인 변수로 취급된다. 즉 신경망 $s_\theta$는 연속적으로 주어지는 노이즈 레벨에 대하여 이미지에 더해진 노이즈 값을 찾아낼 수 있도록 학습된다.
효율적인 학습을 위해서는 커널 $p_{0t}(\mathbf{x}(t)|\mathbf{x}(0))$에 대한 정보가 필요하다. SDE에서 $\mathbf{f}(\mathbf{x},t)$가 affine 변환이면, $p_{0t}(\mathbf{x}(t)|\mathbf{x}(0))$는 정규분포가 되고 평균과 분산도 SDE의 해를 구하는 표준적인 방법으로 구할 수 있다.
3.4 Examples: VE, VP SDEs and Beyond
SMLD의 커널 $p_{\sigma_i}(\mathbf{x}|\mathbf{x}_0)$을 아래와 같이 나타낼 수 있다.
$$ \mathbf{x}_i = \mathbf{x}_{i-1} + \sqrt{\sigma_i^2 - \sigma_{i-1}^2}\mathbf{z}_{i-1}, \;\; i\,=\,1,\dots,N $$
위는 $N$개의 노이즈 스케일을 사용했을 때이고, 만약 $N\rightarrow\infty$를 해주고 연속적인 시간 변수 $t\in[0,1]$을 사용하면, 수열 $\{\sigma_i\}_{i=1}^N$는 함수 $\sigma(t)$가 되고 $\mathbf{z}_i$는 $\mathbf{z}(t)$가 되고 Markov chain $\{\mathbf{x}_i\}_{i=1}^N$는 연속적인 확률과정 $\{\mathbf{x}(t)\}_{t=0}^1$가 된다.
연속적인 $t$에 대해 아래와 같이 SDE로 perturbation process를 나타낼 수 있다.
$$ d\mathbf{x} = \sqrt{\frac{d[\sigma^2(t)]}{dt}}d\mathbf{w} $$
마찬가지로 DDPM의 각 perturbation kernel $\{p_{\alpha_i}(\mathbf{x} | \mathbf{x}_0)\}_{i=1}^N$을 아래와 같이 나타낼 수 있다.
$$ \mathbf{x}_i = \sqrt{1-\beta_i}\mathbf{x}_{i-1} + \sqrt{\beta_i}\mathbf{z}_{i-1},\;\; i\,=\,1,\dots,N $$
SMLD와 똑같이 $N\rightarrow\infty$를 하여 연속적인 $t$에 대해 아래와 같이 SDE로 표현한다.
$$ d\mathbf{x} = -\frac{1}{2}\beta(t)\mathbf{x}dt + \sqrt{\beta(t)}d\mathbf{w} $$
SMLD와 DDPM의 perturbation kernel을 위와 같이 연속적인 $t$에 대하여 SDE로 표현했다. SMLD의 SDE는 $t\rightarrow\infty$에 따라 분산이 매우 커져, Varience Exploding(VE) SDE라 한다. 반면 DDPM의 SDE는 분산이 유지되어, Variance Preserving(VP) SDE라 한다.
VP SDE에서 영감을 받아, 본 논문에서는 아래와 같은 새로운 유형의 SDE를 제시한다.
$$ d\mathbf{x} = -\frac{1}{2}\beta(t)\mathbf{x}dt + \sqrt{\beta(t)(1-e^{-2\int_0^t \beta(s)ds})}d\mathbf{w} $$
VP SDE와 같은 $\beta(t)$를 사용하면 위 SDE의 분산은 모든 시간스텝 $t$에서 VP SDE에 bounded 되어있다. 따라서 위 SDE를 sub-VP SDE라 한다.
세 SDE 모두 drift 계수가 affine이기 때문에 perturbation kernel $p_{0t}(\mathbf{x}(t)|\mathbf{x}(0))$은 정규분포이고 따라서 효율적으로 학습을 진행할 수 있다.
4. Solving the Reverse SDE
4.1 General-Purpose Numerical SDE Solvers
Score-based 모델 $s_\theta$를 학습시킨 후, 샘플 생성을 위해 reverse-time SDE를 풀어야 한다. 연속적인 $t$에 대해 SDE가 설계되었지만, 실제로 샘플링을 연속적으로 할 수 없기 때문에 수치적 근사를 하여 근사 해를 구한다. 일반적으로 알려진 Euler-Maruyama 방식 등을 이용해 SDE를 이산화하여 아래와 같이 계산할 수 있다.
$$ \mathbf{x}_{t+\Delta t} = \mathbf{x}_t + f(\mathbf{x}_t,t)\Delta t + g(t)\sqrt{\Delta t}z $$
또한 저자는 reverse diffusion samplers라는 새로운 샘플링 방식을 제시한다. 이는 forward SDE와 같은 방식으로 reverse SDE를 이산화한다.
4.2 Predictor-Corrector Samplers
본 논문에서는 샘플링 과정에서 Predictor-Corrector(PC) 방식을 사용한다. 각 이산적인 스텝에서 Predictor가 다음 스텝의 이미지를 예측하고 Corrector가 예측의 오차값을 보정해준다.
VE SDE, VP SDE에서 Predictor, Corrector는 각각 아래와 같이 나타난다.
PC sampler의 predictor와 corrector에는 각각 적당한 알고리즘을 넣어주면 된다. 위 예시의 VE SDE의 경우, predictor로는 reverse diffusion sampler를 사용했고, corrector로는 annealed langevin dynamics를 적용했다.
reverse diffusion sampler는 일반적인 상황에서 reverse SDE를 이산적으로 근사하여 푸는 샘플러이다. 자세한 수학적 유도 과정은 논문의 Appendix G에 나와있다.
PC sampler는 SMLD와 DDPM의 샘플링 방식을 일반화한다. SMLD는 identity function을 predictor로, annealed Langevin dynamics를 corrector로 쓴다. DDPM은 ancestral sampling을 predictor로, identity function을 corrector로 쓴다.
아래는 여러가지 샘플링 방식에 대해 FID score를 측정한 결과이다.
위 표에서 확인할 수 있듯이 predictor나 corrector만 사용한 샘플링 방식보다 이들을 동시에 사용한 방식이 더 좋은 품질의 이미지를 생성함을 확인할 수 있다.
4.3 Probability Flow and Connection to Neural ODEs
우리가 section 3.2에서 정립한 reverse-time SDE는 결정론적인 부분과 그렇지 않은 부분으로 나누어져 있다.
$$ d\mathbf{x} = \left[ \mathbf{f}(\mathbf{x},t) - \frac{1}{2}g(t)^2 \nabla_\mathbf{x} \mathrm{log} \, p_t(\mathbf{x}) \right]dt $$
위 ODE는 시간 스텝 $t$에 따라 같은 marginal distribution ${p_t(\mathbf{x})}_{t=0}^T$을 공유한다. 본 논문에서는 이를 probability flow ODE라 부른다. 이는 neural ODE의 한 예시라고 볼 수도 있다.
Neural ODE는 연속적인 $t$에 대하여 hidden state을 미분방정식의 형태로 정의하는 방식이다. 초기 상태가 주어졌을 때 미분방정식의 값을 적분하여 최종상태, 혹은 특정 시점의 상태를 구할 수 있다.
Exact likelihood computation
위 ODE에 neural ODE의 확률 계산 방법을 적용하면 보다 정확한 log likelihood를 계산할 수 있다. 정확한 계산 방식은 논문의 Appendix D.2를 참고하자.
아래 표는 CIFAR-10 데이터셋에서 측정된 여러 모델의 negative log-likelihoods(NLL)과 FID score를 나타낸다.
표의 주요 결과는 아래와 같다.
- 본 논문에서는 DDPM 모델의 NLL을 probability ODE를 이용해 계산했는데, 이전 논문에서 ELBO를 이용해 계산한 값보다 더 좋은 값이 나왔다. (DDPM(L) VS DDPM)
- 똑같은 DDPM 모델을 연속적인 시간 스텝 $t$에 대해 학습시켰을 때 더 좋은 NLL 값을 가진다. (DDPM(L) VS DDPM cont. (VP))
- 이전 DDPM 논문에서 제안된 DDPM 모델보다, 본 논문의 여러 종류의 VP SDE가 더 좋은 NLL 값을 가진다. (DDPM(L) VS DDPM cont. (sub-VP))
- DDPM 모델 구조를 발전시키고, sub-VP SDE를 사용했을 때 가장 좋은 성능을 가졌다. (DDPM++ cont. (deep, sub-VP))
Manipulating latent representations
결정론적인 ODE를 적분하여 임의의 데이터 $\mathbf{x}(0)$를 latent space의 데이터 $\mathbf{x}(T)$로 변형할 수 있다. 이를 이용해 image editing, interpolation 등의 작업을 수행할 수 있다.
Uniquely identifiable encoding
다른 일반적인 invertible 모델들과 달리 본 논문의 모델은 uniquely identifiable하다. 즉, 데이터의 encoding인 $\mathbf{x}_t$는 $\mathbf{x}_0$의 값에 따라 유일하게 결정되며, 이는 forward SDE에 parameter가 없고, probability flow가 결정론적이기 때문이다.
Efficient sampling
위 Neural ODE를 통해 $\mathbf{x}_T$로부터 $\mathbf{x}_0$을 보다 효율적으로 샘플링할 수 있다. 특히 black-box ODE solver를 이용하면 좋은 품질의 이미지를 생성할 수 있고, 효율성과 정확도의 trade-off를 직접 조절할 수 있다. 자세한 과정은 논문의 Appendix D.4에서 볼 수 있다.
4.4 Architecture Improvements
NCSN++ cont, DDPM++ cont는 각각 NCSN, DDPM 모델을 연속적인 $t$에 대하여 VP, VE SDE 방식으로 학습시키고 네트워크의 깊이를 늘린 모델이다. Section 4.2에서의 표에서 볼 수 있듯이 두 모델 모두 기존의 NCSN, DDPM 모델보다 좋은 성능을 보인다.
본 논문의 모델 중 가장 샘플링 품질이 좋은(FID가 낮은) 모델은 NCSN++ cont. (deep, VE)이다. 가장 likelihood가 좋은 (NLL이 낮은) 모델은 DDPM++ cont. (deep, sub-VP)이다.
5. Controllable Generation
$p_t(\mathbf{y}|\mathbf{x}(t))$가 주어져 있다면 아래의 reverse-time SDE를 풀어 $p_0(\mathbf{x}(0)|\mathbf{y})$를 구할 수 있다.
$$ d\mathbf{x} = \left\{ \mathbf{f}(\mathbf{x},t)-g(t)^2 \left[ \nabla_\mathbf{x} \mathrm{log} \, p_t(\mathbf{x}) + \nabla_\mathbf{x} \mathrm{log}\, p_t(\mathbf{y}|\mathbf{x}) \right] \right\}dt + g(t)d\mathbf{w} $$
위 SDE에서 $\nabla_\mathbf{x} \mathrm{log}\, p_t(\mathbf{y}|\mathbf{x})$가 추가되어 조건 $\mathbf{y}$를 만족하는 방향으로 샘플링이 진행된다.
이는 class-conditional generation, image inpainting, colorization 등에 응용될 수 있다.
6. Conclusion
본 논문은 기존의 score-based 생성형 모델을 일반화하여 연속적인 시간 스텝 $t$에 대해 SDE를 정의했다. 또한 reverse-time SDE를 통해 새로운 샘플링 방식을 제시하였다. 또한 결정론적인 neural ODE는 정확한 우도 계산, 빠른 샘플링 등 다양한 이점을 제공한다. 하지만 논문의 방식은 GAN에 비해 샘플링 속도가 아직 느려 더욱 개선이 필요하다.
