[Stable Diffusion] 1. 모든 것의 시작: DDPM과 UNet

그림 1. 이렇게 높은 퀄리티의 이미지를 AI가 그렸다고 하는 건 이제 놀랍지도 않은 일이 됐습니다. (출처: Stable Diffusion)

 

Introduction

    Stable Diffusion은 사실 2022년 8월에 출시되어, 2023년 3월 기준 이미 건드려 볼 사람들은 다 건드려 본 툴입니다. 하지만 여전히 잘 만들어진 web UI나 MidJourney와 같은 앱을 통해서 프롬프트 엔지니어링을 하시는 분들이 대부분일 것이고, 그 기능들의 작동 원리까지 이해하시고 사용하시는 분은 그렇게 많지 않을 것입니다. 사실 작동 원리를 모르더라도 어떻게 프롬프트를 잘 쓰고, 체크포인트를 잘 섞고, 적절한 컨디셔닝을 잘 주는지만 알면 좋은 결과물을 뽑을 수 있기 때문에 이론적인 배경은 크게 중요하지 않을 수도 있습니다.

 

    그럼에도 각 기능의 작동 원리를 충분히 이해하고 있다면 자신이 원하는 결과물을 얻기 위해 주먹구구식으로 노가다하기보다 어느 정도 추론을 해가며 효율적인 창작 활동(?)을 할 수 있을 것입니다. 또한 새로운 기술이 나왔을 때 자신에게 필요한 기술인지 아닌지 빠르게 캐치할 수도 있을 것이고, PyTorch 코딩에 자신 있다면 아예 새로운 기능을 직접 구현할 수도 있겠지요. 그리고 이런 동기들이 없더라도, 기왕 쓰는 거 제대로 알고 써서 나쁠 것 없지 않을까요?

 

    이런 식으로 더 많은 사람들이 Stable Diffusion의 이론적 배경을 이해하고 활용함으로써 이 분야에 조금이라도 더 빨리 발전했으면 좋겠다는 생각에 뒤늦게나마 Stable Diffusion과 관련된 기술들의 이론적 배경을 설명하는 몇 개의 포스트를 준비했습니다. 첫 번째 포스트는 모든 것의 뿌리가 되는 denoising diffusion probabilistic models(DDPM)의 작동 원리에 관한 내용입니다. 이 블로그의 다른 논문 리뷰 포스트처럼 논문에 있는 수식이나 결과물을 하나하나 분석하지는 않고 정말 배경 지식으로 활용할 수 있을 정도의 깊이로만 짚고 넘어갈 예정이니 가볍게 읽으셔도 될 것 같습니다.

 

Denoising Diffusion Probabilistic Models

    딥러닝 기반의 생성모델로는 variational autoencoder, flow-based model, 그리고 그 유명한 GAN 등이 있었습니다. 제각기 다른 방법으로 데이터 생성 과정을 모델링했지만, 근본적인 원리는 동일합니다. 바로 데이터가 가질 수 있는 공간 내에서 실제 데이터들이 존재하는 영역(manifold)을 통계적으로 추정하는 것입니다.

 

    이미지의 경우 데이터가 가질 수 있는 공간이란 가능한 모든 픽셀 조합의 개수라고 할 수 있습니다. $256\times256$ 크기의 RGB 이미지로 한정하자면,  각 픽셀에서는 RGB 값이 각각 0부터 255까지의 값을 가질 수 있으므로 총 $256\times256\times256 = 16777216$개의 조합이 가능하고, 이를 이미지 상의 모든 픽셀에 대해 생각해본다면 무려 $16777216^{256\times256}$ 가지의 조합을 생각할 수 있습니다. 이는 $2^{1572864}$와 같으니까 대략 47만 자리의 수가 되겠군요. 우주에 존재하는 원자의 개수(약 $10^{80}$) 따위는 아득히 초월하는 엄청난 수입니다.

 

그림 2. 이 사진도 $2^{1572864}$ 개의 이미지들 중 하나일 것입니다.

 

    그렇다면 이렇게나 거대한 데이터 공간에서 우리가 생성하고자 하는 이미지들이 차지하는 부분은 얼마나 될까요? 당연히 극히 일부라는 말이 모자랄 정도로 굉장히 적습니다.  [그림 2]의 고양이 사진을 뽑으려면 $2^{1572864}$ 개의 이미지 중 저 이미지 단 하나를 찾아서 뽑아야 한다는 뜻입니다. 이는 모래사장에서 바늘 찾는 것과는 비교가 안 될만큼 어려운 작업이지만 데이터 공간 어딘가에 존재할 거라는 건 분명합니다.

 

    생성 모델은 이렇게 거대한 데이터 공간에서 원하는 이미지가 모여 있는 특정 확률분포를 추정할 수 있도록 학습되고, 그 확률분포에서 이미지를 샘플링함으로써 이미지를 생성할 수 있습니다. 앞서 언급했던 VAE, flow-based model, GAN 같은 경우에는 이러한 과정을 정규분포와 같이 우리에게 잘 알려진 확률분포(prior distribution)를 신경망을 통해 원하는 이미지들이 존재하는 특정 확률분포로 변형해주는 방식으로 구현되었습니다.

 

그림 3. 원자가 공간 상에 균등한 분포로 확산되는 과정은 diffusion model을 이해하는 데 중요한 열쇠가 됩니다. (출처: https://www.mrdubuque.com/home/biodub-my-gifs-to-you-diffusion-osmosis)

 

    그럼 우리가 살펴보고자 하는 diffusion model 같은 경우에는 어떨까요? 이름에서 알 수 있듯이 diffusion model은 '확산'의 원리를 응용합니다. 예를 들자면, 3차원 공간인 자신의 침실을 떠올려 봅시다. 방 안의 어떤 한 위치에 라벤더향 디퓨저를 놓는다고 가정하면 확산의 원리에 의해 라벤더향 분자들이 그림 3처럼 확산되며 방 전체에 균등하게 퍼지게 될 것입니다.

 

   그렇다면, 충분한 시간이 지나 침실 전체에 확산된 분자들 중 하나의 위치를 보고 처음에 디퓨저를 놓았던 위치를 추정하는 것이 가능할까요? 분자들은 매순간 랜덤한 방향으로 확산되기 때문에 불가능해 보이지만, 놀랍게도 확산 과정을 설명하는 시간에 대한 방정식과 확산되는 데 걸린 시간만 정확히 알고 있으면 일정한 오차 이내로 추정할 수 있다는 게 증명되어 있습니다.^[각주:1] Diffusion model은 이러한 통계역학적 특성을 응용하여 데이터 공간(침실) 상의 어느 한 점으로부터 원래의 이미지가 있었던 위치(디퓨저를 놨던 위치)를 추정합니다.

 

그림 4. DDPM의 forward process 예시

 

    이 아이디어를 이미지 생성에  본격적으로 활용한 첫 모델은 Denoising Diffusion Probabilistic Model^[각주:2], 줄여서 DDPM이라고 불리는 모델입니다. 이제 다시 그림 2의 고양이 사진으로 돌아와 DDPM이 어떻게 이미지 생성을 하는지 살펴보도록 합시다.

 

    데이터 공간 상의 어떤 위치에 놓인 고양이 사진이 라벤더향 분자처럼 확산되기 시작하면 원래의 위치와는 전혀 관련 없는 임의의 위치를 향할 것입니다. 같은 방식으로 그림 4처럼 원본 이미지 $\mathrm{x}_{0}$는 사람의 눈에는 단순한 노이즈로밖에 보이지 않는 이미지 $\mathrm{x}_{T}$에 해당하는 위치까지 이동하게 됩니다. 이러한 확산 과정을 diffusion model의 forward process라고 부릅니다.

 

    그런데 이렇게 노이즈가 되어 버린 이미지(라벤더향 분자)를 보고 원본 이미지(디퓨저의 위치)를 추정하기 위해서는 확산 과정을 설명하는 시간에 대한 방정식이 필요하다고 했죠? 그래서 diffusion model은 확산 과정을 아예 다음과 같은 시간에 대한 방정식에 형태로 설정을 하고 시작합니다.

 

$$\mathrm{x}_t = \sqrt{\alpha_{t}}\mathrm{x}_{t-1} + \sqrt{1-\alpha_{t}}\epsilon_{t-1}$$

 

    이 방정식을 시간 스텝 $t=0, 1, \cdots, T$에 대해 순차적으로 계산하면 실제 이미지 $\mathrm{x}_{0}$는 점차적으로 노이즈가 더해지다가 결국에는 완전한 노이즈 $\mathrm{x}_{T}$가 됩니다. 여기서 최대 스텝 수를 나타내는 상수 $T$는 모델마다 다르게 설정되는데, Stable Diffusion의 경우에는 $T=1000$을 사용합니다. 마찬가지로 $t$에 대한 상수 $\alpha_{t}$도 모델마다 다르게 결정되는데, 우선은 항상 $\alpha_{0} = 1$, $\alpha_{T}=0$이 되도록 설정된다고만 아시면 될 것 같습니다. 마지막으로 $\epsilon_{t-1}$은 매 스텝마다 더해지는 평균이 0, 표준편차가 1인 임의의 노이즈입니다.

 

    그런데 $T=1000$이라면 forward process의 끝을 보기 위해 이 식을 1000번이나 계산해야 하는 걸까요? $\epsilon_{t-1}$를 정규분포에서 추출하는 가우시안 노이즈로 설정하면 이 문제가 마법같이 해결됩니다. 정규분포의 가장 강력한 특성 중 하나는 정규분포끼리의 합은 또다른 정규분포가 된다는 것입니다. 가우시안 노이즈도 마찬가지겠죠? 이를 활용해 다음과 같이 여러 개의 forward process step을 한 번에 계산할 수도 있습니다.

 

$$\mathrm{x}_t = \sqrt{\alpha_{t}}\mathrm{x}_{t-1} + \sqrt{1-\alpha_{t}}\epsilon_{t-1} = \sqrt{\alpha_{t}} \left( \sqrt{\alpha_{t-1}}\mathrm{x}_{t-2} + \sqrt{1-\alpha_{t-1}}\epsilon_{t-2} \right) + \sqrt{1-\alpha_{t}}\epsilon_{t-1} \\ = \sqrt{\alpha_{t}\alpha_{t-1}}\mathrm{x}_{t-2} + \sqrt{1-\alpha_{t}\alpha_{t-1}}\bar{\epsilon}_{t-2} = \cdots = \sqrt{\alpha_{t}\cdots\alpha_{1}}\mathrm{x}_0 + \sqrt{1-\alpha_{t}\cdots\alpha_{1}}\epsilon$$

 

    이렇게 원본 이미지 $\mathrm{x}_{0}$와 임의의 가우시안 노이즈 $\epsilon$의 간단한 선형 결합만으로 $t$가 몇이던 상관없이 $\mathrm{x}_{t}$는 한 번에 계산될 수 있습니다. 이런 멋진 특성을 활용하기 위해 Stable Diffusion을 비롯한 거의 모든 diffusion model은 가우시안 노이즈만을 사용합니다. 여기서 $\mathrm{x}_{T}$도 당연히 가우시안 노이즈가 되겠죠?

 

    이제 $\mathrm{x}_{T}$를 다시 $\mathrm{x}_{0}$로 복원하는 reverse process를 위한 모든 준비가 완료되었습니다.  이 부분에서 딥 러닝 기법의 본체인 심층신경망이 나타납니다. Reverse process에서는 신경망 $\epsilon_{\theta}$을 통해 다음과 같은 식으로 $\mathrm{x}_{t}$로부터 $\mathrm{x}_{t-1}$를 복원할 수 있습니다.

 

$$\mathrm{x}_{t-1} = \frac{1}{\sqrt{\alpha_{t}}} \left( \mathrm{x}_{t}-\frac{1-\alpha_{t}}{\sqrt{1-\bar{\alpha}_{t}}}\epsilon_{\theta}(\mathrm{x}_{t}, t) \right) + \sigma_{t}\mathrm{z}$$

 

    위에서 reverse process를 계산하기 위해서는 확산되는 데 걸린 시간도 필요하다고 했는데, 이는 신경망 $\epsilon_{\theta}$의 입력으로 노이즈 낀 이미지 $\mathrm{x}_{t}$와 함께 해당하는 시간 정보 $t$가 필요하기 때문입니다. 이 과정 역시 $t=T, T-1, \cdots, 0$에 대해 순차적으로 계산을 하면 완전한 가우시안 노이즈 $\mathrm{x}_{T}$는 실제 이미지 $\mathrm{x}_{0}$로 복원됩니다. (단, $t=0$일 때는 노이즈 없는 이미지가 나와야 하므로 마지막 항인 $\sigma_{t}\mathrm{z}$가 빠집니다.)

 

그림 5. DDPM의 reverse process를 통한 이미지 생성 예시 (출처: https://minibatchai.com/diffusion/generative-models/text2image/2022/06/17/Diffusion.html)

 

    여기서 $\epsilon_{\theta}$의 역할은 forward process 과정에서 $\mathrm{x}_{t}$을 만들기 위해 $\mathrm{x}_{0}$에 더해졌던 가우시안 노이즈 $\epsilon$을 추정하는 것입니다. 위의 reverse process 식이 의미하는 것은 시간 $t$에서 신경망 $\epsilon_{\theta}$에 의해 추정된 가우시안 노이즈 $\hat{\epsilon} = \epsilon_{\theta}(\mathrm{x}_{t}, t)$를 $\mathrm{x}_{t}$에서 제거하여 원본 이미지의 추정본인 $\hat{\mathrm{x}}_{0}$를 만들고, 다시 $t-1$만큼의 forward process를 진행하여 $\mathrm{x}_{t-1}$을 계산한다는 것입니다. 이 과정을 통해 $\mathrm{x}_{t-1}$를 정확하게 추정하려면 $\hat{\epsilon}$가 정확하게 추정되어야 하므로 diffusion model의 학습 목표는 $\epsilon_{\theta}(\mathrm{x}_{t}, t)$의 값이 forward process의 $\epsilon$에 근접하게 되는 것이라 할 수 있습니다.  

 

    재밌는 건, 이 모델은 노이즈가 더해져 정보가 손상된 이미지를 정확하게 복원하는 방법을 학습했을 뿐이지만, 거기다가 임의의 가우시안 노이즈를 보고 "이건 어떤 이미지에 가우시안 노이즈가 계속 더해지다가 결국 모든 정보를 잃고 완전한 가우시안 노이즈가 된 것이군"이라고도 생각하게끔 학습됐다는 것입니다. 따라서 충분히 학습된 diffusion model의 reverse process에 가우시안 노이즈를 입력해주면 학습 과정에서 배운 대로 자기 딴에 복원 과정을 진행해서 이미지 하나를 생성하게 되는 것입니다. 사용자 입장에서는 노이즈를 넣었더니 랜덤한 이미지 하나가 생성된 것이고, 모델 입장에서는 뭔지 모르겠지만 노이즈를 자기가 생각하는 원본 이미지로 열심히 복원한 셈입니다.

 

    이렇게 기발하면서도 특이한 과정을 거쳐 이미지를 생성한다는 점에서 diffusion model은 다른 생성 모델들과 많은 차이점을 가지고 있습니다. 우선 다른 생성 모델들은 알려진 분포(주로 정규분포)에서 추출한 입력 샘플을 그럴싸한 이미지들이 있을 법한 영역 상의 어떤 지점으로 직접 mapping을 해주는 방식이라면 diffusion model은 입력 샘플을 천천히, 그러나 정밀하게 이동시켜가며 원하는 지점으로 mapping 시킨다고 보면 될 것 같습니다.

 

    그렇다보니 다른 생성 모델에 비해 이미지 하나 생성하는데 굉장히 오랜 시간이 걸린다는 치명적인 단점이 있었습니다. 어느 정도냐면, 흔히 Stable Diffusion으로 주로 생성하는 고해상도 이미지(최소 $512\times512$ 이상)를 DDPM으로 생성하려고 한다면 한 장에 최소 몇 분은 걸릴 정도입니다. 이런 문제점을 극복하여 diffusion model이 우리들의 PC에서도 어느 정도 돌아갈 수 있도록 만들어 준 기술들에 대해서는 다른 포스트들을 통해 차근차근 다뤄보도록 하겠습니다.

 

Implementation

    AUTOMATIC1111의 web UI 기준으로, DDPM 모델은 repositories/stable-diffusion-stability-ai/ldm/models/diffusion/ddpm.py 파일에 DDPM이라는 클래스로 구현되어 있습니다. q_sample이라는 메소드를 통해 forward process를 진행하며, sample이라는 메소드를 통해 reverse process를 진행할 수 있습니다.

 

그림 6. $\epsilon_{\theta}$의 네트워크 구조 (출처: https://arxiv.org/abs/2209.10948)

 

    신경망 $\epsilon_{\theta}$는 repositories/stable-diffusion-stability-ai/ldm/modules/diffusionmodules/openaimodel.py 파일에 UNetModel이라는 클래스로 구현되어 있습니다. 이 신경망은 이름에서 알 수 있듯 UNet^[각주:3]이라는 유명한 네트워크의 구조를 활용하는데, 그림 5처럼 해상도를 점점 낮춰가는 encoder 영역과 점점 높여가는 decoder 영역 사이에 skip connection을 붙임으로써 U자 형태의 구조를 그리는 네트워크입니다. Diffusion model들은 여기에다가 self-attention layer를 특정 위치에 추가해서 사용하는데, 이 self-attention에 관한 내용은 text-conditioning에 대해 다루는 다음 포스트에서 소개하도록 하겠습니다.

 

Conclusion

그림 7. DDPM으로 생성된 $256\times256$ 얼굴 이미지들 (출처: https://arxiv.org/abs/2006.11239)

 

    이렇게 diffusion model을 활용한 이미지 생성 모델의 뿌리라고 할 수 있는 DDPM에 대해 알아봤습니다. 그런데 사실 이 포스트에는 언급이 반드시 되어야 할 것 같지만 한 번도 언급이 되지 않은 내용이 있습니다. 맞습니다. 우리가 쓰는 Stable Diffusion은 기본적으로 text-to-image가 메인인데, 텍스트에 대한 내용이 하나도 없었습니다!

 

    사실 DDPM 논문에서는 텍스트와 같은 conditioning이 하나도 들어가지 않는, unconditional generative model에 대한 내용만을 다뤘습니다. 즉 우리가 원하는 이미지를 생성하도록 통제할 수 없는 모델인 셈입니다. 자연스럽게 이 포스트에서도 diffusion model이 어떻게 우리가 넣는 텍스트대로 이미지를 생성하는지에 대한 내용을 전혀 언급하지 않게 되었는데요, 이 내용은 다음 포스트로 넘기도록 하겠습니다.

 

 

1. Sohl-Dickstein, Jascha, et al. "Deep unsupervised learning using nonequilibrium thermodynamics." International Conference on Machine Learning. PMLR, 2015. [본문으로]
2. Ho, Jonathan, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models." Advances in Neural Information Processing Systems 33 (2020): 6840-6851. [본문으로]
3. Ronneberger, Olaf, Philipp Fischer, and Thomas Brox. "U-net: Convolutional networks for biomedical image segmentation." Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18. Springer International Publishing, 2015. [본문으로]