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 mainly investigates how to further improve the image generation quality of previous work VQGAN, where several modifications are proposed to train a better quantized auto-encoder model. After modeling the discrete tokens with an auto-regressive transformer, we could observe that the generation results of the proposed ViT-VQGAN are really amazing and beat most of previous works in various benchmarks. Meanwhile, ViT-VQGAN also demonstrates its exceptional representation capabilities through unsupervised linear-probing.  To summarize, it seems there are several key factors to obtain a better reconstruction model:  - Project the latent codes into lower dimension, which leads to a larger codebook size and less ‘dead’ codes. - StyleGAN Discriminator. - $\ell_{2}$ normalization on latent variables. - Stronger backbone (ViT-like architecture or decoder with more parameters).