Scaling Image Tokenizers with Grouped Spherical Quantization

🔼 This figure shows how the model scales with the latent dimension and spatial compression factor. The latent dimension (d) is held constant at 16, while the number of groups (G), which control the latent space expansion, is varied. The x-axis represents the latent dimension, while the y-axis shows the spatial compression factor. The plot illustrates the relationship between the latent dimension, the number of groups, and the spatial compression achieved by the model.

read the caption

(b) Scaling behaviour of the latent dimension v.s. spatial compression factor in GSQ; d=16𝑑16d=16italic_d = 16 is fixed while groups G𝐺Gitalic_G increase to expand latent space.

🔼 Figure 1 demonstrates GSQ-GAN’s superior image reconstruction capabilities compared to other state-of-the-art methods. The top panel shows that GSQ-GAN achieves a lower reconstruction FID (Frechet Inception Distance) score, indicating better reconstruction quality, even without employing latent decomposition techniques. Increasing the number of groups (G) during training further improves performance by optimizing the use of latent space. The bottom panel illustrates GSQ-GAN’s scalability. It shows that by increasing the latent capacity, GSQ-GAN maintains high-fidelity reconstruction even at high spatial compression ratios (16x downsampling in this case). This is achieved using significantly fewer training epochs (20) compared to existing methods (over 270 epochs).

read the caption

Figure 1: The top figure shows GSQ-GAN’s reconstruction performance compared to state-of-the-art methods, demonstrating superior results even without latent decomposition. Training with larger G𝐺Gitalic_G, which is more composed of groups, can further optimize the use of latent space, enhancing reconstruction quality. The bottom figure illustrates GSQ-GAN’s efficient scaling behaviour, where expanded latent capacity effectively manages increased spatial compression, thus achieving higher fidelity reconstructions on highly spatial compressed latent. Notably, GSQ-GAN achieves these results with only 20 training epochs on ImageNet at 2562superscript2562256^{2}256 start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT resolution, while methods, such as Luo et al. (2024); Yu et al. (2024b), require over 270 epochs.

🔼 This figure compares the performance of different vector quantization (VQ) methods for training a variational autoencoder (VAE) on images with an 8x spatial compression factor. The methods compared are VQ (Vector Quantization), RVQ (Residual Vector Quantization), FSQ (Finite Scalar Quantization), and the proposed GSQ (Grouped Spherical Quantization). All methods use the same VAE architecture (backbone), latent dimension, and codebook size, ensuring a fair comparison. The graph shows the reconstruction error (rFID) over training steps, illustrating the relative efficiency and effectiveness of each quantization technique. The key difference lies in how the latent vectors are quantized to discrete codebook indices, impacting the quality of the reconstructed images.

read the caption

Figure 2: Comparisons of quantizers for VAE-F8 training. VQ is initialized with uniform distribution; all models have the same backbone, latent dimension, and vocabulary size.

🔼 This figure displays the ablation study results of GSQ-GAN model trained on ImageNet with 256x256 resolution images. It shows the impact of increasing network width and depth, with and without attention blocks, on the model’s reconstruction performance. The results demonstrate that wider and deeper networks, especially those with attention blocks, lead to better reconstruction quality. The x-axis likely represents a measure of network complexity (e.g., number of parameters), and the y-axis shows reconstruction error.

read the caption

Figure 3: GSQ-GAN ablations on wider and deeper networks w/ and w/o attention blocks. Models are trained on 2562superscript2562256^{2}256 start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT resolution on ImageNet.

🔼 This figure shows how the reconstruction quality of the GSQ model changes when varying the latent dimension and vocabulary size at 8x spatial compression. The x-axis represents the latent dimension, while the y-axis shows the reconstruction error (rFID). Multiple lines are plotted, each corresponding to different vocabulary sizes. The figure helps in understanding the impact of increasing latent dimension and vocabulary size on the model’s performance during the compression process. It helps to determine the optimal combination of latent dimension and vocabulary size for a given compression ratio.

read the caption

(a) Scaling of latent dimension and vocabulary size for GSQ at 8×\times× spatial compression.

🔼 Figure 4b shows the relationship between reconstruction quality (rFID), latent dimension, and codebook size. The x-axis represents the codebook size (vocabulary) on a logarithmic scale, and the y-axis represents rFID. Different lines represent different latent dimensions. The figure demonstrates that increasing codebook size generally improves reconstruction quality (lower rFID), and that the optimal codebook size may depend on the latent dimension.

read the caption

(b) Same scaling behaviour as the top figure with vocabulary size in logarithmic scale.

🔼 This figure analyzes the impact of latent dimension and codebook size on the reconstruction quality of the GSQ image tokenizer at 8x spatial compression. The top panel shows that smaller latent dimensions lead to better reconstruction, indicating the latent space is not fully utilized at this compression level. Increasing the codebook size further improves performance. The bottom panel presents the same data but with the vocabulary size shown on a logarithmic scale, highlighting the effectiveness of scaling the codebook size. Importantly, all models in this experiment were trained using G=1, meaning there was no latent space decomposition, making them directly comparable to traditional VQ-based methods. The training dataset used was ImageNet with images of 256x256 resolution.

read the caption

Figure 4: The top figure illustrates the scaling of latent dimension and codebook size for GSQ at 8×\times× spatial compression, where a smaller latent dimension improves reconstruction, suggesting the latent space is not saturated for F8 downsampling. Optimising latent space size further enhances performance. The bottom figure shows the same trend with vocabulary size in logarithmic scale, indicating effective scaling as vocabulary size increases. All models are trained with G=1𝐺1G=1italic_G = 1 and no latent decomposition, making this equivalent to VQ-based methods. All models are trained on ImageNet at 2562superscript2562256^{2}256 start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT resolution.

🔼 Figure 5 demonstrates the effect of increasing latent dimensionality on the reconstruction performance of GSQ-GAN when the spatial compression ratio is 16 (F16). The figure shows that increasing the latent dimension initially improves reconstruction quality, but beyond a certain point (D=16 in this case), adding more dimensions does not provide further gains. This saturation effect highlights the limitations of standard VQ techniques in high-dimensional latent spaces. However, GSQ with latent decomposition (using multiple groups, G >1) successfully scales to much higher latent dimensions (32 and 64) and continues to show significant improvements in reconstruction performance. This highlights the effectiveness of GSQ in managing the complexity of high-dimensional spaces, resulting in improved image generation quality.

read the caption

Figure 5: Latent dimension scaling for GSQ-GAN-F16 training, the latent space is saturated for F16 spatial compression; we expect to enhance reconstruction performance by increasing the latent dimension to increase the latent capacity. Only GSQ with latent decomposition can scale to a higher latent dimension.

🔼 This table presents an ablation study on the impact of codebook auxiliary loss functions (entropy and TCR loss) on the performance of the GSQ-VAE-F8 model. The results show that adding these auxiliary loss functions does not improve performance and even negatively impacts codebook utilization. The GSQ-VAE-F8 model, with the proposed spherical quantization method, consistently achieves near-full codebook usage (close to 100%) without the need for any auxiliary loss. This demonstrates the effectiveness of the proposed GSQ method.

read the caption

Table 2: Ablation of codebook auxiliary loss for GSQ-VAE-F8. Our methods enable the codebook usage to always be full; there is no need to use this auxiliary loss for training.

🔼 This table presents the ablation study results on the impact of using Adaptive Group Normalization (AGN) and Depth2Scale modules within the GSQ-VAE-F8 model. It shows how the reconstruction performance (rFID), Inception Score (IS), Learned Perceptual Image Patch Similarity (LPIPS), Peak Signal-to-Noise Ratio (PSNR), Structural Similarity Index (SSIM), codebook usage, and perplexity (PPL) change when AGN and/or Depth2Scale is included or excluded. This helps to understand the contribution of each module to the overall performance of the image tokenizer model.

read the caption

Table 3: Ablation of using Adaptive Group Norm (AGN) and Depth2Scale for GSQ-VAE-F8.

🔼 This table presents an ablation study on the impact of different perceptual loss functions and their associated weights on the performance of a VAE model (specifically, the VAE-F8 model, as indicated by the table’s title) during training. It examines different configurations by varying weights assigned to perceptual loss (λp) and reconstruction loss (λrec). The results likely show how the balance between these two loss types affects the model’s ability to reconstruct images accurately and maintain perceptual fidelity (i.e., how well the reconstructed images resemble the original images visually). The table likely shows various metrics evaluating reconstruction quality.

read the caption

Table 4: Ablation of perceptual loss and weights for VAE-F8 training. λpsubscript𝜆𝑝\lambda_{p}italic_λ start_POSTSUBSCRIPT italic_p end_POSTSUBSCRIPT and λr⁢e⁢csubscript𝜆𝑟𝑒𝑐\lambda_{rec}italic_λ start_POSTSUBSCRIPT italic_r italic_e italic_c end_POSTSUBSCRIPT are weights of perceptual and reconstruction loss.

🔼 This table presents the results of ablation studies on the hyperparameters of the Adam optimizer used for training the GSQ-VAE-F8 model. Specifically, it investigates the impact of different beta values (β1 and β2) and weight decay on the model’s performance. The table shows the reconstruction performance metrics (rFID, IS, LPIPS, PSNR, SSIM, Usage, and PPL) achieved using various combinations of beta values and weight decay. A key observation is that the codebook usage remains at 100% across all configurations tested, indicating that the chosen hyperparameters effectively utilize the entire codebook.

read the caption

Table 5: Optimizer’s β𝛽\betaitalic_β and weight decay ablations for GSQ-VAE-F8 training. The codebook usage is 100% for all models.

🔼 This table presents an ablation study on different learning rate schedulers used during the training of a Variational Autoencoder (VAE) for image tokenization. The goal is to determine the optimal learning rate schedule for achieving high-quality image reconstruction. Several schedules are compared, including those with and without a warm-up period and with different decay rates. The table shows the impact of each schedule on the reconstruction fidelity (rFID), Inception Score (IS), LPIPS, PSNR, SSIM, codebook usage, and perplexity (PPL). The maximal learning rate used across all experiments was 1e-4. The experiment shows that a constant learning rate provides the best performance, highlighting the tradeoffs between fast convergence and model quality.

read the caption

Table 6: Learning rate scheduler ablations for GSQ-VAE-F8 training, the maximal learning rate is 1e−4superscript𝑒4e^{-4}italic_e start_POSTSUPERSCRIPT - 4 end_POSTSUPERSCRIPT. The codebook usage is 100% for all models.

🔼 This table presents the results of training a GSQ-GAN (Grouped Spherical Quantization Generative Adversarial Network) model on the ImageNet dataset at a resolution of 128x128 pixels. The training involved 80,000 steps, except for the StyleGAN model which was stopped after 1000 steps due to NaN losses during training. The table compares the reconstruction performance (rFID), Inception Score (IS), Peak Signal-to-Noise Ratio (PSNR), Structural Similarity Index (SSIM), Codebook Usage, and Perplexity (PPL) across different discriminator architectures (NLD, SGD, DD) and adversarial loss functions (Vanilla, Hinge, Non-Saturating). This allows for a comprehensive analysis of various components and configurations within the GSQ-GAN model and their effect on performance.

read the caption

Table 7: GSQ-GAN-F8 model trained on 1282 ImageNet, 80k training step. The SGD-GAN model is evaluated at the 1k training step due to the failure of 𝑁𝑎𝑁𝑁𝑎𝑁\mathit{NaN}italic_NaN loss in training.

🔼 This table presents the ablation study on different data augmentation techniques used in the Dino-Discriminator within the GSQ-GAN model. It shows the impact of color augmentation, translation, and cutout on the reconstruction performance (measured by rFID) at two different spatial resolutions: 128x128 and 256x256.

read the caption

Table 8: Ablation on data augmentation in Dino-Discriminator.

🔼 This table presents the results of ablation studies on the hyperparameters of the Adam optimizer and the weight of the adversarial loss in the GSQ-GAN-F8 model training. It shows how different values of Adam’s beta parameters (β1 and β2) and the adversarial loss weight (λadv) affect the model’s performance, as measured by the FID (Fréchet Inception Distance) score. The table helps to determine the optimal hyperparameter settings for training the GSQ-GAN-F8 model.

read the caption

Table 9: Ablation of Adam’s β𝛽\betaitalic_β and adversarial loss weights for GSQ-GAN-F8 training. λa⁢d⁢vsubscript𝜆𝑎𝑑𝑣\lambda_{adv}italic_λ start_POSTSUBSCRIPT italic_a italic_d italic_v end_POSTSUBSCRIPT is the weight of adversarial loss.

🔼 This table presents the results of ablation studies on the batch size and learning rate used during the training of the GSQ-GAN-F8 model. The experiment varied the batch size (256, 512, 768) and learning rate (1e-4, 2e-4, 3e-4) to determine their impact on the model’s performance. The performance is measured by rFID (Reconstruction Fidelity), Inception Score (IS), LPIPS (Learned Perceptual Image Patch Similarity), PSNR (Peak Signal-to-Noise Ratio), SSIM (Structural Similarity Index), Codebook Usage, and Perplexity (PPL). The Dino Discriminator (DD) and the Hinge-NonSaturating (NH) loss combination were used consistently in this experiment.

read the caption

Table 10: Batch size and learning rate ablations of GSQ-GAN-F8 training, with DD-NH discriminator and loss combination.

🔼 This table presents ablation studies on different GAN regularization techniques used in training the GSQ-GAN-F8 model. It shows the impact of weight decay (WD), adaptive adversarial loss weights (AW), gradient clipping (GC), gradient penalty (GP), and LeCAM regularization on the model’s reconstruction performance, measured by rFID. The study also examines the effect of delaying the discriminator updates (warmup) until 20,000 iterations. The results help determine the best regularization strategy for optimal training.

read the caption

Table 11: Ablation studies of GAN’s regularization technologies for GSQ-GAN-F8 training, WD is weight decay, AW is adversarial loss adaptive weight Esser et al. (2021), GC is gradient clip. All modes are trained with gradient clip 2.0 by default, GP is gradient penalty, and LeCAM’s weight is 0.001 if enabled; when warmup is used, the discriminator starts to be updated after 20k iterations.

🔼 This table presents an ablation study on the GSQ-GAN-F8 model, analyzing the impact of data augmentation techniques (color, translation, cutout, resize), attention module integration, and Depth2Scale layers on the model’s performance. The results are measured using FID (Fréchet Inception Distance) and reported for two different spatial resolutions (128x128 and 256x256). The study aims to determine the optimal combination of these components for achieving the best reconstruction quality.

read the caption

Table 12: Ablation of discriminator data augmentation, integration of attention and Depth2Scale for GSQ-GAN-F8 training. D2S is the short for Depth2scale.

🔼 Table 13 presents a comparison of reconstruction performance using different configurations of the Grouped Spherical Quantization (GSQ) method. The table shows how changing the spatial downsampling factor (8x, 16x, 32x), the codebook vocabulary size (8k, 256k), and the group decomposition (Gx d) affects the reconstruction quality, measured by the rFID (Fréchet Inception Distance). GSQ consistently outperforms LFQ (Lookup-Free Quantizer) achieving a 3x lower rFID. The table also highlights the hyperparameters used in GSQ such as the number of groups (G) and the latent dimension of each group (d). One notable experiment (16 x 1*) used a clip normalization instead of the typical l2 normalization.

read the caption

Table 13: Ablation studies of group decomposition with 8, 16 and 32 spatial downsample, vocabulary size is 8k, 256k and 256k respectively. GSQ outperforms LFQ with 3×3\times3 × lower rFID. G𝐺Gitalic_G is the number of groups, and d𝑑ditalic_d is a latent dimension in each group. 16×1∗16superscript116\times 1^{*}16 × 1 start_POSTSUPERSCRIPT ∗ end_POSTSUPERSCRIPT is trained with clip instead of ℓ2subscriptℓ2\ell_{2}roman_ℓ start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT normalization.

🔼 This table presents a comparison of the reconstruction performance (measured by FID score) of various state-of-the-art image tokenizers on the ImageNet dataset. The models are evaluated using images at a resolution of 256x256 pixels. The table allows for a quantitative comparison of the proposed GSQ model against existing methods, considering factors such as the latent dimension, codebook size and other hyperparameters, providing insights into the model’s relative performance and efficiency.

read the caption

Table 14: Reconstruction performance comparison of the proposed model against other state-of-the-art methods on ImageNet (256×256256256256\times 256256 × 256 resolution).

🔼 This table compares various vector quantization (VQ) methods for image tokenization, showing how they relate to the proposed Grouped Spherical Quantization (GSQ) method. It details the key parameters (latent dimension, codebook size, and the use of techniques like fixed codebooks and ℓ2 normalization) of each method, highlighting how they differ from or are special cases of GSQ. This provides a clear overview of existing methods and clarifies the relationship between GSQ and its predecessors.

read the caption

Table 15: The effective configurations of other tokenizers in GSQ’s view.

🔼 Table 16 presents the hyperparameters used during the training phase of the VAE-F8 model. It details the settings for various aspects of the training process, including training parameters (image resolution, number of training steps, gradient clipping, mixed precision, batch size, and exponential moving average beta), model configuration (downsampling factor, hidden channels, channel multipliers, encoder and decoder layer configurations), quantizer settings (embedding dimension, codebook vocabulary size, and group size), codebook initialization and normalization methods, loss function weights (for reconstruction, perceptual loss using LPIPS, and commitment loss), optimizer settings (type of optimizer, base learning rate, learning rate scheduler, weight decay, betas, and epsilon), providing a comprehensive overview of the training setup.

read the caption

Table 16: VAE-F8 Training Hyperparameters

🔼 This table presents a qualitative comparison of image reconstruction results obtained using the VAE-F8 model with different configurations: a baseline model, models incorporating Depth2Scale, Adaptive Normalization, and a combination of both Depth2Scale and Adaptive Normalization. The goal is to show the impact of these techniques on image reconstruction quality. Each image shows the original, and the reconstructions.

read the caption

Table 17: Reconstruction results of the VAE-F8 model (in Section 4.1.2) with ablation of Depth2Scale and Adaptive Normalization.

🔼 Table 18 provides a comprehensive list of hyperparameters used for training the GSQ-GAN model with an 8x spatial downsampling factor. It details both training parameters (e.g., image resolution, number of training steps, batch size) and model configuration settings (e.g., network architecture, latent dimension, codebook size). Furthermore, it specifies the quantization settings, loss weights, optimizer settings, and discriminator configurations used during the GAN training process.

read the caption

Table 18: GAN-F8 Training Hyperparameters

🔼 This table details the architectures of three different discriminators used in the GAN training experiments within the paper: N-Layer, StyleGAN, and DINO. For each discriminator, the table lists the input channels, number of channels, number of layers, and (where applicable) the channel multiplier used in the model architecture. The specific parameters for each discriminator type provide insight into the structural differences used in the training process.

read the caption

Table 19: Discriminator configurations

🔼 This table presents a qualitative comparison of image reconstruction results obtained using different discriminators within the GSQ-GAN-F8 model. The reconstructions of three sample images (a wooden bowl, a rainbow lorikeet, and a library) are displayed for each discriminator configuration. This allows for a visual assessment of the impact of discriminator choice on image generation quality.

read the caption

Table 20: Reconstruction results of the GAN-F8 models (see Section 4.2.1) , trained with different discriminators.

🔼 Table 21 presents the hyperparameters used during the training phase of the GSQ-GAN model with a downsampling factor of 8. It details settings for various aspects of the training process including image resolution, training steps, gradient clipping, batch size, optimizer parameters, model architecture details (e.g., hidden channels, number of layers in encoder and decoder), quantizer settings (latent dimension, codebook size, group number), loss function weights (reconstruction, perceptual, commitment, adversarial), and discriminator-specific settings. The table offers a comprehensive view of the configuration choices employed to optimize GSQ-GAN’s performance.

read the caption

Table 21: GAN-F8 Training Hyperparameters

🔼 This table presents a visual comparison of image reconstruction quality across different latent dimensions used in the GSQ-GAN-F8 model. It shows original images alongside reconstructions generated using latent dimensions of 16 (with 1 group), 32 (with 2 groups), and 64 (with 4 groups). The purpose is to demonstrate the impact of latent dimension scaling on the model’s ability to accurately reconstruct images. Figure 1b in the paper provides details about the specific configurations of each model.

read the caption

Table 22: Scaling latent dimension for GSQ-GAN-F8 model. The models are detailed in fig. 1.

|—|—|—|—| |

|

|

|

| |(a) Orignal images (128 × 128 resolution) |(a) Orignal images (128 × 128 resolution)|(a) Orignal images (128 × 128 resolution)|(a) Orignal images (128 × 128 resolution)| |

|

|

|

| |(b) Reconstruction results by with NLD-NV discriminators|(b) Reconstruction results by with NLD-NV discriminators|(b) Reconstruction results by with NLD-NV discriminators|(b) Reconstruction results by with NLD-NV discriminators| |

|

|

|

| |(c) Reconstruction results by with DD-NH discriminators|(c) Reconstruction results by with DD-NH discriminators|(c) Reconstruction results by with DD-NH discriminators|(c) Reconstruction results by with DD-NH discriminators| |

|

|

|

| |(d) Reconstruction results by with NLD-NV discriminators and β=[0.9,0.99]|(d) Reconstruction results by with NLD-NV discriminators and β=[0.9,0.99]|(d) Reconstruction results by with NLD-NV discriminators and β=[0.9,0.99]|(d) Reconstruction results by with NLD-NV discriminators and β=[0.9,0.99]| |

|

|

|

| |(e) Reconstruction results by with DD-NH discriminators and β=[0.9,0.99]|(e) Reconstruction results by with DD-NH discriminators and β=[0.9,0.99]|(e) Reconstruction results by with DD-NH discriminators and β=[0.9,0.99]|(e) Reconstruction results by with DD-NH discriminators and β=[0.9,0.99]|

🔼 This table presents the results of experiments conducted to analyze how scaling the latent dimension affects the performance of the GSQ-GAN-F16 model. It showcases reconstruction results for different configurations, each varying in latent dimension (D) and number of groups (G). The impact of these changes on the model’s ability to accurately reconstruct images is visually demonstrated through sample images, offering a comparative analysis of the trade-off between latent space complexity and reconstruction quality.

read the caption

Table 23: Scaling latent dimension for GSQ-GAN-F16 model. The models are detailed in fig. 1.

🔼 This table presents results from experiments on scaling the latent dimension of the GSQ-GAN-F32 model. It shows how varying the latent dimension (D) and the number of groups (G), while maintaining a constant overall latent dimensionality (D = d x G), affects the model’s performance. The impact on image reconstruction quality is assessed, and the specific model configurations used in the experiment are referenced in Figure 1 of the paper. The table likely includes metrics such as reconstruction fidelity (rFID), Inception Score (IS), Learned Perceptual Image Patch Similarity (LPIPS), Peak Signal-to-Noise Ratio (PSNR), Structural Similarity Index (SSIM), codebook usage, and perplexity (PPL).

read the caption

Table 24: Scaling latent dimension for GSQ-GAN-F32 model. The models are detailed in fig. 1.

|—|—|—|—| |

|

|

|

| |(a) Orignal images (256 \times 256 resolution) | | | | |

|

|

|

| |(b) GSQ-GAN-F8, D=16, G=1 | | | | |

|

|

|

| |(b) GSQ-GAN-F8, D=32, G=2 | | | | |

|

|

|

| |(b) GSQ-GAN-F8, D=64, G=4 | | | |

🔼 This table presents an ablation study on scaling the latent dimension in the GSQ-GAN-F8 model, specifically focusing on the impact of varying the number of groups (G) while maintaining a fixed latent dimensionality (D=64) and down-sampling factor (F=8). The study examines how changing the group size affects the model’s reconstruction quality, offering insights into the optimal balance between computational efficiency and representational capacity. Details of the models used in this study are provided in Section 4.3.3 of the paper.

read the caption

Table 25: Scaling latent dimension for GSQ-GAN-F8-D64 model. The models are detailed in Section 4.3.3.