[Generative Artificial Intelligence] Score Based Models – 2

Stanford CS236 Deep Generative Models 수업의 자료를 기반으로 생성모델의 기본 개념들을 정리해보고자 한다. (참고 https://deepgenerativemodels.github.io/syllabus.html)

지난 글에서 Score Based Model을 학습시키고, 학습시킨 모델로 생성하는 법까지 다뤘다. 하지만 그대로 적용하면 제대로 생성이 안되고 해결해야할 문제 3가지를 남겨두었다. 이번 lecture에서는 그 문제를 어떻게 해결하여 훌륭한 생성모델로 만들 수 있었는지 다룰 예정이다.

해결해야할 문제는 아래의 3가지였다.

1. 실제 데이터가 저차원에 있을때 score가 정의되지 않는 문제
2. data density가 낮은 지역에서는 score가 제대로 matching되지 않는 문제
3. langevin dynamic이 converge가 매우 느리다는 문제.

사실 위 3가지 문제 모두 크게보면 하나의 solution에 의해 해결된다. 바로 Gaussian perturbation 방법이다. 이미 이전에 denoising score matching에서 나온 방법으로 이미지에 gaussian noise를 더해 noisy한 data의 score를 배우는것이다. 그러면 가우시안 노이즈가 전 차원에 골고루 더해지기 때문에 manifold hypothesis 문제는 해결이 된다.

이때 더해지는 노이즈가 매우 작아서(std=0.0001) 사람이 보기에 보이지 않는 수준이여도 된다고 한다. 아래 왼쪽처럼 일반 CIFAR-10에서도 loss가 발산했었는데 노이즈를 약간만 더해도 오른쪽처럼 잘 converge했다고 한다.

그런데 이렇게 naive한 방법으로는 2번째 문제가 해결이 안된다. low data density 문제인데 아래 그림처럼 data density가 특정 지역에 있을때, data가 많이 존재하는 지역에서는 data score와 estimated score가 매우 정확하게 추정되지만, 데이터가 density가 낮은 지역에서는 score가 실제 data와 다르게 부정확하게 추정된다.

해결법을 설명하기에 앞서 먼저 가우시안 노이즈를 더하는것과 score를 예측하는것의 상관관계를 이해해야한다. 노이즈를 데이터에 많이 더하면 더할수록 데이터가 존재하는 범위는 넓어지게 된다. 하지만 당연히 이미지에 노이즈가 심해지면 실제 clean한 이미지의 진짜 data distribution과는 달라지게 된다. 즉, 노이즈를 많이 더하면, 원래 데이터가 없던 지역에서의 score는 더 정확하게 예측이 가능해지지만, 최종 data 생성 결과는 퀄리티가 떨어지게 된다.

아래 그림을 보면 좀 더 디테일하게 trade off를 알 수 있다. 아래 그림에서 빨간 부분이 에러인데 노이즈가 데이터에 더해지면 더해질수록 실제 data score에는 에러가 증가하지만 estimated score는 정확해진다. 노이즈를 안더하면 실제 data score는 정확하겠지만 데이터가 부족한곳에서 estimated score는 부정확해진다.

이러한 성질을 이용하여, multi-scale noise perturbation이라는 방법으로 이 문제를 해결할 수 있게 된다. 다양한 scale로 노이즈를 더해서 score estimation 학습에 사용하는 것이다. 아래 그림을 보면 조금 더 직관적으로 이해가 될것 같은데 아래 그림에서 검은선이 실제 data density가 높은 지역이라고 했을때 원래는 왼쪽처럼 data가 적은 지역에서는 어느 방향으로 데이터를 보내줘야할지 모른다. 이때 noise를 더해줌으로써 데이터가 적은 지역에서도 score 방향을 알 수가 있게 된다. 하지만 이때도 가운데 그림처럼 특정 지역에 있는 데이터는 어느방향으로 가야할지 모르는데, 다양한 scale의 노이즈로부터 배우면서 오른쪽 그림처럼 데이터가 어디 있어도 실제 data density가 높은 지역의 방향을 알 수 있게되는 것이다.

이제 이렇게 다양한 scale의 노이즈로 학습된 score estimation 네트워크로 생성을 할때 이전과 다르게 Annealed Langevin Dynamic 방법을 사용하게 된다. 노이즈를 더하고 score estimation하는 기존 Langevin Dynamic 방법인데 더하는 노이즈를 schedule을 정해서 처음엔 크게 더하다가 점점 작게 더한다. 이러면 데이터가 적은 부분의 score도 정확하게 estimation하면서도 data quality에 손상 없이 생성을 할 수 있게 된다.

이렇게 noise 레벨을 다르게 더해서 estimation하는데 각각의 noise 레벨에 따라 네트워크를 모두 각각 학습시키면 코스트가 너무 커지기 때문에 아래 그림의 오른쪽처럼 Noise Conditional Score Network (NCSN)을 만들어서 noise 레벨은 condition으로 주고 네트워크 하나로 모든 noise 레벨에 대해 estimation을 할 수 있는 네트워크를 만들어 학습시켜 사용한다.

노이즈 레벨이 condition으로 들어가서 바뀐 식이다. 기존 denoising score matching loss식에서 z라는 Normal distribution 노이즈를 \( \sigma \)로 스케일을 조절해서 data x에 더하여 입력으로 사용하는 식이다. 앞에 weight function으로 \( \lambda(\sigma_i) \)도 사용되는데 이에 대한 설명은 뒤에서 얘기할 예정이다.

이제 이 noise scale을 어떻게 조절할지에 대해 여러 팁들도 설명해준다.

먼저 maximum scale과 가장 minimum scale을 정하는 법에 대해 설명해주는데, 가장 큰 scale은 datapoint들 중에서 maximum pairwise distance 정도로 설정하면 langevin chain을 적용할때 특정 데이터 포인트에서 다른 포인트로 이동할때 쉽게 이동할 수 있게 도와준다고 한다. 그리고 minimum은충분히 작아야 최종 sample 퀄리티가 떨어지지 않는다고 한다.

그리고 scheduling은 서로 scale끼리 겹치는 부분이 충분히 있어야 좋다고 하고 아래 나와 있듯이 같은 비율로 scale이 줄어들도록 만든다고 한다.

그리고 앞에서 얘기했던 loss 식에 들어가는 weight function \( \lambda(\sigma_i) \)은 \( \sigma_i^2\) 로 설정한다. 노이즈가 크면 클수록 loss에 가중치가 크게 들어가는 개념이다. 그렇게 넣으면 아래의 최종식처럼 바뀌는데 원래 score를 예측하던 네트워크에 \( \sigma \)를 곱하면서 \( \epsilon \) 으로 바꿔 표현하면 \( \epsilon \)은 노이즈 z를 예측하는 네트워크가 된다.

최종적으로 이 noise conditional score network를 학습시키는 방법은 아래와 같다. 먼저 training data에서 sample을 뽑고 설정된 노이즈 스케줄 중에서 랜덤으로 스케줄을 정해서 training data에 더한뒤 network에 노이즈 스케줄 정보와 함께 입력으로 넣어준다 그리고 아래와 같은 loss function으로 학습을 시키면 score estimation network가 학습된다.

이후 annealed langevin dynamic 방법으로 데이터생성을 하면 굉장히 훌륭한 퀄리티의 이미지들을 만들어낼 수 있다.

지금까지 방식은 노이즈 schedule을 미리 정해서 discrete하게 더해가면서 학습시키는데 이러한 noise schedule을 continuous하고 inifinite한 level로 schedule을 할수도 있다. 아래 보면 t=0일때 분포가 일반 데이터 분포인데 t가 증가하면서 가우시안 perturbation이 가해지면서 t=T일때 분포는 가우시안 분포로 변한다.

이러한 과정을 수식적으로 stochastic differential equation (SDE)로 풀어낼 수가 있다.

forward과정은 t=0 –> T로 노이즈를 더해주는 과정이고

reverse 과정은 t=T –> 0으로 노이즈에서 데이터분포로 변환되는 과정이다.

이를 식으로 표현하면 아래와 같이 된다.

forward과정은 단순히 노이즈를 더해주는 과정이라 단순한데 reverse과정은 위와 같이 score function과 noise가 더해지는 과정이 같이 들어가 있다. score estimation 네트워크가 학습되어있으면 위와 같은 process로 가우시안 분포에서 시작해서 실제 데이터 분포로 데이터를 reverse SDE과정으로 변환시킬수가 있게 된다.

여기서 Score estimation 네트워크를 학습시키는 방법은 아래와 같다. 이전에 \( \sigma \)로 레벨을 정해서 더해주던 노이즈가 t에 대한 식으로 바뀐것뿐이다. 학습시킨 score estimation 네트워크를 적용해 reverse SDE에 적용해서 가우시안 분포에서 시작해서 이미지가 학습된 분포의 이미지로 변환되는것이다.

위의 SDE방식은 계속 노이즈가 더해지기때문에 trajectory를 보면 일직선으로 부드럽게 이동하지 않고 노이지하게 움직인다. 여기서 노이즈를 더해주는 부분을 빼면 SDE가 아니라 ODE (Ordinary Differntial Equation)으로 풀수가 있게 된다. 아래 그림에서 흰색 trajectory가 ODE에 의한 trajectory인데 부드럽게 변환되는것을 볼 수 있다.

그리고 ODE는 노이즈 분포에서 데이터분포로의 변환을 아래의 식으로 풀 수 있다. 아래 식을 보면 SDE와 다르게 노이즈가 더해지지 않기 때문에 Score estimation network 모델이 한번 학습되면 같은 노이즈에서 변환되는 데이터가 항상 같은 결과가 나온다.

그래서 SDE와 ODE는 약간 다른 특성을 지니는데, SDE는 노이즈가 더해지는 과정덕분에 생성되는 결과에 variation이 생긴다. 또한 생성과정중에 약간의 에러가 발생해도 노이즈가 더해지면서 수정이 된다. 하지만 생성하는데 필요한 계산량이 많다. ODE는 노이즈가 더해지지 않기 때문에 생성과정에 variation이 없어진다. 따라서 다양한 데이터를 생성하기는 어려워진다. 하지만 SDE에 비해 연산량이 적고 더 정확한 데이터 생성을 할 수 있다.

DDIM이라는 방식은 위의 방식들을 조금 조합해서 time step을 크게 나누고 각 구간에서 integrator (ODE solver)로 ODE를 풀어서 sampling을 하면 훨씬 적은 timestep으로 생성을 할 수 있게 되고 성능이 매우 훌륭하다고 한다.

마지막으로 text나 다른 condition을 써서 생성을 컨트롤하는 방법이다. 기존의 score를 측정하는 식에서 condition으로 y가 들어오고 이를 오른쪽 아래의 식처럼 bayes’ rule을 써서 풀면 기존 score estimation모델과 추가로 forward모델이 있으면 최종 score를 예측할수 있게 된다. forward model같은 경우에는 아래같은 경우에는 classifcation같은 네트워크 (이미지가 주어졌을때 class예측)을 사용할수 있고 그 아래 예시같은 경우에는 이미지 captioning 모델 등을 사용할 수 있다.