Title: Accelerating Relative Entropy Coding with Space Partitioning
Venue: NeurIPS 2024
Accepted: yes
Keywords: Relative Entropy Coding; Neural Compression; Information Theory
Abstract: Relative entropy coding (REC) algorithms encode a random sample following a target distribution $Q$, using a coding distribution $P$ shared between the sender and receiver. Sadly, general REC algorithms suffer from prohibitive encoding times, at least on the order of  $2^{D_{\text{KL}}[Q||P]}$, and faster algorithms are limited to very specific settings. This work addresses this issue by introducing a REC scheme utilizing space partitioning to reduce runtime in practical scenarios. We provide theoretical analyses of our method and demonstrate its effectiveness with both toy examples and practical applications. Notably, our method successfully handles REC tasks with $D_{\text{KL}}[Q||P]$ about three times greater than what previous methods can manage, and reduces the bitrate by approximately 5-15\% in VAE-based lossless compression on MNIST and INR-based lossy compression on CIFAR-10, compared to previous methods, significantly improving the practicality of REC for neural compression.

Accelerating Relative Entropy Coding with Space Partitioning

Abstract

Relative entropy coding (REC) algorithms encode a random sample following a target distribution Q , using a coding distribution P shared between the sender and receiver. Sadly, general REC algorithms suffer from prohibitive encoding times, at least on the order of 2 D KL r Q || P s , and faster algorithms are limited to very specific settings. This work addresses this issue by introducing a REC scheme utilizing space partitioning to reduce runtime in practical scenarios. We provide theoretical analyses of our method and demonstrate its effectiveness with both toy examples and practical applications. Notably, our method successfully handles REC tasks with D KL r Q || P s about three times greater than what previous methods can manage, and reduces the bitrate by approximately 5-15% in V AE-based lossless compression on MNIST and INR-based lossy compression on CIFAR-10, compared to previous methods, significantly improving the practicality of REC for neural compression.

1 Introduction

Let's consider a two-party communication problem where the sender wants to transmit some data X to the receiver. A widely used approach is transform coding (Ballé et al., 2020), where X is first transformed to a discrete variable Z and entropy coded to achieve an optimal codelength on average. However, directly finding a discrete Z is difficult in many scenarios. For example, in lossy image compression, X represents some image, and Z represents the latent embedding output by an encoder network in a model akin to the variational auto-encoder (Kingma and Welling, 2013). A common solution to obtain a discrete Z is to first learn a continuous variable and quantize it (Ballé et al., 2017).

However, perhaps surprisingly, there is also a way to directly handle a continuous Z in this pipeline. Specifically, instead of a deterministic value, the sender transmits a random sample following a posterior Z ' Q Z | X 1 . These algorithms are referred to as relative entropy coding (REC, Flamich et al., 2020), and also known as channel simulation or reverse channel coding (Theis and Yosri, 2022). Li and El Gamal (2018) showed that the codelength of such an algorithm is upper-bounded by the mutual information 2 between X and Z plus some logarithmic and constant overhead:

REC has a clear advantage over quantization: quantization is a non-differentiable operation, while REC directly works for continuous variables and eases the training of some neural compression models, which highly rely on gradient descent. Also, Theis and Agustsson (2021) exemplified that stochastic encoders can be significantly better than their deterministic counterpart if we target realism.

However, the encoding time of REC algorithms is at least typically on the order of 2 D KL r Q } P s , which is prohibitively long in practice. 3 While there are several works on accelerating REC (Flamich et al., 2022; Flamich and Theis, 2023; Flamich et al., 2024), so far they only work on highly limited problems. In fact, in practice, the common option to employ REC is to segment Z into the concatenation of independent blocks, denoted as Z ' r Z 1 , Z 2 , ¨ ¨ ¨ , Z K s , where the coding cost of each block Z k is approximately κ bits. This strategy reduces the runtime to O p K 2 κ q . However, by Equation (1), the logarithmic and constant overhead only becomes negligible if I r X ; Z i s is sufficiently large. For example, for the runtime to be feasible, κ is set between 16 to 20 in Havasi et al. (2019); Guo et al. (2024); He et al. (2023). In this case, the overhead will typically constitute 40% to 50% of the total mutual information, resulting in sub-optimal compression performance.

Our work's primary aim is to reduce the complexity of REC runtime for more practical settings. Specifically, our contributions are

We propose a faster REC framework based on space partitioning. Equivalently, our method can be viewed as introducing a search heuristic to a standard REC algorithm, which can significantly reduce the algorithm's runtime when chosen appropriately. Furthermore, we provide theoretical analysis, showing that our method achieves a close codelength to Equation (1) with an extra cost ϵ that is negligible for some commonly used distributions in neural compression.

We show that, interestingly, using different space partitioning and different search heuristics only influences the runtime but not the (upper bound on) codelength. Following this, we discuss two cases: 1) encoding exact samples from the target and 2) encoding approximate samples to further reduce the computational complexity.

We draw attention to a hitherto unused fact for designing fast, practical relative entropy coding schemes: when the sender wishes to communicate vector-valued random variate Z | X to the receiver, we may assume that they share knowledge of the dimension-wise mutual information I r Z i ; X s for each dimension i , as opposed to just the total I r Z ; X s . Incorporating this information into our space partitioning scheme allows us to construct faster relative entropy coding algorithms.

We conduct experiments on synthetic examples and neural codecs, including a VAE-based lossless codec on MNIST (LeCun and Cortes, 1998) and INR-based lossy codecs on CIFAR-10 (Krizhevsky et al., 2009). We demonstrate that our method can handle blocks with larger κ using much fewer samples and reduces the bitrate by approximately 5-15% compared to previous methods.

2 Preliminary

In this paper, we focus on accelerating the Poisson functional representation (PFR; Li and El Gamal, 2018) and ordered random coding (ORC; Theis and Yosri, 2022), and hence we discuss these in detail below. However, we note that our space partitioning scheme is also applicable to other relative entropy coding (REC) algorithms (Flamich and Theis, 2023; Flamich et al., 2024; Flamich, 2024).

Relative entropy coding (REC). Consider a two-party communication problem, where the sender wants to transmit some data X to the receiver. However, instead of encoding X directly, the sender first transforms X into a representation Z that they encode. In REC, we allow Z to be stochastic with Z ' Q Z | X . Then, the goal of REC is to encode a single , random realization Z ' Q Z | X with the assumption that the sender and receiver share a coding distribution P and have access to common randomness S . In practice, the latter can be achieved by a shared pseudo-random number generator (PRNG) and seed. Given these assumptions, the optimal coding cost is given by H r Z | S s , as the code cannot depend on X since the receiver doesn't know it. Surprisingly, Z can be encoded very efficiently, as Li and El Gamal (2018) show that

Note that I r X ; Z s and hence H r Z | S s can be finite even if Z is continuous and H r Z s is infinite. Next, we describe a concrete scheme with which we can encode Z at such an efficiency.

Poisson functional representation (PFR). To encode a sample from the target distribution Q with density q using the coding distribution P with density p , PFR (Li and El Gamal, 2018) starts by

…

…

Figure 1: An illustrative comparison between the standard REC algorithm and REC with space partitioning. We illustrate the prior P 's density in blue and Q 's density in orange. (a) In a standard REC algorithm, we may draw numerous samples (colored in red) before identifying one that aligns well with Q (colored in green). The majority of these samples do not directly contribute to the desired result. (b) In the method we propose, we first divide the search space into smaller grids and then reweight each grid. This amounts to adjusting the prior P to a search heuristic P 1 , which can align better with Q . The samples from P 1 will thus be more relevant to Q , potentially reducing the runtime.

drawing a random sequence Z 1 , Z 2 , ¨ ¨ ¨ from P using the public random state S . Furthermore, the sender draws a sequence of random times T 1 , T 2 , ¨ ¨ ¨ as follows:

Next, letting r ' q { p be the density ratio, the sender calculates τ n Ð T n { r p Z n q for each sample and returns N ˚ Ð arg min i P N t τ i u . Using the theory of Poisson processes, it can be shown that Z N ˚ ' Q as desired (Maddison, 2016). Additionally, while the minimum is taken over all positive integers, in practice, N ˚ can be found in finite steps if r is bounded. In fact, in expectation, PFR will halt after 2 D 8 r Q } P s steps (Maddison, 2016). We summarize this process in Algorithm 1.

Ordered random coding (ORC). Unfortunately, PFR's random runtime can be a significant drawback. In practice, we may want to set a limit on the number of iterations to ensure a consistent and manageable runtime at the cost of some bias in the encoded sample. To this end, Theis and Yosri (2022) proposed ordered random coding (ORC). Specifically, in ORC with N candidates, rather than calculating T n by Equation (3), we first draw N i.i.d. sample from Exp p 1 q , and sort them in ascending order as ˜ T 1 1 ď ˜ T 1 2 ď ¨ ¨ ¨ ď ˜ T 1 N . We then set T n ' ˜ T 1 n for each n ' 1 , 2 , ¨ ¨ ¨ , N . In practice, instead of generating the T n -s in O p N log N q time by sorting, Theis and Yosri (2022) suggest an iterative procedure similar to Equation (3) with O p N q time complexity:

3 Relative Entropy Coding with Space Partitioning

In this section, we describe our proposed algorithm and analyze its codelength. To motivate our method, recall that in PFR, the sender draws a random sequence from P , and examines each sample's density ratio. We can interpret this process as a random search in P 's support, aiming to find a point that has a relatively high density ratio between Q and P . However, when Q concentrates only within a small region of P 's support, most of the search does not contribute to the final outcome. Thus, can we instead quickly navigate the search towards the region where Q is concentrated?

For one-dimensional distributions Q and P , the answer is affirmative, as we can perform a branchand-bound search by partitioning the 1D space on the fly (Maddison et al., 2014; Flamich et al., 2022). Unfortunately, it is unclear how to generalize these adaptive partitioning strategies to spaces with dimension greater than one. Instead, in Section 3.1, we propose to partition the space in advance according to a rule that the sender and receiver agree on such that we can carry out the search fast enough in practical problems and retain an efficient codelength.

3.1 Coding Scheme

Given a shared coding distribution P and a target Q , our algorithm proceeds as follows:

The sender and receiver agree on a partition of P 's support consisting of J bins t B 1 , ¨ ¨ ¨ , B J u with equal probability mass, i.e. P p B j q ' 1 { J for j ' 1 , . . . , J . As we will discuss later, these

Algorithm 1 Encoding of standard PFR

Algorithm 2 PFR with Space Partitioning

bins can overlap, but to aid understanding for now, it may be helpful to imagine the space is partitioned by non-overlapping bins.

According to the target distribution Q , the sender selects a categorical distribution for the bins, with event probabilities defined as π p j q for j ' 1 , 2 , ¨ ¨ ¨ , J . We can view this categorical distribution as a reweighting of each bin to adjust the coding distribution P . The coding distribution adjusted by this reweighting can be better aligned with Q . We will discuss the choice of π later.

Then, the sender starts to draw and examine samples iteratively, similar to the PFR algorithm. However, instead of drawing samples directly from P , the sender first samples a bin from π and then draws a sample from the prior restricted to this bin. Specifically, at each iteration n :

(a) the sender first draws a bin index, j n , according to the distribution π ;

(b) the sender then draws a sample Z n ' P | B j n . This sample is generated with the random state S j n associated with this bin. It's critical that each bin has a unique random state to ensure different bins have different random sample sequences. In practice, this can be easily achieved by setting the PRNG's seed of the bin B j to j , for j ' 1 , 2 , ¨ ¨ ¨ , J . Note that given these random states, the value of Z n is uniquely determined by the bin index j n and its sample index within this bin, which we denote by ˜ n j n . For the sake of clarity, we will call n the global sample index and ˜ n j n the local sample index hereafter;

(c) the sender examines the sample by calculating τ n :

and keeps track of N ˚ Ð arg min i ' 1 , 2 , ¨¨¨ ,n t τ i u . Here, T n is also obtained by Equation (3);

(d) finally, the sender checks the following stopping criterion:

If the criterion is met, halt and return j N ˚ and ˜ n N ˚ ; otherwise, proceed to the next iteration.

We detail the above coding scheme in Algorithm 2. As a comparison, we describe the standard PFR algorithm in Algorithm 1, and highlight their difference in red. We also illustrate these two algorithms from a high-level point of view in Figure 1. It is easy to verify that Algorithm 2 is equivalent to running standard PFR with an adjusted prior P 1 , whose density is defined as

Therefore, the encoded sample is guaranteed to be Q -distributed. By choosing a sensible π , the adjusted prior P 1 can be more closely aligned with Q than the original P , potentially decreasing the runtime. From this point forward, we will refer to P 1 as the search heuristic .

Note that the receiver does not need to be aware of π to decode the sample. Instead, after receiving the bin index j N ˚ , the receiver can construct the random sequence from P | B j N ˚ with seed S j N ˚ . Then, taking the ˜ n N ˚ -th sample in this sequence, the receiver successfully retrieves the desired sample.

3.2 Codelength of the two-part code

In our proposed algorithm, the sender needs to transmit a two-part code that includes both the bin index and the local sample index. This is distinct from the standard PFR, where the sender transmits only the sample index, hence potentially raising concerns about the codelength. Indeed, we find the two-part code can introduce an extra cost, as outlined in the following theorem:

Theorem 3.1. Let a pair of correlated random variables X , Z ' P X , Z be given. Assume we perform relative entropy coding using Algorithm 2 and let j ˚ denote the bin index and ˜ n ˚ the local sample index returned by the algorithm. Then, the entropy of the two-part code is bounded by

We prove Theorem 3.1 in Appendix C.1. Note, that when I r X ; Z s is sufficiently large, the term log 2 p I r X ; Z s ´ log 2 J ` E X r ϵ s ` 1 q ` 4 will be negligible. However, without further information on Q and P , it is difficult to assert how small E X r ϵ s is. We can view ϵ as the extra cost introduced by the space partitioning algorithm. Fortunately, for commonly used distributions in neural compression, including Uniform and Gaussian, we have the following conclusion under reasonable assumptions:

Proposition 3.2 (Bound of ϵ for Uniform and Gaussian) . Assume setting J ď 2 D KL r Q } P s when running Algorithm 2 for each Q Z | X . Then, for Uniform Q and P , if Q ! P (i.e., Q is absolute continuous w.r.t P ), we have ϵ ' 0 ; for factorized Gaussian Q and P , if Q has smaller variance than P along each axis, we have ϵ ď 0 . 849 a D KL r Q } P s , and E X r ϵ s ď 0 . 849 a I r X ; Z s .

We prove Proposition 3.2 in Appendix C.2. Also, we highlight that the conclusion for Gaussian in Proposition 3.2 is derived by considering the worst case when P has the same variance as Q along all dimensions. In practice, this worst case can barely happen since in neural compression P represents the prior, and Q represents the posterior, which will be more concentrated than the prior. Empirically, we find this ϵ -cost yields no visible influence on the codelength (e.g., Figure 2b in Section 5).

3.3 Generality of our Space Partitioning Algorithm

We highlight that the conclusions in Section 3.2 are independent of the partitioning strategy and π . This allows us to use different π without the need to revisit the bound of the codelength.

More interestingly, we do not need to partition the space with non-overlapping bins. This is because the density of P 1 , as stated in Equation (7), can be directly interpreted as a mixture of priors with J components, where π p j q represents the mixture weights. There are no restrictions preventing this mixture model from having overlapping components. We provide more explanation and discussion on overlapping components in Appendix B.1. As an extreme case, we can even have entirely overlapping bins. Notably, this scenario coincides with the 'parallel threads' version of REC proposed by Flamich (2024). Concretely, Flamich (2024) proposed to initiate several threads for a single REC task and run them in parallel on various machines, thereby decreasing the runtime.

Furthermore, the concept of space partitioning and its codelength, as stated in Theorem 3.1, extends beyond the scope of Poisson functional representation algorithms. Here, we state the codelength by applying our proposed space partitioning method to greedy Poisson rejection sampling (GPRS) (Flamich, 2024) and leave its application to other REC algorithms for future work.

Theorem 3.3. Let a pair of correlated random variables X , Z ' P X , Z be given. Assume we perform relative entropy coding using GPRS with space partitioning and let j ˚ denote the bin index and ˜ n ˚ the local sample index returned by the algorithm, and ϵ be as in Equation (9) . Then,

We prove Theorem 3.3 in Appendix C.3. Additionally, since we can view the 'parallel threads' version of GPRS as a special case of our proposed method, Theorem 3.3, and hence Proposition 3.2, offer alternative bounds for Theorem 3.5 in Flamich (2024).

4 Exemplifying the choice of Partitioning Strategy and π

In the above sections, we do not specify the partitioning strategy and π . In this section, we will exemplify their choices. To keep our discussion simple, we take the assumption that both Q and P are fully-factorized distributions. This covers a majority of neural compression applications with relative entropy coding, including both V AE-based (Flamich et al., 2020) and INR-based codecs (Guo et al., 2024; He et al., 2023). Under this assumption, we adopt a simple partitioning strategy to split space with axis-aligned grids, which allows us to draw samples and evaluate the density ratio easily.

Now, let's focus on π . We consider two scenarios: 1) the sender aims to encode a sample that exactly follows Q , and 2) the sender encodes a sample following Q approximately for a more manageable computational cost. As we will see, the optimal choice of π differs in these cases.

We note that the latter garners more interest in practical neural compression. This is because we can always construct an example where the standard PFR algorithm does not terminate while the first sample from the prior already has little bias. As an example, take q p z q ' N p z | 0 . 001 , σ 2 q , and p p z q ' N p z | 0 , 1 q . When σ Õ 1 (approaching 1 from below), the expected runtime for PFR diverges: 2 D 8 r Q || P s Ñ8 . Unfortunately, our proposed space partitioning approach can do little in this case. This is because, for a small ϵ -cost, as stated in Proposition 3.2, the number of partitions J should be less than 2 D KL r Q } P s , which reduces to 1 when σ Õ 1 . For multi-dimensional (factorized) Q and P , this issue will occur whenever it arises in any single dimension. Making things even worse, this example is pervasive in neural compression due to numerical inaccuracies. Therefore, limiting the number of candidate samples is more practical than running the PFR algorithm until the stopping criterion is met. Consequently, we will mainly study the non-exact case in the following, but for the sake of completeness, we first discuss the exact scenario.

4.1 Exact Sampler

When encoding a sample following Q exactly, we hope to minimize the expected runtime by adjusting the search heuristic P 1 . Since we can view our proposed algorithm as running standard PFR with P 1 , its expected runtime is 2 D 8 r Q || P 1 s . Thus, we have the following constrained optimization problem:

The solution turns out to be intuitive:

However, sampling from this π is generally challenging. Fortunately, if Q and P are both factorized and we partition the space with axis-aligned grids, then π factorizes dimensionwise into a product of categorical distributions. Also, this choice of π simplifies the evaluation of the stopping criterion in Equation (6). Specifically, r 1 max simplifies to r 1 max ' Z ¨ J , where Z ' ř J j ' 1 π p j q denotes the normalization constant. This constant is shared in both the maximum density ratio r 1 max and the density ratio for individual samples. We thus can omit it when evaluating the stopping criterion. We detail the procedure in Algorithm 3 in Appendix A.1.

4.2 Non-exact Sampler

If we only need to encode a sample following Q approximately, we can apply our space partitioning strategy to ordered random coding (ORC). In this case, we hope to reduce the bias with the same sample size or reduce the sample size for the same bias. To determine the optimal choice of π , we state the following corollary on ORC's bias. This is a corollary of Theis and Yosri (2022, Lemma D.1) and Chatterjee and Diaconis (2018, Theorem 1.2). We present the proof in Appendix C.4.

Corollary 4.1 (Biasness of sample encoded by ORC) . Given a target distribution Q and a prior distribution P , run ordered random coding (ORC) to encode a sample. Let ˜ Q be the distribution of encoded samples. If the number of candidates is N ' 2 D KL r Q } P s` t for some t ě 0 , then

Conversely, supposing that N ' 2 D KL r Q } P s´ t for some t ě 0 , then

This corollary tells us that when running ORC with target Q and prior P , if the density ratio between Q and P is well concentrated around its expectation, choosing N « 2 D KL r Q } P s candidates is both sufficient and necessary to encode a low-bias sample in terms of total variation (TV) distance . Recall that our proposed algorithm can be viewed as running ORC with the search heuristic P 1 as the prior. We, therefore, want to choose π to minimize the KL-divergence between the target Q and the search heuristic P 1 . The optimal π turns out to be surprisingly simple:

This is because D KL r Q } P 1 s ' D KL r Q } P s ´ ř J j ' 1 Q p B j q log 2 π p j q ` ř J j ' 1 Q p B j q log 2 P p B j q , and the cross-entropy p´ ř J j Q p B j q log 2 π p j qq takes its minimum when π matches Q . In practice, to sample from this categorical distribution, we can simply draw z ' Q and find the bin it belongs to. However, choosing π p j q ' Q p B j q complicates the calculation of r 1 max in Equation (6). Fortunately, in ORC, we do not need to check the stopping criterion, so this complication does not pose an issue. We formalize this new ORC algorithm in Algorithm 4 in Appendix A.2.

Algorithm 4 still leaves three questions unanswered: first, we need to determine the number of partitions J . As a finer partition allows better alignment between Q and P 1 , we pick J ' 2 t D KL r Q } P s u , the maximum value for which Proposition 3.2 provides a bound on the extra coding cost ϵ . Note that this will require the receiver to know t D KL r Q } P s u . As we will demonstrate in Section 5, in practice, we can achieve this by either encoding the KL using negligible bits (e.g., in neural compression with VAE) or enforcing the KL budget during optimization (e.g., in neural compression with INRs).

Second, as we partition the space using axis-aligned grids, we need to determine the number of bins assigned to each axis . To explain why this choice matters, we consider an example where Q and P share the same marginal distribution along a specific axis, and the space is partitioned solely by dividing this axis; in this case, we have P 1 ' P , leading to no improvement in runtime. Fortunately, in most neural compression applications, the sender and receiver can also have access to an estimation of the mutual information I d along each axis d from the training set. If the mutual information is non-zero along one axis, on average, Q and P will not have the same marginal distribution along this axis. Based on this observation, we suggest partitioning the d -th axis into approximately 2 n d intervals where n d ' D KL r Q } P s ¨ I d L ř D d 1 ' 1 I d 1 ; see Appendix B.2 for further discussion, including the derivation, a toy example illustration, and ablation studies on this.

Third, we determine how many candidate samples we need to draw from P 1 to ensure the encoded sample has low bias. Recall that our proposed algorithm can be viewed as running ORC with P 1 as the prior, we thus require a sample size « 2 D KL r Q } P 1 s according to Corollary 4.1. When the total number of partitions J ' 2 D KL r Q } P s , we find D KL r Q } P 1 s ' ´ ř J j ' 1 Q p B j q log 2 Q p B j q . We do not have an analytical form for this value, but we can estimate it by samples from Q . In fact, we empirically find a small number of samples to be sufficient for an accurate estimator.

Why choose the TV distance as the measure of approximation error in Corollary 4.1? Indeed, TV distance is not the only possible metric and the choice should align with the intended application. We choose total variation as it fits well with the task in our experiments: the one-shot nature of data compression for human consumption. We will explain this in two parts:

Control of the TV distance in the latent space implies control in data space: naturally, we wish to assess reconstruction quality in data space, but we use REC only to encode latent variables from which we reconstruct the data. However, since the total variation satisfies the data processing and triangle inequalities, if our generative model approximates the true data distribution with δ total variation error and we use an ϵ -approximate REC algorithm, then encoding the latents incurs no more than ϵ ` δ total variation error in the data space (Flamich and Wells, 2024).

TV distance captures a reasonable notion of realism. Imagine two distributions Q (e.g., the ground truth data distribution) and ˜ Q (e.g., the data distribution learned by our compressor). Let x 0 ' Q, x 1 ' ˜ Q and let B ' Bern p 1 { 2 q be a fair coin toss. Then, the probability that any observer can correctly decide which distribution z B was sampled from, i.e., correctly predict the value of B given x B , is at most 1 { 2 ¨ p 1 ` D TV r ˜ Q,Q sq (Nielsen, 2013; Blau and Michaeli, 2018). Thus, in the context of approximate REC, the TV distance bounds the accuracy with which any observer can tell the compressed and reconstructed data apart from the original.

5 Experiments

We now verify our proposed algorithm with three different experiments, including synthetic toy examples, lossless compression on MNIST with VAE, and lossy compression on CIFAR-10 with INRs. We include details of the experiment setups in Appendix D.

Toy Experiments. We explore the effectiveness of our algorithm on 5D synthetic Gaussian examples. We run PFR with space partitioning (Algorithm 3) to encode exact samples and ORC with space partitioning (Algorithm 4) for approximate samples. We show the results in Figure 2 and Figure 3, and also include standard PFR and ORC for comparison. We can see that our space partitioning algorithm reduces the runtime by up to three orders of magnitude while maintaining codelength when encoding exact samples, and requires a much smaller sample size to achieve a certain bias (quantified by maximum mean discrepancy (MMD, Smola et al., 2007)) when encoding approximate samples.

Lossless Compression on MNIST with VAE. As a further proof of concept, we apply our methods to losslessly compress MNIST images (LeCun and Cortes, 1998). We train a VAE following Flamich et al. (2024), employ REC (specifically, ORC with our proposed space partitioning) to encode the latent embeddings, and entropy-encode the image. The KL divergence of entire latent embeddings averages over 90 bits. Unfortunately, even after employing our proposed space partitioning algorithm, the new KL divergence D KL r Q } P 1 s exceeds our manageable size. We hence randomly divide the latent dimensions into smaller blocks. Recall that, in our proposed approach, the sender and receiver need to share t D KL r Q } P s u so that they can partition the space into the same J ' 2 t D KL r Q } P s u bins. However, this value varies across different images. Therefore, we estimate the distribution of t D KL r Q } P s u for each block from the training set, and entropy code it for each test image.

We evaluate the performance achieved by t 2 , 4 u blocks and different sample sizes in Table 1, with the theoretically optimal codelength and the results by GPRS, a REC algorithm that is faster in 1D space but incurs coding overhead dimension-wise (Flamich, 2024). Notably, our algorithm's codelength is only 2% worse than the theoretically optimal result and about 6% better than GPRS's. We also investigate the reasons for overhead. Compared to the ELBO, our method incurs overhead in three ways: (a) the cost to encode t D KL r Q } P s u ; (b) the overhead from encoding the latent embeddings by REC; and (c) the overhead when encoding the target image, which arises from the bias in encoding the latent embeddings, as ORC encodes only approximate samples. We can see our algorithm achieves extremely low bias while maintaining a relatively small overhead caused by (a) and (b).

Lossy Compression on CIFAR-10 with INRs. We apply our methods to a more practical setting: compressing CIFAR-10 images with RECOMBINER (He et al., 2023), an implicit neural representation (INR)-based codec. The authors of RECOMBINER originally encoded INR weights using a block size of D KL r Q } P s ' 16 bits. We empirically find that our proposed method can handle a block size of D KL r Q } P s ' 48 bits while maintaining D KL r Q } P 1 s within a manageable range, approximately 12-14 bits. To further reduce the bias of the encoded sample, we opt to use 2 16 samples

5

10

15

D

[

Q

||

P

]

10 2

10 4

10 6

Steps

standard PFR

Space partition PFR

(a) Runtime w.r.t

D

8

r

Q

||

P

s

.

2.5

5.0

7.5

10.0

I

[

X

;

Z

]

0

10

20

Codelength

standard PFR

Space partition PFR

I

[

X

;

Z

]

(b) Codelength w.r.t

I

r

X

;

Z

s

.

2

4

6

8

DKL

[

Q

||

P

]

0

200

400

runtime

(#samples)

MMD < 0.006

MMD < 0.009

MMD < 0.012

Standard ORC

Space partition ORC

1.0

1.5

2.0

2.5

3.0

Bitrate (bit per pixel)

27.5

30.0

32.5

PSNR (dB)

w/o finetuning

0.84dB

0.32 bpp

Space partition ORC (block size=48 bits)

Standard ORC (block size=16 bits)

Theorical RD

1.0

1.5

2.0

2.5

3.0

3.5

Bitrate (bit per pixel)

27.5

30.0

32.5

35.0

w. finetuning

0.36dB

0.21 bpp

for each block. Besides, unlike in the V AE case, where the KL for each block varies across different test images, here, we follow He et al. (2023) to enforce the KL of all blocks to be close to 48 bits when training the INR for each test image. This eliminates the need to encode t D KL r Q } P s u .

Additionally, He et al. (2023) enhanced their results by fine-tuning Q for the not-yet-compressed weights after encoding each block, which can achieve a more complicated posterior distribution Q in an auto-regressive manner. However, the effectiveness of fine-tuning is closely tied to the number of blocks. As we have reduced the number of blocks in our approach, fine-tuning becomes less effective. Therefore, we present results both with and without fine-tuning in Figure 4. We also present the theoretical RD curve (without fine-tuning) to evaluate how close we are to the theoretical bound.

As we can see, compared with standard ORC using a block size of D KL r Q } P s ' 16 bits, our proposed algorithm with the block size of D KL r Q } P s ' 48 bits reduces the codelength by approximately 10% with fine-tuning and 15% without. This gain is due to three reasons: 1) when the KL of each block is larger, the ϵ -cost, the algorithmic and constant overhead in Equation (8) becomes negligible; 2) the ´ log 2 J term in Equation (8) also reduces the algorithmic overhead in our method's codelength comparing with Equation (1); and 3) as the new KL divergence D KL r Q } P 1 s is around 12-14 bits and we choose a sample size of 2 16 , we eliminate most of the bias in our encoded samples. Moreover, the algorithm's codelength remains within 2% of the theoretical optimal result. As a further comparison, we present the performance of other compression baselines in Figure 12, where we can see that our proposed algorithm makes RECOMBINER more competitive.

6 Related Works

There are several efforts to accelerate REC algorithms. Flamich et al. (2022, 2024); Flamich (2024) leveraged the idea of partitioning in 1D space, achieving an impressive runtime in the order of O p D 8 r Q || P sq or O p D KL r Q } P sq . However, these fast approaches only work for 1D distributions. Besides, Flamich and Theis (2023) proposed bits-back quantization (BBQ), which encodes a sample with time complexity linear in dimensionality. However, BBQ assumes P and Q to be uniform distributions in two hypercubes with their edges aligned, which may not be practical in many applications. The algorithm most similar to ours is hybrid coding (Theis and Yosri, 2022), which employs dithered quantization followed by a sampling procedure, an approach that resembles a special case of our proposed algorithm. However, hybrid coding relies on the assumption that the support of Q is contained within a hypercube, restricting its practical applicability.

7 Conclusions and Limitations

In this work, we propose a relative entropy coding (REC) scheme based on space partitioning, which significantly reduces the runtime in practical settings. We provide both theoretical analysis and experimental evidence supporting our method. Being among the few successful attempts to accelerate REC for practical settings, we firmly believe that our contributions will broaden the application of REC methods and inspire further research.

However, our proposed algorithm still faces several limitations at the current stage. First, although our method and conclusion apply to general partitions, in practice, we are largely limited to handling axis-aligned grids. Second, in our experiments, we require the mutual information to be factorized per dimension and shared between the sender and receiver, which restricts our algorithm's utility, for example, in non-factorized cases (Theis et al., 2022) or one-shot REC tasks (Havasi et al., 2019). A potential avenue for future research involves employing a mixture of priors to form 'partitions' without hard boundaries as discussed in Section 3.3. Additionally, by considering the prior as a convolution of two distributions, we may potentially create infinitely many partitions. However, managing the codelength in such cases may raise new challenges.

8 Acknowledgments

We would like to thank Zongyu Guo and Sergio Calvo-Ordoñez for their proofreading and insightful feedback on the manuscript. JH was supported by the University of Cambridge Harding Distinguished Postgraduate Scholars Programme. JMHL and JH acknowledge support from a Turing AI Fellowship under grant EP/V023756/1; GF acknowledges funding from DeepMind.

References

Ballé, V. Laparra, and E. P. Simoncelli. End-to-end optimized image compression. In 5th International Conference on Learning Representations, ICLR 2017 , 2017.

Ballé, D. Minnen, S. Singh, S. J. Hwang, and N. Johnston. Variational image compression with a scale hyperprior. In International Conference on Learning Representations , 2018.

Ballé, P. A. Chou, D. Minnen, S. Singh, N. Johnston, E. Agustsson, S. J. Hwang, and G. Toderici. Nonlinear transform coding, 2020.

Bellard. BPG image format. https://bellard.org/bpg/ , 2014. Accessed: 2023-09-27.

Blau and T. Michaeli. The perception-distortion tradeoff. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 6228-6237, 2018.

Chatterjee and P. Diaconis. The sample size required in importance sampling. The Annals of Applied Probability , 28(2):1099-1135, 2018.

Cheng, H. Sun, M. Takeuchi, and J. Katto. Learned image compression with discretized gaussian mixture likelihoods and attention modules. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 7939-7948, 2020.

Dupont, H. Loya, M. Alizadeh, A. Goli´ nski, Y. W. Teh, and A. Doucet. Coin++: Neural compression across modalities, 2022.

Flamich. Greedy poisson rejection sampling. Advances in Neural Information Processing Systems , 36, 2024.

Flamich and L. Theis. Adaptive greedy rejection sampling. In 2023 IEEE International Symposium on Information Theory (ISIT) , pages 454-459. IEEE, 2023.

Flamich and L. Wells. Some notes on the sample complexity of approximate channel simulation. In First 'Learn to Compress' Workshop@ ISIT 2024 , 2024.

Flamich, M. Havasi, and J. M. Hernández-Lobato. Compressing images by encoding their latent representations with relative entropy coding. Advances in Neural Information Processing Systems , 33:16131-16141, 2020.

Flamich, S. Markou, and J. M. Hernández-Lobato. Fast relative entropy coding with a* coding. In International Conference on Machine Learning , pages 6548-6577. PMLR, 2022.

Flamich, S. Markou, and J. M. Hernández-Lobato. Faster relative entropy coding with greedy rejection coding. Advances in Neural Information Processing Systems , 36, 2024.

Guo, G. Flamich, J. He, Z. Chen, and J. M. Hernández-Lobato. Compression with bayesian implicit neural representations. Advances in Neural Information Processing Systems , 36, 2024.

Havasi, R. Peharz, and J. M. Hernández-Lobato. Minimal random code learning: Getting bits back from compressed model parameters. In International Conference on Learning Representations , 2019.

He, G. Flamich, Z. Guo, and J. M. Hernández-Lobato. Recombiner: Robust and enhanced compression with bayesian implicit neural representations. In The Twelfth International Conference on Learning Representations , 2023.

JVET. VVC offical test model. https://jvet.hhi.fraunhofer.de , 2020. Accessed: 2024-0305.

P. Kingma and J. Ba. Adam: A method for stochastic optimization, 2017.

P. Kingma and M. Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images.(2009), 2009.

LeCun and C. Cortes. The mnist database of handwritten digits. 1998. URL http://yann. lecun.com/exdb/mnist/ .

T. Li and A. El Gamal. Strong functional representation lemma and applications to coding theorems. IEEE Transactions on Information Theory , 64(11):6967-6978, 2018.

J. Maddison. A poisson process model for monte carlo. Perturbation, Optimization, and Statistics , pages 193-232, 2016.

J. Maddison, D. Tarlow, and T. Minka. A* sampling. Advances in neural information processing systems , 27, 2014.

Mentzer, E. Agustsson, M. Tschannen, R. Timofte, and L. Van Gool. Practical full resolution learned lossless image compression. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2019.

Nielsen. Hypothesis testing, information divergence and computational geometry. In International Conference on Geometric Science of Information , pages 241-248. Springer, 2013.

Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research , 12:2825-2830, 2011.

R. Schwarz, J. Tack, Y. W. Teh, J. Lee, and J. Shin. Modality-agnostic variational compression of implicit neural representations. In Proceedings of the 40th International Conference on Machine Learning , pages 30342-30364, 2023.

Smola, A. Gretton, L. Song, and B. Schölkopf. A hilbert space embedding for distributions. In International conference on algorithmic learning theory , pages 13-31. Springer, 2007.

Theis and E. Agustsson. On the advantages of stochastic encoders, 2021.

Theis and N. Yosri. Algorithms for the communication of samples. In International Conference on Machine Learning , pages 21308-21328. PMLR, 2022.

Theis, T. Salimans, M. D. Hoffman, and F. Mentzer. Lossy compression with gaussian diffusion. arXiv preprint arXiv:2206.08889 , 2022.

Townsend, T. Bird, and D. Barber. Practical lossless compression with latent variables using bits back coding. In 7th International Conference on Learning Representations, ICLR 2019 , volume 7. International Conference on Learning Representations (ICLR), 2019.

A Algorithms

A.1 Practical PFR Algorithm with Space Partitioning

Algorithm 3 Practical encoding procedure of PFR with dimension-wise space partitioning

Input: Fully-factorized target Q ' Q r 1 s ˆ Q r 2 s ˆ¨¨¨ ˆ Q r D s , shared and fully-factorized coding distribution P ' P r 1 s ˆ P r 2 s ˆ¨¨¨ ˆ P r D s , and shared random states S j for each partition j . Output: Bin index j ˚ , local sample index ˜ ˚ .

n # initialize: τ ˚ Ð8 , t 0 Ð 0 , j ˚ Ð 0 , ˜ n ˚ Ð 0 . Partition d -th dimension into J r d s intervals B r d s 1 , B r d s 2 , ¨ ¨ ¨ , B r d s J r d s , s.t., J Ð ś D d ' 1 J r d s . ˜ n j Ð 0 for j ' 1 , 2 , ¨ ¨ ¨ , J. r 1 max Ð 1 . Ź Note that we omit the normalization factor and J # pre-process for fast sampling along each dimension: for d ' 1 , 2 , ¨ ¨ ¨ , D do for i ' 1 , 2 , ¨ ¨ ¨ , J r d s do π r d s i Ð sup z P B r d s i ´ q r d s p z q p r d s p z q ¯ . end for Z Ð ř J r d s i ' 1 π r d s i . Define π r d s ' Categorical ´ π r d s 1 { Z , π r d s 2 { Z , ¨ ¨ ¨ , π r d s J r d s { Z ¯ . end for # run PFR: for n ' 1 , 2 , ¨ ¨ ¨ do # sample time: Sample ∆ t n ' Exp p 1 q . t n Ð t n ´ 1 ` ∆ t n . # sample a bin by sampling intervals per dimension: j n Ð 0 . Ź j n keeps track of the bin index. ℓ n Ð 1 . Ź ℓ n keeps track of the likelihood of the bin being sampled. for d ' 1 , 2 , ¨ ¨ ¨ , D do k ' π r d s . ℓ n Ð ℓ n ¨ π r d s k . if d ' 0 then j n Ð k . else j n Ð k ` j n ¨ J r d ´ 1 s . end if end for # simulate a sample with PRNG: ˜ n j n Ð ˜ n j n ` 1 . z n Ð PRNG p P | B j n , S j n , ˜ n j n q . Ź by inverse transform sampling when both of the partitioning and P are axis-aligned. # update τ ˚ : τ n Ð ℓ n ¨ t n ¨ p p z n q{ q p z n q . Ź Note that we omit the normalization factor and J . if τ n ď τ ˚ then τ ˚ Ð τ n , j ˚ Ð j n , ˜ n ˚ Ð ˜ n j n . end if # check stopping criterion: if t n { r 1 max ą τ ˚ then break end if end for

A.2 Practical ORC Algorithm with Space Partitioning

Algorithm 4 Practical encoding procedure of ORC with Space Partitioning

Input: Target Q , shared coding distribution P , number of candidate samples N , and shared random states S j for each partition j .

Output: Bin index j ˚ , local sample index ˜ n ˚ .

# sample time:

# simulate a sample with PRNG:

z

'

Q

, and find

j

n

z

P

B

j

n

.

Ź by inverse transform sampling when both of the partitioning and P are axis-aligned.

# update

τ

˚

:

end if end for

A.3 Decoding Algorithm

Algorithm 5 Dncoding procedure of PFR/ORC with Space Partitioning

Input: Bin index j ˚ , local sample index ˜ n ˚ , and shared random states S j for each partition j . Output: Sample z .

Partition space in J 'bins'. z Ð PRNG p P | B j ˚ , S j ˚ , ˜ n ˚ q .

B Additional Results and Discussions

B.1 Elucidating Non-overlapping Bins

As we discussed in the main text, the partitioning does not need to form non-overlapping grids. In fact, there are only two constraints for the partitions: (1) the density of the distributions for all 'bins', without being weighted by π , should sum to p p z q at each z , and (2) integrating the density of z in each 'bin' should yield the same constant 1 { J . As long as we follow these two constraints, we can create any kind of partition, no matter whether they overlap or not. In Section 3.3, we explain this by viewing Equation (7) as a mixture model. In this section, we provide another explanation from an intuitive point of view.

First, we assume the reader agrees that we can always create non-overlapping bins as non-overlapped grids. Then, to create overlapping bins, we can add an auxiliary axis x aux to the original space Ω , forming an augmented space as x aux ˆ Ω . We define the prior and target in the auxiliary axis as P aux and Q aux, and define the prior and target in the augmented space as P aug ' P aux ˆ P , Q aug ' Q aux ˆ Q . This augmentation will not influence the codelength since we can always ensure the augmented P aug and Q aug to have identical marginal densities along the auxiliary axis. We illustrate this augmented

Ź Note the difference from the PFR algorithm.

space in Figure 5a. We can view partitions in the original space as non-overlapping ones in an augmented space, as illustrated in Figures 5b to 5d.

In this augmented space, the partitions show follow these two constraints: (1) without considering π , if we marginalize our the auxiliary axis, P aug should revert to the original P ; (2) we require all bins to have the same probability under the prior P aug. The second constraint comes from the first step in Section 3.1, which is necessary for the proof of the codelength later in Equation (38). If we only consider the original space, these two constraints correspond to the ones aforementioned at the beginning of this section.

Ω

𝑥aux

𝑝

𝑞

𝑞aug

𝑝aug

a

ugmented space = Ω

×

𝑥aux

(a) The augmented space.

Ω

𝑥aux

𝑝

𝑞

𝑞aug

𝑝aug

a

ugmented space = Ω

×

𝑥aux

(b) Non-overlapping bins.

Ω

𝑥aux

𝑝

𝑞

𝑞aug

𝑝aug

a

ugmented space = Ω

×

𝑥aux

(c) Fully-overlapping bins.

Ω

𝑥aux

𝑝

𝑞

𝑞aug

𝑝aug

a

ugmented space = Ω

×

𝑥aux

(d) General partitions.

B.2 Choosing the Number of Intervals Assigned to Each Axis

First, we elaborate on why this question matters. We take the following example: Q and P are 2D Gaussians, and they have the same marginal distribution along the first dimension, denoted as x 1 , but have different marginal distributions along the second dimension, denoted as x 2 .

Now, we partition the space x 1 ˆ x 2 by axis-aligned grids, i.e., the boundaries of grids are parallel to the axes. If we partition the space by solely dividing x 1 , we will find that the prior and the posterior have the same probability for each bin, i.e., P p B j q ' Q p B j q , @ j . In this case, choosing

π p j q ' Q p B j q as discussed in Section 4.2, we can write the density of P 1 (Equation (7)) as

which means that the adjusted search heuristic P 1 is the same as the original prior P , and thus we will have no improvements in the runtime. Figure 6a visualize this example.

On the other hand, if we partition the space by dividing x 2 , as shown in Figure 6b, the search heuristic P 1 will align with Q better than P , and hence we can reduce the runtime. In practice, it is not always feasible to partition in this way. Instead, we may partition the space by dividing both x 1 and x 2 , as shown in Figure 6c. In this case, we will still have reduced runtime. This is because

¯

We recognize the term ´ ř J j ' 1 Q p B j q log 2 Q p j q ´ ř J j ' 1 Q p B j q log 2 P p B j q as a KL divergence and hence will be non-negative. Whene there exists B j , s.t. Q p B j q ‰ P p B j q , the KL term ´ ř J j ' 1 Q p B j q log 2 Q p j q ´ ř J j ' 1 Q p B j q log 2 P p B j q ¯ will always be positive to ensure D KL r Q } P 1 s ď D KL r Q } P s and hence guarantee a reduction in the runtime.

𝑥2

𝑥1

𝑝

𝑞

(a)

𝑥2

𝑥1

𝑝

𝑞

(b)

𝑥2

𝑥1

𝑝

𝑞

(c)

From this example, we conclude that even when the total number of partitions is fixed, different partition strategies will lead to different runtime. What we want to avoid is the case in Figure 6a. Fortunately, in most neural compression applications, we have access to an estimation of the mutual information I d along each axis d from the training set. For the reader who is not familiar with the term, we can view I d as the average KL divergence along d -th axis. If the mutual information is non-zero along one axis, on average, Q and P will not have the same marginal distribution along this axis. Therefore, we can partition the d -th axis into approximately 2 ^ p I d ¨ D KL r Q } P s { ř D d 1 ' 1 I d 1 q intervals to avoid the worst case in Figure 6a.

B.2.1 Examples and Ablations on Choosing the Number of Intervals Assigned to Each Axis

Qualitative visualization . We first visualize the approximation error of using ORC with different partitioning strategies on a toy problem. Specifically, we create a 5D Gaussian target, and we set dimension 1 to have 0 mutual information (we call it a collapsed/uninformative dimension). We compare three partition strategies: only partitioning the collapsed dimension, randomly assigning intervals to dimensions, and assigning intervals according to mutual information (our proposed strategy). We also include standard ORC for reference. We run the four settings with the same number of samples (20) and repeat each setting 5000 times. We show the histogram of 5000 encoded results in Figure 7. As we can see, assigning intervals according to mutual information works best, whereas partitioning only the collapsed dimension yields almost the same results as standard ORC. This verifies our discussion above.

2

1

0

1

0

2

density

Standard ORC

2

1

0

1

0

2

Only partition collapsed dim

2

1

0

1

0

2

Randomly assign intervals per dim

2

1

0

1

0

2

Assign intervals according to MI (ours)

Empirical density of the encoded samples (average acorss 5000 runs)

target density

Ablation study . We now provide ablation studies showing how partition strategies will influence performance. We run the ablation on the Gaussian toy example and the CIFAR-10 compression experiments. As we can see, our proposed partition strategy is always better than randomly assigning intervals per dim. For CIFAR-10, we also compare the results by first removing uninformative dimensions (mutual information 0), and then randomly assigning intervals to other dimensions. We find this is only slightly worse than our proposed partition strategy. This not only further verifies our discussion above in Appendix B.2, but also indicates that our algorithm is not very sensitive to how we construct the intervals as long as we avoid partitioning along uninformative dimensions.

s

4

6

8

10

12

14

16

D

[

Q

||

P

]

10 1

10 2

10 3

10 4

10 5

10 6

Steps

Space Partitioning PFR (assign intervals according to MI, ours)

Space Partitioning PFR (randomly assign intervals per dim)

Standard PFR

(a) PFR's runtime w.r.t

D

8

r

Q

||

P

s

1

2

3

4

5

6

7

8

DKL

[

Q

||

P

]

0

50

100

150

200

250

300

350

Sample Size

Space Partitioning PFR (assign intervals according to MI, ours)

Space Partitioning PFR (randomly assign intervals per dim)

Standard ORC

MMD < 0.006

MMD < 0.009

MMD < 0.012

(b) ORC's runtime for a certain bias w.r.t

D

KL

r

Q

||

P

1.0

1.5

2.0

2.5

3.0

3.5

Bitrate (bit per pixel)

28

30

32

34

36

PSNR (dB)

Space partition ORC (assign intervals according to MI, ours)

Space partition ORC (random assign interval per dim)

Space partition ORC (remove collapsed dim, then random assign interval per dim)

Standard ORC

theorical RD

C Proofs and Derivations

C.1 Proof of Theorem 3.1

Restate Theorem 3.1 for easy reference:

Theorem 3.1. Let a pair of correlated random variables X , Z ' P X , Z be given. Assume we perform relative entropy coding using Algorithm 2 and let j ˚ denote the bin index and ˜ n ˚ the local sample index returned by the algorithm. Then, the entropy of the two-part code is bounded by:

Proof. The proof will be organized as follows: first, we will prove an upper bound of E r ˜ n ˚ | X ' x s for the one-shot case, i.e., the expectation of ˜ n ˚ given a realization of the data X . Then, we will take the expectation over X to achieve an upper bound for the average case E r ˜ n ˚ s . Finally, we will consider encoding ˜ n ˚ using a Zipf distribution to calculate the upper bound given in Theorem 3.1.

Derive E r ˜ n ˚ | X ' x s . Assume that the posterior distribution Z | X ' x is Q . We run PFR with space partitioning according to Algorithm 2 using a coding distribution P to encode a sample from Q . As we discussed in the main text, we can view our proposed algorithm as running standard PFR with an adjusted prior P 1 , whose density is defined as

Therefore, we follow similar arguments as Li and El Gamal (2018) to prove the codelength.

We first assume the minimal τ is τ ˚ , the corresponding sample is z ˚ , and the global sample index is n ˚ . WLG, let's assume the sample falls in bin j ˚ and its local sample index in this bin is ˜ n ˚ .

According to Li and El Gamal (2018), we have

Equation (28) is obtained from Appendix A in Li and El Gamal (2018) and replacing p by p 1 . Additionally, we have

This is because: if we define 'success' as a sample falling in the bin j ˚ , then the probability that the sample in the bin j ˚ with global sample index n ˚ has a local sample index ˜ n ˚ is equivalent to the probability that there are ˜ n ˚ ´ 1 success in the first n ˚ ´ 1 Bernoulli trials.

Therefore, according to the property of Poisson distribution, we have

Therefore, we have

The last line is because we require P p B 1 q ' P p B 2 q ' ¨ ¨ ¨ ' P p B J q ' 1 { J . We can see that Equation (38) does not depend on either the specific space partitioning strategy or the choice of π .

Based on Taylor expansion of log 2 p 1 ` x q at 0 , `8 , we have

Therefore,

E log 2 r ˜ n ˚ | X ' x s ď E z ' Q log 2 ' q p z q Jp p z q ` 1 ȷ (40) ď E z ' Q log 2 ' ✶ " q p z q p p z q ď J *ˆ q p z q Jp p z q ˙ log 2 e ȷ ` E z ' Q log 2 ' ✶ " q p z q p p z q ą J *ˆ log 2 q p z q Jp p z q ` Jp p z q q p z q log 2 e ˙ȷ (41) ' E z ' Q log 2 ' ✶ " q p z q p p z q ą J *ˆ log 2 q p z q Jp p z q ˙ȷ ` log 2 e ¨ E z ' Q log 2 ' ✶ " q p z q p p z q ď J *ˆ q p z q Jp p z q ˙ ` ✶ " q p z q p p z q ą J *ˆ Jp p z q q p z q ˙ȷ (42) ď E z ' Q log 2 ' ✶ " q p z q p p z q ą J *ˆ log 2 q p z q Jp p z q ˙ȷ ` log 2 e (43) The last line is because

We can further simplify Equation (43) as follows:

where we denote

Derive E r ˜ n ˚ s . Now, we consider the average case by taking the expectation over X . We have

Derive entropy. Following the maximum entropy argument similar to Li and El Gamal (2018), we have

Finally, we need log 2 J bits to encode the index j ˚ . Therefore, the two-part code's codelength is upper-bounded by

We also note that, if we do not take the expectation over X , we can also achieve a similar bound for the one-shot channel simulation task targeting Q using a coding distribution P :

This ensures that the codelength is consistent for each realization of X and guarantees a good encoding performance both on average and for each certain data observation in practice.

C.2 Proof of Proposition 3.2

To prove Proposition 3.2, we first prove the following Lemma:

Lemma C.1. The ϵ -cost defined in Equation (9) monotonically increases w.r.t J .

Proof. By setting r ' log 2 q p z q { p p z q which is a new random variable, we rewrite the ϵ -cost as

where we also write J ' log 2 J , which is a monotonically increasing function of J . Take derivative w.r.t. J , we have

which finishes the proof.

Therefore, to prove Proposition 3.2, we only need to look at the worst case, where J ' 2 D KL r Q || P s . We now prove for Uniform and Gaussian, respectively:

Uniform. For uniform distribution, WLG, assume P ' U p 0 , 1 q n , and Q ' U p A q , where A Ď p 0 , 1 q n . In this case, we have D KL r Q || P s ' ´ log 2 µ p A q , and q p z q p p z q ' 1 µ p A q , @ z P A , where µ is the Lesbague measure. Therefore,

On the other hand, the indicator function ensures the integrand is non-negative, and hence ϵ ě 0 . Therefore, ϵ ' 0 .

Factorized Gaussian. To prove for Gaussian, we first prove the following lemma:

Lemma C.2. When J ' 2 D KL r Q || P s , the ϵ -cost defined in Equation (9) is upper-bounded by

where r is the RV defined by r ' log 2 q p z q { p p z q , z ' Q .

Proof. First, we know

and we recognize

When J ' 2 D KL r Q || P s , we have

Therefore, we have

Writing r ' log 2 q p z q { p p z q , we have

Lemma C.2 means that the ϵ -cost is closely related to the variance of the log-density ratio between the target and the prior. We highlight that this conclusion holds generally and can be used to prove the bound of ϵ for all distribution families. But for now, we only focus on Gaussian, and leave more extensions to future exploration.

According to Lemma C.2, if we want to bound ϵ , we only need to bound Var r r s . To achieve this, we state the following Lemma:

Lemma C.3. For 1D Gaussian Q and P , with density q ' N p µ q , σ 2 q q , p ' N p µ p , σ 2 p q , and σ q ď σ p , defining random variable r by r ' log 2 q p x q { p p x q , x ' Q , if D KL r Q || P s is fixed to a constant, then Var r r s monotonically increases w.r.t σ q .

and thus

Reparameterizating y ' p x ´ µ q 1 ´ v q{ ? v , we have

The variance of r is given by

Taking derivative, we have

In summary, Lemma C.2 tells us if we want to upper-bound ϵ , we only need to upper-bound Var r r s ; while Lemma C.3 further tells us if we want to upper-bound Var r r s , we only need to look at the worst case, when Q and P share the same variance.

Following this, we can finally finish our prove:

WLG, assume the density of Q and P are given by

where d is the dimensionality. In this case, we have

and

Therefore,

Finally, taking expectation over X , and by Jensen's Inequality, we have

which finishes the proof for the Gaussian case in Proposition 3.2.

C.3 Proof of Theorem 3.3

We first restate Theorem 3.3 for easy reference:

Theorem 3.3. Let a pair of correlated random variables X , Z ' P X , Z be given. Assume we perform relative entropy coding using GPRS with space partitioning and let j ˚ denote the bin index and ˜ n ˚ the local sample index returned by the algorithm, and ϵ be as in Equation (9) . Then,

Proof. For those who are not familiar with GPRS, please note the notation used by Flamich (2024) is slightly different from those defined in this paper. We define the sample we encode as z ˚ , while in GPRS paper, it is denoted by ˜ X . We define the smallest τ as τ ˚ , while in GPRS paper, Flamich (2024) name it as the first arrival and denote it by ˜ T . We define the global sample index of the encoded sample by n ˚ , while Flamich (2024) denote it by N . In the following proof, since we will mainly modify the proof of Flamich (2024), we follow their notations unless otherwise stated.

We still follow the structure similar to the proof in Appendix C.1. Specifically, we first derive the upper bound of log 2 n ˚ conditional on a certain X ' x . Then, we take expectation over X . Finally, we derive the entropy by the same maximum entropy argument.

Derive E r log 2 n ˚ | X ' x s . Our proof starts from Equation (59) on Page 17 of the GPRS paper. It says that N ´ 1 follows a Poisson distribution given ˜ T ' t , with mean:

where ˜ µ is defined by Flamich (2024) as the average number of points 'under the graph' up to time t (See Appendix A on Page 13 of the GPRS paper for a detailed definition).

WLG, assume the encoded sample ˜ X is in the j ˚ -th bin. Then, according to the property of Poisson distribution, n ˚ ´ 1 is also Poisson-distributed, with mean π p j ˚ q ¨ p t ´ ˜ µ p t qq , where n ˚ is the local sample index in the j ˚ -th bin.

Therefore, we can modify Equation (133) on Page 21 of the GPRS paper as follows:

By Equation (132) on Page 21 of the GPRS paper, we have

where p 1 is the density of our adjusted prior, defined in Equation (25).

Therefore,

Hence, we have

Taking this back to Equation (104), we have

where the last follows the same principle as Equation (137) on Page 21 of the GPRS paper. Now, by the law of iterated expectations,we find

According to Equation (64) on Page 17 of the GPRS paper, E t r ˜ µ p t qs ' 1 . Therefore, we have

Surprisingly, this coincide with Equation (40) in the proof for PFR up to a constant. Therefore, following exactly the proof for PFR, we have

where ϵ is defined in the same way as the PFR case.

Derive E r log 2 ˜ n ˚ s . Following exactly the proof for PFR, we have

where ϵ is defined in the same way as the PFR case.

Derive entropy. Finally, following the same maximum entropy argument, we have

Do not forget we need log 2 J bits to encode the index j ˚ . Therefore, the two-part code's codelength is upper-bounded by

which finishes the proof.

C.4 Proof of Corollary 4.1

Corollary 4.1. Given a target distribution Q and a prior distribution P running Ordered random coding (ORC) to encode a sample from Q . Let ˜ Q be the distribution of encoded samples. If the number of candidates is N ' 2 D KL r Q || P s` t for some t ě 0 , then

Conversely, suppose that N ' 2 D KL r Q || P s´ t for some t ě 0 , then

Proof. Theis and Yosri (2022) prove Equation (124). Therefore, we focus on Equation (125).

First, we note that, as proved by Theis and Yosri (2022), the sample returned by ORC has the same distribution as the sample returned by Minimal random coding Havasi et al. (2019) and hence inherits all the conclusions of importance sampler by Chatterjee and Diaconis (2018).

Then, we define set A as also, define

Note, that

Assume given a random state S , the set of all candidate samples t z 1 , ..., z N u is uniquely determined. We therefore denote

According to Chatterjee and Diaconis (2018), we have

Therefore, we have

on the other hand, we have

Therefore, we have

D Experiment Details

D.1 Synthetic Gaussian Toy Examples

D.1.1 Generation of the toy examples

We explore the effectiveness of our proposed algorithm on 5D synthetic Gaussian toy examples. Specifically, we assume p p x q ' N p x | 0 , diag p σ 2 qq , and q p z | x q ' N p z | x , diag p ρ 2 qq . The density of prior P is given by p p z q ' N p 0 , diag p σ 2 ` ρ 2 qq . Therefore, for a certain realization of σ , ρ , we have the following REC task: transmitting a sample following Q whose density is q p z | x q ' N p z | x , diag p ρ 2 qq with the help of the prior P with density p p z q ' N p 0 , diag p σ 2 ` ρ 2 qq . To generate multiple toy examples, we randomly draw σ , ρ ' U p 0 , 1 q 5 . The dimension-wise mutual information I r X , Z s is 1 2 log 2 p p σ 2 ` ρ 2 q { ρ 2 q .

D.1.2 Space partitioning details

In this toy example, we naively assume the receiver knows both the dimension-wise mutual information and the entire KL divergence between Q and P (but dimension-wise KL is only available to the sender). In practice, we can achieve this by either encoding the KL divergence using negligible bits (for VAE in Section 5) or enforcing the KL budget during optimization (for INRs in Section 5).

˜

¸

Then we partition the d -th axis into approximately 2 I d ¨ D KL r Q } P s ř D d 1 ' 1 I d 1 intervals, where I d ' 1 2 log 2 p p σ 2 d ` ρ 2 d q { ρ 2 d q is the mutual information along d -th dimension. To achieve this, we adopt the following steps:

Initialize : we first initialize the number of intervals per dimension to 1, i.e., J r d s Ð 1 , @ d ' 1 , 2 , ¨ ¨ ¨ , D ; and we also initialize a vector of mutual information r I 1 , I 2 , ¨ ¨ ¨ , I d s .

Iterate : we iteratively increase the number of intervals by two in the dimension that currently has the highest mutual information. Specifically, we find d ˚ Ð arg max d t I 1 , I 2 , ¨ ¨ ¨ , I d u , and set J r d ˚ s Ð 2 J r d ˚ s .

Update and Repeat : after increasing J r d ˚ s , we decrement both the mutual information of that dimension and the total KL divergence by 1. We repeat this process until KL is exhausted.

D.1.3 Generation of the plots

After this, we run PFR with space partitioning (Algorithm 3) to encode exact samples and ORC with space partitioning (Algorithm 4) for approximate samples. We show the results in Figure 2 and Figure 3, and also include standard PFR and ORC for comparison. To plot Figure 2, we sample 200 different σ , ρ pairs, and repeat the encoding procedure 50 times with different seeds for each realization of σ , ρ . We calculate the negative log-likelihood of the 50 encoded indices using Zipf distribution P p N q9 N ´p 1 ` 1 { ζ q as a proxy of the codelength. To find the optimal value of ζ , we optimize the log-likelihood of the 50 indices with scipy.optimize.minimize . Finally, we average and find the IQR across these 50 repeats, and smooth the curves at the end for better visualization. In real-world applications, one may raise the concern that we need to transmit this ζ , which may lead to another overhead in the codelength. However, as we will see later in our V AE and INR cases, we can estimate the value of ζ from the training set (or a very small subset of the training set) and share it between the sender and receiver.

To plot Figure 3, we first sample 200 different σ , ρ pairs. For each realization of σ , ρ , we run ORC with a range of sample sizes: 2 t 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 u , and repeat the encoding procedure 500 times with different seeds for each sample size. Then we estimate maximum mean discrepancy (MMD, Smola et al., 2007) between the 500 encoded samples and real target samples with RBF kernel 5 . Since different target Q can lead to different MMD scales, we normalize the Q to standard isotropic Gaussian (and also scale and shift the encoded samples correspondingly) before calculating MMD. Note, that the absolute value MMD means little here. What we should care about is the relative comparison. After this process, we achieve a scatter plot. Finally, for clear visualization, we fit a

Gaussian Process regressor with scikit-learn package (Pedregosa et al., 2011) to estimate the mean and IQR.

D.2 Lossless Compression on MNIST with VAE

D.2.1 VAE architecture and training details

We follow (Flamich et al., 2024; Townsend et al., 2019) for the VAE architecture. Specifically, both encoder and decoder are 2-layer MLP. We set the hidden size to 400 and the latent size to 100, with ReLU activation function. We use fully factorized Gaussian as the latent distribution and Beta-Binomial as the likelihood distribution. The encoder maps the input flattened image to two 100-D vectors, corresponding to the mean and std in Gaussian distribution. We use softplus to ensure the standard deviation vector is positive. The decoder maps the 100-D latent vector back to two 784-D vectors, corresponding to two parameters in Beta-Binomial distribution. We also use softplus to ensure positivity. We use Adam (Kingma and Ba, 2017) with a learning rate of 0.001 as the optimizer, to train the V AE with a batch size of 1,000 for 1,000 epochs.

D.2.2 Space partitioning details

We estimate the mutual information I d along each dimension d in the latent space (i.e., averaging over the KL divergence across the entire training set), and then we partition the d -th axis into approximately 2 ˜ I d ¨ t D KL r Q } P s u ř D d 1 ' 1 I d 1 ¸ intervals. We use the same strategy as Appendix D.1.2, but also found simply rounding 2 ˜ I d ¨ t D KL r Q } P s u ř D d 1 ' 1 I d 1 ¸ to the nearest integer works very well. Here, we apply the floor function to the KL divergence since we need to transmit this value to the receiver.

To ensure our space partitioning strategy can reduce the runtime, we need mutual information to match well with the KL divergence for each test image along all dimensions. We empirically check this in Figure 10. As we can see, the test image KL divergence concentrated around mutual information. Also, if the mutual information is zero for some dimension, the test KL divergence will also be zero.

0

20

40

60

80

100

Dimension

0

2

4

6

8

10

12

14

bits

KL divergences

Mutual Information

D.2.3 Encoding details

We losslessly encode a test image with a code with four components: (1) t D KL r Q } P s u for each block; (2) the local sample index for each block; (3) the bin index for each block and (4) the image conditional on the encoded latent embedding. We discuss these 4 components in the following:

20

30

40

50

60

70

80

90

100

KL divergence (floored)

0.00

0.01

0.02

0.03

0.04

0.05

0.06

Density

For the first component, we estimate the distribution of t D KL r Q } P s u for each block from the training set, and use the negative log-likelihood as a proxy for the codelength. To provide an intuition on the distribution of t D KL r Q } P s u , we take the 2-block case as an example and visualize the histogram of t D KL r Q } P s u for each block in Figure 11.

For the second component, we estimate the codelength by the negative log-likelihood of Zipf distribution P p N q9 N ´p 1 ` 1 { ζ q . We find the optimal ζ value for each block from only 50 training images and find it generalizes very well to all test images. Note, that we find a different ζ for each block. This is because, as we can see from Figure 11, the random splitting will not ensure the two blocks to have the same KL divergence distribution and hence may lead to different optimal ζ .

The codelength for the bin index is simply t D KL r Q } P s u for each block.

For the last component, we estimate the codelength by the negative log-likelihood of Beta-Binomal distribution with the parameters predicted by the decoder.

One concern is that using negative log-likelihood as a proxy may underestimate the codelength. However, since we use the same proxy across all settings, this doesn't raise an issue in comparison.

D.3 Lossy Compression on CIFAR-10 with INRs

Following the methodology described by He et al. (2023), we model the latent weights and latent positional encodings (collectively denoted by h ) with a fully factorized Gaussian distribution q h . We also learn the prior p h using 15,000 images from the CIFAR-10 training set. During the testing and encoding stage, He et al. (2023) suggested using 16-bit blocks. Specifically, He et al. (2023) split h into K blocks, h ' r h r 1 s , ¨ ¨ ¨ , h r K s s , and trained q h for each test image x by optimizing the following β -ELBO:

where ˆ x h is the reconstructed image with parameter h , and β k -s are adjusted dynamically to ensure D KL r q h r k s || p h r k s s « 16 . In our experiments, we find that our proposed methods can safely accommodate larger block sizes of 48 bits. To ensure a fair comparison, we follow He et al. (2023)'s setting on CIFAR-10 exactly, except for the block sizes and the REC algorithm. Specifically, we modify the block sizes and the REC algorithm in codes at https: //github.com/cambridge-mlg/RECOMBINER (MIT License), and keep other settings unchanged.

As for the theoretical RD curve in Figure 4, we do not split h into blocks, and directly optimize

The β is set to the β obtained during training and kept the same for all test samples since we do not need to ensure any bit budget for a theoretically ideal REC codec. Then we directly draw a sample from q h to calculate PSNR, since we assume a theoretically ideal REC codec can encode a sample from the target exactly. As for the rate, we average D KL r q h || p h s across all test images to estimate I r h , X s , and calculate the codelength by I r h , X s ` log 2 p I r h , X s ` 1 q ` 4 .

D.3.1 Space partitioning details

To apply our proposed REC algorithm, we also need to estimate the mutual information I d along each dimension d in h . To achieve this, we randomly select 1,000 training images, and learn their posterior distributions q h by Equation (143). We also adjust β k to ensure D KL r q h r k s || p h r k s s « 48 on these 1,000 training images. After this, we average their KL divergence per dimension as an estimation of the dimension-wise mutual information. Having the estimation of I d , we use the same strategy as Appendix D.1.2 to partition the space during test time.

D.3.2 Encoding details

During the encoding stage, we simply apply Algorithm 4 to encode each block. We use Zipf distribution P p N q9 N ´p 1 ` 1 { ζ q to encode the local sample index, and then spend 48 bits to encode the bin index. We find the optimal ζ value from only 10 training images. Different from the VAE case, we find using the same ζ value for all blocks works well. Also, to ensure a fair comparison with other codecs, we do not estimate the codelength by log-likelihood but encode the index using torchac package (Mentzer et al., 2019) and count the length of the obtained bitstream.

E Licenses

F More Results

In Figure 12, we compare RECOMBINER (He et al., 2023) with our proposed algorithm with its original performance and also present the RD curves of other codecs for reference. Specifically, our baselines include classical codecs: (JVET, 2020), BPG (Bellard, 2014), JPEG2000; VAE-based codecs: Ballé et al. (2018), Cheng et al. (2020); and INR-based codecs: COIN++ (Dupont et al., 2022), VC-INR (Schwarz et al., 2023) and COMBINER (Guo et al., 2024).

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Bitrate (bit per pixel)

24

26

28

30

32

34

36

38

40

PSNR (dB)

BGP

JPEG2000

Ballé et al. (2018)

COIN++

COMBINER

VC-INR

VTM

RECOMBINER (He et.al (2023), block size=16 bits)

RECOMBINER (ours, block size=48 bits)

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We propose a framework to accelerate REC algorithm. We clearly state our main contribution in the last part of the introduction.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]

Justification: We discuss limitations in the last paragraph of the Section Conclusions and Limitations.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We prove the theorems, corollaries, and propositions in Appendix C.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We describe our algorithm in details in Algorithm 3 and Algorithm 4 and include how to reproduce the experiments in Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: All data we used is public. Our main contribution is the new REC algorithm and its algorithm to implement following the pseudo-codes.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We briefly describe the experiments and discuss results in Section 5. We include full experimental settings in Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: For the toy example, we report IQRs. For lossless compression using V AE in Section 5, we include the standard deviation across multiple running. For lossy compression using INRs, we do not report the error bars. This is because all the baselines do not report this in the RD curve, and we follow the convention. And since there are 10,000 test images, the results should be relatively reliable already.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [No]

Justification: Our main contribution is the new REC algorithm, which can easily run on any CPU. We, in practice, run the scripts on many CPUs in parallel to get the error bar across multiple runs, but this is not needed.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We do not believe our work has ethical concerns.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA] .

Justification: We do not believe our work has societal impacts.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA] .

Justification: We do not believe the paper poses such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA] .

Justification: The paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA] .

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors provide a formalization on how to introduce search heuristics for channel-simulation (sometimes called \"relative entropy coding\" in the ML literature). The encoder and decoder agree a priori on a binning scheme that divides the support of the prior/public distribution, which is used to control the search during encoding.\n\nDisclaimer: while I'm familiar with the basic literature (e.g. REC, PFRL, ORC), I haven't kept up with the literature since 2022.\n\nI'm more than willing to change my scores if the authors point out I missed something or didn't fully understand the contributions.\n\nSTRENGTHS:\n- The paper is nicely written. This is notably a complicated topic to understand and I feel the authors made quite a didactic effort.\n- The excess code-length is bounded for the distributions found in practice for NTC.\n\nWEAKNESSES:\n- Line 21: I feel this statement is misleading. In the first scenario (lines 14 to 20), the actual value taken on by the random variable is what is being transmitted. Meanwhile, in channel-simulation, all we care about is that the received quantity be distributed according to $Q_{Z | X}$. These are fundamentally different problems. The comparison would make some sense if, in the first scenario, the decoder is necessarily stochastic, as decoding can be seen as sampling from some posterior distribution over the data conditioned on the latent.\n\nQUESTIONS:\n- Line 66: Why is it important for the state $S$ to be random? I don't see this being used anywhere.\n\n- Line 213: The non-exactness here is with respect to $Q$. REC-like algorithms are usually formulated to be used in the latent space, i.e., the encoder wants to send a sample from $Q_{Z | X}$, but what we care a about from a practical standpoint is to transmit the data $X$ according to some rate-distortion trade-off. Shouldn't this discussion be formulated around the data space (maybe this is too hard of a problem?)?\n\n- Regarding the binning prior $\\pi$. The target distribution is a function of $X$, i.e., $Q_{Z | X}$. Isn't it hard to pick a good $\\pi$ given this distribution is constantly changing?\n\nLIMITATIONS:\n- Line 267: The results are computationally intensive, and thus impractical, even for MNIST. Requiring the value of mutual information to be known by encoder/decoder seems a bit much (the authors mention this in the limitations section).", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-06-17"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n**Global disclaimer:** I am very unfamiliar with the topic of the paper. I did my best to try and read the literature and understand as much as I could but my input may be very limited.\n\n**Summary:**\nThe paper focus on relative entropy coding (REC) algorithms and propose to circumvent a major pitfall which is the prohibitive encoding time. Indeed the commonly used PFR algorithm requires to sample points until finding a good enough alignment with the target distribution. Efficient algorithms only exist in dimension 1. The authors propose to refine the algorithm by partitioning the search space into bins that will better drive the sampling from the coding to the target distribution.\nThe author then theoretically prove that their approach by is founded (Theorem 3.1), and derive precise codelength bounds for the space partition algorithm. They also discuss a generic partitionning strategy.\nFinally, a benchmark highlight the benefits of the space partition approach for compression against the standard PFR algorithm on toy example and real datasets.\n\nSTRENGTHS:\nThe paper is very didactic, well written and is quite nice to read. The proposed approached is well structured and detailed (sections 3.1). The theory exposed in the paper is sound and the proof seem correct to me.\nThe benchmark is well made and honest.\nFinally the appendix is very helpful, I particularly appreciated appendix B that shed some light on the partition strategy.\n\nWEAKNESSES:\nThe idea is finally quite elementary: dividing the space once for all into bins. The proofs are \"elementary\" (still technical) refinements from the classical PFR theory (starting point is often [Li and El Gamal, 2018] and [ Flamich, 2024]).\n\nQUESTIONS:\n- How the bounds from Theorem 3.1 do compare with the bound from [Li and El Gamal, 2018, Theorem 1] in the same setting, i.e. when $J=1$? More generally, could we recover a bound such as Proposition 3.2 but with a dependence in $J$.\n- The number of intervals for each axis is justified in Appendix B, but could the author give more details about how these interval are constructed?\n- The proposed approach to the construction of the bins follow depends on the chosen representation. Would it be possible to automatically find a transformation prior to the bining that would provide better results while keeping the same strategy?\n- Is there a way to update $\\pi$ along the algorithm, using e.g. the Gumbel trick and gradient based method?\n- It would be interesting to add an ablation study on the number of intervals per axis.\n- Table 1 shows some improvements in number of steps but it would be interesting to have the same thing in time. Could you elaborate about Figure 3, I do not understand what is meant by runtime here.\n\nLIMITATIONS:\nThe approach is not fully generic and despite very good success on some examples, it could fail on other as clearly explained in Appendix B.\n\nTypos: the punctuation missing in the maths of the appendix\n\n**Proposed score:**\nMy rating is based on the fact that the paper is sound, the maths are correct but the scope of the idea is quite simple. I do believe that this work is well presented, interesting and worth publication. On the other hand I am not familiar with this field. As such, I propose a score of 5/10 (a weak accept) but I am very open to discussion with the authors and the other reviewers.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a space partioning technique to speed up  relative entropy coding.\nIn entropy coding the sender first transforms X into a representation Z that they encode.\nThis particular step can done  by Poisson functional representation (PFR).\nUnfortunately, PFR’s random runtime can be a significant drawback.\n\nTo circumvent this issue,  the author propose a space partitioning scheme to reduce runtime in practical scenarios.\nThis can be seen (sic) as a search heuristic to a standard REC algorithm\n\nThe author then proceed to give a theoretical analysis of the method and support their improvement with numerical results.\n\nSTRENGTHS:\nThe paper adresses an important practical problem with many important application.\nThe algorithm and theoretical results are stated clearly.\nThe context and related work is clear to the reviewer (who is not an expert in the field).\nThe numerics show clearly the advantage/improvements  yielded by the technique.\n\nWEAKNESSES:\nAlthough the proposed technique is \"simple and efficient\", it lacks discussion whether the binning technique impacts the final result (besides codelength).\n\nQUESTIONS:\n- it is unclear to the reviewer how the grid should be set.  \n-it is unclear to the reviewer how the choice of prior for the binning distribution influences the algorithm.\n-  the weakness section raises a question that I would like to see adressed:  how is the binning technique impacting the final result (besides codelength).\nI agree that the uniform distribution seems like a natural choice, but it would be good to discuss the choice of prior here\n\nI will happily increase my grade if those questions are properly adressed.\n\nLIMITATIONS:\nThe authors recognize some clear limitations of their approach in the conclusion section.\nIn particular, the approach proposed by the authors is too limited and does not extend to all possible practical cases.\nThe reviewer agrees with those limitations, although still believes they should not be an obstacle for publication. as this work is a \"first step\" in this direction.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}