Learning Optimal Lattice Vector Quantizers for End-to-end Neural Image Compression

• W1. [Omission of integral operation]
  Thanks for pointing out the lack of rigor in our math notation. Hope you can understand that it is only an innocent omission of an integral symbol. We will write out, in the final version, the corrected probability mass function (PMF) over the lattice quantized vector:
  $p_{m} = \int_{V(m)} \sum_{k=1}^{K} \mathbf{\Phi}_k \mathcal{N} (x; \mu_k, \Sigma_k) dx$

  $p_{m} = \prod_{i=1}^n \int_{V(m^i)} \sum_{k=1}^{K} \mathbf{\phi_k}^i \mathcal{N} (x^i; \mu_k^i, \sigma_k^i) dx^i$ where $V(m^i) = \lbrace x^i | Q(x^i)=m^i \rbrace$.  

• W2. [Rate estimation of traditional LVQ and GVQ]
  The rate estimation for traditional lattice VQ follows the method described in LVQ-VAE [1]. For general VQ, the rate estimation method follows the approach outlined in NVTC [2]. For more implementation details, please refer to their respective papers.
   

• W3. [Effectiveness of orthogonal constrai]
  Sorry, we should have made this point more clearly.
  The orthogonal constraint is needed to reduce the gap between lattice vector quantization during training and inference. As stated in Section 3.3, because lattice vector quantization is not differentiable (and thus cannot be directly integrated into end-to-end training), we propose a differentiable alternative using Babai’s Rounding Technique (BRT). BRT estimates a vector that is sufficiently close to, but not necessarily exact, the closest lattice point. However, our experiments revealed that the BRT-estimated lattice vector may be far away from the exact closest lattice point if the cell shape of the optimized LVQ is arbitrary and unconstrained. This discrepancy causes an inconsistency between training and inference, if the learnt lattice vector quantizer is employed in the inference stage.
  As analyzed in [3], the BRT-estimated lattice point will match the closest lattice point if the basis vectors of the generator matrix are mutually orthogonal. For this reason, we propose to impose orthogonal constraints on the basis vectors of the generator matrix, enhancing the accuracy of the BRT-estimated lattice point and reducing the gap between lattice vector quantization during training and inference.
  The performance without the orthogonal constraint will drop by about 4.2% for BD-rate. We report the detailed performance numbers with and without the orthogonal constraint in the following table.
  We will make the above clarifications in the final version.
    | Entropy Model | Bmshj2018 w/ orthogonal constraint | Bmshj2018 w/o orthogonal constraint | SwinT-ChARM w/ orthogonal constraint | SwinT-ChARM w/o orthogonal constraint | | ---- | ---- | ---- | ---- | ---- |
  | Factorized (w/o context) | -22.60% | -18.2% | -12.44% | -8.09% | | Checkerboard | -11.94% | -7.94% | -6.37% | -2.25% | | Channel-wise Autoregressive | -10.61% | -6.34% | -5.61% | -1.22% | | Spatial-wise Autoregressive | -8.31% | -4.20% | -2.31% | +2.08% |  

• We provide the RD-curves of different methods in the rebuttal PDF. It can be seen that the proposed optimal lattice vector quantization consistently improves the existing compression models in both low and high bitrate cases. However, the performance gain is more pronounced at lower bitrates.
   

[1]. Kudo, S., Bandoh, Y., Takamura, S., & Kitahara, M. LVQ-VAE: End-to-end Hyperprior-based Variational Image Compression with Lattice Vector Quantization. OpenReview 2023.
[2]. Feng, R., Guo, Z., Li, W., & Chen, Z. NVTC: Nonlinear vector transform coding. CVPR 2023.
[3]. Babai L. On Lovász’ lattice reduction and the nearest lattice point problem. Combinatorica, 1986.

Dear Reviewer,

The authors have posted an author response here. Could you go through the rebuttal and update your opinion and rating as soon as possible?

Your AC

Thanks for your feedback. We are glad to clarify your further concerns regarding the entropy model design and orthogonal constraint.
 

• RD curves
  The RD curves for the SOTA model LIC-TCM follow a similar trend to the bmshj2018 model, although the gains are smaller at low bitrates. Unfortunately, we are unable to upload images or files at this stage, so we cannot provide the RD curves right now. However, we will include these RD curves and discuss them in future revisions.
   

• Entropy model and orthogonal constraint
  You are correct in understanding that we factorized the distribution of quantized latent vectors. However, the integration area is not a simple hypercube; it’s more accurately described as a hyperrectangle, with different lengths for each dimension. We recognize that approximating the integral over such an area cannot precisely match the true region of LVQ. This approximation is a necessary compromise during training, as integrating over a dynamically changing region of arbitrary shape and size is impractical (at least to us). While this does introduce a gap between training and inference, the orthogonal constraint helps mitigate this by constraining the LVQ shape to a rotated hyperrectangle, though some discrepancy remains due to rate estimation over the non-rotated hyperrectangle.
   
  During inference, we calculate the integration over the arbitrarily shaped Voronoi polyhedron using Monte Carlo integration, a numerical integration method involving random sampling, as also employed in LVQ-VAE. It is noteworthy that the rate estimation for traditional lattice VQ follows the same approach in our experiments.
   

References

[1]. Kudo, S., Bandoh, Y., Takamura, S., & Kitahara, M. LVQ-VAE: End-to-end Hyperprior-based Variational Image Compression with Lattice Vector Quantization. OpenReview 2023.

Dear Reviewer,

We hope that our rebuttal has effectively addressed your concerns. We are grateful for the time and effort you’ve dedicated to reviewing our paper. Your valuable feedback and suggestions have been very helpful, and we will certainly incorporate them into our future revisions. Thank you once again for your time.

Thanks for your replying. I have carefully read all the review and rebuttal. I still have concern about the entropy model design and orthogonal constraint. Additionally, the rebuttal only provides RD curves for the bmshj2018 model. Could you also provide the RD curves for the SOTA model LIC-TCM mentioned in Table 1?

Based on the revised PMF function and your response W2 to Reviewer B5jh, it seems that you factorized the distribution of quantized latent vectors instead of the distribution of integer indices during training. This implies that you approximate the shape of integration area from Voronoi polyhedron to hypercube. Such an approximation would lead to incorrect rate estimation results. Only the scalar quantization could perfectly match the hypercube integration area. LVQ with orthogonal matrix would rotate and scale the hypercube, not to mention LVQ with other matrices. Moreover, in rebuttal you claim that orthogonal constraint is used to reduce train-test mismatch of BRT-estimated lattice point and exact lattice point. The test-time LVQ would further increase the gap between the approximated PMF and actual PMF. I am curious to know how you calculate the integration over the Voronoi polyhedron area during inference.

I would prefer to keep the score unless these concerns are fully addressed.

• W1. [Separated evaluation]
  We provide the results which are separately evaluated in the following table (BD-rate over uniform scalar quantizer). It can be observed that the performance trend is consistent across the separate and combined evaluations.
    | Entropy Model | Bmshj2018 Kodak | Bmshj2018 CLIC | SwinT-ChARM Kodak | SwinT-ChARM CLIC | | ---- | ---- | ---- | ---- | ---- |
  | Factorized (w/o context) | -20.61% | -23.15% | -11.16% | -13.12% | | Checkerboard | -10.53% | -12.49% | -5.28% | -6.74% | | Channel-wise Autoregressive | -10.32% | -11.21% | -4.33% | -6.12% | | Spatial-wise Autoregressive | -7.12% | -8.97% | -1.98% | -2.65% |  
• W2. [R-D curves]
  We provide the RD-curves of different methods in the rebuttal PDF. It can be seen that the proposed optimal lattice vector quantization consistently improves the existing compression models in both low and high bitrate cases. However, the performance gain is more pronounced at lower bitrates.
   
• W3. [Characteristics of generator matrix]
  Thanks for your great suggestion. We analyze the deformation of the learned generator matrix by measuring the ratio of the longest to the shortest basis vectors, which turned out to be 6.8, indicating extreme asymmetry in the optimized lattice structure. Additionally, the angle between the longest and the shortest basis vectors is approximately 60 degrees. These characteristics of the learned generator matrix indicate that the cell shape of the optimized LVQ is significantly different from that of a uniform scalar quantizer.
   
• Q1. [Effectiveness of orthogonal constraint]
  Sorry, we should have made this point more clearly.
  The orthogonal constraint is needed to reduce the gap between lattice vector quantization during training and inference. As stated in Section 3.3, because lattice vector quantization is not differentiable (and thus cannot be directly integrated into end-to-end training), we propose a differentiable alternative using Babai’s Rounding Technique (BRT). BRT estimates a vector that is sufficiently close to, but not necessarily exact, the closest lattice point. However, our experiments revealed that the BRT-estimated lattice vector may be far away from the exact closest lattice point if the cell shape of the optimized LVQ is arbitrary and unconstrained. This discrepancy causes an inconsistency between training and inference, if the learnt lattice vector quantizer is employed in the inference stage.
  As analyzed in [1], the BRT-estimated lattice point will match the closest lattice point if the basis vectors of the generator matrix are mutually orthogonal. For this reason, we propose to impose orthogonal constraints on the basis vectors of the generator matrix, enhancing the accuracy of the BRT-estimated lattice point and reducing the gap between lattice vector quantization during training and inference.
  The performance without the orthogonal constraint will drop by about 4.2% for BD-rate. We report the detailed performance numbers (BD-rate over uniform scalar quantizer) with and without the orthogonal constraint in the following table.
  We will make the above clarifications in the final version.
    | Entropy Model | Bmshj2018 w/ orthogonal constraint | Bmshj2018 w/o orthogonal constraint | SwinT-ChARM w/ orthogonal constraint | SwinT-ChARM w/o orthogonal constraint | | ---- | ---- | ---- | ---- | ---- |
  | Factorized (w/o context) | -22.60% | -18.2% | -12.44% | -8.09% | | Checkerboard | -11.94% | -7.94% | -6.37% | -2.25% | | Channel-wise Autoregressive | -10.61% | -6.34% | -5.61% | -1.22% | | Spatial-wise Autoregressive | -8.31% | -4.20% | -2.31% | +2.08% |  

[1]. Babai L. On Lovász’ lattice reduction and the nearest lattice point problem. Combinatorica, 1986.

• W1. [Compared with other VQ based image compression methods]
  A good suggestion. We have addressed your concern by adding new comparison results with the following vector quantization-based image compression methods: SHVQ (NeurIPS ’17) [1] and McQUIC (CVPR ’22) [2]. From the provided table (BD-rate over uniform scalar quantizer), it can be observed that the proposed learnable LVQ method not only outperforms these VQ-based image compression methods in terms of BD-rate but also demonstrates superior performance in computational complexity. As shown, our method achieves lower BD-rates, i.e., better compression performance. At the same time, the computational complexity of our approach is significantly lower than SHVQ and McQUIC, making it more practical, especially on resource-constrained devices.
    | Quantizers | Bmshj2018 Factorized | Bmshj2018 Checkerboard | SwinT-ChARM Factorized | SwinT-ChARM Checkerboard | Inference time | | ---- | ---- | ---- | ---- | ---- | ---- | | Optimized LVQ | -20.49% | -10.71% | -10.68% | -5.89% | ~40ms | | SHVQ [1] | -8.23% | -5.18% | -6.55% | -1.92% | ~328ms | | McQUIC [2] | -15.11% | -7.93% | -7.24% | -3.18% | ~771ms |  
• W2. [Performance at different bitrates]
  To address your concern, we provide the rate-distortion curves of different quantization methods in the rebuttal PDF. It can be observed that the proposed optimal lattice vector quantization consistently improves the existing compression models in both low and high bitrate cases. However, the performance gain is more pronounced at lower bitrates.
   
• W3. [Impact of $\lambda_2$ and the loss term L]
  To address this, we have conducted ablation studies to thoroughly examine how these two hyperparameters affect the performance of our method.
  • For $\lambda_2$, we employed a search strategy to identify the optimal value that yields the best performance. Our findings indicate that both excessively large and excessively small values of $\lambda_2$ can lead to performance degradation. When $\lambda_2$ is too large, it overly constrains the shape of the optimized lattice codebook, resulting in a sub-optimal lattice structure. Conversely, when $\lambda_2$ is too small, it loosens the orthogonal constraint, causing a significant gap between lattice vector quantization during training and inference (Please refer to the global rebuttal for more details).
  • Regarding the loss term L, our initial experiments utilized the MSE metric. To further investigate, we have now included a study using the SSIM metric as the loss term L. The results from this study indicate that our conclusions drawn from using the MSE metric remain valid when using the SSIM metric. Specifically, the performance trends and the effectiveness of our proposed method are consistent across both metrics (Please refer to the rebuttal PDF for the R-D curves of SSIM metric).
     
• Q1. [Details of inference time]
  Sorry for missing the details. The quantization times are not measured for quantizing a single token; instead, they are measured for quantizing a $512 \times 512 \times 1$ feature map. The GPU utilization is approximately 25%, 28%, and 42% for Uniform Scalar Quantization (USQ), Lattice Vector Quantization (LVQ), and General Vector Quantization (GVQ), respectively.
   
• Q2. [Metric for BD-rate] Yes, the BD-rates are measured in PSNR.
   

[1]. Agustsson, E., Mentzer, F., Tschannen, M., et al. Soft-to-hard vector quantization for end-to-end learning compressible representations. NeurIPS 2017.
[2]. Zhu, X., Song, J., Gao, L., et al. Unified multivariate gaussian mixture for efficient neural image compression. CVPR 2022.

Dear Reviewer,

We hope the additional clarifications appropriately addressed your further questions. We appreciate the time and effort you’ve dedicated to reviewing our paper. Your valuable feedback and suggestions have been immensely helpful, and we will certainly incorporate them into our future revisions. Thank you again for your time.

• W1. [General vector quantization]
  We build the general vector quantization model based on NVTC (CVPR'23) [1].
   
• W2. [Entropy model]
  Sorry for the discrepancy. The entropy model is optimized over the lattice quantized vectors and used to code the same vectors into the bitstream. The idea of coding the integer combination coefficients is still under investigation.
   
• W3. [Limited gain with autoregressive model]
  We acknowledge this limitation and have discussed it in the paper. However, as well known, context models, especially autoregressive ones, tend to slow down the coding speed dramatically. Our method is meant to build LIGHTWEIGHT models for deployment on edge devices. As these devices have severe resource constraints, the model efficiency and simplicity become paramount. Forgoing computationally intensive context models offers a practical solution for real-time applications on such devices.
   
• W4. [Toy example plots]
  We appreciate the value of visualizing the lattice structure generated by the proposed method. However, plotting the lattice structures in more than three dimensions remains a challenge; in fact, we have not seen such drawings in any prior literature on the topic.
   
• Q1. [Optimized LVQ + factorized model vs. autoregressive model]
  We have included R-D curves of different methods in the rebuttal PDF to address your concerns. It can be observed that the proposed learnable LVQ method coupled with the factorized entropy model can approach the performance of the autoregressive entropy model, particularly at low bitrates. For high bitrates, the autoregressive entropy model still has an edge. Given this, we believe that the proposed LVQ method can serve as a viable alternative to the autoregressive entropy model, especially in scenarios where low bitrate performance is crucial and computational efficiency is a priority. The factorized entropy model, being less computationally intensive, combined with our LVQ method, offers an attractive trade-off between performance and efficiency.
   

[1]. Feng, R., Guo, Z., Li, W., & Chen, Z. NVTC: Nonlinear vector transform coding. CVPR 2023.

Dear Reviewer,

The authors have posted an author response here. Could you go through the rebuttal and update your opinion and rating as soon as possible?

Your AC

Dear Reviewer,

We hope the rebuttal has appropriately addressed your concerns. Thank you for the time and effort you have invested in reviewing our paper. Your insightful feedback and suggestions have been extremely beneficial, and we will make sure to include them in our future revisions. We appreciate your continued support and time.