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

들어가기 전

 제로샷 : 학습 데이터가 아예 없어도 유연한 패턴인식을 할 수 있는 걸 의미한다. 제로샷 학습은 Transfer Learning에서 발전했음. 

 

MS-COCO : Microsoft사에서 만든 4세 아이가 쉽게 인식할 수 있는 91개 개체 유형의 사진이 포함된 총 250만개 레이블로 표시된 인스턴스 데이터셋으로 91개의 클래스와 328,000장의 이미지가 포함되어 있고 object 수는 2,500,000여 개

 

CUB-200 : 북미 새 200종 11,788개의 이미지를 포함하는 데이터셋으로 각 이미지에 대한 5개의 fine-grained 설명을 포함

 

fine-grained : 세밀한 class들을 분류하는 것

 

Reparametrization : VAE에서 encdoer의 출력값 mu&sigma를 그대로 decoder에 반영하지 않고  e를 출력에 더해주는 형태를 취해 z값에 encoder가 미치는 관계를 찾을 수 있게 하는 것(미분 가능한 모델로 표현 가능) {N(mu, sigma) + e, e~N(0,1)}

 

KL Divergence : 두 확률 분포의 다름의 정도를 나타내는 Divergence의 한 방법

 

ELB : 관찰한 P(z|x)가 다루기 힘든 분포를 이루고 있을 때 이를 조금 더 다루기 쉬운 분포인 Q(x)로 대신 표현하려 하는 과정에서 두 분포 (P(z|x)와 Q(x))의 차이 (KL Divergence)를 최소화 하기 위해 사용된다. D_KL(qφ(y, z | x), pψ(y, z))로 표현. https://hugrypiggykim.com/2018/09/07/variational-autoencoder%EC%99%80-elboevidence-lower-bound/

 

reconstruction error : ELB의 첫 번째 항으로 VAE가 AE처럼 reconstruction을 잘 할 수 있게 만들어주는 error

 

Gumbel-softmax : 미분불가능한 이산확률변수(범주형 변수)를 이산적인 상태를 표현하며 오차역전파알고리즘에 사용할 수 있게하는 함수

 

log-laplace distribution : laplace 분포는 가우시안 정규분포의 변형되 형태로, 정규분포는 종모양인 것에 비해 뾰족한 형태를 띔.

 

receptive field : convolutional layer에서 feature map 하나의 값이 이전 층 이미지에서 담당하는 범위. 층이 깊어질수록 커짐

 

autoregressive Transformer : 오리지날 트랜스포머 모델의 decoder를 사용하며, 이전 모든 토큰들에 대해 다음 토큰을 예측하는 방식으로 사전학습 된다. 대표적으로 GPT모델이 있다.

 

1. Abstract

  기존의 text-to-image 생성은 전통적으로 고정된 데이터셋에서의 모델링의 가정에 초점이 맞추어져있었다. 이 가정들은 훈련 시 복잡한 구조나 보조 손실함수, 또는 객체 부분 라벨이나 세그멘테이션 마스크같은 부가 정보가 포함되어야한다. 우리는 단일 데이터 스트림으로 텍스트와 이미지를 자동회귀 모델링하는 transfomer에 기초하여 이 문제에 대해 간단한 접근방식을 설명한다. 충분한 데이터와 스케일로, 우리의 접근방식은 제로샷 평가 방식에서 기존의 도메인 모델에 경쟁력을 갖는다.

 

1. Introduction

  현대 머신러닝의 text-to-image 생성 접근방식은 Image cation으로 확장시킬 때 생성모델도 새로운 시각적 장면을 생성할 수 있음과 관련된 연구에서 시작되었다. 이후 연구로 Gan을 사용하는 것이 recurrent variational auto-encoder을 사용하는 것보다 이미지 충실도가 더 향상된다고 입증되었다. Reed et al. (2016b)는 이 시스템이 인식 가능한 속성으로 객체를 생성하는 것 뿐만 아니라 제로샷을 held-out 카테고리로 일반화할 수 있다.

 

 이와는 별개로 Nguyen at el.(2017)는 품질을 크게 향상시킨 조건부 이미지 생성을 위한 energy-based 프레임워크를 제안했다. 이 접근방식은 사전 훈련된 판별모델을 포함하며, 그들은 MS-COCO에서 사전 훈련된 캡션 모델에 적용될 때 텍스트-이미지 생성을 수행할 수 있음을 보여준다. 더 최근엔, Cho et al. (2020)은 cross-modal masked language model의 사전훈련 Input을 최적화하는 방법을 제안했다. 계속해서 시각적 충실도에서 의미있는 향상이 있었지만, 샘플들은 객체 왜곡, 비논리적 객체 배치 또는 전경과 배경 요소의 부자연스러운 합성과 같은 어려움을 겪을 수 있다.

 

 최근의 발전에서 Large-scale 생성 모델이 추가 개선할 수 있는 방향을 제시했다. 명확하게, 컴퓨팅, 모델 크기 및 데이터를 신중하게 확장할 때 autogressiv transformers는 인상적인 텍스트와 이미지, 그리고 음성같은 몇몇의 도메인에서 인상적인 결과를 성취했다.

 

 그에 비해, text-to-image 생성은 일반적으로 MS-COCO and CUB-200와 같은 비교적 작은 데이터셋에서 평가되었다. 데이터셋 사이즈와 모델의 사이즈가 현재의 접근방식에서 제한적인 요소가 될까? 이 논문에서 , 우리는 12-billion 파라미터를 가진 autogressive transformer에 250million개의 image-text쌍을 학습시켜 자연어를 통해 제어 가능한 이미지의 유연하고 충실도가 높은 생성 모델이 생성된다는 것을 보여준다.

 결과는 MS-COCO 데이터셋에서  제로샷으로 훈련 라벨을 사용하지 않고 높은 퀄리티의 이미지 생성 달성했다. 이것은 human evaluators의 데이터셋에 학습된 이전의 작업들보다 선호된다. 우리는 또한 image-to-image translation같은 복잡한 task들에서도 가장 기본적인 작업을 수행할 수 있음을 발견했다.

 

2. Method

Figure 1

 

 우리의 목표는 text와 image를 단일 스트림 데이터로 사용해 transformer를 자동회귀적으로 모델링하는 것이다. 픽셀을 그대로 이미지 토큰으로 사용하는 것은 높은 메모리 용량을 필요로할 수 있다. Likelihood 목표는 픽셀의 짧은 범위 내 종속성을 우선으로 모델링하는 경향이 있고, 이는 모델이 우리가 객체를 시각적으로 인식할 수 있게 하는 저빈도의 구조대신 고빈도의 디테일을 포착하려는데 용량이 사용되게 할 수  있다.

 우리는 이 문제들을 두 단계에 걸친 훈련 과정을 통해 해결한다. ( (Oord et al., 2017; Razavi et al., 2019):와 비슷함)

•  Stage 1 : 256X256RGB 이미지를 32X32사이즈의 이미지 토큰으로 압축시키기 위해  dVAE를 학습시킨다.  이 때 각각의 값들은 8192개의 가능한 값들을 추정할 수 있다.  (-> 확률분포에서 샘플링하는 값이 8192인가?)  이것은 시각적 퀄리티의 큰 저하없이 Transformer의  context size를 줄여준다. (Figure 1)
•  Stage 2 :  우리는 32 × 32 = 1024 이미지 토큰과 최대 256개의 BPE으로 인코딩된 텍스트 토큰을 연결하고, 텍스트 및 이미지 토큰에 대한 결합확률분포를 모델링하기 위해 autogressive Transformer를 훈련한다.

 전체적인 절차는 이미지 x, 캡션 y 및 인코딩된 RGB Image token z에 대한 모델 분포의 결합우도에 대한 ELB(evidence lower bound)를 최대화하는 것으로 볼 수 있다. 다음과 같은 lower bound를 산출할 수 있는 pθ,ψ (x, y, z) = pθ (x | y, z)pψ (y, z)를 사용해 이 분포를 모델링하기로 했다.

 

dall-e의 ELB

• qφ는 x^2이 주어진 dVAE의 encdoer로부터 생성된  32X32 Image token(z)의 분포를 나타낸다.
• pθ는 Image token이 주어진 dVAE의 decoder로부터 생성된 RGB image(x)들의 분포를 나타다.
• pψ는 Transformer로 모델링된 text(y)와 image(x0의 결합확률분포를 나타낸다.

 경계는 β = 1에 대해서만 유지되지만 실제로는 더 큰 값을 사용하는 것이 도움이 된다.다음 하위 섹션에서는 두 단계에 대해 자세히 설명한다.

 

  2.1 Stage one : Learning the visual codebook

   학습의 첫 단계에서, 우리는 이미지에서만 dVAE를 학습시키는 것에 해당하는 φ와 θ에 대해 ELB를 최대화한다. 우리는 초기 pψ의 사전분포를 K=8192 codebook 벡터들에 대해 이산균일분포로 설정하고, qφ는 인코더에 의해 32 x 32 그리드 출력에서 동일한 공간 위치에 있는 8192 로짓에 의해 매개변수화된 카테고리분포로 설정했다. 

  ELB는 최적화하기 어려워졌다 : qψ가 이산분포이고, 이것을 최대화하기위해 reparameterization gradient를 사용할 수 없다. Ord 외(2017); Razavi 외(2019)는 직선 추정기와 결합된 온라인 클러스터 할당 절차를 사용하여 이를 해결한다. 우리는 대신 Gumbel-softmax완화를 사용하여 qθ에 대한 기대치를 θ → 0에 따라 이완이 엄격해지는 qθ에 대한 기대치로 대체한다. pθ에 대한 평가는 log-laplace distribution을 사용한다. relaxed ELB는 지수적으로 가중되는 반복 평균에 의해 최대화 된다. 부록 A.2는 하이퍼 파라미터에 대한 완전한 설명을 제공하며, 다음 내용은 안정된 학습을 위한 중요한 부분이다.

• 특정한 annealing은 완화된 temperature과 step size를 스케쥴링한다. 우리는 τ에서 1/16까지 어닐링하면 검증 relaxed ELB와 검증 true ELB 사이의 갭을 qφ가 아닌 qφ^τ로 좁힐 수 있다는 것을 발견했다.
• 1X1 합성곱 층의 encoder 마지막과 decoder 시작 부분에 사용하는 것이다.  우리는 relaxation 근처의 합성곱을 위한 receptive field의 사이즈를 줄이는 것이 true ELB에 가깝게 생성하도록 하는 것을 확인했다.
• encdoer 및 decoder의 resblock(residual block)에서 나가는 활성화에 작은 상수를 곱하면 초기에 안정적인 학습이 보장된다

우리는 또한 KL 가중치를 β = 6.6으로 증가시키면 codebook 사용이 개선되고 궁극적으로 교육이 끝날 때 reconstruction error가 더 작아진다는 것을 발견했다.

 

2.2 Stage Two : Learning the Prior

 두 번째 단계에서 φ와 θ를 고정시키고, ψ에 대한 ELB를 최대화함으로써 text와 image token에 대한 prior distribution을 학습한다. pψ는 120억 개의 Sparse Transformer 파라미터로 표현된다. text-image 쌍이 주어지면 text는 소문자로 변형해 16,384크기의 vocabulary size로 최대길이 256의 토큰으로 BPE을 거치고, Image는 vocabulary size 8192로 32X32=1024의 token으로 encoding한다. Image token은 gumbel noise의 추가없이 dVAE의 encoder logits으로부터 argmax sampling을 통해 추출한다. 끝으로, image와 text token은 concat되어 single-stream-data로 autoregressively하게 모델링된다.

 여기 쓰인 Transformer는 64개의 self-attention층을 통해 image가 모든 text token을 엿볼 수 있다. 여기서 세 가지 다른 종류의 self-attention mask가 쓰인다. text-to-text는 standard-casual-mask, image-to-image는 row, column 또는 convolutional attend mask가 쓰인다. 

 text token의 마지막과 image token의 시작부분의 padding position을 어떻게 할지 명확하진 않지만  text caption의 길이를 256으로 제한했다. 우리는 256개의 text token 각각에 대해 개별적으로 special text token을 학습시키기로 선택했다. 이 token은 사용할 수 있는 text token이 없는 경우에만 사용된다. 우리는 이로 인해 유효성 손실은 더 높지만 분포 외의 caption에서 더 나은 성능을 보인다는 것을 발견했다. 

 우리는 batch size로 text와 image token의 cross-entropy loss를 정규화했다. 우리는 image를 모델링하는데 더 중점을 두었기 때문에, text의 loss에 1/8의 가중치를, image의 loss에는 7/8의 가중치를 두었다. 

ZeroWe normalize the cross-entropy losses for the text and image

 

2.6 Sample Generation

 Pre-trained constrastive model을 사용하여 Transformer에서 추출한 sample의 순위를 재조정했다. caption과 후보 image가 주어지면 constrastive model은 image와 caption이 얼마나 잘 일치하는지에 따라 점수를 할당한다.

 

 위 이미지는 상위 k개의 이미지를 선택하는 샘플 N의 수를 늘리는 효과를 보여준다. 언급이 없는 한, 정성적 및 정량적 결과에 사용되는 모든 sample은 temperature reduction 없이 생성되고, N = 512를 사용해 rerank했다.

 

 

 

---

 

 VAE의 MLE와 ELB 적용 과정부터 다시 훑고 왔는데 이해가 안될 뻔(ㅜ)했다. Method section의 2.5 ~ 2.6 Part는 data collection과 GPU에 관련된 부분이라 생략했다. 컴퓨팅 파워를 고려한 모델의 설계도 중요하다고 생각해 기회가 되면 공부를 하고 다시 보려고 한다