Post

[논문 리뷰] DALL-E, Zero-Shot Text-to-Image Generation

Posted
By Luna Lee
17 min read



Opan AI ICML 2021

Introduction

최근 다양한 방법들을 이용한 이미지 생성을 위한 다양한 모델들이 개발되었다(GAN, VAE, Energy-based model 등). 이미지 캡션을 활용하여 text-image 생성을 수행하는 시도 또한 진행되었다. 하지만 object distortion(artifact)이나 배경-물체 balending과 같은 문제에 취약한 문제가 있다.

최근에는 large-scale generative model이 다양한 가능성을 열어주고 있다. 특히 autoregressive transformer를 기반으로 한 생성모델은 인상적인 결과를 보여주었다. 하지만 text-image generation은 일반적으로 MS-COCO 및 CUB-200과 같은 상대적으로 작은 데이터 세트에서 진행되어왔다. 본 논문에서 저자는 인터넷에서 수집한 2억 5천만 개의 image-text 쌍에 대해 120억 개의 매개변수 autoregressive transformer를 학습시켰다. 자연어를 통해 제어할 수 있는 유연하고 충실도가 높은 이미지 생성 모델을 제안한다.




Method

본 논문의 목표는 text와 image token을 single stream으로 autoregressive 하게 모델링하는 transformer를 학습하는 것이다. 그러나 픽셀을 직접 image token으로 사용하려면 high-resolution 이미지에 대해 과도한 메모리가 필요하다. 또한 Likelihood objective는 픽셀들간의 short-range dependency를 우선시하는 경항이 있어 모델의 capacity가 low-frequency structure보다는 high-frequency detail에 대부분 소모되는 문제가 있다.

DALL-E_1.png

본 논문에서는 이러한 문제를 해결하기 위해 two-stage 학습 절차를 사용한다.

Stage 1. dVAE(Discrete Variation Autoencoder)를 이용하여 큰 품질 저하 없이 256×256 RGB 이미지를 32×32 이미지 토큰 그리드로 압축

Stage 2. 256개의 BPE 인코딩 텍스트 토큰을 32 × 32 = 1024개의 이미지 토큰과 concat, Autoregressive transformer를 이용하여 text와 image token의 joint distribution 학습



전체 과정은 ELB(evidence lower bound)식에 대한 maximize 문제로 볼 수 있다. 여기서 ELB는 image x, caption y, token z에 대한 모델 분포의 joint likelihood에 대한 식이다. 여기서 분포는 factorization으로 $p_{\theta, \psi}(x, y, z) = p_\theta(x|y,z)p_\psi(y,z)$ 와 같이 나타낼 수 있고 이를 정리하면 아래와 같은 식이 나온다.

DALL-E_2.png


결국 우리가 구해야하는 것:

  • $q_\phi$: dVAE Encoder → RGB 이미지를 이용하여 32 X 32 이미지 token생성, 이 token에 대한 분포를 나타냄
  • $p_\theta$: dVAE Decoder → 이미지 token을 이용하여 RGB 이미지를 생성, 이 RGB 이미지에 대한 분포를 나타냄
  • $p_\psi$: transformer → 이미지 및 텍스트 token을 이용하여 joint distribution 모델링, 이 joint distribution을 나타냄




1. Stage One: Learning the Visual Codebook

(여기서도 VQ-VAE의 codebook 개념을 사용했다고 보면 된다)

학습의 첫 번쨰 단계에서는 $\phi$ 와 $\theta$에 대해 ELB를 maximize한다. 즉, 이미지만을 사용하여 dVAE의 Encoder와 Decoder를 학습한다. 저자들은 transformer의 initial prior $p_\psi$를 K=8192(codebook size)인 codebook vector에 대한 uniform categorical 분포로 설정하고, $q_\phi$는 Encoder output 32 × 32 grid에 대해 8192 logit을 가지는 categorical distribution으로 설정했다(32 × 32 × 8192).


< dVAE Encoder >

Transformer $p_\psi$를 discrete한 분포로 설정했기 때문에, ELB식을 reparameterization gradient로 optimize할 수 없게 되었다. 저자들은 이러한 문제를 해결하기 위해서 gumbel-softmax relaxation 방법을 사용하여 $q_\phi$에 대한 expectation을 $q^\tau_\phi$에 대한 expectation으로 변경했다.

Gumbel-softmax trick

미분 불가능한 함수를 미분 가능하게 변경하여 backprop이 가능하게 하기 위해 사용하는 방법.

Categorical 분포에 대해 미분가능하도록 Reparameterization을 수행하여 gradient를 흐르도록 사용하는 방법 중 하나이다.

Categorical 분포에서 sampling을 수행할 때, Gumbel noise를 추가한 뒤 softmax해주는 방법이다. 유사하게는 Gumbel-max(Gumbel noise를 추가한 뒤 argmax 추출)가 있다(하지만 argmax를 구하는 과정 역시 미분불가능하기 때문에 continuous하게 softmax를 사용한다).

DALL-E_3.png

참조: https://kaen2891.tistory.com/81




< dVAE Decoder >

$p_\theta$의 likelihood는 log-laplace distribution으로 평가된다.

Appendix A.3 일반적으로 VAE를 학습할 때 L1, L2 reconstruction loss를 사용한다.
그 이유는 위의 수식 1에서 $\ln p_\theta(x|y,z)$를 Laplace, Gaussian 분포라고 가정하기 때문이다. 하지만, 픽셀은 제한된 간격 내에 있는데 비해 이러한 분포는 전체 Real line에 존재하기 때문에 모델링에 불일치가 있다. 일정량의 likelihood는 허용되는 픽셀 값 범위 내를 벗어나게 되어있다.

따라서 본 논문의 저자는 제한된 구간에서도 지원되는 Laplace 분포의 변형을 제시했다. 해당 분포는 (0,1) 사이에서 정의되고 logit-Laplace distribution이라고 부르고, 해당식을 dVAE의 Decoder 학습의 reconstruction loss에 이용한다.



Relaxed ELB 식에 대한 optimization은 Adam으로 수행된다. 하지만 안정적인 학습을 위해 특정 annealing schedules이 필요하다고 한다. 또한 Encoder의 끝과 Decoder의 시작에 1×1 convolution을 추가하여 relaxation 주위의 receptive field size를 줄였을 때 ELB의 일반화를 더 잘 이끌어낸다고 밝혔다.




2. Stage Two: Learning the Prior

두 번째 단계에서는 $\phi$ 와 $\theta$를 fix하고 $\psi$에 대해 ELB를 최대화하여, 이미지와 텍스트 token에 대한 prior 분포를 학습한다. $p_\psi$ 는 120억개의 파라미터를 가진 sparse transformer에 해당한다.

text-image pair가 주어지면 다음과 같은 encoding을 수행한다.

  • Text: BPE로 encoding. 256 token, vocabulary size 16,384.
  • Image: dVAE로 encoding, 32 × 32 = 1024 tokens, vocabulary size 8192. (여기서 image tokens은 dVAE encoder logits으로부터 argmax sampling만 취함, gumbel noise X)

⇒ 마지막으로 텍스트와 이미지 token은 concat되고 single-stream으로 모델링됨(autoregressively).


< Transformer >

Transformer는 Decoder only model로, 각 이미지 token은 모든 텍스트 token에 attend 가능하다(64개의 self-attention layer로 구성, 이 중 하나에서 attention). 모델에는 세 종류의 self-attention mask가 사용되었다.

  • text-to-text attention: standard causal mask 사용

  • image-to-image attention: row, column, convolutional attention mask 사용 가능

    DALL-E_4.png

    (a) 각각의 이미지 token이 현재 위치(픽셀) 이전의 5개(raster order) 이미지 token(픽셀)에 대해 attend

    (b) GPU 활용을 위해 (b)를 (c)로 변경, image 상태에 대해서만 transpose 수행

    (c) 고정된 이전의 열(column)에 대해서 attend, text와 같은 데이터에 유용하다고 함 (데이터에 구조와 주기성이 약한 경우)

    (d) 일반적인 3X3 convolutional attention pattern, 마지막 self-attention 레이어에만 사용됨


텍스트 캡션의 길이를 256으로 제한했지만 마지막 텍스트 token과 이미지 token사이의 padding 위치에 무엇을 넣어야하는지는 명확하지 않았다고 한다. 한가지 옵션으로 self-attention 작업에서 이부분의 token에 대한 logit을 $-∞$로 설정하는 것이다. 하지만 논문에서는 256개의 텍스트 위치 각각에 대해 별도로 special padding token을 학습했다.

마지막으로 텍스트 token과 이미지 token에 대한 cross-entropy loss는 각각 1/8, 7/8(1:7)의 비율로 설정하였다. 생성 모델의 특성상 image modeling에 초점을 맞추기 위해 위와 같이 설정했다고 한다.




3. Data Collection

  • Conceptual Captions 데이터셋(330만개의 text-image pair, MS-COCO 확장본)으로 1.2B 모델 학습
  • 모델 파라미터를 12B로 늘리기 위해 2억 5천만 개의 text-image pair를 수집하여 JFT-300M과 유사한 규모의 데이터셋 생성했다. YFCC100M의 filtered subset, MS-COCO validation 일부, Wikipedia의 text-image pair가 포함되었다.




4. Mixed-Precision Training

  • GPU 메모리를 절약하고 처리량을 높이기 위해 대부분의 파라미터를 16-bit precision을 사용.
  • activation checkpointing 사용, resblock에서 backward pass 중 activation recompute

하지만 학습이 잘 되지 않았고 16-bit gradient에서 발생한 underflow가 원인이라고 분석하였고, 이를 해결하기 위해 per-resblock gradient scaling 방법을 사용하였다.

< per-resblock gradient scaling >

  • 이전 연구와 유사하게, 이전 resblock에서 다음 resblock으로 이동할 때 activation gradient의 norm이 단조 감소한다는 것을 발견했다.
  • 모델이 깊어지고, 넓어질수록 뒤쪽의 resblock에 있는 activation gradient는 16-bit 형식의 최소 단위보다 작아질 수 있다.
  • 결과적으로 gradient가 0으로 반올림되는데, 이러한 현상을 underflow라고 한다.
  • 이러한 underflow를 제거하면, 모델이 안정적으로 수렴하여 학습할 수 있다는 것을 발견함.

DALL-E_5.png

따라서 위의 그림과 같이 각각의 resblock에 대해 별도의 gradient scale을 사용한다.




5. Distributed Optimization

12B 파라미터 모델은 16-bit precision을 사용할때 약 24GB의 메모리를 소비하여 16GB인 V100의 GPU 메모리를 초과한다. 따라서 parameter sharding을 사용하여 이 문제를 해결한다. 이 방법은 compute-intensive operation을 사용하여 intra-machine 통신의 latency를 거의 없앨 수 있다고 한다.

DALL-E_6.png

모델을 학습하는데 사용되는 클러스터에서의 머신 간의 대역폭은 동일한 머신 내의 GPU 대역폭보다 작다. 이로 인해서 학습 중 발생하는 머신간 gradient average cost로 인한 병목현상을 해결하기 위해 저자는 PowerSGD를 사용하여 gradient를 압축하여 cost를 대폭 줄였다.

< Power SGD >

  • 별도의 버퍼를 할당하는 대신 backpropagation 중에 error 버퍼에 gradient를 누적하여 메모리를 절약
  • Error 버퍼를 0으로 만드는 인스턴스를 최소화(예: mixed-precision backpropagation중의 nonfinite 값이 나올 때, checkpoint로부터 학습 재시작할 때)
  • 입력에 identity matrix의 작은 배수를 추가하고 Gram-Schmidt 대신 Householder orthogonalization를 사용하여 수치적 안정성을 개선
  • Custom 16-bit floating point(부동 소수점) 형식을 사용하여 underflow 방지



6. Sample Generation

DALL-E_7.png

Pretrain된 contrastive model을 사용하여 transformer에서 추출한 샘플에 대해 rerank를 수행했다. 후보 이미지와 캡션이 주어지면, contrastive 모델은 이미지가 캡션과 얼마나 잘 매치가 되는지를 기반으로 score를 할당한다. 위의 그림 6은 상위 k개 이미지를 선택하는 샘플 수 N을 증가시켰을 때의 효과를 보여준다.




Experiments

1. Quantitative Results

Zero-shot 평가를 위해 AttnGAN, DM-GAN, DF-GAN과 비교했다.

DALL-E_8.png

MS-COCO 데이터셋에 대한 결과이다. 생성 이미지 품질을 이전 작업들과 비교하기 위한 결과이다.


DALL-E_9.png

이미지 품질에 대한 평가을 위한 human evaluation을 수행했다. 이에 대한 결과는 위의 그래프와 같다.


DALL-E_10.png

FID, IS 비교 결과는 위와 같다. (a)는 MS-COCO, (b)는 CUB, (c)는 MS-COCO에 대해 constrastive 모델과의 reranking에 사용되는 샘플 크기가 증가함에 따른 결과이다.



2. Data Overlap Analysis

제거할 이미지를 결정하기 위해 CLIP에서 사용한 중복 제거 절차를 사용했다. 이 작업을 위해 Contrastive 모델을 따로 학습하고, 이 모델을 사용하여 각각의 validation 이미지에 대해서 가장 가까운 이미지를 찾는다. 그 다음 학습 데이터에서 가장 가까운 일치를 보이는 데이터 순으로 정렬한다. 그 뒤엔 직접 결과를 검사한 후, false negative rate를 최소화하도록 설계된 conservative threshold를 manually 선택하여 제거할 이미지를 결정한다.



3. Qualitative Findings

DALL-E_11.png

또한 저자들은 논문의 모델이 예상하지 않았던 방식으로 generalize 능력을 학습했다는 것을 발견했다. (a) 그림에서 결과를 확인할 수 있다. 이 결과는 높은 수준의 추상화로 특이한 개념을 구성하는 능력을 개발했음 의미한다.

(b), (c) 결과를 통해 모델이 combinatorial generalization 성능 또한 가지고 있다는 것을 알 수 있다. 캡션 내 단어들의 관계에 대한 이해가 가능하다.

또한 (d) 결과를 통해 모델이 zero-shot image-to-image translation 또한 가능하다는 것을 알 수 있다.



AI, Paper Review
Multi-modal Generation Zero-shot Transformer
This post is licensed under CC BY 4.0 by the author.