핫한 Text to Image 생성 모델 - Stable Diffusion

안녕하세요 콥스랩(COBS LAB)입니다.

오늘 소개해 드릴 논문은 ‘Stable Diffusion’입니다.

해당 내용은 유튜브 ‘딥러닝 논문 읽기 모임' Stable Diffusion’ 영상 스크립트를 편집한 내용으로, 영상으로도 확인하실 수 있습니다. (영상링크:https://youtu.be/7fBQDaJkcSU)

 

최근에 Text to Image 생성 모델이 많은 관심을 받고 있습니다. DALL-E나 Image GAN 같은 모델들이 있고 그중에서 가장 핫한 Stable Diffusion에 대해서 리뷰를 시작해 보도록 하겠습니다.

Stable Diffusion 같은 경우에는 다른 모델들과 다르게 컴퓨터 사용 리소스가 대폭 줄어서 4GB 이하의 브이랩에서도 돌릴 수 있습니다. 따라서 현재 API들이 굉장히 많이 나오고 있는 모델이라고 할 수가 있습니다.

 

텍스트 생성 샘플들을 몇 개 있습니다. 예를 들어서 현재 A zombie in the style of Picasso라는 문장에서 이미지를 생성해 내는 예시입니다. 주어진 문장에 맞게 사람이 이해해도 Semantic 하게 납득이 가는 이미지를 잘 생성해 낼 것을 확인을 할 수가 있습니다.

 

먼저 Stable Diffusion에 들어가기 전에 그전에 있었던 Auto Encoder나 VAE의 개념을 조금 알고 가야 됩니다.

 

VAE 같은 경우에는 이미지와 같은 고차원 데이터를 encoding을 통해서 저 차원 space로 맵핑을 시킨 다음에 그것들을 다시 decoder를 이용해서 생성을 시켜내는 구조라고 할 수 있습니다. Encoder를 이용해서 이미지에 고차원을 저 차원으로 줄이는 개념은 사실상 주성분 분석과 같은 PCA와 같은 성분 분해의 개념과 굉장히 유사하다고 볼 수가 있습니다. 따라서 이러한 의미 정보 담는 차원의 공간을 latent space에 들어가는 변수들을 latent variable이라고 부르게 됐습니다.

 

decoder 같은 경우에는 latent variable을 받아서 다시 이미지를 복원시키는 역할을 하고 있습니다.

 

사용을 할 때 정규 분포 방식을 이용해서 z를 샘플링하는 것이 아니라 샘플링하는 과정에서 noise를 더해줘서 미분 가능한 함수를 바꿔준 후에 그 함수를 학습시키는 방식으로 유도를 해냈습니다.

 

Auto Encoder의 objective를 유도해 낸 공식입니다. 수식의 유도에 대해서는 너무 깊게 들어가지 않고 의미론적인 차원에서 검토하는 방식으로 진행을 하도록 하겠습니다.

 

Auto Encoder에서 궁금한 적은 결국은 X를 만들어내는 p와 같은 함수를 찾아내는 것입니다. 앞에서 보셨다시피 Z로 한번 압축을 한 다음에 이 Z를 전제로 했을 때 X를 만들어내는 조건부 확률을 최대화하는 방식으로 학습을 진행을 하게 됩니다.

 

다시 말해서 최종적인 VAE 로스 같은 경우 objective function 같은 경우에는 Z를 베이스로 했을 때 X를 만들어내는 함수를 유도해 내는 것이 그 목표라고 할 수가 있습니다. 따라서 Z로부터 X가 나타날 확률을 최대화하는 것으로 Maximum Likelihood Estimation을 하게 되면서 이를 cross entropy와 같은 방식으로 표현을 하게 됩니다.

 

오른쪽에 있는 KL divergence 공식 같은 경우에는 latent space를 만들 때 함수를 이미 알고 있고 그리고 X에서 Z를 만들 함수를 알고 있으므로 두 함수의 분포차는 최소한의 바운드가 되는 lower bound가 되겠습니다.

 

그래서 Regularization 텀으로 이 공식을 사용을 한 것을 확인을 할 수가 있습니다.

 

다음은 본격적으로 Diffusion 모델의 process에 대해서 설명을 하도록 하겠습니다.

 

앞서 Auto Encoder가 noise를 한 번에 더하는 방식으로 noise를 더 했다면 Diffusion 같은 경우에는 noise를 더 하는 과정 자체를 스텝을 몇 번을 나눠서 하게 됩니다.

 

보시는 바와 같이 각각의 스텝이 하나의 함수가 되어서 이 스텝이 진행이 될 때마다 원본의 이미지가 좀 더 Gaussian 분포에 가까워지는 과정을 Diffusion process라고 부르게 됩니다. 따라서 함수 식을 보시게 되면 t-1 시점에 X 가 t 시점으로 변화할 수 있는 forward 함수를 Markov chain으로 형성한 것이 전체적인 process의 정의라고 볼 수가 있습니다. 이를 Markov Diffusion kernel이라고 부르게 됩니다.

 

그리고 여기서 β라는 파라미터를 사용하게 되는데 이 β t 파라미터 같은 경우에는 t0 시점부터 t시점까지의 변환을 몇 번이나 길게 가져가게 되느냐, 스텝을 얼마나 촘촘하게 가져가느냐를 결정해 주는 파라미터라고 볼 수가 있습니다.

 

앞에서 process 했던 함수들은 이미 알고 있는 함수기 때문에 궁금한 함수는 그 함수들의 inverse 함수에 초점을 맞추고 그 학습을 진행을 하게 되고 다시 말해서 아까의 그림에서는 t-1 시점에서 t 시점을 추론했다면 inverse 함수는 반대로 t 시점에서 t-1 시점을 추론을 하게 됩니다.

 

보시는 그림 같은 경우에는 앞이 noise 이미지고 이 noise 이미지에서 다시 원본 이미지를 만들어내는 함수를 추론해 내는 것이 목표라고 할 수가 있습니다. 따라서 여기서는 파라미터로 β가 주어지면 됩니다. 이 β같은 경우에는 다시 말해서 noise를 주입하는 스텝이 매우 작다면 그 reverse function 또한 비슷한 분포를 따를 수 있다는 전제를 깔고 가는 게 됩니다. 따라서 이러한 전제를 깔고 가게 되면이 β t를 결정하는 것이 reverse process를 얼마나 잘 추론해 낼 수 있느냐에 대해서 중요한 역할을 한다라고 합니다.

 

앞에서 얘기드렸던 noise에서 시 원본 이미지를 추론해 내는 과정 Xt에서 t-1 시점을 추론해 내는 함수가 목표라고 할 수 있습니다.

 

해당 그림에서는 생성 모델로 유명한 GAN과 Auto Encoder 그리고 Diffusion 모델의 전체적인 아키텍처를 정리해 놓은 표입니다.

 

Auto Encoder 같은 경우에는 원본이미지가 들어왔을 때 Encoder를 통해서 latent variable로 표현해 내고 latent variable을 다시 decoder를 통해서 이미지를 추론해 내는 전체적인 과정을 겪게 되고 여기서 Encoder에서 decoder까지 이르는 과정을 Diffusion 모델들은 스텝을 나눠서 진행을 하게 됩니다.

 

Stable Diffusion는 아키텍처 중심의 논문이라고 이해를 했습니다.

 

다시 말해서 VAE의 Auto Encoder noise 과정과 Diffusion 모델의 noise 과정을 섞은 아키텍처를 가지고 있습니다. 지금 보시는 바와 같이 X Input값에서 Encoder를 해서 latent variable를 만들어내는 것까지는 Auto Encoder의 과정과 유사하다고 할 수 있습니다. latent vector상에서 모델이 했던 대로 forward process를 통해 가지고 noise들을 주입을 하게 됐습니다. 그 주입을 다시 decoder로 풀어내는 과정은 Auto Encoder와 동일합니다.

 

더 자세한 아키텍처의 내용은 뒤에서 다루도록 하겠습니다.

 

논문에 실려 있는 아키텍처의 모양이 있습니다. 앞서 설명한 것과 같이 pixel space 상에 encoding을 해서 latent vector들을 뽑아냈을 때 Diffusion process 그리고 뽑아낸 후에 latent variable들을 다시 Denoising 해내는 decoder로 옮기는 과정은 다시 스텝바이스텝으로 나눠줘서 각각의 스텝은 UNet구조의 함수로 이루어지게 됩니다.

 

이러한 전체적인 아키텍처를 가지고 생성 모델의 역할을 해야 되는 것을 볼 수가 있습니다. 그래서 보시는 Conditioning 같은 경우에는 예를 들어서 텍스트나 이미지나 다른 도메인의 데이터들이 들어왔을 때 데이터들을 각각의 스텝에서 attention 메커니즘을 통해서 학습을 하게 됩니다.

 

먼저 Stable Diffusion의 KEY POINT입니다.

 

Denosing 과정에서 Auto Encoder를 사용을 하게 되고 pixel 공간이 아닌 latent space에서 Denosing을 진행하면서 엄청난 양의 GPU나 computing cost가 들어갔던 포인트들을 개선을 시켰고, 모델의 복잡성들을 감소시켰고 cross attention을 사용함으로써 다른 도메인의 모델들을 함께 사용 가능하게 된 점이 있습니다.

 

논문의 introduction입니다.

 

Diffusion 모델들이 컴퓨팅 코스트를 많이 필요로 해서 그러한 것들을 개선하는데 포인트를 맞췄습니다. 예를 들어서 가장 많은 연산을 필요로 하는 Diffusion 모델들 중에 하나를 골라보면 거의 몇백대의 V100을 이용해서 수일간 학습시켜야 되는 케이스도 있고 이 케이스 같은 경우에는 50000개의 샘플들을 A100으로 추론해도 5일씩 걸리는 학습과 추론에 너무 많은 컴퓨팅 코스트를 필요로 하는 단점들이 존재를 했습니다.

그림은 다른 생성 모델들과 이미지 합성 그러니까 이미지 복원 task를 했을 때 성적을 보여주고 있습니다.

 

보시는 그림이 원본 이미지고 각각의 패치가 생성을 해낸 결과라고 해석을 할 수가 있습니다. DALL-E가 생성해 내는 이미지, VQGAN이 생성해 낸 이미지, 그리고 Stable Diffusion이 생성해 낸 이미지입니다.

 

여기서는 FID 매트릭을 사용을 하고 있습니다. FID 매트릭은 뒤에서 설명을 할 텐데 낮을수록 좋은 매트릭이라고 해석을 하시면 될 것 같습니다. 그리고 논문에서는 less aggressive 한 down sampling 다시 말해서 이미지에 크기를 크게 down sampling 하지 않고도 효율적으로 task를 수행해 냈다, 좋은 성적으로 task를 수행해 내었다는 점을 계속해서 강조를 하고 있습니다.

 

앞에서 말씀드린 평가 매트릭인 FID 같은 경우에는 Inception V3 모델을 feature extractor를 사용을 해서 실제 이미지와 생성된 이미지 사이에서 Activation Map들을 모두 추출해 냈습니다. 이 추출해 내는 전체 Activation Map에 대해서 분포를 만들게 되고 이 분포를 Wassertein-2 distance를 이용을 해서 추정을 한 값이라고 합니다.

 

논문의 process 파트입니다.

 

논문에서는 Auto Encoder를 먼저 학습을 시키게 되고 Auto Encoder를 학습을 시키고 process를 쭉 하면서 Auto Encoder 외에 스텝마다 UNet 아키텍처를 또 학습을 시키게 되는 구조라고 보시면 됩니다. 따라서 이런 모델들을 latent Diffusion 모델, 다시 말해서 latent 공간에서 학습을 진행하는 Diffusion 모델이라고 정의를 하겠다고 밝히고 있습니다.

 

앞에서 사용했던 Encoder 같은 경우에는 마치 Trasnfer learning을 하듯이 정확히 동치 되는 개념이라고 보기엔 조금 힘들지만 Transfer learning 하듯이 multiple 한 다른 task에 반복적으로 사용을 할 수가 있다고 합니다. 재활용을 할 수 있다고 하고 따라서 뒤편에 나오는 appendix들을 보면 여기서 보이시는 다양한 task들의 해당 Encoder를 써서 진행을 했기 때문에 빠른 속도로 많은 task들의 실험을 해볼 수 있었다고 말을 하고 있습니다.

 

그래서 task들로는 unconditional image synthesis. 다시 말해서 이미지 복원입니다.

 

inpainting 같은 경우에는 이미지를 가렸을 때 다시 복원시키는 super resolution task에 적용을 시켰고 computational cost를 많이 낮추었다고 표현을 하고 있고 다른 도메인의 embedding들에 대해서 noise 이미지와 함께 cross attention을 사용을 하였다고 합니다.

 

논문에서 키포인트로 잡았던 latent space로 옮겨서 연산을 한다라는 포인트는 결국은 pixel space 보다 low dimension 한 space의 영상이 컴퓨팅 코스트를 줄이는데 훨씬 유리하다는 점이랑 모델의 추론능력 상으로 봐서 Semantic 한 포인트들에 좀 더 집중할 수 있다는 것을 내세웠습니다. UNet을 사용한 것은 UNet 구조가 제공하는 inductive bias를 사용할 수 있기 때문에 UNet을 사용했다고 밝히고 있습니다.

 

모델에서 사용을 했던 attention 메커니즘을 좀 더 자세하게 살펴보겠습니다.

 

zT가 생각한 noise가 오염된 이미지라고 생각을 하시게 되면 되고 이게 query값으로 들어가게 되고 입력을 Input을 넣고 싶은 text embedding을 key 값으로 가져가게 됐습니다.

 

따라서 두 개의 가중치를 dot product 해서 계산을 한 다음에 softmax 값으로 가중치 형태로 끌어내고 이것들을 다시 텍스트 이미지에 dot product 하는 전형적인 cross attention 메커니즘을 사용을 했다고 합니다.

 

Loss는 두 가지를 사용을 했습니다.

 

perceptual loss와 patch based adversarial objective를 사용을 했습니다.

 

perceptual loss 같은 경우에는 네트워크상에서 feature map이 나오게 되는데 이 feature map마다 거리 계산을 했다고 보시면 될 것 같습니다.

 

patch based adversarial objective는 전체적인 이미지를 한 번에 맞냐 틀렸나를 통해서 스코어를 측정하는 것이 아니라 패치단위로 True/False를 판별하는 방식이라고 합니다. 이 방식을 사용하게 되면 사실성을 살릴 수가 있습니다. 그러니까 L1이나 L2 로스처럼 pixel 단위 로스를 사용했을 때 나타날 수 있는 Blurriness 현상을 많이 완화하고 있다고 합니다.

 

다음은 실험 결과입니다. 사용을 했던 데이터셋들을 보시면은 CelebAHQ나 FFHQ Churches, Beds, ImageNet 등의 데이터셋을 사용을 했습니다.

 

미국의 할리우드 셀럽들을 모아놓은 사진이고 고해상도 이미지 사진으로 많이 샘플로 쓰이는 사진입니다.

 

 표는 CelebAHQ입니다. 위 사진을 복원 task를 했을 때 unconditional 복원이라고 하면 이미지 합성이라고 하면 그냥 복원입니다. 복원 테스트를 했을 때 FID 기준으로 SOTA을 달성했다고 하고 있습니다. 256 x 256 기준입니다.

 

다음에 Conditional latent Diffusion입니다.

 

Text to 이미지 합성에 대해서는 지금 FID상으로도 지금 latent 모델이 괜찮은 스코어를 보여주고 있는데 여기서 말하는 포인트는 결국 파라미터 개수입니다. 거의 40억 개에 달하는, 60억 개에 달하는 파라미터들을 10억 개로 줄였다는 것이 그 포인트라고 할 수 있습니다.

 

그래서 inference 속도에 의해서 되게 유리하게 됐다는 장점을 가지고 가는 것 같습니다.

 

다음 task들은 샘플로 가져와 본 task들입니다. 이것들은 여기서 지금 Semantic 한 박스의 레이블을 주게 되면 공간에 이런 식의 이미지를 생성을 해준다고 합니다.

 

다음에 inpainting task입니다. ground truth가 있고 Input이 있을 때 이것들을 이렇게 대칭적으로 학습시키게 되면 Input값이 다 비게 들어와도 빈칸들을 메꿔주는 형태의 task입니다.

 

이 task 같은 경우에는 배경에 대한 Semantic Segmentation를 가져가고 Segmentation 클래스를 바꿔주는 것들로 이미지 합성을 해낸 케이스입니다.

 

이상으로 latent Stable Diffusion 모델 발표를 마치도록 하겠습니다.