Paper: Finite Scalar Quantization: VQ-VAE Made Simple (ICLR-2024)

Reviewer comment (Evaluation):
FSQ's expressive capability is slightly weaker than VQ, and its performance is slightly worse with small codebooks.

Author response:
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.

Reviewer comment (Evaluation):
The paper lacks analysis on the speedup and compression rate brought by the quantization of FSQ.

Author response:
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).

Reviewer comment (Experiments):
FSQ also faces difficulties in finding proper quantization hyperparameters (dimensions and number of quantization levels).

Author response:
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)

Reviewer comment (Evaluation):
The authors should add some intuitive numerical comparisons to demonstrate the inference speedup brought by FSQ, and emphasize that FSQ does not need the parameters for the codebook.

Author response:
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).

Reviewer comment (Novelty):
The main weakness lies in the originality/novelty of the method which appears to be a compressive autoencoder (e.g., [1,2]) applied to generative tasks. As the Method section demonstrates, there are very few (no) technical/theoretical contributions.

Author response:
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.”

Reviewer comment (Evaluation):
The results are still (mostly slightly) worse than VQ.

Author response:
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.

Reviewer comment (Evaluation):
The method obviously lacks compression performance since the quantized vectors' indices are insufficient/irrelevant in this case.

Author response:
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).

Reviewer comment (Evaluation):
There is an obvious difference between the FSQ and VQ-VAE functionals (rounding vs L2 distance). It would be interesting to see the difference in the quality of the representations (e.g., same class samples should be close(r with FSQ?) in their representations) and not only regarding the optimized metrics.

Author response:
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.

Reviewer comment (Theory):
Following 4., it is not clear how the codebook is seen as implicit and not simply defined by the hypercube's samples.

Author response:
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.

Reviewer comment (Evaluation):
How do you explain that extending the model capacity (Sec 3.3) does not result in improvement, especially since the method is still worse than VQ-VAE?

Author response:
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.

Reviewer comment (Reproducibility):
The relation between VQ-VAE and FSQ is confusing. Specifically, I do not understand how the FSQ model is trained. Is the FSQ model trained from scratch using the VAE architecture, or is the FSQ model transformed from a pretrained VQ-VAE model in some way. 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.

Author response:
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).

Reviewer comment (Reproducibility):
As I mentioned in the weakness part, is FSQ used in MaskGIT and UViM trained from scratch? If yes, what is the model architecture and the other hyperparameters like training epochs? And what is the reconstruction and generation ability of FSQ used as a VAE?

Author response:
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.

Reviewer comment (Experiments):
There is another concurrent work named Lookup Free Quantization [1], which also aims to compress the codebook. Their method directly let the dimension of codebook features become zero, which can also support a larger vocabulary size. Can the authors compare the proposed method with theirs?

Author response:
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.