Approaching Rate-Distortion Limits in Neural Compression with Lattice Transform Coding

Thank you for taking the time to read the paper and provide valuable feedback! Please find our responses to your comments in the following point-by-point response, which refers to the revised manuscript.

Response to weaknesses:

W1. We have improved the clarity of some of the writing, such as more detailed comparison of the methods for image compression, clarifying some of the motivation for nested LQ and block LQ, the role of lattice quantizer dimension, and hyperparameters such as $\lambda$, as suggested by other reviewers as well. Regarding Sec. 3, we have added additional details at the end of the section to reflect the larger gap between NTC and optimal performance as the dimension grows, to further support the suboptimality of scalar quantization.

W2. Thanks for your suggestion, and we have clarified the block LTC section (4.3) accordingly. To clarify, neural compression has traditionally been designed for images, which naturally results in a one-shot compressor design (compresses one source vector at a time). In information theory, it is known that compressing multiple i.i.d. realizations simultaneously achieves optimality and outperforms one-shot compressor designs. This can be useful in practice for certain settings such as dataset compression, e.g. when one wishes to compress a dataset of scientific measurements or audio samples, where each sample is independent, like the Physics and Speech datasets. In the neural compression literature, no such block coding scheme currently exists. Our proposed block LTC method provides a low-complexity framework to realize block coding, by building on LTC to apply LQ across iid latent vectors. To put things in perspective, the LTC framework recovering VQ and achieving optimality for vector source in a one-shot sense is the major contribution of our work, and block LTC is an extension of our LTC framework, showing that one can approach rate-distortion limits in block coding.

W3. Regarding complexity, the total complexity for a neural compressor is determined by the transform, likelihood computation, quantization of latents, and entropy coding. On images, we use the same transforms and autoregressive entropy model as the Cheng2020 model, and it is well-known that the complexity of neural compressors is dominated by the autoregressive entropy model [1] and not quantization. For quantization, the complexity of solving CVP is not an issue. For each lattice, CVP can be computed at any rate in a fixed number of operations, and is significantly faster than the codebook search in VQ, which is exponential in the rate. We do not evaluate entropy coding in this paper as it is not the focus of our paper, and we leave this to future work; as mentioned in Sec. 5, the rates reported are the cross-entropies of the learned entropy model. However, we do not expect this to increase complexity significantly for the lattices used in this work. For example, for the $E_8$ lattice used for image compression (Fig. 7), the effective codebook size will be around $2^(8*\text{bpp})$, which is a codebook size commonly seen in reduced-complexity VQ methods such as NVTC [2], and there are also 8x less symbols to code compared to scalar quantization. Additionally, Sec. 4.2 discusses the use of nested LQ, which provides a low-complexity solution to encoding with lattices. Regarding how LTC is applied for image compression, we have added details in the manuscript. To summarize, we apply lattices across the channel dimension of the latent tensor of the Cheng2020 model, which is usually either 128 or 192, depending on rate. For the $D_n^*$ lattice, it is defined for arbitrary dimension, so we can directly apply it to the channel dimension. For the E8 and Leech lattices, they are only defined for 8 and 24 dimensions respectively, so we use product lattices. For the E8 lattice, this results in 16 or 24 E8 lattices, and for the Leech lattice, this results in 5 or 8 Leech lattices (for the former, we use a channel dimension of 120 instead of 128). We have updated the experimental details in the revised draft to incorporate these details.

References:

[1] Y. Yang, S. Mandt, and L. Theis. An introduction to neural data compression. arXiv preprint arXiv:2202.06533, 2022.

[2] Runsen Feng, Zongyu Guo, Weiping Li, and Zhibo Chen. Nvtc: Nonlinear vector transform coding. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6101–6110, 2023.

Thank you for taking the time to read the paper and provide valuable feedback! Please find our responses to your comments in the following point-by-point response, which refers to the revised manuscript.

Response to weaknesses:

W1. Regarding training, fine-tuning a pre-trained model (such as a lower rate one) can indeed be done for LTC. Architecturally, the parameters that are learned are the same (i.e., transforms and entropy model); the only difference is the choice of quantization, which is not learned. On inference: we acknowledge the reviewer’s comments that spatial auto-regression is time consuming. However, the use of lattice quantization is independent of the choice of entropy model. The Cheng2020 architecture is not our contribution; we merely use it to compare how scalar quantization and lattice quantization perform for that particular architecture. The LTC framework can indeed be used with other NTC architectures, such as those with more refined entropy models like those you mentioned. As mentioned in Sec. 5.1, and the new runtime tables in Appendix B (Table 2), we compare just the effect of lattice quantization for runtime efficiency during inference on Cheng2020, and show that using the $E_8$ lattice (with dimension 8) incurs a 1.23x increase in runtime over NTC, whereas the Leech lattice (with dimension 24) incurs a 2.68x increase. Thus, the use of lattice quantization is not significantly more computationally expensive than using scalar quantization with the spatial autoregression.

W2. Thanks for catching the typos and missed reference. We have updated the paper accordingly.

Thank you for taking the time to read the paper and provide valuable feedback! Please find our responses to your comments in the following point-by-point response.

Response to weaknesses:

W1. We believe the value of this observation should be taken in context of the broader question of whether NTC can recover VQ (i.e., optimal) performance. This is the original question our paper seeks to answer, and is both central and important to the neural compression and information theory communities. The answer to this question is rather nuanced, as NTC consists of many components (transforms, quantizer, and learned entropy model), not just the transforms, each of which could be a contributing or limiting factor to achieving optimality. It may also depend on the source to compress. As mentioned in the introduction, prior work [2, 3] has identified certain cases in which NTC can indeed recover VQ; namely, when the source lies on a 1-d manifold. There, no component of NTC restrains it from achieving optimality. Our work examines this question for general sources, when the source lies on higher-dimensional manifolds, as is the case of iid scalar sequences. In such cases, the fact that NTC cannot recover VQ is likely due to the highly discontinuous mapping between different tessellations the transforms would need to recover, as you mentioned. The value or surprise of this observation by itself is debatable and subjective; however, in the context of our original question, this observation is more significant, as it identifies the exact reason and component determining why NTC cannot recover VQ in the more general settings we consider. A priori, we do not believe it is clear (a) if all NTC components can work together to result in optimality, or (b) if there are one or more components preventing optimality, and if so, whether they can be identified. Furthermore, we use this finding to propose the LTC framework, which successfully achieves VQ regions as an end-to-end system.

W2. Thank you for bringing up this point. We agree that the discussion of $\lambda$ is sparse in the original manuscript, and have updated the paper accordingly around Eq. (1) and in the experimental setup. To clarify, $\lambda$ is set to a fixed value during training, which results in a learned compressor operating at a single R-D trade-off point. To get the whole trade-off for a model (all R-D plots are done this way), $\lambda$ is swept across a variety of values, with each resulting compressor’s R-D performance plotted as a point. This is the same way any variable-rate compressor is optimized, whether it is ECVQ, NTC, or LTC, and is the same way it is done in [1]. A larger $\lambda$ encourages the transforms to use more codebook vectors in the latent space (and therefore more rate). This is shown, for example, in Figs. 19-21. As an example, on iid scalar sequences, the $\lambda$ values are swept over 0.5, 1, 1.5, 2, 4, 8, resulting in the markers shown on the R-D plots.

W3. Thanks for bringing this point to our attention. In Appendix B, we have compared a LTC model with a trained NTC model whose scalar quantizer is replaced by a lattice quantizer at inference time. The performance of this scheme is largely dependent on the source at hand (as one might expect). Fig. 11 shows that on the Gaussian source, this scheme performs the same as training LTC from scratch. This is perhaps not unexpected, as the transforms learned by NTC and LTC for the i.i.d. Gaussian source effectively perform a scaling of the source, and so interchanging them has no effect. We use the same Flow entropy model for fair comparison. In general, however, this is not expected to be the case, as the analysis transform determines which region of the lattice gets used; for more complex sources, there may be more sophisticated ways the transform determines how the lattice is used. For example, in Fig. 12, we demonstrate the same scheme on Kodak images with the Cheng2020-LTC-$\Lambda_{24}$ model. Here, we see that while using a Leech lattice quantizer with the trained NTC model improves the performance over NTC by a BD-rate of -12.984%, training the model with the lattice quantizer helps boost the performance further to a BD-rate of -16.708%. This result implies that one can avoid training from scratch when a good set of transforms and entropy model are already trained from a NTC architecture, and instead potentially fine-tune the NTC transforms and entropy model with the chosen lattice for best performance. In the future, when better transforms and entropy models are developed, one can also train with the lattice from scratch to get the best performance, as the training times between NTC and LTC (from scratch) are on the same scale (weeks), so there is not a significant difference in training time.

Thank you for taking the time to read the paper and provide valuable feedback! Please find our responses to your comments in the following point-by-point response.

Response to weaknesses:

W1. We acknowledge that synthetic sources are a significant part of our experiments. We do work with real-world data, ranging from scientific and audio to images, with significant gains achieved for image compression. The primary goal of our work is to understand the fundamental limitations of NTC, and how to improve the performance to approach optimality. If optimality is achieved, then there is no more room for improvement. Such a fundamental analysis requires synthetic sources where optimality is known, and all causes for potential gains or suboptimality can be isolated. While real-world data may be more realistic, there are many factors that may contribute to differences in R-D performance, which precludes a fundamental analysis. Our set of experiments demonstrate the suboptimality of NTC (requiring synthetic data), ability of LTC to approach the optimal (requiring synthetic data), and the superiority of LTC over NTC (shown on both synthetic and real-world data). In addition, our experiments on real-world data do cover a broad range of data modalities, ranging from scientific and audio data to images, demonstrating the broad applicability of LTC to improve upon NTC in practical scenarios.

W2. Thank you for your suggestion. We have added a BD-rate table in Appendix B, with both VTM and Cheng2020-NTC as an anchor, and additionally revised the text to provide precise comparisons between methods. To summarize, (i) while NVTC achieves a -2.919 BD-rate gain over VTM, Cheng2020 with Leech is significantly better, with a -11.895% BD-rate gain, and (ii) the increasing BD-rate gains of Cheng2020-LTC with lattice dimension further supports evidence that the rate-distortion performance is significantly determined by the quantization choice in the latent space, and that the NTC model is able to leverage the enhanced packing efficiency of higher-dimensional lattices.

W3. This is now addressed by the more precise comparisons provided in the previous point. We agree that good gains achieved with Cheng2020 are an important conclusion. We hope that the BD-rate table and discussions help highlight these findings in addition to the fundamental analysis provided in this paper, such as understanding of where and when NTC methods are suboptimal, and why and how LTC resolves them. Those conclusions are also valuable because NTC/LTC can be applied to many modalities (not just images), so such an understanding impacts a broader range of research than just image compression.

Response to questions:

Q1. Thank you for your suggestion. We added some more details on the practical aspects of lattice quantizers in the revised manuscript, in the implementation details in the appendix. For the low-to-moderate dimension Euclidean sources (e.g., iid scalar sequences, Physics, and Speech), a single lattice of dimension up to 24 is used. Thus a single call to the lattice quantizer’s CVP algorithm is required to compute the quantized latent, as well as a CVP call for each Monte Carlo sample for the entropy model. For images, a product lattice is applied along the channel dimension of the latent tensor, which has shape C x H x W. There are (C/n) lattices applied H x W times, where n is the lattice dimension. We do not evaluate entropy coding in this paper as it is not the focus of our paper, and we leave this to future work; as mentioned in Sec. 5, the rates reported are the cross-entropies of the learned entropy model. However, we do not expect this to increase complexity significantly for the lattices used in this work. For example, for the E_8 lattice used for image compression (Fig. 7), the effective codebook size will be around 2^(8*bpp), which is a codebook size commonly seen in reduced-complexity VQ methods such as NVTC [1], and there are also 8x less symbols to code compared to scalar quantization. Additionally, Sec. 4.2 discusses the use of nested LQ, which provides a low-complexity solution to encoding with lattices.

References:

[1] Runsen Feng, Zongyu Guo, Weiping Li, and Zhibo Chen. Nvtc: Nonlinear vector transform coding. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6101–6110, 2023.

Response to questions:

Q1. When the source dimension does not equal that of the source, there are several options, depending on the model architecture and source dimensions. For Euclidean vector sources, such as Speech and Physics, we map the source to a latent dimension of $d_y$, and use a single lattice of dimension $d_y$. If $d_y$ is greater than 24, which we do not use in this paper for Speech and Physics, then a product lattice could be used. For images, since the source is now a tensor, the latent space is also a tensor of shape C x H x W, where C is the channel dimension, and H, W are the spatial dimensions. As mentioned in the paper, we apply LQ to the channel dimension. C ranges from 128 to 320, depending on the model architecture; we thus use product lattices for E8 and Leech. The $D_n^*$ lattice is defined for any dimension, so it does not require a product lattice. We have clarified this in the paper in the experimental details.

Q2. The goal of block LTC is to design the lowest complexity architecture that will capture gains from compressing iid blocks simultaneously. There are two options: (i) block LTC as it is currently defined, and (ii) a block LTC method that jointly maps the iid vector sequence to a larger latent space. For (i), since the source is a iid sequence of vectors, they can all share a common latent space, which only requires g_an and g_s to be applied sample-wise. Now we have a iid sequence of latent vectors $y_i$. If we apply LQ with each $y_i$ vector, this will be no different from one-shot coding, so that is why LQ is applied across the $n$-length blocks (see response to weakness 2). The other option (ii) would be to not use a sample-wise $g_a$ and $g_s$, and instead use transforms that jointly map the iid block to a much larger latent space of $n*d_y$. Then, the iid structure is likely lost and we could perhaps get more gains by applying LQ in other ways. However, such a $g_a$ and $g_s$ would not scale well with $d$ or $n$ as they are implemented with MLPs for the sources we use block LTC for (i.e., parameter count grows $O((dn)^2)$). With idea (i), the g_a transforms have $O(d^2)$ parameters, and the $c_a$, $c_s$ transforms have $O(n^2)$ parameters. Note that option (ii) is quite similar conceptually to how the image compression models work, except convolutional layers are used to have a more efficient parameter count, the image’s pixels (which can be thought of as a sequence of vectors) are not iid, and the lattice is indeed applied across vectors that should be correlated in the latent space. Future work could explore option (ii) to achieve further gains in block coding. We hope this discussion clarifies your question.

Q3. Generally speaking, different rate regimes may require a different number of latent dimensions to be used for NTC and LTC. If the number of latent dimensions used by the transforms does not match the lattice quantizer dimension, this can result in suboptimal performance as the tessellation in source space may not be optimal. This is demonstrated for the Speech and Physics results, where different latent space dimensions (with corresponding lattices) result in different rate regimes where LTC performance “peaks”. If one takes the convex hull of all the R-D points, it results in a curve that outperforms NTC uniformly over all rate regimes. Since the integer lattice of NTC can be seen as a product lattice of a 1-d lattice, NTC performance will not differ as long as the latent dimension is larger than the latent dimensions used by the transforms. For iid sequences and images, the transforms effectively use all of the latent dimensions at all rate regimes used, so the LTC models there do not need to align the lattice dimension with rate regime for best performance. We have updated the revised manuscript in Sec. 5.1 with these discussions.

References:

[1] Johannes Ballé, Philip A. Chou, David Minnen, Saurabh Singh, Nick Johnston, Eirikur Agustsson, Sung Jin Hwang, and George Toderici. Nonlinear transform coding. IEEE Journal of Selected Topics in Signal Processing, 15(2):339–353, 2021. doi: 10.1109/JSTSP.2020.3034501.

Thank you for taking the time to read the paper and provide valuable feedback! Please find our responses to your comments in the following point-by-point response.

Response to weaknesses:

W1. To operationalize the R-D curves, the interpretation would resemble that for NTC methods (e.g., see Fig. 10 in [1]). That is, one defines a discrete PMF over the quantized codebook vectors via integrating a parameterized continuous entropy model; for LTC, this is defined in Eq. 5. The continuous entropy model would be available at both the encoder and decoder to provide likelihoods for entropy encoding and decoding. To ensure reproducibility of the likelihood values in Eq. 5 at both the encoder and decoder, a shared random seed could be either predetermined or transmitted as part of the bitstream.

W2. Thanks for the feedback, and we have clarified the writing of nested LQ (Sec. 4.2) and block LTC (Sec 4.3) in the text. We hope the following comments help resolve your questions. Regarding nested LQ, the high-level idea is to use the coarse lattice to define a bounded region of the fine lattice. The codebook vectors of the fine lattice, contained in the said region, are the ones we wish to use to encode the source via index coding. This happens when the majority of the latent mass is placed in the coarse lattice cell centered at the origin. Regarding block LTC, the idea is for $g_a$ and $g_s$ (which are shared across the block) to map each source vector to a lower-dimensional latent space of $d_y$ dimensions, which is often sufficient; for example, the Sawbridge is mapped from 1024 dimensions to 1 or 2 dimensions. We now have a sequence of latent i.i.d. vectors. If we directly apply LQ across the i.i.d. slices, this would not result in an optimal compression, as it is equivalent to applying LQ directly to i.i.d. sequences, which is suboptimal (mentioned in the ablation study, Sections 5.3 and B.1) for any source that does not have a uniform density. The $c_a$, $c_s$ transforms help shape the lattice regions to match the statistics of the iid latent source; for the i.i.d. vector sources in this paper, a small MLP is sufficient, similar to the i.i.d. scalar sequences in Sec. 5.1.

Thanks for pointing out the typo. Line 458 should indeed say “d=16 and 33”. We have updated the paper accordingly.