[Paper Review] DALL-E : Zero-Shot Text-to-Image Generation

Zero-Shot Text-to-Image Generation 논문 링크

Abstract

• text-to-image generation은 전통적으로 고정된 데이터셋 안에서 학습하면서 최적의 모델링을 하는 것을 기본으로 해왔다.
• 이와 같은 전제는 학습 시 복잡한 구조, 보조 함수들, 기타 여러 정보들을 필요로 한다.
• 우리는 text와 image 토큰들을 하나의 스트림으로 concat한 데이터를 autoregressive하게 모델링한 transformer를 통해 이 태스크를 보다 단순하게 접근한다.
  • autoregressive → 대개의 언어모델의 학습 방식을 떠올리면 된다.
• 충분한 데이터와 스케일을 적용한다면 이 모델은 zero-shot방식으로 평가하더라도 domain-specific한 모델들과 견줄만한 성능을 보인다.

1. Introduction

• 이미지 화질이 높아지고 있지만 결과 이미지들은 물체 왜곡, 불가능한 위치에 놓인 물체, 앞뒤에 위치한 물체간의 부자연스러운 조합 등의 문제가 아직도 빈번히 발생한다.
• large-scale generative model을 통한 최근의 발전은 위와 같은 문제에 대한 해결책을 제시한다. 모델 크기, 데이터, 연산 등이 적절히 조절될 때, autoregressive transformers는 text, images, audio 등의 도메인들에 대해 유의미한 결과를 낼 수 있다고 한다.
• 한편 그동안 text-to-generative model은 상대적으로 작은 크키의 데이터셋(MS-COCO, CUB-200)을 사용해왔다.
• 본 논문에서는 2억 5천만 개의 image-text 쌍에 대해 120억 개의 파라미터를 가진 autoregressive transformer로 학습 → 자연어로 컨트롤이 가능한, 높은 품질의 유연한 이미지 생성모델을 만들 수 있었다.
• 위와 같은 시스템은 MS-COCO 데이터셋에 대해 별도의 레이블도 사용하지 않은 채 zero-shot으로 높은 품질의 이미지 생성을 달성하였다.
• 이전의 결과들과 비교하여 human evaluators에게 90%이상의 선호도를 얻었다.
• 또한 복잡한 task에서도 잘 작동했다. 예를 들면 image-to-image translation을 가장 기초적인 수준에서도 충분히 해냈다.
  • 이전의 모델들은 image-to-image translation을 수행하기 위해서 task-specific한 맞춤형 모델을 구축했어야했다.
  • 이제는 오히려 대규모 단일 생성 모델의 능력으로 가능해졌다.

2. Method

• 목표 : 텍스트와 이미지 토큰들을 single stream의 데이터로 만들어서, 이를 autoregressive하게 모델링함으로써 transformer를 학습시키는 것
• 하지만 픽셀들을 그대로 이미지 토큰으로 사용할 시 → 고해상도의 이미지를 처리하기 위해 과도한 메모리 사용 발생
  • 가능도 목적 함수는 픽셀 간 서로 가까운 범위에서의 dependency 모델링을 우선시하는 경향이 있음 → 구체적인 특징은 무시되더라도 이미지의 큰 특징을 구분해주는 전반적인, 즉 거시적인 이미지 추출이 더 중요함에도 불구하고, 국소적인 이미지 특징 추출에 더 중점을 두어 모델의 대부분의 역량을 이에 모두 소진하게 됨
• 이러한 이슈들을 커버하기 위해 본 논문에서는 두 단계의 학습 과정을 적용한다.
• stage 1
  • dVAE(descrete VAE)를 학습 : 256 X 256 rgb 이미지를 32 X 32 grid 이미지 토큰으로 압축, 각 토큰은 8192개의 값 중 하나를 가진다고 상정. → 이는 transformer의 context size를 상당히 줄여준다. visual quality가 크게 downgrade되지 않는 선에서 192배 가량 줄일 수 있음
    • 192 : 8×8×3
    • 8 = 256/32(256 X 256 rgb 이미지를 32 X 32 grid 이미지 토큰으로 압축) : 각각 W, H
    • 3 : rgb이므로 3
• stage 2
  • 최대 256개의 BPE-encoded text tokens를 1024(32x32)개 image tokens와 concat한다.
  • text & image tokens의 결합 분포를 모델링하는 autoregressive transformer를 학습시킨다.
    
    
• 전반적인 과정은 images x, 텍스트 형태의 captions y, encoded rgb image의 tokens z의 (모델 분포에 대한)joint likelihood의 ELB(Evidence lower bound)를 최대화하는 것이라 볼 수 있다.
• 아래와 같은 인수분해를 사용해 하한을 구하는 분포를 모델링할 수 있다.
  • pθ,ψ​(x,y,z)=pθ​(x∣y,z)pψ​(y,z)
  • lnpθ,ψ​(x,y)⩾Ez∼qφ​(z∣x)​(lnpθ​(x∣y,z)−βDKL​(qφ​(y,z∣x),pψ​(y,z)))
    • qφ​ : 주어진 rgb이미지 x를 dVAE 인코더를 통해 생성한 32 X 32 이미지 토큰의 분포를 나타낸다. (* y가 p(x∣z)에 대해 독립적이라 가정)
    • pθ​ : 주어진 이미지 토큰을 dVAE 디코더를 통해 생성한 rgb이미지의 분포를 나타낸다.
    • pψ​ : transformer로 모델링한 텍스트와 이미지 토큰의 결합 분포를 나타낸다.

 

2.1 Stage One : Learning the Visual Codebook

• ELB를 ϕ, θ 에 대해 최대화한다. ϕ, θ는 images에 대해서만 dVAE를 학습하는 것에 해당(아직 text는 사용 x)
  • 초기의 사전 확률 pψ​를 K=8192인 codebook vector에 대한 uniform categorical distribution(균등 카테고리 분포)으로 세팅한다.
    • code book : 임베딩 벡터를 담은 fixed-size table. 여기서는 8192개의 element를 담고 있는 table이다.
• 인코더의 output인 32 X 32 grid에 대해 각 position에서 8192개(각 position에서 8192개의 값을 가질 수 있음)의 logits에 의해 parameterize된 categorical distribution을 qϕ​로 삼는다.
• ELB 최적화 굉장히 어려움 : qψ​가 discrete 분포(이산분포)이기 때문에 이를 최대화하는 데에 reparameterization gradient를 사용할 수 없음
• 기존 연구에서는 이로 인해 online cluster assignment procedure와 결합된 straight-through estimator를 사용하기도 했다.
• 대신 gumbel-softmax relaxation을 사용해 qϕ​에 대한 기댓값을 qϕ​τ에 대한 기댓값으로 대체한다.

💡 Gumbel-Softmax

1) sampling을 하고 싶은데, neural network에서 backpropagation이 불가능하다. 이를 해결하기 위해 Gumbel-Max Trick을 사용하여 backpropagation이 흐르도록 해주자

2) argmax를 사용하였더니 backpropagation이 흐르지 않는다. 이를 어떻게 해결할까? softmax를 취하여 해결함과 동시에, continuous하게 relaxation을 하기 위해 temperature τ를 쓰자

• relaxation : 분포를 보다 연속적으로, 부드럽게 완화(?)
• temperature τ : τ가 0에 가까워질수록 tight(hard해진다고 표현하기도)해진다. 즉 덜 relax해진다.

 

• pθ​의 가능도는 log-laplace distribution을 활용해 evaluate
• relaxed ELB는 Adam으로 최대화(exponentially weighted iterate averaging-지수 가중 반복 평균을 사용). 이때 안정적인 학습을 위해 다음과 같은 3 가지가 중요하다는 것을 발견하였다:
  • relaxation temperature과 step size에 대한 특정 annealing schedule : T를 1/16으로 annealing 하면 relaxed validation ELB가 실제 validation ELB와 거의 유사해진다는 것을 발견했다.
  • 인코더의 마지막과 디코더의 시작점에 1x1 컨볼루션 사용하기 : relaxation 주변의 컨볼루션에서 receptive field 크기를 줄였을 때 실제 ELB로 더 잘 일반화한다.
  • 인코더와 디코더 resblock에서 나가는 activation에 작은 상수 곱하기 : 시작 부분에서 학습이 안정적
• KL weight를 β=6.6으로 증가시킨 것이 codebook의 효율을 높히고, 학습의 끝에 다다라서는 궁극적으로 더 낮은 reconstruction error를 보였다.

 

2.2 Stage Two : Learning the Prior

• 두번 째 단계에서 ϕ와 θ를 고정한 후 ψ에 대한 ELB를 최대화함으로써 text와 image tokens의 사전 분포를 학습한다.
• pψ​는 120억 개의 parameter를 갖는 sparse transformer를 사용한다.
  
  
• text-image pair에 대해, 소문자로 된 캡션 text(최대 256개 토큰들(with vocabulary size 16,384)을 사용)을 BPE 인코딩한다.
• 1024개 토큰(with vocabulary size 8192)로 이미지를 인코딩한다.
• 이미지 토큰은 gumbel-noise 추가 없이 dVAE 인코더 logits로부터 argmax 샘플링했다.
• 위와 같은 일련의 과정을 통해 하나의 데이터 스트림으로서 text, image tokens를 연결하여 autoregressive하게 모델링하게 된다.
  
  
• transformer는 각 image token이 64개 self-attention layers 전체의 모든 text tokens에 attend할 수 있는 디코더 모델이다. 이 중 어느 한 layer에서 attention이 일어난다.
• 3가지 self-attention mask가 모델에 쓰였다.
  • text-to-text attention : standard causal mask
  • image-to-image attention : row, column, or convolutional attention mask
    
    
• 하나의 방법은 셀프어텐션 계산 시 이 토큰들에 대한 로짓값을 음의 무한대로 설정하는 것이다.
• Conceptual Captions에 대한 예비 실험에서, 해당 방법은 높은 validation loss가 나왔지만, 한번도 학습하지 않은 캡션에 대해서는 좋은 성능을 냈다. 즉 일반화 성능이 높았다는 것이다.
  
  
• 본 연구는 한 배치 사이즈에 속한 개수로 text, image tokens의 cross-entropy를 정규화했다.
• 본 논문은 image modeling에 중점을 두었기에 text에 대한 loss를 1/8배, image에 대한 loss를 7/8배로 하여 사용했다.
• 지수 가중 반복 평균을 사용한 Adam으로 최적화했다.
• 오버피팅 없이 수렴하는 양상을 보였다.

 

2.3 Data Collection

• 대략 12억개 정도의 파라미터를 가진 예비 실험에서의 모델이 Conceptual Captions에 수행되었다.
• 330만개의 text-image pair를 갖는, 확장된 MS-COCO dataset에 수행되었다.
• 모델의 파라미터를 120억개 정도로 늘리기 위해 JFT-300M과 비슷한 스케일의 데이터셋을 구축했습니다(2억 5천개의 text-image pair).
  • Conceptual Captions와 YFCC100M의 정제된 서브셋을 포함하는 데이터셋이다.
  • MS-COCO 데이터셋의 validation image 일부가 포함되었으나 caption은 전혀 쓰이지 않았다.

 

2.4 Mixed-Precision Training

• float16-precision & float32-precision을 번갈아 사용하여 GPU 메모리를 효율적으로 사용

 

2.5 Distributed Optimization

• 각 GPU로 병렬 학습 → 연산량을 줄임
  • PowerSGD를 사용하여 그레디언트를 압축함으로써 연산량을 줄일 수 있다.
• (기타 각종 팁(?)이 포함됨)

 

2.6 Sample Generation

• 트랜스포머로부터 샘플링한 결과물을 pretrained contrastive model을 사용하여 rerank한다.

3. Experiment

3.1 Quantitative Result

• DF-GAN과의 비교에서 realism(현실에 있을 법함)과 accuracy(캡션 구현 정확도) 모두, 본 논문의 모델(더구나 zero-shot으로 평가했음에도 불구하고)이 압도적으로 높은 평가를 받았다.
• MS-COCO와 CUB에 대해 모두 여타 모델들과 견줄만한 스코어를 보여준다. (FID & IS Score 모두 낮을수록 좋은 것.)
• CUB 데이터셋에서의 낮은 스코어 : CUB 데이터셋의 도메인이 특이하여 zero-shot으로 접근하는 것이 fine-tuning의 접근 방식 대비 성능이 떨어진 것으로 추측.

 

3.2 Data Overlap Analysis

• 수작업으로 false negative 비율을 최소화하는 신뢰구간을 확보했다.

 

3.3 Qualitative Findings

• 본 모델은 일반화 성능을 갖춘 모델이다.
• 아주 기초 수준의 능력을 높은 수준의 특이하고 추상적인 컨셉을 구현하는 수준으로 끌어올린다.
• 자연어로 조절할 수 있는 zero-shot image-to-image translation이 가능하다.
• image operations, style transfer 등의 변환이 가능하다.
• 매우 기초적인 수준의 object detection 또한 가능하다.

4. Conclusion

• large scale의 autoregressive transformer 기반의 text-to-image generation에 대한 단순한 접근을 제안했다.
• 본 연구 결과는 zero-shot performance와 single generative model 등의 관점에서 scale이 일반화 성능을 개선하는 데에 유의미함을 시사한다.