cs236 13-14장 Score-Based Models

Generative model 복습을 위해 CS236 강의를 듣고 정리해보고자 한다.

피피티는 아래 페이지를 참고하면 된다.

https://deepgenerativemodels.github.io/

 

※ PPT의 내용 정리와 더불어 같이 보면 좋을 자료들을 정리했습니다. 강의를 보고 이해한대로 작성했기 때문에 부정확한 내용이 포함되어 있을 수 있음을 알려드립니다. 또한 참고한 모든 블로그와 유튜브는 출처(Reference)에 있습니다.

 

Introduction

앞선 강의에서 우리는 data probability distribution $p(x)$를 어떻게 추정할 지를 중심으로 모델을 살펴보았다. $p(x)$를 나타내고자 한 모델로는 Autoregressive models, Flow models, Variational autoencoers, Energy-based models 등이 있다. 또는 sampling process에 중점을 두고 GAN과 같이디자인된 모델들도 있다.

 

이번 강의에서는 probability function의 gradient인 score function에 대해 알아보고자 한다. score function은 다음과 같이 정의된다.

Score function은 아래와 같은 그림들로 설명이 가능하며, probability density와 score의 관계는 electric field와 electric potential의 관계로 생각해볼 수 있다.

Recap on energy-based models

Energy-based models에 대해 다시 복습해보자. 우리의 neural net $f_{\theta}(x)$를 통해 $p_{\theta}(x)$를 아래와 같이 정의하며, 학습하기 위해 contrastive divergence algorithm을 사용한다. 하지만 contrastive divergence algorithm은 sampling을 하는데 MCMC 기법을 사용해야 하며, MCMC 기법은 하나의 샘플을 생성하는데 많은 시간이 걸린다고 알려져있다.

따라서 다른 방법으로 학습하기 위해 Score matching 기법이 등장하였으며, Score matching 기법은 score function을 이용해 다음과 같은 objective function을 optimize한다. 참고로 score function에는 더 이상 partition function $Z(\theta)$가 포함되지 않는다는 사실에 주목하자.

score function을 probability density의 gradient가 아닌 vector filed of gradient로 바로 정의한다면 Score-based model은 Energy-based model보다 좀 더 일반화된 모델임을 알 수 있다.

Score estimation

$p_{data}(x)$에서부터 i.i.d 샘플을 추출해 score function $s_{\theta}(x)$을 추정하는 과정을 score estimation이라고 부른다.

Score estimation을 위해 가장 먼저 생각해볼 수 있는건 Score matching을 이용하는 방법이다. Fisher divergence로 목적함수를 표현하면 다음과 같다.

식을 변형하면 다음과 같이 표현가능하다.

여기서 Jacobian의 trace term이 문제가 되는데, 이 term을 계산하기 위해서는 $x$의 차원에 비례하는 연산횟수가 필요하다. 높은 차원의 이미지 $x$를 다루는 생성형 모델 입장에서 학습하는데 시간이 오래 걸린다는 치명적인 단점을 가지고 있다. (이를 Score matching is not scalable이라고 강의에서는 표현한다.)

 

Jacobian의 trace term 을 계산하지 않고 우회하는 방법은 크게 두 가지가 있다. 첫번째는 Denoising Score matching 이고, 두번째는 Sliced Score estimation이다. 우선 Denoising Score matching 방법론부터 살펴보자.

 

Denoising Score matching(Vincent, 2011)

※ 해당 논문은 이 포스팅에 따로 정리했습니다.

noise를 추가한 $q_{\sigma}$ 분포를 생각해보자.($q_{\sigma}$는 marginal distribution, $q_{\sigma} (_|_)$는 perturbation kernel이라고 부른다)

만약 $\sigma$의 값이 작다면 $q_{\sigma}$와 $p$는 상당히 비슷한 분포를 가질 것이다. (작은 noise만을 추가해줬기 때문이다)

$q_{\sigma}$에 대해서 Fisher divergence를 생각하면 Jacobian의 trace term 없이 깔끔한 식을 얻을 수 있다. (증명은 더보기를 참고하세요!)

여기서 $q_{\sigma}(\hat{x} | x)$이 Gaussian 분포를 따르므로 constant term을 빼고 아래와 같이 쓸 수 있다.

증명)

위 식을 아래와 같이 고치면 $s_{\theta}(\hat{x})$이 추가된 noise를 측정하는 방향으로 학습된다는 점도 알 수 있다.

Denoising score matching(DSM)의 학습 과정을 정리하면 아래와 같다.

DSM의 장점) 

• Jacobian의 trace term이 없어 초고해상도 데이터에 대해서도 효율적으로 학습 가능

 

DSM의 단점) 

• clean data의 score를 추정하는 것이 아니기 때문에 불확실성 내포
• $\sigma$가 0으로 수렴할 경우 loss의 variance가 기하급수적으로 증가(증명은 더보기를 참고하세요!)

증명)

DSM을 Tweedie's formula의 관점에서도 생각할 수 있다. * 공부 필요

DSM은 noise를 포함한 data set에 대해 score matching을 진행하고 있기 때문에 부정확한 score estimation이 단점이다. 이러한 단점을 극복하기 위해 고안된 방법론이 Sliced Score estimation이다.

Sliced Score estimation

※ 해당 논문은 이 포스팅에 따로 정리했습니다.

Sliced Score estimation은 두 개의 벡터가 같다면 projection된 두 개의 벡터도 동일해야 한다는 아이디어에서 출발한다. 임의의 벡터 방향으로 투영된 score vectors를 비교하므로써 score model을 학습하는 방법으로, 2019년 Song et al.에 의해 제안되었다.

 

Sliced Score estimation에서의 objective는 sliced Fisher divergence라고 부른다. 벡터 $v$는 확률분포 $p_v$를 따르는 임의의 벡터로 sliced Fisher divergence 는 투영된 data의 score과 score 모델의 L2 norm을 측정한다.

다만 우리는 data의 score를 알지 못하므로 DSM에서 한 것처럼 Integreation by Parts를 진행하여 위 objective는 다음과 같은 식으로 변형해야 한다. 

위 식은 jacobian을 포함하고 있지만, Jacobian-vector product form으로 나타내져 있어 상당히 효율적으로 계산할 수 있다.(scalable but slightly slower than denoising score matching)

Computing Jacobian-vector products is scalable

 

위와 같은 과정을 거쳐 계산은 두 번의 backpropagation만을 필요로 한다. 

Sliced score matching(SSM)의 학습 과정을 정리하면 아래와 같다.

 

지금까지 우리는 Score estimation을 위한 두 가지 방법론(Denoising Score Matching & Sliced Score Matching)에 대해 공부해보았다. 이제 우리에게 적절한 score는 있다.

 

그렇다면 샘플링은 어떻게 해야할까?

 

 

From scores to samples

random으로 초기화된 많은 data points들이 주어져 있다고 가정하자. 우리는 이 random points들을 sample 분포와 최대한 비슷하게 움직이고 싶다. 가장 먼저 생각해 볼 수 있는 방법은 각 random points들을 score를 따라 움직여 주는 것이다. 이 방법은 효과적일 것이라고 생각하기 쉬우나, 많은 점들이 몇몇의 특정한 점으로 모여들기만 할 뿐 분포를 따리지는 않을 것이다. 이때 사용할 수 있는 방법론이 바로 Langevin MCMC이다.

 

Langevin MCMC

Langevin MCMC 는 random points들을 움직일 때 score에 gaussian noise를 추가해서 random points들을 움직인다. Langevin dynamics sampling 방법은 다음과 같다.

결론

참고로 Score-based generative model의 결과는 다음과 같다.(좋지 않다)

왜 score-based generative model은 잘 작동하지 않는가

(1) Manifold hypothesis

매니폴드 가설에 의하면 현실 세계의 data point가 데이터의 차원보다 훨씬 작은 저차원의 공간에 놓여있다. 예를 들어, data point는 3차원 공간에 있지만, 실제로는 2차원의 뫼비우스 띠 위에 존재할 수 있다.

 

데이터가 저차원 매니폴드에 국한되면 data point의 score는 정의되지 않는다. 아래의 그림과 같이 데이터 분포가 원형으로 집중되어 있는데 원이 점점 얇아진다면 score vector는 점점 커질 것이며, 원이 무한히 얇아진다면 score vector 역시 무한히 커질 것이다. 이것은 저차원 매니폴드 구조에서 score vector가 well-defined되지 않음을 의미한다.

(2) Challenge in low data density regions

매니폴드 가설도 영향이 있지만 score matching과 Langevin MCMC는 실제로 잘 작동하지 않는데, 그 이유는 데이터의 밀도가 낮은 영역에서 score estimation이 부정확하기 때문이다. 데이터는 확률이 큰 부분에 한정되어 분포하고 있으며, 이때문에 밀도가 낮은 영역에서는 score estimation이 크게 부정확하다. 

 

이러한 이유가 발생하는 근본적인 이유는 score matching objective가 data point와 score model의 차이만을 고려하기 때문이다.

 

(3) slow mixing of Langevin dynamics between data modes

아래와 같이 data distribution이 완전히 분리된 구성을 생각해보자. 각각의 가중치는 $\pi$, $1-\pi$라고 하자. 각각의 지점에서 얻은 score는 $\pi$에 의존하지 않는데, Langevin dynamics는 score만을 사용하기 때문에 $\pi$를 복구할 수 없다.

사실 Langevin dynamics에서는 연속한 distribution에 대해서만 작동하므로 위에서 만든 예제는 엄밀히 말하면 틀린 예제다. 그렇지만 분리된 구간이 아주 약간 연결되어 있다고 가정하면 Langevin dynamics는 모드의 가중치를 회복하는데 오랜 시간이 걸릴 것이다.(이것을 slow mixing이라고 표현한다.)

실제 샘플은 왼쪽과 같이 우측 상방에 샘플이 많은 반면, Langevin dynamics로 생성한 샘플은 좌측 하방과 우측 상방에 같은 양에 샘플을 생성한다.

해결방안

우선 (1)번 문제를 데이터에 noise를 추가하므로써 해결 가능하다. 노이즈를 추가 하지 않은 데이터의 score funciton은 매니폴드 가설에 의하면 발산한다. 하지만 데이터에 아주 작은 양의 노이즈를 추가하면 score function이 발산하지 않고 잘 정의되기 때문에 score-based model이 잘 작동할 수 있다.  

 

그럼에도 불구하고, (2)번 문제를 해결하는 것은 어렵다. 데이터의 품질(data quality)과 낮은 밀도의 샘플이 있는 구역(the size of low data-density region)은 일종의 trade-off 관계를 가지고 있기 때문이다. 또한 낮은 밀도의 샘플이 있는 구역이 넓을수록 score estimation의 정확도는 낮아지기 때문에 데이터의 품질(data quality)과  score estimation의 정확도(the accuracy of score estimation)도 역시 trade-off 관계를 가지고 있다.

 

아래 그림을 살펴보자.

가운데 그림과 오른쪽 그림에서의 붉은 부분이 에러를 의미한다. 위 그림은 노이즈를 추가하지 않았을 때의 score를 나타낸 것이다. 오른쪽 그림에서 낮은 밀도의 샘플들이 존재하는 공간은 에러가 큰 것을 알 수 있다. 

 

위와 같이 데이터에 점점 더 큰 노이즈를 추가할수록 낮은 밀도의 샘플들이 존재하는 공간이 줄어들면서 score estimation의 정확도가 올라가는 모습을 볼 수 있다. 하지만 perturbed data score는 original data score과 확연히 다른 것을 알 수 있는데 좋은 데이터 퀄리티와 score estimation, 이 두 마리의 토끼를 모두 잡기 위해 다양한 노이즈의 레벨을 고려하는 Multi-scale Noise Perturbation 방법이 고안되었다. 

 

Multi-scale Noise Perturbation

Multiscale Noise Perturbation에 녹아있는 직관은 다음과 같다.

 

만약 노이즈를 추가하지 않는다면 낮은 밀도의 샘플이 있는 구역에서는 score estimation의 정확도가 매우 낮을 것이며, 이는 Langevin dynamics를 수행함에 있어서 high data density region으로 가는 방향을 score가 제대로 제시하지 못한다는 의미이다. 하지만 너무 큰 노이즈를 추가하면, Langevin dynamics가 높은 데이터 밀도 영역에서 멀리 떨어져 있을 때는 좋은 방향 정보를 제공하지만, 높은 데이터 밀도 영역에 가까워질수록 세밀한 방향 정보를 제공받지 못한다. 

 

이러한 문제를 해결하기 위해 여러개의 노이즈를 사용해 여러 방면에서의 score function에 대한 정보를 수집해 high data density region으로의 방향성을 제공하고자 한다.  

 

Annealed Langevin Dynamics

Annealed Langevin Dynamics에서는 다양한 레벨의 noise로 만들어진 perturbed data distribution에 대해 Langevin dynamics를 실행한다. 처음에는 가장 perturbed된 데이터를 통해 Langevin dynamic를 실행하고, 그 결과물을 initial sample로 하여 그 다음 큰 노이즈에 의해 perturbed된 데이터로  Langevin dynamic를 실행한다. 이러한 일련의 과정을 통해 마지막으로 가장 덜 perturbed된 데이터 분포로부터 샘플을 생성할 수 있다. (강의 ppt 14쪽 애니메이션 참고)

Noise conditional Score Networks

Multi-scale Noise perturbation을 수행하기 위해 각각의 perturbed data distribution에서 샘플을 만들고, score estimtation을 진행해야 한다. 이 때 각각의 score model을 사용해 학습하게 되면 상당히 비싸므로, 노이즈를 하나의 input으로 선택하고 conditional score network를 사용하는 방법을 고안했다. 이 모델을 우리는 Noise Conditional Score Network(NCSN)이라고 부른다.

NCSN에서 score mathcing loss는 denoising score matching loss를 가져와 가중평균 으로 아래와 같이 표현할 수 있다.

여기서 $\lambda$는 $\lambda({\sigma}_i) = {{\sigma}_i}^2$으로 정의하여 아래와 같이 식을 전개할 수 있다.

 

(1) 노이즈의 범위 설정

노이즈의 범위는 최대 노이즈(${\sigma}_1$)이 각각의 데이터 쌍의 최대 거리를, 최소 노이즈(${\sigma}_L$)이 sufficiently small한 임의의 숫자로 설정한다. 이 때 ${\sigma}_1$을 두 데이터 쌍의 최대 거리로 설정하므로써, Langevin chain이 하나의 데이터 포인트를 다른 데이터 포인트로 용이하게 옮기는 것을 가능하게 하여 다양한 샘플 생성을 돕도록 한다.

 

(2) 노이즈의 스케일 설정

노이즈의 범위 안에서 나머지 노이즈들을 설정할 때 각각의 노이즈들이 geometric sequence가 되도록 설정하는 것을 권장한다. 이렇게 설정하므로써 노이즈들이 충분이 밀집되고, 인접한 노이즈들이 서로에게 충분히 가까움을 보장할 수 있다. 

NCSN의 학습과정은 아래와 같다.

정리본

NCSN에서는 naive Score matching + Langevin MCMC를 사용했을때와 다르게 상당히 좋은 성능을 보임을 알 수 있다.

 

Understanding Score-baed model via SDEs

이번엔 score-based generative model을 SDE 관점에서 바라보자.

 

왼쪽 그림은 4개의 명작을 2 *2 모양으로 묶은 것으로 위와 같은 그림에 노이즈를 추가해 마지막에는 Gaussian noise가 되게끔 노이즈를 계속 추가하였다. 이러한 과정은 시간 $t$가 지남에 따라 $X_t$에 노이즈가 추가되는 형태인 일종의 SDE로 치환해 이해할 수 있다.

 

우선 여기서 노이즈만을 추가했으므로 아래와 같이 drift term 없이 diffusion term만 있는 상황을 생각해보자.

Forward SDE : $dX_t = \sigma(t) dW_t$

 

앞에서의 과정을 시간만 반대로 흐르게 진행하면 노이즈에서 샘플을 생성하게끔 된다.

그리고 우리가 앞서 정의한 Forward SDE : $dX_t = \sigma(t) dW_t$는 다음과 같이 쓸 수 있다.

위 식이 의미하는 바는 reverse-time SDE를 풀면 샘플을 생성할 수 있다는 의미이다!

 

SDE Score-based generative model(NCSN++)

NCSN에서는 score function이 noise의 scale인 $\sigma$의 함수였다는 사실을 상기시켜보자. Brownian motion에서의 분산은 누구인가? 바로 시간 $t$이다! 따라서 우리는 이제 score function이 시간 $t$에 의존하는 함수임을 알 수 있다. 고로 score-based model 역시 아래와 같이 쓸 수 있다.

또한 objective 역시 time-dependent score-based model을 대입해 아래와 같이 나타낼 수 있다.

Reverse SDE를 풀기 위해서는 다양한 방법이 있지만 가장 쉬운 Euler-Maruyama approach를 생각해볼 수 있다. 이는 ODE에서의 Euler scheme과 매우 유사한 방법론으로, 매우 작은 시간 $dt$를 $\triangle t$로, $dW_t$를 $z ~ N(0, |\triangle t| I)$로 나타내 SDE의 수치적 계산을 가능하게 한다.

* 여기서  $\triangle t$는 reverse-time SDE이므로 음수이다.

위와 같은 방법으로 학습한 score-based model은 이제 time-dependent하다!

 

그렇다면 이번엔 샘플링을 어떻게 할 지 생각해보자. SDE Score-based generative model에서는 predictor-corrector samplers 기법을 선택하였다. 이름에서 알 수 있듯, 이 방법론은 predictor와 corrector가 있으며, predictor는 일반적으로 잘 알려져 있는 SDE Solver(ex : Euler-Maruyama method)를, corrector는 score-based MCMC approach(ex : Langevin dynamics)등을 선택할 수 있다.

predictor-corrector samplers 기법에서는 reverse-time SDE를 따라 먼저 predictor가 rough하게 예측하고, 이를 corrector가 fine-tuning하는 식으로 각각의 스텝이 전개된다. 시간이 지남에 따라(0에 가까워짐에 따라) noise가 점점 작아지므로 sample이 unperturbed data distribution으로부터 얻어질 수 있다.

 

이 방법은 annealed Langevin dynamics와 닮았지만, predictor가 먼저 rough하게라도 예측하므로써 Langevin dynamics가 각각의 노이즈들에 대해 smooth transition을 하게끔 도와준다.

 

결과는 상당히 성공적이며, 처음으로 score-based model에서 고해상도 이미지를 성공적으로 생성해낸다.

Converting the SDE to an ODE

Fokker-Planck equation을 이용하면 SDE를 ODE로 변경할 수 있다.(어쩌다보니 Fokker-Planck equation도 포스팅해버렸다)

위 식을 reverse-time SDE로 생각하고, ODE를 change of variable formula를 통해 풀면 다음과 같은 식을 얻을 수 있는데, 이는 일종의 continuous-time normalizing flow model이 된다. (FFJORD)

DDIM