Accelerating Relative Entropy Coding with Space Partitioning

Authors: jiajun  he
gergely  flamich
josé miguel hernández-lobato
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 19:54:30

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.

Main paper: 
 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:IrX; Zs `log 2 pIrX; Zs `1q `Op1q.(1)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 DKLrQ } 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 " rZ 1 , Z 2 , ¨¨¨, Z K s, where the coding cost of each block Z k is approximately κ bits. This strategy reduces the runtime to OpK2 κ q. However, by Equation (1), the logarithmic and constant overhead only becomes negligible if IrX; 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 IrZ i ; Xs for each dimension i, as opposed to just the total IrZ; Xs. Incorporating this information into our space partitioning scheme allows us to construct faster relative entropy coding algorithms. 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:EQUATIONNext, letting r " q{p be the density ratio, the sender calculates τ n Ð T n {rpZ n q for each sample and returns N ˚Ð arg min iPN 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 D8rQ } 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 Expp1q, 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 OpN log N q time by sorting, Theis and Yosri (2022) suggest an iterative procedure similar to Equation (3) with OpN q time complexity: T 0 " 0, T n Ð T n´1 `N{pN ´n `1q∆T n , ∆T n " Expp1q, n " 1, 2, ¨¨¨, N3 Relative Entropy Coding with Space PartitioningIn 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.We explore the effectiveness of our proposed algorithm on 5D synthetic Gaussian toy examples. Specifically, we assume ppxq " N px|0, diagpσ 2 qq, and qpz|xq " N pz|x, diagpρ 2 qq. The density of prior P is given by ppzq " N p0, diagpσ 2 `ρ2 qq. Therefore, for a certain realization of σ, ρ, we have the following REC task: transmitting a sample following Q whose density is qpz|xq " N pz|x, diagpρ 2 qq with the help of the prior P with density ppzq " N p0, diagpσ 2 `ρ2 qq. To generate multiple toy examples, we randomly draw σ, ρ " Up0, 1q5 . The dimension-wise mutual information IrX, Zs is 1 2 log 2 p pσ 2 `ρ2 q {ρ 2 q. 
coding scheme: 
 Given a shared coding distribution P and a target Q, our algorithm proceeds as follows:1. The sender and receiver agree on a partition of P 's support consisting of J bins tB 1 , ¨¨¨, B J u with equal probability mass, i.e. P pB j q " 1 {J for j " 1, . . . , J. As we will discuss later, these # initialize:τ ˚Ð 8, t 0 Ð 0, N ˚Ð 0 . r max Ð sup z ! qpzq ppzq) .# run PFR:for n " 1, 2, ¨¨¨do Sample ∆t n " Expp1q; t n Ð t n´1 `∆t n .# simulate a sample with PRNG:z n Ð PRNGpP, S, nq 4 # update τ ˚: τ n Ð t n ¨ppz n q{qpz n q. if τ n ď τ ˚then τ ˚Ð τ n , N ˚Ð n. end if # check stopping criterion: if t n {r max ą τ ˚then break end if end forAlgorithm 2 PFR with Space Partitioning Input: Q, P , and random states tSju J j"1 . Output: Bin index j ˚, local sample index ñ˚.# initialize: τ ˚Ð 8, t 0 Ð 0, j ˚Ð 0, ñ˚Ð 0. Partition space in J "bins", s.t. P pB j q " 1{J.Select categorical distribution π. ñj Ð 0 for j " 1, 2, ¨¨¨, J. r 1 max Ð maxj"1,¨¨¨,J ! sup zPB j ! qpzq ppzq P pB j q πpjq)). # run PFR:for n " 1, 2, ¨¨¨do Sample ∆t n " Expp1q; t n Ð t n´1 `∆t n .# simulate a sample with PRNG:j n " π; ñjn Ð ñjn `1. z n Ð PRNGpP | Bj n , S jn , ñjn q. # update τ ˚: τ n Ð J ¨πpj n q¨t n ¨ppz n q{qpz n q. if τ n ď τ ˚then τ ˚Ð τ n , j ˚Ð j n , ñ˚Ð ñjn . end if # check stopping criterion: if t n {r 1 max ą τ ˚then break end if end forbins can overlap, but to aid understanding for now, it may be helpful to imagine the space is partitioned by non-overlapping bins.2. According to the target distribution Q, the sender selects a categorical distribution for the bins, with event probabilities defined as πpjq 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.3. 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 | Bj n . This sample is generated with the random state S jn 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 ñjn . For the sake of clarity, we will call n the global sample index and ñjn the local sample index hereafter; (c) the sender examines the sample by calculating τ n :τ n Ð T n {r 1 , r 1 -qpZ n q{p 1n , p 1 n -πpj n qppZ n q{P pB jn q " J ¨πpj n q ¨ppZ n q, (5) 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:T n {r 1 max ą τ ˚, r 1 max -max j"1,2,¨¨¨,J ! sup zPBj qpzqP pB j q ppzqπpjq ) (6)If the criterion is met, halt and return j N ˚and ñ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 asEQUATION)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 | Bj N ˚with seed S j N ˚. Then, taking the ñN ˚-th sample in this sequence, the receiver successfully retrieves the desired sample.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 ñ˚t he local sample index returned by the algorithm. Then, the entropy of the two-part code is bounded byEQUATIONEQUATIONWe prove Theorem 3.1 in Appendix C.1. Note, that when IrX; Zs is sufficiently large, the term log 2 pIrX; Zs ´log 2 J `EX rϵs `1q `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 DKLrQ } 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 rQ } P s, and E X rϵs ď 0.849 a IrX; Zs.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).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 πpjq 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 ñt he local sample index returned by the algorithm, and ϵ be as in Equation (9). Then, Hrj ˚, ñ˚s ď IrX; Zs `EX rϵs `log 2 pIrX; Zs ´log 2 J `EX rϵs `1q `6.(10)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 VAE-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 qpzq " N pz|0.001, σ 2 q, and ppzq " N pz|0, 1q. When σ Õ 1 (approaching 1 from below), the expected runtime for PFR diverges: 2 D8rQ||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 DKLrQ } 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. 
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 D8rQ||P 1 s . Thus, we have the following constrained optimization problem:EQUATION)The solution turns out to be intuitive:EQUATIONHowever, 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 πpjq 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.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 DKLrQ } P s`t for some t ě 0, thenEQUATION)Conversely, supposing that N " 2 DKLrQ } P s´t for some t ě 0, thenEQUATION)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 DKLrQ } 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:πpjq " QpB j q, j " 1, 2, ¨¨¨, J.This is because D KL rQ } P 1 s " D KL rQ } P s ´řJ j"1 QpB j q log 2 πpjq `řJ j"1 QpB j q log 2 P pB j q, and the cross-entropy p´ř J j QpB j q log 2 πpjqq 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 πpjq " QpB 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 tDKLrQ } P su , 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 tD KL rQ } P su. 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 wheren d " D KL rQ } P s ¨Id 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 DKLrQ } P 1 s according to Corollary 4.1. When the total number of partitions J " 2 DKLrQ } P s , we find D KL rQ } P 1 s " ´řJ j"1 QpB j q log 2 QpB 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:1. 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) . 2. 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 " Bernp1{2q 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 ¨p1 `DTV r Q, Qsq (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. 
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 rQ } 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 tD KL rQ } P su so that they can partition the space into the same J " 2 tDKLrQ } P su bins. However, this value varies across different images. Therefore, we estimate the distribution of tD KL rQ } P su for each block from the training set, and entropy code it for each test image.We evaluate the performance achieved by t2, 4u 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 tD KL rQ } P su; (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) Table 1 : Lossless compression performance on MNIST test set with different REC settings. We include the theoretical results and the results by GPRS (Flamich, 2024) for reference. We repeat each setting 5 times and report the mean and std. *: κ represents the average KL divergence between target Q and the search heuristic P 1 . We estimate it by MC estimator and average across all test images. We empirically find κ « 14 when using 2 blocks and κ « 7 when using 4 blocks. :: Given that κ is relatively small, we can employ more than t2 κu samples to further reduce the bias. We hence report the results achieved by t2 κ`2 u samples, which remains manageable in both cases. We also provide the theoretical RD curve for an ideal REC algorithm (i.e., assuming we can encode an exact sample in a single block, whose codelength is calculated by Equation ( 1)). Notably, our method's performance is already very close to this theoretical result.for each block. Besides, unlike in the VAE 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 tD KL rQ } P su.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 rQ } P s " 16 bits, our proposed algorithm with the block size of D KL rQ } 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 rQ } 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. 
related works: 
 There are several efforts to accelerate REC algorithms. Flamich et al. (2022 Flamich et al. ( , 2024)) ; Flamich (2024) leveraged the idea of partitioning in 1D space, achieving an impressive runtime in the order of O pD 8 rQ||P sq or O pD KL rQ } 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. 
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.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 VAE 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. 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). 
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] Justification: See Appendix E.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. 
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. 
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. 
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. 
 
 Appendix: 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 r1s ˆQr2s ˆ¨¨¨ˆQ rDs , shared and fully-factorized coding distribution P " P r1s ˆP r2s ˆ¨¨¨ˆP rDs , and shared random states S j for each partition j. Output: Bin index j ˚, local sample index ñ˚.# initialize:τ ˚Ð 8, t 0 Ð 0, j ˚Ð 0, ñ˚Ð 0. Partition d-th dimension into J rds intervals B rds 1 , B rds 2 , ¨¨¨, B rds J rds , s.t., J Ð ś D d"1 J rds . ñ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: Ź by inverse transform sampling when both of the partitioning and P are axis-aligned. # update τ ˚: Partition space in J "bins". z Ð PRNGpP | B j ˚, S j ˚, ñ˚q .for d " 1, 2, ¨¨¨, D do for i " 1,τ n Ð J ¨QpB jn q¨t n ¨ppz n q{qpz n q. if τ n ď τ ˚then τ ˚Ð τ n , j ˚Ð j n ,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 ppzq 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 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. Figure 5 : Elucidating the generality of space partitioning. Ω represents the original space. We use dashed lines to represent the boundaries of the partitions. (a) We can add an auxiliary axis x aux to the original space, forming an augmented space. 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. We also require P aux and Q aux to be the same, so that D KL rQ aug ||P aug s is the same as the original KL D KL rQ } P s. Dividing the augmented space into non-overlapping bins will lead to non-overlapping bins or overlapping bins in the original space. For example, as shown in (b), dividing the augmented space into non-overlapping bins whose boundaries are parallel to the auxiliary axis results in the standard non-overlapping bins in the original space Ω. As shown in (c), dividing the augmented space into non-overlapping bins whose boundaries are orthogonal to the auxiliary axis results in fully overlapping bins in the original space Ω. Also, as in (d), the augmented space can be divided in an arbitrary manner, leading to generally overlapping bins in the original space Ω.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 ˆx2 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 pB j q " QpB j q, @j. In this case, choosing πpjq " QpB j q as discussed in Section 4.2, we can write the density of P 1 (Equation ( 7)) as p 1 pzq " ř J j"1 1tz P B j u πpjqppzq P pB j q (16)EQUATION" ppzq (20) 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 becauseD KL rQ } P 1 s " D KL rQ } P s ´J ÿ j"1QpB j q log 2 πpjq `J ÿ j"1QpB j q log 2 P pB j q (21)" D KL rQ } P s ´˜J ÿ j"1 QpB j q log 2 Qpjq ´J ÿ j"1QpB j q log 2 P pB j q ¸(22)We recognize the term ´řJ j"1 QpB j q log 2 Qpjq ´řJ j"1 QpB j q log 2 P pB j q ¯as a KL divergence and hence will be non-negative. Whene there exists B j , s.t. QpB j q ‰ P pB j q, the KL term ´řJ j"1 QpB j q log 2 Qpjq ´řJ j"1 QpB j q log 2 P pB j q ¯will always be positive to ensure D KL rQ } P 1 s ď D KL rQ } P s and hence guarantee a reduction in the runtime.𝑥 2 𝑥 1 𝑝 𝑞 (a) 𝑥 2 𝑥 1 𝑝 𝑞 (b) 𝑥 2 𝑥 1 𝑝 𝑞 (c)Figure 6 : An example to explain why the number of intervals assigned to each axis matters when we partition the space with axis-aliged grids. We use dashed lines to represent the boundaries of the partitions. In this example, Q and P have the same marginal along x 1 . In (a), we partition the space by solely dividing x 1 . This will make the search heuristic P 1 equal to P , leading to no improvements in the runtime. In (b), we partition the space by dividing x 2 . This will make the search heuristic P 1 align with Q better than P (i.e., D KLrQ||P 1 s ă D KLrQ||P s ), reducing the runtime. In (c), we partition the space by dividing both x 2 and x 1 . This will still reduce the runtime, but is not as efficient as (b). Please note that while the plot looks similar to Figure 5 , they represent different concepts.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 ^pI d ¨DKLrQ } P s { ř D d 1 "1 I d 1 q intervals to avoid the worst case in Figure 6a .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. 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. 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 Figure 9 : Rate-distortion of RECOMBINER (on 100 CIFAR-10 test images) by our space-partitioning algorithm using different partition strategies: randomly assigning intervals to dimensions; first removing axis with MI « 0 and then randomly assigning intervals; assigning intervals according to MI (our proposed strategy). We include standard ORC and theoretical RD for reference. Our proposed partition strategy is better than randomly assigning intervals per axis. Surprisingly, the results achieved by first removing uninformative dimensions (MI « 0) and then randomly assigning intervals to other dimensions are only slightly worse than our proposed partition strategy. This indicates that our algorithm is not very sensitive to how we construct the intervals, as long as we avoid partitioning along uninformative dimensions.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 ñ˚t he local sample index returned by the algorithm. Then, the entropy of the two-part code is bounded by: Hrj ˚, ñ˚s ď IrX; Zs `EX rϵs `log 2 pIrX; Zs ´log 2 J `EX rϵs `1q `4, (23)EQUATION)Proof. The proof will be organized as follows: first, we will prove an upper bound of Er ñ˚| X " xs for the one-shot case, i.e., the expectation of ñ˚g iven a realization of the data X. Then, we will take the expectation over X to achieve an upper bound for the average case Er ñ˚s . Finally, we will consider encoding ñ˚u sing a Zipf distribution to calculate the upper bound given in Theorem 3.1.Derive Er ñ˚| X " xs. 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 asEQUATION)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 ñ˚.According to Li and El Gamal (2018) , we haveτ ˚" Expp1q (26) n ˚´1 " Poissonpλq (27) λ ď τ ˚qpz ˚q p 1 pz ˚q (28)Equation ( 28) is obtained from Appendix A in Li and El Gamal (2018) and replacing p by p 1 .Additionally, we haveEQUATION)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 ñ˚i s equivalent to the probability that there are ñ˚´1 success in the first n ˚´1 Bernoulli trials.Therefore, according to the property of Poisson distribution, we haveEQUATIONTherefore, we have E log 2 r ñ˚| X " xs " E j ˚Ez ˚"Q| B j ˚Eτ ˚E ñ˚| j ˚,z,τ ˚rlog 2 p ñ˚´1 `1qsJensen's ď E j ˚Ez ˚"Q| B j ˚Eτ ˚"log 2 E ñ˚| j ˚,z,τ ˚p ñ˚´1 `1q ‰ (32)eq. ( 30) ď E j ˚Ez ˚"Q| B j ˚Eτ ˚rlog 2 pπpj ˚qλ `1qseq. ( 28)EQUATIONJensen'sEQUATIONeq. ( 25)EQUATIONEQUATION)EQUATIONThe last line is because we require P pB 1 q " P pB 2 q " ¨¨¨" P pB 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 p1 `xq at 0, `8, we haveEQUATIONTherefore, EQUATIONď E z"Q log 2 " 1 " qpzq ppzq ą J * ˆlog 2 qpzq Jppzq ˙ȷ `log 2 e (43)The last line is becauseEQUATIONWe can further simplify Equation ( 43) as follows:E log 2 r ñ˚| X " xs ď D KL rQ||P s ´log 2 J `EQ " 1 " log 2 qpzq ppzq ď log 2 J * ˆlog 2 J ´log 2 qpzq ppzq ˙ȷ `log 2 e (47)" D KL rQ||P s ´log 2 J `ϵ `log 2 e (48)where we denoteϵ " E z"Q " 1 " log 2 qpzq ppzq ď log 2 J * ˆlog 2 J ´log 2 qpzq ppzq ˙ȷ (49) " E z"Q " max " log 2 J ´log 2 qpzq ppzq , 0 *ȷ (50)Derive Er ñ˚s . Now, we consider the average case by taking the expectation over X. We have 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 :Er ñ˚s ď E X D KL rQ Z|X ||P Z s ´log 2 J `EXHr ñ˚s `log 2 J ď D KL rQ||P s `ϵ `log 2 pD KL rQ||P s ´log 2 J `ϵ `1q `4.(57)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.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 qpzq {ppzq which is a new random variable, we rewrite the ϵ-cost asEQUATION)"E r r1 tr ď J u pJ ´rqs (59)" ż J ´8 pprqpJ ´rqdr (60)where we also write J " log 2 J, which is a monotonically increasing function of J. Take derivative w.r.t. J , we haved ş J ´8 pprqpJ ´rqdr dJ " ppJ qpJ ´J q `ż J ´8 pprqdr " ż J ´8 pprqdr ě 0 (61)which finishes the proof.Therefore, to prove Proposition 3.2, we only need to look at the worst case, where J " 2 DKLrQ||P s . We now prove for Uniform and Gaussian, respectively:Uniform. For uniform distribution, WLG, assume P " Up0, 1q n , and Q " UpAq, where A Ď p0, 1q n . In this case, we have D KL rQ||P s " ´log 2 µpAq, and qpzq ppzq " 1 µpAq , @z P A, where µ is the Lesbague measure. Therefore,EQUATIONOn 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 DKLrQ||P s , the ϵ-cost defined in Equation ( 9) is upper-bounded byEQUATION)where r is the RV defined by r " log 2 qpzq {ppzq, z " Q.Proof. First, we know 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.E z"Q " D KL rQ||PAccording to Lemma C.2, if we want to bound ϵ, we only need to bound Varrrs. 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 qpxq {ppxq, x " Q, if D KL rQ||P s is fixed to a constant, then Varrrs monotonically increases w.r.t σ q .Proof. WLG, let's assume p " N p0, 1q, µ q ą 0, D KL rQ||P s " δ. Denoting v " σ 2 q , we haveµ q " b 2δ `logpσ 2 q q ´σ2 q `1 " a 2δ `logpvq ´v `1 (78)and thusr " 1 2v `p1 ´vqx 2 ´2µ q x `µ2 q ´v `vµ 2 q `v2 ˘( 79)EQUATIONReparameterizating y " px ´µq 1´v q{ ? v, we haveEQUATIONThe variance of r is given byEQUATIONTaking derivative, we haveEQUATIONIn summary, Lemma C.2 tells us if we want to upper-bound ϵ, we only need to upper-bound Varrrs; while Lemma C.3 further tells us if we want to upper-bound Varrrs, 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 byEQUATIONwhere d is the dimensionality. In this case, we have 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 ñt he local sample index returned by the algorithm, and ϵ be as in Equation (9). Then, Hrj ˚, ñ˚s ď IrX; Zs `EX rϵs `log 2 pIrX; Zs ´log 2 J `EX rϵs `1q `6.EQUATIONµ i z i ´n ÿ i"1 µ 2 i 2 ¸" N ˜log 2 e n ÿ i"1 µ 2 i 2 , log 2 2 e n ÿ i"1 µ 2 i ¸(95) Therefore, ϵ ď 1 2 a Varrrs " g f fProof. 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 Erlog 2 n ˚|X " xs. 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:EQUATION)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 πpj ˚q ¨pt ´μptqq, 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:EQUATION" log 2 rπpj ˚q ¨pt ´μptqq `1sEQUATIONBy Equation ( 132) on Page 21 of the GPRS paper, we havet ď qpxq p 1 pxqPr T ě ts (105)where p 1 is the density of our adjusted prior, defined in Equation (25).Therefore,EQUATION" ¨π pj ˚q P pB j ˚qqpxq ¨π pj ˚q ppxqPr T ě ts (107)EQUATIONHence, we haveEQUATION" qpxq {pJ ¨ppxqq Pr T ě ts `1 (110)EQUATIONTaking this back to Equation ( 104), we haveEQUATIONwhere the last follows the same principle as Equation ( 137) on Page 21 of the GPRS paper. Now, by the law of iterated expectations,we findEQUATIONeq. ( 113)EQUATIONEQUATIONAccording to Equation ( 64) on Page 17 of the GPRS paper, E t rμptqs " 1. Therefore, we haveEQUATIONSurprisingly, this coincide with Equation (40) in the proof for PFR up to a constant. Therefore, following exactly the proof for PFR, we haveEQUATION)where ϵ is defined in the same way as the PFR case.Derive Erlog 2 ñ˚s . Following exactly the proof for PFR, we haveEQUATION)where ϵ is defined in the same way as the PFR case.Derive entropy. Finally, following the same maximum entropy argument, we have Hr ñ˚s ď IrX; Zs ´log 2 J `EX rϵs `2 ¨log 2 e `log 2 pIrX; Zs ´log 2 J `EX rϵs `2 ¨log 2 e `1q `1 (120) ď IrX; Zs ´log 2 J `EX rϵs `log 2 pIrX; Zs ´log 2 J `EX rϵs `1q `1 `2 ¨log 2 e `log 2 p2 ¨log 2 e `1qď IrX; Zs ´log 2 J `EX rϵs `log 2 pIrX; Zs ´log 2 J `EX rϵs `1q `6Do not forget we need log 2 J bits to encode the index j ˚. Therefore, the two-part code's codelength is upper-bounded by Hr ñ˚s `log 2 J ď IrX; Zs `EX rϵs `log 2 pIrX; Zs ´log 2 J `EX rϵs `1q `6which finishes the proof.C.4 Proof of Corollary 4.1Corollary 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 DKLrQ||P s`t for some t ě 0, thenEQUATIONConversely, suppose that N " 2 DKLrQ||P s´t for some t ě 0, thenD TV r Q, Qs ě 1 ´2´t{2 ´Pz"Q ˆlog qpzq ppzq ď D KL rQ||P s ´t 2 ˙(125)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 asEQUATIONNote, thatEQUATIONAssume given a random state S, the set of all candidate samples tz 1 , ..., z N u is uniquely determined. We therefore denoteEQUATIONAccording to Chatterjee and Diaconis (2018) , we haveEQUATIONE S r1tI N,S " 1us " PrI N,S " 1s ě 1 ´2´t{2 .(131) Therefore, we haveQpAq " E Qrf s (132)EQUATIONě 1 ´2´t{2 .(137) on the other hand, we haveQpAq " P z"Q rz P As " P z"Q ˆlog qpzq ppzq ď D KL rQ||P s ´t 2 ˙(138)Therefore, we haveEQUATIONě QpAq ´QpAq (141)ě 1 ´2´t{2 ´Pz"Q ˆlog qpzq ppzq ď D KL rQ||P s ´t 2 ˙(142) D Experiment Details D.1 Synthetic Gaussian Toy Examples D.1.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˜Id ¨DKL rQ } P s ř D d 1 "1 I d 1 ¸intervals, where I d " 1 2 log 2 p pσ 2 d `ρ2d 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 rds Ð 1, @d " 1, 2, ¨¨¨, D; and we also initialize a vector of mutual information rI 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 tI 1 , I 2 , ¨¨¨, I d u, and set J rd ˚s Ð 2J rd ˚s.• Update and Repeat: after increasing J rd ˚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.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 pN q9N ´p1`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 VAE 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: ,3,4,5,6,7,8,9,10u , 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) 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 VAE with a batch size of 1,000 for 1,000 epochs.2 t1,2We 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 approximately2 ˜Id ¨tD KL rQ } P su ř D d 1 "1 I d 1¸intervals. We use the same strategy as Appendix D.1.2, but also found simply rounding 2˜Id ¨tD KL rQ } P su ř 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. We can see the KL divergences concentrate around mutual information. And if the mutual information is zero for some dimension, the test KL divergence will also be zero. This ensures our space partitioning strategy can reduce the runtime. Here, we take the case where we split the latent embeddings into 2 blocks as an example. Different colors represent different blocks. We note that the difference between these 2 blocks is solely due to randomness in the splitting.We losslessly encode a test image with a code with four components: (1) tD KL rQ } P su 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:• For the first component, we estimate the distribution of tD KL rQ } P su 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 tD KL rQ } P su, we take the 2-block case as an example and visualize the histogram of tD KL rQ } P su for each block in Figure 11 . • For the second component, we estimate the codelength by the negative log-likelihood of Zipf distribution P pN q9N ´p1`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 tD KL rQ } P su 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.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 " rh r1s , ¨¨¨, h rKs s, and trained q h for each test image x by optimizing the following β-ELBO:EQUATIONwhere As for the theoretical RD curve in Figure 4 , we do not split h into blocks, and directly optimize L " β ¨DKL rq h ||p h s `Eq h rDistortionpx h , xqs .(144) 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 rq h ||p h s across all test images to estimate Irh, Xs, and calculate the codelength by Irh, Xs `log 2 pIrh, Xs `1q `4.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 rq h rks ||p h rks 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.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) . (He et.al (2023) , block size=16 bits) RECOMBINER (ours, block size=48 bits) 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. • 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.We note that REC's decoding is very fast, so in this paper, runtime will refer exclusively to encoding time.PRNGpP, S, nq means simulating the n-th sample in the random sequence from P by PRNG with seed S.We use the codes at https://github.com/jindongwang/transferlearning/blob/master/code/ distance/mmd_numpy_sklearn.py (MIT License)https://github.com/cambridge-mlg/RECOMBINERhttps://github.com/jindongwang/transferlearning/blob/master/code/distance/mmd_ numpy_sklearn.pyhttps://github.com/fab-jul/torchac 


Decision: This paper proposes an heuristic to improve relative entropy coding algorithms, and validates its method numerically.

The reviewers unanimously agreed that this well-written paper effectively addresses important questions: I recommend it for acceptance.

Reviews: review: paper_summary: paper_summary: The paper proposes a space partioning technique to speed up  relative entropy coding.
In entropy coding the sender first transforms X into a representation Z that they encode.
This particular step can done  by Poisson functional representation (PFR).
Unfortunately, PFR’s random runtime can be a significant drawback.

To circumvent this issue,  the author propose a space partitioning scheme to reduce runtime in practical scenarios.
This can be seen (sic) as a search heuristic to a standard REC algorithm

The author then proceed to give a theoretical analysis of the method and support their improvement with numerical results.
strength_weakness: strength_weakness: Although the proposed technique is "simple and efficient", it lacks discussion whether the binning technique impacts the final result (besides codelength).
questions: questions: - it is unclear to the reviewer how the grid should be set.  
-it is unclear to the reviewer how the choice of prior for the binning distribution influences the algorithm.
-  the weakness section raises a question that I would like to see adressed:  how is the binning technique impacting the final result (besides codelength).
I agree that the uniform distribution seems like a natural choice, but it would be good to discuss the choice of prior here

I will happily increase my grade if those questions are properly adressed.
limitations: limitations: The authors recognize some clear limitations of their approach in the conclusion section.
In particular, the approach proposed by the authors is too limited and does not extend to all possible practical cases.
The 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.
paper_summary: paper_summary: **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.

**Summary:**
The 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.
The 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.
Finally, a benchmark highlight the benefits of the space partition approach for compression against the standard PFR algorithm on toy example and real datasets.
strength_weakness: strength_weakness: The 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]).
questions: questions: - 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$.
- 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?
- 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?
- Is there a way to update $\pi$ along the algorithm, using e.g. the Gumbel trick and gradient based method?
- It would be interesting to add an ablation study on the number of intervals per axis.
- 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.
limitations: limitations: The approach is not fully generic and despite very good success on some examples, it could fail on other as clearly explained in Appendix B.

Typos: the punctuation missing in the maths of the appendix

**Proposed score:**
My 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.
paper_summary: paper_summary: The 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.

Disclaimer: while I'm familiar with the basic literature (e.g. REC, PFRL, ORC), I haven't kept up with the literature since 2022.

I'm more than willing to change my scores if the authors point out I missed something or didn't fully understand the contributions.
strength_weakness: strength_weakness: - 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.
questions: questions: - Line 66: Why is it important for the state $S$ to be random? I don't see this being used anywhere.

- 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?)?

- 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?
limitations: limitations: - 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).

Comments: title: Response to authors
Thank you for your response.
Thank you for your review! Please consider our response
Thank you for your review! Please consider our response
Thank you for your review! Please consider our response
comment: I read the response, as well as the other reviews, and took time to understand the author's responses to all raised points. Overall I am pleased with the author's answers, some points I missed during the review have been clarified.

**Weakness: the idea and proofs are elementary**
I fully agree with the author's response, thank you for the clarifications. I would encourage the author to make point (c) maybe a little bit clearer in the introduction of the paper.

I believe the two added figures in Appendix B2 fill an important gap in the original manuscript and shed some light on the proposed approach.

I think I understood the paper, the ideas, and the approach the author are proposing and I think this work is strong and worth acceptance.

As such I raise my grade to 8.
I'm satisfied with the responses provided by the authors.
I'm happy about the clarification and apologize for missing any details. 
I updated my score to reflect the discussion.
I'm satisfied with the responses in general and have updated my score accordingly. Thank you for the detailed response.

> However, we’d also like to note that the difference between using [...]

Agreed. To be clear, I in no way was trying to make a technical distinction at the practical level, but just thought it might be easier for a reader outside the field to digest.

> We do not necessarily agree that the VAE results are computationally intensive and impractical. [...]

This is a fair point: within the REC family this paper significantly improves upon the number of candidates needed, which relates to computational complexity. I guess my comment was geared towards REC methods in general, but I believe that isn't a fair way to review a paper.

> We agree that this is not a totally general assumption. [...]

I guess this makes sense. It would be nice if this discussion was in the paper, even if in the appendix (but referenced in the main text), as it seems a bit arbitrary at first glance. Maybe an example where this is already assumed to be true, even if implicitly by calculating, for example, cross-correlations in the gaussian setting.
We thank the reviewer once again for their effort in the reviewing process. As there are only a few working days left in the discussion period, we would like to ask if our response has satisfied the reviewer’s concerns. If so, we kindly invite the reviewer to consider raising their score. If any concerns remain, we are happy to discuss them further here.
We thank the reviewer once again for their effort in reviewing our paper and for their constructive review. As there are only a few working days left in the discussion period, we would like to ask if our response has satisfied the reviewer’s concerns. If so, we kindly invite the reviewer to consider raising their score. If any concerns remain, we are happy to discuss them further here.
We thank the reviewer once again for the effort put into reviewing our paper. As there are only a few working days left in the discussion period, we would like to ask if our response has satisfied the reviewer’s concerns. If so, we kindly invite the reviewer to consider raising their score. If any concerns remain, we are happy to discuss them further here.