Not All Image Regions Matter: Masked Vector Quantization for Autoregressive Image Generation

Mengqi Huang, Zhendong Mao, Quan Wang, Yongdong Zhang

Introduction

Deep generative models of images have received significant improvements over the past few years and broadly fall into two categories: likelihood-based models, which include VAEs, flow-based, diffusion models and autoregressive models, and generative adversarial networks (GANs), which use discriminator networks to distinguish samples from generator networks and real examples. Compared with GANs, likelihood-based models’ training objective, i.e., the negative log-likelihood (NLL) or its upper bound, incentives learning the full data distribution and allows for detecting overfitting.

Among the likelihood-based models, autoregressive models have recently attracted increasing attention for their impressive modeling ability and scalability. Recent autoregressive image generation follows the two-stage generation paradigm, i.e., the first stage learns a codebook in the latent space for image reconstruction and the second stage completes the image generation in the raster-scan order by autoregressive models based on the learned codebook. Since codebook learning in the first stage defines the discrete image representation for the next autoregressive modeling, a high-quality codebook is the key to generate high-quality images. Several recent works focus on improving the codebook learning in the first stage, e.g., VQGAN introduces adversarial loss and perceptual loss. ViT-VQGAN introduces a more expressive transformer backbone. RQ-VAE introduces the residual quantization to reduce the resolution of the latent space. In general, the essence of existing codebook learning is the modeling of all local region information (i.e., an $8\times 8$ or $16\times 16$ patch) of images in the dataset, without distinguishing their different perceptual importance.

In this study, we point out that existing codebook learning exists gaps with classical image coding theory, the basic idea of which is to remove redundant information by perceiving the importance of different regions in images. The image coding theory reveals that an ideal image coding method should only encode images’ perceptually important regions (i.e., which cannot be restored if missing) while discarding the unimportant ones (i.e., which can be restored by other image regions even if missing). The neglect of considering such perceptual importance in existing works poses problems in two aspects, as illustrated in Figure 1(a): (1) the existence of this large amount of repetitive and redundant information brings redundancy to the learned codebook, which further makes the autoregressive model in the next stage overly focus on modeling this redundancy while overlooking other important regions and finally degrades generation quality. (2) the redundancy makes the autoregressive model need to predict more (redundant) quantized codes to generate images, which significantly increases the training cost and decreases the generating speed. Although the effectiveness and efficiency of image coding theory have been widely validated, how to introduce this idea into codebook learning remains unexplored.

The key of applying image coding theory to codebook learning is to distinguish important image parts from unimportant ones correctly. Considering that the essential difference between these two sets lies in whether they can be restored if missing, we found that this distinction can be realized through the mask mechanism, i.e., the masked part is important if it cannot be faithfully restored, and otherwise unimportant. Based on the above observation, we thereby propose a novel two-stage generation paradigm upon the mask mechanism to relieve the model from modeling redundant information. Specifically, we first propose a Masked Quantization VAE (MQ-VAE) with two novel modules, i.e., an adaptive mask module for adaptively masking redundant region features before quantization, and an adaptive de-mask module for adaptively recovering the original grid image feature map to faithfully reconstruct original images after quantization. As for the adaptive mask module, it incorporates a lightweight content-aware scoring network that learns to measure the importance of each image region feature. The features are then ranked by the importance scores and only a subset of high-scored features will be quantized further. As for the adaptive de-mask module, we design a direction-constrained self-attention to encourage the information flow from the unmasked regions to the masked regions while blocking the reverse, which aims to infer the original masked region information based on unmasked ones. Thanks to the adaptive mask and de-mask mechanism, our MQ-VAE removes the negative effects of redundant image regions and also shortens the sequence length to achieve both effectiveness and efficiency.

Moreover, since different images have different important regions, the position of quantized codes in the feature map also dynamically changed. Therefore, we further propose Stackformer for learning to predict the combination of both codes and their corresponding positions. Concretely, the proposed Stackformer stacks a Code-Transformer and a Position-Transformer, where the Code-Transformer learns to predict the next code based on all previous codes and their positions, and the Position-Transformer learns to predict the next code’s position based on all previous codes’ positions and current code.

With our method, as shown in Figure 1(b), the codebook learning only includes the important regions, e.g., the structural regions, since unimportant ones like the background can be restored even if missing. And therefore the autoregressive model in the second stage could focus on modeling these important regions and brings better generation quality.

In a nutshell, we summarize our main contributions as:

Conceptually, we point out that existing codebook learning ignores distinguishing the perceptual importance of different image regions, which brings redundancy that degrades generation quality and decreases generation speed.

Technically, (i) we propose MQ-VAE with a novel adaptive mask module to mask redundant region features before quantization and a novel adaptive de-mask module to recover the original feature map after quantization; (ii) we propose a novel Stackformer to predict the combination of both codes and their corresponding positions.

Experimentally, comprehensive experiments on various generations validate our effectiveness and efficiency, i.e., we achieve 8.1%, 2.3%, and 18.6% FID improvement on un-, class-, and text-conditional state-of-the-art at million-level parameters, and faster generation speed compared to existing autoregressive models.

Related Work

Autoregressive models for image generation have recently attracted increasing research attention and have shown impressive results among various generation tasks. Early autoregressive models directly optimizing the likelihood of raw image pixels, e.g., Image-GPT trains a transformer to autoregressively predict pixels’ cluster centroids, which could only generate images with a maximum resolution of $64\times 64$ . presents the Vector Quantized Variational Autoencoder (VQVAE), which learns images’ low-dimension discrete representation and models their distribution autoregressively. VQ-VAE2 extends this approach using a hierarchy of discrete representations. VQGAN further improves the perceptual quality of reconstructed images using adversarial and perceptual loss. ViT-VQGAN introduces a more expressive transformer backbone. RQ-VAE uses Residual Quantization to iteratively quantizes a vector and its residuals and represent the vector as a stack of tokens. Although vector quantization has become the fundamental technique for modern visual autoregressive models, the critical removing redundancy in codebook learning has not been explored yet, which becomes a critical bottleneck.

2 Masked Modeling

Masked modeling is popular among both natural language processing and computer vision. BERT randomly masks a portion of the input sequence and trains models to predict the missing content. In the computer vision domain, the ViT studies masked patch prediction for self-supervised learning. BEiT proposes to predict discrete tokens. Most recently, MaskGIT also used the masking strategy for VQ-based image generation. However, our proposed method differs from MaskGIT in two aspects: (1) Our primary motivation for the masking strategy applied in the proposed MQ-VAE in the first stage aims to learn a more compact and effective vector quantization (VQ) itself by masking perceptual unimportant regions, while MaskGIT uses masking strategy in the second stage to better use a learned VQ. (2) The mask in our proposed MQ-VAE is learned and adaptively changed according to different image content, while the mask in MaskGIT is randomly sampled for the mask-and-predict training. In conclusion, to the best of our knowledge, this is the first time that masked modeling has been applied for vector quantization.

Methodology

We propose a novel two-stage framework with MQ-VAE and Stackformer for autoregressive image generation, as illustrated in Figure 2. MQ-VAE only masks redundant region features to prevent redundant codes and Stackformer stacks two transformers to autoregressively predict the next code and its position. In the following, we will first briefly revisit the formulation of vector quantization and then describe our proposed method in detail.

2 Stage 1: MQ-VAE

Existing methods quantize each feature vector of $\mathbf{Z}$ without distinguishing their different perceptual importance and thus bring redundancy in the learned codebook, which not only degrades the generation quality but also decreases the generation speed. To relieve the model from this redundancy, we propose MQ-VAE with two novel modules, i.e., the adaptive mask module for adaptively masking redundant region features before vector quantization and adaptive de-mask module for adaptively recovering the original grid image feature map after vector quantization.

The larger score $s_{l}$ is, the more important the region feature $\boldsymbol{z}_{l}$ is. Then the region features are sorted in descending order according to the predicted scores. The sorted region features and their scores are denoted as $\{\boldsymbol{z}^{{}^{\prime}}_{l}\}$ and $\{s^{{}^{\prime}}_{l}\}$ respectively, where $l=1,...,L$ . To enable the learning of $f_{s}$ , the predicted scores are further multiplied with the normalized region features as modulating factors. We select the top $N$ scoring vectors as the important region features,

Here, $\hat{\mathbf{Z}}$ denotes the selected important region features set, and $\hat{\mathbf{P}}$ denotes the corresponding position set that represents the position of each selected region feature in the original 2D feature map. The selected number $N=\alpha\times L$ , where $\alpha$ is a constant fractional value. The mask ratio is defined as $1-\alpha$ . This design also enables a flexible trade-off between the image generation speed and image generation quality, which we will discuss in experiments. After obtaining $\hat{\mathbf{Z}}$ , we further apply the quantization function $\mathcal{Q}$ to each of them and obtain the quantized important region features set $\hat{\mathbf{Z}}^{\boldsymbol{q}}$ as well as its code matrix $\hat{\mathbf{M}}$ .

Adaptive De-mask Module. After quantization, we fill the quantized features $\hat{\mathbf{Z}}^{\boldsymbol{q}}$ back into the original 2D feature map according to $\hat{\mathbf{P}}$ , while other masked positions are filled with a uniformly initialized learnable mask code embedding, as shown in Figure 2(a). Directly inputting filled grid features to the decoder $D$ could bring sub-optimal reconstruction results since the mask code embedding here only serves as the placeholders that contain little information. Therefore, we further propose the adaptive de-mask module, which applies a novel direction-constrained self-attention to encourage the information flow from unmasked regions to the masked ones while blocking the reverse. Such a design allows the model to utilize the unmasked region features to infer the masked ones while also preventing the masked regions to have negative impacts on the unmasked ones since they are less informative.

Our adaptive de-mask module is implemented as $H$ identical sub-modules, where each consists of a direction-constrained self-attention block and a Resnet block. The direction-constrained self-attention is mathematically formed as (Resnet block is omitted for simplicity):

3 Stage 2: Stackformer

The perceptual important regions of different images vary. Therefore the positions of quantized codes in the feature map also dynamically change along with the image content. As a result, our proposed MQ-VAE formulates an image as both the code sequence $\hat{\mathbf{M}}$ and the code position sequence $\hat{\mathbf{P}}$ . To simultaneously learn the combination of the codes and their positions, we propose a novel Stackformer, which stacks a Code-Transformer and a Position-Transformer. The Code-Transformer learns to predict the next code based on all previous steps’ codes and their positions, while the Position-Transformer learns to predict the next code’s position based on all previous steps’ positions and current code. Directly treating the importance descending order sequence $\hat{\mathbf{M}}$ and $\hat{\mathbf{P}}$ as the inputs are natural, but the dramatic position changes of adjacent code could make the network hard to converge. For example, the position of the first code may be in the upper left corner of the image, while the position of the second code may be in the lower right corner of the image. Therefore, we further propose to use the raster-scan order to rearrange both sequences to deal with the converge problem.

Mathematically, taking the raster-scan code and code position sequence $(\overline{\mathbf{M}},\overline{\mathbf{P}})=\text{rearrange}(\hat{\mathbf{M}},\hat{\mathbf{P}})$ , Stackformer learns $p(\overline{\mathbf{M}},\overline{\mathbf{P}})$ , which is autoregressively factorized as:

Code-Transformer takes the sum of code embeddings $\boldsymbol{e}_{c}(\cdot)$ and code position embedding $\boldsymbol{e}_{p}(\cdot)$ as inputs:

where $N_{c}$ is the condition length. For the unconditional generation, we add a $<$ sos $>$ code at the start of the code and code position sequence. For conditioning, we append class or text codes to the start of the code sequence and the same length of $<$ sos $>$ code to the code position sequence. We further add an extra learned absolute position embedding to $\mathbf{U}_{c}$ to form the final input, which makes the network aware of the absolute position of the sequence. After processing by Code-Transformer, the output hidden vector $\mathbf{H}_{c}$ encodes both code and their position information and is used for the next code prediction. The negative log-likelihood (NLL) loss for code autoregressive training is:

Position-Transformer takes the sum of Code-Transformer’s output hidden vector $\mathbf{H}_{c}$ and an extra code embedding as input:

Here $\mathbf{U}_{p}$ is the input for Position-Transformer and the information of current code is included in $\boldsymbol{e}_{c}(\overline{\mathbf{M}}_{[N_{c}+1:N_{c}+N]})$ . The design idea behind this is that when predicting the next code’s position, the model should not only be aware of previous steps’ codes and their position information but also should be aware of current code information. The negative log-likelihood (NLL) for position autoregressive training is:

Training $\&$ Inference. The total loss for training Stackformer is defined as:

The inference procedure is illustrated in Algorithm 1, where we take the unconditional generation as an example and conditional generation can be derived accordingly.

Experiment

Benchmarks. We validate our model for unconditional, class-conditional, and text-conditional image generation tasks on FFHQ, ImageNet, and MS-COCO benchmarks respectively, with $256\times 256$ image resolution.

Metrics. Following previous works, the standard Frechet Inception Distance (FID) is adopted for evaluating the generation and reconstruction quality (denoted as rFID). Inception Score (IS) is also used for class-conditional generation on the ImageNet benchmark. FID and IS are calculated by sampling 50k images. rFID is calculated over the entire test set.

Implemented details. The architecture of MQ-VAE exactly follows except for the proposed mask and de-mask modules, with the codebook size of $K=1024$ . For the de-mask module, the sub-module number $H=8$ . For the Stackformer, we implement two settings: a small version uses $18$ transformer encoder blocks for the Code-Transformer and another $6$ transformer encoder blocks for the Position-Transformer with a total of $307$ M parameters, and a base version uses $36$ transformer encoder blocks for the Code-Transformer and another $12$ transformer encoder blocks for the Position-Transformer with a total of $607$ M parameters to further demonstrate our scalability. The generation results are reported with a 25% mask ratio at $32\times 32$ resolution feature map using eight RTX-3090 GPUs. Top-k and top-p sampling are used to report the best performance. More details could be found in the supplementary.

2 Comparison with state-of-the-art methods

Unconditional generation. We first compare with million-level state-of-the-art autoregressive models in Table 1. Our model significantly outperforms other autoregressive models with the same parameters (307M). With more parameters, we further increase the FID from 6.84 to 5.67, which demonstrates our scalability. We also compare with other types of unconditional state-of-the-art and large-scale big models in Table 3, where we also achieve top-level performance. Our qualitative unconditional generation results are shown on the left of Figure 3.

Class-conditional generation. We first compare with million-level state-of-the-art in Table 2. We achieve the best FID score compared to all types of models including the recent Mo-VQGAN and RQ-Transformer. We also compare our million-level model with existing billion-level big models in Table 4, where we also achieve top performance with fewer parameters and is only inferior to ViT-VQGAN big model. Our qualitative class-conditional generation results are shown on the right of Figure 3.

Text-conditional generation. We compare with existing text-conditional state-of-the-art without extra web-scale data or pretrained models on MS-COCO for fair comparison in Table.5. We achieve 18.6% FID improvement.

3 Ablations

We conduct ablations on $16\times 16$ resolution feature map using four RTX-3090 GPUs for saving computation resources and the experimental trends are all consistent with $32\times 32$ resolution feature map of the main results.

Ablations of adaptive mask module. As shown in Table 6, our proposed learned adaptive mask mechanism significantly outperforms the random one, which quantitatively validates that the adaptive mask module enables the model to learn perceptually important regions.

Our model with 10% and 25% mask radio has only slightly inferior reconstruction compared with VQGAN, but achieves a significant improvement in generation quality, which indicates that the effectiveness of focusing autoregressive models on modeling important regions. When we further increase the mask radio to 50% and 75%, the final generation quality drops, we believe the reason lies that an improper high mask radio will inevitably mask some important regions that greatly decrease the reconstruction results and hinder autoregressive modeling.

The redundancy of the existing learned codebook can be verified from two aspects: i) the PCA of the learned codebook in Figure 5, where each point is a code and closer codes have more similar semantics. We show many codes in VQGAN’s codebook overlap, which indicates these codes have nearly the same semantics and are thus redundant. The redundancy increase (more overlaps) when VQGAN uses more code to represent images (smaller downsampling factor $f$ ). The redundancy is largely alleviated in our MQ-VAE. ii) in Table 6, a higher codebook usage indicates more “useful” codes in the codebook and thus less redundant. VQGAN has a lower usage compared with our MQ-VAE.

We visualize the training and validation curves of VQGAN and Stackformer in Figure 6(a). Previous autoregressive models always suffer from the overfitting problem while our Stackformer successfully gets rid of it, which indicates the better generalization of our masked discrete representation and our model.

We compare the sampling speed on a single RTX-1080Ti in Figure 6(b). Compared with VQGAN, our 25% mask radio model achieves 32.72% quality improvement and 15.45% speed improvement, while Our 50% mask radio model achieves 26.67% quality improvement and 61.1% speed improvement. Therefore, our design enables a flexible trade-off between speed and quality.

Finally, We visualize the learned mask in Figure 4, with a 75% mask ratio on $32\times 32$ resolution feature map, which validates that our proposed adaptive mask mechanism successfully learns to preserve the perceptual important image regions, i.e., the structural and edge regions of objects.

Ablations of adaptive de-mask module. In Table 7, we show that MQ-VAE outperforms VQGAN and the stronger baseline (“VQGAN*”), which validates our effectiveness. We could conclude that the proposed direction-constrained self-attention and the mask updating mechanism both improve the reconstruction and generation quality.

Conclusion

In this study, we point out that the existing two-stage autoregressive generation paradigm ignores distinguishing the perceptual importance of different image regions, which brings redundancy that not only degrades generation quality but also decreases generation speed. We propose a novel two-stage generation paradigm with MQ-VAE and Stackformer to relieve the model from redundancy. MQ-VAE incorporates the adaptive mask module to mask redundant region features before quantization and the adaptive de-mask module to recover the original feature map after quantization. Stackformer then efficiently predict the combination of both codes and their positions. Comprehensive experiments on various types of image generation tasks validate the effectiveness and efficiency of our method.

Acknowledgments

This work is supported by National Natural Science Foundation of China under Grants 62222212 and U19A2057, Science Fund for Creative Research Groups under Grant 62121002.