Faster Relative Entropy Coding with Greedy Rejection Coding
Authors: Gergely Flamich
Stratis Markou
José Miguel Hernández-Lobato
Venue: NeurIPS 2023 poster (2023)
Keywords: Compression
Learnt Compression
Relative Entropy Coding
Information Theory
TL;DR: We introduce Greedy Rejection Coding, a Relative Entropy Coding algorithm whose two variants, GRCS and GRCD, achieve optimally fast runtimes and codelengths, improving upon A* coding, the current state-of-the-art method.

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
Gergely Flamich' Department of Engineering University of Cambridge gf3320cam. ac ac.uk
Stratis Markou" Department of Engineering University of Cambridge em6260cam. ac. uk
José Miguel Hernández Lobato Department of Engineering University of Cambridge jmh233@cam ac.uk
Abstract
Relative entropy coding (REC) algorithms encode a sample from a target distribu- tion Q using proposal distribution P using few bits 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-algoriee of Harsha et al. (2007) 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 IR with unimodal density ratio dQ dQ/PP, the expected runtime of where 3 4.82, and its expected
GRCS is upper bounded by 3DkL P] codelength is optimal. This makes GRCS the first REC algorithm with guaranteed optimal runtime for this class of distributions, up to the multiplicative constant B. This significantly improves upon the previous trate-of-the-aa method, A coding (Flamich et al. 2022). Under the same assumptions, we experimentally observe and conjecture that the expected runtime and codelength of GRCD upper bounded by DKL[Q] P]
[Q] + (1). Finally, we evaluate GRC in a variational autoencoder- based compression pipeline on MNIST, and show that modified ELBO and index-compression method can further improve compression efficiency.
Introduction and motivation
Over the past decade, the development of excellent deep generative models (DGMs) such as varia- tional 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 range of domains, such as lossless and lossy image or video compression (Zhng et al. 2021 Mentzer et al. 2020 2022).
Transform coding. Most learnt compression algorithms are transform coding methods: they first map datum to a latent variable using a learnt transform, and encode it using entropy coding (Balle) al. 2020). Entropy coding assumes discrete variables while the latent variables in DGMs are typically continuous, most transform coding methods quantize the latent variable prior to entropy coding. Unfortunately, quantization has zero derivative almost everywhere. Thus, state-of-the-ee DGMs trained with gradient-based optimisation must resort to some continuous approximation to
"Equal contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

GRCS
Rejection sampling
GRCD
Perturb-and maximiee
GRCG
AS coding
AD coding
On-sample partitioning
Dyadic partitioning
Global A*
Global partitioning
Figure 1: An illustration of the relations between the variants of GRC, introduced in this work, and the variants of A' coding, Algorithms in purple are introduced in this work. The algorithms of Harsha et al. (207)) and Li & El Gamal (2018) are equivalent to GRCG and Global A coding respectively.
quantisation during training and switch to hard quantisation for compression. Previous works have argued that using quantisation within learnt compression restrictive or otherwise harmful, and that method which naturally interfaces with continuous latent variables is needed (Havasi 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. Z020). and formally, requires that we encode a random sample from 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 DKL Q Plmany bits on average and, notably, allow Q and P to be continuous (L & El Gamal 2018 Flamich et al., 2022 Flamich & Theis 2023 Flamich 2023) Thus, can forgo quantization in compression pipeline and learn our translorms end end using gradient descent and the reparameterization trick. Morcover, REC 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] Shahmin et al. 2023) and compression
has fundamental advantages
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 overhcads. 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 DKL Q P < oc, which ensures that it produces finite code. However, unfortunately, Agustsson & Theis 2020) showed that without additional assumptions on Q and P. the expected runtime of any algorithm that simulates Q-distributed sample using P must scale least exponen- tially in DKr. [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 the algorithm simulates a -iistriibueed sample with ||Q @llTv < 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 which is optimal since any REC algorithm has an expected runtime of at least D However, Howeer, it requires both Q and P to be uniform distributions, which limits its applicability. Recently, Flamich et al. (022) 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.x QI... Unfortunately, this runtime is suboptimal and is not always practically fast, since Dxx (P)PP) can be arbitrarily large for fixed DKL. P). Further, as discussed in Flamich et al. (022)) 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. p we develop GRCS and GRCD, twO variants of GRC that partition the sample space to dramatically speed up termination. Figure Tauueeeillsstaaa iillutrrate 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 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 IR, 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/dPP is unimodal, the expected runtime and codelength of GRCS is O(DKL (QI)))) This significantly improves upon the (Q1== 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/ /QPPPP is uni- modal, the expected runtime and codelength of GRCD are [ Pl. This also signifi- cantly improves over the Do [Q] P] empirically observed runtime of AD" coding.
We implement a compression pipeline with VAEs, using GRC to compress MNIST images. We propose 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 T is stricter than the one given by |Flaiic et al.) (202), 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 RR be a set of pairs of distributions (Q.P) over (X, such that DKL [P] < co and P be the set of all distributions P such that (Q, P) R for some distribution Q. Let S = (Si, (2...) be a publicly available sequence of independent and fair coin tosses, with corresponding probability space (S F,F) and let C = {0, 1} be the set of all finite binary sequences. A REC algorithm is pair of functions enc : R x S + C and dec : C x P Xx S - X. such that for each (Q.P) E R and S ~ PP, the outputs of the encoder C enc(Q, P,S) and the decoder X dec(P. (P, C, S) satisfy
X X~Q and Es[[C] = DKL[Q[IP] + Ollog ( + 1)),
(1)
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. (007)'s ''s REJ-SAMPLER, by drawing inspiration from A coding (Flamich et al. 2022). Specifically, A coding generalises Li & El Gamal (018)) 's Poisson functional representation by introducing a partitioning scheme to speed up the algorithm's 's termination. In analogous fashion, GRC generalises Harsha al. (007) by also introducing partitioning schemes to speed up termination and achieve optimal runtimes. Here we discuss relevant algorithms, building up GRC in section E
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 to arbitrary partitioning schemes (see definition 5 in section 3. The generalisation arbitrary probability spaces relies on
3

---

Q
Q
Q
P
Figure 2:
Example run of Harsha et al. (2007). for a pair of continuous Q and P over [o. 1]. The green and red regions correspond acceptance and rejection regions each step. Here the algorithm rejects the first two samples and accepts the third one, terminating at the third step.
the Radon-Nits Nikodym derivative dQ/dP, which is guaranteed exist since < P by definition When Q and P both have densities, dQ/dP coincides with the density ratio.
At each step, the algorithm draws sample from P and performs an accept-reject step, illustrated in fig. 2 If it rejects the sample, 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 To to be the
on X, and recursively for N, set:
(S) Ta(S) + As+1(S),
AS+ (2)1 (p)))
(2)
ta(x)
def dTa
dP (z),
(x) def min dP
dQ
ta(a),, (1
(3)
Ba+1(x) def 0.d+1
Xa ~ P, Ua Uniform(0,
(4)
1- Ta(a)
for all X.SEE The algorithm terminates at the first occurrence of Us S Ba+1 (Xa). The Ta measure corresponds to the mass that has been ruled off up to and including the d rejection: T (X), (x and (2(')) are the sums of the blue and green masses in the left, middle and right plots of fig. 2 respectively. The As measure corresponds to the acceptance mass at the d step: (x)), (2 and A2(X') are the masses the green regions in the left, middle and right plots of fig. respectively Lastly, la, oa are the Radon-Nikodym derivatives i.e., roughly speaking, the densities, of Ta. Ad with respect to P, and Ba+1(Xa) is the probability of accepting the sample Xa-
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, Fllaiich Theis (2023) showed expected runtime is 2Dx [QIIP), where Doo [Q]] supaex log(d//dP)(r is the Rényi co-divergence. Unfortunately. this runtime is prohibitively slow in most practical cases.
REC with Poisson & Gumbel processes. Li & El Gamal (018) introduced REC algorithm based on Poisson processes, referred as Poisson Functional Representation (PFR). PFR assumes that dQ/dP is bounded above, and relies on the fact that (Kingman 1992), if T2 are the ordered arrival times of homogeneous Poisson process on R- and Xn P. then
{ {T. dP JQ
N arg min nEN
(Xm)
XN XN~Q.
(5)
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.,, X., until it solves the - 1. The decoder can recover y from (P,C. S)). by drawing samples from 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 2D [1 Maddison 2016).
search problem in eq. (). and returns = XN, C
Fast REC requires additional assumptions. These algorithms' slow runtimes are perhaps unsur- prising considering Agustsson & Theis s result, which shows under the computational hardness assumption RP * NP that without making additional assumptions Q and P. there is no REC

---

Algorithm 1 Harsha et al. rejection algorithm; equivalent to GRC with global partition
Algorithm 2 GRC with partition process Z; differ- ences to Harsha et al. algorithm shown in green
Require: Target Q. proposal I', space X
Require: Tarett Q. proposal I', space X, partition Z
I: d 0, <-
1: d <- 0. 0
2:
2: To 1, f
3: while True do
3: while True do
4: Xd+1 P
4:
X1, Pls./P(Sa)
5: Us+1 Uniform(0,) 1)
5:
Ula Uniform(O,
6: Ba+1 AcceptProb(Q, Xa+1,
6:
3tg AcceptProb Xtg Ta)
if U&+1 S Ba+1 then
if if I 8841 or d Dax then
8:
return X441,
8:
return X1a: Ia end if
9:
end if
9:
10:
10:
PartitionProb(Q Zza, Zza41) ba Bernuuli(p)
12:
12:
Ia+1 <- 2Ta ba and Sa+1 f- 7lgs1
13:
Ta+1 +- RuledDutMass(Q. P,Ta)
13:
Ta+1 + RuledOutMasa(( (Q, P, Ta. S4+1)
14:
d + +++
14:
d+d+
15:
end while
15:
end while
algorithm whose expected runtime scales polynomially in PI Therefore, in order achieve faster runtimes, 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 satisfied by many deep latent variable models in practice and (2) 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 partitioning scheme, which splits the sample space X' in nested partitioning subsets, to speed up the solution of eq. ( Drawing inspiration from this, our proposed algorithm generalises eqs. fig. D. introducing partitioning processes (definition to speed up the algorithm's 's termination.
in analogous manner (see
Definition (Partitioning process).
partitioning process is process Z N+ such that
Z1 = X, Zzn n 22a+1 = 0, 7.2 U Zznt1 Zn-
In other words, partitioning process Z is an infinite binary tree-structured process, where the root node is Z1 X' and has index 1, and each node Zn with index n is partitioned by its two children nodes Zzn. Z2n+1 with indices 2n and 2n + 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 (023)) introduces a rejection sampler based on Poisson processes, called Greedy Poisson Rejection Sampling (GPRS), which can be used as REC algorithm. Similar to GRC and A coding, GPRS partitions the sample space speed up the convergence to the accepted sample. Furthermore, 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 certain ODE, while our method does not, making potentially more favourable practice. We believe establishing closer connection between GPRS and GRC is a promising future research direction.
3 Greedy Rejection Coding
Generalising Harsha al. (007)). In this section we introduce Greedy Rejection Coding (GRC; d5 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 applicable to arbitrary REC problems, including REC with continuous distributions. Second, similar to A coding, GRC can be combined with arbitrary partitioning processes, allowing achieve optimal runtimes given additional assumptions the REC problem, and an appropriate choice of partitioning process.
5

---

Q
P
=I
Z21
Z21 +1
S2=Zi,
Sample accept reject
(b) Partition & sample by [0
(c) Sample & accept reject
@
P
O
Si Z1
Z2n
Z
S2 = 24,
(d) Sample & accept reject
(c) Partition & sample b2 E {0,
(f) Sample & accept or reject
Figure 3: Illustrations the two variants of GRC considered in this work. (a) (c) show GRC with the on-sample partitioning process (GRCS). (d) (f) show GRC with the dyadic partition process (GRCD). GRC interleaves accept-reject steps with partitioning steps. the former, it draws a sample and either accepts rejects it. In the latter, it partitions the sample space and randomly chooses one of the partitions, ruling out large parts the sample space and speeding up termination.
3.1 Algorithm definition
Overview. Before specifying GRC, summarise its operation with accompanying illustration. On a high level, GRC interleaves accept-reject steps with partitioning steps, where the latter are determined by partitioning process. Specifically, consider the example in figs. 3d where Q and P are distributions over X' = =0, 1]. and Z is the partitioning process defined by
Zn = L = Zzn [L, M), Zzn+1 = [M, R), where M = (L + R)/2.
(7)
In each step d = 1,2,.... GRC maintains a heap index Ia of an infinite binary tree, and an active subset S, = Zr C of the sample space, initialised as Io and S Zi = respectively.
Accept-reject step. In each step, GRC draws a sample from the restriction of P to Sa, namely (()a) and either accepts or rejects it. If the sample is accepted, the algorithm terminates. Otherwise, GRC performs a partitioning step shown in fig. 3d
Partitioning step. In each partitioning step, GRC partitions S. VI, into Zar and Z2ra+1. as specified by the partitioning process Z. It then samples a Bernoulli random variable by, whose outcomes have probabilities proportional to the mass of which has not been accounted for, up to and including step d, within the partitions Z2,, and Z2l.+1 respectively. In fig. Be these two masses correspond to the purple and orange areas, and the algorithm has sampled by Last, GRC updates the heap index to la+1 2 + ba and the active subset to Sa+1 Z GRC proceeds by interleaving accept-reject and partitioning steps until acceptance occurs.
Algorithm specification. The aforementioned algorithm can be formalised terms of probability measures over arbitrary spaces and arbitrary partitioning processes. Above, algorithms D and describe Harsha et al. rejection sampler and our generalisation of it, respectively. For the sake of keeping the exposition lightweight, we defer the formal necsure-theorede definition of GRC to the appendix (see definition 5 in appendix I. and refer to algorithm as a working definition here.
Comparison to Harsha et al. While algorithms/ I and are similar, they differ in two notable ways. First, rather than drawing sample from P. GRC draws sample from the restriction of tt an active subset S2 Za C X, namely Pls./ P(Sa). Second, GRC updates its active subset S Za at each step, setting to one of the children of Za, namely either Zzd or Z2441, by drawing by Bernoulli, and setting Z24+by. This partitioning mechanism, which does not appear in algorithm yields

---

different variant of GRC for each choice of partitioning process Z. In fact, as shown in Proposition
below, algorithm I is a special case of GRC with S X for all d. See appendix [A.2 for the proof.
Proposition (Harsha et al. (2007) is a special case of GRC), Let Z be the global partitioning process over Σ, defined as
Z = X, Zzn = Zn. 1 = 9, for all = 1,2,.....
(8)
Harsha et al.) 2007 is equivalent to GRC using this Z and setting C D instead of C = TD. We refer to this al algorithm as Global GRC, or GRCG for short.
Partitioning processes and additional assumptions. While Proposition I shows that Harsha et al. algorithm is equivalent to GRC with a particular choice of Z. a range of other choices of Z is possible, and this is where can leverage additional structure. In particular, we show that when Q and P are continuous distributions over R with a unimodal density ratio dQ (QPPP. we can dramatically speed will consider the on-sample and dyadic
up GRC with an appropriate choice of Z. In particular, partitioning processes from Flamich et al. (022). given in Definitions S and
Definition& 3 (On-sample partitioning process). Let X = R (-00, and P a continuous distribution. The on mn-sample partitioning process is defined as
Zn (a, b). a. b E X = Zzn = Xn), Zzn+1 [X,, b], where Xn. P Plz.//(Zy,
In other words, in the on-sample partitioning process, Zn are intervals of R, each of which is partitioned into sub-intervals Z2n and Z by splitting at the sample Xn drawn from Pl Plz.,/ /7.,,P(24 (Z.). We refer to GRC with on-sample partitioning as GRCS.
Definition 4 (Dyadic partitioning process), Let X' = R (- oo and P a continuous distribution.
The dyadic partitioning process is defined as
Zn [a, b), a, b E = Zzn = [a. c), Zzn+11 [c, b), such that P(Zza) P((zen+1):
Similar to on-sample partitioning, in the dyadic process Zn are intervals of R. However, in the dyadic process, Zn is partitioned into sub-intervals Zzn and Z2n+1 such that P(Zza) = P(Zza+1). 1). We refer GRC with the dyadic partitioning process as GRCD.
GRC with a tunable codelength, Flamich et al. (2022) presented a depth- variant of AD" coding, DAD' coding, in which the codelength CC can be provided as tunable input the algorithm. Fixed-codelength REC algorithms are typically approximate because they introduce bias in their samples, but nevertheless useful in certain contexts, such for coding 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 Dmax (see algorithm 2 Setting D algorithm 2 corresponds to exact GRC, while setting Dmax < co corresponds to depth-limited GRC. In appendix D we provide detailed analysis of depth-limited GRC.
in
3.2 Theoretical results
Correctness of GRC. Below, we present three sets of assumptions on Q, P and Z, and in theorem we show that fulfilling any of them is sufficient to ensure the correctness of GRC.
D
Assumption 1. GRC has a finite ratio mode, i.e. dx/dP(x) < M for all x E for some M 0.
Assumption is the most generally applicable, as it does not restrict the sample space. Assumption | holds for GRCG, GRCS and GRCD, so long as dQ/dl dP is bounded. While this assumption is very general, some cases we may want to consider Q, P with unbounded dQ/dP IdP. To this end, we show that it can be replaced by alternative assumptions, such assumptions 2 and S
Assumption 2. GRC is dingle-brancee i.e. for each d, ba = 0 or ba = almost surely.
GRC with the global partitioning process (eq. g satisfies assumption 2 In addition, Q and are distributions over IR and dQ/dPP is unimodal, GRCS also satisfies assumption P
Assumption 3. X' N and GRC has nicely-shrinknng bounds, i.e. IP-almost surely the following holds: for each x E X which is in a nested sequence of partitions x E 21 2 ... Zr. 2 with P(Zua) - kis + xc, there exist Y,, T1, T.. R2 such that
0, Zhe C Br and P(Zk, 2 (x)(
(9)
where Bx)x) denotes the open ball of radius r centered on T. We recall that P. is the measure associated with the sequence of publicly available coin flips S.
7

---

Intuitively, Z is nicely shrinking if its branches shrink in roughly uniform way over A. As the most important example. when Q and are distributions over R. GRCD satisfies assumption B Now. Theorem 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
Theorem 1 (Correctness of GRC), Suppose Q. P and Z satisfy any one of assumptions 1o Then, algorithm terminates with probability L, and its returned sample X has law X ~~ Q.
Expected runtime and codelength of GRCS. Now we turn the expected runtime and codelength of GRCS. Theorem 2 shows that the expected codelength of GRCS is optimal, while Theorem establishes that its runtime order-optimal. We present the proofs of the theorems in appendix
Theorem (GRCS codelength). Let Q and be distributions over with DRL [ < co and dQ/dP unimodal. Let Z. be the on-sample partitioning process, and X its returned sample. Then,
H[ S DKI(Q1) P] + log (DRr. [Q]] 0(1).
(10)
Theorem 3 (GRCS runtime). Let Q and P be distributions over IR with DEL Q| < co 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 3 = 2/ log(4/3) z 4.82 we have
[[D] S D + O(1).
(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, have pair of correlated random random variables X, Y ~ Px.y- During one round of communication, the encoder receives Y Py and needs encode a sample X Pxiy from the posterior using Px as the proposal distribution. Thus, RCC can be thought of as the average-cass version of REC, where the encoder sets Q + PxIy and P + Px. In this case, when the conditions of theorem 2 hold for every
(Pxix: Px) pair, in appendix C we show that the coefficient of the log-factor in eq. (10) can be improved that the average-case bound becomes I[X; Y + log([[X; Y] + 1) + O(I), where is the mutual information between X and Y.
[[ = Ey [
substantially improves upon the runtime of A" coding, which is the current fastest REC algorithm with similar assumptions, In particular, AS* coding has O [ Pl) expected runtime, which can be arbitrarily larger than that of GRCS. Remarkably, the runtime of GRCS is optimal up to the multiplicative factor B
GRCS runtime is order-optimal. Theorem 3
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 experiments under fair and reproducible conditions. Our code is available at htttps://githut com/ http:/bbrddeemml/ffassssseeeeeee atwaass--eeccee wineewwithgrrr
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 shows the results of our synthetic experiments.
Partitioning processes improve the runtime of GRC. First, we observe that, assuming that dQ/ /dP unimodal, introducing the on-sample the dyadic partitioning process speeds up GRC dramatically. In particular, fig. shows that increasing the infinity divergence Dxc | P] (for a fixed DKL PP) does not affect the runtimes of GRCS and GRCD, which remain constant and small. This is remarkable speed-up over the exponential expected runtime 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 shows that increasing the infinity divergence Doc [[] while keeping the
8

---

Runtimes with Gaussian Q and P
Codelengths with Gaussian Q and P
150
+0-
Codelength a
# steps
100
10
50
5-
0
10
20
30
40
50
4
8 10
D
Figure 4: Comparison between GRC and A* coding on synthetic REC problems with Gaussian Q and P. Left: we fix DkI P] 3 and vary Doo [Q|| measuring the number of steps taken by each algorithm. Right: we fix Dxe [Q [] codelengths produced by each algorithm. Reported codelengths do not include additional logarithmic overhead terms. Results are averaged over 4 x 10 different random seeds for each datapoint. We have included error-bars in both plots but these are too small see compared to the plot scales.
D [|] + 2 and vary DKL. DKI[[Q||)) [Q]]P] plotting the
KL divergence DRL [Q] P] fixed, increases the runtime of both AS* and AD" coding, while the runtimes of GRCS and GRCD remain constant. More generally, for fixed KL divergence, the infinity divergence can be arbitrarily large 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 in learnt compression. Here, we implement a proof-of-concee lossless neural compression pipeline using a VAE with a factorized Gaussian posterior on MNIST and take the architecture used by Townsend et al. (018).. To compress an image Y, we encode 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 DRa.(g(Xa (Xal x)Ilp(Xa)) bits coding cost for latent dimension d, this incurs an overhead of log( PxL (x ())lp(X))) [x)]+1] 1 + (1)) bits, analogously to how a symbol code, like Huffman coding, incurs constant overhead per symbol (MacKay) 2003 However, since log(1 1 + z) se x when x Rs O, 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. 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 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):
Ex=e(X ( X)]
Σ log (Dxul((X.) +X)) (12)
Rogular ELBO
Logarithmic overbhed per dimension
Coding the latent indices. As the final step during the encoding process, we need a prefix code to encode the heap indices La returned by GRCD for each d. Without any further information, the best we can do is use Flias 8 coding (Elias) 1975), which, assuming our conjecture on the expected runtime of GRCD holds, yields an expected codelength of x 2og x XX] +) OL). However, we can improve this if we can estimate Bloog La] for each d: it can be shown, that the maximum entropy distribution a positive integer-valued random variable with under a constraint on the expectation on its logarithm is (())A) x n with - Ellog La] + 1. In this case, entropy coding La using this distribution yields improves the expected codelength to I[Y; X + log(I[Y;; +0))) ]
Experimental results. We trained our VAE with L E {20, 50, 100} latent dimensions optimized using the negative ELBO and its modified version in Equation (2). and experimented with encoding the heap indices of GRCD with both 8 and ( coding. We report the results of our in Table I on the MNIST test set bits per pixel. In addition to the total coding cost, we report the negative BLOO per pixel, which is the fundamental lower bound the compression efficiency of REC with each

---

TRAINING OBJECTIVE
TOTAL BPP WITH CODING
TOTAL BPP
NEG. ELBO
OVERHEAD BPP
# LATENT
WITH CODING
PER PIXEL
WITH CODING
20
1.472 0.004
1.482 0.004
1.391 0.004
0.091 = 0.000
ELBO
50
1.511 0.003
1.530 0.003
1.357 = 0.003
0.172 = 0.000
100
1.523 0.003
1.600 1 0.003
1.362 = 0.003
0.238 - 0.000
20
1.470 0.004
1.478 0.004
1.393 + 0.004
0.085 = 0.000
MODIFIED ELBO
50
1,484 + 0.008
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 small VAE with
different latent space sizes, optimized using either the EL 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 á 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 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 Harsaa 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 -oon-of-concept learnt compression experiment on MNIST using VAEs. We demonstrated that principled modification to the ELBO and entropy coding GRCD's indices using 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* and AD" 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 non-negligible cost. Thus, an important direction is to investigate whether fast REC algorithms for multivariate distributions can be devised, to circumvent this challenge.
Author Contributions
SM suggested that single-branch greedy rejection coding (Flamich & Theis 2023 can be extended arbitrary partitioning processes. GF discovered GRCS and SM discovered GRCD. SM provided a proof for Theorem I and GF provided proofs for Theorems 2 and B SM carried out the synthetic experiments and GF carried out the VAE experiments. SM drafted the majority 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.
7
Acknowledgements
GF acknowledges funding from DeepMind. SM acknowledges funding from the Qualcomm In- novation Fellowship and the Vice Chancellor's & George and Marie Vergottis scholarship of the Cambridge Trust.
10

---

References
Firikur 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. IEEE International Symposium Information Theory, 2023.
Gergely Flamich, Marton Havasi, and José Miguel Hemáández-Lobato. Compressing images by en- coding their latent representations with relative entropy coding. 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. Compres- sion with Bayesian implicit neural representations. In Advances in Neural Information Processing Systems, 2023.
Prahladh Harsha, Rahul Jain, David McAllester, and Jaikumar Radhakrishnan. The communica- tion complexity of correlation. In Twenty-Second Anmual IEEE Conference on Computational Complexity (C)C'07). IEEE, 2007.
Marton Havasi, Robert Peharz, and José Miguel Heráá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 a3Xiv: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 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 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 genera- tive 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: video compression transformer. In Advances in Neural Information Processing Systems, 2022.
Fernando Pérez-Cruz. Kullback-Leible 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--ffiieen Laplace mecha- nism for differential privacy via random quantization. arXiv preprint arXiv:2309.0699: 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 Learning Representations, 2018.
James Townsend, Thomas Bird, Julius Kunze, and David Barber. HiLLoC: lossless image com-
pression 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 Information Theory, 1985.
12


Reviews:
[{"review_id": "D8XadA9x8c", "reviewer": "Reviewer_t3AN", "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.", "questions": "- Can the author explain why the bound mentioned in line 274 is better than the bound in eq. 10?\n\n- 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.\n\n- There's a typo in line 321: \"GRC can become significant when the KL over Hence...\"\n\n", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The proposed methods achieve the optimal runtime up to a constant order while guaranteeing the optimal expected codelength.\n\n- 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*.\n\n- Comparisons to previous results, technical novelties, and limitations are thoroughly discussed.\n\n- The paper also provides plenty of high-level discussions and visualizations of their results.\n\n", "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer T3AN for their time reviewing our paper and valuable feedback. We were happy to hear they found our theoretical and experimental results as strong points of this work. We were also pleased that they thought our exposition, discussion of related work and limitations, and vizualisations were of a good standard. T3AN also raised a point on the intuition behind our suggested modifications and a few questions we address below.\n\n__Intuition on the effectiveness of our modifications:__ T3AN commented that we could improve the paper's readability by adding some intuitions about why our proposed modifications help speed up the runtime.\n\nWe agree with this point and will amend our paper based on the argument we outline below:\n\nFirst, let us explain the motivation for making the particular modifications that we chose in this paper. We already know from a result by Augustsson & Theis (discussed in our paper) that we cannot avoid exponential worst-case runtimes in REC without further assumptions in $P$ and $Q$. Therefore, we need additional assumptions to achieve a practically fast algorithm. Flamich et al. 2022 showed that assuming $Q$ and $P$ over the reals and a unimodal density ratio $dQ / dP$, and using a sample splitting scheme results in a dramatic speedup and a computationally tractable algorithm compared to a global partitioning scheme (i.e. compared to no partitioning). The success of this approach motivated us to examine if we can extend similar partitioning schemes to rejection-based coders and quantify the speedup we can obtain under the same assumptions.\n\nWe now explain why these modifications improve the runtime of the algorithm. In regular GRCG (figure 2), the probability of termination (= (green area) / (green area + white area)) becomes smaller as the algorithm progresses. Intuitively, this happens because most of the proposal $P$ proposes samples in regions with no remaining mass under $Q$. These samples will never get accepted, causing long runtimes. If we can make the algorithm somehow adaptive so that $P$ adapts to the shape of the white area, we can reduce wasteful samples and speed up termination. We can achieve this using the sample-splitting process.\n\nNote that if $dQ / dP$ is unimodal, then when GRC rejects a sample, all the mass that has not been ruled out (white area) lies either to the left or to the right of the rejected sample (see Figure 2; left subfigure: all the mass lies to the right of sample #1; middle subfigure: all the mass lies to the left of sample #2). Therefore, we can split the sample space into two parts and draw our next sample from the restriction of $P$ to the part containing all the remaining white mass. This increases the termination probability at each step, remarkably improving the algorithm’s runtime to be order-optimal. This notion extends to other adaptive partitioning schemes, such as the dyadic partitioning process. In this case, we have to choose the partition in a randomized way (see Algorithm 2) to ensure the samples are unbiased.\n\n__We will add this further intuition in our camera-ready version.__\n\n__Improved bound:__ The reason why the bound in line 274 improves on eq 10 is as follows. First, in the average case, we assume a publicly available joint distribution $P_{X, Y}$ over the correlated variables $X$ and $Y$. Given a private sample $Y \\\\sim P_Y$, the encoder aims to communicate a sample from the conditional $X \\\\sim P_{X | Y}$. Therefore in this case, if for a given $Y$, we set $P = P_X$ and $Q = P_{X | Y}$ and use a GRC, we can communicate a sample $X \\\\sim P_{X | Y}$ using at most\n\n$$D_{KL}[Q || P] + 2 \\\\log(D_{KL}[Q || P] + 1) + \\\\mathcal{O}(1) = D_{KL}[P_{X | Y} || P_X] + 2 \\\\log(D_{KL}[P_{X | Y} || P_X] + 1) + \\\\mathcal{O}(1)$$\n\nbits on average. Further, averaging out over $Y \\\\sim P_Y$, we obtain an average-case communication cost of at most\n\n$$\\\\mathbb{I}[X; Y] + 2 \\\\log(\\\\mathbb{I}[X; Y] + 1) + \\\\mathcal{O}(1)$$\n\nbits. Lastly, we can use $\\\\zeta$-coding (which is not possible in the single-shot REC setting) to further improve the coefficient in front of the logarithm, to improve the communication cost to at most\n\n$$\\\\mathbb{I}[X; Y] + \\\\log(\\\\mathbb{I}[X; Y] + 1) + \\\\mathcal{O}(1).$$\n\nbits. We hope this clarifies the question raised by T3AN and welcome further clarifying questions. We will add a discussion of this point in our camera-ready version.\n\n__Runtimes:__ Although the modified algorithm involves a few additional steps, these are far outweighed by the reduction in the overall number of steps they yield. Specifically using our JAX implementation, choosing $Q = \\\\mathcal{N}(0.5, 0.8^2)$ and $P = \\\\mathcal{N}(0, 1)$ and drawing $10^4$ samples with each method, yields wallclock times of $3.24$ sec, $0.90$ sec and $0.49$ for GRCG, GRCS and GRCD respectively (where for GRCG we have used algorithm 1 instead of algorithm 2 to avoid unnecessary computations). The difference in runtimes can be made arbitrarily larger by picking problems with a fixed KL but increasing infinity divergence. Note that although the modified algorithms require slightly more computation per step, they are much faster because they take far fewer steps. The exact runtimes above are implementation-dependent though we believe our implementation is relatively well-optimized and fair between the methods.\n\n__Typo:__ Thank you for pointing this out. We have updated this in our working draft.\n\n__Conclusion:__ We thank T3AN for their time and valuable feedback. We found that their review was positive, and we hope our rebuttal has addressed the points they raised. If T3AN is satisfied with our response, we would like to invite them to consider increasing our score."}, {"review_id": "fUH7rGDE8y", "reviewer": "Reviewer_ULav", "summary": "This 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.", "questions": "Please see the Weaknesses.", "limitations": "This work does not pose any ethical concerns.", "rating": 5, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. **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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer ULAV for their time reviewing our paper. We are happy they found our theoretical and experimental results as strong points. Reviewer ULAV also raised some concerns, which we address below.\n\n__Dependency on CDF:__ ULAV commented that since GRC requires access to the CDF of the target $Q$ and proposal $P$, this may “limit the algorithm's practicality in situations where the CDF is difficult or computationally expensive to calculate.” \n\nIn many state-of-the-art communication pipelines, $P$ and $Q$ have simple forms, and their CDFs can be evaluated cheaply:\n\n1. In data compression, e.g., using VAEs, $P$ and $Q$ are typically chosen to be uniform, Gaussian, or Laplace.\n2. In model compression, both $P$ and $Q$ are often Gaussian.\n3. In federated learning, $P$ and $Q$ are also typically Gaussian.\n\nThe CDFs of most base distributions in learnt compression are easy to evaluate. Existing pipelines __use tractable distributions to evaluate their training objectives__, and we expect future ones will do so too. Therefore, assuming access to the CDF is not a constraining assumption.\n\n__Applicability to multivariate distributions:__ ULAV commented that GRC assumes $P$ and $Q$ over the reals, which might limit its effectiveness when decomposing joints into univariate conditionals due to the additional coding overheads. We do not entirely agree with this comment:\n\n1. GRC directly applies to multivariate $P$ and $Q$. However, we have yet to find appropriate assumptions and a partitioning process that simultaneously achieves optimal runtimes and codelengths in this setting. While this problem is challenging, we believe more suitable partitioning schemes may exist, so we do not view this as a fundamental limitation of GRC.\n2. The provably fast instances of A* coding, the current state-of-the-art, also assume univariate $P$ and $Q$. Therefore, this does not limit the applicability of GRC any more than the current state-of-the-art.\n3. We have provided practical approaches to mitigate overhead codelengths, showing these can be as low as 6% of the total coding cost. We regard these approaches as proofs-of-concept and believe that future work could further reduce these overheads.\n\nOverall, we believe this is neither a prohibitive nor a fundamental drawback of GRC, and we do not think our work should be viewed unfavourably on these grounds.\n\n__Lack of generalisability to other datasets:__ Reviewer ULAV commented that it is unclear whether our approach would generalize to more complex datasets or tasks.\n\nWe stress that our main contributions are theoretical rather than empirical and note that applying REC to complex learned compression pipelines is well outside the scope of this work. Moreover, we do not expect the performance of REC would differ on larger-scale problems: Flamich et al. (2020) studied applying REC for high-resolution image compression using VAEs and found that their method was competitive with the state-of-the-art at the time. \n\nNevertheless, our runtime and codelength bounds apply directly to a range of existing and future pipelines. These give quantitative guarantees about the performance of GRC as a function of the target and proposal, which are aspects of the model alone. This is one of the appeals of GRC (and generally REC): given a model’s performance on some data, one already has a good idea of the performance of the entire pipeline.\n\n__Overhead mitigation measures:__ One issue raised by ULAV is that our proposed methods for mitigating overhead costs could add complexity to the algorithm and that it is “not yet fully clear how much improvement these can bring in diverse scenarios.” We would like to respond with the following points:\n\n1. Our modified ELBO objective is a simple, efficient, and principled method: it modifies the objective to account for the overhead codes during training, reducing overheads and incurring negligible computational overheads in practice.\n2. $\\\\zeta$-coding is a widely known method already applied in REC (see Li and El Gamal, 2017). Using $\\\\zeta$-coding is a straightforward, computationally cheap procedure that does not significantly complicate the overall compression pipeline, especially considering the improvements it yields.\n3. As demonstrated in our experiments, as the number of dimensions in $Q$ and $P$ increases, the performance gains from using the methods above also increase. This suggests our methods would be highly beneficial for larger latent variable models across a range of scenarios.\n\nWe believe the aforementioned methods are simple, principled, computationally cheap, effective in practice, and arguably of interest to related work. Further, they bear minimal complexity, which is far outweighed by their improvements.\n\nEfficiency analysis: ULAV commented that our paper “does not provide a detailed computational efficiency analysis.” We think this comment somewhat misrepresents our work. There are two measures of computational efficiency: computational and memory complexity. Runtime is synonymous with computational complexity, and we have provided near-optimal bounds for the expected runtime of GRCS. Further, at each step, GRC updates all variables in-place (Algorithm 2), so its memory cost is O(1). We will clarify this in our camera-ready version.\n\nCould the reviewer elaborate on what additional computational analysis they would like to see? If there are no further points of analysis they can think of, we kindly ask them to reconsider this point.\n\n__Conclusion:__ We thank ULAV for their time and feedback. Overall we found that some comments did not accurately represent our work. We hope our response has clarified these points and sufficiently addressed their concerns. If ULAV thinks we have adequately addressed their concerns, we would like to invite them to consider raising their score."}, {"review_id": "Iv9vCVTNUp", "reviewer": "Reviewer_BHnG", "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. \n\nDisclaimer - The reviewer is not the best person to evaluate this paper. They have no experience in this area. \n\n", "questions": "Look at weakness sections", "limitations": "The authors have highlighted some limitations.", "rating": 5, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper is reasonably well write and is able to explain the challenges very well.\n\n* The proofs by and large look correct, however the reviewer can not comment on their novelty\n\n*The authors have some experiments.\n", "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.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank reviewer BHNG for their time to review our work and their feedback. We were happy to hear that the reviewer thought our paper was well-written, discussed relevant challenges well, and was satisfied by the correctness of our proofs. The only point of weakness they mentioned was that they would have liked to see how our approach would perform in even higher dimensions and different models. We would like to respond to this point below.\n\n__Performance in higher dimensions:__ On one hand, it would certainly be interesting and valuable to assess the performance of GRC with larger latent variable models within realistic state-of-the-art compression pipelines. On the other hand, in our view, the main contribution of our paper is a theoretical one and consists of (1) devising GRC, (2) proving that it terminates and yields unbiased samples, and (3) achieving optimal compression codelengths and near-optimal runtimes under realistic assumptions. Our experiments with VAE-based compression constitute a proof-of-concept experiment to validate the efficacy of GRC in a controlled setting rather than to integrate it within a realistic compression pipeline with high-dimensional latent variable models.\n\nDeveloping a realistic compression pipeline with large latent variable models using GRC is an exciting direction of empirical work. However, it is beyond the scope of our current investigation, which is chiefly theoretical. For this reason, we think our work should not be viewed negatively on this basis.\n\n__Conclusion:__ We thank reviewer BHNG for the time taken to review our paper. We hope we have adequately addressed their concern regarding compression with higher-dimensional distributions. We note that while reviewer BHNG found no substantial issues with our paper and was overall satisfied by the quality of our exposition and the correctness of our results, they have recommended our paper with a somewhat conservative acceptance score. In light of this, we would like to invite reviewer BHNG to consider raising their score if they are satisfied with and convinced by our response to their review, as well as the other reviews and subsequent discussions."}, {"review_id": "0iiIT7B5te", "reviewer": "Reviewer_tpKF", "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.  \n", "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. \n\n 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? \n\n", "limitations": "n/a", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\nThe paper is relatively easy to follow.\n", "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. \n\nThe differences between the present algorithm and prior work on A* coding must be explained more clearly.\n\nIt 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.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank TPKF for their time reviewing our paper and their feedback. We were pleased to hear TPFK thought our work was a meaningful contribution and our paper well-written and relatively easy to follow. TPKF raised a few points which we would like to address below.\n\n__Sketch proof:__ We agree the paper would benefit from sketch proofs. We would also like to bring Appendix B to the author’s attention, where we provide an outline of the correctness proof (Theorem 1). We will edit that section to explain how we use assumptions 1, 2, and 3. Here, we provide some intuition as to why GRCS improves over the runtime of GRCG and A* coding and will include this with a sketch proof in our camera-ready version (see also our response to reviewer T3AN).\n\nIn GRCG (see Figure 2), the termination probability at each step (given that termination has not yet occurred) becomes smaller, resulting in long runtimes. This is because most proposed samples occur in regions with no remaining mass (white area) under $Q$ and are always rejected. The idea behind GRCS is to make the proposal distribution adapt to the remaining mass. We can achieve this, e.g., with the sample-splitting process.\n\nIf $dQ / dP$ is unimodal, then when GRCS rejects a sample, all the mass that has not been ruled out (white area) lies either to the left or right of the rejected sample. Therefore, we can split the sample space into two parts and draw our next sample from the restriction of $P$ to the part containing the remaining white mass. This adapts the proposal’s shape, increasing the termination probability at each step and dramatically improving the algorithm’s runtime. In Appendix C, we formalize this argument and show that the adaptive proposal of GRCS shrinks, homing into the remaining mass exponentially quickly, resulting in far shorter runtimes than GRCG and A* coding (see also the “differences to A* coding” section below).\n\n__Differences to A* coding:__ We thank TPKF for pointing out that the differences between GRC and A* coding could be explained further. Figure 1 illustrates the relationships between the algorithms pictorially. Following TPKF’s advice, we will add a supplementary section explaining these differences in more detail in our camera-ready version and highlight some below.\n\nWhile both algorithms use partitioning schemes to speed up termination, A* coding is based on Gumbel processes and uses a branch-and-bound search and termination condition. On the other hand, GRC is based entirely on rejection sampling. In fact, it is partly due to the different termination conditions that GRC achieves faster runtimes: GRC terminates once it accepts a sample (i.e., once it produces a good candidate sample), while A* coding must keep going even after it has produced a good candidate until it has sufficiently reduced an upper bound used in its termination condition. This upper bound depends on $D_{\\\\infty}[Q || P]$, which can be arbitrarily large for a fixed $D_{KL}[Q || P]$. For this reason, A* coding can take much longer to terminate compared to GRC, which can be much faster.\n\n__Integrating equation 2:__ Equation 2 can be integrated efficiently so long as we have access to the CDF of $P$ and $Q$. This CDF is tractable for a wide range of distributions used in learnt compression, including Gaussians, Laplace, uniforms, and many others. While we highlight the access to a CDF as an important requirement in our paper, this is not a significant limitation since most state-of-the-art learnt compression pipelines use tractable distributions such as the above. In our submitted code, we have included Laplace and uniform GRC samplers and executable examples under `notebooks/sampling.ipynb`.\n\n__Q: assumptions 2 & 3:__ Note that assumptions 2 & 3 are only used in Theorem 1 to prove the correctness of the algorithms. They are not used in our proof for the runtime of GRCS (Theorem 3). GRCS achieves faster runtimes than A* coding due to the algorithms’ different constructions and termination conditions. We have outlined how this results in faster runtimes for GRC in the “sketch proof” and “differences to A* coding” sections above.\n\n__Q: codelength bounds:__ The reason the factor of 2 does not appear in eq. 274 is as follows. In the average case, we can use $\\\\zeta$-coding (which is not possible in the single-shot REC setting) to reduce the prefix code’s length further. Li and El Gamal showed that this could improve the prefix codelength, eliminating the factor of 2 from the bound. We appreciate that this point might not have been clear and will add a clarifying comment to our camera-ready version.\n\nRegarding the constant overhead, note that the quantity involving the 11 bits is an _upper bound_ rather than an equality for the exact codelength. This constant term is likely not incurred in practice: it results from our analysis and can likely be improved. However, the primary purpose of this bound is to show the codelength is proportional to $D_{KL}[Q || P]$. Please also note that the results shown in Figure 4 (right) “do not include additional logarithmic overhead terms,” as explained in the figure caption (see eq. 194). We chose this presentation to have comparable results with A* coding (Flamich et al.). We hope this clarifies the above point.\n\n__Conclusion:__ We thank T3PN for their time and feedback. T3PN generally viewed our work favourably and suggested useful ways to improve our paper. We have highlighted our planned changes to our paper toward this end. If our response and suggested changes satisfy T3PN, we would like to invite them to consider increasing our score.\n"}]

Author rebuttal:
Dear Reviewers and ACs,

We would like to thank all four reviewers for the time they took reviewing our paper and for their valuable feedback. We found that the feedback from all four reviewers was generally positive. Most reviewers listed our theoretical contributions and the experimental validation of our approach among the main strengths of our work. They also thought that our paper was generally well-written and our exposition to be of a good standard.

Some of the reviewers also raised a few points of concern and some questions including, for example, the tractability of the CDFs used within GRC and technical clarifications on our derived bounds. In our view, we thoroughly addressed their points of concern in our response, justifying why these are not significant limitations of our method and clarifying their questions. The reviewers also gave useful suggestions to further improve our paper, such as adding more sketch proofs and discussions to provide additional intuition about our algorithms. In response, we provided an overview of concrete changes we plan to make in our manuscript for this purpose. We think these changes are beneficial and thank the reviewers for suggesting them.

In conclusion, we think the reviews for our work were overall positive. We hope that we have addressed the reviewers’ points of concern and their questions to their satisfaction and welcome further feedback and discussion on our work.



Thank you,

the authors of paper 5337


Meta-review: This paper presents a new algorithm for efficient compression of real-valued data, along with theoretical analysis and preliminary empirical validation. The authors are suggested to work on the presentation of their methods: the reviewers seemed overall to find the paper somewhat difficult to parse. However overall the paper was deemed to present a valuable contribution to an important problem.
Decision: Accept (poster)