Finite Scalar Quantization: VQ-VAE Made Simple

Thanks for the review, and the kind words in the Strenghts section!

Regarding the capability difference, and small codebook differences: We agree, and we note this in the text (Sec 5.1, “Codebook size correlates with Reconstruction FID for FSQ”). However, we also note that FSQ has more favorable scaling behavior (Fig 3), which we see as a very promising direction for future work: we note that 1024 (2^10) is not very large, compared to language modeling, where people routinely use 32k+ codewords. In general, we think that approaches relying on VQ mostly did not further scale the codebook because VQ becomes unstable as one increases the codebook, but a priori, if one can afford a big enough model in the second stage, it should be desirable to increase the codebook. We see this empirically for FSQ in Fig. 3, where scaling the codebook improves sampling FID, without requiring a larger stage 2 model.

Regarding speed (weakness 2 and question 1): we note that we do not claim inference speedups from FSQ in the paper. We note that FSQ is only faster when calculating indices (since VQ requires to compare a vector with every possible vector in the codebook, it has O(|C|) cost of computing an index, whereas FSQ has O(d) cost, as we just need to run a single dot product (see Sec. A.1)). For inference, when we sample from the transformer and then feed the result to the VAE decoder, both methods have the same theoretical runtime. We validated this with a small benchmark added to Section A.5 in the updated PDF.

Regarding the compression rate, we see that FSQ can utilize the codebook better than VQ (Fig 3, d), and thus uses more bits when used as a compressor. This results in better sampling FID, since the transformer can create high information representations (Fig 3, b).

Regarding the number of quantization levels, we found the rule to use L >= 5 to work well across tasks, and report good choices for d and L for different target codebook sizes in Table 1. We note that the principle of using L>=5 is validated empirically, see Sec A.4.1 Trade-off study, where we found that codebooks with L<5 also work well, but L>=5 performs a bit better (the difference is <1 FID). Intuitively, we think this is because of the straight through estimation (STE) that we use: for this estimation, we approximate the gradient of the round operation with the gradient of the identity function. This approximation is very bad for L=2, when the quantizer becomes a step function, but gets better and better as we increase L, with diminishing returns as L further increases. We added this intuition to A.4.1, see updated PDF (changes in pink)

Thanks for the review, and the strengths you point out.

Regarding novelty: first, we note the link to compressive autoencoders in the paper (Sec 1, 2) and do not claim novelty on the scalar quantizer itself. However, as we note in the general reply, our contribution lies in considering the implied product codebook of FSQ: “We’d like to emphasize that prior to this paper, no one was thinking of using finite scalar quantization (FSQ) as a drop-in for VQ in generative image models or for dense vision tasks. Our key insight is to consider the implicit product codebook created by scalar quantization if you bound each channel, which allows training large generative transformers like on VQ representations (i.e., by bounding each of the $d$ channels to $L$ levels, we get a product codebook of $L^d$ codewords). While it might seem obvious after reading our paper, a priori, it is very surprising that FSQ obtains such competitive results compared to VQ -- despite not using any of its auxiliary losses and tricks.”

Regarding comparison to VQ: we note that in Fig. 3, we see that FSQ scales well to high-bit codebooks. Arguably, this is where the future research lies, as small codebooks are becoming less and less important. We think that people focused on eg 1024 codewords, because VQ did not scale beyond that. Our contribution lies in showing that our simple approach can be highly competitive even though VQ is highly optimized and much more challenging to train.

Regarding “The method obviously lacks compression performance since the quantized vectors' indices are insufficient/irrelevant in this case.”: we note that in this work we do not target compression, which typically operates at much higher bitrates (note that our codebooks of e.g. 10 bits (1024 codewords) correspond to 0.03 bits per pixel without any entropy coding). One can increase the codebook size and eventually obtain a competitive compression method, similar to what was studied in the compression literature (see “Neural Compression” references in Section 2). To improve rate-distortion performance, one probably wants to leverage entropy coding to get an even more compressed bitstream. However, we note that we target generative downstream tasks, and we see in Fig 3d that FSQ can utilize the codebook better than VQ. This is desirable for the tasks we study, and results in better sampling FID (Fig 3b).

Q1) Thanks for this suggestion, this also made us curious. We calculated the representations for 3 images: 2 from class A, and 1 from class B. When calculating the MSE between the two images in class A in FSQ representation space, we obtain 1.5e5. When calculating the MSE between one image from A and the image from B, we obtain 1.4e5. For VQ, the same story emerges, we get 1.67e5 for in-class, and 1.53e5 for between-class. In both cases, see that images from the same class do not lie closer in the representation space. This is in line with what we already report in Section A.3., where we concluded that the representations of both FSQ and VQ mostly store medium level textures. Inspired by your comment, we added a few results to Section A.3:

• Fig 10, which shows the codewords used for two different imagenet classes when encoding 20 examples. As it turns out, the codewords seem to represent textures more than semantic concepts: both models use most codewords to represent very different classes. This explains that we do not see a close relationship when calculating MSE above.
• Fig 11, where we show what happens when adding offsets to the FSQ channels. Again, only low level texture details change.

Overall, both VQ and FSQ seem to store mostly medium level texture detail. Note that both approaches store 1 codeword per 16x16 pixel patch.

Q2) we call the codebook "implicit" because it is defined by the product of the per-channel codebook (eg for L=3, d=2, the per channel codebook is {-1, 0, 1}, and the implied codebook is given by {(-1, -1), (-1, 0), (-1, 1), …}). As you point out, this could be implemented by using VQ with a fixed codebook set to the hyper cube, but in our approach, we never have to construct this codebook: during encoding, the index into the implicit codebook can be calculated with a single dot product (see Sec. A.1), and we thus call it implicit.

Q3) This is likely because in the experiment we ran, the capacity we add is comparatively small, e.g., VQ codebooks for MaskGIT store 1024 * 256 parameters, which is <0.1% of the 227M parameters of the entire model (VAE + transformer). In general, it should help to add a substantial amount of parameters in exchange for more memory and compute requirements.

Thanks for the review, and for the kind words. Regarding the weaknesses and questions:

W1) Apologies for the confusion around training, what is happening is that we train what you could call a “FSQ-VAE” instead of a VQ-VAE. Then we train transformers in a second stage. Your understanding thus is correct (“Since MaskGIT is a two stage method, where the first stage is to train a discrete tokenizer, I will regard FSQ as a model trained from scratch.”) -- We updated the PDF with the following: “We note that like in VQ-VAE based approaches, we train FSQ models in two stages: We first train a FSQ-VAE, and then a transformer on frozen FSQ-VAE representations.” (Sec 3, change highlighted in pink).

W2) The principle of using L>=5 is validated empirically, see Sec A.4.1 Trade-off study, where we found that codebooks with L<5 also work well, but L>=5 performs a bit better (the difference is <1 FID). Intuitively, we think this is because of the straight through estimation (STE) that we use: for this estimation, we approximate the gradient of the round operation with the gradient of the identity function. This approximation is very bad for L=2, when the quantizer becomes a step function, but gets better and better as we increase L, with diminishing returns as L further increases. We added this intuition to A.4.1, see updated PDF (changes in pink).

Q1) Indeed, our FSQ-VAEs are trained from scratch. We use virtually the same hyper parameters as the baseline VQ based methods, relying on the public code. We clarified this in the updated draft in A.4.3 and the newly added A.4.4. With the respective GitHub repos and our code in Sec A.1., our paper can be reproduced. Regarding the reconstruction ability: thanks for this suggestion, we added Fig 13 to the paper (in the supplementary material). Both VQ and FSQ behave very similarly in terms of reconstructions, we see FSQ is slightly more consistent with the input in terms of color.

Q2) Thanks for pointing out this concurrent work. We note that the authors still rely on an additional entropy maximization loss (Equation 5 in their paper), whereas FSQ does not use any auxiliary losses. It will be interesting to see if future work sees a benefit from adding an entropy maximization loss to our method.

Thanks for the detailed explanation. The illustration on the two-stage training address my concern. I love the potential of this work to scale up the size of the codebook. Thus I will raise my score from 5 to 8.

Additionally, there is another work named MAGE [1] which also relies on discrete tokenization to reach supreme classification performance and SOTA generation ability among MIM methods. MAGE is a variant of MAE [2] where the input and prediction target of MAGE are both discrete tokens. I believe this work will shed some lights on the improvement of MAGE.

[1] MAGE: MAsked Generative Encoder to Unify Representation Learning and Image Synthesis. CVPR 2023. [2] Masked Autoencoders Are Scalable Vision Learners. CVPR 2022.

Thanks for the review, and thorough summary, your understanding of the method is correct.

The main concern of yours seems to be “if the method is so simple and works so well, why has it not been studied before”? As we wrote in the general reply, we think that this is because a crucial insight was missing: "Our key insight is to consider the implicit product codebook created by scalar quantization if you bound each channel, which allows training large generative transformers like on VQ representations (i.e., by bounding each of the $d$ channels to $L$ levels, we get a product codebook of $L^d$ codewords)." We updated the Intro with a clarifying sentence.

Regarding limitations: We believe FSQ-VAEs can work as good or better as VQ-VAEs in any task involving a VAE, and we showed 4 different vision tasks in the paper (generation, colorization, depth estimation, semantic segmentation). As we note at the start of Sec. 3, intuitively, FSQ works well here because the generality of a learned codebook is “absorbed” into the encoder transform, which can morph the data so that it is well quantized by FSQ.

Regarding tanh: We explored using sine instead of tanh but found it to perform worse. There might however be future work that shows a better bounding function than tanh.