Image and Video Tokenization with Binary Spherical Quantization

1. VQ in image reconstruction

We will follow up with the results once the training is done. We report our reproduced VQ with 18 bits on ImageNet 256x256 below (BSQ numbers are from Table 1)

Although the parameter number in ViT is larger than CNN-based approaches, the measured FLOPs are significantly smaller than CNNs. We compare SD-VAE 1.x and BSQ-ViT in the table and observe that ViT requires 36% FLOPs compared to CNN in SD-VAE.

2. VQ or LFQ-based method

We add VQ and LFQ results below. For LFQ, we use the checkpoint from OpenMAGVIT2, an open-source implementation of MAGVIT2.

3. Image generation with the same number of steps.

We added the result with the same number of steps. We see a slight decrease in FID. The Inception Score decreases more significantly. The reason is that during masked language modeling, we divide each BSQ token into two sub-groups. This increases the effective sequence length, requiring a larger number of decoding steps to generate more details.

4. Proof of Eq 13

The sketch of deriving $d_{max} (\mathbf{u}, \mathbf{\hat{u}})$ is as follows. For an arbitrary vector $\mathbf{u}$ residing on the unit hypersphere ($|\mathbf{u}|_2 = 1$), its quantized counterpart $\mathbf{\hat{u}} = \frac{1}{\sqrt{L}} \mathrm{sign}(\mathbf{u})$ also resides on the same unit hypershpere. We can write the quantization error as $d(\mathbf{u}, \mathbf{\hat{u}}) = \sqrt{2 - 2 \mathbf{u}\cdot\mathbf{\hat{u}}}$. The maximum of $d(\mathbf{u}, \mathbf{\hat{u}})$ is equivalent to finding the minimum of $\mathbf{u} \cdot \mathbf{\hat{u}} = \sum_i u_i \hat{u}_i$. Note that $\hat{u}_i = \frac{1}{\sqrt{L}} \mathrm{sign}(u_i)$. This leads to $u_i \hat{u}_i = \frac{1}{\sqrt{L}} | u_i |$ and $\mathbf{u} \cdot \mathbf{\hat{u}} = \frac{| \mathbf{u} |_1} {\sqrt{L}}$. Minimizing $\mathbf{u} \cdot \mathbf{\hat{u}}$ is then equivalent to finding an $\ell_2$ unit vector within the minimal $\ell_1$ norm, which is known to be achieved only if the vector $\mathbf{u}$ is a one-hot vector and the minimal $\ell_1$ norm is 1.

This leads to $d_{max} (\mathbf{u}, \mathbf{\hat{u}})$ being $\sqrt{2 - \frac{2}{\sqrt{L}}}$ .

We have added the proof in the appendix. Please refer to the rebuttal revision on Line 936.

We thank the reviewer for spotting the typos. We have fixed all of them in the rebuttal revision, highlighted in red underlined text.

1. How to address different resolutions.

In our reconstruction experiments, the input for image/video reconstruction experiments is either 128x128 or 256x256. Therefore, we fix the image resolution to be 128 or 256 and train two models from scratch separately. For image/video compression experiments, we partition the image/video into multiple 128x128 crops and apply our BSQ-ViT(128x128) on all crops. For cases where the edge is not divisible by 128 (e.g. UVG-1080p), we pad the last row/column to be 128x128 with repeat padding.

It is also possible to adapt the model with 128x128 inputs to accept 256x256 inputs by linearly interpolating the patch embedding layer and fine-tuning the model with the new input size, e.g. like FlexViT (https://arxiv.org/pdf/2212.08013).

2. LPIPS comparison.

This is a great suggestion! We include LPIPS scores for image and video compression benchmarks.

We can see that our method outperforms JPEG2000 and WebP significantly. We also show three compressed images from the Kodak dataset at https://ibb.co/5GH1V1J. We see the compressed images generated by BSQ-ViT preserve more details. On the MCL-JCV (640x360) Video compression benchmark, our method also yields a much lower LPIPS score than H.264 and H.265 under the same BPP.

3. Minor issues.

Thanks for spotting these issues! (1) should be correct. (2) and (3) are typos and we have fixed them in the rebuttal revision.

4. Arithmetic Coding for image compression.

This is a great point! We train an auto-regressive transformer to model the conditional distribution for arithmetic coding. The average BPP reduces from 0.2812 to 0.2073, a 26% reduction (when $L=18$). This result is already close to the VTM result reported at https://github.com/InterDigitalInc/CompressAI/blob/v1.2.6/results/image/kodak/vtm.json. We will update the number once the auto-regressive transformer finishes training. We update the rate-distortion curve in the Appendix Figure 12.

5. Intra-frame information in video compression.

Video compression uses intra-frame information in two aspects. First, the Transformer-based autoencoder with blockwise casual mask models intra-frame information via self-attention within each block. The block causal masks ensure that each token can attend to all other tokens in the same frame as well as tokens in all previous frames. Second, the auto-regressive transformer for arithmetic coding jointly models the intra-frame and inter-frame information.

1. Rate-Distortion curve.

This is a great suggestion! We provide a more detailed rate-distortion (RD) comparison at https://ibb.co/GpBV47b. The figure can also be found in the Appendix Figure 12. From the RD curve, we can see that BSQ consistently improves MS-SSIM with increasing bpp. VQ-AE shows diminishing gain at higher bpp. We ascribe this to the difficulty of training VQ when the codebook size increases. As per Reviewer PNuG's comments, we run an auto-regressive Transformer to compress the image tokens with arithmetic coding. This further leads to a ~26% reduction of bpp and the final result is very close to the VTM result reported at https://github.com/InterDigitalInc/CompressAI/blob/v1.2.6/results/image/kodak/vtm.json. We will add more data points in the final revision.

2. Testing settings for ablation studies.

(1) We use the same encoder-decoder while only changing the quantization bottleneck. We train all models on ImageNet-1k 128x128 from scratch.

(2) An alternative way is to train an AE first and learn the quantizer on top of the encoder while freezing the encoder and decoder. We will follow up with the result once the training is done. We observe separating training encoder-decoder and bottlenecks does not yield satisfying results for all quantization methods (reconstruction loss shown in the table below).

The implicit codebook in BSQ and LFQ employs geometric prior of the embedding at the bottleneck, whereas training an Auto-Encoder does not. On the contrary, the explicit codebook in VQ is freely learnable and has more parameters ($O(2^L\times d)$ vs. $O(L\times d)$ in BSQ and LFQ), leading to slightly better performance. Nevertheless, none of the methods are comparable to the default end-to-end training setup ("e2e" in the table). It is hard for an autoencoder (without a quantization bottleneck) to learn low-variance latent space without adding inductive biases [1,2]. Therefore, we believe it makes more sense to train an auto-encoder with a quantization bottleneck in an end-to-end manner which is the default setup in our ablation.

[1] Kingma and Welling. "Auto-encoding variational bayes." ICLR 2014. [2] Higgins, et al. "beta-vae: Learning basic visual concepts with a constrained variational framework." ICLR 2017.

3. Computation complexity and encoding/decoding time for comparison

We left the comparison of computation costs in the appendix due to limited space. Please refer to Table 9. Due to the simplicity of the Transformer-based encoder and decoder, our method runs faster than VCT. For H.264, we use FFmpeg’s API and only report the overall FPS.

1. Effectiveness of BSQ within a CNN-based model.

BSQ is effective within a CNN-based autoencoder. We use the CNN-based encoder-decoder similar to SD-VAE. The only architectural change is that the encoder outputs an $L$-dim vector which is then $\ell_2$-normalized and fed into binary quantization without an external projection layer ($\mathrm{z} \rightarrow \mathrm{v}$) in ViT. We will post the final result once the training is done. Update: We train a BSQ-ViT and BSQ-CNN with a shorter schedule (200K iterations). The results are shown below. We can see that BSQ-CNN works comparably well with BSQ-ViT, validating the effectiveness of BSQ within a CNN-based model.

2. BSQ without this causal masking.

We show VQ without blockwise causal masking in the last but fourth line in Table 2. We denote it with “non-BC” ViT for short.

We will add BSQ without blockwise causal masking in the final revision.

3. Subjective evaluation metrics of image compression.

We show three compressed images from the Kodak dataset at https://ibb.co/5GH1V1J. The figure has also been added in the appendix (Figure 11)

4. BSQ vs VQ.

In the ideal scenario, VQ works better than BSQ in that it leverages the entire high-dimensional space while BSQ only leverages the hypercube corners. However, in reality, VQ is difficult to train when the codebook size increases, as widely acknowledged in previous literature and shown in our ablation (Table 5). We also provide a more detailed rate-distortion (RD) comparison at https://ibb.co/GpBV47b. The figure can also be found in the Appendix Figure 12. We can see VQ-AE falls behind BSQ-AE when bpp increases. In addition, given the same bottleneck of $L$ bits, VQ's parameter space is $O(d\times 2^L)$ while BSQ is only $O(d\times L)$.

5. Reproducing LFQ.

When we submitted the manuscript we re-implemented the LFQ following the parameters in the original paper. To ensure the results are reproducible, in the discussion period, we additionally implemented LFQ using OpenMAGVIT2's implementation (https://github.com/TencentARC/Open-MAGVIT2). We use their recommended hyper-parameters. Since the quantization loss uses a different set of weights, we compare the reconstruction loss at https://ibb.co/PQNHWtv. We can see the reconstruction loss of LFQ converges much slower. This is consistent with our original ablation in the submission that LFQ is difficult to train within the ViT-based autoencoder architecture due to unbounded quantization error. We will report the final revision. Update: We report the best result we achieved during the discussion period compared to the one we got when we submitted the manuscript. PSNR, SSIM, and LPIPS score improve while rFID remains similar. Neither case is comparable to the proposed BSQ.

Thank you for the follow-up comments!

1. We include the result of LFQ+CNN for fair comparison below. With the same CNN backbone, BSQ achieves better PSNR, SSIM, LPIPS, and comparable FID scores.

2. We have added the LPIPS scores in the revised manuscript as requested by Reviewer PNuG. We further compute the VMAF (v0.6.1) metric. We post both numbers below.

Our method outperforms H.264 and H.265 in terms of LPIPS but underperforms in terms of VMAF. We look at the results and observe that our method works better on videos with slow motion but falls behind on videos with large motions. We uploaded two representative videos at https://file.io/c7WKp9Gacb45. The reason is probably because we fine-tune the video auto-encoder on UCF-101 only, whose diversity and quality limit the performance.