Accelerating Relative Entropy Coding with Space Partitioning

We thank the reviewer for their detailed and insightful review. We are delighted that the reviewer recognizes the difficulty of this task and appreciates the contribution of our work. Below, we respond to the reviewer's concerns and questions. Should the reviewer find our answers satisfactory, we kindly invite them to consider raising their score.

Weakness: Line 21 in the Introduction is misleading, quantization and REC are fundamentally different

We thank the reviewer for raising this concern. Upon revisiting the paragraphs the reviewer mentions, we agree that our argument can easily appear as an apples-to-oranges comparison. Hence, to clarify: the point we wanted to make here is that quantization makes end-to-end learning difficult due to its non-differentiability. Hence, we advocate for REC, which swaps out quantization for some continuous perturbation as a means of limiting the latent representation’s information content, as the latter integrates well with end-to-end training thanks to the reparameterization trick. We will clarify the introduction in our camera-ready version based on this discussion.

However, we’d also like to note that the difference between using 1) quantization + entropy coding or 2) relative entropy coding for lossy compression is not necessarily as big as the reviewer is claiming. The analysis of quantization often relies on treating the quantization error as uniform noise. This approximation has been remarkably successful in the study of quantization and in practice, as it is the foundation of Balle et al.’s neural compression methods as well. Furthermore, Agustsson and Theis (2020) show that quantization-based schemes can be turned into relative entropy coding-based schemes (with uniform target distributions) with a simple modification.

Regarding Questions:

1. Why is it important for the state $S$ to be random?

Thanks for raising this point. By random state $S$, we mean the common randomness between communication parties, a standard requirement by all channel simulation schemes. We will clarify our terminology in the camera-ready version of the paper.

2. Shouldn't this discussion be formulated around the data space?

This is a great question! While we do not discuss this in our paper, this concern is precisely why we use total variation to measure the approximation error, as bounding the error in latent space gives us an immediate guarantee in data space. More formally, denote the decoder by $f$, the true target distribution as $Q$ and the approximate target $\tilde{Q}$. Furthermore, assume $D_{TV}[Q \Vert \tilde{Q}] \leq \epsilon$. By the data processing inequality, we have $D_{TV}[f_{\*}Q \Vert f_{\*}\tilde{Q}] \leq\epsilon$, where $f_{\*}Q$ denotes the pushforward measure. For more discussion on this topic, please Section III in Flamich and Wells (2024); we will update the camera-ready version to include this discussion.

3. Isn't it hard to pick a good $\pi$ given this distribution is constantly changing?

No, thankfully is not difficult!

• first, the partition is fixed across all data points, and hence, this operation is amortized and not expensive;

• second, for ORC (Sec 4.2), to sample from eq (15), we only need to first sample from $Q$ and find the bin in which the sample lies. Therefore, we do not need extra operations for pre-processing; for PFR (Sec 4.1), we need to calculate $\pi$ in eq(12) per dimension for each data point before sampling. But this is still much cheaper compared to sampling;

• additionally, we highlight that the receiver does not need to be aware of $\pi$ to reconstruct the sample. Therefore, while the $\pi$ we use differs for different data points, there is no extra cost in transmitting.

Regarding Limitations:

The results are computationally intensive, and thus impractical, even for MNIST.

We do not necessarily agree that the VAE results are computationally intensive and impractical. We use $2^{16}$ samples to encode a block that requires around $2^{50}$ samples by previous methods. With our method, we have already largely reduced the computation cost and made REC more practical, though we definitely believe that this can be further improved.

Requiring the value of mutual information to be known by encoder/decoder seems a bit much.

We agree that this is not a totally general assumption. However, it is not too much for a practical neural compression application. In most neural compression applications, including data compression and (potentially) federated learning, we have access to training data on which we can easily estimate the dimensionwise mutual information. Also, compared with assumptions for previous faster methods (either assuming 1D or uniform/Gaussian with a known scale per dim), our assumption is already a lot more general.

References

Agustsson, E., & Theis, L. (2020). Universally quantized neural compression. Advances in neural information processing systems, 33, 12367-12376.

Flamich, G., & Wells, L. (2024). Some Notes on the Sample Complexity of Approximate Channel Simulation. arXiv preprint arXiv:2405.04363.

We thank the reviewer once again for their effort in the reviewing process. As there are only a few working days left in the discussion period, we would like to ask if our response has satisfied the reviewer’s concerns. If so, we kindly invite the reviewer to consider raising their score. If any concerns remain, we are happy to discuss them further here.

We are thankful for the reviewer's time and detailed review and are delighted that the reviewer found our paper clear and appreciated our technical and theoretical contributions. Below, we respond to the reviewer’s questions. We are happy to discuss any further concerns the reviewer might have. Should we have addressed the reviewer’s concerns and questions, we kindly invite them to raise their score.

1. it is unclear to the reviewer how the grid should be set:

We discuss how we set the bins in our manuscript from line 241 to line 249; we will signpost this discussion better in the camera-ready version.

In short, in most neural compression applications, we can compute the mutual information per dimension $I_d$ from the training set. Specifically, for each training data $X$, we have access to its posterior $Q_{Z|X}$ during training, and we can calculate the KL divergence between $Q_{Z|X}$ and the prior $P_Z$ along each dimension. The dimension-wise MI $I_d$ is estimated by averaging the dimension-wise KL across all training data. Knowing the mutual information per dimension $I_d$, we partition the d-th axis into approximately $2^{n_d}$ intervals, where $n_d = D_{KL}[Q||P]\cdot I_d / \sum_{d’} I_{d’}$. We discuss the motivation for this choice in Appendix B.2. We also provide the implementation details from line 621 to line 630 on Page 24.

2. how the binning technique impacts the final result (besides codelength):

In Appendix B.2, we discuss how different binning strategies influence the results. For a further illustration of the effect of the chosen grid on the efficiency of the algorithm, please see Figure 1 in the attached rebuttal PDF.

Essentially, if we only partition the space by dividing the axes where the prior $P$ and the target $Q$ have the same marginal, we will not reduce the runtime. However, luckily, as we discussed above, in most neural compression applications, we can compute the mutual information per dimension $I_d$ on the training set. If $I_d> 0$, it tells us, on average, that the prior $P$ and the target $Q$ will not have the same marginal along the d-th axis. Therefore, we propose to partition the d-th axis into approximately $2^{n_d}$ intervals, where $n_d = D_{KL}[Q||P]\cdot I_d / \sum_{d’} I_{d’}$.

To help the reviewer further understand this argument, we use a toy example for visualization in our attached rebuttal PDF. In Figure 1 in our rebuttal PDF, we create a toy 5D Gaussian target, and we set dimension 1 to have 0 mutual information (we call it a collapsed/uninformative dimension). We compare our proposed method with 3 different binning strategies: only partitioning the collapsed dimension; randomly assigning the number of intervals per dimension; and partitioning the d-th axis into approximately $2^{n_d}$ intervals, where $n_d = D_{KL}[Q||P] I_d / \sum_{d’} I_{d’}$. We also include standard ORC for reference. We run the 4 settings with the same number of samples (20) and repeat each setting 5000 times. The error can be seen from the histogram of 5000 encoded results. We can see that partitioning according to mutual information yields the best results, whereas partitioning only the collapsed dimension yields almost the same results as standard ORC.

3. how the choice of prior for the binning distribution influences the algorithm:

We assume the reviewer is asking how the results will be influenced when each bin has different probabilities under the prior. We answer this question below, and we are happy to address any further questions if we have misunderstood the reviewer's question.

Setting all bins to have the same probability mass under the prior is the basic requirement for our algorithm to adhere to the codelength outlined in Theorem 3.1 and Proposition 3.2. We note this in lines 104 and 414.

This can be understood intuitively: under an ideal Bayesian treatment, the average Q_{Z|X} across all data X should equal the prior $P_Z$. Also, we note that the partitions are shared across all data X in the dataset. Therefore, ideally, we will expect the average probability (average is taken over all randomness and all X) that the encoded sample falls into any bin should be the same. We can also see the reason from our proof of Theorem 3.1. Specifically, we replace all $P(B_j)$ in eq(37) by $1/J$ in eq(38), which simplifies the expression and is essential for the subsequent derivation.

In Summary: our algorithm and theory require all bins to have the same probability under the prior. Theorem 3.1 and Proposition 3.2 will hold for all binning strategies that satisfy this requirement. However, there are many (actually, infinitely many) binning strategies that satisfy this requirement. Different choices of binning strategies can have different runtimes or different biases given the same runtime budget. To this end, we propose to partition the d-th axis into approximately $2^{n_d}$ intervals, where $n_d = D_{KL}[Q||P] I_d / \sum_{d’} I_{d’}$. We have discussed the motivation and implementation details in our manuscript and provided an additional visualization to aid understanding in our rebuttal PDF(Fig 1) in the global response.

We thank the reviewer once again for the effort put into reviewing our paper. As there are only a few working days left in the discussion period, we would like to ask if our response has satisfied the reviewer’s concerns. If so, we kindly invite the reviewer to consider raising their score. If any concerns remain, we are happy to discuss them further here.

We thank the reviewer for their dedicated time and constructive review, which we feel largely enhances the quality of our paper. While the reviewer may not be familiar with this field, we believe they understand most of our manuscript and method well.

We are delighted that the reviewer found our paper well-written and recognized its soundness. Below, we respond to the reviewer's concerns, and we are happy to discuss any further questions the reviewer might have. Should the reviewer find our answers satisfactory, we kindly invite them to consider raising their score.

Weakness: the idea and proofs are elementary: While we agree that the main idea behind our method is simple, we respectfully disagree that it is a weakness on three grounds. (a) ideas that are simple and work well in practice have the largest potential to be impactful in the field as they are easy to understand and replicate. This is the case for our method: our ideas improved performance by many orders of magnitude in terms of the largest KL values that many future REC algorithms could consider. (b) the theory justifying this idea required new insight; mainly around deriving the overhead term $\epsilon$ that previous works did not deal with or consider. (c) the general insight that we can use the dimensionwise mutual information to devise schemes to improve compression speed is a significant contribution to the field, and we believe that any future scheme that does not use this additional information will be either limited in speed (as all previous schemes are) or codelength.

Regarding Questions:

1. How Theorem 3.1 compare with the bound from [Li and El Gamal, 2018, Theorem 1] when $J=1$

when $J=1$, our algorithm is equivalent to PFR, and therefore, we end up with the same bound as [Li and El Gamal, 2018, Theorem 1]. This can be seen from eq(38) on page 17: when $J=1$, eq (38) becomes $E\log_2[ \frac{q(z)}{p(z)}+1]$. We can then follow the proof of [Li and El Gamal, 2018, Appendix p13].

Could we recover a bound with a dependence in $J$?

Please note that $\epsilon$ in Theorem 3.1 already expresses a fine-grained dependence on $J$. Furthermore, Proposition 3.2 also depends on $J$ in the sense that it is derived by the extreme case $J = 2^{KL}$ and holds for all $J \leq 2^{KL}$. It is an interesting question if there is a bound that interpolates $J=1$ and $J = 2^{KL}$. In our paper, we did not explore this because, for our method, we always want to set $J$ as large as possible to maximize the gain in speed.

2. Could the author give more details about how these intervals are constructed?

As we discussed in our manuscript from line 241 to line 249, we partition the d-th axis into approximately $2^{n_d}$ intervals, where $n_d = D_{KL}[Q||P] I_d / \sum_{d’} I_{d’}$. We provide the implementation details of constructing these bins from line 621 to line 630.

In fact, we empirically found that the methods are NOT very sensitive to how we construct the intervals, as long as we avoid partitioning along uninformative dimensions. Please see the rebuttal PDF (Fig 3) for further empirical verification.

3. Would it be possible to automatically find a transformation prior to the binning that would provide better results while keeping the same strategy:

It is not possible to design such a transformation. Note that to avoid sending the transformation (which can be expensive), it should be shared across all data points X. Therefore, if we want to find a transformation to the prior, it is equivalent to asking if we can find a better prior that provides better results on average.

4. Is there a way to update $\pi$ along the algorithm

Since our algorithm will work for any $\pi$, it is possible to update it. However, there is no point in doing so as the optimal form of $\pi$ is already given in eq (12) and eq (15), which we can use at negligible computational cost.

5. It would be interesting to add an ablation study on the number of intervals per axis

Thank you for this suggestion. We added ablation studies on toy examples and INR experiments. Please find the results in the PDF we uploaded to global rebuttal. We will add these plots to our camera-ready version. We now explain these studies, and please also see the global rebuttal for a more detailed description:

• We first provide a qualitative visualization in Figure 1 in our rebuttal comparing 3 different binning strategies: only partitioning the collapsed dimension; randomly assigning the number of intervals per axis; and assigning intervals per axis according to MI (our suggested strategy). We also include standard ORC for reference. As we can see, our suggested strategy works best, whereas partitioning only the collapsed dimension yields almost the same results as standard ORC. This verifies our discussion on how the partitioning strategies will influence the runtime in Appendix B.2 in our manuscript.

• We then provide quantitive ablations on toy examples (Fig 2) and INR experiments (Fig 3). We can see our suggested partition strategy (partitioning according to MI) is always better than randomly assigning intervals per axis. Also, we found if we first remove uninformative dimensions (mutual information $\approx$ 0), and then randomly assign intervals to other dimensions, the performance is only slightly worse than our proposed partition strategy. This indicates that our algorithm is not very sensitive to how we construct the intervals, as long as we avoid partitioning along uninformative dimensions.

6. it would be interesting to have Tab 1 in time

The sample size is linearly proportional to the runtime, so we only provide this in Fig 3 and Tab 1. The sample size is a better metric as it is implementation-agnostic, unlike wallclock time.

I do not understand what is meant by runtime in Fig 3

This figure shows ‘how many samples in ORC we need for a certain bias’. The y-axis shows the sample size to achieve a certain bias (measured by MMD).

We thank the reviewer once again for their effort in reviewing our paper and for their constructive review. As there are only a few working days left in the discussion period, we would like to ask if our response has satisfied the reviewer’s concerns. If so, we kindly invite the reviewer to consider raising their score. If any concerns remain, we are happy to discuss them further here.