[Paper Review] Zero-Shot Text-to-Image Generation

1. Introduction

 

Recent advances fueled by large-scale generative models suggest a possible route for further improvements. Specifically, when compute, model size, and data are scaled carefully, autoregressive transformers have achieved impressive results in several domains such as text, image, and audio.

 

최근 연구에 따르면, 1) 계산량, 2) 모델 크기 그리고 3) 데이터 크기를 scaling하여 자기회귀 트렌스포머 모델(e.g., GPT-2)를 학습시키면 대부분 분야(e.g., 텍스트, 이미지, 오디오)에서 높은 성능 달성한다.

 

In this work, we demonstrate that training a 12-billion parameter autoregressive transformer on 250 million image-text pairs collected from the internet results in a flexible, high fidelity generative model of images controllable through natural language. 

 

이번 연구에서는, "2억 5천만개 이미지-텍스트 쌍 데이터 세트로 120억 파리미터 트렌스포머를 학습시키면, 자연어로 수준 높은 이미지 생성이 가능하다"는 것을 증명해보일 것이다.

 

2. Method

Our goal is to train a transformer autoregressively model the text and image tokens as a single stream of data.

 

이번 연구 목표는 텍스트-이미지 토큰을 자기회귀적으로 모델링하는 트렌스포머(e.g., GPT-2)를 학습시키는 것이다. 즉, $p(x, y) = \prod_{i=1}^N p(x_i | x_{<i}, y)$와 같이 모델링된 트렌스포머를 학습시키는 것이다.

 

However, using pixels directly as image tokens would require an inordinate amount of memory for high-resolution images. Likelihood objectives tend to prioritize modeling short-range dependencies between pixels, so much of the modeling capacity would be spent capturing high-frequency details instead of the low-frequency structure that makes object visually recongizable to us.

 

픽셀을 이미지 토큰으로 사용하면 2가지 문제점이 발생한다. 1). 메모리 비용이 비싸다. 2). Likelihood 함수는 "이미지 구조"(low-frequency)보단 가까운 픽셀 관계성 즉, "이미지 디테일"(low-frequency)을 학습하는데 보다 많은 모델링 자원을 사용한다.

 

We address these issues by using a two-stage training procedure, similar to (Oord et., 2017\ Razavi et at., 2019): Stage 1.    We train a discrete variational autoencoder (dVAE) to compress each $256 \times 256$ RGB image into a $32 \times 32$ grid of image tokens, each element of which can assume 8192 possible values. This reduces the context size of the transformer by a factor of 192 without a large degration in visual quality. Stage 2.    We concatenate up to 256 BPE-encoded text tokens with the $32 \times 32 = 1024$ image tokens, and train an autoregressive transformer to model the joint distribution over the text and image tokens.

 

이 문제들을 2단계 학습 과정으로 해결한다. 1) $256 \times 256$개 RGB 픽셀 이미지를 $32 \times 32$개 잠재 변수로 변환해주기 위해 dVAE 모델을 학습시킨다. 각 잠재 변수는 8192가지 값 중 하나가 된다. 이 모델은 시각적 품질의 큰 저하 없이 이미지 토큰 개수를 줄여준다. 2). 텍스트 토큰과 이미지 토큰을 합쳐 자기회귀 트렌스포머를 학습시켜 이미지-텍스트 토큰의 결합 분포($p(x, z)$)를 모델링한다.

 

글쓴이 뇌피셜

픽셀을 잠재 변수로 변환하면서 1) 이미지 토큰 개수가 줄어들어 메모리 비용을 감소시켰다. 2) 뿐만 아니라 잠재 변수는 픽셀들 간의 "이미지 구조(low-frequency)" 특성을 추출하기에 트렌스포머가 "이미지 구조"를 학습할 때 보다 많은 모델링 자원을 할당할 수 있다.

 

The overall procedure can be viewed as maximizing the evidence lower bound (ELBO) on the joint likelihood of the model distribution over images $x$, captions $y$, and tokens $z$ for the encoded RGB image. We model this distribution using the factorization $p_{\theta, \psi} = p_\theta(x | y, z)p_\psi(y, z)$, which yields the lower bound
$\ln p_{\theta, \psi}(x, y) \geqslant \mathbb E_{z \sim q_\phi(z|x)} \left( \ln p_\theta(x | y, z) - \beta D_{KL}\left(q_\phi (y, z | x), p_\psi(y, z)\right)\right), \quad (1)$ where:
- $q_\phi$ denotes the distribution over the $32 \times 32$ image tokens generated by the dVAE encoder given the RGB image $x^2$;
- $p_\theta$ denotes the distribution over the RGB images generated by the dVAE decoder given the image tokens; and 
- $p_\psi$ denotes the joint distribution over the text and image tokens modeled by the transformer.

 

전반적인 학습 과정은 이미지 $x$, 텍스트 $y$, 이미지 토큰 $z$의 결합 분포 $p(x, y, z)$ ELBO를 최대화하는 것이다. 수식을 $p(x, y, z) = p(x|y, z)p(y, z)$와 같이 factorization(=인수분해)했으며, 이는 "식 (1)"를 도출해 냈다.

 

글쓴이 뇌피셜

처음 자기회귀 트렌스포머 모델을 통해 $p(x, y)$ 모델링하려 했다. 하지만 위 2가지 한계점이 존재한다. 이 문제점들을 $256 \times 256$개 RGB 픽셀을 $32 \times 32$개 잠재 변수($z$)로 변환해주는 dVAE 모델로 극복했다. 이로 우리가 구해야 하는 결합 분포는 $p(x, y, z)$가 되었다. 이 결합 분포를 $p(x|y, z)p(y, z)$로 fractorization했다. $p(x|y, z)$는 dVAE 디코더이며 $p(y, z)$는 자기회귀 트렌스포머다.

 

글쓴이 궁금증

$\ln p_{\theta, \psi}(x, y) \geqslant \mathbb E_{z \sim q_\phi(z|x)} \left( \ln p_\theta(x | y, z) - \beta D_{KL}\left(q_\phi (y, z | x), p_\psi(y, z)\right)\right)$이 어떻게 도출되었는지 모르겠다.

 

2.1. Stage One: Learning the Visual Codebook

In the first of training, we maximize the ELBO with respect to $\phi$ and $\theta$, which corresponds to training a dVAE on the images alone. We set the initial prior $p_\psi$ to the uniform categorical distribution over the $K = 8192$ codebook vectors, and $q_\phi$ to be categorical distributions parameterized by the 8192 logits at the same spatial position in the $32 \times 32$ grid output by the encoder.

 

첫 학습 단계에서는, ELB의 $\phi$와 $\theta$를 학습한다. 다시 말해 dVAE를 학습시킨다.

사전 분포 $p_\psi$를 "$K = 8192$ codebook 벡터들의 균등 이산 분포"로 가정한다.

그리고 $q_\phi$를 "8192개 로짓에 의해 생성된 (or 매개변수화된)" 이산 분포로 가정한다. 이 연구에서는 $32 \times 32$개 잠재 변수를 사용한다. 즉, $32 \times 32$개 이산 분포를 사용한다.

 

The ELB now becomes difficult to optimize: as $q_\phi$ is a discrete distribution, and we cannot use the reparameterization gradient to maximize it. Oord et at. (2017); Razavi et al. (2019) address this using an online cluster assignment procedure coupled with the straight-though estimator (Bengio et al., 2013). We instead use the gumbel-softmax relaxation (jang et al., 2016; Maddison et al., 2016), replacing the expectation over $q_\phi$ with one over $q_\phi^\tau$, where the relaxation become tight as the temperature $\tau \rightarrow 0$. The likelihood for $p_\theta$ is evaluated using the log-laplace distribution.

 

ELB를 최적화하기 어렵다. $q_\phi$가 이산 분포라 reparameterization을 통해 미분을 할 수 없기 때문이다. (자세히 말하자면, VAE에서 $q_\phi$가 연속 확률 분포라 $\epsilon$으로 $z$를 eparameterization하여 $q_\phi$ 평균을 미분했다.)

VQ-VAE에서 디코더 입력($z_q(x)$) 기울기를 그대로 가져와 인코더 출력($z_e(x)$) 기울기로 사용했다. (VQ-VAE의 재구성 손실 함수$p_\theta(x|z_q(x))$에 $q_\phi$ 평균$\left(=\mathbb E_{z \sim q_\phi(z|x)}\right)$이 빠진 이유는 $q(z=k|x) = \begin{cases} 1 \quad \text{for} \quad k = \arg\min_j \lVert z_e(x) - e_j \rVert_2 \\ 0 \quad \text{otherwise} \end{cases}$이기 때문이다.)

하지만 이 연구에서는 gumbul-softmax relaxation을 사용한다. 다시 말해, $q_\phi$ 평균을 $q_\phi^\tau$ 평균으로 대체한다. $\left(y_i = \frac{e^\frac{g_i + \log{(q(e_i|x))}}{\tau}}{\sum_{j=1}^k e^\frac{g_j + \log{(q(e_j|x))}}{\tau}}, \quad z = \sum_{j=1}^k y_je_j\right)$ 이때 $\tau$가 $0$으로 갈수록 팽팽해집니다. 

$p_\theta$는 log-laplace 분포를 사용합니다.

 

The relaxed ELB is maximized using Adam (Kingma & Ba, 2014) with exponentially weighted iterate averaging.

 

2.2. Stage Two: Learning the Prior

 

In the second stage, we fix $\phi$ and $\theta$, and learn the prior distribution over the text and image tokens by maximizing thhe ELB with respect to $\phi$. Here, $p_\phi$ represented by a 12-billion parameter sparse transformer.

 

두번째 학습 단계에서는, ELB의 $\phi$와 $\theta$를 고정시킨 후, $\psi$ 즉, 사전 분포를 학습한다. 다시 말해 120억 파라미터 자기회귀 트렌스포머를 학습시킨다.

 

Given a text-image pair, we BPE-encode the lowercased caption using at most 256 tokens wiith vocabulary size 16384, and encode the image using $32 \times 32 = 1024$ tokens with vocabulary size 8192.

 

텍스트-이미지 쌍을 주면 256개 텍스트 토큰을 "16384개 vocabulary"로 BPE 인코딩하고 $32 \times 32$개 이미지 토큰을 "8192개 codebook"으로 인코딩한다.

 

The transformer is decoder-only model in which each image token can attend to all text tokens in any one of its 64 self-attention layers.
There are three different kinds of self-attention masks used in the model. The part of the attention masks corresponding to the text-to-text attention is the standard causal mask, and the part of the image-to-image attention uses either a row, column, or convolutional attention mask.

 

이번 연구에서는, 트렌스포머 디코더로 구성된 모델을 사용한다. 그리고 이 모델은 3가지 self-attention 마스킹을 사용한다. 1) 텍스트-이미지: 각 이미지 토큰은 모든 텍스트 토큰을 attend한다. 2) 텍스트-텍스트: 표준 마스킹 기법을 사용한다. 3) 이미지-이미지: 행, 열, 또는 합성곱 마스킹 기법을 사용한다. 

 

We limit the length of a text caption to 256 tokens, though it is not totally clear what to do for the "padding" positions in between the last text token and the start-of-image token.
we opt to learn a special padding token separately for each of the 256 text positions.

 

우리는 텍스트 토큰 개수를 256개로 한정했지만, 텍스트 마지막 토큰과 이미지 시작 토큰 사이 구간, 즉 padding 구간을 어떻게 해야 할지 명확하지 않다. 이번 연구에서는, 특별한 padding 토큰 256개를 학습하는 걸로 결정했다.

 

We normalize the cross-entropy losses for the text and image tokens by the total number of each kind in a batch of data. Since we are primarily interested in image modeling. we multiply the cross-entropy loss for text by 1/8 and the cross-entropy loss for the image 7/8.

 

이미지와 텍스트 토큰에 대한 cross-entropy 손실값을 정규화한다. 이때, 이번 연구는 이미지 토큰에 초점을 두기에, 텍스트 손실 값에 1/8 가중치를 곱하고 이미지 손실 값에 7/8 가중치를 곱한다.

 

The objective is optimizied using Adam with exponentially weighted iterate averaging.

 

2.3. Data Collection

To scale up to 12-billion parameters, we created a dataset of a similar scaled to JFT-300M by collecting 250 million text-images pairs from the internet. This dataset does not include MS-COCO, but does include Conceptual Captions and a filtered subset of YFCC100M. As MS-COCO was created from the latter, our training data includes a fraction of MS-COCO validation images. We control for this in the quantitative presented in Section 3 and find that it has no appreciable bearing on the results.

 

120억 파라미터로 모델 크기를 확장하기 위해, JFT-300M와 유사한 크기의 데이터 세트를 생성했다. 이 데이터 세트는 인터넷에서 2억 5천개 텍스트-이미지 쌍 데이터를 수집했다. 이 데이터 세트에는 Conceptual Captions와 YFCC100M 일부분을 포함한다.

MS-COCO이 최신 데이터 세트이기 때문에 MS-COCO 검증 데이터 일부분을 학습 데이터로 포함시킨다. Section 3에서 "학습 데이터에 포함된 MS-COCO 데이터 크기"를 조절하여, 이것(MS-COCO 데이터 크기)이 성능에 큰 영향이 없다는 것을 발견했다.

 

2.4. Mixed-Precision Training

To save GPU memory and increase throughput, most paramters, Adam moments, and activations are stored in 16-bit precision. We also use activation checkpointing and recompute the activations within the resblocks during the backward pass. Getting the model to train in 16-bit precision past one billion parameters, without diverging, was the most challenging part of this project. 

 

GPU 메모리를 줄이고 처리량을 높이기 위해, GPU에 할당되는 텐서들(e.g., 파라미터, Adam momentom, activation)을 16비트 정밀도로 저장했다. 뿐만 아니라 activation checkpoint을 사용해 GPU 메모리 사용랑을 줄였다.

10억 파라미터 이상 모델을 16비트 정밀도로 발산 없이 학습하는 과정이 이번 연구에서 가장 어려운 부분이였다.

 

We believe the root cause of this instability to be underflow in the 16-bit gradients.

 

학습 불안정성 근본 원인을 16비트 미분값의 underflow라 생각한다.

 

Our fix, which is shown in Figure 4, involves using a separate "gradient scale" for each resblock in the model.

 

그림4와 같이 각 resblock마다 "gradient scale"을 적용시켜 underflow문제를 해결했다.

 

2.6. Sample Generation

Similar to Razavi et al. (2019), we rerank the samples drawn from the transformer using a pretrained contrastive model. Given a caption and candidate image, the contrastive model assigns a score based on how well the image matches the caption. Figure 6 shows the effect of increasing the number of samples $N$ from which we select the top $k$ images.

 

사전 학습된 대조 모델을 통해 자기회귀 트렌스포머에서 나온 샘플 순위를 변경한다. 자세히 말하자면, 텍스트와 후보 이미지가 주어지면, 대조 모델은 후보 이미지가 텍스트와 얼마나 어울리는지 점수를 매긴다. 사진6에서 볼 수 있듯이, 샘플 개수를 늘어나면 결과값이 좋아진다.

 

3. Experiments

3.1. Quantitative Results

We evaluate our model zero-shot by comparing it to three prior approaches: AttnGAN, DM-GAN, and DF-GAN, the last of which reports the best Inception Score and Frechet Inception Distance on MS-COCO. 

 

이전 연구 방법들(e.g., AttnGAN, DM-GAN 그리고 DF-GAN)과 DALL-E를 비교할 것이다. 뿐만 아니라 DALL-E의 최고 점수를 발표할 것이다.

 

Figure 9(a) shows that our model also obtains an FID score on MS-COCO within 2 points of the best prior approach, despite having never been trained on the captions. Our training data incorporates a filtered subset of YFCC100M, and we found that it includes about 21% of the images in the MS-COCO validation set from a de-duplication procedure described in the next section. To isolate this effect, we compute the FID statistics for the validation set both with these images (solid lines) and without them (dashed lines), finding no significant change in the results.

 

MS-COCO로 학습을 안했음에도 불구하고, MS-COCO에 대해 DALL-E가 SOTA모델과 비슷한 FID점수를 얻었다.

중복제거 과정에서 "MS-COCO 검증 데이터 세트 21%가 학습 데이터에 포함되어 있음"을 발견했다. 하지만 MS-COCO 검증 데이터 중복 유무에 따른 모델 성능 차이는 크게 없다.

 

Training the transformer on the tokens from the dVAE encoder allows us to allocate its modeling capacity to the low-frequency information that makes images visually recognizable to us. However, it also disadvantages the model, since the heavy compression renders it unable to produce high-frequency details. To test the effect of this on the quantitative evaluations, we compute the FID and IS in Figure 9(a) after applying a Gaussian filter with varying radius to both the validation images and samples from the models. The gap between our approach and others tends to widen as the blur radius is increased. 

 

(dVAE 인코더에서 생성된) 토큰으로 학습하면, 트렌스포머는 "이미지 구조(low-frequency)"을 학습하는데 많은 모델링 자원을 할당할 수 있다.  하지만 이는 모델 성능에 부정적 영향을 줄 수 있다. 큰 압축으로 인해 "이미지 디테일(high-frequency)"을 학습하는데 매우 적은 모델링 자원을 사용할 수 있기 때문이다.

사실 여부 확인을 위해, 다양한 radius의 가우시안 필터를 검증 이미지와 모델 샘플에 적용해봤다. 그 결과 radius가 증가할수록, DALL-E와 다른 모델 사이의 성능 차이가 점점 증가한다.

 

글쓴이 궁금증

"radius가 증가할수록, DALL-E와 다른 모델 사이의 성능 차이가 점점 증가한다."가 "압축이 덜하면 점수가 떨어지는 건지" 아니면 "압축이 덜하면 점수가 올라가는 건지" 모르겠다.

 

Our model fares significantly worse on the CUB dataset, for which there is a nearly 40-point gap in FID between our model and the leading prior approacch (Figure 9(b)). 
We speculate that our zero-shot approach is less likely to compare favorably on specialized distributions such as CUB. We believe that fine-tuning is a promising direction for improvement, and leave this investigation to future work. 

 

CUB 데이터 세트에서는 매우 좋지 않은 결과를 얻었다. 

제로샷 접근 방법은 CUB와 같은 speicalized distribution와 궁합이 좋지 않다. 하지만 파인튜닝을 통해 성능 향상을 기대해 볼 수 있다.

 

Finally, Figure 9(c) shows clear improvements in FID and IS for MS-COCO as the sample size used for reranking with the contrastive model is increased. This trend continues up to a sample size of 32, after which we observe diminshing returns.

 

그림9(c)에서 보이는 것과 같이, 샘플 크기를 증가 시키면 성능 향상을 나타난다. 자세히 말하자면, 샘플 크기 32개까지는 성능 향상이 되지만, 그 이후로는 성능 감소가 된다.

 

 

3.3. Qualitative Findings

When given the caption "a tapir made of accordion...", the model appears to draw a tapir with an accordion for a body, or an accordion whose keyboard or bass are in the shape of a tapir's trunk or legs. This suggests that it has developed a rudimentary ability to compose unusual concepts at high levels of abstraction.

 

 DALL-E는 "생소한 개념을 높은 수준의 추상화로 구사할 수 있는 능력"이 있다. (e.g., 그림2(a))

 

Our model also appears to be capable of combinatorial generalization, such as when rendering text (Figure 2b) or when probed on sentences like "an illustration of a baby hedgehog in a christmas sweater walking a dog" (Figure 2c).

 

DALL-E는 "조합 일반화" 능력도 보여준다. 즉, 복잡한 문장도 그림으로 구사할 수 있다. (e.g., 그림2(c), 그림2(b))

 

To a limited degree of reliability, we also find our model to be capable of zero-shot image-to-image translation controllable by natural language (Figure 2d).
This works with several other kinds of transformations, including image operation (e.g., changing the color of the image, converting it to grayscale, or flipping it upside-down) and style transfer (e.g., drawing the cat on a greeting card, a postage stamp, or a cell phone case).

 

뿐만 아니라, 자연어를 통해 image-to-image 변환도 가능하다. (e.g., 그림2(d)) 자세히 말하자면, 여려 종류의 image-to-image 변환이 가능하다. (e.g., 이미지 색깔 변환, 위 아래 전환, 스타일 변환)

 

4. Conclusion

We investiage a simple approach for text-to-image generation based on an autoregressive transformer, when it is executed at scale. 

 

이번 연구에서는, text-to-image 생성에 대한 간단한 접근법(=초거대 자기회귀 트렌스포머)을 제안한다.

 

We find that scale can lead to improved generalization, both in term of zero-shot performance relative to to previous domain-specific approaches, and in terms of the range of capabilities that emerge from a single generative model.

 

이번 연구에서는, "scaling(e.g., 모델 크기, 데이터 크기, 계산량)을 통해 1). 학습 과정(=제로샷) 및 2). 모델(=단일 생성 모델) 일반화가 가능하다는 것"을 보여주었다.