DALL-E는 2021년 OpenAI에서 발표한 Image Generation 모델입니다. 초현실주의 미술가인 Salvador Dali와 WALL-E를 재치있게 합성한 이름을 가지고 있습니다. 세상에 존재하지 않는 것( 이를테면, 아보카도 모양의 의자 )에 대한 설명으로부터 높은 퀄리티의 이미지를 생성해내는 것으로 전세계적 유명세를 얻고 있습니다!
Introduction
Text-to-Image Generation Task를 위한 데이터셋은 이미지와 그에 맞는 캡션을 부여해야 하기 때문에 구축에 많은 비용이 듭니다. 이에 많은 연구자들은 이미 라벨링이 되어있는 오픈 데이터(MS-COCO 혹은 CUB-200)으로 학습 데이터를 고정한 뒤, 복잡한 구조의 모델과 부가적인 Loss, 혹은 Segmentation Mask 같은 추가 정보를 활용하는 방식으로 성능의 향상을 꾀하였습니다.
본 논문은 아주 심플하게 충분히 큰 Transformer에 충분히 많은 데이터를 학습시킵니다. 그 결과, Domain-specific한 모델들과 Zero-shot으로 경쟁할 수 있는 멋진 모델이 만들어졌습니다!
Method
저자들은 Transformer 모델에 Text와 Image를 하나의 스트림에서 학습할 수 있게 하고자 하였습니다. 하지만 이를 위해 이미지 데이터를 픽셀 단위로 처리하는 것은 고화질 이미지에서 요구될 메모리를 생각하면 바람직하지 못합니다. 또한 작은 단위에서 Likelihood를 목적으로 학습하다보면 이미지를 좁은 부분 밖에 보지 못하는 Short-range Dependency가 발생합니다. 다시 말해, 이미지를 인식하는데 크게 도움이 되지 않는 부분을 위해 많은 자원이 낭비되는 것입니다. 이에 저자들은 Two-stage Training Procedure를 통해 이 문제를 해결하고자 하였습니다.
Stage 1
먼저 dVAE(discrete Variational AutoEncoder)를 훈련합니다. dVAE는 256 x 256 이미지를 32 x 32의 Image Tokens로 변환하는 역할을 합니다. 각 Token들은 8192개의 Codeword가 되어 Visual Codebook에 업데이트됩니다. 이를 이 방법을 통해 이미지의 큰 손실 없이 Transformer 학습에 필요한 Context를 192배 줄일 수 있습니다. 필연적으로 이미지는 손상되나 그 정도가 심하지 않다는 것을 아래 그림에서 확인할 수 있습니다.
Comparison of original images (top) and reconstructions from the discrete VAE (bottom). The encoder downsamples the spatial resolution by a factor of 8. While details (e.g., the texture of the cat’s fur, the writing on the storefront, and the thin lines in the illustration) are sometimes lost or distorted, the main features of the image are still typically recognizable. We use a large vocabulary size of 8192 to mitigate the loss of information.
단, dVAE가 Codebook에서 Image Tokens을 Discrete하게 처리한다는 특성을 가지고 있기 때문에, 일반적인 방법으로는 Gradient가 흐를 수 없어 Back Propagation이 불가능합니다. 이에 저자들은 Gumbel Softmax를 적용하였습니다. Gumbel Softmax는 Gumbel-max Trick과 Softmax의 결합으로 tau값에 따라 One-hot과 Uniform을 오가는 값으로 Reparameterization할 수 있게 하는 함수입니다. tau가 0에 가까울수록 One-hot에 가까운 값이 나오고 tau가 커질수록 Uniform에 가까워집니다.
실제 학습에서는 tau 값을 처음 150,000 Iteration동안 1에서 1/16까지 선형적으로 감소시켰습니다.
Stage 1 학습엔 총 64개의 16GB NVIDIA V100 GPU를 사용하였고 GPU당 Batch Size는 8을 사용했습니다. 결과적으로 512 Batch Size를 사용한 셈이며 총 3,000,000번의 Iteration동안 학습하였습니다.
Stage 2
후에 이미지에 대한 캡션을 BPE(Byte Pair Encoding)로 쪼갭니다. 최대 256개의 Text Tokens와 32 x 32 = 1024개의 Image Tokens를 Concatenate하여 Stream을 완성합니다. 이 Stream은 Transformer에 입력으로 사용되어 Image와 Text의 Joint Distribution을 학습하게 됩니다.
Transformer은 Decoder만 사용하며 64개의 Self Attention Layer를 갖습니다. 각 Layer는 62개의 Attention Head를 갖습니다. 독특한 점은 Attention Masking에 여러 종류의 Mask를 사용했다는 점입니다. Text-to-Text에는 기본적인 Causal Mask를 사용하고 Image-to-Image에는 Row, Column, Convolution Mask를 각각 정의하였습니다.
Illustration of the three types of attention masks for a hypothetical version of our transformer with a maximum text length of 6 tokens and image length of 16 tokens (i.e., corresponding to a 4 × 4 grid). Mask (a) corresponds to row attention in which each image token attends to the previous 5 image tokens in raster order. The extent is chosen to be 5, so that the last token being attended to is the one in the same column of the previous row. To obtain better GPU utilization, we transpose the row and column dimensions of the image states when applying column attention, so that we can use mask (c) instead of mask (b). Mask (d) corresponds to a causal convolutional attention pattern with wraparound behavior (similar to the row attention) and a 3 × 3 kernel. Our model uses a mask corresponding to an 11 × 11 kernel.
그림의 (c)의 경우는 GPU 메모리를 위해 이미지를 Transpose 했을 경우에 적용되는 Column Attention Mask이고 (d)는 Convolution 형태로 Attend할 수 있게 하는 Mask입니다. 각 Mask들은 일련의 규칙에 따라서 적용됩니다.
- 맨 마지막 Attention Layer에는 Convolution Mask가 적용됩니다.
- 나머지 [1, 63] 범위의 Attention Layer Index 중,
index - 2 % 4 == 0일 경우 Column Mask를 적용합니다.
- 그 외에는 Row Mask를 적용합니다.
위와 같은 규칙에 따라 Mask를 부여하면 첫 4층의 레이어에 대해서 " 1: Row, 2: Column, 3: Row, 4: Row " 의 순서가 됩니다.
또, Attention을 할 때에 Padding 부분에 -inf를 더하여 Softmax 결과를 0으로 만드는 것이 일반적이나, 저자들의 Stream에서 Text Padding 부분은 Last Text Token과 Start Image Token의 사이에 위치하기에 그 방법이 Clear하지 않다고 말합니다. 이에 저자들은 공백을 학습하는 Special Padding Token을 정의하였습니다. 이는 Loss 값 자체는 증가하지만 캡션이 Out of Distribution일 경우 더 나은 성능을 보여줍니다.
Stage 2 학습엔 총 1024개의 16GB NVIDIA V100 GPU를 사용하였고 GPU당 Batch Size는 1을 사용했습니다. 결과적으로 1024 Batch Size를 사용한 셈이며 총 430,000번의 Iteration동안 학습하였습니다.
전체 학습 프로세스는 이미지 x와 캡션 y, 토큰 z에 대한 ELBO(Evidence Lower Bound)를 최대화하는 것으로 볼 수 있습니다. 따라서 최종 목적 함수는 아래와 같이 계산됩니다.
여기서 Beta는 1일 때 Lower Bound가 성립하지만, 실제로는 더 큰 값을 사용하는게 좋은 것을 확인하였다고 합니다. 실제 학습에서는 처음 5000 Iteration동안 0부터 6.6까지 값을 증가시키며 학습하였습니다.
Data Collection
Text-to-Image Generation에서 주로 사용되는 데이터셋인 MS-COCO와 CUB-200은 각각 330K Pairs, 6033 Pairs로 ImageNet이 약 14M Images를 제공하는 것에 비하면 굉장히 작은 수준입니다.
이에 저자들은 Conceptual Captions(3.3M), YFCC100M(100M)과 Wikipedia로부터 250M Pairs를 수집하였습니다. 수집은 무분별하게 진행하지 않고 아래 기준에 맞게 필터링하였습니다.
- 캡션이 너무 짧은 경우
- 캡션이 영어로 작성되지 않은 경우
- 캡션이 사진의 촬영 날짜 같은 상용구인 경우
- 이미지의 Ratio가 [1/2, 2]를 벗어나는 경우
Mixed-Precision Training
저자들은 대다수의 파라미터와 Adam Moments, Activation을 16-bit Precision으로 저장하여 GPU 메모리를 절약하고 Throughput을 늘렸습니다. 약 10억 개의 파라미터를 가진 모델을 16-bit Gradient로 훈련시키려면 Underflow를 피하기 어렵고 이를 해결하는 것이 도전이었다고 말합니다.
ResBlock에서 생기는 Gradient는 ResBlock에서 멀어질수록 단조롭게(Monotonically) 감소하게 되는데, 모델이 넓어지고 깊어지면 결국엔 Gradient가 16-bit의 최소 지수보다 작아질 수 있습니다. 이 때 Gradient가 0으로 반올림되는 현상이 발생하는데 이를 Underflow라 칭합니다.
이를 해결하기 위해 위 그림과 같은 프로세스로 Gradient Scaling을 진행합니다.
Distributed Optimization
16-bit Precision으로 저장된 모델은 24GB 정도의 메모리를 차지하는데. NVIDIA V100 GPU의 16GB를 초과하는 크기입니다. 이에 저자들은 파라미터 샤딩(Sharding)을 적용하여 모델을 다수의 GPU에 올리고 기기 간에 통신에 요구되는 Latency를 거의 완전하게 없앨 수 있었습니다.
Sample Generation
저자들은 DALL-E와 함께 발표한 CLIP이라는 모델을 활용하여 DALL-E가 생성한 이미지에 점수를 부여했습니다. 해당 점수를 지표로 Top-k 이미지를 선정할 수 있게 되었고, Sample을 많이 Generation할수록 좋은 이미지를 얻을 확률이 높아지게 되었습니다. 위 그림에서 Sample 수에 따른 이미지의 퀄리티 차이를 확인할 수 있습니다.
Experiments
Quantitative Result
위 그림은 MS-COCO 데이터셋에서 AttnGAN, DM-GAN, DF-GAN을 Zero-shot DALL-E와 비교한 결과를 나타냅니다. MS-COCO 데이터를 전혀 본 적이 없음에도 뛰어난 성능을 확인할 수 있습니다.
저자들은 추가적으로 DF-GAN과 DALL-E가 생성한 이미지에 대해 Human Evaluation을 진행하였고 결과는 좌측과 같습니다. Evaluation은 아래와 같은 문제지를 푸는 방식으로 진행되었습니다.
단, CUB Dataset에 대해서는 좋지 못한 성능을 보였습니다. DALL-E는 웹으로부터 수집한 다양한 Domain의 데이터로 학습이 되었으나 CUB Dataset은 오로지 새라는 Domain에 치중되어 있는 Specialized Distribution이기 때문이라고 저자들은 설명합니다.
Qualitative Findings
저자들은 DALL-E로부터 예상하지 못한 능력을 발견했다고 말합니다.
" 아코디언으로 만든 테이퍼 " 나 " BACKPROP 이라고 적힌 네온 사인 " 등을 캡션으로 사용했을 때 실제로 그 이미지를 그려주는 것인데, 이는 두 개념을 합성하는 기초적인 능력이 개발되었음을 시사합니다. 또, 이미지 토큰의 절반과 그에 대한 캡션을 제공함으로써 Image-to-Image Translation을 수행할 수 있는 것도 확인되었습니다. (d)에서 위의 고양이 사진을 토대로 아래에 스케치를 그려낸 것이 그 예입니다.Conclusion
저자들은 Image Generation을 아주 심플하게, 하지만 여태 시도된 적 없었던 Scale로 실험을 했고 그 결과 기존의 Domain-specific한 모델들을 Zero-shot으로 능가하는 의미있는 Improvement를 이끌었습니다. 하지만 아직은 생성된 이미지의 퀄리티가 Improve될 여지가 있고, 모델의 크기가 굉장히 크기 때문에 이를 일반화하는 것이 향후에 연구될 수 있겠습니다.
Share article