[Paper Review] DALL-E 2: Hierarchical Text-Conditional Image Generation with CLIP Latents
DALL-E 2는 텍스트 설명을 바탕으로 이미지를 생성하는 text-to-image 모델입니다. DALL-E의 후속 모델이며, 2022년에 발표되었습니다. 선행 연구들은 주로 텍스트로부터 이미지를 바로 생성했습니다. 그런데 텍스트가 이미지를 100% 완벽하게 설명하는 경우는 거의 없기 때문에, 이런 방식으로 생성한 이미지의 품질에는 한계가 있을 수밖에 없습니다.
이와 달리 DALL-E 2는 이미지 생성 과정을 2단계로 나누어 텍스트 설명(caption)을 입력받아 이미지 embedding을 생성한 다음, 이미지 embedding을 조건으로 이미지를 생성합니다. 조건부 분포 $p(\text{image_embedding} \vert \text{text_caption})$를 명시적으로 모델링한 다음 $p(\text{image} \vert \text{image_embedding})$을 모델링하는 것입니다. 여기에서 이미지 embedding을 생성하는 데는 CLIP 모델을 사용하고, 이미지를 생성하는 데는 Diffusion 모델을 사용했습니다. 이 과정을 통해 이미지의 다양성을 향상시키고, 사진처럼 사실적이면서 텍스트 설명에 부합하는 이미지를 생성할 수 있었다고 합니다. 이제 DALL-E 2의 모델 구조, 학습 방법과 성능에 대해 좀 더 자세히 알아보겠습니다.
1. 모델 학습
DALL-E 2 architecture
앞서 설명한 것처럼 DALL-E 2는 CLIP과 Diffusion 모델을 함께 사용하는 모델입니다. 모델 학습과 이미지 생성 과정에서 CLIP과 Diffusion 모델이 어떻게 사용되는지 자세히 살펴보겠습니다.
CLIP 모델 pre-training
DALL-E 2의 학습에 앞서, 대규모의 (텍스트, 이미지) 쌍 데이터셋을 사용해 CLIP을 pre-training합니다. CLIP의 학습 목표는 이미지와 텍스트를 동일한 다차원 공간에 매핑하여, 서로 관련된 이미지와 텍스트의 embedding은 가깝게, 관련 없는 embedding들은 멀게 위치하도록 만드는 것입니다. 그 결과로 텍스트와 이미지를 embedding으로 변환하는 텍스트 인코더와 이미지 인코더가 완성됩니다.
DALL-E 2는 모델 학습의 첫 번째 단계에서 이 인코더들을 가져와 사용하는데, 위 이미지에서 text encoder와 img encoder 부분에 해당합니다. (텍스트, 이미지) 쌍으로 이루어진 데이터셋에서 텍스트와 이미지를 각각 텍스트 인코더와 이미지 인코더에 입력하여 embedding $z_t$와 $z_i$를 얻습니다.
Prior 학습
CLIP으로 이미지와 텍스트의 embedding을 얻은 이후, DALL-E 2의 학습은 두 가지의 단계 prior와 decoder로 나누어 진행됩니다. 먼저 prior $p(z_i \vert y)$는 텍스트 $y$를 입력받아 CLIP 이미지 embedding $z_i$를 생성하는 모델입니다. 저자들은 두 가지의 prior 모델 Autoregressive(AR) prior와 Diffusion prior를 모두 실험했는데요. 결과적으로 Diffusion prior가 더 좋은 성능을 보였다고 합니다.
Autoregressive(AR) prior
PCA를 통해 CLIP 텍스트 embedding $z_t$의 차원을 축소(1024→319)한 다음, Transformer로 이미지 embedding $z_i$를 예측합니다. 이미지 embedding과 텍스트 embedding의 유사도 $z_i \cdot z_t$를 추가적인 정보로 사용했는데요. $z_i \cdot z_t$를 몇 개의 이산적인 구간으로 나누어 양자화해서 입력 sequence 맨 앞에 추가했습니다.
Diffusion prior
Diffusion prior도 마찬가지로, Transformer로 이미지 embedding $z_i$를 예측합니다. 하지만 Diffusion 모델이기 때문에 입력 sequence와 학습 방식에서 차이가 있습니다. 입력 sequence는 (텍스트 토큰, CLIP 텍스트 embedding $z_t$, timestep embedding, 노이즈가 주입된 이미지 embedding)으로 구성되어 있습니다.
\[L_{\text{prior}} = \mathbb{E}_{\,t \sim [1,T],\; z_i^{(t)} \sim q_t} \Bigl[ \,\|\,f_\theta\!\bigl(z_i^{(t)},\, t,\, y\bigr) - z_i \|^2 \Bigr]\]AR prior는 다음 토큰이 무엇인지 분류하는 것이 목적이므로, Cross-entropy loss를 계산하지만, Diffusion prior는 모델이 예측한 이미지 embedding과 원본 이미지 embedding $z_i$와의 MSE loss를 계산합니다. DDPM처럼 이미지 embedding에 주입된 노이즈 $\varepsilon$를 예측하는 버전도 실험했지만, 이미지 embedding을 직접 예측했을 때 성능이 더 좋았다고 합니다.
Decoder 학습
디코더 $p(x \vert z_i, y)$는 prior 모델에서 예측한 CLIP 이미지 embedding $z_i$와 timestep embedding을 입력으로 받아 실제 이미지 $x$를 생성하는 Diffusion 모델입니다. 텍스트 $y$는 decoder의 성능을 높이기 위해 선택적으로 사용할 수 있습니다.
Diffusion prior와는 달리, U-Net을 사용해 실제 이미지 $x$에 주입된 노이즈 $\varepsilon$을 예측합니다.
\[L_{\text{decoder}} = \mathbb{E}_{\,t \sim [1,T],\; \varepsilon} \Bigl[ \,\|\,\hat{\varepsilon}_\theta(x_t, t, z_i, z_t) - \varepsilon \|^2 \Bigr]\]prior와 decoder를 쌓으면 텍스트 $y$가 주어졌을 때 이미지 $x$를 생성하는 생성 모델 $p(x \vert y)$가 만들어집니다.
\[P(x \vert y)=P(x \vert z_i,y)=P(x \vert z_i,y)P(z_i \vert y)\]2. 모델 성능
좌측 상단과 우측 상단의 원본 이미지를 입력해서 DALL-E 2로 이미지를 변형한 결과입니다. CLIP으로 의미 정보나 스타일 요소만 남긴 이미지 embedding을 만들고, 이를 이용해 Diffusion으로 새로운 이미지를 만들기 때문에 의미는 유지하면서 픽셀 단위 표현을 자유롭게 변형한 것을 확인할 수 있습니다.
좌측 그림에서는 시계, 황야, 초현실적인 분위기 같은 의미적 요소는 그대로 유지되고, 시계 배치, 하늘 색감이나 구름 모양같은 디테일한 요소들은 변형되었습니다. 우측 그림에서도 로고 구조라는 핵심 의미는 유지되면서 배경의 그라데이션이나 로고 굵기, 비율 등은 다양하게 변형되었습니다.
위 이미지들은 DALL-E 2 모델로 이미지를 생성한 결과입니다. 2021년에 발표된 DALL-E보다 4배 더 높은 해상도로 더욱 사실적이고 정확한 이미지를 생성합니다.