PaperlayerView exact revision on arXiv

Source-linked AI summary

Finite Scalar Quantization: VQ-VAE Made Simple

Fabian Mentzer, David Minnen, Eirikur Agustsson, Michael Tschannen

arXiv:2309.15505v2cs.CVcs.LG

TL;DR

VQ-based representations are difficult to optimize because larger codebooks become underutilized. The paper introduces FSQ, which bounds and rounds a few dimensions to form an implicit product codebook, and reports competitive performance across image generation and dense vision tasks.

  • Problem

    VQ codebooks become underutilized as their size increases, motivating a simpler formulation that removes auxiliary losses while preserving high utilization and drop-in compatibility.

  • Method

    FSQ projects representations to a few dimensions, bounds each channel to fixed values, and rounds them to create an implicit product codebook.

  • Results

    FSQ serves as a drop-in replacement across MaskGIT and UViM tasks, with metric reductions of only 0.5–3% and nearly complete codebook usage in most models.

  • Takeaways & Limitations

    FSQ provides expressive discrete representations without codebook collapse or the auxiliary machinery used by VQ.

Abstract

from arXiv · show

We propose to replace vector quantization (VQ) in the latent representation of VQ-VAEs with a simple scheme termed finite scalar quantization (FSQ), where we project the VAE representation down to a few dimensions (typically less than 10). Each dimension is quantized to a small set of fixed values, leading to an (implicit) codebook given by the product of these sets. By appropriately choosing the number of dimensions and values each dimension can take, we obtain the same codebook size as in VQ. On top of such discrete representations, we can train the same models that have been trained on VQ-VAE representations. For example, autoregressive and masked transformer models for image generation, multimodal generation, and dense prediction computer vision tasks. Concretely, we employ FSQ with MaskGIT for image generation, and with UViM for depth estimation, colorization, and panoptic segmentation. Despite the much simpler design of FSQ, we obtain competitive performance in all these tasks. We emphasize that FSQ does not suffer from codebook collapse and does not need the complex machinery employed in VQ (commitment losses, codebook reseeding, code splitting, entropy penalties, etc.) to learn expressive discrete representations.

1 INTRODUCTION

The paper proposes FSQ as a simpler drop-in replacement for VQ, addressing underutilized codebooks while preserving the functional setup and supporting multiple architectures and tasks.

  • VQ commonly suffers from underutilized codebooks as codebook size increases, motivating techniques such as reinitialization and stochastic formulations.
  • FSQ bounds and rounds a few scalar representation dimensions, creating an implicit codebook whose size is the product of per-dimension value counts.For d=3 and L=3, the implied codebook contains 27 vectors.
  • FSQ uses the straight-through estimator and reconstruction loss to spread information across quantization bins, achieving codeword usage without auxiliary losses.
  • FSQ serves as a drop-in replacement for VQ in MaskGIT image generation and UViM depth estimation, colorization, and panoptic segmentation, with metric reductions of only 0.5–3%.The comparison spans convolutional and transformer-based autoencoders and several transformer configurations.
  • The paper analyzes VQ–FSQ trade-offs, finding that FSQ supports large codebooks with nearly complete usage and that VQ provides little benefit over FSQ in the reported settings.

2 RELATED WORK

Prior work developed VQ-VAE improvements and alternative quantization schemes, while neural compression established bounded and unbounded scalar quantization for image representations.

  • VQ-VAE introduced commitment loss and EMA, and later work used soft EM, random restarts, stochastic formulations, and multiscale variants to improve codebook learning.
  • Residual quantization refines codes with quantized residuals, while product quantization factors a codebook into products of smaller codebooks.
  • Neural compression commonly uses unbounded scalar quantization with entropy constraints, while bounded scalar quantization has represented images with high fidelity and under extreme compression.Reported configurations include d=16, L=5 and d=5, L=5.

3 METHOD

FSQ replaces learned nearest-neighbor codebooks with a low-dimensional fixed grid formed by bounding and rounding channels, while retaining compatibility with VQ-based networks.

  • FSQ uses a fixed grid in a much lower-dimensional latent space, with encoder and decoder capacity absorbing the nonlinear partitioning handled by VQ.
  • FSQ computes ˆz = round(f(z)), where each channel takes L values and the implicit product codebook has size |C| = L^d.The codewords can be enumerated as integers, allowing FSQ to replace VQ after adapting adjacent layer dimensions.
  • The straight-through estimator propagates gradients through rounding, implemented with a stop-gradient operation.
  • FSQ is controlled by the channel count d and per-channel levels L = [L_1, ..., L_d], with multiple configurations able to approximate a target codebook size.
  • Using L_i ≥ 5 for every channel is a heuristic that performs well across the considered tasks.
  • FSQ with the same codebook size has fewer parameters than VQ because it does not learn a |C| · d codebook and typically uses fewer latent dimensions.For |C|=2^12=4096 and VQ dimension d=512, the learned codebook contains 2M parameters; the cited FSQ example uses d=5.

4 EXPERIMENTS

The experiments evaluate FSQ and VQ within MaskGIT image generation and UViM dense prediction, using matched codebook configurations and task-specific metrics. They examine representation quality, modeling complexity, classifier-free guidance, and qualitative sample quality.

  • MaskGIT: MaskGIT trains a convolutional VQ-GAN autoencoder for reconstruction before a masked transformer predicts quantized representations.Inference progressively replaces MASK tokens using prediction confidence until all tokens are uncovered.
  • UViM: UViM applies a shared architecture to panoptic segmentation, depth estimation, and colorization, with task-specific weights and optional input context.Its second-stage transformer predicts quantized dense labels conditioned on the task input.
  • Evaluation: 128 × 128 ImageNet experiments track Reconstruction FID, Sampling FID, Codebook Usage, and Compression Cost while sweeping codebook size and FSQ configurations.Reconstruction FID evaluates quantized autoencoder reconstructions, whereas Sampling FID evaluates representations sampled by the Stage II transformer.
  • MaskGIT evaluation: For 256 × 256 ImageNet MaskGIT, the study sweeps classifier-free guidance weights and compares FSQ and VQ using Sampling FID, precision, and recall, with ADM as a reference.The ADM reference is shown horizontally because its CFG weight is not directly comparable with the FSQ and VQ guidance parameter.
  • Results: FSQ and VQ produce very comparable MaskGIT metrics and visual samples, while UViM experiments report PQ, RMSE, and FID-5k across the three dense prediction tasks.The UViM setup uses 4096 VQ codewords and FSQ levels [7, 5, 5, 5, 5].

5 RESULTS

FSQ matches VQ closely across image generation and UViM dense-prediction tasks while using codebooks more effectively and avoiding auxiliary codebook machinery. Scaling FSQ improves reconstruction and sampling behavior until modeling complexity limits further gains.

  • Trade-off study: FSQ reconstruction FID keeps improving with codebook size, while VQ reaches its minimum at 2^11 codes as codebook usage begins decreasing.At low codebook sizes, VQ marginally outperforms FSQ.
  • UViM: FSQ obtains competitive but marginally worse results than VQ across UViM’s depth estimation, colorization, and panoptic segmentation tasks.Table 2 reports metrics averaged over three runs with standard deviations.
  • Codebook utilization: More than 2^15 FSQ codewords are utilized for a codebook of size 2^16, whereas VQ falls below 50% usage above 2^11 and cannot utilize more than 2^10 codewords for larger codebooks.FSQ achieves this usage without auxiliary algorithms.
  • Trade-off study: FSQ sampling FID saturates at about 2^12 codewords as representation compression cost continues increasing and makes transformer modeling harder.Within each model family, higher compression cost is associated with worse sampling FID.
  • MaskGIT: FSQ and VQ achieve very comparable FID, precision, recall, and visual samples on 256×256 ImageNet with MaskGIT.Increasing the codebook size further did not produce additional gains for these models.
  • UViM: Removing UViM context lowers PQ for both quantizers, but the FSQ-based model degrades less.Disabling VQ codebook splitting instead causes significantly worse RMSE and 0.78% codebook usage on NYU Depth.

A.1 CODE

The implementation converts bounded, quantized channel values into compact integer indices and reconstructs codes from those indices. It uses straight-through rounding so quantization remains trainable.

  • Initialization: The constructor stores per-channel levels and builds cumulative-product bases for mixed-radix code indexing.The codebook size is the product of the levels, and indices span the resulting codebook.
  • Bounding and quantization: The bounding function uses tanh, level-dependent scaling, and offsets to constrain each channel before quantization.The implementation accounts for both odd and even numbers of levels.
  • Bounding and quantization: Quantization rounds the bounded representation and renormalizes it to [-1, 1] while applying straight-through gradients.The inverse scaling restores the code values from the normalized representation.
  • Index conversion: Codes are mapped to integer indices by scaling values and summing them with the mixed-radix basis, with an inverse routine recovering codes from indices.This provides the explicit enumeration of FSQ’s implicit product codebook.

A.2 ADDITIONAL UVIM VISUALS

Additional UViM visualizations cover segmentation, colorization, and representation stitching. The figures also examine randomly sampled and spatially repeated FSQ and VQ codes.

  • Segmentation and colorization: Figure 7 visualizes panoptic segmentation in its first two rows and colorization in its last two rows.
  • Representation stitching: Figure 8 compares pixel-space stitching with latent-space stitching from FSQ-GAN and VQ-GAN representations.The GAN decoder maps sharp representation-space transitions to smooth pixel-space transitions.
  • Code analysis: Figure 9 samples FSQ and VQ codes from marginal histograms and creates representations using each of the three most common codes across all spatial locations.

A.3 VISUALIZING VQ AND FSQ REPRESENTATIONS

Visualizations suggest that FSQ- and VQ-based autoencoders produce similar decoded image structures and blend encoded representations smoothly. The analysis further implies that final RGB content depends more on combinations of codes and decoder weights than on individual codes alone.

  • Both FSQ-GAN and VQ-GAN produce similar “soup of patches” when decoding sampled marginal representations.The representations are sampled from marginal histograms over the ImageNet validation set.
  • Both decoders smoothly blend representations stitched from real images when decoding into RGB space.
  • Individual codes do not appear to learn highly abstract concepts; combinations of codes and decoder weights determine the final RGB image.

A.4.1 TRADEOFF STUDY

The study explores FSQ quantization-level configurations for MaskGIT and finds that using at least five levels per dimension is a robust heuristic. Masking-ratio lower-bounding is also used to stabilize Stage II training.

  • Quantization-level configurations: Li ≥5 for every dimension is a simple heuristic that performs well across the considered tasks.The study explores multiple quantization-level configurations and tabulates configurations for common target codebook sizes.
  • Training setup: The experiments train Stage I for 100 epochs and Stage II for 200 epochs, using batch size 256 in both stages.
  • Masking-ratio stabilization: Masking more than 0.45S tokens at every training step stabilizes MaskGIT Stage II training.The default cosine schedule can produce steps with only one masked token, so the masking ratio is lower-bounded to ensure NM > 0.45S.
  • Quantization-level configurations: Figure 10 varies quantization levels per channel, with color and marker indicating the smallest Li used for each model.
Loading 2309.15505v2…