SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Comparison to Harsha et al. While algorithms 1 and 2 are similar, they differ in two notable ways. First, rather than drawing a sample from  $P$ , GRC draws a sample from the restriction of  $P$  to an active subset  $S_{d} = Z_{d} \subseteq \mathcal{X}$ , namely  $P|_{S_d} / P(S_d)$ . Second, GRC updates its active subset  $S_{d} = Z_{d}$  at each step, setting it to one of the children of  $Z_{d}$ , namely either  $Z_{2d}$  or  $Z_{2d + 1}$ , by drawing  $b_{d} \sim$  Bernoulli, and setting  $Z_{2d + b_d}$ . This partitioning mechanism, which does not appear in algorithm 1, yields a
different variant of GRC for each choice of partitioning process  $Z$ . In fact, as shown in Proposition 1 below, algorithm 1 is a special case of GRC with  $S_{d} = \mathcal{X}$  for all  $d$ . See appendix A.2 for the proof.
Proposition 1 (Harsha et al. 2007) is a special case of GRC). Let  $Z$  be the global partitioning process over  $\Sigma$ , defined as
$$
Z _ {1} = \mathcal {X}, \quad Z _ {2 n} = Z _ {n}, \quad Z _ {2 n + 1} = \emptyset , \quad f o r a l l n = 1, 2, \dots . \tag {8}
$$
Harsha et al. (2007) is equivalent to GRC using this  $Z$  and setting  $C = D^{*}$  instead of  $C = I_{D^{*}}$ . We refer to this algorithm as Global GRC, or GRCG for short.
Partitioning processes and additional assumptions. While Proposition 1 shows that Harsha et al.'s algorithm is equivalent to GRC with a particular choice of  $Z$ , a range of other choices of  $Z$  is possible, and this is where we can leverage additional structure. In particular, we show that when  $Q$  and  $P$  are continuous distributions over  $\mathbb{R}$  with a unimodal density ratio  $dQ / dP$ , we can dramatically speed up GRC with an appropriate choice of  $Z$ . In particular, we will consider the on-sample and dyadic partitioning processes from Flamich et al. (2022), given in Definitions 3 and 4.
Definition 3 (On-sample partitioning process). Let  $\mathcal{X} = \mathbb{R}\cup \{-\infty ,\infty \}$  and  $P$  a continuous distribution. The on-sample partitioning process is defined as
$$
Z _ {n} = [ a, b ], a, b \in \mathcal {X} \Longrightarrow Z _ {2 n} = [ a, X _ {n} ], Z _ {2 n + 1} = \left[ X _ {n}, b \right], w h e r e X _ {n} \sim P \left| _ {Z _ {n}} / P \left(Z _ {n}\right). \right.
$$
In other words, in the on-sample partitioning process,  $Z_{n}$  are intervals of  $\mathbb{R}$ , each of which is partitioned into sub-intervals  $Z_{2n}$  and  $Z_{2n + 1}$  by splitting at the sample  $X_{n}$  drawn from  $P|_{Z_n} / P(Z_n)$ . We refer to GRC with on-sample partitioning as GRCS.
Definition 4 (Dyadic partitioning process). Let  $\mathcal{X} = \mathbb{R}\cup \{-\infty ,\infty \}$  and  $P$  a continuous distribution. The dyadic partitioning process is defined as
$$
Z _ {n} = [ a, b ], a, b \in \mathcal {X} \Rightarrow Z _ {2 n} = [ a, c ], Z _ {2 n + 1} = [ c, b ], \text {s u c h t h a t} P (Z _ {2 n}) = P (Z _ {2 n + 1}).
$$
Similar to on-sample partitioning, in the dyadic process  $Z_{n}$  are intervals of  $\mathbb{R}$ . However, in the dyadic process,  $Z_{n}$  is partitioned into sub-intervals  $Z_{2n}$  and  $Z_{2n + 1}$  such that  $P(Z_{2n}) = P(Z_{2n + 1})$ . We refer to GRC with the dyadic partitioning process as GRCD.
GRC with a tunable codelength. Flamich et al. (2022) presented a depth-limited variant of AD\* coding,  $\mathrm{DAD^{*}}$  coding, in which the codelength  $|C|$  can be provided as a tunable input to the algorithm. Fixed-codelength REC algorithms are typically approximate because they introduce bias in their samples, but are nevertheless useful in certain contexts, such as for coding a group of random variables with the same fixed codelength. GRCD can be similarly modified to accept  $|C|$  as an input, by limiting the maximum steps of the algorithm by  $D_{\mathrm{max}}$  (see algorithm 2). Setting  $D_{\mathrm{max}} = \infty$  in algorithm 2 corresponds to exact GRC, while setting  $D_{\mathrm{max}} < \infty$  corresponds to depth-limited GRC. In appendix D we provide detailed analysis of depth-limited GRC.
3.2 Theoretical results
Correctness of GRC. Below, we present three sets of assumptions on  $Q$ ,  $P$  and  $Z$ , and in theorem [1], we show that fulfilling any of them is sufficient to ensure the correctness of GRC.
Assumption 1. GRC has a finite ratio mode, i.e.  $dQ / dP(x) < M$  for all  $x \in \mathcal{X}$ , for some  $M > 0$ .
Assumption  $\mathbb{I}$  is the most generally applicable, as it does not restrict the sample space. Assumption  $\mathbb{I}$  holds for GRCG, GRCS and GRCD, so long as  $dQ / dP$  is bounded. While this assumption is very general, in some cases we may want to consider  $Q, P$  with unbounded  $dQ / dP$ . To this end, we show that it can be replaced by alternative assumptions, such as assumptions  $\mathbb{I}$  and  $\mathbb{S}$ .
Assumption 2. GRC is single-branch, i.e. for each  $d$ ,  $b_d = 0$  or  $b_d = 1$  almost surely.
GRC with the global partitioning process (eq. 8) satisfies assumption 2. In addition, if  $Q$  and  $P$  are distributions over  $\mathbb{R}$  and  $dQ / dP$  is unimodal, GRCS also satisfies assumption 2.
Assumption 3.  $\mathcal{X} \subseteq \mathbb{R}^N$  and GRC has nicely-shrinking bounds, i.e.  $\mathbb{P}$ -almost surely the following holds: for each  $x \in \mathcal{X}$  which is in a nested sequence of partitions  $x \in Z_1 \supseteq \dots \supseteq Z_{k_d} \supseteq \dots$  with  $P(Z_{k_d}) \to 0$  as  $k_d \to \infty$ , there exist  $\gamma, r_1, r_2, \ldots \in \mathbb{R}_{>0}$  such that
$$
r _ {d} \rightarrow 0, Z _ {k _ {d}} \subseteq B _ {r _ {d}} (x) a n d P \left(Z _ {k _ {d}}\right) \geq \gamma P \left(B _ {r _ {d}} (x)\right), \tag {9}
$$
where  $B_r(x)$  denotes the open ball of radius  $r$  centered on  $x$ . We recall that  $\mathbb{P}$  is the measure associated with the sequence of publicly available coin flips  $S$ .
Intuitively,  $Z$  is nicely shrinking if its branches shrink in a roughly uniform way over  $\mathcal{X}$ . As the most important example, when  $Q$  and  $P$  are distributions over  $\mathbb{R}$ , GRCD satisfies assumption 3. Now, Theorem 1 shows that if any of the above assumptions hold, then GRC terminates almost surely and yields unbiased samples from  $Q$ . We provide the proof in appendix 2.
Theorem 1 (Correctness of GRC). Suppose  $Q, P$  and  $Z$  satisfy any one of assumptions [1] to [3]. Then, algorithm [2] terminates with probability 1, and its returned sample  $X$  has law  $X \sim Q$ .
Expected runtime and codelength of GRCS. Now we turn to the expected runtime and codelength of GRCS. Theorem 2 shows that the expected codelength of GRCS is optimal, while Theorem 3 establishes that its runtime is order-optimal. We present the proofs of the theorems in appendix C.
Theorem 2 (GRCS codelength). Let  $Q$  and  $P$  be distributions over  $\mathbb{R}$  with  $D_{\mathrm{KL}}[Q \| P] < \infty$  and  $dQ / dP$  unimodal. Let  $Z$  be the on-sample partitioning process, and  $X$  its returned sample. Then,
$$
\mathbb {H} [ X | Z ] \leq D _ {\mathrm {K L}} [ Q \| P ] + 2 \log \left(D _ {\mathrm {K L}} [ Q \| P ] + 1\right) + \mathcal {O} (1). \tag {10}
$$
Theorem 3 (GRCS runtime). Let  $Q$  and  $P$  be distributions over  $\mathbb{R}$  with  $D_{\mathrm{KL}}[Q \| P] < \infty$  and  $dQ / dP$  unimodal. Let  $Z$  be the on-sample partitioning process and  $D$  the number of steps the algorithm takes before accepting a sample. Then, for  $\beta = 2 / \log(4/3) \approx 4.82$  we have
$$
\mathbb {E} [ D ] \leq \beta D _ {\mathrm {K L}} [ Q \| P ] + \mathcal {O} (1). \tag {11}
$$
Improving the codelength of GRCD. In theorem [2] we state the bound for the REC setting, where we make no further assumptions on  $Q$  and  $P$ . However, we can improve the bound if we consider the reverse channel coding (RCC) setting (Theis & Yosri 2022). In RCC, we have a pair of correlated random random variables  $X, Y \sim P_{X,Y}$ . During one round of communication, the encoder receives  $Y \sim P_{Y}$  and needs to encode a sample  $X \sim P_{X|Y}$  from the posterior using  $P_{X}$  as the proposal distribution. Thus, RCC can be thought of as the average-case version of REC, where the encoder sets  $Q \gets P_{X|Y}$  and  $P \gets P_{X}$ . In this case, when the conditions of theorem [2] hold for every  $(P_{X|Y}, P_{X})$  pair, in appendix [C] we show that the coefficient of the log-factor in eq. (10) can be improved so that the average-case bound becomes  $\mathbb{I}[X;Y] + \log (\mathbb{I}[X;Y] + 1) + \mathcal{O}(1)$ , where  $\mathbb{I}[X;Y] = \mathbb{E}_{Y \sim P_{Y}}[D_{\mathrm{KL}}[P_{X|Y} \| P_{Y}]]$  is the mutual information between  $X$  and  $Y$ .
GRCS runtime is order-optimal. Theorem [3] substantially improves upon the runtime of  $\mathrm{A}^*$  coding, which is the current fastest REC algorithm with similar assumptions. In particular,  $\mathrm{AS}^*$  coding has  $\mathcal{O}(D_{\infty}[Q \| P])$  expected runtime, which can be arbitrarily larger than that of GRCS. Remarkably, the runtime of GRCS is optimal up to the multiplicative factor  $\beta$ .
We conducted two sets of experiments: one on controlled synthetic REC problems to check the predictions of our theorems numerically, and another using VAEs trained on MNIST to study how the performance of GRC-based compression pipelines can be improved in practice. We conducted all our experiments under fair and reproducible conditions. Our code is available at https://github.com/cambridge-mlg/fast-rec-with-grc.
4.1 Synthetic Experiments
Synthetic REC experiments. First, we compare GRCS and GRCD, against  $\mathrm{AS^{*}}$  and  $\mathrm{AD^{*}}$  coding, on a range of synthetic REC problems. We systematically vary distribution parameters to adjust the difficulty of the REC problems. Figure 4 shows the results of our synthetic experiments.
Partitioning processes improve the runtime of GRC. First, we observe that, assuming that  $dQ / dP$  is unimodal, introducing the on-sample or the dyadic partitioning process speeds up GRC dramatically. In particular, fig. 4 shows that increasing the infinity divergence  $D_{\infty}[Q\| P]$  (for a fixed  $D_{\mathrm{KL}}[Q\| P]$ ) does not affect the runtimes of GRCS and GRCD, which remain constant and small. This is a remarkable speed-up over the exponential expected runtime of GRCG.
GRC is faster than  $\mathbf{A}^*$  coding. Further, we observe that GRC significantly improves upon the runtime of  $\mathbf{A}^*$  coding, which is the fastest previously known algorithm with similar assumptions. In particular, Figure 4 shows that increasing the infinity divergence  $D_{\infty}[Q\| P]$ , while keeping the
Figure 4: Comparison between GRC and  $\mathbf{A}^*$  coding on synthetic REC problems with Gaussian  $Q$  and  $P$ . Left: we fix  $D_{\mathrm{KL}}[Q\| P] = 3$  and vary  $D_{\infty}[Q\| P]$ , measuring the number of steps taken by each algorithm. Right: we fix  $D_{\infty}[Q\| P] = D_{\mathrm{KL}}[Q\| P] + 2$  and vary  $D_{\mathrm{KL}}[Q\| P]$ , plotting the codelengths produced by each algorithm. Reported codelengths do not include additional logarithmic overhead terms. Results are averaged over  $4\times 10^{3}$  different random seeds for each datapoint. We have included error-bars in both plots but these are too small to see compared to the plot scales.
KL divergence  $D_{\mathrm{KL}}[Q \| P]$  fixed, increases the runtime of both  $\mathrm{AS}^*$  and  $\mathrm{AD}^*$  coding, while the runtimes of GRCS and GRCD remain constant. More generally, for a fixed KL divergence, the infinity divergence can be arbitrarily large or even infinite. In such cases,  $\mathrm{A}^*$  coding would be impractically slow or even inapplicable, while GRCS and GRCD remain fast.
4.2 Compression with Variational Autoencoders
Compressing images with VAEs and REC. One of the most promising applications of REC is in learnt compression. Here, we implement a proof-of-concept lossless neural compression pipeline using a VAE with a factorized Gaussian posterior on MNIST and take the architecture used by Townsend et al. (2018). To compress an image  $Y$ , we encode a latent sample  $X$  from the VAE posterior  $q(X \mid Y)$  by applying GRCD dimensionwise after which we encode the image  $Y$  with entropy coding using the VAE's conditional likelihood  $p(Y \mid X)$  as the coding distribution. Unfortunately, in addition to the  $D_{\mathrm{KL}}[q(X_d \mid Y) \| p(X_d)]$  bits coding cost for latent dimension  $d$ , this incurs an overhead of  $\log(D_{\mathrm{KL}}[q(X_d \mid Y) \| p(X_d)] + 1) + O(1)$  bits, analogously to how a symbol code, like Huffman coding, incurs a constant overhead per symbol (MacKay, 2003). However, since  $\log(1 + x) \approx x$  when  $x \approx 0$ , the logarithmic overhead of GRC can become significant compared to the KL divergence. Hence, we now investigate two approaches to mitigate this issue.
Modified ELBO for REC. A principled approach to optimizing our neural compression pipeline is to minimize its expected codelength. For bits-back methods (Townsend et al., 2018, 2019), the negative ELBO indeed expresses their expected codelength, but in REC's case, it does not take into account the additional dimensionwise logarithmic overhead we discussed above. Thus, we propose to minimize a modified negative ELBO to account for this (assuming that we have  $D$  latent dimensions):
$$
\underbrace {\mathbb {E} _ {X \sim q (X \mid Y)} [ - \log p (Y \mid X) ] + D _ {\mathrm {K L}} [ q (X \mid Y) \| p (X) ]} _ {\text {R e g u l a r E L B O}} + \sum_ {d = 1} ^ {D} \underbrace {\log \left(D _ {\mathrm {K L}} [ q (X _ {d} \mid Y) \| p (X _ {d}) ] + 1\right)} _ {\text {L o g a r i t h m i c o v e r h e a d p e r d i m e n s i o n}}. \tag {12}
$$
Coding the latent indices. As the final step during the encoding process, we need a prefix code to encode the heap indices  $I_{d}$  returned by GRCD for each  $d$ . Without any further information, the best we can do is use Elias  $\delta$  coding (Elias 1975), which, assuming our conjecture on the expected runtime of GRCD holds, yields an expected codelength of  $\mathbb{I}[Y;X] + 2\log (\mathbb{I}[Y;X] + 1) + \mathcal{O}(1)$ . However, we can improve this if we can estimate  $\mathbb{E}[\log I_d]$  for each  $d$ : it can be shown, that the maximum entropy distribution of a positive integer-valued random variable with under a constraint on the expectation on its logarithm is  $\zeta (n|\lambda)\propto n^{-\lambda}$ , with  $\lambda^{-1} = \mathbb{E}[\log I_d] + 1$ . In this case, entropy coding  $I_{d}$  using this  $\zeta$  distribution yields improves the expected codelength to  $\mathbb{I}[Y;X] + \log (\mathbb{I}[Y;X] + 1) + \mathcal{O}(1)$ .
Experimental results. We trained our VAE with  $L \in \{20, 50, 100\}$  latent dimensions optimized using the negative ELBO and its modified version in Equation (12), and experimented with encoding the heap indices of GRCD with both  $\delta$  and  $\zeta$  coding. We report the results of our in Table 1 on the MNIST test set in bits per pixel. In addition to the total coding cost, we report the negative ELBO per pixel, which is the fundamental lower bound on the compression efficiency of REC with each
TRAINING OBJECTIVE	# LATENT	TOTAL BPP WITH ζ CODING	TOTAL BPP WITH δ CODING	NEG. ELBO PER Pixel	OVERHEAD BPP WITH δ CODING
ELBO	20	1.472 ± 0.004	1.482 ± 0.004	1.391 ± 0.004	0.091 ± 0.000
50	1.511 ± 0.003	1.530 ± 0.003	1.357 ± 0.003	0.172 ± 0.000
100	1.523 ± 0.003	1.600 ± 0.003	1.362 ± 0.003	0.238 ± 0.000
MODIFIED ELBO	20	1.470 ± 0.004	1.478 ± 0.004	1.393 ± 0.004	0.085 ± 0.000
50	1.484 ± 0.003	1.514 ± 0.003	1.373 ± 0.003	0.141 ± 0.000
100	1.485 ± 0.003	1.579 ± 0.003	1.373 ± 0.003	0.205 ± 0.000
Table 1: Lossless compression performance comparison on the MNIST test set of a small VAE with different latent space sizes, optimized using either the ELBO or the modified ELBO in eq. (12). We report the bits per pixel (BPP) attained using different coding methods, averaged over the 10,000 test images, along with the standard error, using GRCD. See section 4.2 for further details.
VAE. Finally, we report the logarithmic overhead due to  $\delta$  coding. We find that both the modified ELBO and  $\zeta$  coding prove beneficial, especially as the dimensionality of the latent space increases. This is expected, since the overhead is most significant for latent dimensions with small KLs, which becomes more likely as the dimension of the latent space grows. The improvements yielded by each of the two methods are significant, with  $\zeta$  coding leading to a consistent  $1 - 7\%$  gain compared to  $\delta$  coding and the modified objective resulting in up to  $2\%$  gain in coding performance.
Summary. In this work, we introduced Greedy Rejection Coding (GRC), a REC algorithm which generalises the rejection algorithm of Harsha et al. to arbitrary probability spaces and partitioning processes. We proved the correctness of our algorithm under mild assumptions, and introduced GRCS and GRCD, two variants of GRC. We showed that the runtimes of GRCS and GRCD significantly improve upon the runtime of  $\mathrm{A}^*$  coding, which can be arbitrarily larger. We evaluated our algorithms empirically, verifying our theory and conducted a proof-of-concept learnt compression experiment on MNIST using VAEs. We demonstrated that a principled modification to the ELBO and entropy coding GRCD's indices using a  $\zeta$  distribution can further improve compression efficiency.
Limitations and Further work. One limitation of GRC is that, unlike  $\mathrm{A}^*$  coding, it requires us to be able to evaluate the CDF of  $Q$ . While in some settings this CDF may be intractable, this assumption is satisfied by most latent variable generative models, and is not restrictive in practice. However, one practical limitation of GRCS and GRCD, as well as  $\mathrm{AS}^*$  and  $\mathrm{AD}^*$ , is that they assume target-proposal pairs over  $\mathbb{R}$ . For multivariate distributions, we can decompose them into univariate conditionals and apply GRC dimensionwise, however this incurs an additional coding overhead per dimension, resulting in a non-negligible cost. Thus, an important direction is to investigate whether fast REC algorithms for multivariate distributions can be devised, to circumvent this challenge.
6 Author Contributions
SM suggested that single-branch greedy rejection coding (Flamich & Theis, 2023) can be extended to arbitrary partitioning processes. GF discovered GRCS and SM discovered GRCD. SM provided a proof for Theorem  $\boxed{1}$  and GF provided proofs for Theorems  $\boxed{2}$  and  $\boxed{3}$ . SM carried out the synthetic experiments and GF carried out the VAE experiments. SM drafted the majority of the paper, while GF wrote section 4.2 and parts of sections 2 and 3.2. They contributed to the editing of the paper equally. JMH supervised and steered the project.
GF acknowledges funding from DeepMind. SM acknowledges funding from the Qualcomm Innovation Fellowship and the Vice Chancellor's & George and Marie Vergottis scholarship of the Cambridge Trust.
Eirikur Agustsson and Lucas Theis. Universally quantized neural compression. In Advances in Neural Information Processing Systems, 2020.
Johannes Ballé, Philip A. Chou, David Minnen, Saurabh Singh, Nick Johnston, Eirikur Agustsson, Sung Jin Hwang, and George Toderici. Nonlinear transform coding. IEEE Journal of Selected Topics in Signal Processing, 2020.
Rewon Child. Very deep VAEs generalize autoregressive models and can outperform them on images. In International Conference on Learning Representations, 2020.
Richard M. Dudley. Real analysis and probability. CRC Press, 2018.
Nelson Dunford and Jacob T. Schwartz. Linear operators, part 1: general theory, volume 10. John Wiley & Sons, 1988.
Peter Elias. Universal codeword sets and representations of the integers. IEEE Transactions on Information Theory, 1975.
Gergely Flamich. Greedy Poisson rejection sampling. In Advances in Neural Information Processing Systems, 2023.
Gergely Flamich and Lucas Theis. Adaptive greedy rejection sampling. In IEEE International Symposium on Information Theory, 2023.
Gergely Flamich, Marton Havasi, and José Miguel Hernández-Lobato. Compressing images by encoding their latent representations with relative entropy coding. In Advances in Neural Information Processing Systems, 2020.
Gergely Flamich, Stratis Markou, and José Miguel Hernández-Lobato. Fast relative entropy coding with A* coding. In International Conference on Machine Learning, 2022.
Zongyu Guo, Gergely Flamich, Jiajun He, Zhibo Chen, and José Miguel Hernández-Lobato. Compression with Bayesian implicit neural representations. In Advances in Neural Information Processing Systems, 2023.
Prahladh Harsha, Rahul Jain, David McAllester, and Jaikumar Radhakrishnan. The communication complexity of correlation. In Twenty-Second Annual IEEE Conference on Computational Complexity (CCC'07). IEEE, 2007.
Marton Havasi, Robert Peharz, and José Miguel Hernández-Lobato. Minimal random code learning: Getting bits back from compressed model parameters. In International Conference on Learning Representations, 2018.
Jiajun He, Gergely Flamich, Zongyu Guo, and José Miguel Hernández-Lobato. RECOMBINER: Robust and enhanced compression with Bayesian implicit neural representations. arXiv preprint arXiv:2309.17182, 2023.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in neural information processing systems, 2020.
Durk P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. In Advances in neural information processing systems, 2016.
John F. C. Kingman. Poisson Processes. Oxford Studies in Probability. Clarendon Press, 1992.
Cheuk Ting Li and Abbas El Gamal. Strong functional representation lemma and applications to coding theorems. IEEE Transactions on Information Theory, 2018.
David J. C. MacKay. Information theory, inference and learning algorithms. Cambridge university press, 2003.
Chris J. Maddison. Poisson process model for Monte Carlo. *Perturbation, Optimization, and Statistics*, 2016.