MoVQGAN Architecture: Enhanced VQGAN with MaskGIT

• MoVQGAN-based Architecture is a two-stage generative framework that uses multichannel quantization and spatially conditional normalization to enhance reconstruction fidelity.
• It integrates a VQ-VAE encoder-decoder with a MaskGIT-based Transformer for parallel discrete token prediction, reducing sampling steps dramatically.
• The design mitigates repeated artifacts seen in standard VQGANs and achieves high photorealistic synthesis quality through innovative training protocols.

MoVQGAN (Modulated Vector-Quantized Generative Adversarial Network) defines a two-stage image generation framework that introduces spatially conditional normalization and multichannel quantization for improved reconstruction fidelity and sample diversity. Addressing limitations in standard VQ-based pipelines—namely, repeated-pattern artifacts due to codebook collapse—MoVQGAN enables effective photorealistic image synthesis via a combination of architectural and training protocol innovations. At the prior modeling stage, it leverages Masked Generative Image Transformer (MaskGIT) for highly parallelizable discrete token prediction, accelerating image sampling while maintaining high output quality (Zheng et al., 2022).

1. Pipeline Overview and Core Motivations

MoVQGAN operates in two stages: (i) a quantizer stage using a VQ-VAE-style encoder-decoder for vector-quantized latent representation of images, and (ii) a prior modeling stage using a bidirectional Transformer to generate latent index maps for image synthesis.

Stage 1: The encoder $\mathcal{E}_\psi$ maps $x\in\mathbb{R}^{256\times256\times3}$ to a continuous latent $\hat z\in\mathbb{R}^{16\times16\times256}$. Multichannel quantization is applied, generating a discrete latent index tensor $s\in\{0,\dots,1023\}^{16\times16\times4}$. The decoder $\mathcal{G}_\theta$ reconstructs the image.

Stage 2: A MaskGIT-based Transformer is trained to model the prior distribution over discrete index tensors. By conditioning on masked versions of $s$, the Transformer predicts latent tokens in parallel, enabling dramatically reduced sampling steps.

Key motivations relative to standard VQGAN include:

• Mitigating repeated “tile” artifacts from codebook index reuse.
• Injection of spatially variant information through Spatially Conditional Normalization (SCN).
• Expanding representational capacity via multichannel quantization (from $K$ to $K^c$ possible codes).
• Faster and more flexible generative prior sampling using MaskGIT over standard autoregressive models.

2. Encoder and Multichannel Quantization Methodology

The encoder follows the VQGAN/Taming Transformers structure, employing Gated ReLU activations and GroupNorm (encoder only). The input sequence of convolutions and residual blocks reduces spatial resolution, outputting $\hat z\in\mathbb{R}^{16\times16\times256}$.

Multichannel Quantization proceeds by splitting the 256-dimensional feature at each spatial location into $c=4$ chunks of 64 dimensions, quantizing each independently against a shared codebook $x\in\mathbb{R}^{256\times256\times3}$0. For each chunk $x\in\mathbb{R}^{256\times256\times3}$1 at location $x\in\mathbb{R}^{256\times256\times3}$2,

$x\in\mathbb{R}^{256\times256\times3}$3

and $x\in\mathbb{R}^{256\times256\times3}$4.

The discrete index tensor $x\in\mathbb{R}^{256\times256\times3}$5 encodes the quantized state per spatial location. Multichannel quantization increases the expressive capacity without scaling the codebook size, making $x\in\mathbb{R}^{256\times256\times3}$6 compositional codes possible at constant memory and compute.

Codebook updates follow the standard VQ-VAE commitment and codebook loss formulation: $x\in\mathbb{R}^{256\times256\times3}$7 with $x\in\mathbb{R}^{256\times256\times3}$8.

3. Spatially Conditional Normalization (SCN) and Decoder Advances

A noted problem in standard VQ-based generators is that decoder activations for locations sharing a codebook index become visually repetitive, leading to tiling artifacts. SCN addresses this by injecting per-location modulation signals derived from the quantized latent.

For each decoder block $x\in\mathbb{R}^{256\times256\times3}$9 (in the first three blocks), the pre-activation $\hat z\in\mathbb{R}^{16\times16\times256}$0 is normalized via GroupNorm: $\hat z\in\mathbb{R}^{16\times16\times256}$1 SCN then applies learned per-position affine transforms from up-/downsampled $\hat z\in\mathbb{R}^{16\times16\times256}$2: $\hat z\in\mathbb{R}^{16\times16\times256}$3 where $\hat z\in\mathbb{R}^{16\times16\times256}$4 and $\hat z\in\mathbb{R}^{16\times16\times256}$5 are $\hat z\in\mathbb{R}^{16\times16\times256}$6 convs mapping $\hat z\in\mathbb{R}^{16\times16\times256}$7 to $\hat z\in\mathbb{R}^{16\times16\times256}$8 scale and shift parameters.

The decoder otherwise follows the VQGAN upsampling stack, but only the first three blocks substitute GroupNorm with SCN.

4. Prior Modeling via MaskGIT

Conventional transformer priors factorize $\hat z\in\mathbb{R}^{16\times16\times256}$9 and require sequential token prediction, with $s\in\{0,\dots,1023\}^{16\times16\times4}$0 for $s\in\{0,\dots,1023\}^{16\times16\times4}$1 maps. MaskGIT instead models $s\in\{0,\dots,1023\}^{16\times16\times4}$2, with $s\in\{0,\dots,1023\}^{16\times16\times4}$3 masking a random subset $s\in\{0,\dots,1023\}^{16\times16\times4}$4 of tokens.

Training: At each step, $s\in\{0,\dots,1023\}^{16\times16\times4}$5 of tokens are randomly masked (with $s\in\{0,\dots,1023\}^{16\times16\times4}$6), and the transformer minimizes the cross-entropy over masked positions.

Inference: Starting with all locations masked, at each of $s\in\{0,\dots,1023\}^{16\times16\times4}$7 iterations (e.g., $s\in\{0,\dots,1023\}^{16\times16\times4}$8), the transformer predicts logits for all masked tokens in parallel, the most confident $s\in\{0,\dots,1023\}^{16\times16\times4}$9 are selected to fill, and the process is repeated until all positions are filled. This reduces parallel sampling from 1024 to fewer than 12 steps.

5. Losses and Training Protocols

Stage 1 (Quantizer+Decoder):

• Pixelwise reconstruction: $\mathcal{G}_\theta$0
• Perceptual (VGG feature) loss: $\mathcal{G}_\theta$1
• Adversarial loss: $\mathcal{G}_\theta$2
• VQ codebook and commitment losses per 2.3 above

The total loss is: $\mathcal{G}_\theta$3 with weights $\mathcal{G}_\theta$4, $\mathcal{G}_\theta$5, $\mathcal{G}_\theta$6.

Stage 2 (MaskGIT Prior):

• Negative log-likelihood on masked tokens: $\mathcal{G}_\theta$7

Training utilizes batch sizes of 48 and 64 for stages 1 and 2, respectively; Adam optimizer with $\mathcal{G}_\theta$8, and a base learning rate of $\mathcal{G}_\theta$9.

6. Experimental Evaluation and Ablation Analysis

Quantitative and sample quality results on FFHQ and ImageNet highlight substantial gains relative to vanilla VQGAN.

Model/Dataset	PSNR	SSIM	LPIPS	rFID	FID (Stage 2)	IS (Stage 2)
VQGAN/FFHQ	22.24	0.664	0.1175	4.42	11.4	—
MoVQGAN/FFHQ	26.72	0.821	0.0585	2.26	8.78	—
VQGAN/INet	19.47	0.521	0.1950	6.25	—	78.3
MoVQGAN/INet	22.42	0.673	0.113	1.12	7.22	130.1
MaskGIT/INet	—	—	—	—	6.18	182.1

MoVQGAN eliminates repeating artifacts observed in past models such as standard VQGAN (“grass-tile” effects, as shown in Fig. 3 of (Zheng et al., 2022)). It also exhibits efficient sampling—8 MaskGIT inference steps versus up to 1024 steps for AR baselines.

Ablations:

• Fourier features combined with SCN and multichannel quantization yield the best rFID (2.26) and FID (8.78).
• Increasing codebook size $s$0 from 1024 to 16,384 produces negligible rFID improvements in MoVQGAN, but has much more impact in vanilla VQGAN; this suggests code capacity in MoVQGAN is not a limiting factor.
• Increasing number of channels $s$1 improves PSNR but can degrade Stage 2 FID as sampling complexity increases; $s$2 yields the best trade-off.

7. Architectural and Implementation Details

• Hardware: 4×Tesla V100 GPUs
• Encoder: Input $s$3, down to $s$4
• Quantization: Codebook size $s$5, code dimension 64, $s$6 channels per spatial location
• Decoder: Up-convolutional stack, three initial blocks with SCN, all blocks otherwise match VQGAN
• Stage 2 Transformer: 24 layers, 16 heads, $s$7, $s$8
• MaskGIT: 8 sampling steps (FFHQ), 12 steps (ImageNet), with progressive random masking schedule

The architecture enables high compression rates, improved quality metrics, and practical inference time for high-resolution image synthesis (Zheng et al., 2022).