Stochastic differential equation
Published on April 22, 2023 by JunYoung
SDE Score based model Neural ODE AI Deep learning
17 min READ
오늘 리뷰할 논문은 Score based generative modeling을 continous variable SDE로 풀어낸 논문이며, diffusion based approach 중 가장 유명한 DDPM과 더불어 디퓨전 기초 논문이라고 불리는 녀석이다(근데 난이도로 봐서는 기초는 아닌 것 같음). DDPM과 결을 달리하는 부분은 sampling 방식에서 numerical하게 미분 방정식을 푼다는 점이고 단계적 예측에 의존하는 DDPM이나 DDIM과는 다르다고 볼 수 있다. 이 논문의 장점은 score based approch를 사용하는 다양한 방법들을 일반화하고 다양하게 변화시켜 보다 효과적인 generative sampling에 대해 디스커션한다는 것이다.
DDPM보다는 수식이 훨씬 어렵기도 하고 확률 미분 방정식은 풀어내는 과정에서 constraint에 따라 해결할 수 있는 범위가 크게 바뀌기 때문에 조금 더 복잡하게 느껴진다. 그래도 해당 논문을 이해해야만 단계적 샘플링에 국한된 DDPM, DDIM과 같은 방식에서 벗어날 수 있는 사고가 가능하다고 느꼈다.
예를 들어 어떠한 데이터셋 $x \sim \mathcal{X}$에 대해 marginal 확률 분포인 $p(\mathcal{\mathcal{X}})$이 있다고 생각해보자. 샘플링하기 쉬운 임의의 분포 $z \sim \mathcal{Z}$로부터 시작하여 실제 데이터와 유사한 데이터 $\hat{x} \sim q(\mathcal{\hat{X} \vert \mathcal{Z}})$를 만들기 위해서는 두 분포 사이의 연관성을 찾아야하고, 이러한 task를 딥러닝으로 해결하고자 하는 것이 deep generative modeling의 개요이다.
대표적인 방법으로는 likelihood의 maximum lower bound(ELBO)를 사용하여 intractable한 확률 분포를 간접적으로 alignment하는 방법이 있으며(Variational Autoencoder),
학습 단계에서부터 임의로 샘플링한 노이즈 $z$에 대해 거짓 이미지(생성된 이미지)를 만드는 디코더와 실제 이미지를 구분하는 인코더의 적대적 학습을 통해 sampling quality를 올리는 GAN(Generative Adversarial Network)이 있다.
이외에도 역연산이 가능한 분포 간의 관계를 flow로 정의하고, 이를 통해 latent와 data 분포 사이의 관계를 찾는 flow based modeling도 있다.
그런 모델과는 다르게 score based model은 어떤 노이즈가 주어졌을때, 이를 기준으로 실제 데이터셋을 만들어낼 수 있는 방향(gradient)를 학습하여, 노이즈로부터 점진적으로 데이터 샘플을 만드는 process를 정의하게 된다.
Score에 대한 설명은 DDPM 게시글에서도 잠깐 다뤘었지만 다시 언급하자면 확률 밀도 함수에서의 기울기이다. 예컨데 원래 데이터의 연속 변수에 대한 확률 분포 $p(x)$가 있다면, $t$번째 step을 기준으로 하는 noised sample $x_t$이 그럴 듯한 샘플을 생성하게끔 움직여야하는 방향은 $p(x)$의 gradient에 해당될 것이고, 이를 energy based approach와 연관짓게 된다면 log likelihood의 gradient를 score로 정의할 수 있게 된다.
[ \nabla_x \log p(x) ]
바로 요 방법들을 사용한 생성 모델 중 혜성처럼 등장한 애들이 SMLD(Score matching with Langevin Dynamics)와 DDPM(Denoising diffusion probablistic Model)이다. 사실상 두 방법들 모두 일반화시키게 되면 같은 미분 방정식에 계수만 달라지는 꼴이 되므로 같은 문제를 푼다고 할 수 있다(이 부분에 대해서는 수식 증명에서 밝힌다).
그렇다면 대체 디퓨젼이라는 모델의 메커니즘과 확률 미분 방정식을 어떻게 연관지어 생각할 수 있을까? 예컨데 위와 같은 그림을 생각해보자. DDPM이나 SMLD에서 데이터로부터 노이즈를 만드는 과정은 간단하게도 아주 작은 diffusion을 더해줌으로써 성립한다. 이를 discretized($x_0,~x_1,~\cdots,~x_T$)해서 노이즈를 더했던 것과 다르게 만약 distribution이 점차 노이즈로 변해가는 과정이 일련의 연속 시간에 대해 정의될 수 있다면, 이는 Ito SDE의 솔루션으로 모델링될 수 있다.
사실 이 부분에 대한 내용은 뒤에서 DDPM 및 SMLD의 process와 연관지어 증명하면서 같이 이해해야하기 때문에 앞부분에서는 그냥 그렇다 생각하고 넘어갔다.
[ dx = f(x,~t)dt + g(t)dw ]
확률 미분 방정식도 미분 방정식이 가지는 여러 형태 중 하나이므로, 변화율($dx, dt, dw$)에 대한 관계로 문제가 정의된다. 미분 방정식의 해는 $x(t)$에 대해 해당 변화율이 그리는 trajectory(궤도)로 이해할 수 있다. 물론 확률 미분 방정식은 정해진 solution $x(t)$에 추가로 stochastic term $dw$가 붙는다는 점에서 차이가 있다(일단 얘는 무시하고 넘어가자).
$dw$는 시간에만 의존하는($x$의 실제 trajectory와는 무관) diffusion term이고 $dt$는 $x$의 실제 trajectory를 따라가는 방향을 정의하게 된다. 결국 solution $x(t)$는 특정 시점 $t$가 주어졌을 때, 해당 위치에서의 $x(t)$를 찾을 수 있는 함수의 꼴로 정의되고, 이를 위의 검은색 실선과 같이 함수의 형태로 표현할 수 있다.
하지만 이를 반대로 생각하면 solution을 정의하는 trajectory는 방향에 따라 달라질 수 있기 때문에 이번에는 역으로 $x_T$에서 $x_0$을 따라가면서 $x(t)$를 예측하는 새로운 문제를 정의해볼 수 있다.
[ dx = \left(f(x,~t) -g^2(t)\mathbf{\nabla_x \log p_t(x)}\right) dt + g(t) d\bar{w} ]
식을 잘 보게되면 diffusion term인 $g(t)dw$는 forward SDE와 동일하고 drift term인 $dt$가 적용되는 궤도의 방향을 forward에서 reverse로 바꿔주는 과정에 score term이 들어가는 것을 볼 수 있다(수식에서 진하게 표시된 부분).
결론부터 말하자면 forward SDE와 reverse SDE는 출발점과 도착점이 다르지만 서로 같은 궤도를 그려야하며, 만약 우리가 forward SDE(궤도의 형태)를 미리 알고 있다면 reverse SDE에게 솔루션을 제공해줄 수 있다는 것이다. 즉 score를 예측하는 것이 곧 reverse SDE를 푸는 과정이 되고, 따라서 SDE solver 형태의 모든 샘플러가 생성에 관여할 수 있다는 일반화로 이어진다. 이 논문에서는 PC(Predictor-Corrector) sampler와 deterministic sampler 두 가지를 소개한다. 이름에서도 볼 수 있듯이PC sampler의 경우 score based model을 활용하여 보다 퀄리티 좋은 샘플링을 목적으로 하며, deterministic sampler는 랜덤한 부분들을 모두 제거했기 때문에 빠른 샘플링이나 latent manipulation을 목적으로 한다는 점에서 차이가 있다.
SMLD이 뭔지 제대로 짚지 않고 넘어왔는데 이게 뭐냐면 NCSN(Noise Conditional Score Network)이다. NCSN이랑 DDPM 모두 score based generative model이면서 각각 SDE problem을 방법론으로 서로 다르게 발전시킨 논문들이다. 이에 대해 살펴보기 위해 SMLD와 DDPM에서 어떤 식으로 score estimator를 학습했고, 어떻게 샘플링했는지 살펴보도록 하자.
$x$에 대해 $\sigma^2$의 variance를 가지는 perturbation kernel이 더해졌다고 생각해보자. 해당 kernel은 noise 분포인 $p_\sigma(\tilde{x} \vert x) := \mathcal{N}(\tilde{x}; x, \sigma^2I)$를 만족한다. 점차적으로 커지는 noise scale인 $\sigma_{\min} = \sigma_1 < \sigma_2 < \cdots < \sigma_N = \sigma_{\max}$ 가 있다고 생각해보자. 이때 가장 작은 noise인 $\sigma_{\min}$에 대해서는 $p_{\sigma_{\min}}(x) \approx p_\text{data}(x)$를 만족할 수 있을 정도라고 보면 되고, 가장 큰 noise인 $\sigma_{\max}$에 대해서는 $p_{\sigma_{\max}}(x) \approx \mathcal{N}(x;0,\sigma_{\max}^2I)$를 만족할 수 있을 정도라고 생각하면 된다.
Denoising score matching의 개요는 직접 $\nabla_x \log p_\text{data}(x)$를 예측할 수 없기 때문에(실제 분포는 intractable) 이에 근접하게끔 아주 작은 노이즈가 더해진 noised sample $\tilde{x}$의 conditioned score를 예측하게 된다($\nabla_{\tilde{x}} \log p_{\sigma_i}(\tilde{x} \vert x)$). 가장 심한 노이즈부터 시작해서 원래의 데이터와 거의 유사한 수준의 데이터까지 다양한 노이즈 단계에 대해 score estimator $s_\theta(x, \sigma)$는 다음과 같은 objective를 가진다.
[ \theta^\ast = \underset{\theta}{\arg\min} \sum_{i=1}^N \sigma_i^2 \mathbb{E}_{p_\text{data}(x)} \mathbb{E}_{p_{\sigma_i}(\tilde{x} \vert x)} \left( \parallel s_\theta(\tilde{x}, \sigma_i) - \nabla_{\tilde{x}} \log p_{\sigma_i}(\tilde{x} \vert x) \parallel_2^2 \right) ]
충분한 데이터셋과 모델 capacity에 대해 학습된 score-based model은 샘플링 과정에 필요한 score를 잘 예측하게끔 학습된다. 샘플링은 조금 복잡하지만 Annealed Langevin dynamics를 사용했기에 이중 loop라고 생각하면 된다. 각 노이즈 단계 $\sigma_i$에 대해 $N$부터 $1$까지 점차 노이즈의 농도를 낮춰가며 각 noise kernel에 대해 총 $M$번의 Langevin MCMC를 진행한다. 즉 $N \times M$ 만큼의 process라고 생각하면 된다. 예컨데 $\sigma_i$ 만큼의 노이즈가 더해진 분포에 대해서,
[ x_i^m = x_i^{m-1} + \epsilon_i s_{\theta^\ast}(x_i^{m-1}, \sigma_i) + \sqrt{2\epsilon_i}z_i^m ]
gradient ascent의 learning rate라 볼 수 있는 step size $\epsilon_i > 0$와 랜덤한 standard normal distribution vector $z_i^m \sim \mathcal{N}(0, I)$에 대해 위의 프로세스를 총 $M$번 진행한다. $M$번 만큼의 sampling이 끝나게 되면 농도가 낮춰진 노이즈 단계 $\sigma_{i-1}$에 대해 다시 $M$번의 MCMC를 진행한다($x_{i-1}^0 = x_i^M$).
[ x_{i-1}^m = x_{i-1}^{m-1} + \epsilon_{i-1} s_{\theta^\ast}(x_{i-1}^{m-1}, \sigma_{i-1}) + \sqrt{2\epsilon_{i-1}}z_{i-1}^m ]
for sigma in sigma_N, ..., sigma_1:
for step in 1, ..., M:
noise = N(0, I) # Random noise
x = x + epsilon * model(x, sigma) + sqrt(2*epsilon)*noise
충분히 큰 $M$(sampling 수)과 아주 작은 $\epsilon$(step size)를 통해 생성한 샘플인 $x_1^M$은 다시 원래 데이터의 분포를 따라가게 된다.
[ p_{\sigma_{\min}}(x) \approx p_{\text{data}}(x) ]
DDPM에서는 noise scale을 단계적으로 증가시킨 variance schedule(sequence)인 $\beta_i$를 정의한다.
[ 0 < \beta_1,~\beta_2,~\cdots,~\beta_N < 1 ]
각 training data point(이미지) $x_0 \sim p_\text{data}(x)$에 대해, 앞서 정의한 variance schedule에 대해 perturbation kernel를 다음과 같이 Markov Chain에 따라 다음과 같이 정의하게 된다.
[ p(x_i \vert x_{i-1}) = \mathcal{N}(x_i;~\sqrt{1-\beta_i}x_{i-1},~\beta_i I) ]
따라서 이를 모든 sequence에 대해 적용하면 다음과 같은 marginal을 획득할 수 있다.
[ p(x_i \vert x_0) = \mathcal{N}(x_i; \sqrt{\bar{\alpha}_i}x_0,~(1-\bar{\alpha}_i)I),~\bar{\alpha}_i := \prod_{j=1}^i (1-\beta_j) ]
또한 해당 표현법은 앞서 SMLD의 예시와 같이perturbed data distribution으로 표현할 수 있다 ($p_{\bar{\alpha}_i}(\bar{x} \vert x)$가 위에서 획득한 marginal을 의미함).
[ p_{\bar{\alpha}_t}(\tilde{x}) = \int p_\text{data}(x)p_{\bar{\alpha}_i}(\tilde{x} \vert x) dx ]
또한 DDPM의 경우에는 노이즈가 더해지는 각 process에서 변수의 variance가 $1$로 고정되기 때문에 $x_N \sim \mathcal{N}(0, I)$를 만족하게 된다. DDPM에서는 reverse process를 다른 형태로 표현했지만 SMLD와 같이 score matching 형태로 reverse process의 variational Markov chain을 표현하면 다음과 같다.
[ p_\theta (x_{i-1} \vert x_i) = \mathcal{N}(x_{i-1};~\frac{1}{\sqrt{1-\beta_i}}(x_i + \beta_i s_\theta(x_i,~i)), \beta_iI) ]
그리고 이에 따라 evidence lower bound도 다른 weighted summation을 가지게 된다. 아래의 식을 보게되면 SMLD에서 전개했던 식과 상당히 닮아있는 것을 볼 수 있는데, weighted summation의 계수로 작용하는 $1-\bar{\alpha}_i$가 $i$번째 perturbation kernel인 $p_{\bar{\alpha}_i}(\tilde{x} \vert x)$의 variace인 것을 알 수 있다. 앞서 SMLD에서도 계수가 perturbation kernel의 variance인 $\sigma_i^2$였던 것을 생각하면 두 식의 공통점을 찾을 수 있다.
보다 엄밀하게 말하자면 noised score estimation인 $\nabla_x \log p(\tilde{x} \vert x)$에 대해 $1/\mathbb{E}\left[ \parallel \nabla_x \log p(\tilde{x} \vert x) \parallel_2^2 \right]$에 비례하게 된다.
[ \theta^\ast = \underset{\theta}{\arg\min} \sum_{i=1}^N (1-\bar{\alpha}_i) \mathbb{E}_{p_\text{data}(x)} \mathbb{E}_{p_{\bar{\alpha}_i}(\tilde{x} \vert x)} \left( \parallel s_\theta(\tilde{x}, i) - \nabla_{\tilde{x}} \log p_{\bar{\alpha}_i}(\tilde{x} \vert x) \parallel_2^2 \right) ]
따라서 해당 loss term을 만족하는 $\theta^\ast$에 대해서 역으로 Markov chain process를 통해 샘플링을 진행한다.
[ x_{i-1} = \frac{1}{\sqrt{1-\beta_i}} (x_i + \beta_i s_{\theta^\ast}(x_i, i)) + \sqrt{\beta_i}z_i,~i=N,~N-1,~\cdots,~1 ]
이러한 샘플링 방법을 저자들은 ancestral sampling(조금씩 거슬러 올라오는 샘플링 과정을 reverse process인 $p_\theta$에 대해 명명)이라고 부른다. 사실상 DDPM의 loss 식은 위와 구조를 다르게 하지만(아래 참고) 실제로 DDPM이 score estimation model을 학습시키는 과정이라는 것을 보여주기 위해 위와 같이 수식화함.
[ \mathbb{E}_{x_0,\epsilon} \left( \frac{\beta_t^2}{2\sigma_t^2 \alpha_t(1-\bar{\alpha}_t)} \parallel \epsilon - \epsilon_\theta(\sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon, t) \parallel^2\right) ]
결국 말하고자 하는 것은 기존의 score-based generation 접근법이던 SMLD나 DDPM 모두 inference(forwarding)이 SDE로 표현되었던 것처럼, 역과정도 마찬가지로 score matching을 통해 SDE를 풀어가는 것으로 표현될 수 있다는 것이다.
기존 방식들을 보면 data를 여러 noise scale로 구성된 perturbation kernel를 사용한 연구들이라고 요약할 수 있다. 이 논문에서는 score based generation 방법을 미리 정해진 noise scale(discrete value)가 아닌 모든(continuous value) noise scale에 일반화시키고자 하며,
이는 곧 SDE를 연속 변수 방정식으로 풀고자 하는 목표로 작용한다(아래 그림 참고).
앞서 디퓨전 프로세스의 forward와 reverse 모두를 SDE로 표현할 수 있다고 했는데, 이를 다시 언급하면 다음과 같다.
Diffusion의 모델링의 가장 기본은 우리가 가지고 있는 임의의 i.i.d. 데이터셋인 $x(0) \sim p_0$에 continous/discrete time 변수 $t \in [0, T]$로 인덱싱되는 diffusion process를 기반으로 가우시안 노이즈에 가까운 $x(T) \sim p_T$를 만드는 것이다. 데이터 분포 $p_0$로부터prior 분포 $p_T$를 만드는 diffusion process는 앞서 소개했던 것과 같이 확률 미분 방정식의 solution으로 모델링할 수 있다.
[ dx = f(x,~t)dt + g(t) dw ]
$w$는 Brownian motion이라고도 불리는 standard Wiener process를 의미한다. $f(x, t)$는 $x, t$에 대한 drift coefficient이며 $g(t)$는 diffusion coefficient이다. SDE는 $f(x, t)$ 및 $g(t)$가 state와 time에 대해 모두 Lipshitz를 만족한다는 가정 하에(미분 값이 bounded되어 있다고 할 때) 한정된 변화율 사이에서 unique strong solution을 가질 수 있다(참고 링크).
[ \begin{aligned} &\vert \nabla_{x,~t} f(x,~t) \vert \le \epsilon_1 \newline &\vert \nabla_{t} g(t) \vert \le \epsilon_2 \end{aligned} ]
Prior를 알고 있다면(사전에 정의를 했다면) prior sample인 $x(T) \sim p_T$을 샘플링할 수 있다. 예컨데 우리가 prior를 다변수 가우시안으로 정의했다면 가우시안 샘플링을 통해 $x_T$를 구할 수 있다. 그리고 증명을 통해 diffusion SDE의 역과정 또한 SDE임을 다음과 같이 수식화할 수 있다(참고 링크).
[ dx = \left(f(x,~t) -g^2(t)\mathbf{\nabla_x \log p_t(x)}\right) dt + g(t) d\bar{w} ]
$\bar{w}$는 표현만 다를 뿐 앞서 forward process에서 봤던 Wiener process인 $w$와 동일하다고 보면 된다. 앞서 noise를 점차 만들어가는 과정은 아주 작은 단위로 $t$가 증가하면서 미분 방정식의 궤도를 그리지만, reverse 식에서는 아주 작은 단위의 $t$가 감소하면서 미분 방정식의 궤도를 그려나간다. 만약 위의 식에서 marginal distribution $p_t(x)$에서 score를 알 수 있다면 임의의 노이즈 샘플 $x_T$을 통해 $x_0$를 구할 수 있게 된다.
결국 샘플링을 위해서는 score를 예측해야한다는 것인데, 실질적으로 원래 데이터셋 분포의 marginal에 대한 score를 구할 수 없기 때문에 score-matching에 기반한 time-dependendent model인 $s_\theta(x, t)$를 앞서 본 SMLD와 DDPM의 loss term과 같이 최적화할 수 있다. 다만 앞서 본 SMLD와 DDPM은 연속되지 않은 perturbation sequence $N$에 대해 계산을 했다면, 이를 연속 변수 $t$에 대해 일반화하게 되면 다음과 같다.
[ \theta^\ast = \underset{\theta}{\arg\min} \mathbb{E}_t\left(\lambda(t) \mathbb{E}_{x(0)} \mathbb{E}_{x(t) \vert x(0)} \left( \parallel s_\theta(x(t), t) - \nabla_{x(t)} \log p_{0t}(x(t) \vert x(0)) \parallel_2^2 \right)\right) ]
앞서 언급했던 것과 같이 $\lambda(t) \propto 1/\mathbb{E} \left[ \parallel \nabla_{x(t)} \log p_{0t}(x(t) \vert x(0)) \parallel_2^2 \right]$인 값으로 고르게 된다. 물론 굳이 이 식과 같이 score matching을 하는 방식이 denoising score matching이 아니더라도 다른 방법들 중 하나인 sliced score matching이나 finite-difference score matching이어도 상관이 없다.
SMLD와 DDPM의 noise perturbation 방식은 SDE를 discrete하게 바꾼 것이라고 볼 수 있다. $N$개의 noise scale이 있을 때, SMLD의 Markov chain에 기반한 forward process는 다음과 같다.
[ x_i = x_{i-1} + \sqrt{\sigma_i^2 - \sigma_{i-1}^2} z_{i-1},~i = 1,~\cdots,~N ]
해당 식에서 $N$이 극한으로 증가한다면($\infty$) noise scale은 연속 시간에 대한 함수 $\sigma(t)$로 표현할 수 있으며 마찬가지로 standard gaussian variable $z_{i-1}$ 또한 $z(t)$로 나타낼 수 있다. 이때의 Markov chain을 continuous stochastic process $x(t)$로 바꾼다면,
[ x(t+\Delta) = x(t) + \sqrt{\sigma(t+\Delta)^2 - \sigma(t)^2}z(t) ]
$\Delta$가 매우 작다는 가정 하에 테일러 1차 근사를 통해 미분 방정식으로 근사할 수 있다. 일반적인 오일러 메소드라고 보면 된다.
[ dx = \sqrt{\frac{d\sigma^2(t)}{dt}} dw ]
마찬가지로 DDPM의 경우에는
[ x_i = \sqrt{1-\beta_i} x_{i-1} + \sqrt{\beta_i}z_{i-1},~i=1,~\cdots,~N ]
위와 같은 forward process를 따르기 때문에 $N$이 극한으로 증가한다면,
[ x(t+\Delta) = \sqrt{1-\beta(t+\Delta) \cdot \Delta}x(t) + \sqrt{\beta(t+\Delta)\cdot \Delta}z(t) ]
위와 같다. 이때 $\beta$는 원래 scheduling된 각 Markov process에서의 노이즈이기 때문에 미소 단위의 variance 변화에 대해 $\Delta$가 추가로 곱해지는 형태가 된다. 역시 오일러 메소드를 통해 미분 방정식으로 근사하면
[ dx = -\frac{1}{2} \beta(t) x dt + \sqrt{\beta(t)} dw ]
위와 같이 정리된다. 논문에서는 SMLD의 미분 방정식은 variance가 계속 증가하는 VE(Variance Exploding) SDE이고 DDPM의 미분방정식은 variance가 유지되는 VP(Variance Preserving) SDE라고 명명하였다.
이때, drift coefficient와 drift coefficient에 대해 affine 형태를 가지기 때문에 variance 미분 방정식으로 발전시킬 수 있다(다음 식을 사용함).
[ \frac{d\Sigma_{\text{VP}}(t)}{dt} = \beta (t) (I - \Sigma_{\text{VP}}(t)) ]
참고로 위와 같은 수식 전개가 가능한 것은 gaussian 분포를 가정하고 있고 non-linear case가 아니기 때문이다. 해당 ODE를 풀게 되면 $x(t$ )의 covariance function을 구할 수 있다.
[ \Sigma_\text{VP} (t) = I+\exp\left(\int_0^t -\beta(s) ds\right)(\Sigma_\text{VP}(0) - I) ]
이 식은 모든 $t$에 대한 variance $\Sigma_\text{VP}$가 $\Sigma_{\text{VP}}(0)$을 기준으로 bounded 된다는 사실이다. 예컨데 DDPM에서와 같이 $\Sigma_\text{VP}(0) = I$라면 모든 $t$에서 $I$로 유지된다. 그리고 저자가 밝힌 새로운 sub-VP SDE는 다음과 같다.
[ dx = -\frac{1}{2} \beta(t) x dt + \sqrt{\beta(t) (1 - \exp\left(-2\int_0^t \beta(s) ds\right))}dw ]
그리고 같은 방식으로 Covariance에 대한 ODE solution을 구하면 새로운 미분 방정식에 대한 covariance를 구할 수 있다.
[ \Sigma_\text{sub-VP}(t) = I+\exp\left(-2\int_0^t \beta(s) ds \right)I + \exp\left(-\int_0^t \beta(s) ds \right) (\Sigma_\text{sub-VP}(0)-2I) ]
sub-VP SDE는 사실상 VP SDE를 upper bound로 가진다고 생각하면 좋다. Diffusion proces는 $dx = f(x,~t)dt + g(t)dw$ 형태의 모든 식이 가능하기 때문에 처음부터 discrete diffusion으로 접근한 것이 아닌 continuous function의 관점으로 접근했기 때문에 해당 SDE를 제시할 수 있었다고 판단된다. sub-VP의 경우 특정 실험에서 좋은 결과를 보였다고 한다(likelihood).
지금까지 총 3개의 서로 다른 SDE를 소개했는데, 각각의 perturbation kernel을 가우시안 형태의 marginal distribution으로 나타내면 다음과 같다.
[ p_{0t}(x(t) \vert x(0)) = \begin{cases} \mathcal{N}(x(t); x(0), (\sigma^2(t)-\sigma^2(0))I),& \text{(VE SDE)} \newline \mathcal{N}(x(t);x(0)e^{-\frac{1}{2}\int_0^t \beta(s)ds}, I-Ie^{-\int_0^t \beta(s) ds}),&\text{(VP SDE)} \newline \mathcal{N}(x(t);x(0)e^{-\frac{1}{2}\int_0^t \beta(s)ds}, \left(1-e^{-\int_0^t \beta(s) ds}\right)^2I),&\text{(sub-VP SDE)} \end{cases} ]
sub-VP SDE의 전개 방식이 조금 복잡해서 어렵게 느껴지는데, 단순히 원래 DDPM에서 노이즈를 만들었던 방식이
[ p(x_i \vert x_0) = \mathcal{N}(x_i; \sqrt{\bar{\alpha}_i}x_0,~(1-\bar{\alpha}_i)I),~\bar{\alpha}_i := \prod_{j=1}^i (1-\beta_j) ]
다음과 같았다면, variance 부분을 quadratic하게 만들어준 것과 같다.
[ p(x_i \vert x_0) = \mathcal{N}(x_i; \sqrt{\bar{\alpha}_i}x_0,~(1-\bar{\alpha}_i)^2I),~\bar{\alpha}_i := \prod_{j=1}^i (1-\beta_j) ]
미분 방정식에서 score 예측 모델을 통해각 time point에 대해 $s_\theta(\cdot)$만 알아낼 수 있다면 어쩌구저쩌구 이론에 의해 reverse-SDE를 구성할 수 있다는 것은 앞에서 입아프게 짚고 넘어왔다. Reverse SDE를 numerical하게 풀어내어 $x_0$을 찾는 과정이 결국 샘플링하는 것과 같다.
SDE를 numerical하게 풀어내는 것은 일종의 solution이 되는 함수의 형태(trajectory)를 예측하는 것과 같다. SDE를 풀어내는 방법으로는 Euler-Maruyama나 stochastic Runge-Kutta 방법 등등(사실 뭔지는 잘 모른다) 다양하게 있고, 암튼 말하고자 하는 것은 score predictor만 있다면 어느 형태의 SDE solver든 사용해서 sample을 만들어낼 수 있게 된다.
DDPM과 같은 ancestral sampling 방법은 특이 케이스인데, SDE가 조금만 달라지더라도 적용할 수 없다는 문제가 있다(일반화가 어렵다). 따라서 이를 좀 완화하기 위해 reverse diffusion sampler를 제안했고(아래 식과 같음), 해당 방법이 단순히 SMLD/DDPM 모델의 ancestral sampling보다 약간 더 좋은 성능을 보였다고 한다.
[ x_i = x_{i+1} - f_{i+1}(x_{i+1}) + G_{i+1}G_{i+1}^\top s_{\theta^\ast}(x_{i+1}, i+1) + G_{i+1}z_{i+1} ]
아무래도 보다 0에 가까운 미소 단위의 $\Delta$(time 변화)에 대해 정의할수록 ancestral sampling과 실제 SDE solution과의 차이가 줄어들게 되는데, 기존 ancestral sampler는 SDE의 원형을 고려하지 않기 때문인 듯하다($p_\theta(x_i \vert x_{i+1})$에 의존).
일반적으로 SDE를 풀어내는 것은 score prediction만 진행하게 되는 과정인데, 여기에 추가로 MCMC approach를 더해주게 된다(Langevin dynamics). 방법을 간단하게 numerical SDE가 다음 time step의 solution을 예측하면, 해당 point에서 score estimation을 토대로 correction sampling(조정 작업)을 진행하는 것이다. Numerical SDE solver가 predictor 역할을 수행하고 score model이 corrector 역할을 수행한다.
각각 알고리즘(Predictor만 사용/Corrector만 사용/둘 다 사용)에 대해 SMLD, DDPM 모델을 적용해봤을 때, PC를 같이 사용하는 것이 predictor만 사용하는 것이나 corrector를 사용하는 것보다 샘플링 효율성이 상당히 증대되었다고 한다.
Score-based model을 활용하여 SDE를 numerical하게 풀어내는 것에 사용하게 된다. 그런데 사실 모든 diffusion process는 marginal likelihood $p(x_{0:T})$를 동일하게 가지는 ODE를 찾을 수 있다. 예컨데 원래의 diffusion SDE가 다음과 같다면,
[ dx = f(x, t)dt + g(t)dw ]
이 SDE와 동일한 marginal likelihood를 가지는 ODE는
[ dx = \left( f(x, t) - \frac{1}{2}g(t)^2 \nabla_x \log p_\theta(x) \right) dt ]
아래와 같이 나타내어진다. 해당 부분에 대한 증명은 너무 복잡해서 아직 이해는 못했다… 아무튼$\tilde{f}(x, t) = f(x, t) - \frac{1}{2}g(t)^2 \nabla_x \log p_\theta(x)$에 대해 Wiener process term인 $dw$를 없앤 ODE
[ dx = \tilde{f}(x, t)dt ]
는 원래의 SDE와 동일한 marginal likelihood를 가진다. 즉, sampling을 했을때 가지는 결과가 같다는 것이다. 이럴 경우 DDPM에서와 같이 강제로 discrete한 likelihood를 어림짐작하는 것이 아니라(아래 참고),
[ p_\theta(x_0 \vert x_1) = \prod_{i=1}^D \int_{\delta_{-}(x_0^i)}^{\delta_+(x_0^i)} \mathcal{N}(x; \mu_\theta^i (x_1, 1), \sigma_1^2) dx ]
[ \delta_+(x) = \begin{cases} \infty,&\text{if }x=1 \newline x+\frac{1}{255},&\text{if }x<1 \end{cases},~~\delta_-(x) = \begin{cases} -\infty,&\text{if }x=-1 \newline x-\frac{1}{255},&\text{if }x>-1 \end{cases} ]
실제 likelihood를 계산할 수 있게 된다.
[ \log p_0(x(t)) = \log p_T(x(T)) + \int_0^T \nabla \cdot \tilde{f}_\theta(x(t),t)dt ]
그리고 ODE의 경우 $dw$(랜덤한 부분)을 아예 무시할 수 있기 때문에 특정 datapoint와 $x(0)$와 경로 상의 $x(T)$가 $1$대 $1$ mapping된다는 장점이 있다. 이는 마치 flow based model이나 neural ODE와 같은 invertible model로 생각할 수 있는데, 이럴 경우 latent representation을 사용하여 image editing이나 interpolation 등등이 손쉽게 가능해진다. 마찬가지로 각 이미지가 유일한 latent로 mapping되기 때문에 forward SDE를 일종의 encoder로도 생각할 수 있다(실제로 forwarding 과정은 파라미터 학습과 전혀 무관하다).
마지막으로 neural ODE는 stochastic한 부분이 사라지므로 빠른 샘플링이 가능하다.
Continous SDE를 해결하는 과정에서 또다른 장점 중 하나는 샘플을 생성하는 것 뿐만 아니라 조건부 확률인 $p_0(x(0) \vert y)$에 대한 샘플을 생성할 수 있다는 것이다. 이때, 각 time step의 class 확률 $p_t(y \vert x(t))$가 있어야한다.
[ dx = \left( f(x, t) - g(t)^2(\nabla_x \log p_t (x) + \nabla_x \log p_t(y \vert x)) \right)dt + g(t)d\bar{w} ]
클래스 conditional probability는 class guidance 논문인 Diffusion beats GAN과 동일한 방식으로 noised sample에 대한 supervised learning을 통해 인코더를 학습시킨다. 나머지 방법은 Appendix에 있는데 딱히 디테일하게 볼 필요는 없을 것 같아서 패스했다.
