SMLD와 DDPM을 score-based generative model의 category로 묶고, 이를 SDE관점에서 풀어낸 논문
0. Abstract
단계적으로 noise를 추가하지 않고 연속적으로 noise를 추가하는 SDE 식을 보이고, 그 SDE에 해당하는 reverse SDE를 통해 prior distribution(noise 분포)로부터 noise를 제거해 image를 생성할 수 있다. 이 때 reverser SDE는 score를 사용하는데 이 score는 기존 score based model이 score를 estimate하도록 학습하게 하여 얻을 수 있다. 그 이후에 reverser SDE는 잘 알려진 euler-maruyama 같은 SDE solver를 사용해서 풀면 image를 generation할 수 있다.
1. Introduction
noise를 연속적으로 추가해 data를 noise로 망가뜨리고 그 reverse를 학습하여 image를 생성하는 두 가지 probibility generative model이 있다. (SMLD, DDPM)
Score Matching with Langevin Dynamics (SMLD)
Score(시간 t에 대한 data ditribution인 p_t(x)의 log 및 x에 대한 gradient)를 각 noise scale에 대해 예측할 수 있도록 학습하고 annealed langevin dynamics (MCMC) 방법을 사용해 noise를 순차적으로 제거하며 data를 생성해낸다.
Langebin dynamic는 score를 가지고 데이터를 sampling 하는 방법을 말한다. 하지만 데이터가 적은 space에서는 gradient 값을 정확히 구할 수 없어 제대로 score를 학습하지 못하는 문제가 발생한다. 이를 해결하기 위해 단계적으로 noise를 줘서 각 noise에서의 score를 학습시킨다. sampling을 할 때는 큰 noise부터 작은 noise로 가며 latent image를 각 noise에서 converge하고 넘어가도록 한다. 이를 해결하기 위해 단계적으로 noise를 줘서 각 noise에서의 score를 학습시킨다. sampling을 할 때는 큰 noise부터 작은 noise로 가며 latent image를 각 noise에서 converge하고 넘어가도록 한다.
MCMC: target 확률분포(stationary distribution)에서 sampling한 것 처럼 sampling을 하는 것. marcov chain을 기반해 알고 있는 분포(x0)로부터 sampling한 data를 연쇄적으로 변형해 target 분포의 data로 만든다. (step이 많을수록 target 분포의 data와 가까워진다. Monte carlo: 수치를 얻기 위한 simulation, 무한한 시도를 통해 알 수 있는 결과를 유한한 시도의 결과로 추정
Denoising Diffusion Probabilistic Modeling (DDPM)
Denoising step의 확률적 분포를 학습해서 sampling을 하는 방법이다. 가우시안 분포를 따른다는 가정으로 확률 모델을 설계하고 이를 간소화해 noise(ε)만 예측하여 denoising하여 sampling한다. 이때 ε를 구하는 것은 state space가 continuous하다면 결국 score를 구하는 것과 마찬가지(가우시안 분포에 로그를 취하는 것이기 때문에 score에 상수를 곱한 것이 ε가 됨)이기 때문에 DDPM역시 score based generative model이라고 볼 수있다.
따라서 SMLD, DDPM 모두 score model이라는 범주로 볼 수 있고, step을 이산적이 아닌 연속적으로 고려하면 SDE로 이 문제를 볼 수 있게 된다.
forward SDE는 반드시 reverser SDE가 존재하기 때문에 SDE의 reverser를 유도할 수 있고 여기서 알 수 없는 score 값만 model로 추정해 구해내면 reverser SDE를 SDE solver를 활용해 sampling을 할 수 있게된다.
Key contributions
1. Flexible sampling & likelihood computation:
SDE로 문제를 정의하기 때문에 아무 SDE solver를 갖다 써도 되는 유연함이 있다.
Predictor-Corrector (PC) sampling 기법을 도입해서 sampling quality를 올렸다. SDE solver로 predict를 하고 갖고 있는 score를 활용하여 MCMC 방법을 사용하여 image를 더 target 분포의 중심에 가깝게 변경한다. SDE의 방법과 score model의 방법을 통합함으로써 성능을 높혔다.
SDE를 ODE로 p_t(x)를 그대로 사용하면서 변경할 수 있다는 것을 증명했다. 따라서 SDE기반으로 학습한 score를 reverse ODE에 적용해서 sampling할 수 있게되고 이는 deterministic하기 때문에 많은 이점을 준다. (fast sampling, uniquely identifiable encoding, exact likelihood computation)
2. Controllable generation:
training에 사용되지 않은 정보를 conditioaning해서 생성할 수 있다. conditional reverse SDE의 score는 unconditional reverser SDE의 score로 쉽게 얻을 수 있기 때문이다. 이는 class conditional generation, inpainting, colorization 등 다양한 작업을 가능하게 해준다.
3. Unified framework:
score based model을 SDE의 표현으로 통합했다. 다양한 SDE를 시도해볼 수 있으며 여기서는 sub-VP SDE를 도입해 좋은 성능을 냈다.
Objective (ε에 대한 식을 score-based와 유사하게 변형)Ancestral sampling (i: noise step)
DDPM의 식과 조금 다르다. α_tilda를 그냥 α로 표기하고 Σ_θ를 그냥 β로 사용했다. p(x)를 가우시안 분포 식으로 풀어 ε을 s로 표현하도록 했다.
3. Score-based generative modeling with SDEs
Multiple noise scale를 추가하는 score-based model의 아이디어를 infinite noise scale로 늘려 일반화한다. 이에 따라 noise 를 추가하고 제거하는 과정을 SDE로 모델링하여 시간 t에 대한 perterbed data distribution이 이 SDE를 따르도록 한다.
ODE에서는 pt(x) 분포가 p0(x) 분포에 대해 deterministic하게 구해지고, SDE에서는 확률적으로 구해짐 (데이터 하나(x0)의 시점에서 보면 ODE의 xt는 deterministic하고 SDE의 xt는 확률분포)
3.1 Perterbing data with SDEs
연속적인 시간 t에 대한 diffustion process x(t)를 구성해야 한다. x(0)은 data의 분포(p_0)를 따르고, x(T)는 noise 분포 (p_T: 보통 가우시안)를 따른다. 이 때 diffusion process x(t)는 아래의 SDE의 solution으로 모델링할 수 있다.
f: drift coefficient, g; diffusion coefficient
t시점 x(t)가 따르는 확률분포를 p_t(x)로 표기하고, 위의 SDE form은 어떻게 noise를 추가하는지에 따라 다양한 형태로 표현할 수 있다.
3.2 Generating samples by reversing the SDE
Diffusion process x(t)의 reverse 역시 diffusion process이고 그 reverse diffusion process는 아래의 reverse SDE의 solution이라는 것이 밝혀져 있다.
(5)번 식의 reverse, (reverse SDE)
3.3 Estimating scores for the SDE
위의 reverser SDE를 풀기 위해서는 score값을 알아내야한다. 직접 p_t(x)를 알 수 없으니 model로 추정한다. 하지만 p_t(x(t))를 구할 수 없으므로 갖고 있는 데이터인 x(0)에 대한 조건부로 p_t(x(t) | x(0))를 근사(denoising score matching)한다. (sliced score matching 등 다른 score matching 방법을 사용해도 된다.) (target 분포 p_t(x)는 알 수 없지만 갖고 있는 데이터 x0에 대해 p_t(x(t) | x(0))의 값은 계산할 수 있다 --> 해당 값을 근사하도록 모델을 학습시킨다.)
SDE의 score를 학습하기 위한 objective ((1), (3) 식의 일반화된 버전)
충분한 data가 있다면 s_θ*(x, t)는 score값과 유사해질 것이다.
3.4 Examples: VE, VP, sub-VP SDEs
3.1에서 얘기한 것처럼 SDE를 다양한 형태로 설계할 수 있다. SMLD, DDPM 역시 서로 다른 SDE를 discretization한 형태이다.
SMLD의 perterbation kernel p_σi(x | x0)은 아래와 같은 markov chain을 따르는 p_σi(x_i | x_i-1)의 분포라고 할 수 있다.
(i번째에 x0에서 σi의 noise를 추가하니까 1step은 이전의 noise와의 차이를 추가하는 꼴)
SMLD의 forward step의 distribution (markov chain으로 perterbation kernel 정의)
여기에서 N이 무한대로 가면, x_i는 continuous stochastic(diffusion) process가 되고 아래의 SDE의 solution이 된다.
SMLD의 SDE
DDPM도 유사하게 SDE로 나타낼 수 있다.
DDPM의 forward step의 distribution (markov chain으로 perterbation kernel 정의)DDPM의 SDE
SMLD의 경우에는 noise가 계속 더해지는 형태로 variance가 발산하고, DDPM은 x_i-1도 조정되며 noise가 더해지면서 variance가 일정하다. 이러한 특성을 따서 Variance Exploding(VE), Variance Perserving(VP) SDE라고 카테고리를 나눈다. 더 나아가서 VP를 변형시켜 variance가 발산은 하지만 VE보다는 항상 작은 hybrid SDE를 제안한다.
sub-VP SDE
VE, VP, sub-VP SDE 모두 affin drift coefficient를 갖기 때문에 perterbation kernel p_σi(x | x0)은 가우시안 분포이고, 이는 (7)번 식의 계산을 간단하게해 학습을 효율적으로 하게 해준다.
4. Solving the revserse SDE
Score-based model을 SDE로 변경하고 일반화된 denoising score matching 방법으로 score에 대한 objective를 세워 score에 대한 모델 s_θ*를 얻었다. 이제는 reverser SDE를 numetical solver와 score 값을 사용해 image를 sampling하면 된다.
4.1 General purpose numetical SDE solvers
Numerical solver들은 approximate trajectory를 제공한다. 즉, x(t)의 값을 확률적(noise를 추가함으로써 값은 하나만 나옴)으로 알아낼 수 있게 해준다. SDE를 discretization해주는 다양한 solver들을 자유롭게 reverser SDE에 적용할 수 있다. (논문에서는 reverser diffusion sampler를 제안하는데 아래의 probability flow ODE가 더 좋아서 설명 생략)
4.2 Predictor-Corrector samplers
일반적인 SDE를 푸는 것처럼 reverser SDE를 따르는 것에서 더 나아가서 얻어낸 score 값을 score-based model에서 사용했던 것처럼 MCMC approach를 추가로 적용할 수 있다.
각 step에서 SDE solver로 reverser SDE를 풀어서 next x_t를 얻고(predictor), score-based MCMC로 그 x_t의 분포의 가장 높은 부분으로 가도록 x_t를 수정한다(corrector).
이런 방법은 SMLD와 DDPM의 sampling 방법을 일반화 한 것으로, SMLD의 annealed langevin dynamics는 corrector를 사용하는 것과 같고, DDPM은 ancester sampling은 predictor만 사용하는 것과 같다. 반대로 말하면 이 논문의 방법론을 사용함으로써 SMLD는 다음 step의 x_t를 확률적으로 알아내는 ancestral sampling 기법을 가미하는 것이고, DDPM은 score를 구해 x_t를 p_t(x)의 확률분포에서 가장 확률(pdf의 값이)이 높은 상태의 x_t로 변경하는 langevin dynamics를 가미하게 되는 것이다.
4.3 Probability flow & connection to neural ODEs
모든 diffusion process(SDE의 solution)는 동일한 marginal probability densities p_t(x)를 갖는 deterministic process(ODE의 solution)가 존재한다. 따라서 reverse SDE와 p_t(x)를 공유하는 ODE를 구해낼 수 있고, 설계한 SDE의 coefficient(f, g)와 score-based model로 구한 score 값을 가지고 ODE를 풀어낼 수 있다. 즉, reverse SDE를 p_t(x)를 공유하는 ODE로 변경해 deterministic하게 sampling을 할 수 있게 되는 것이다.
probability flow ODE
이렇게 함으로써 몇가지 이점을 얻을 수 있다.
Exact likelihood computation을 계산할 수 있어 기존보다 더 정확한(낮은) test 결과를 얻을 수 있다.
Manipulating latent representations: ODE로 x(T)를 결정적으로 x(0) decode할 수 있기 때문에 x(T)에 대한 manipulation으로 x(0)에 변화를 줄 수 있다. (editting, interpolation 등)
Uniquely identifiable encoding: score값이 정확하다면 SDE와 ODE는 trajectory가 동일하기 때문에 encoding역시 uniquely identify하게 할 수 있다.
Efficient sampling: ODE solver를 사용함으로써 더 좋은 quality image를 더 빠른 step으로 얻을 수 있었다.
4.4 Architecture improvement
(7)번 식처럼 objective를 continuos하게 바꾸고 model의 depth를 더 키우니 좋은 성능이 나더라
5. Controllable generation
지금까지 얘기한 p0(x0)에서 data를 sampling하는 것을 얘기했다. 더 나아가서 p_t(y | x(t))를 알 수 있다면 학습된 모델을 가지고 p0(x(0) | y)같이 y라는 condition을 준 상태로도 sampling을 할 수 있다.
reverse SDE for conditional sampling
∇log p_t(y | x(t))를 다른 모델로 학습할 수도 있겠지만 휴리스틱하게 값을 계산할 수 있다. 따라서 위의 식을 추가적인 학습 없이 계산할 수 있고, 이에 따라 conditional sampling이 가능해진다.
