Vector-quantized Image Modeling with Improved VQGAN

Pretraining language models with next-token prediction on massive text corpora has delivered phenomenal zero-shot , few-shot , transfer learning and multi-tasking capabilities on both generative and discriminative language tasks . Motivated by this success , we explore a Vector-quantized Image Modeling ( VIM ) approach that involves pretraining a Transformer to predict rasterized image tokens autoregressively . The discrete image tokens are encoded from a learned Vision-Transformerbased VQGAN ( ViT-VQGAN ) . We first propose multiple improvements over vanilla VQGAN from architecture to codebook learning , yielding better efficiency and reconstruction fidelity . The improved ViT-VQGAN further improves vectorquantized image modeling tasks , including unconditional , class-conditioned image generation and unsupervised representation learning . When trained on ImageNet at 256 × 256 resolution , we achieve Inception Score ( IS ) of 175.1 and Fréchet Inception Distance ( FID ) of 4.17 , a dramatic improvement over the vanilla VQGAN , which obtains 70.6 and 17.04 for IS and FID , respectively . Based on ViT-VQGAN and unsupervised pretraining , we further evaluate the pretrained Transformer by averaging intermediate features , similar to Image GPT ( iGPT ) . This ImageNet-pretrained VIM-L significantly beats iGPT-L on linear-probe accuracy from 60.3 % to 73.2 % for a similar model size . ViM-L also outperforms iGPT-XL which is trained with extra web image data and larger model size . 1 INTRODUCTION . Natural language processing ( NLP ) has recently experienced dramatic improvements from learning general-purpose representations by pretraining language models on unlabeled text corpora . This strategy has produced large performance gains for a wide range of natural language generation ( NLG ) and natural language understanding ( NLU ) tasks ( Dai & Le , 2015 ; Radford et al. , 2018 ; 2019 ; Brown et al. , 2020 ) . Conceptually , generative pretraining models the data density P ( X ) in a tractable way , with the hope of also helping discriminative tasks of P ( Y |X ) ( Lasserre et al. , 2006 ) ; importantly , there are no limitations on whether the signals are from the language domain or others , such as vision . In computer vision , in contrast , most recent unsupervised or self-supervised learning research focuses on applying different random augmentations to images , with the pretraining objective to distinguish image instances ( Chen et al. , 2020b ; He et al. , 2020 ; Chen et al. , 2020d ; Grill et al. , 2020 ; Chen et al. , 2020c ; Caron et al. , 2021 ) . The quality of learned representation relies on manually chosen augmentations , such as random brightness , cropping , blurring , and others . Chen et al . ( 2020a ) explored GPT-style ( Radford et al. , 2018 ) generative pretraining on images by autoregressively predicting pixels without incorporating knowledge of the 2D structure . Each pixel is represented as a 9-bit value created by clustering ( R , G , B ) pixel values , using k-means with k=512 . Unfortunately , this color encoding does not scale to typical image resolutions as it entails very long sequences to represent the image ( e.g. , 224× 224 resolution leads to 50,176 tokens per image ) , and this demands much more memory and computation for training , compared to language models . As a result , iGPT ’ s maximum resolution is 64 × 64 for image recognition at scale—which severely limits its representation capabilities . Remarkable image generation results have been achieved by pre-quantizing images into discrete latent variables and modeling them autoregressively , including VQVAE ( Oord et al. , 2017 ) , DALLE ( Ramesh et al. , 2021 ) and VQGAN ( Esser et al. , 2021 ) . In these approaches , a convolution neural network ( CNN ) is learned to auto-encode an image and a second stage CNN or Transformer is learned to model the density of encoded latent variables . These have been proved effective for image generation , but few studies have evaluated the learned representation in discriminative tasks ( Ramesh et al. , 2021 ; Esser et al. , 2021 ) . We explore an approach we refer to as Vector-quantized Image Modeling ( VIM ) and apply it to both image generation and image understanding tasks . VIM follows a two-stage approach : • Stage 1 : Image Quantization . Given an image of resolution 256×256 , a VisionTransformer-based VQGAN encodes it into 32×32 discretized latent codes where the codebook size is 8192 . We propose multiple improvements–from architecture to codebook learning–to VQGAN ( Esser et al. , 2021 ) . The resulting ViT-VQGAN is more efficient and improves reconstruction fidelity in terms of pixel-wise reconstruction metrics , Inception Score ( IS ) and Fréchet Inception Distance ( FID ) . ViT-VQGAN is trained end-to-end on image-only data with combined objective functions of logit-laplace loss , ` 2 loss , adversarial loss and perceptual loss ( Johnson et al. , 2016 ; Zhang et al. , 2018 ) . • Stage 2 : Vector-quantized Image Modeling . We train a Transformer model to predict rasterized 32×32 = 1024 image tokens autoregressively , where image tokens are encoded by a learned Stage 1 ViT-VQGAN . For unconditional image synthesis or unsupervised learning , we pretrain a decoder-only Transformer model to predict the next token . For class-conditioned image synthesis , a class-id token is prepended before the image tokens . To evaluate the quality of unsupervised learning , we average the intermediate Transformer features and learn a linear head to predict the logit of the classes ( a.k.a. , linear-probe ) . We show that one key component for improving both image generation and image understanding with VIM is to have a better image quantizer with respect to both computational efficiency and reconstruction quality . An efficient quantizer can speed up Stage 2 training , where random augmentations are applied first to an image , followed by the encoder of image quantizer to obtain the input tokens . Moreover , an image quantizer with better reconstruction quality can reduce information loss compared with the original image in pixel space , which is critical for image understanding tasks . The evaluations of our proposed ViT-VQGAN and VIM are studied with three aspects . ( 1 ) We evaluate the image quantizer based on reconstruction quality metrics including ` 1 distance , ` 2 distance , log-laplace distance , as well as Inception Score ( IS ) and Fréchet Inception Distance ( FID ) of reconstructed images . ( 2 ) We evaluate the capabilities of the learned quantizer for unconditional or class-conditioned image synthesis based on FID and IS , and compare with other methods . ( 3 ) We rely on linear-probe accuracy to evaluate representations with the common intuition that good features should linearly separate the classes of downstream tasks . 2 RELATED WORK . Image Synthesis . Image generation has received much attention with the progress of deep generative models , including Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ; Zhang et al. , 2019b ) , Variational Autoencoders ( VAEs ) ( Kingma & Welling , 2014 ; Vahdat & Kautz , 2020 ) , Diffusion Models ( Song & Ermon , 2019 ; Dhariwal & Nichol , 2021 ) and Autoregressive Models ( van den Oord et al. , 2016 ; Parmar et al. , 2018 ) . Unlike many autogressive methods which generate sequence directly in pixel space , VQVAE ( van den Oord et al. , 2017 ; Razavi et al. , 2019 ) decomposes the image generation process into two stages : the first stage trains a vector quantized autoencoder with image reconstruction objective to convert an image into a shorter sequence of discrete tokens . Then the second stage learns an autoregressive model , e.g. , PixelSNAIL ( Chen et al. , 2018 ) , to model the underlying distribution of token sequences . Driven by the effectiveness of VQVAE and progress in sequence modeling ( Vaswani et al. , 2017 ; Devlin et al. , 2019 ) , many approaches follow the two-stage paradigm . DALL-E ( Ramesh et al. , 2021 ) improves token prediction in second stage by using Transformers ( Vaswani et al. , 2017 ) , resulting in a strong text-to-image synthesis model . VQGAN ( Esser et al. , 2021 ) further uses adversarial loss and perceptual loss ( Johnson et al. , 2016 ; Zhang et al. , 2018 ) to train a better autoencoder in the first stage to synthesize greater detail in images . Image Recognition with Generative Pretraining . Many image generation models ( Goodfellow et al. , 2014 ; Kingma & Welling , 2014 ; Radford et al. , 2016 ; Donahue et al. , 2017 ; Higgins et al. , 2017 ) have been studied for their capabilities in representation learning . However , their performance is usually not superior to competing self-supervised approaches that solve auxiliary classification tasks ( Noroozi & Favaro , 2016a ; Gidaris et al. , 2018a ; van den Oord et al. , 2018 ) . BigBiGAN ( Donahue & Simonyan , 2019a ) first demonstrated that a generation-based model can match other self-supervised methods in representation learning on ImageNet . iGPT ( Chen et al. , 2020a ) uses the autoregressive objective to learn a giant transformer that directly predicts pixel values , producing even more competitive results . Compared to iGPT , our method first tokenizes the original image into discrete image tokens and then trains a transformer to predict them . As a result , our approach obtains comparable results with smaller model and less data . Similar to our method in predicting image tokens , BEiT ( Bao et al. , 2021 ) follows pre-training scheme of BERT Devlin et al . ( 2019 ) by learning to recover randomly masked image tokens with a bidirectional transformer . Unlike BEiT , we explore vector-quantized image modeling for image generation in addition to image recognition . 3 VECTOR-QUANTIZED IMAGES WITH VIT-VQGAN . The Vector-quantized Variational AutoEncoder ( VQVAE ) ( van den Oord et al. , 2017 ) is a CNNbased auto-encoder whose latent space is a matrix of discrete learnable variables , trained end-to-end via straight-through estimation . Esser et al . ( 2021 ) introduce VQGAN , a model which improves upon VQVAE by introducing an adversarial loss produced by a discriminator . Below , we introduce further improvements to VQGAN that boost efficiency and enhance reconstruction quality . 3.1 VQGAN WITH VISION TRANSFORMERS . The core network architectures used by both VQVAE and VQGAN to encode and reconstruct images are CNNs . VQGAN introduces transformer-like elements in the form of non-local attention block ( Zhang et al. , 2019a ) , allowing it to capture distant interactions with fewer layers . We propose taking this approach one step further by replacing the CNN encoder and decoder with Vision Transformer ( ViT ) ( Dosovitskiy et al. , 2020 ) , as shown in Figure 1 . Given sufficient data ( for which unlabeled image data is plentiful ) we find that ViT-VQGAN is less constrained by the inductive priors imposed by convolutions . Furthermore , ViT-VQGAN yields better computational efficiency on accelerators , and produces higher quality reconstructions , as shown in Table 1 . The encoder of ViT-VQGAN first maps 8×8 non-overlapping image patches into image tokens , followed by Transformer blocks , encoding a 256×256 resolution image into a 32×32=1024 token sequence . The decoder performs the inverse operation , mapping each image token from latent variables back to 8×8 image patches and regrouping them into a 256×256 image ( see Figure 1 ) . At the output of transformer blocks , we apply a two-layer feed-forward network with a tanh activation layer in the middle . No activation is applied at the output of ViT-VQGAN encoder or decoder . We find that this simple approach yields high quality reconstructions without any noticeable grid artifacts .

Summary: This paper proposes vision transformer based VQGAN whose encoder and decoder are implemented as transformers rather than standard CNNs. It provides few improvements that demonstrates quantitatively better results in multiple datasets in unconditional / conditional image generation. It also shows how the proposed model can perform well as an unsupervised representation learning framework.