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

## Abstract

Abstract Neural compression has brought tremendous progress in designing lossy compressors with good rate-distortion (RD) performance at low complexity. Thus far, neural compression design involves transforming the source to a latent vector, which is then rounded to integers and entropy coded. While this approach has been shown to be optimal in a one-shot sense on certain sources, we show that it is highly sub-optimal on i.i.d. sequences, and in fact always recovers scalar quantization of the original source sequence. We demonstrate that the sub-optimality is due to the choice of quantization scheme in the latent space, and not the transform design. By employing lattice quantization instead of scalar quantization in the latent space, we demonstrate that Lattice Transform Coding (LTC) is able to recover optimal vector quantization at various dimensions and approach the asymptotically-achievable rate-distortion function at reasonable complexity. On general vector sources, LTC improves upon standard neural compressors in one-shot coding performance. LTC also enables neural compressors that perform block coding on i.i.d. vector sources, which yields coding gain over optimal one-shot coding.

## 1 Introduction

Neural compression, especially nonlinear transform coding (NTC), has made significant progress in advancing practical data compression performance on real-world data at low complexity (Ballé et al., ; Yang et al., ). In this work, we seek to understand when and how neural compressors may be optimal in terms of the rate-distortion trade-off. Obtaining such an understanding can provide insights on if and how neural compressors can be further improved, as well as why they perform well.

To investigate this, it is important to understand the underlying principles behind NTC. Shown in Fig. 2, NTC operates on a source realization $\bm{x}\sim P_{\bm{x}}$ by first transforming it to a latent vector $\bm{y}$ via an analysis transform $g_{a}$. The latent vector is then scalar quantized to $\hat{\bm{y}}$ via $Q$, which is typically implemented via entry-wise rounding to integers. Then, a synthesis transform $g_{s}$ transforms $\hat{\bm{y}}$ into the reconstruction $\hat{\bm{x}}$. Additionally, an entropy model $p_{\hat{\bm{y}}}$ estimates likelihoods of the quantized latent vectors $\hat{\bm{y}}$, providing a coding rate estimate via its entropy. Together, $g_{a}$, $g_{s}$, and $p_{\hat{\bm{y}}}$ are learned to minimize a rate-distortion loss,

$$ $\min_{g_{a},g_{s},p_{\hat{\bm{y}}}}\operatorname{\mathbb{E}}_{\bm{x}}\left[- \log p_{\hat{\bm{y}}}(\hat{\bm{y}})\right]+\lambda\cdot\operatorname{\mathbb{E }}_{\bm{x}}[\operatorname{\mathsf{d}}(\bm{x},\hat{\bm{x}})],$ (1) $$

where $\bm{y}=g_{a}(\bm{x})$, $\hat{\bm{y}}=Q(\bm{y})$, and $\hat{\bm{x}}=g_{s}(\hat{\bm{y}})$. NTC is considered a one-shot coding scheme of $\bm{x}$, because it compresses one source realization at a time.

Figure: Figure 1: Nonlinear Transform Coding (NTC) recovers scalar quantization (ECSQ) on Gaussian i.i.d. sequences of any dimension $n$. Lattice Transform Coding (LTC) recovers optimal vector quantization (VQ) in each dimension, and approachces the rate-distortion function $R(D)$ without needing an exponential codebook search.
Refer to caption: x1.png

A recent line of work has investigated how close neural compressors perform in comparison to information-theoretic limits, which define the optimal performance any compressor may achieve. Yang and Mandt () and Lei et al. () estimate the rate-distortion function $R(D)$, which is the fundamental limit of lossy compression, for real-world data sources. Their results show that there exists a gap between NTC performance and $R(D)$. In general, the gap between NTC performance (or any one-shot code) and $R(D)$ can be attributed to (i) sub-optimality of NTC with regards to the optimal one-shot performance, and/or (ii) the lack of block coding (multi-shot coding), since $R(D)$ is achievable only under asymptotic block-lengths (Cover and Thomas, ).

The authors in (Wagner and Ballé, ; Bhadane et al., ) demonstrate that on stylized sources, NTC achieves the optimal one-shot coding performance. These sources are nominally high-dimensional, but have a low-dimensional representation, which we refer to as the latent source. The authors show that $g_{a},g_{s}$ can successfully extract the latent source, optimally quantize it, and map back to the original space, providing an optimal one-shot scheme as a whole.

However, most of the sources they analyze have a one-dimensional uniform latent source $U\sim\mathrm{Unif}([a,b])$. While these are sufficient to analyze the role of $g_{a}$ and $g_{s}$, it does not provide insights on the role of the quantizer $Q$. This is because a uniform source $U$ can be optimally quantized using uniform scalar quantization, which is exactly what NTC uses in the latent space. For sources with higher-dimensional latent sources, where uniform scalar quantization is not necessarily optimal, it is unclear whether NTC can still provide optimal one-shot coding schemes.

An information-theoretic framework relevant to this question is that of companding (Bucklew, ). Companding defines a function $f:\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}$, known as the compressor, which is applied to the source $\bm{x}\in\mathbb{R}^{n}$. $f(\bm{x})$ is then uniform lattice quantized, and passed to $f^{-1}$, known as the expander. Clearly, NTC can be thought of as an instantiation of a compander where $g_{a},g_{s}$ serve the role of $f,f^{-1}$. It is known that optimal companding is governed by the lattice’s packing efficiency and not $f$ (Linder et al., ). Uniform *scalar* quantization is equivalent to the integer lattice, which is sub-optimal in dimensions greater than 1.

This suggests that the results of one-shot optimality of NTC (Wagner and Ballé, ) may not hold in general when the latent source has dimensionality greater than 1. Surprisingly, this implies that when $\bm{x}$ contains an i.i.d. scalar sequence on its entries, NTC may not be able to compress it optimally. In this case, the latent source’s dimensionality is that of $\bm{x}$. The same implication holds for general vector sources, such as those from real-world data, where the latent source may reduce dimensionality but still have dimension greater than 1. In these cases, companding suggests the use of lattice quantization (LQ) in order to improve the one-shot coding performance of $\bm{x}$. Additionally, LQ provides a potential method to perform block coding of i.i.d. sequences of $\bm{x}$ in a scalable way, in order to bridge the gap between optimal one-shot coding and the rate-distortion function.

Figure: Figure 2: Nonlinear Transform Coding (Ballé et al., ). Lattice Transform Coding replaces scalar quantizer $Q$ with a lattice quantizer $Q_{\Lambda}$.
Refer to caption: x2.png

In this paper, we first provide evidence that NTC compresses such sources sub-optimally in Sec. 3. These results agree with those in the companding literature, and motivate the use of lattice quantization in the latent space as a potential way to improve performance, which we refer to as Lattice Transform Coding (LTC). LTC brings several challenges that require changes compared to NTC training, such as efficient lattice quantization algorithms, dealing with the non-differentiable quantization operation, integration of the entropy model over lattice regions, and accurate likelihood modelling. LTC is able to achieve optimal coding schemes for i.i.d. scalar sequences, while avoiding the exponential complexity of a direct codebook search under vector quantization (VQ); see Fig. 1. On general vector sources, LTC improves one-shot coding performance, as well as enable block coding of i.i.d. vector sequences. Our contributions are as follows.

- 1)
We first demonstrate the inability of current NTC architectures to optimally compress i.i.d. scalar sequences. We show that this is due to the choice of scalar quantization in the latent space, and not the transform design.
- 2)
We propose Lattice Transform Coding (LTC), and show that it is able to optimally compress i.i.d. sequences, matching the performance of optimal vector quantization, yet maintain reasonable complexity. We discuss various design choices regarding quantization, and entropy modelling that help achieve optimal performance. We demonstrate the coding gains achieved are due to the packing efficiency of the lattice in the latent space.
- 3)
On general vector sources, we demonstrate LTC’s ability to improve one-shot coding schemes, as well as develop block coding schemes that compress i.i.d. blocks of vector sources, which are able to approach the rate-distortion bound.

## 2 Related Work

#### Rate-Distortion Theory.

Given a source $X\sim P_{X}$, the rate-distortion function (Cover and Thomas, )

$$ $R(D):=\min_{P_{\hat{X}|X}:\operatorname{\mathbb{E}}[\operatorname{\mathsf{d}}( X,\hat{X})]\leq D}I(X;\hat{X})$ (2) $$

describes the fundamental limit of compressing $X$. It is achievable in the limit of large block-lengths, where the rate and distortion are measured per-sample.

For a general vector source $\bm{x}$, its optimal one-shot coding scheme can be found via entropy-constrained vector quantization (ECVQ). ECVQ will approach $R(D)$ when $\bm{x}$ consists of i.i.d. sequences of $X$ and rate and distortion are divided by the sequence length. ECVQ directly enumerates quantization centers (codewords) in the space of $\bm{x}$, and minimizes an entropy-distortion objective,

$$ $\min_{\{\bm{c}_{i}\}_{i=1}^{\infty}}\operatorname{\mathbb{E}}_{\bm{x}}[-\log P (e(\bm{x}))]+\lambda\cdot\operatorname{\mathbb{E}}_{\bm{x}}[\operatorname{ \mathsf{d}}(\bm{x},e(\bm{x}))],$ (3) $$

where $e(\bm{x}):=\operatorname*{arg\,min}_{\{\bm{c}_{i}\}_{i=1}^{\infty}}
\operatorname{\mathsf{d}}(\bm{x},\bm{c}_{i})$. Since the solution is unknown in closed form, ECVQ is solved numerically by a generalized Lloyd’s algorithm (Chou et al., ).

Although ECVQ is optimal, an issue that precludes its wider use is its high complexity; the number of codebook vectors required is exponential in the rate. This complexity issue is one reason why NTC is appealing, since scalar quantization in the transformed space is used, whose complexity is linear in the dimension of $\bm{y}$.

#### Rate-Distortion Estimation.

The work of Yang and Mandt () and Lei et al. () demonstrate a gap between NTC performance and estimates of $R(D)$ on real-world sources, and potentially attributes it to the lack of large block-length. We note that a comparison to $R(D)$ depends on the source and its relevant regimes on the rate-distortion trade-off. If an entire image is considered as a single sample, then typical regimes of interest include rates on the order of $10^{5}$ to $10^{6}$ bits-per-sample(^1^11The bits-per-pixel (bpp) usually ranges between 0 and 2.). From Li and Gamal (), the gap between optimal one-shot coding and $R(D)$ is no more than $\log(R(D)+1)+5$. Thus, the gap becomes insignificant relative to the coding rate at such large rates; to compress a $256\times 256$ image patch at a bpp of 1, the gap corresponds to $\approx 0.00032$ bpp, which is much smaller than that reported in Yang and Mandt (). Therefore, for sources such as images, the performance gap can mostly be attributed to the sub-optimality of NTC as a one-shot coding scheme. For sources that have lower rate regimes of interest, the gap to $R(D)$ may be attributed to lack of large block-length as well.

#### Analysis of NTC.

The principles of nonlinear transform coding (NTC) are described in Ballé et al. (). An analysis of NTC was performed in Wagner and Ballé (); Bhadane et al. (), who consider synthetic sources and verify when NTC can recover the optimal (one-shot) rate-distortion performance. The focus is on determining when the analysis and synthesis transforms are able to recover the low-dimensional latent source; when this is possible, optimality is achieved empirically. However, the sources analyzed have a one-dimensional uniform latent source which uniform scalar quantization is optimal for; it is unknown whether sources with more complex latent sources can still be compressed optimally by NTC.

#### Vector Quantization in NTC.

Although NTC is based on VAEs with scalar quantization in the latent space, several works have attempted to apply vector quantization in the latent space instead. Notable works include VQVAE (van den Oord et al., ), which performs a direct codebook search in the latent space and can be used as a compressor. PQVAE (El-Nouby et al., ) evaluates this approach with a product quantization scheme to improve scalability, demonstrating competitive performance on image compression benchmarks. Finally, NVTC (Feng et al., ) applies entropy-constrained vector quantization (ECVQ) algorithms (Chou et al., ) in the latent space, performing a direct codebook search. We show that LTC is able to recover the performance of ECVQ while avoiding a direct codebook search in the latent space, which is infeasible for large rates and dimensions.

#### Lattice Quantization.

A lattice quantizer $Q_{\Lambda}$ (Conway and Sloane, ) consists of a countably infinite set of codebook vectors $\Lambda$ given by $\bm{u}\bm{G}$, where $\bm{G}\in\mathbb{R}^{n\times n}$ is the generator matrix of the lattice, and $\bm{u}$ are integer-valued row vectors. The Voronoi region of a lattice is given by $V(\bm{0}):=\{\bm{x}\in\mathbb{R}^{n}:Q_{\Lambda}(\bm{x})=\bm{0}\}$, where $Q_{\Lambda}(\bm{x}):=\operatorname*{arg\,min}_{\bm{v}\in\Lambda}\|\bm{x}-\bm{v
}\|_{2}^{2}$ finds the closest lattice vector to $\bm{x}$ in $\Lambda$. In lattice theory, the objective is to find lattices that minimize the normalized second moment (NSM)

$$ $L_{\Lambda}:=\frac{1}{n\mathrm{Vol}(V(\bm{0}))^{1+\frac{2}{n}}}\int_{V(\bm{0}) }\|\bm{x}\|_{2}^{2}d\bm{x},$ (4) $$

which is a scale-invariant measure of the lattice’s packing efficiency. Recent work describing lattices which have the lowest known NSM can be found in (Agrell and Allen, ; Lyu et al., ; Agrell et al., ).

Another challenge is the quantization (also referred to as decoding) operation itself $Q_{\Lambda}(\bm{x})$, which requires solving the closest-vector problem (CVP); this is NP-hard in general. For many of the best-known lattices in dimensions up to 24 ($n=2$: hexagonal; $n=4$: $D_{n}^{*}$,; $n=8$: Gosset $E_{8}$; $n=16$: $\Lambda_{16}$ Barnes-Wall; and $n=24$: $\Lambda_{24}$ Leech) efficient algorithms to solve CVP are known (Conway and Sloane, ). For general lattices, CVP can be solved using the Schnorr-Euchner (SE) decoder (Agrell et al., ), but solving exact CVP using SE is much slower compared to custom lattice decoders, and is not as easily parallelizable across batches for GPU implementation.

In this paper, we use the CVP algorithms outlined in Conway and Sloane (, Ch. 20). For the Barnes-Wall lattice $\Lambda_{16}$ in dimension 16, we use the algorithm outlined in Conway and Sloane (). Finally, the Leech lattice $\Lambda_{24}$ in dimension 24 has the slower method described in Conway and Sloane () and the improved version in Conway and Sloane ().

#### Companding.

Companding has typically been analyzed in the (asymptotically) large rate, small distortion regime. Optimal scalar quantization can be implemented using scalar companding (Bennett, ). For the vector case, vector companders cannot implement optimal VQ (Gersho, ; Bucklew, ; ), except for a very small class of sources. Despite this, it is still possible optimal companding may perform near-optimal VQ.

This is shown in Linder et al. (), who show that the optimal compander can arbitrarily approach the rate-distortion function in the limit of large dimensions and rates:

$$ $\displaystyle\lim_{D\rightarrow 0}\frac{1}{n}\left(H(Q_{D,f})-R_{\bm{x}}(D)\right)$ $\displaystyle=\frac{1}{2}\log(2\pi eL_{\Lambda}),$ (5) $$

where $R_{\bm{x}}(D)$ is the rate-distortion function of $\bm{x}\sim P_{\bm{x}}$, $H(Q_{D,f})$ is the entropy of the optimal compander at distortion $D$, and $L_{\Lambda}$ is the NSM for a lattice $\Lambda$ in $n$ dimensions. As the authors note, if $G_{n}$ is the minimal NSM of any lattice in $n$ dimensions, $\lim_{n\rightarrow\infty}G_{n}=\frac{1}{2\pi e}$, such that (5) converges to 0 at a rate $O\left(\frac{\log n}{n}\right)$. This result indicates that the gap to $R_{\bm{x}}(D)$ is determined by the packing efficiency of the lattice and not $f$; for the integer lattice, the NSM is $1/12$, which yields a gap of $\approx 0.255$ bits per dimension. This indicates that it is impossible for NTC (which uses the integer lattice) to be optimal with respect to the rate-distortion function.

The companding literature does not address non-asymptotic rate regimes, which NTC models typically operate at. In this paper, we provide experimental evidence that the gains of vector companding still hold for non-asymptotic rates.

## 3 NTC is Sub-Optimal on i.i.d. Sequences

Let $P_{S}\in\mathcal{P}(\mathbb{R})$ be a one-dimensional distribution, and suppose that the entries of $\bm{x}$ are i.i.d. samples $S_{1},\dots,S_{n}\stackrel{{\scriptstyle\textrm{i.i.d.}}}{{\sim}}P_{S}$. We first choose $P_{S}=\mathcal{N}(0,1)$, so that $\bm{x}\sim\mathcal{N}(\bm{0},\bm{I}_{n})$. This is because i.i.d. Gaussians are the most difficult source to compress, and have no low-dimensional latent structure for the transforms. We train an NTC model following Ballé et al. () using a factorized entropy model for the distribution of the latent $\bm{y}$. We use the objective in (1), with $\operatorname{\mathsf{d}}(\bm{x},\hat{\bm{x}})=\|\bm{x}-\hat{\bm{x}}\|_{2}^{2}$, and normalize the entropy and distortion terms by $n$ to compare with the rate-distortion function of $S$, as well as entropy-constrained vector quantization (ECVQ) which is the optimal quantizer for any fixed $n$.

Shown in Fig. 1, NTC with $n=1$ is able to recover ECVQ with $n=1$, also known as entropy-constrained scalar quantization (ECSQ). This is perhaps expected, since the task of the transforms is to scale the 1-D Gaussian to be optimally quantized on an integer grid in 1-D, and then scale the quantized latent to form the reconstruction. As expected, a gap between this (optimal) one-shot performance of coding $S$ and $R(D)$ of $S$ exists.

For $n=2$, however, NTC is unable to recover ECVQ, and always achieves the same performance as that of ECSQ. To see this more clearly, we define the normalized rate-redundancy of a compressor $\mathsf{c}$ as

$$ $\mathcal{E}_{\mathsf{c}}(D):=\frac{R_{\mathsf{c}}(D)-R(D)}{R(D)},$ (6) $$

where $R_{\mathsf{c}}(D)$ is the rate achieved by $\mathsf{c}$ at distortion $D$, and $R(D)$ is the rate-distortion function of $S$. Intuitively, $\mathcal{E}_{\mathsf{c}}(D)$ measures the fraction of additional bits per sample required to code a sequence relative to $R(D)$.

Figure: Figure 3: $\mathcal{E}_{\mathsf{c}}(D)$ of NTC for $n=1,2$.
Refer to caption: x3.png

Figure: (a) NTC.
Refer to caption: extracted/5464271/figures/quantizers/square_original.png

Shown in Fig. 3, we see that the $\mathcal{E}_{\mathsf{c}}(D)$ for ECVQ at $n=2$ is lower than that of ECSQ. However, NTC at $n=2$ is unable to outperform $\mathcal{E}_{\mathsf{c}}(D)$ of ECSQ, which is the lower bound of $\mathcal{E}_{\mathsf{c}}(D)$ at $n=1$. Visualizing the quantization regions and centers on the 2-D Gaussian distribution (Fig. 3(a)) brings more clarity as to why this is the case. It is clear that NTC is merely implementing a scalar quantization of the two dimensional input, perhaps with a rotation, which is equivalent to applying ECSQ to the entries of $\bm{x}$ independently. If we compare to the optimal ECVQ in $n=2$, shown in Fig. 3(b), we see that the optimal quantization regions are hexagonal, not square.

In our experiments, we varied the depth, width, nonlinearities, biases, and even increased the latent dimension to a higher-dimensional space, all to no avail. Setting $n>2$ did not help either. All models recovered the performance of ECSQ at best.
Thus, it appears that $g_{a}$ and $g_{s}$ parameterized as neural networks are unable to generate optimal quantization regions in the original space when scalar quantization is used in the latent space. This experimentally verifies results in companding (Linder et al., ).

## 4 Lattice Transform Coding

We seek a compressor that can provide better performance than scalar quantization yet avoid the high complexity of vector quantization. A lattice quantizer $Q_{\Lambda}$ provides an intermediate solution that manages both complexity and space-filling efficiency.
We first discuss how to optimize the rate-distortion loss under lattice quantization in the latent space, which brings several challenges. Then we the role of the entropy model in LTC. Lastly, we discuss Block-LTC, which uses LTC to implement a block coding scheme for i.i.d. vector sources.

### 4.1 Rate-Distortion Loss Under Lattice Quantization

Typically, the objective in (1) is replaced with a proxy loss since quantization operations are not differentiable. We investigate both the straight-through estimator (STE) and dithered quantization with uniform noisy proxy. The STE approach can use (1) directly as an objective during training, but may lead to training instabilities. The uniform noise proxy has stabler training, but requires a rate-distortion proxy loss of (1) using a connection to dithered quantization.

We assume that the lattice quantizer $Q_{\Lambda}$ under consideration has an efficient algorithm to find the closest vector; see the appendix for lattice details. We also assume that the lattice’s generator matrix $\bm{G}$ has been normalized to be unit volume, i.e. $\mathrm{Vol}(V(\bm{0}))=\sqrt{\det(\bm{G}\bm{G}^{\top})}=1$. If a non-unit volume generator matrix $\bm{G}^{\prime}$ is given, one can simply make it unit volume via $\bm{G}=(\det(\bm{G}^{\prime}\bm{G}^{\prime\top}))^{-\frac{1}{2k}}\bm{G}^{\prime}$, where $k$ is the number of rows of $\bm{G}^{\prime}$.

#### Distortion term.

We use straight-through estimation (STE) for the distortion term $\operatorname{\mathbb{E}}_{\bm{x}}[\operatorname{\mathsf{d}}(\bm{x},g_{s}(Q_{
\Lambda}(g_{a}(\bm{x}))))]$. Specifically, we compute the quantized output $\hat{\bm{y}}=Q_{\Lambda}(g_{a}(\bm{x}))$ during the forward pass via

$$ $\hat{\bm{y}}=\bm{y}+\mathtt{stop_gradient}(Q_{\Lambda}(\bm{y})-\bm{y}).$ (7) $$

#### Rate term.

For the rate term, we typically have some parameterized density model $p_{\bm{y}}$ for the continuous latent $\bm{y}$ (more details on this in the next section). This induces a PMF over the quantized latents $p_{\hat{\bm{y}}}$, given by

$$ $p_{\hat{\bm{y}}}(\hat{\bm{y}})=\int_{V(\hat{\bm{y}})}p_{\bm{y}}(\bm{y})d\bm{y},$ (8) $$

where $V(\hat{\bm{y}}):=\{\bm{y}:Q_{\Lambda}(\bm{y})=\hat{\bm{y}}\}$. In NTC, where $Q_{\Lambda}$ implements the integer lattice, (8) is easy to integrate when $p_{\bm{y}}$ is assumed factorized (Ballé et al., ). In that case, (8) can be written in terms of the CDF of $p_{\bm{y}}$, which is one reason why the factorized $p_{\bm{y}}$ is typically parameterized by its CDF in NTC. However, for lattice quantization, solving the integral in (8) is not as straightforward.

As a result, we compute (8) via Monte-Carlo estimate,

$$ $\displaystyle p_{\hat{\bm{y}}}(\hat{\bm{y}})$ $\displaystyle=\int_{V(\hat{\bm{y}})}p_{\bm{y}}(\bm{y})d\bm{y}=\int_{V(\hat{\bm {y}})}\frac{p_{\bm{y}}(\bm{y})}{p_{\bm{u}}(\bm{y})}p_{\bm{u}}(\bm{y})d\bm{y}$ (9) $\displaystyle=\mathrm{Vol}(V(\hat{\bm{y}}))\cdot\mathbb{E}_{\bm{u}\sim\mathrm{ Unif}(V(\hat{\bm{y}}))}[p_{\bm{y}}(\bm{u})]$ (10) $\displaystyle=\mathbb{E}_{\bm{u}^{\prime}\sim\mathrm{Unif}(V(\bm{0}))}[p_{\bm{ y}}(\hat{\bm{y}}+\bm{u}^{\prime})],$ (11) $$

where $p_{\bm{u}}$ is the uniform density over $V(\hat{\bm{y}})$, $\mathrm{Vol}(V(\hat{\bm{y}}))=1$ since we assumed unit-volume lattices, and the last equality holds because for a lattice quantizer $Q_{\Lambda}(\bm{x}+\bm{v})=Q_{\Lambda}(\bm{x})+\bm{v}$, $\forall\bm{v}\in\Lambda$. From Conway and Sloane (), we can sample $\bm{u}^{\prime}\sim\mathrm{Unif}(V(\bm{0}))$ by computing $\bm{u}^{\prime}=\tilde{\bm{u}}-Q_{\Lambda}(\tilde{\bm{u}})$, where $\tilde{\bm{u}}=\bm{s}\bm{G}$, and $\bm{s}\sim\mathrm{Unif}([0,1])$. Thus, (11) can be estimated via a sample mean. With the ability to reliably estimate $p_{\hat{\bm{y}}}(\hat{\bm{y}})$, we can also use STE for backpropagation of the rate term $\operatorname{\mathbb{E}}_{\bm{x}}[-\log p_{\hat{\bm{y}}}(\hat{\bm{y}})]$ by using (7).

#### Dithering.

While STE could be used in both cases, prior works have achieved benefits from using the dithering approach (Ballé et al., ). In this case, a random dither $\bm{u}^{\prime}\sim\mathrm{Unif}(V(0))$ is applied prior to quantization, with the resulting rate $\operatorname{\mathbb{E}}_{\bm{x},\bm{u}^{\prime}}[-\log p_{\hat{\bm{y}}}(Q_{
\Lambda}(\bm{y}-\bm{u}^{\prime}))]$. This rate is equivalent in the integer lattice to $\operatorname{\mathbb{E}}_{\bm{x},\bm{u}^{\prime}}[-\log p_{\bm{y}+\bm{u}^{
\prime}}(\bm{y}+\bm{u}^{\prime})]$ (Ballé et al., ). The density $p_{\bm{y}+\bm{u}^{\prime}}$ is a convolved density with likelihoods given by

$$ $\displaystyle p_{\bm{y}+\bm{u}^{\prime}}(\bm{y}^{\prime})$ $\displaystyle=\int_{V(0)}p_{\bm{y}}(\bm{y}^{\prime}-\bm{u}^{\prime})p_{\bm{u}^ {\prime}}(\bm{u}^{\prime})d\bm{u}$ (12) $\displaystyle=\operatorname{\mathbb{E}}_{\bm{u}^{\prime}\sim\mathrm{Unif}(V(0) )}[p_{\bm{y}}(\bm{y}^{\prime}-\bm{u}^{\prime})],$ (13) $$

which can be estimated via Monte-Carlo.

To apply this for our lattice quantization case during training, rather than provide the STE of the quantized output $\hat{\bm{y}}$ to the Monte-Carlo likelihood estimate, we provide the noise-added continuous latent $\bm{y}+\bm{u}^{\prime}$, where $\bm{u}^{\prime}$ is uniformly distributed over the Voronoi region. In our experiments, we observed that STE was less stable than the dithering approach, but had better performance at low rates, as will be shown in Sec. 5.

### 4.2 Entropy Models

Previously, the rate (under non-dithered quantization) has been referred to as $\operatorname{\mathbb{E}}_{\bm{x}}[-\log p_{\hat{\bm{y}}}(\hat{\bm{y}})]$, which is actually the cross-entropy between the PMF $P(\hat{\bm{y}}):=\Pr_{\bm{x}}(Q_{\Lambda}(g_{a}(\bm{x})=\hat{\bm{y}})$ on the centers $\hat{\bm{y}}$ induced by $\bm{x}$ and $p_{\hat{\bm{y}}}(\hat{\bm{y}})$, which is some PMF that is learned. Thus, what we call the rate is actually an upper bound on the true entropy of $P(\hat{\bm{y}})$, i.e.

$$ $\displaystyle\operatorname{\mathbb{E}}_{\bm{x}}[-\log P(\hat{\bm{y}})]$ $\displaystyle\leq\operatorname{\mathbb{E}}_{\bm{x}}[-\log P(\hat{\bm{y}})]+D_{ \mathsf{KL}}(P(\hat{\bm{y}})||p_{\hat{\bm{y}}}(\hat{\bm{y}}))$ $\displaystyle=\operatorname{\mathbb{E}}_{\bm{x}}[-\log p_{\hat{\bm{y}}}(\hat{ \bm{y}})],$ (14) $$

with equality if the learned PMF matches the true PMF exactly, $p_{\hat{\bm{y}}}(\hat{\bm{y}})=P(\hat{\bm{y}})$. Although the authors in Wagner and Ballé () do not discuss the effect of the entropy model, their results on the Sawbridge imply that the learned PMF indeed can match the true (optimal) PMF exactly, resulting in no sub-optimality due to the choice of entropy model. However, this cannot be always expected, and in the case of lattice quantization, we will show that the choice of entropy model does play a role in whether optimality can be achieved.

#### Complexity of Entropy Estimation.

One may ask why the true entropy is not directly used. Similar to how exponential complexity is what prevents one from directly using ECVQ on general sources, it is also what prevents direct use of the true entropy $H(\hat{\bm{y}})=\operatorname{\mathbb{E}}_{\bm{x}}[-\log P(\hat{\bm{y}})]=
\operatorname{\mathbb{E}}_{\hat{\bm{y}}}[-\log P(\hat{\bm{y}})]$ as a rate estimate. As mentioned previously, if operating at rate of $R$ bits per dimension, at dimension $n$, there will be on the order of $2^{nR}$ centers in the latent space. A naive plug-in estimate of the entropy will hence require at least $2^{nR}$ samples. Even the best known entropy estimators (Wu and Yang, ) will require $O\left(\frac{k}{\log k}\right)$ samples for an alphabet of size $k$, which is not enough of an improvement for our purposes since $k$ is on the order of $2^{nR}$. Interestingly, cross-entropy upper bounds have no such complexity issues and provide accurate estimates at large rates without requiring exponential samples (McAllester and Stratos, ). Hence, finding good entropy models that are able to match the true entropy closely is highly important.

#### Modeling General Densities.

In NTC, $p_{\bm{y}}$ is assumed factorized across its components, i.e., $p_{\bm{y}}(\bm{y})=\prod_{i=1}^{n}p_{\bm{y}_{i}}(\bm{y}_{i})$, following Ballé et al. (). This choice reduces modeling the density to modeling a several scalar densities, parameterized by the CDF, which is easy to integrate over the square lattice. Since we are now using a more general Monte-Carlo-based integration scheme for providing the likelihoods of $\hat{\bm{y}}$, we are no longer tied to having to parameterize $p_{\bm{y}_{i}}$ by its CDF. Instead, we propose to parameterize the PDF $p_{\bm{y}}$ in a fully multivariate fashion using a normalizing flow (Kobyzev et al., ). In particular, let $p_{\bm{z}}$ be $\mathcal{N}(\bm{0},\bm{I})$. Then, we model $\bm{y}=\phi(\bm{z})$ as the output of some invertible neural network $\phi$. The density of $\bm{y}$ is then

$$ $p_{\bm{y}}(\bm{y})=p_{\bm{z}}(\phi^{-1}(\bm{y}))|\det\mathbf{J}\phi^{-1}(\bm{y })|.$ (15) $$

Hence, we use (15) as our flow-based entropy model. Since (15) estimates the likelihoods directly, it is more efficient to perform the Monte-Carlo integration via (15) than the factorized model, since the factorized model requires an additional differentiation of the CDF to provide likelihoods.

### 4.3 Vector i.i.d. Sequences

LTC can be directly applied to any vector source $\bm{x}\in\mathbb{R}^{d}$. In this case, LTC operates as a one-shot code for compressing $\bm{x}$. When $\bm{x}$ consist of i.i.d. elements $S\sim P_{S}$ on its entries, the optimal one-shot compression of $\bm{x}$ approaches the rate-distortion limit of $S$. As mentioned in the introduction, when $\bm{x}$ is a general (correlated) vector source with a low-to-moderate rate regime of interest, we can also expect coding gains when multiple i.i.d. copies $\bm{x}_{1},\dots,\bm{x}_{n}\sim p_{\bm{x}}$ are compressed together, whose optimal performance should approach $R_{\bm{x}}(D)$ when $n$ grows large. Note that we we retain $n$ as the block-length parameter, but use $d$ as the dimension of each realization of $\bm{x}$.

To realize such a compressor, we propose Block-LTC (BLTC), which applies LTC across the i.i.d. latent variables of $\bm{x}$, shown in Fig. 5.

Figure: Figure 5: Block-LTC (BLTC) applies LTC in the latent space for vector i.i.d. sequences.
Refer to caption: x4.png

Specifically, the analysis transform $g_{a}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d_{y}}$ maps each vector of the block to its latent variable, $\bm{y}_{i}=g_{a}(\bm{x}_{i})$, which has $d_{y}$ dimensions. This produces a latent i.i.d. vector sequence $\bm{y}_{1},\dots,\bm{y}_{n}$. Then, $d_{y}$ separate LTC models $\{c_{a}^{(i)},c_{s}^{(i)},p_{\hat{\bm{y}}}^{(i)}\}_{i=1}^{d_{y}}$ are applied to across the block, where $c_{a}^{(i)},c_{s}^{(i)}$ are the analysis/synthesis transforms of LTC. The $i$-th LTC model is applied to the latent slice $[\bm{y}_{1,i},\dots,\bm{y}_{n,i}]^{\top}$, which has i.i.d. scalar entries. The resulting reconstructed latents $\hat{\bm{y}}_{i}$’s are passed to the synthesis transform $g_{s}:\mathbb{R}^{d_{y}}\rightarrow\mathbb{R}^{d}$, such that $\hat{\bm{x}}_{i}=g_{s}(\hat{\bm{y}}_{i})$.

## 5 Results

We analyze the three settings of interest: (i) i.i.d. scalar sequences, (ii) general vector sources, and (iii) i.i.d. vector sequences. In addition, we provide an ablation study demonstrating the effects of each component in LTC (lattice choice, quantization training scheme, density model, transform choice).

### 5.1 Experimental Setup

The objective function used is the rate-distortion lagrangian (1), with STE used for the distortion term. For the rate term, we found varying results for dithered quantization (13) and STE (11) depending on the source; these are specified for each experiment individually. Additional implementation details can be found in Appendix A.

#### Lattices.

We choose several commonly used lattices that are known to be the best in their dimension (Agrell and Allen, ): hexagonal, $D^{*}_{n}$, $E_{8}$ and $\Lambda_{24}$ (Leech) lattices. These lattices all have efficient algorithms to solve the closest vector problem (CVP), which computes $Q_{\Lambda}(\bm{y})$. The algorithms used are referenced in Sec. 2.

#### Transforms.

For the scalar i.i.d. sequences and general vector sources, we choose the transforms $g_{a}$ and $g_{s}$ to be MLPs with two hidden layers, a hidden dimension of 100, and softplus nonlinearities after each layer (except the last). For the vector i.i.d. sources, where the lattice transform coding is applied in the latent space, we apply additional companding MLPs of the same architecture for the $c_{a}^{(i)}$ and $c_{s}^{(i)}$ transforms.

#### Entropy Models.

We use the RealNVP (Dinh et al., ) normalizing flow for the density model $p_{\bm{y}}$, with 5 coupling layers and isotropic Gaussian prior. The rates plotted (unless mentioned otherwise) are given by the cross-entropy upper bound (14) of the density model. We use 4096 samples for the Monte-Carlo estimate of the PMF in (11) for all lattices. This was sufficient for most sources and lattices; see Appendix. B for a discussion on this topic.

### 5.2 Compression of Scalar i.i.d. Sequences

We consider i.i.d. sequences of several sources, including Gaussian with $L_{2}$ distortion, Laplacian with $L_{1}$ distortion, and a 1-dimensional marginal of the Banana source (Ballé et al., ) with $L_{2}$ distortion. The 1-d Banana marginal is used to test a skewed distribution, since both Gaussian and Laplacian sources are symmetric. In this setting, $n$ is the dimension of both $\bm{x}$ and the latent space $\bm{y}$, and also represents the block-length with regards to the rate-distortion function. Here, one-shot coding of $\bm{x}$ is equivalent to block coding of the scalar marginal source. We first discuss the 2-dimensional Gaussian case, which provides visualization of what LTC achieves before expanding to more general dimensions. We found that using STE for the rate term performed better in this setting (see Sec. 5.3).

#### 2-dimensions.

Figure: (a) Original space.
Refer to caption: extracted/5464271/figures/quantizers/hex_original.png

In 2-dimensions, we choose $Q_{\Lambda}$ to be the hexagonal lattice, which is known to be optimal (Conway and Sloane, ), and use the NF density model. As shown in Figs. 5(a), LTC produces quantization regions in the original source space that closely match that of optimal ECVQ in Fig. 3(b). The latent space (Fig. 5(b)) effectively uses the hexagonal lattice, compared to NTC’s use of the integer lattice (Fig. 13). The analysis transform $g_{a}$ determines which source realizations $\bm{x}$ get quantized to a specific lattice point in the latent space (colored regions); the synthesis transform $g_{s}$ maps the lattice points to their corresponding reconstructions in the original space (black dots). Regarding rate-distortion performance, the rate-redundancy matches that of ECVQ in 2-dimensions.

#### n 𝑛 n italic_n -dimensions.

Here, we choose $Q_{\Lambda}$ to be the best-known lattices in $n$-dimensions, including $A_{2}$ (hexagonal), $D_{4}^{*}$, $E_{8}$, and $\Lambda_{24}$. As $n$ grows, ECVQ performance with approach the rate-distortion function of the source, which for the Gaussian (squared error distortion) is given by $R(D)=\frac{1}{2}\log\left(\frac{\sigma^{2}}{D}\right),0\leq D\leq\sigma^{2}$
and the Laplacian (absolute error distortion) is given by $R(D)=-\log(\lambda D),0\leq D\leq 1/\lambda$.
as provided by Cover and Thomas (); Equitz and Cover (). For the Banana 1-d marginal, we use Blahut-Arimoto (Blahut, ; Arimoto, ) to estimate $R(D)$ numerically since it is unknown in closed-form.

Figure: Figure 7: Rate-redundancy of Gaussian i.i.d. sequences.
Refer to caption: x5.png

Figure: (a) Laplace ($L_{1}$ distortion).
Refer to caption: x6.png

We plot $\mathcal{E}_{\mathsf{c}}(D)$ against $D$ and $R$. On the Gaussian source, we see that in Fig. 7, LTC performs very close to optimal ECVQ in dimensions $n=2,4,8$ using the hexagonal, $D^{*}_{n}$, and $E_{8}$ lattice respectively. For larger $n$, ECVQ is too costly to compute, but we see that LTC with the Leech lattice for $n=24$ continues the trend, further reducing the rate-redundancy. We observe similar performance gains for the Laplacian in Fig. 14; we do not show the rate-redundancy for Banana-1d since its rate-distortion function does not have a closed-form.

Since each of the lattice quantizers used have efficient CVP algorithms, LTC succeeds in larger dimensions and rates, whereas ECVQ could not even be run due to its complexity that is exponential in dimension and rate. Note that ECVQ performance provides a lower bound on $\mathcal{E}_{\mathsf{c}}(D)$ for any fixed $n$, which indicates that LTC is successfully leveraging space-filling efficiency in higher dimensions.

A standard rate-distortion figure is shown in Figs. 1,8, where LTC outperforms NTC by a constant rate reduction for a fixed SNR, in all three sources. This provides numerical evidence that a companding performance gap corresponding to the lattice’s packing efficiency such as (5) also holds for non-asymptotic rates.

### 5.3 Ablation Study

Figure: (a) LTC with $A_{n}$ and $D^{*}_{n}$ lattices.
Refer to caption: x8.png

We analyze (i) the lattice choice, (ii) the training objective used, and (iii) the effect of density models on the entropy estimate of $\hat{\bm{y}}$. For the lattice choice, only using the best-known lattices in each dimension (in terms of NSM) yielded near-ECVQ performance. If sub-optimal lattices are used, this was not the case. Fig. 8(a) shows the $A_{n}$ and $D^{*}_{n}$ lattices used for $n=8$; they perform about the same as ECVQ for $n=4$, but still yield coding gain over NTC.

For the training objective used, we compare using dithered quantization with STE for the rate term (i.e., using likelihoods from (13) and (11) respectively). Shown in Fig. 8(b), for low rates, the reported cross-entropies from dithering are much larger than that of STE, but converge to the STE performance as rate increases. One possible reason is that the convolved density in (13) over-smooths the true density of $\bm{y}$ in the low-rate regime, leading to sub-optimality.

For the entropy models, we compare the factorized density with the normalizing flow density in Fig. 16. Both models use STE during training. Interestingly, when computing the true entropy of the quantized latent $H(\hat{\bm{y}})$, both models yielded near-optimal performance (curve labeled $H(\hat{\bm{y}})$). However, the cross-entropy upper bound of the factorized density was not as tight as that of the normalizing flow. The normalizing flow has a near-zero gap between $H(\hat{\bm{y}})$ and the cross entropy.

### 5.4 One-Shot Compression of Vector Sources

In general, $\bm{x}$ is a vector source with non-i.i.d. entries. NTC and LTC function as one-shot codes for $\bm{x}$, and we use $d$ to refer to the dimensionality of $\bm{x}$, and $d_{y}$ the latent dimensionality. This setting encompasses one-shot compression of the Sawbridge (Wagner and Ballé, ) as well as the full 2-d Banana source (Ballé et al., ). These works demonstrate that NTC already recovers optimal one-shot coding on these sources, so there is no gain possible over NTC for one-shot codes. Instead, we consider real-world sources such as the Physics and Speech datasets (Yang and Mandt, ), where $d=16$ and $33$ respectively. It is unknown if NTC is one-shot optimal on these sources.

Figure: Figure 10: One-shot compression of the Speech source.
Refer to caption: x10.png

For NTC, $\bm{x}$ is mapped to a $d$-dimensional latent space. For LTC, we map $\bm{x}$ to a latent space with dimension of the lattice used. For these sources, we found training instabilities with STE, so dithering is used instead. At higher rates, we also observed stabler training with the factorized density model compared to the normalizing flow; more details in appendix. For Speech, we observe coding gain for rates above $\approx 7$ bits with the $D_{4}^{*}$ and $E_{8}$ lattices, shown in Fig. 10. Shown in Fig. 17, coding gain is observed on Physics for LTC in different rate regimes. LTC with the hexagonal lattice performs best at lower rates, but the performance drops off as rate increases. At higher rates, LTC with 4 latent dimensions performs the best with $D_{4}^{*}$ lattice.

### 5.5 Block Compression of i.i.d. Vector Sources

For a general vector source $\bm{x}$, gains over the optimal one-shot coding performance can be achieved by coding i.i.d. blocks of $\bm{x}$. We apply BLTC to the Sawbridge and Banana sources, for which NTC was shown to be one-shot optimal. For these i.i.d. vector sources, we again use STE for the BLTC models.

Figure: (a) Sawbridge.
Refer to caption: x11.png

For the Sawbridge, the optimal one-shot performance is given by $H(D)$ (Wagner and Ballé, , Thm. 2), and $R(D)$ is estimated using NERD (Lei et al., ). Shown in Figs. 10(a), 15, we verify that BLTC bridges the gap between $H(D)$ and $R(D)$. In this case, BLTC uses a latent dimension $d_{y}=1$, and the nominal dimension is $d=1024$. For the Banana source, we estimate $R(D)$ using Blahut-Arimoto, and again, LTC bridges the gap between NTC performance and $R(D)$. Here, BLTC uses a latent dimension of $d_{y}=2$ and $d=2$. Thus, we verify on two sources that BLTC can outperform optimal one-shot codes due to the use of block coding. Finally, for real-world data, on the Physics source we additionally see coding gains with BLTC, in Fig. 18. For Physics, optimal one-shot coding is unknown (improved one-shot coding was achieved via LTC in the previous section); however, block compression is able to outperform NTC.

## 6 Conclusion

We show that NTC is sub-optimal on i.i.d. sequences due to its use of the integer lattice. By leveraging better lattices in higher dimensions, LTC compresses i.i.d. sequences nearly optimally, improves one-shot coding performance, and enables the design of i.i.d. vector block compressors.

## 7 Acknowledgements

We sincerely thank Lucas Theis for discussions which helped develop many ideas in this work.

## Appendix A Experimental Details

The code for these experiments can be found at [https://github.com/leieric/lattice-transform-coding](https://github.com/leieric/lattice-transform-coding). In all experiments, we use a batch size of 64 and train until convergence using Adam. For the i.i.d. scalar sequences, we found that using CELU nonlinearities with no biases in the transforms sometimes helped improve training stability.

#### Source generation.

The Gaussian and Laplacian sources are straightforward to simulate. The Banana source is taken from the Tensorflow Compression library [Ballé et al., ]. To generate the Banana-1d marginal, we simply take the first dimension of the 2-d Banana realizations. The Sawbridge source [Wagner and Ballé, ] is a continuous random process

$$ $X(t)=t-\mathbbm{1}\{t\geq U\},\quad 0\leq t\leq 1,$ (16) $$

where $U\sim\mathrm{Unif}([0,1])$. Following Wagner and Ballé [], we uniformly sample 1024 time steps between 0 and 1, which creates 1024-dimensional vectors. For these sources, we sample $10^{7}$ samples as our training dataset. Finally, the Physics and Speech sources are taken from [https://github.com/mandt-lab/RD-sandwich](https://github.com/mandt-lab/RD-sandwich).

#### Efficiency.

The main computational bottleneck that LTC has is the use of Monte-Carlo integration for (11). Sampling uniformly random vectors from the Voronoi region of a lattice requires a quantization step, and hence solving the CVP. Since we use 4096 sampled vectors per batch, this means CVP needs to be solved on all 4096 vectors. The lattices chosen above all have GPU-friendly implementations that can parallelize across the input vectors. On NVIDIA RTX 5000 GPUs with 16 GB memory, LTC training until convergence took at most a few hours for the Speech and Physics sources, and minutes for the i.i.d. scalar sequences.

## Appendix B Additional Figures

#### Monte-Carlo Samples.

For lattices of lower dimensions, we found that setting the number of Monte-Carlo samples $N_{\text{int}}$ for approximation of (11) to be 4096 was sufficient to recover near-ECVQ performance. For the Leech lattice of dimension 24, however, better performance could be recovered by increasing $N_{\text{int}}$ at the cost of increased GPU memory requirements; see Fig. 12.

Figure: Figure 12: Effect of $N_{\text{int}}$ on the Leech lattice, for i.i.d. Gaussian sequences.
Refer to caption: x13.png

We see that setting $N_{\text{int}}=1024$ results in worse performance than ECVQ for $n=8$ at low rates. Setting $N_{\text{int}}=4096$ improves the performance, and nearly approaches that of $N_{\text{int}}=15000$.

Figure: Figure 13: NTC quantization regions in the latent space (integer lattice).
Refer to caption: extracted/5464271/figures/quantizers/square_latent.png

Figure: Figure 14: Rate-redundancy of Laplacian i.i.d. sequences with L1 distortion.
Refer to caption: x14.png

Figure: Figure 15: Rate-redundancy of Sawbridge i.i.d. sequences with L2 distortion. Here, we measure rate-redundancy with respect to $H(D)$, since a closed-form expression for $H(D)$ is known, but not for $R(D)$.
Refer to caption: x15.png

Figure: Figure 16: Ablation study on entropy models. Cross-entropy (CE) is plotted for both factorized prior (FP) and normalizing flow (Flow) density models, as well as the true entropy of the quantized latent $H(\hat{\bm{y}})$. The cross-entropy from the normalizing flow is a much tighter upper bound on $H(\hat{\bm{y}})$ compared to that of the factorized density.
Refer to caption: x16.png

Figure: Figure 17: One-shot compression of the Physics source.
Refer to caption: x17.png

Figure: Figure 18: Performance of BLTC on i.i.d. sequences of Physics realizations, where $n$ is the block-length (number of realizations compressed).
Refer to caption: x18.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. I'm confused by how LQ is applied when the source dimension does not equal the lattice dimension. It seems the lattice dimension used in this paper only goes up to 24, but the source dimension can be 33 (Speech) or even higher (Kodak). Is it based on some kind of product lattice then?


2. Related to above, in the block LTC scheme (Sec 4.3), what is the motivation for grouping columns of the transform coefficients matrix into length-n blocks? Given that the data $x$ and the resulting latent tensor $y$ is often high dimensional, can we do block coding on length-n blocks of $y$ in some other ways? I would expect more gains if the entries in each block are **not** i.i.d.

3. Any idea why the R-D performance is sometimes worse when a bigger latent space ($d_y$) is used, all else being equal (e.g., the $d_y=4$ curve is worse than $d_y=2$ at low rates in the Physics Figure 7)? The usual observation in NTC is that the R-D performance can be improved by increasing the latent dimensionality when the latter is too small (esp. towards the low distortion regime), but further increasing the latent dimensionality beyond a sufficient capacity generally does not harm the R-D performance.
soundness: 3
strengths: I consider the originality of the paper to be moderately high. Although the ideas of lattice coding, nested quantization, block coding, etc. have been known for a while, and lattice coding itself even explored in the contemporary work of Kudo et al. (2023), this paper brings these ideas together and gives novel insights for the performance gap between the current NTC methods and the theoretical optimum of R(D). 

The paper is well-written and mostly easy to follow (although see issues below), with insightful explanations and experiments. 

I consider the contribution of this paper timely and significant. Unlike much of current research on NTC which is in architectural / entropy modeling improvements, this paper examines a fundamentally different component of NTC and makes a compelling case of the advantage of lattice quantization over scalar quantization.  Achieving the theoretical advantage of block coding / VQ has been a non-trivial goal of lossy compression, and this paper offers valuable perspectives and a potential path forward.
summary: This paper proposes lattice transform coding (LTC), which uses lattice quantization beyond the conventional scalar quantization in neural lossy compression. The paper offers evidence that learned transforms in conventional nonlinear transform coding (NTC) are insufficient in overcoming sub-optimal quantization in the latent space, and shows that LTC with more efficient space-filling lattices indeed bring improvements in R-D performance. Besides new training and rate estimation techniques for LTC, the paper also proposes nested LTC to lower the practical computation cost of entropy coding, as well as block LTC to perform block coding in order to close the gap to the information-theoretic optimal performance (R(D)). 
Experiments on various synthetic and real-world data sources support the effectiveness of the proposed LTC approach and give insight to the suboptimality of current NTC approaches.
weaknesses: 1. A significant difficulty of lattice coding seems to be the complexity of implementing entropy coding. Nested lattice quantization was proposed but appears to come with a significant penalty to R-D performance. Therefore most of the R-D curves for LTC in this paper are theoretical. I don't expect this problem to be solved in this paper, but I think it'll be useful to give insights for how to operationalize the reported R-D curves, in particular considerations for deriving the discrete entropy model from the normalizing flow prior in a way that is reproducible by the decoder. 

2. Some of the the explanations around nested LQ and block LQ were a bit hard to follow. For example, the first two sentences on lines 334 - 335 seemed to come out of nowhere ( "The goal is for the majority of the latent mass to be placed in the Voronoi region of Λc. This happens when QΛc (yf ) = 0."). Similarly, line 375 introduces a new pair of "analysis/synthesis transforms of LTC", denoted $c_a^{(i)}, c_s^{(i)}$ --- why are they necessary, given that we have already operate in the transformed space using $g_a, g_s$? And how should we choose  $c_a^{(i)}, c_s^{(i)}$ if they are indeed necessary? 

Also, a minor typographical error: there seems to be inconsistency between "n" and "d": line 368 explains "Note that we we retain n as the block-length parameter, but use d as the dimension of x." So should line 458 write "d = 16 and 32"?

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: It would be interesting to see more information about the practical aspects of using lattice quantizers during training and for actual coding.
soundness: 3
strengths: The authors correctly provide good evidence that the suboptimality of scalar quantization remains even with learned neural transformations.
Showing it on toy problems helps to make it clearer, but is not good at indicating how much worse on practical cases. 

The paper presents examples of application to other data, and a series of experimental tests are used to answer those question on those data sources.
summary: This paper covers the application of lattice vector quantization to neural lossy compression.

The problem of how to address the suboptimality of scalar quantization for lossy compression has been studied for decades, and even today most practical methods use scalar quantization, and instead rely on other techniques to improve compression.

One important difference is that conventional methods normally rely on linear transform, nonadaptive methods, while this paper novelty is that it studies neural compression that uses nonlinear learned transformations. Numerical experiments demonstrate better compression results.
weaknesses: The experiments are well planned with good comparisons, but somewhat limited. Maybe too much space is used to cover the synthetic cases, and could be better used with real practical data.

It would also be better to have more detailed analysis of the results in those cases. Comments like “…with the Leech significantly outperforming VTM and NVTC,” are too vague and not accurate.

It looks like there are good gains with the modifications to the Cheng2020 method, which may be the most important conclusion of the paper, and that should be more carefully reported.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. Clarification surrounding the choice of hyperparameter $\lambda$ discussed in "Weaknesses" would be appreciated.
2. Experimental results: LTC performs worse than NTC at extremely low bit rates. This is surprising since the benefits of vector/lattice quantization should be more pronounced at such rates.
soundness: 3
strengths: 1. The submission's attempt at marrying nonlinear transform coding (NTC) with lattice quantization (LQ) as a way of sidestepping computationally expensive vector quantization is highly relevant to the learned compression research community.
summary: 1. This submission extends scalar-quantized non-linear transform coding of signals (Ballé et al., 2020) using lattice quantization and coding. A classical result from signal compression is that vector quantization can provide packing gains over scalar quantization especially when applied directly on correlated (i.e. non-decorrelated) signals, at the cost of higher encoding complexity. The authors propose to overcome the complexity of vector quantization using lattice quantization, which can be considered a constrained form of vector quantization. Experiments (on images and speech) were conducted to measure the gain from lattice quantization over the scalar quantization baseline.
weaknesses: 1. The fact that integer quantization of latent variables of NTC does not induce hexagonal quantization cells in the source domain—this is not a surprising or unexpected result, as non-linear transforms still define continuous / diffeomorphic mappings between the two domains. A transform that maps squares (scalar quantization cells in the latent domain) to hexagons (vector quantization cells in the source domain) would be highly discontinuous around cell boundaries, not a mapping that is readily learned by a neural transform in the first place.
2. Eq. (1) is central to this work, yet, no description of the Lagrange multiplier / rate–distortion tradeoff $\lambda$ appears in the manuscript. Wouldn't the choice of $\lambda$ have an impact on the characteristics/performance of the learned non-linear transform? Was an effort made to reparameterize $\lambda$ as in the work of Ballé et al. (2020)? Discussion around this is lacking in the current version of the manuscript.
3. LTC (of latent variables) inducing near-hexagonal quantization cells in the source domain (in the 2D case) seems to be due more to the diffeomorphic / continuous nature of the learned non-linear transform itself. This would mean that one could take the non-linear transform of Ballé et al. (2020) and lattice-quantize the latent variables post-hoc, and achieve similar results to what this work is proposing. The benefits of incorporating the actual lattice structure in the rate term of the optimization should be discussed.
4. Experimental results fail to clearly demonstrate the advantage of LQ over SQ. Could this be an optimization issue? For example, was $\lambda$ was fixed to some value to hurt the performance of the non-linear transform itself?

-------
4. The current presentation seems somewhat unbalanced, with the authors' actual work only starting on page 6. While this is not inherently a weakness per se, this does mean that the authors were perhaps not able to fully explore / describe the actual contributions of their work in the remaining 4–5 pages.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: 1. Is it possible to further reduce the computational complexity of LTC models for training and inference? Authors can try the following things to reduce concern.
   - Provide more detailed computational complexity analysis, comparing LTC to state-of-the-art NTC methods beyond just Cheng2020-NTC.
   - Discuss potential approaches to reduce training and inference time for LTC, such as whether fine-tuning techniques used for NTC could be adapted.
   - Comment on how the complexity of LTC scales with different lattice dimensions and source dimensionality.
   - Address whether there are any trade-offs between computational complexity and rate-distortion performance for LTC.
2. Figure 5 does not show the learned lattice of LTC in the regions of lower probability density. How does LTC perform in these regions? Please provide visualizations or quantitative analysis of LTC's performance in lower probability density regions, and discuss any potential implications for overall compression performance.
3. LTC provides a construction on the quantization that is quite different from previous work. Then, how do the nonlinear transforms affect the rate-distortion performance of the LTC models? Please provide an ablation study or analysis that isolates the effects of the nonlinear transforms on LTC's rate-distortion performance.
soundness: 4
strengths: 1. The idea of using lattice quantization in neural compression has potential whose benefits cannot be achieved by improving transforms and the entropy model.
2. The theoretical analysis and experimental evaluation presented in this paper are convincing. The authors have carried out extensive experiments on both synthetic sources and on higher dimensional real sources.
3. The paper is well written and easy to follow.
summary: The authors show the suboptimality of existing neural compression models for simple multidimensional sources, such as 2-D Gaussians. They point out that the scalar quantization used in current neural compression models can lead to the suboptimality. They then propose a solution, *i.e.*, lattice transform coding, with lattice quantization in latent space. Extensive experiments verify the superiority of the proposed lattice transform coding.
weaknesses: 1. The real-world application of the proposed Lattice Transform Coding (LTC) may be limited by the computational complexity of training and inference. 

   - *On training.* The authors show that training an LTC model per rate point for image compression takes about 10 days. By comparison, training of nonlinear transform coding (NTC) models can be reduced by fine-tuning the lower rate model from a pre-trained higher rate model. I'm not sure if such tricks can also speed up the training process of LTC.
   - *On inference.* The authors provide a comparison of the NTC method (Cheng2020-NTC) with the proposed LTC method (Cheng2020-LTC-E8), which shows that the computational complexity of LTC is comparable to NTC. However, Cheng2020-NTC is a nonlinear transform coding method using spatial auto-regression in entropy modeling, which is notoriously time consuming. The method is further refined in subsequent work, such as wavefront coding, checkerboard context, and channel-wise auto-regression. Notably, the proposed lattice transform coding can be even more computationally expensive than spatial auto-regression. This can be a fatal flaw in applications.

2. Typos:

   - Line 470. The citation of Vimeo-90k should be realized by a `\citep` command instead of `\citet`.
   - Line 472. The citation of CompressAI is not included in the reference.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: See weaknesses.
soundness: 3
strengths: The observation that neural transforms alone cannot overcome scalar quantization limitations is a valuable insight.
The use of dithered noise and Monte Carlo estimation for rate estimation is also noteworthy.
summary: This paper highlights the limitations of scalar quantization in neural compression and introduces lattice transform coding as a solution. Given the challenges in rate estimation within this approach, the authors propose a Monte Carlo-based method to address it. To manage complexity, they further introduce a nested lattice structure. Experimental results are presented using both synthetic datasets and real-world datasets, including speech and the Kodak image dataset.
weaknesses: - The presentation could be improved for clarity. All figures are too small, and Section 3 (on the suboptimality of scalar quantization) would benefit from additional detail beyond the currently limited Figure 2. 
- Additionally, the distinction in block LTC needs clarification.
- Lattice coding encoding complexity remains an issue (except for $n=24$). More specifics on how this scheme was applied to the Kodak image dataset would be helpful.

Meta-review: The authors discuss a lossy compression method based on transform coding. In contrast to the rectangular lattices used in the typical nonlinear transform coding (NTC) setup, the authors propose to use higher dimensional lattices and demonstrate that even in the context of nonlinear transforms, significant gains can be achieved.

The main limitation of the proposed method is in the practical implementation, which is less straight-forward compared to NTC, as noted by the reviewers (e.g., complexity of the entropy model, determining probabilities of each quantization bin in an efficient way). The authors should address these limitations in the final version of the manuscript.

Decision: Accept (Spotlight)