Title: Faster Relative Entropy Coding with Greedy Rejection Coding
Venue: NeurIPS 2023
Accepted: yes
Keywords: Compression

Learnt Compression

Relative Entropy Coding

Information Theory
Abstract: Relative entropy coding (REC) algorithms encode a sample from a target distribution $Q$ using a proposal distribution $P$ using as few bits as possible. Unlike entropy coding, REC does not assume discrete distributions and require quantisation.
As such, it can be naturally integrated into communication pipelines such as learnt compression and differentially private federated learning. Unfortunately, despite their practical benefits, REC algorithms have not seen widespread application, due to their prohibitively slow runtimes or restrictive assumptions. In this paper, we make progress towards addressing these issues. We introduce Greedy Rejection Coding (GRC), which generalises the rejection sampling-based algorithm of Harsha et al. (2007) to arbitrary probability spaces and partitioning schemes. We first show that GRC terminates almost surely and returns unbiased samples from $Q$, and then focus on two variants of GRC, namely GRCS and GRCD. We show that for continuous $Q$ and $P$ over $\mathbb{R}$ with unimodal $dQ/dP$, the expected runtime of GRCS is upper bounded by $\beta D_{KL}(Q||P) + \mathcal{O}(1)$ where $\beta \approx 4.82$, and its expected codelength is optimal. This makes GRCS the first REC algorithm with guaranteed optimal runtime for this class of distributions, up to the multiplicative constant $\beta$. This significantly improves upon the previous state-of-the-art method, A* coding (Flamich et al., 2022). Under the same assumptions, we experimentally observe and conjecture that the expected runtime and codelength of GRCD are upper bounded by $D_{KL}(Q||P) + \mathcal{O}(1)$. Finally, we evaluate GRC in a compression pipeline with variational autoencoders on MNIST, and show that a modified training objective and a codelength-compression method can further improve compression efficiency.

Faster Relative Entropy Coding with Greedy Rejection Coding

Abstract

Relative entropy coding (REC) algorithms encode a sample from a target distribution Q using a proposal distribution P using as few bits as possible. Unlike entropy coding, REC does not assume discrete distributions or require quantisation. As such, it can be naturally integrated into communication pipelines such as learnt compression and differentially private federated learning. Unfortunately, despite their practical benefits, REC algorithms have not seen widespread application, due to their prohibitively slow runtimes or restrictive assumptions. In this paper, we make progress towards addressing these issues. We introduce Greedy Rejection Coding (GRC), which generalises the rejection based-algorithm of Harsha et al. (2007) to arbitrary probability spaces and partitioning schemes. We first show that GRC terminates almost surely and returns unbiased samples from Q , after which we focus on two of its variants: GRCS and GRCD. We show that for continuous Q and P over R with unimodal density ratio dQ/dP , the expected runtime of GRCS is upper bounded by β D KL [ Q ‖ P ] + O (1) where β ≈ 4 . 82 , and its expected codelength is optimal. This makes GRCS the first REC algorithm with guaranteed optimal runtime for this class of distributions, up to the multiplicative constant β . This significantly improves upon the previous state-of-the-art method, A* coding (Flamich et al., 2022). Under the same assumptions, we experimentally observe and conjecture that the expected runtime and codelength of GRCD are upper bounded by D KL [ Q ‖ P ] + O (1) . Finally, we evaluate GRC in a variational autoencoderbased compression pipeline on MNIST, and show that a modified ELBO and an index-compression method can further improve compression efficiency.

1 Introduction and motivation

Over the past decade, the development of excellent deep generative models (DGMs) such as variational autoencoders (VAEs; Vahdat & Kautz, 2020; Child, 2020), normalising flows (Kingma et al., 2016) and diffusion models (Ho et al., 2020) demonstrated great promise in leveraging machine learning (ML) for data compression. Many recent learnt compression approaches have significantly outperformed the best classical hand-crafted codecs across a range of domains, such as lossless and lossy image or video compression (Zhang et al., 2021; Mentzer et al., 2020, 2022).

Transform coding. Most learnt compression algorithms are transform coding methods: they first map a datum to a latent variable using a learnt transform, and encode it using entropy coding (Ballé et al., 2020). Entropy coding assumes discrete variables while the latent variables in DGMs are typically continuous, so most transform coding methods quantize the latent variable prior to entropy coding. Unfortunately, quantization has zero derivative almost everywhere. Thus, state-of-the-art DGMs trained with gradient-based optimisation must resort to some continuous approximation to

AS

∗

coding

Global A

∗

AD

∗

coding

GRCD

GRCS

GRCG

On-sample partitioning

Global partitioning

Dyadic partitioning

Perturb-and-maximize

Rejection sampling

quantisation during training and switch to hard quantisation for compression. Previous works have argued that using quantisation within learnt compression is restrictive or otherwise harmful, and that a method which naturally interfaces with continuous latent variables is needed (Havasi et al., 2018; Flamich et al., 2020; Theis & Agustsson, 2021; Flamich et al., 2022).

Relative entropy coding. In this paper, we study an alternative to quantizing and entropy coding the latent representations: we consider randomly perturbing them and encoding the perturbed representations instead. This problem is called relative entropy coding (REC; Havasi et al., 2018; Flamich et al., 2020), and formally, it requires that we encode a random sample from a target Q (the distribution of the perturbed latent representations) using a coding distribution P and some publicly available randomness S (e.g. a publicly shared PRNG seed). Remarkably, there exist algorithms which encode a target sample using only approximately D KL [ Q ‖ P ] -many bits on average and, notably, allow Q and P to be continuous (Li & El Gamal, 2018; Flamich et al., 2022; Flamich & Theis, 2023; Flamich, 2023). Thus, we can forgo quantization in our compression pipeline and learn our transforms end-to-end using gradient descent and the reparameterization trick. Moreover, REC has fundamental advantages over quantization in lossy compression with realism constraints (Theis &Agustsson, 2021; Theis et al., 2022) and has a range of other applications, such as differentially private compression for federated learning (Shah et al., 2022; Shahmiri et al., 2023) and compression with Bayesian implicit neural representations (Guo et al., 2023; He et al., 2023).

Limitations of existing REC algorithms. While algorithms for solving REC problems already exist, most of them suffer from limitations that render them impractical. These limitations fall into three categories: 1) prohibitively long runtimes, 2) overly restrictive assumptions, or 3) excessive coding overheads. In this work, we study and make progress towards addressing these limitations.

General-purpose REC algorithms. The minimal assumption made by any REC algorithm on Q and P is that D KL [ Q ‖ P ] < ∞ , which ensures that it produces a finite code. However, unfortunately, Agustsson & Theis (2020) showed that without additional assumptions on Q and P , the expected runtime of any algorithm that simulates a Q -distributed sample using P must scale at least exponentially in D KL [ Q ‖ P ] in the worst case, which is impractically slow for most practical problems. This result holds even in the approximate sense, i.e. the exponential runtime persists even if the algorithm simulates a ˜ Q -distributed sample with || Q -˜ Q || TV < 1 / 12 .

Faster algorithms with additional assumptions. On the other hand, there exist algorithms which make additional assumptions in order to achieve faster runtimes. For example, dithered quantization (Ziv, 1985; Agustsson & Theis, 2020) achieves an expected runtime of D KL [ Q ‖ P ] , which is optimal since any REC algorithm has an expected runtime of at least D KL [ Q ‖ P ] . However, it requires both Q and P to be uniform distributions, which limits its applicability. Recently, Flamich et al. (2022) introduced A ∗ coding, an algorithm based on A ∗ sampling (Maddison et al., 2014) which, under assumptions satisfied in practice, achieves an expected runtime of D ∞ [ Q ‖ P ] . Unfortunately, this runtime is suboptimal and is not always practically fast, since D ∞ [ Q ‖ P ] can be arbitrarily large for fixed D KL [ Q ‖ P ] . Further, as discussed in Flamich et al. (2022) this runtime also comes at a cost of an additional, substantial, overhead in codelength, which limits the applicability of A ∗ coding.

Our contributions. In this work, we address some of these limitations. First, we propose greedy rejection coding (GRC), a REC algorithm based on rejection sampling. Then, inspired by A* coding (Flamich et al., 2022), we develop GRCS and GRCD, two variants of GRC that partition the sample space to dramatically speed up termination. Figure 1 illustrates the relations between GRC and its variants with existing algorithms. We analyze the correctness and the runtime of these algorithms and, in particular, prove that GRCS has an optimal codelength and order-optimal runtime on a wide class of one-dimensional problems. In more detail, our contributions are

We introduce Greedy Rejection Coding (GRC), which generalises the algorithm of Harsha et al. (2007) to arbitrary probability spaces and partitioning schemes. We prove that under mild conditions, GRC terminates almost surely and returns an unbiased sample from Q .

We introduce GRCS and GRCD, two variants of GRC for continuous distributions over R , which adaptively partition the sample space to dramatically improve their convergence, inspired by AS ∗ and AD ∗ coding (Flamich et al., 2022), respectively.

We prove that whenever dQ/dP is unimodal, the expected runtime and codelength of GRCS is O ( D KL [ Q ‖ P ]) . This significantly improves upon the O ( D ∞ [ Q ‖ P ]) runtime of AS ∗ coding, which is always larger than that of GRCS. This runtime is order-optimal, while making far milder assumptions than, for example, dithered quantization.

We provide clear experimental evidence for and conjecture that whenever dQ/dP is unimodal, the expected runtime and codelength of GRCD are D KL [ Q ‖ P ] . This also significantly improves over the D ∞ [ Q ‖ P ] empirically observed runtime of AD ∗ coding.

We implement a compression pipeline with VAEs, using GRC to compress MNIST images. We propose a modified ELBO objective and show that this, together with a practical method for compressing the indices returned by GRC further improve compression efficiency.

2 Background and related work

Relative entropy coding. First, we define REC algorithms. Definition 1 is stricter than the one given by Flamich et al. (2022), as it has a stronger condition on the the expected codelength of the algorithm. In this paper, all logarithms are base 2, and all divergences are measured in bits.

Definition 1 (REC algorithm) . Let ( X , Σ ) be a measurable space, let R be a set of pairs of distributions ( Q,P ) over ( X , Σ ) such that D KL [ Q ‖ P ] < ∞ and P be the set of all distributions P such that ( Q,P ) ∈ R for some distribution Q . Let S = ( S 1 , S 2 , . . . ) be a publicly available sequence of independent and fair coin tosses, with corresponding probability space ( S , F , P ) and let C = { 0 , 1 } ∗ be the set of all finite binary sequences. A REC algorithm is a pair of functions enc : R × S → C and dec : C × P × S → X , such that for each ( Q,P ) ∈ R and S ∼ P , the outputs of the encoder C = enc ( Q,P,S ) and the decoder X = dec ( P, C, S ) satisfy

where | C | is the length of the string C . We call enc the encoder and dec the decoder.

In practice, S is implemented with a pseudo-random number generator (PRNG) with a public seed. In the remainder of this section, we discuss relevant REC algorithms, building up to GRC in section 3.

Existing REC algorithms. While there are many REC algorithms already, they suffer from various issues limiting their applicability in practice. Our proposed algorithm, Greedy Rejection Coding (GRC), is based on and generalises Harsha et al. (2007)'s REJ-SAMPLER, by drawing inspiration from A ∗ coding (Flamich et al., 2022). Specifically, A ∗ coding generalises Li & El Gamal (2018)'s Poisson functional representation by introducing a partitioning scheme to speed up the algorithm's termination. In an analogous fashion, GRC generalises Harsha et al. (2007) by also introducing partitioning schemes to speed up termination and achieve optimal runtimes. Here we discuss relevant algorithms, building up to GRC in section 3.

REC with rejection sampling. In this work, we generalize the rejection sampler introduced by Harsha et al. (2007). While they presented the algorithm for discrete Q and P originally, we generalise it to arbitrary probability spaces in this section and further extend it to arbitrary partitioning schemes (see definition 5) in section 3. The generalisation to arbitrary probability spaces relies on

the Radon-Nikodym derivative dQ/dP , which is guaranteed to exist since Q /lessmuch P by definition 1. When Q and P both have densities, dQ/dP coincides with the density ratio.

At each step, the algorithm draws a sample from P and performs an accept-reject step, as illustrated in fig. 2. If it rejects the sample, it rules out part of Q corresponding to the acceptance region, adjusts the proposal to account for the removed mass, and repeats until acceptance. More formally, define T 0 to be the zero-measure on X , and recursively for d ∈ N , set:

for all x ∈ X , S ∈ Σ . The algorithm terminates at the first occurrence of U d ≤ β d +1 ( X d ) . The T d measure corresponds to the mass that has been ruled off up to and including the d th rejection: T 1 ( X ) , T 2 ( X ) and T 3 ( X ) are the sums of the blue and green masses in the left, middle and right plots of fig. 2 respectively. The A d measure corresponds to the acceptance mass at the d th step: A 1 ( X ) , A 2 ( X ) and A 3 ( X ) are the masses of the green regions in the left, middle and right plots of fig. 2 respectively. Lastly, t d , α d are the Radon-Nikodym derivatives i.e., roughly speaking, the densities, of T d , A d with respect to P , and β d +1 ( X d ) is the probability of accepting the sample X d .

To encode the accepted sample X , enc outputs the number of rejections C that occurred before acceptance. To decode X from C, dec draws C +1 samples from P, using the same PRNG seed as the encoder, and returns the last sample. While this algorithm is elegantly simple and achieves optimal codelengths, Flamich & Theis (2023) showed its expected runtime is 2 D ∞ [ Q ‖ P ] , where D ∞ [ Q ‖ P ] = sup x ∈ X log( dQ/dP )( x ) is the Rényi ∞ -divergence. Unfortunately, this runtime is prohibitively slow in most practical cases.

REC with Poisson & Gumbel processes. Li & El Gamal (2018) introduced a REC algorithm based on Poisson processes, referred to as Poisson Functional Representation (PFR). PFR assumes that dQ/dP is bounded above, and relies on the fact that (Kingman, 1992), if T n are the ordered arrival times of a homogeneous Poisson process on R + and X n ∼ P , then

Therefore, PFR casts the REC problem into an optimisation, or search, problem, which can be solved in finite time almost surely. The PFR encoder draws pairs of samples T n , X n , until it solves the search problem in eq. (5), and returns X = X N , C = N -1 . The decoder can recover X N from ( P, C, S ) , by drawing N samples from P , using the same random seed, and keeping the last sample. While, like the algorithm of Harsha et al. (2007), PFR is elegantly simple and achieves optimal codelengths, its expected runtime is also 2 D ∞ [ Q ‖ P ] (Maddison, 2016).

Fast REC requires additional assumptions. These algorithms' slow runtimes are perhaps unsurprising considering Agustsson & Theis's result, which shows under the computational hardness assumption RP = NP that without making additional assumptions on Q and P , there is no REC

/negationslash

Algorithm 1 Harsha et al.'s rejection algorithm; equivalent to GRC with a global partition

Algorithm 2 GRC with partition process Z ; differences to Harsha et al.'s algorithm shown in green

algorithm whose expected runtime scales polynomially in D KL [ Q ‖ P ] . Therefore, in order achieve faster runtimes, a REC algorithm must make additional assumptions on Q and P .

A ∗ coding. To this end, Flamich et al. (2022) proposed: (1) a set of appropriate assumptions which are satisfied by many deep latent variable models in practice and (2) a REC algorithm, referred to as A ∗ coding, which leverages these assumptions to achieve a substantial speed-up over existing methods. In particular, A ∗ coding generalizes PFR by introducing a partitioning scheme, which splits the sample space X in nested partitioning subsets, to speed up the solution of eq. (5). Drawing inspiration from this, our proposed algorithm generalises eqs. (2) to (4) in an analogous manner (see fig. 1), introducing partitioning processes (definition 2) to speed up the algorithm's termination.

Definition 2 (Partitioning process) . A partitioning process is a process Z : N + → Σ such that

In other words, a partitioning process Z is an infinite binary tree-structured process, where the root node is Z 1 = X and has index 1, and each node Z n with index n is partitioned by its two children nodes Z 2 n , Z 2 n +1 with indices 2 n and 2 n +1 . We refer to this system of indexing as heap indexing . In section 3 we present specific choices of partitioning processes which dramatically speed up GRC.

Greedy Poisson Rejection Sampling. Contemporary to our work, Flamich (2023) introduces a rejection sampler based on Poisson processes, called Greedy Poisson Rejection Sampling (GPRS), which can be used as a REC algorithm. Similar to GRC and A* coding, GPRS partitions the sample space to speed up the convergence to the accepted sample. Furthermore, a variant of GPRS also achieves order-optimal runtime for one-dimensional distribution pairs with a unimodal density ratio. However, the construction of their method is significantly different from ours, relying entirely on Poisson processes. Moreover, GPRS requires numerically solving a certain ODE, while our method does not, making it potentially more favourable in practice. We believe establishing a closer connection between GPRS and GRC is a promising future research direction.

3 Greedy Rejection Coding

Generalising Harsha et al. (2007). In this section we introduce Greedy Rejection Coding (GRC; definition 5), which generalises the algorithm of Harsha et al. (2007) in two ways. First, GRC can be used with distributions over arbitrary probability spaces. Therefore, it is applicable to arbitrary REC problems, including REC with continuous distributions. Second, similar to A ∗ coding, GRC can be combined with arbitrary partitioning processes, allowing it to achieve optimal runtimes given additional assumptions on the REC problem, and an appropriate choice of partitioning process.

(a) Sample & accept or reject

(b) Partition & sample

b

1

∈

{

0

,

1

}

(c) Sample & accept or reject

(d) Sample & accept or reject

(e) Partition & sample

b

1

0

,

1

(f) Sample & accept or reject

∈

{

}

3.1 Algorithm definition

Overview. Before specifying GRC, we summarise its operation with an accompanying illustration. On a high level, GRC interleaves accept-reject steps with partitioning steps, where the latter are determined by a partitioning process. Specifically, consider the example in figs. 3d to 3f, where Q and P are distributions over X = [0 , 1] , and Z is the partitioning process defined by

In each step d = 1 , 2 , . . . , GRC maintains a heap index I d of an infinite binary tree, and an active subset S d = Z I d ⊆ X of the sample space, initialised as I 0 = 1 and S 1 = Z 1 = X respectively.

Accept-reject step. In each step, GRC draws a sample from the restriction of P to S d , namely P | S d /P ( S d ) , and either accepts or rejects it. If the sample is accepted, the algorithm terminates. Otherwise, GRC performs a partitioning step as shown in fig. 3d

Partitioning step. In each partitioning step, GRC partitions S d = Z I d into Z 2 I d and Z 2 I d +1 , as specified by the partitioning process Z . It then samples a Bernoulli random variable b d , whose outcomes have probabilities proportional to the mass of Q which has not been accounted for, up to and including step d , within the partitions Z 2 I d and Z 2 I d +1 respectively. In fig. 3e, these two masses correspond to the purple and orange areas, and the algorithm has sampled b d = 1 . Last, GRC updates the heap index to I d +1 = 2 I d + b d and the active subset to S d +1 = Z I d +1 . GRC proceeds by interleaving accept-reject and partitioning steps until an acceptance occurs.

Algorithm specification. The aforementioned algorithm can be formalised in terms of probability measures over arbitrary spaces and arbitrary partitioning processes. Above, algorithms 1 and 2 describe Harsha et al.'s rejection sampler and our generalisation of it, respectively. For the sake of keeping the exposition lightweight, we defer the formal measure-theoretic definition of GRC to the appendix (see definition 5 in appendix A.1), and refer to algorithm 2 as a working definition here.

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 ⊆ 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 2 d or Z 2 d +1 , by drawing b d ∼ Bernoulli, and setting Z 2 d + 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 = for all d . See appendix A.2 for the proof.

X

Proposition 1 (Harsha et al. (2007) is a special case of GRC) . Let Z be the global partitioning process over Σ , defined as

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 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 X = R ∪ { -∞ , ∞ } and P a continuous distribution. The on-sample partitioning process is defined as

In other words, in the on-sample partitioning process, Z n are intervals of R , each of which is partitioned into sub-intervals Z 2 n and Z 2 n +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 X = R ∪ { -∞ , ∞ } and P a continuous distribution. The dyadic partitioning process is defined as

Similar to on-sample partitioning, in the dyadic process Z n are intervals of R . However, in the dyadic process, Z n is partitioned into sub-intervals Z 2 n and Z 2 n +1 such that P ( Z 2 n ) = P ( Z 2 n +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, 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 max (see algorithm 2). Setting D max = ∞ in algorithm 2 corresponds to exact GRC, while setting D max < ∞ 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 ∈ X , for some M > 0 .

Assumption 1 is the most generally applicable, as it does not restrict the sample space. Assumption 1 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 2 and 3.

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 R and dQ/dP is unimodal, GRCS also satisfies assumption 2.

Assumption 3. X ⊆ R N and GRC has nicely-shrinking bounds, i.e. P -almost surely the following holds: for each x ∈ X which is in a nested sequence of partitions x ∈ Z 1 ⊇ · · · ⊇ Z k d ⊇ . . . with P ( Z k d ) → 0 as k d →∞ , there exist γ , r 1 , r 2 , ... ∈ R > 0 such that

where B r ( x ) denotes the open ball of radius r centered on x. We recall that 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 X . As the most important example, when Q and P are distributions over 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 B.

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 ∼ 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 R with D KL [ Q ‖ P ] < ∞ and dQ/dP unimodal. Let Z be the on-sample partitioning process, and X its returned sample. Then,

Theorem 3 (GRCS runtime) . Let Q and P be distributions over R with D KL [ Q ‖ P ] < ∞ 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 β = 2 / log(4 / 3) ≈ 4 . 82 we have

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 ∼ P X,Y . During one round of communication, the encoder receives Y ∼ P Y and needs to encode a sample X ∼ 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 ← P X | Y and P ← 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 I [ X ; Y ] + log( I [ X ; Y ] + 1) + O (1) , where I [ X ; Y ] = E Y ∼ P Y [ D 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 A ∗ coding, which is the current fastest REC algorithm with similar assumptions. In particular, AS ∗ coding has O ( D ∞ [ 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 β .

4 Experiments

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 AS ∗ and 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 ∞ [ Q ‖ P ] (for a fixed D 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 A ∗ coding. Further, we observe that GRC significantly improves upon the runtime of A* coding, which is the fastest previously known algorithm with similar assumptions. In particular, Figure 4 shows that increasing the infinity divergence D ∞ [ Q ‖ P ] , while keeping the

KL divergence D KL [ Q ‖ P ] fixed, increases the runtime of both AS ∗ and 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, 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 | Y ) by applying GRCD dimensionwise after which we encode the image Y with entropy coding using the VAE's conditional likelihood p ( Y | X ) as the coding distribution. Unfortunately, in addition to the D KL [ q ( X d | Y ) ‖ p ( X d )] bits coding cost for latent dimension d , this incurs an overhead of log( D KL [ q ( X d | 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 ) ≈ x when x ≈ 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):

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 δ coding (Elias, 1975), which, assuming our conjecture on the expected runtime of GRCD holds, yields an expected codelength of I [ Y ; X ] +2 log( I [ Y ; X ] +1)+ O (1) . However, we can improve this if we can estimate 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 ζ ( n | λ ) ∝ n -λ , with λ -1 = E [log I d ] + 1 . In this case, entropy coding I d using this ζ distribution yields improves the expected codelength to I [ Y ; X ] + log( I [ Y ; X ] + 1) + O (1) .

Experimental results. We trained our VAE with L ∈ { 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 δ and ζ 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

VAE. Finally, we report the logarithmic overhead due to δ coding. We find that both the modified ELBO and ζ 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 ζ coding leading to a consistent 1 -7% gain compared to δ coding and the modified objective resulting in up to 2% gain in coding performance.

5 Conclusion and Future Work

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 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 ζ distribution can further improve compression efficiency.

Limitations and Further work. One limitation of GRC is that, unlike 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 AS ∗ and AD ∗ , is that they assume target-proposal pairs over 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 1 and GF provided proofs for Theorems 2 and 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.

References

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.

Chris J. Maddison, Daniel Tarlow, and Tom Minka. A* sampling. In Advances in Neural Information Processing Systems , 2014.

Fabian Mentzer, George Toderici, Michael Tschannen, and Eirikur Agustsson. High-fidelity generative image compression. In Advances in Neural Information Processing Systems , 2020.

Fabian Mentzer, George Toderici, David Minnen, Sergi Caelles, Sung Jin Hwang, Mario Lucic, and Eirikur Agustsson. VCT: A video compression transformer. In Advances in Neural Information Processing Systems , 2022.

Fernando Pérez-Cruz. Kullback-Leibler divergence estimation of continuous distributions. In IEEE International Symposium on Information Theory , 2008.

Walter Rudin. Real and Complex Analysis . McGraw-Hill, 1986.

Abhin Shah, Wei-Ning Chen, Johannes Balle, Peter Kairouz, and Lucas Theis. Optimal compression of locally differentially private mechanisms. In International Conference on Artificial Intelligence and Statistics . PMLR, 2022.

Ali Moradi Shahmiri, Chih Wei Ling, and Cheuk Ting Li. Communication-efficient Laplace mechanism for differential privacy via random quantization. arXiv preprint arXiv:2309.06982 , 2023.

Lucas Theis and Eirikur Agustsson. On the advantages of stochastic encoders. In Neural Compression Workshop at ICLR , 2021.

Lucas Theis and Noureldin Yosri. Algorithms for the communication of samples. In International Conference on Machine Learning , 2022.

Lucas Theis, Tim Salimans, Matthew D. Hoffman, and Fabian Mentzer. Lossy compression with Gaussian diffusion. arXiv preprint arXiv:2206.08889 , 2022.

James Townsend, Thomas Bird, and David Barber. Practical lossless compression with latent variables using bits back coding. In International Conference on Learning Representations , 2018.

James Townsend, Thomas Bird, Julius Kunze, and David Barber. HiLLoC: lossless image compression with hierarchical latent variable models. In International Conference on Learning Representations , 2019.

Arash Vahdat and Jan Kautz. NVAE: A deep hierarchical variational autoencoder. In Advances in neural information processing systems , 2020.

Shifeng Zhang, Ning Kang, Tom Ryder, and Zhenguo Li. iFlow: Numerically invertible flows for efficient lossless compression via a uniform coder. In Advances in Neural Information Processing Systems , volume 34, pp. 5822-5833, 2021.

Jacob Ziv. On universal quantization. IEEE Transactions on Information Theory , 1985.

Peer reviews:
comments: SUMMARY:
The paper presents a coding scheme for greedy rejection sampling by Harsha, which builds upon the existing A* coding framework. The trick here is to add a partition operation between each reject/accept steps, which allows a more efficient scheme under certain assumptions. Overall, the work contributes meaningfully to the field of rejection sampling.  Minor adjustments to the figure margins are needed for improved readability.

STRENGTHS:
The proposed method and analysis demonstrate the potential for achieving faster runtime under certain assumptions about the proposal and target distributions.  The experiment shows the benefits of the method.
The paper is relatively easy to follow.

WEAKNESSES:
While the paper is generally well-written, the lack of sketch proof hinders the ease of understanding why the proposed coding scheme works or how the assumption matters. Including a sketch proof would greatly enhance the understanding of the presented approach. 

The differences between the present algorithm and prior work on A* coding must be explained more clearly.

It would be beneficial if the author could comment on the complexity of integrating equation (2) with distributions beyond the Gaussian, thus expanding the scope of the analysis.

QUESTIONS:
Can the authors explain the implication of assumption 2&3? Comparing the A* coding, I do not follow how GRCS and GRCD can outperform A* coding in terms of runtime complexity. 

 Line 274: I do not see where in the appendix factor 2 (in the log term) does not appear in the coding length. Also, 11 bits in O(1) seem to be longer than the bound by Li and El Gamal but A* coding and GRC seem to be on par in Figure 4 (right). Can the author provide some additional justifications for this?

LIMITATIONS:
n/a
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-02

comments: SUMMARY:
The authors propose a greedy rejection coding based algorithm for performing compression. The authors claim that GRC an ERC based compression method is the first method which does not make unrealistic assumptions, is guaranteed to finish in fixed amount of time. They generalize the algorithm to arbitrary partition spaces and provides guarantees to finish in limited runtime. 

Disclaimer - The reviewer is not the best person to evaluate this paper. They have no experience in this area.

STRENGTHS:
* The paper is reasonably well write and is able to explain the challenges very well.

* The proofs by and large look correct, however the reviewer can not comment on their novelty

*The authors have some experiments.

WEAKNESSES:
It would have been good to see how their method works at even higher dimension., can they compress say other models using their method.

QUESTIONS:
Look at weakness sections

LIMITATIONS:
The authors have highlighted some limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a new algorithm, Greedy Rejection Coding (GRC), for Relative Entropy Coding (REC). REC is an alternative to entropy coding and quantization, which naturally interfaces with continuous variables and has applications in data compression. Current REC algorithms have faced challenges like slow runtimes or restrictive assumptions, hindering their widespread use. The new GRC method, along with its two variants GRCS and GRCD, generalizes a previous rejection sampling-based algorithm to arbitrary probability spaces and partitioning schemes. The authors demonstrate that GRC terminates almost surely, producing unbiased samples, and that GRCS has an optimal expected runtime and codelength for a wide range of one-dimensional problems. The effectiveness of GRC is showcased in a compression pipeline with variational autoencoders on the MNIST dataset, where it outperforms in compression efficiency.\n\nSTRENGTHS:\n1. The paper introduces a new approach, Greedy Rejection Coding (GRC), to the field of rate-distortion compression (REC). The authors build on the rejection algorithm from Harsha et al., extending it to arbitrary probability spaces and partitioning processes. Theoretical results show that the GRC method is accurate and terminates successfully under certain general assumptions. \n\n2. The authors validate their ideas with a proof-of-concept implementation using Variational Autoencoders (VAEs) on the MNIST dataset. They compare them against previous methods like A* coding. Their experiments validate the theoretical underpinnings of the algorithms and show improvements over prior work.\n\n3. The paper deals with the challenge of coding overhead in compression problems. The authors propose a principled modification to the Evidence Lower Bound (ELBO) and entropy coding of GRCD's indices using a ζ distribution as innovative ways to improve compression efficiency.\n\nWEAKNESSES:\n1. **Dependency on CDF:** A major limitation of the proposed Greedy Rejection Coding (GRC) approach is that it requires the ability to evaluate the cumulative distribution function (CDF) of the distribution Q. This could potentially limit the algorithm's practicality in situations where the CDF is difficult or computationally expensive to calculate.\n\n2. **Applicability to Multivariate Distributions:** The paper mentions that the current approaches assume target-proposal pairs over the real line. This might limit the effectiveness of the proposed algorithm when dealing with multivariate distributions, as they are decomposed into univariate conditionals, resulting in additional coding overhead per dimension.\n\n3. **Lack of Generalizability to Other Datasets:** The empirical validation of the approach is based on synthetic data and one real-world dataset (MNIST). How well the proposed approach would generalize to more complex datasets or tasks is unclear.\n\n4. **Overhead Mitigation Measures:** The proposed methods to reduce the overhead, such as a modified ELBO and entropy coding using a ζ distribution, could add complexity to the algorithm. Moreover, it's not yet fully clear how much improvement they can bring in diverse scenarios.\n\n5. **Efficiency Analysis:** While the paper discusses the runtime of the proposed algorithm compared to existing methods, it does not provide a detailed computational efficiency analysis. Therefore, it might be unclear how the proposed algorithm would perform under different computational resources.\n\nQUESTIONS:\nPlease see the Weaknesses.\n\nLIMITATIONS:\nThis work does not pose any ethical concerns.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 1, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-17"}

comments: SUMMARY:
This paper proposes new Greedy Rejection Coding methods that generalize to continuous probability space and arbitrary partitioning schemes. The new methods produce correct samples with the expected code length of $O(D_{KL}[Q||P])$ where $Q$ is the target distribution and $P$ is the proposal distribution. The paper then evaluates the new algorithms and verifies the improvements in runtime in practical experiments.

STRENGTHS:
- The proposed methods achieve the optimal runtime up to a constant order while guaranteeing the optimal expected codelength.

- The experiment results also verify the theory. The runtime of proposed methods (in terms of the number of steps) stays constant when $D_\infty[Q||P]$ varies and is always smaller than the runtimes of previous methods AS* and AD*.

- Comparisons to previous results, technical novelties, and limitations are thoroughly discussed.

- The paper also provides plenty of high-level discussions and visualizations of their results.

WEAKNESSES:
- I think some intuitions about why the modifications help improve the runtime and also why the authors decide to make that particular modification can be added to the main text to improve the readability of the paper.

QUESTIONS:
- Can the author explain why the bound mentioned in line 274 is better than the bound in eq. 10?

- What is the wall-clock runtime of the proposed methods? It's clear that the proposed methods significantly improve the runtime in the number of steps but I'm curious if some of the extra steps in the algorithms could cause overhead in terms of wall-clock runtime/step.

- There's a typo in line 321: "GRC can become significant when the KL over Hence..."

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-26