I MPROVING L OSSLESS C OMPRESSION R ATES VIA M ONTE C ARLO B ITS -B ACK C ODING
A BSTRACT
Latent variable models have been successfully applied in lossless compression with the bits-back coding algorithm.
However, bits-back suffers from an increase in the bitrate equal to the KL divergence between the approximate posterior and the true posterior.
In this paper, we show how to remove this gap asymptotically by deriving bits-back schemes from tighter variational bounds.
The key idea is to exploit extended space representations of Monte Carlo estimators of the marginal likelihood.
Naively applied, our schemes would require more initial bits than the standard bits-back coder, but we show how to drastically reduce this additional cost with couplings in the latent space.
We demonstrate improved lossless compression rates in a variety of settings.
I NTRODUCTION
In principle, lossless data compression can be improved with better generative models for approximating the data generating distribution.
From the panoply of generative models, latent variable models are particularly attractive for compression applications, and have facilitated some of the most successful learned compressors for large scale natural images (see Yang et al.
(2020) for lossy, and Townsend et al.
(2020) for lossless), because they are ﬂexible and typically easy to parallelize.
The bits-back coding algorithm (Hinton & Van Camp, 1993) makes lossless compression with latent variable models tractable and efﬁcient.
However, since bits-back uses a variational approximation q ( z | x ) to the true posterior p ( z | x ) , it adds roughly D KL ( q ( z | x ) (cid:107) p ( z | x )) bits to the model’s optimal rate.
This seems unimprovable for a ﬁxed q if approximating p ( z | x ) is difﬁcult or expensive.
In this paper, we show how to remove (asymptotically) the D KL gap of bits-back schemes for (just about) any ﬁxed q .
Our method is based on recent work that derives tighter variational bounds using Monte Carlo estimators of the marginal likelihood (e.g., Burda et al., 2015).
The idea is that q and p can be lifted into an extended latent space (e.g., Andrieu et al., 2010) such that the D KL over the extended latent space goes to zero and the overall bitrate approaches − log p ( x ) .
For example, our simplest extended bits-back method, based on importance sampling, introduces N identically distributed latents
z i and a categorical random variable that picks from z i to approximate p ( z | x ) .
We also deﬁne extended bits-back schemes based on more advanced Monte Carlo methods (AIS, Neal, 2001; SMC, Doucet et al., 2001).
Adding O ( N ) latent variables increases the initial bit cost of our methods if naively applied.
We introduce a novel technique to reduce this cost to O (log N ) for some of our coders using couplings in latent space.
So, our coders extend bits-back onto a time-bitrate tradeoff with a negligible additional initial bit cost.
We test our methods in a variety of settings and explore the factors that affect the rate savings in our settings.
B ACKGROUND
Lossless compression with latent variable models can be practically achieved using the bits-back coding algorithm (Hinton & Van Camp, 1993), in particular the recently proposed efﬁcient coding method BitsBack with Asymmetric Numeral System (BB-ANS, Townsend et al., 2019).
We introduce it brieﬂy here and include a detailed introduction is in Appendix A.
ANS is an efﬁcient and near optimal entropy coder.
ANS stores data in a stack-like ‘message’ structure; encode pushes symbols onto the message and decode pops symbols from the message.
ANS has an important property: it is LIFO (i.e., the encoded symbols will be decoded in the opposite order), which makes it compatible with the bits-back algorithm.
For our purposes, the ANS message can be thought of as a store of randomness.
Given a message m with enough bits, regardless of m ’s provenance, we can decode from m using any distribution p .
This will return a random symbol x and remove roughly − log p ( x ) bits from m .
Conversely, we can encode a symbol x onto m with p , which will increase m ’s length by roughly − log p ( x ) .
See Fig.
4 in Appendix A.
A latent variable model is speciﬁed in terms of a joint distribution p ( x, z ) between a discrete observation random variable (or symbol) x ∈ S and a latent discrete random variable z ∈ S (cid:48) .
We assume that the joint distribution of latent variable models factorizes as p ( z ) p ( x | z ) and both p ( z ) and p ( x | z ) are tractable.
However, computing the marginal p ( x ) = (cid:80) z p ( z ) p ( x | z ) is often intractable.
This fact means that we cannot directly encode x onto m .
A naive strategy would be for the sender to pick some z ∈ S (cid:48) , and encode ( x, z ) using p , which would require − log p ( x, z ) bits; however, this involves communicating the symbol z , which is redundant information.
[Equation]
[Table]
[Table]
[Table]
/zC/
h0/
[Table]
[Table]
=">AAACF3icbVDLSgNBEOyNrxhfUY+
[Table]
+
[Table]
/skFSsXD4Eg3I6z49CzgNc6oNlI1my/7VNeln3jIQg/bpExa5OZHvOS63r4h52stck2uqs56NZuz83YXZJg4fZIrwMHX7fLD4X41+1H2Qhb7GGgmqFIlx450JaFScyawnSnHCiPKGvQMS4YG1EdVSbp/tcmqUTxSC6WpQJOu+nsiob5STd81nZ2D1aDXEf/
[Table]
//MqoWqe1SLmBaHSoccfNUOOlI+GDaEGE5Qo3tUEE8H0rYi0scBE6R4TugR7OvIsKR5l7OPMSV63gWCMOOzAHqTBhlPIwgXkoAAE7uABnuHFuDeejFfjbTwaMyY72/
AHxvs3XkehQw==</latexit>
BB-ANS gets a better bitrate by compressing sequences of symbols in a chain and by having the sender decode latents z from the intermediate message state, rather than picking z ; see Fig. 1.
Suppose that we have already pushed some symbols onto a message m .
BB-ANS uses an approximate posterior q ( z | x ) such that if p ( x, z ) =
then q ( z | x ) = 0 .
To encode a symbol x onto m , the sender ﬁrst pops z from m using q ( z | x ) .
Then, they push ( x, z ) onto m using p ( x, z ) .
The new message m (cid:48) has approximately − log p ( x, z )+log q ( z | x ) more bits than m .
m (cid:48) is then used in exactly the same way for the next symbol.
The per-symbol rate saving over the naive method is − log q ( z | x ) bits.
However, for the ﬁrst symbol, an initial message is needed, causing a one-time overhead.
Thus, we deﬁne the net bitrate to be the expected increase in message length per symbol and the initial bits to be the number of bits needed to initialize the message.
The distribution of the latent z , popped with q ( z | x ) , is approximately a sample from q (Townsend et al., 2019).
Then net bitrate of BB-ANS is the (negative) ‘evidence lower bound’ (ELBO)
[Equation]
which assumes z ∼ q ( z | x )
Thus, thereafter we refer to vanilla BB-ANS as BB-ELBO.
M ONTE C ARLO B ITS -B ACK C ODING
The net bitrate of bits-back is ideally the negative ELBO, which seems difﬁcult to improve without ﬁnding a better q .
However, the ELBO may be a loose bound on log p ( x ) .
Recent work in variational inference shows how to bridge the gap from the ELBO to the marginal log-likelihood with tighter variational bounds (e.g., Burda et al., 2015; Domke & Sheldon, 2018), motivating the question: can we derive bits-back coders from those tighter bounds and approach the model’s optimal rate?
In this section we provide an afﬁrmative answer with a framework called Monte Carlo bits-back coding (McBits).
We ﬁrst motivate our framework with simple examples.
Details are in Appendix B.
3. B ITS -B ACK I MPORTANCE S AMPLING
The simplest of our McBits coders is based on importance sampling (IS).
IS samples N particles
z i ∼ q
( z i | x ) i.i.d.
and uses the importance weights p ( x, z i ) /q
( z i | x ) to estimate p ( x ) .
The corresponding variational bound (IWAE, Burda et al., 2015) is the log-average importance weight:
[Equation]
IS provides a consistent estimator of p ( x ) .
If the importance weights are bounded, the left-hand side of equation (2) converges monotonically to − log p ( x ) (Burda et al., 2015).
Surprisingly, equation (2) is actually the ELBO between a different model and a different approximate posterior on an extended space (Andrieu et al., 2010; Cremer et al., 2017; Domke & Sheldon, 2018).
In particular, consider an extended latent space S (cid:48)
N × { 1 .. N } that includes the conﬁgurations of the N particles { z i } Ni =1 and an index j ∈ { 1 .. N } .
The left-hand side of equation (2) can be re-written as the (negative) ELBO between a pair of distributions P and Q deﬁned over this extended latent space, which are given in Alg.
3 in Appendix B.4.
Brieﬂy, given x , Q samples N particles z i ∼ q
z i | x ) i.i.d.
and selects one of them by sampling an index j with probability ˜ w j ∝ p ( x, z j ) /q
( z j | x ) .
The distribution P pre-selects the special j th particle uniformly at random, samples its value
z j ∼ p ( z j ) from the prior of the underlying model, and samples x ∼ p ( x | z j ) given z j .
The remaining z i ∼ q
( z i | x ) for i (cid:54) =
j are sampled from the underlying approximate posterior.
Because P and Q use q for all but the special j th particle, most of the terms in the difference of the log-mass functions cancel, and all that remains is equation (2).
See Appendix B.4.
[Table]
+laIO0ggrHSRWV2CNRl5mlT3itZ+8eBUtwHBCBmwCbZAAVjgEJTACSiDCsDgDjyBF/Bq3BvPxpvxPhpNGeOdDfAHxsc33L2iEw==</latexit> encode j with Cat(1 /N ) <latexit sha1_base64="wJwUgpFeCQ4W+BesI2a+WAp2Des=">AAACHXicbVDLSgNBEOyN7/iKehRkMAjxEndF0aPgxZOoGA0kIcxOJsmY2QczvUpY1ps/4cXf8OjFgyIevIjf4E84eSiaWDBQU9VNd5cbSqHRtj+s1Mjo2PjE5FR6emZ2bj6zsHimg0gxXmCBDFTRpZpL4fMCCpS8GCpOPVfyc7e13/HPL7nSIvBPsR3yikcbvqgLRtFI1cxW2aPYVF7MfRbUeELK5IJcf4tXApvJz2+fYpJzNg7Xq5msnbe7IMPE6ZPsHhx/3qzcnxxVM2/lWsAij/vIJNW65NghVmKqUDDJk3Q50jykrEUbvGSoTz2uK3H3uoSsGaVG6oEyz0fSVX93xNTTuu25prKzqB70OuJ/XinC+m4lFn4Yobm+N6geSYIB6URFakJxhrJtCGVKmF0Ja1JFGZpA0yYEZ/DkYXK2mXe28tvHJg0CPUzCMqxCDhzYgT04gCMoAINbeIAneLburEfrxXrtlaasfs8S/IH1/gWdZaW4</latexit>
IS)
Once we identify the left-hand side of equation (2) as a negative ELBO over the extended space, we can derive a bits-back scheme that achieves an expected net bitrate equal to equation (2).
We call this the Bits-Back Importance Sampling (BB-IS) coder, and it is visualized in Fig.
To encode a symbol x , we ﬁrst decode
N particles z i and the index j with the Q process by translating each ‘sample’ to ‘decode’.
Then we encode x , the particles z i , and the index j jointly with the P process by translating each ‘sample’ to ‘encode’ in reverse order .
By reversing P at encoding time, we ensure the receiver decodes with P in the right order.
Ideally BB-IS’s asymptotic net bitrate is close to the left-hand side of equation (2), which converges to the marginal log-likelihood (Burda et al., 2015).
Ultimately, as N → ∞ , it reaches the cross-entropy (model’s optimal rate).
Unfortunately, BB-IS requires roughly − log ˜ w j − (cid:80)
Ni =1 log q ( z i | x ) ∈ O ( N ) initial bits, because each decoded random variable z i needs to remove some bits from m .
To avoid this, we design Bits-Back Coupled Importance Sampling (BB-CIS), which achieves a net bitrate comparable to BB-IS while reducing the initial bit cost to O (log N ) by coupling the latent variables.
BB-CIS achieves this by decoding a single common random number which is shared by the z i and designing an encoding process that matches the modiﬁed decoding process.
Details in Appendix B.5.
3. G ENERAL F RAMEWORK
Monte Carlo bits-back coders generalize BB-IS and BB-CIS, which are built from extended latent space representations of Monte Carlo estimators of the marginal likelihood.
Let ˆ p N ( x ) be a positive unbiased estimator that can be simulated with O ( N ) random variables, i.e., E
[ˆ p N ( x )]
= p ( x ) .
A variational bound on the log-marginal likelihood can be derived from ˆ p N ( x ) by Jensen’s inequality, i.e., − E [log ˆ p N ( x )] ≥ − log p ( x ) .
If ˆ p N ( x ) is strongly consistent in N and − log ˆ p N ( x ) satisﬁes a uniform integrability condition, then − E
[log ˆ p N ( x )]
→ − log p ( x ) (Maddison et al., 2017).
This framework captures recent efforts on tighter variational bounds (Burda et al., 2015; Maddison et al., 2017; Naesseth et al., 2018; Le et al., 2018; Domke & Sheldon, 2018; Caterini et al., 2018).
As with BB-IS and BB-CIS, the key step is to identify an extended latent space representation of ˆ p N ( x ) .
Let Z ∼ Q
( Z | x ) be a set of random variables (often including those needed to compute ˆ p N ( x ) ).
If there exists a ‘target’ distribution P ( x, Z ) over x and Z with marginal p ( x ) such that
[Equation]
then the McBits coder, which decodes Z with Q ( Z | x ) and encodes ( x, Z ) with P ( x, Z ) , will achieve a net bitrate of − log ˆ p N ( x ) .
If the log estimator converges in expectation to − log p ( x ) , then ideally D KL ( Q ( Z | x ) (cid:107) P ( Z | x )) → 0 and the McBits coder will achieve the model’s optimal rate.
The challenge is to identify Z , Q , and P .
While Monte Carlo estimators of p ( x ) get quite elaborate, many of them admit such extended latent space representations (e.g., Neal, 2001;
Andrieu et al., 2010; Finke, 2015; Domke & Sheldon, 2018).
These constructions are techniques for proving the unbiasedness of the estimators, but in McBits they become bits-back schemes themselves.
In Appendix B, we derive Bits-Back Sequential Monte Carlo (BB-SMC) from SMC (Andrieu et al., 2010) and Bits-Back Annealed Importance Sampling (BB-AIS) from AIS (Neal, 2001).
We also derive a coupled variant of BB-SMC and a BitSwap (Kingma et al., 2019) variant of BB-AIS to reduce the initial bit cost. In Appendix C, we discuss their
16 32 64 128 N 6 .6 .
N e t B it r a t e ( b it s / s y m )
EntropyBB-ELBOBB-ISBB-CIS
BB-AISBB-AIS-BitSwap
[Figure]
12  32 N
T o t a l B it r a t e ( b it s / s y m )
BB-ELBOBB-ISBB-AIS BB-CIS&BB-AIS-BitSwap
[Figure]
1 2 4 8 16 32 64 128 256 N
N e t B it r a t e ( b it s / s e q )
EntropyBB-ELBOBB-ISBB-SMC
[Figure]
Letters
[Figure]
[Figure]
[Table]
computational cost.
All coders require O ( N ) time, but the IS- and SMC-based coders are amenable to parallelization over particles (see Fig. 11).
E XPERIMENTS
We ﬁrst studied the empirical properties of McBits coders on synthetic data and then highlighted the superiority of McBits coders in a transfer setting on image compression.
Details and more experiments, including musical piece and lossy image compression, are in Appendix D.
Empirical properties We assessed the convergence properties and initial bit cost of our McBits coders on synthetic data.
All coders were evaluated using the true data generating distribution with a uniform approximate posterior, ensuring a large mismatch with the true posterior.
First, a dataset of 5000 symbols was generated i.i.d.
from a mixture model with alphabet sizes 64 andfor the observations and latents, respectively.
The net bitrates of BB-IS, BB-CIS, and BB-AIS converged to the entropy (optimal rate) as N increased, as shown in Fig. 3a.
The indistinguishable gap between BB-CIS and BB-IS illustrates that particle coupling did not lead to a deterioration of net bitrate.
BB-AIS-BitSwap did not converge, likely due to the dirty bits issue (as discussed in Appendix D.1).
We then quantiﬁed the initial bit cost by computing the total bitrate (i.e, the net bitrate plus initial bits per symbol) after the ﬁrst symbol.
As shown in Fig.
3b, it increased linearly with N for BB-IS and BB-AIS, but remained ﬁxed for BB-CIS and BB-AIS-BitSwap.
Our second experiment was with a dataset of 5000 symbol subsequences generated i.i.d.
from a small hidden Markov model (HMM).
We used 10 timesteps with alphabet sizes 16 and 32 for the observations and latents, respectively.
BB-ELBO, BB-IS, and BB-SMC were evaluated.
The net bitrates of BB-IS and BB-SMC converged to the entropy, but BB-SMC converged much faster, illustrating the effectiveness of resampling particles for compressing sequential data (Fig. 3c).
Superiority in transfer setting We assessed BB-IS in a transfer learning setting on practical image compression.
We used an alphanumeric extension of MNIST, called EMNIST (Cohen et al., 2017), dynamically binarized following Salakhutdinov & Murray (2008), and a VAE with 1 stochastic layer as in Burda et al.
We trained models on the standard EMNIST-Letters and EMNIST-MNIST splits and evaluated compression performance on both test sets.
BB-IS achieved greater rate savings than BB-ELBO when transferred to an out-of-distribution setting (Table 1).
This illustrates that BB-IS may be particularly useful in more practical compression settings where the data distribution is different from that of the training data.
A D ETAILED B ACKGROUND A. A SYMMETRIC N UMERAL S YSTEMS
The goal of lossless compression is to ﬁnd a short binary representation of the outcome of a discrete random variable x ∼ p
d ( x ) in a ﬁnite symbol space x ∈ S .
Achieving the best possible expected length, i.e., the entropy H ( p d ) of p d , requires access to p d , and typically a model probability mass function (PMF) p ( x ) is used instead.
In this case, the smallest achievable length is the cross-entropy of p relative to p d ,
H ( p d , p ) = − (cid:80) x p d ( x ) log p ( x ) 1 .
See MacKay (2003); Cover & Thomas (1991) for more detail.
Asymmetric numeral systems (ANS) are model-based coders that achieve near optimal compression rates on sequences of symbols (Duda, 2009).
ANS stores data in a stack-like ‘message’ data structure, which we denote m , and provides an inverse pair of functions, encode p and decode p , which each process one symbol x ∈ S at a time: encode
[Equation]
Encode pushes x onto m , and decode pops x from m (cid:48) .
Both functions require access to routines for computing the cumulative distribution function (CDF) and inverse CDF of p .
If n symbols, drawn i.i.d.
from p d , are pushed onto m with p , the bitrate (bits/symbol) required to store the ANS message approaches H ( p d , p ) + (cid:15) for some small error (cid:15) (Duda, 2009; Townsend, 2020).
The exact sequence is recovered by popping n symbols off the ﬁnal m with p .
For our purposes, the ANS message can be thought of as a store of randomness.
Given a message m with enough bits, regardless of m ’s provenance, we can decode from m using any distribution p .
This will return a random symbol x and remove roughly − log p ( x ) bits from m .
Conversely, we can encode a symbol x onto m with p , which will increase m ’s length by roughly − log p ( x ) .
If a sender produces a message m through some sequence of encode or decode steps using distributions p
i , then a receiver, who has access to the p i , can recover the sequence of encoded symbols and the initial message by reversing the order of operations and switching encodes with decodes.
When describing algorithms, we often leave out explicit references to m , instead writing steps like ‘ encode x with p ( x ) ’.
See Fig. 4.
A. B ITS -B ACK C OMPRESSION WITH ANS
The class of latent variable models is highly ﬂexible, and most operations required to compute with such models can be parallelized, making them an attractive choice for model-based coders.
Bits-back coders, in particular Bits-Back with ANS (BB-ANS, Townsend et al., 2019), specialize in compression using latent variable models.
A latent variable model is speciﬁed in terms of a joint distribution p ( x, z ) between x and a latent discreterandom variable taking value in a symbol space z ∈ S (cid:48) .
We assume that the joint distribution of latent variable models factorizes as p ( z ) p ( x | z ) and that the PMFs, CDFs, and inverse CDFs, under p ( z ) and p ( x | z ) are tractable.
However, computing the marginal p ( x ) = (cid:80) z p ( z ) p ( x | z ) is often intractable.
This fact means that we cannot directly encode x onto m .
A naive strategy would be for the sender to pick some z ∈ S (cid:48) , and encode ( x, z ) using p , which would require − log p ( x, z ) bits; however, this involves communicating the symbol z , which is redundant information.
BB-ANS gets a better bitrate, by compressing sequences of symbols in a chain and by having the sender decode latents z from the intermediate message state, rather than picking z ; see Fig. 1.
Suppose that we have already pushed some symbols onto a message m .
BB-ANS uses an approximate posterior q ( z | x ) such that if p ( x, z ) =
then q ( z | x ) = 0 .
To encode a symbol x onto m , the sender ﬁrst pops z from m using q ( z | x ) .
Then they push ( x, z ) onto m using p ( x, z ) .
The new message m (cid:48) has approximately [Table] [Table] /kzXlyXpx352PeuuLkM0fkD5zPH0n0lq4=</latexit> totalbitrate <latexit sha1_base64="
[Table]
=">AAACA3icbVDLSgMxFM34rPU16k43wSK4KjNF0WXBjcsK9gGdoWTS2zY0kxmSO0IpBTf+ihsXirj1J9z5N6btLLT1QOBwzr1JzolSKQx63rezsrq2vrFZ2Cpu7+zu7bsHhw2TZJpDnScy0a2IGZBCQR0FSmilGlgcSWhGw5up33wAbUSi7nGUQhizvhI9wRlaqeMeBxwUghaqTzFBJoOgGAnUDKHjlryyNwNdJn5OSiRHreN+Bd2EZ7G9kEtmTNv3UgzHTKPgEibFIDOQMj5kfWhbqlgMJhzPMkzomVW6tJdoexTSmfp7Y8xiY0ZxZCdjhgOz6E3F/
[Table]
+Sc+
[Table]
/Ou/MxH11x8p0j8gfO5w/3mJef</latexit> entropy <latexit sha1_base64="NQLYIspRmiem9gIjuqgnMZn+lHA=">AAAB+XicbVBLSwMxGMzWV62vVY9egkXwIGVXFL1Z8OKxgn1Au5Rsmm1D81iSbGFZevZPePGgiFf/iTfvHv0D3sy2PWjrQGCY+b5kMmHMqDae9+EUlpZXVteK66WNza3tHXd3r6FlojCpY8mkaoVIE0YFqRtqGGnFiiAeMtIMh9e53xwRpakUdyaNScBRX9CIYmSs1HXdDkdmoHhGhFEyTsddt+xVvAngIvFnpHz1fU+/Tj5Jreu+d3oSJ9xegBnSuu17sQkypAzFjIxLnUSTGOEh6pO2pQJxooNsknwMj6zSg5FU9ggDJ
[Table]
crossentropy approx.
[Table]
+gjUjUA05SCGI2VCISnKGV+u5pj4NC0EIN6UBonNBQoOm7Fa/qzUFXiV+QCinQ6LtfvUHCs9jO4pIZ0/W9FIOcaRRcwrTcywykjI/ZELqWKhaDCfL5+lN6YZUBjRJtn0I6V3935Cw2ZhKHtjJmODLL3kz8z+tmGN0EuVBphqD44qMokxQTOstidjBwlBNLGNfC7kr5iGnGbSCmbEPwl09eJa1a1b+qeve1Sp0WcZTIGTknl8Qn16RO7kiDNAknOXkmr+TNeXJenHfnY1G65hQ9J+QPnM8fTKGVlw==</
[Figure]
− log p ( x, z ) + log q ( z | x ) more bits than m .
m (cid:48) is then used in exactly the same way for the next symbol.
The per-symbol rate saving over the naive method is − log q ( z | x ) bits.
However, for the ﬁrst symbol, an initial message is needed, causing a one-time overhead.
A. T HE B ITRATE OF B ITS -B ACK C ODERS
When encoding a sequence of symbols, we deﬁne the total bitrate to be the number of bits in the ﬁnal message per symbol encoded; the initial bits to be the number of bits needed to initialize the message; and the net bitrate to be the total bitrate minus the initial bits per symbol, which is equal to the expected increase in message length per symbol.
As the number of encoded symbols grows, the total bitrate of BB-ANS will converge to the net bitrate.
One subtlety is that the BB-ANS bitrate depends on the distribution of the latent z , popped with q ( z | x ) .
In an ideal scenario, where m contains i.i.d.
uniform Bernoulli distributed bits, z will be an exact sample from q .
Unfortunately, in practical situations, the exact distribution of z is difﬁcult to characterize.
Nevertheless, Townsend et al.
(2019) found the ‘evidence lower bound’ (ELBO)
[Equation]
which assumes z ∼ q ( z | x ) , to be an accurate predictor of BB-ANS’s empirical compression rate; the effect of inaccurate samples (which they refer to as ‘dirty bits’) is typically less than 1% .
So, in this paper we mostly elide the dirty bits issue, regarding (5) to be the net bitrate of BB-ANS, and hereafter we refer to vanilla BB-ANS as BB-ELBO.
Taking the expectation under p d , BB-ELBO achieves a net bitrate of approximately H ( p d , p ) + E x ∼ p
d [ D KL ( q ( z | x ) (cid:107) p ( z | x ))] .
B M ONTE C ARLO B ITS -B ACK C ODERS B. N OTATION
• I
( A ) is the indicator function of the event A , i.e., I ( A )
= 1 if A holds and 0 otherwise.
• P ( A ) is the probability of the event A .
B. A SSUMPTIONS
• We assume that, if q ( z | x ) > 0 , then p ( x, z ) > 0 .
B. G ENERAL F RAMEWORK
Procedure Encode( symbol x , message m ) decode Z with Q ( Z | x ) encode x and
Z with P ( x, Z ) return m (cid:48)
Procedure Decode( message m ) decode x and Z with P ( x, Z ) encode
Z with Q ( Z | x ) return x , m (cid:48)
As discussed in the main body, the McBits coders can asymptotically drive the KL gap to zero, as illustrated in Fig.
Given the extended space representation of an unbiased estimator of the marginal likelihood, we can derive the general procedure of McBits coders as Alg. 1.
The decode procedure could be easily derived from the encode procedure by simply reverse the order and switch ‘encode’ with ‘decode’, as in Alg. 2.
Therefore, we do not tediously present the decode procedure for each McBits coder in the following subsections.
Process Q ( Z | x ) sample { z i } Ni =1 ∼ (cid:81)
Ni =1 q
( z i | x ) compute ˜ w i ∝ p ( x,z i )
q
( z i | x ) sample j ∼ Cat ( ˜ w j ) return { z i } Ni =1 , j
Process P ( [Table] , Z )
[Table]
return x, { z i } Ni =1 , j
Procedure Encode( symbol x , message m ) decode { z i } Ni =1 with (cid:81) Ni
=1 q ( z i | x ) decode j with Cat ( ˜ w j ) encode { z i } i (cid:54) =
j with (cid:81) i (cid:54) =
j q ( z i | x ) encode x with p ( x | z j ) encode z j with p ( z j ) encode j with Cat(1 /N )
return m (cid:48)
B.
B ITS -B ACK I MPORTANCE S
AMPLING (BB-IS)
Importance sampling (IS) samples N particles z i ∼ q ( z i | x ) i.i.d.
and uses the average importance weight to estimate p ( x ) .
The corresponding variational bound (IWAE, Burda et al., 2015) is the logaverage importance weight which is given in equation (2).
For simplicity, we denote importance weight as w i = p
( x, z i ) /q
( z i | x ) and normalized importance weight as ˜ w i = w i / (cid:80) i w i .
The basic importance sampling estimator of p ( x ) is given by
[Equation]
[Equation]
Extended latent space representation
The extended latent space representation of the importance sampling estimator is presented in Alg. 3, from which we can derive the proposal and target distribution as followed:
[Equation]
Now note,
[Equation]
which exactly gives us the IS estimator.
BB-IS Coder Based on the extended latent space representation, the Bits-Back Importance Sampling (BB-IS) coder is derived in Alg. 4 and visualized in Fig.
6a.
The expected net bit length for encoding a symbol x is:
[Equation] [Equation]
Ignoring the dirty bits issue, i.e., z i ∼ q
( z i | x ) i.i.d., the expected net bit length exactly achieves the negative IWAE bound.
[Table]
[Table]
+DhokgRbswv6PeEVlO
[Table]
/olqA4IRMmATbIECMMEBKIJjUAJlgMEdeAIv4FW71561N
[Table] + [Table] /naqKQ</latexit>
[Table]
<latexit
[Table]
/wDC/WvfVkvVpv/dIxa9CzAr9gfXwDA4Ol4g==</latexit>
[Figure]
B. B ITS -B ACK C OUPLED I MPORTANCE S AMPLING
Unfortunately, the BB-IS coder requires roughly − log ˜ w j − (cid:80)
Ni =1 log q ( z i | x ) ∈ O ( N ) initial bits.
The reason is that each decoded random variable needs to remove some bits from m .
Can this be avoided?
Here, we design Bits-Back Coupled Importance Sampling (BB-CIS), which achieves a net bitrate comparable to BB-IS while reducing the initial bit cost to O (log N ) .
BB-CIS is based on a reparameterization of the particles z i as deterministic functions of coupled uniform random variables.
The method is a discrete analog of the inverse CDF technique for simulating non-uniform random variates (Devroye, 2006).
Speciﬁcally, suppose that the latent space S (cid:48) is totally ordered, and the probabilities of q are approximated to an integer precision r > 0 , i.e., for all z ∈ S (cid:48)
[Equation]
(11) where q z is an integer.
Note that this assumption is also required for ANS, so this is not an additional assumption.
Deﬁne the unnormalized cumulative distribution function F q of q as well as the following related objects:
[Equation]
[Equation]
Notice that | U ( z ) | = q z .
Thus, if u ∼ unif { 0 .. 2 r − 1 } , then
[Equation]
Therefore we can reparameterize the sampling process of z as uniform sampling over { 0 .. 2 r − 1 } followed by the deterministic mapping F − 1 q .
This is a classical approach in non-uniform random variate generation, specialized to this discrete setting.
This is visualized in Fig.
Coupled importance sampling (CIS) samples the latent variables by coupling their underlying uniforms.
In particular, a coupled set of particles z i with marginals q ( z i | x ) can be simulated with a common random number by sampling a single uniform u and setting z i = F − 1 q ( T i ( u )) for some functions T i : { 0 .. 2 r − 1 } → { 0 .. 2 r − 1 } .
Intuitively, T i maps u to the uniform u i underlying z i . To ensure that z i
[Table]
[Figure]
Process Q ( Z | x ) sample u∼ unif { 0 , . . .
, 2 r − 1 } for i = 1 , . . . , N do assign u i
= T i ( u 1 ) assign z
i
=
F − 1 q ( u i ) compute ˜ w i ∝ p ( x, z i ) /q
( z i | x ) sample j ∼ Cat( ˜ w i ) return { z i } N i =1 , { u i } N i =
1 , j Process P ( x, Z ) sample j ∼ Cat(1 /N )
sample
z j ∼ p ( z j ) sample x ∼ p ( x | z j ) sample u j ∼ unif { u : F − 1 q ( u ) =
z j } for i (cid:54) =
j do assign u i
= T i ( [Equation]
[Equation]
( u j )) assign z i = F − 1 q ( u i ) return x, { z i [Equation] Ni =1 ,
[Equation] u i [Equation] =1 , j have the same marginal, we require that T i are bijective functions.
Precisely, if u ∼ unif { 0 .. 2 r − 1 } , then
[Equation]
Thus, if u ∼ unif { 0 .. 2 r − 1 } , then T i ( u ) d = u and F − 1 q ( T i ( u )) ∼ q ( z | x ) .
For example, simple bijective operators T i can be deﬁned as applying ﬁxed “sampling shift” ¯ u i ∈ { 0 , 1 , . . .
, 2 r − 1 } to u :
[Equation]
For simplicity, we deﬁne T 1 to employ the zero shift ¯ u 1 = 0 , i.e., T 1 ( u ) = u .
We now have the deﬁnitions that we need to analyze the extended space construction for coupled importance sampling.
Let u 1 ∼ unif { 0 .. 2 − 1 } .
As with IS, we denote importance weight as w i = p ( x, F − 1 q ( T i ( u 1 ))) /q
( F − 1 q ( T i ( u 1 ))
| x ) and the normalized importance weight as ˜ w i = w i / (cid:80) i w i .
The coupled importance sampling estimator of p ( x ) is given
[Equation]
Extended latent space representation
The extended latent space representations of the coupled importance sampling estimator is presented in Alg.
This extended space construction is novel, and we believe it may be useful for deriving other coupling schemes that reduce initial bit consumption.
The Q and P processes have the following probability mass functions.
For convenience, let U = { 0 .. 2 r − 1 } and z i =
F − 1 q ( T i ( u 1 )) . [Equation] decode u 1 with unif { 0 , 1 , . .
, 2 r − 1 } assign u i
= T i ( u ) and z i =
F − 1 q ( u i ) for i ∈ { 1 , . .
, N } decode j with Cat ( ˜ w j ) encode u j with unif { u : F − 1 q ( u ) =
z j } encode x with p ( x | z j ) encode z j with p ( z j ) encode j with Cat(1 /N )
return m (cid:48) where we used the fact that q z j = 2 r q ( z j | x ) =
| U ( z j ) | .
Because each T i is a bijection, we have that the range of T i ( u 1 ) is all of U , as u 1 ranges over U .
Thus,
[Equation]
Moreover, for any ( u, z ) ∈ U × S (cid:48) ,
I ( u ∈ U ) 2 r
I (cid:0)
[Equation] [Equation]
=
[Equation]
(cid:1) =
[Equation]
( u
[Equation]
( z
[Equation]
= q ( z
[Equation]
( u
[Equation]
[Equation]
[Equation]
| x ) .
Thus,
Q ( { z i
[Equation]
Ni =
[Equation]
, { u i
[Equation]
=
[Equation]
[Equation]
=
[Equation]
[Equation]
U
[Equation]
r
N (cid:89) i =
I ( T i ( u 1 ) =
u i )
N (cid:89) i =
I (cid:0) F − 1 q ( u i
[Equation]
From this it directly follows that
[Equation]
which is exactly the CIS estimator.
BB-CIS coder Based on the extended latent space representation, the Bits-Back Coupled Importance Sampling (BB-CIS) coder is derived in Alg. 6 and visualized in Fig.
6b.
The expected net bit length for encoding a symbol x is:
[Equation]
Here we explain BB-CIS in more detail to give intuition about the derivations of BB-CIS and its extended latent space representation.
BB-CIS uses couplings in a latent decoding process that saves initial bits.
It decodes u 1 with unif { 0 .. 2 r − 1 } , sets z i = F − 1 q ( T i ( u 1 )) , and decodes the index j of the special particle
z j with Cat( ˜ w j ) .
This reduces the initial bit cost to r − log ˜ w j ∈ O (log N ) .
However, the challenge is showing that a net bitrate comparable to BB-IS is still achievable under such a reparameterization.
The coupled latent decoding process needs to be matched with an appropriate encoding process for x .
Suppose that we ﬁnished encoding x , as with BB-IS, by encoding { z i } i (cid:54) =
j with q ( z i | x ) i.i.d.
and encoding ( x, z j , j ) with p ( x, z j ) /N .
The net bitrate would be
[Equation]
This is clearly worse than BB-IS.
The culprits are the N − 1 latents that BB-IS pushes onto m , which is wasteful, because they are deterministically coupled.
Fundamentally, the encoding process of BB-IS is not balanced with the initial bit savings of our coupled latent decoding process.
The solution is to design an encoding process over { z i } Ni =1 , { u i } Ni =1 , and x , which exactly matches the initial bit savings of the coupled decoding process.
The idea is to encode just ( u j , x, z j , j ) , which is enough information for the receiver to reconstruct all other variables.
The key is to design the encoding for u j .
Encoding u j with a uniform on { 0 .. 2 r − 1 } is unnecessarily wasteful, because z j restricts the range of u j .
It turns out, that the best we can do is to encode u j with unif { u : F − 1 q ( u ) =
z j } .
Finally ( x, z j , j ) are encoded with p ( x, z j ) /N .
This gives the BB-CIS coder and its expected net bitrate exactly achieves equation (24) which is comparable to that of BB-IS.
The convergence of BB-CIS’s net bitrate to − log p ( x ) will depend on the T i .
There are many possibilities, but one consideration is that both sender and receiver must share the T i (or a procedure for generating them).
Another consideration is that the number of particles should never exceed 2 r , N ≤ 2 r .
This is because we can execute numerical integration with a budget of 2 r particles.
Assuming T i ( u )
=
i for i ∈ { 0 .. 2 r − 1 } :
[Equation]
We brieﬂy mention two possibilities:
Let T i ( u ) =
( u + k i
mod 2 r where k i ∼ unif { 0 .. 2 r − 1 } i.i.d. generated by a pseudorandom number generator where the sender and receiver share the seed.
This will enjoy a convergence rate similar to BB-IS, because { ( u + k i ) mod 2 r } Ni =1 will be i.i.d.
uniform on { 0 .. 2 r − 1 } for u ∼ unif { 0 .. 2 r − 1 } .
This is used in our experiments.
Inspired by the idea of numerical integration, let T i ( u ) =
( u + k i
mod 2 r where k i is the i th element of a random permutation of { 0 .. 2 r − 1 } .
With this scheme, BB-CIS is performing numerical integration on a random permutation of { 0 .. 2 r − 1 } .
The rate at which the net bitrate converges to − log p ( x ) could clearly be made worse by inefﬁcient permutations.
Randomizing the order may help avoid such inefﬁcient permutations.
We did not experiment with this.
B. B ITS -B ACK A NNEALED I MPORTANCE S AMPLING
Annealed importance sampling (AIS) generalizes importance sampling by introducing a path of intermediate distributions between the tractable base distribution q ( z | x ) and the unnormalized target distribution p ( x, z ) .
For AIS with N steps, for i ∈ { 0 , . . .
, N } , deﬁne annealing distributions
[Equation]
[Equation]
Note that π 0 ( z ) = f 0 ( z ) =
q ( z | x ) .
Then deﬁne the MCMC transition operator T i that leave the intermediate distribution
π i invariant and its reverse ˜ T
i :
[Equation] [Equation]
[Equation]
Note, ˜ T i is a normalized distribution.
AIS samples a sequence of latents { z i } Ni =1 from base distribution and the MCMC transition kernels (see the Q process in Alg. 7) and obtains an unbiased estimate of p ( x ) as the importance weight over the extended space.
[Equation]
Process Q ( Z | x ) sample z∼ q ( z 1 | x ) =
π 0 ( z 1 ) for i = 1 , . . .
, N − 1 do sample z
i +1 ∼ T i (
z i +1 |
z i ) return { z i } Ni =
Process P ( x, Z ) sample z N ∼ p
( z N ) for i = N − 1 , . . .
, 1 do sample
z i ∼ ˜ T i ( z i
| z i +1 )
sample x ∼ p ( x | z N ) return x, { z i } Ni =
Procedure Encode( symbol x , message m ) decode z 1 with q ( z 1 | x ) = π 0 ( z 1 ) for i = 1 , . . .
, N − 1 do decode z
i +1 with T
i
( z i +1 |
z i ) encode x with p ( x | z N ) for i = 1 , . . .
, N − 1 do encode z i with ˜ T
i ( z i
| z i +1 )
encode z N with p ( z N ) return m (cid:48)
The corresponding variational bound is the log importance weight:
[Equation]
Extended latent space representation
The extended space representation of AIS is derived in (Neal, 2001), as presented in Alg.
The extended latent variables contain all the latents { z i } Ni =1 .
Brieﬂy, given x , the distribution Q ﬁrst samples z 1 from the base distribution q ( z 1 | x ) and then sample { z i } Ni =2 sequentially from the transition kernel T i ( z i +1 |
z i ) .
The distribution P ﬁrst samples
z N from the prior p ( z N ) , then samples { z i } 1 i = N − 1 from the reverse transition kernel ˜ T i ( z i
| z i +1 )
in the reverse order, and ﬁnally samples x from p ( x | z N ) .
The proposal distribution and the target distribution can be derived as followed:
[Equation]
Now note
[Equation]
[Equation]
which exactly gives us the AIS estimator.
Note that we have used equation (30) above.
BB-AIS coder Based on the extended space representation of AIS, the Bits-Back Annealed Importance
Sampling coder is derived in Alg. 8 and visualized in Fig.
8a.
The expected net bit length for encoding a symbol x is:
[Equation]
Ignoring the dirty bits issue, i.e., { z i } N i =1 ∼ Q ( { z i } N i =1 | x ) , the expected net bit length exactly achieves the negative AIS bound. encode [Table]
<latexit sha1_base64="6MOGo8j2Jxq8zOvzwjge1cGg6yo=">AAACGHicbZDJSgNBEIZr4h63qEdBGoOglzgjih4DXjwmYkwgCaGnp2Ka9Cx296hxjDcfwYtPInjxoIhXbz6DL2FnEdx+aPj5qoqu+t1IcKVt+91KjYyOjU9MTqWnZ2bn5jMLi8cqjCXDEgtFKCsuVSh4gCXNtcBKJJH6rsCy297v1ctnKBUPgyPdibDu05OANzmj2qBGZrPmU92SfuIhCz3skhohlw2HXH/xc65b3Wtyum7o1cVGI5O1c3Zf5K9xhiabh+LHzcr9YaGReat5IYt9DDQTVKmqY0e6nlCpORPYTddihRFlbXqCVWMD6qOqJ/3DumTNEI80Q2leoEmffp9IqK9Ux3dNZ29d9bvWg//VqrFu7tUTHkSxxoANPmrGguiQ9FIiHpfItOgYQ5nkZlfCWlRSpk2WaROC8/vkv+Z4K+ds53aKJg0CA03CMqzCOjiwC3k4gAKUgMEtPMATPFt31qP1Yr0OWlPWcGYJfsh6+wTYqaMo</latexit> [Table] <latexit sha1_base64="N1TzDsdtcP59xi4Uex43Py+mvuU=">AAACJnicbVDLSgNBEOz1GeMr6lGQQRH0EnaDohch4MWjSqJCEpbZ2Y4ZnH0wM6sk [Table]
=<
[Table]
</
[Table]
[Table]
+jnKTRUf29E2JWy4zp6cnCtnPQG4n9eKVT1s0rEvCBUOvXoo3rIkfLRoCJUY4ISxTuaYCKYvhWRJhaYKF1kXJdgT0aeJvnDlH2cOsnqNhCMEINt2IUk2HAKabiADOSAwAP04RlejEfjyXg13kajM8Z4Zwv
[Table]
[Table]
</
[Table]
sha
[Table]
+cG+ev828YHXEeZ5bgBZz7BwAvrxg=</latexit> encode z 3 with p ( z 3 ) <latexit sha1_base64="0Jy4CqmHZTdc82GrJJFywLMzZUA=">AAACFXicbVDJSgNBFHwTtxi3qEdBGoMQQcKMC3oMePGYiFkgCUNPp5M06Vno7lHiMLn5A178lVw8KOJV8OY3+BN2FkETCxqKqvfoV+UEnEllmp9GYm5+YXEpuZxaWV1b30hvbpWlHwpCS8Tnvqg6WFLOPFpSTHFaDQTFrsNpxeleDP3KDRWS+d616gW04eK2x1qMYKUlO31Yd7HqCDeiHvGbNEZ1hO7sY9T/0W+Z6sT9IKvFAzudMXPmCGiWWBOSyUPx6353cFWw0x/1pk9Cl3qKcCxlzTID1YiwUIxwGqfqoaQBJl3cpjVNPexS2YhGqWK0r5UmavlCP0+hkfp7I8KulD3X0ZPDW+W0NxT/82qhap03IuYFodKpxx+1Qo6Uj4YVoSYTlCje0wQTwfStiHSwwETpIlO6BGs68iwpH+Wsk9xpUbeBYIwk7MAeZMGCM8jDJRSgBAQeYADP8GI8Gk/Gq/E2Hk0Yk51t+APj/RupcqID</latexit>
[Figure]
a symbol x with BB-AIS incurs a very large initial bit cost that scales like O ( N ) .
This can be signiﬁcantly reduced by applying the BitSwap trick.
The BB-AIS coder has several practical issues.
First is the computational cost.
Since the intermediate distributions are usually not factorized over the latent dimensions, it is very inefﬁcient to encode and decode the latents with entropy coders for high dimensional latents.
Second is the increased initial bit cost as with BB-IS.
More precisely, the initial bits that BB-AIS requires is as followed:
[Equation]
which scales like O ( N )
However, because the structure of BB-AIS is very similar to the hierarchical latent variable modelss with Markov structure in (Kingma et al., 2019), the BitSwap trick can be applied with BB-AIS to reduce the initial bit cost.
The BB-AIS coder with BitSwap is derived in Alg. 9 and visualized in 8b.
In particular, note that the latents { z i } Ni =1 in BB-AIS are encoded and decoded both in the forward time order, we can interleave the encode/decode operations such that z i are encoded with z i + as soon as z i +1 are decoded.
Thus there are always at least − log ˜ T
i ( z i
| z i +1 )
available for decoding z i +2 at the next step, and the initial bit cost is bounded by:
[Equation]
[Equation] Although BitSwap helps to reduce the initial bit cost, we ﬁnd that it suffers more from the dirty bits issue than the naive implementation and affects the expected net bit length (see Fig. ??
for an empirical study).
Therefore, using BB-AIS with BitSwap leads to a trade-off between net bitrate distortion and inital bit cost that depends on the number
B. B ITS -B ACK S EQUENTIAL M ONTE C ARLO
Sequential Monte Carlo (SMC) is a particle ﬁltering method that combines importance sampling with resampling, and its estimate of marginal likelihood typically converges faster than importance sampling for sequential latent variable models C´erou et al.
(2011); B´erard et al.
Suppose the observations are a sequence of T random variables x T ∈ S T , where x t := x 1: t .
Sequential latent variable models introduce a sequence of T (unobserved) latent variables z T ∈ S (cid:48)
T associated with x T , where z t :
= z 1: t .
We assume the joint distribution p ( x T , z T ) can be factored as:
[Equation]
where f and g are (generalized) transition and emission distributions, respectively.
For the case t = 1 , f ( z 1 | ∅ , ∅ ) =
µ ( z 1 ) reduces to the prior distribution and g ( x 1 | ∅ , z 1 ) only conditions on z 1 .
We also assume the approximate posterior q ( z
T | x T ) can be factorized as:
[Equation]
For the case t = 1 , q ( z 1 | x 1 , ∅ ) only conditions on x 1 .
SMC maintains a population of particle states
Z t :=
z 1: N 1: t .
Here we use subscript for indexing timestep t and superscript for indexing particle i .
In this appendix we describe the version of SMC that resamples at every iteration.
Adaptive resampling schemes are possible (Doucet et al., 2001), but we omit them for clarity.
At each time step t , each particle independently samples an extension z it ∼ q
( z it | x t , τ it − 1 ) where τ it − 1 is the ancestral trajectory of the i th particle at timestep t (will be described in details later).
The sampled extension z it is appended to τ it − 1 to form the trajectory of i th particle as ( τ it − 1 , z it ) .
Then the whole population is resampled with probabilities in proportion to importance weights deﬁned as followed:
[Equation]
In particular, each particle samples a ‘parent’ index A it ∼ Cat( ˜ w it ) .
At the next timestep t + 1 , the i th particle ‘inherits’ the ancestral trajectory of the A it
th particle.
More precisely,
[Equation]
Thus, at timestep t +1 , given the parent index A
it , as well as the collection of particle states Z t and previous ancestral indices
A t − 1 :
= A 1: N 1: t − 1 , one can trace back the whole ancestral trajectory τ it by recursively applying equation (46).[-assign B it − 
 =-]
[-A B ik +1 k assign τ it − 
-]
For brevity and consistency, we provide the pseudocode (Alg. 10) and visualization (Fig. 9) for tracing back the ancestral trajectory τ
it − 1 for the i th particle at timestep t .
In particular, we ﬁrst obtain the ancestral index B i k at each previous timestep k by using the fact
[Equation]
Thus we can obtain a series of ancestral indices ( B i 1 , . . .
, B it − 1 ) in the backward time order and then obtain the ancestral trajectory by indexing Z t − 1 .
Note that Alg.
10 is only for helping understand the
SMC algorithm, in practice, typically the ancestral trajectories are book-kept and updated along the way using equation (47).
SMC obtains an unbiased estimate of p ( x T ) using the intermediate importance weights: [Equation]
[Equation]
The corresponding variational [Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
[Table]
is
[Equation]
{+Procedure TraceBack( Z t − 1 , A t − 2 , A it − 1 ) assign B it − 
+}
=
A it − 1 for k = t − 2 , . .
. , 1 do assign B ik =
{+A B ik +1 k assign τ it − 
+}
=
( z B i 1 1 , z B i 2 2 , . . .
, z B it − 1 t − 1 ) return τ it −<latexit sha1_base64="ZUbYK2UYOOCh3R1DVU/1FkE2WQ0=">AAACFXicbVC7TsMwFHXKq5RXgZHFokViQFVSicdYxMJYJPqQmhA5jtOaOnZkO0hR1J9g4VdYGECIFYmNv8F9DNByJNtH59wr33uChFGlbfvbKiwtr6yuFddLG5tb2zvl3b22EqnEpIUFE7IbIEUY5aSlqWakm0iC4oCRTjC8GvudByIVFfxWZwnxYtTnNKIYaSP55ZNLjonSEjFornuCtZAZrLqBYKHKYvPkrkbpyD+9c6p+uWLX7AngInFmpAJmaPrlLzcUOI0J15ghpXqOnWgvR1JTzMio5KaKJAgPUZ/0DOUoJsrLJ1uN4JFRQhgJaQ7XcKL+7shRrMYjmsoY6YGa98bif14v1dGFl1OepJpwPP0oSk0AAo4jgiGVJgeWGYKwpGZWiAdIIqxNkCUTgjO/8iJp12vOWc2+qVcacBZHERyAQ3AMHHAOGuAaNEELYPAInsEreLOerBfr3fqYlhasWc8++APr8wev+58E</latexit>
Ancestral trajectory
[Table]
z
[Table]
<
[Table]
sha
[Table]
="
[Table]
="
[Table]
[Table]
7VVZRIQl0S8Ug2fawoZ4K6mmlOm7GkOPQ5bfiDy7HfuKVSsUjU9DCmXoh7ggWMYG0k965Tuil38gW7aE
[Table]
[Table]
[Table]
[Table]
[Table]
[Figure]
space representation
The extended space representation of SMC is derived in (Andrieu et al., 2010), as presented in Alg.
The extended latent space variables contain the particle states Z T , ancestral indices A T − 1 , and a particle index j used to pick one special particle trajectory.
Brieﬂy, the Q distribution samples Z T and A T − 1 in the same manner as SMC, with an additional step that samples the particle index j with probability proportional to the importance weight w T .
The P distribution selects the ancestral indices of the special particle ( B 1 , . . .
, B T ) uniformly, and samples the special particle trajectory z ∗ T =
( z B 1 1 , . . .
, z B T T ) and the observations x T jointly with the underlying model distribution.
The remaining particle states and ancestral indices are sampled with the same distribution as Q .
Note that the special particle trajectory z ∗ T =
TraceBack ( Z T , A T − 1 , j ) .
From Alg.
11, the proposal distribution and the target distribution and be derived as followed:
[Equation]
, A T − 1 , j | x T ) =
˜ w jT
N (cid:89) i = q
( z i 1 | x 1 )
T (cid:89) t =
N (cid:89) i = ˜
w A it − 1 t − 1 q ( z it | x t , τ it − 1 ) (50)
P ( x T , Z T , A T − 1 , j ) =
1 N T µ ( z B 1 1 ) g
( x 1 | z B 1 1 )
T (cid:89) t = f
( z B t t | x t − 1 , τ B t t − 1 ) g
( x t | x t − 1 , ( τ B t t − 1 , z B t t ))
(cid:89) i (cid:54) =
Bq ( z i 1 | x 1 )
T (cid:89) t
= (cid:89) i (cid:54) =
B t ˜ w A it − 1 t − 1 q ( z it | x t , τ it − 1 ) (51) =
N T p ( x T , z ∗ T ) (cid:89) i (cid:54) =
Bq ( z i 1 | x 1 )
T (cid:89) t
= (cid:89) i (cid:54) =
B t ˜ w A it − 1 t − [Equation] ( z it | x t , τ it − 1 ) . Process Q ( Z | x T ) for t = 1 , . .
, T do if t (cid:54) =
1 then for i = 1 , . . .
, N do sample A it − 1 ∼ Cat( ˜ w t − 1 ) assign A t − 1 = (cid:2)
A t − 2 , A 1: N t − 1 (cid:3) for i = 1 , . .
. , N do assign τ it − 1 =
TraceBack
( Z t − 1 , A t − 2 ,
A it − 1 ) sample z it ∼ q
( z it | x t , τ it − 1 ) assign Z t = (cid:2)
Z t − 1 , z 1: N t (cid:3) sample j ∼ Cat( ˜ w T ) return Z
T , A T − 1 , j
Process P ( x T , Z ) for t = 1 , . . .
, T do sample B t ∼ Cat(1 /N )
if t (cid:54) = 1 then assign A B t t − 1 =
B t −for i (cid:54) =
B t do sample A it − 1 ∼ Cat( ˜ w t − 1 ) assign A t − 1 = (cid:2)
A t − 2 , A 1: N t − 1 (cid:3) assign τ B t t − 1 =
TraceBack ( Z t − 1 , A t − 2 , B t − 1 ) sample z B t t ∼ f
( z B t t | x t − 1 , τ B t t − 1 ) sample x t ∼ g ( x t | x t − 1 , ( τ B t t − 1 , z B t t )) for i (cid:54)
=
B t do assign τ it − 1 =
TraceBack ( Z t − 1 , A t − 2 , A it − 1 ) sample z it ∼ q
( z it | x t , τ it − 1 ) assign Z t = (cid:2)
Z t − 1 , z 1: N t (cid:3) assign j = B T return x T , Z T , A T − 1 , j Now note,
[Equation]
which exactly gives us the SMC estimator.
Note that we have used equation (47) and B T =
j in the above derivation. BB-SMC coder Based on the extended space representation of SMC, the Bits-Back Sequential Monte Carlo coder is derived in Alg.
12 and visualized in Fig.
Intuitively, BB-SMC ﬁrst decodes all the
[Table]
<latexit sha1_base64="VKv5eec8npVsvkyQKCQat944wCE=">AAAB6HicbZDJSgNBEIZr4hbjFpebl8YgeAozoujNgAc9JpANkiH0dGqSNj0L3T1CHPIEXjwo4lU8+zDePPomdpaDRn9o+Pj/KrqqvFhwpW3708osLC4tr2RXc2vrG5tb+e2duooSybDGIhHJpkcVCh5iTXMtsBlLpIEnsOENLsd54xal4lFY1cMY3YD2Qu5zRrWxKtVOvmAX7YnIX3BmULh4u
[Table] [Table]
<
[Table] >
Encode other extended latents with the same distributions as Q <latexit sha1_base64="cbRp2i2V6/sBSTh3TQFxh5u0iig=">AAACKXicbVBNSwMxEM3W7/pV9ehlsBU8ld2i6LEggscWrC3UUrLZqQ1mN0syq5bi3/HiX/
[Table] [Table]
++ [Table] [Table]
+SLSIovd/0Jxa9uBn1JnyA1JofAuf5FZTLm44pfYdjBxi21nOHZ6B7uOiaCnjXsJwZj9OTHksbWDOHSdMae [Table]
[Figure]
Procedure Encode( symbol x T , message m ) for t = 1 , . .
, T do if t (cid:54) =
1 then for i = 1 , . . .
, N do decode A
it − 1 with Cat( ˜ w t − 1 ) assign A t − 1 = (cid:2)
A t − 2 , A 1: N t − 1 (cid:3) for i = 1 , . .
. , N do assign τ it − 1 =
TraceBack
( Z t − 1 , A t − 2 ,
A it − 1 ) decode z it with q ( z it | x t , τ it − 1 ) assign Z t = (cid:2)
Z t − 1 , z 1: N t (cid:3) decode j with Cat( ˜ w T ) set the ancestral lineage B T = j and B t =
A B t +1 t for t =
T − 1 , . .
, 1 for t =
T, . .
. , 1 do for i (cid:54) =
B t do encode z it with q ( z it | x t , τ it − 1 ) encode x t with g ( x t | x t − 1 , ( τ B t t − 1 , z B t t )) encode z B t t with f ( z B t t | x t − 1 , τ B t t − 1 ) if t (cid:54)
= 1 then for i (cid:54)
=
B t do encode A it − 1 with Cat( ˜ w t − 1 ) encode B t with Cat(1 /N )
return m (cid:48) extended latent variables
Z = { Z T , A T − 1 , j } in the forward time order.
Then it picks a special particle and traces its trajectory z ∗ T backward in time.
The distribution of the special trajectory can be seen as a non-parametric approximation of the true posterior.
The special trajectory z ∗ T is used to encode the sequential observations x T and is itself encoded with with the model distribution.
The special trajectory’s ancestral indices are encoded using a uniform distribution.
All other extended latent variables are encoded back with the same distribution as decoding and in the backward time order.
The expected net bit length for encoding a symbol x T is:
[Equation]
Ignoring the dirty bits issue, the expected net bit length exactly achieves the negative FIVO bound. Note that our BB-SMC scheme can be easily extended to adaptive resampling setting where the particles are only resampled when a certain resampling criteria is satisﬁed. In adaptive schemes, the importance weights accumulate multiplicatively over time (between resampling events).
One typical adaptive resampling criteria is applying resampling if the effective sample size (ESS) of the particles drops below N/ 2 . Details can be found in (Doucet et al., 2001). Adaptive resampling is possible, because the same resampling decisions made at encode time by the sender can be exactly recovered by the receiver. The encoding process of the sender produces the SMC state in the forward direction of time.
{+The receiver also reconstructs the the SMC state in the forward direction of time. Thus, if the receiver has the same resampling criteria as the sender, they can recover exactly the SMC forward pass of the sender. Thus, we can make two major modiﬁcations to incorporate adaptive resampling to BB-SMC. First is that the parent indices A it − 1 are only decoded (and then encoded back) when the resampling criteria is satisﬁed at each timestep t , otherwise each particle inherits itself, i.e., set A it − 
+}
= i .
Second is that when resampling is not performed, the importance weights need to be accumulated and used for resampling next time. After resampling, the accumulated importance weights are reset to uniform for all particles. BB-CSMC coder As with BB-IS, BB-SMC also suffers from a increased initial bits cost equal to − log Q ( Z T , A T − 1 , j ) that scales like O ( NT ) , in contrast to the O ( T ) initial bit cost of BB-ELBO.
Procedure Encode( symbol x T , message m ) for t = 1 , . .
, T do if t (cid:54) =
1 then decode v t − 1 with unif { 0 , · · · , 2 r − 1 } assign p t − 1 =
Cat( ˜ w t − 1 ) for i = 1 , . .
. , N do assign v it −
= R it − 1 ( v t − 1 ) assign A it −
= F − 1 p t − 1 ( v it − 1 ) assign A t − 1 = (cid:2)
A t − 2 , A 1: N t − 1 (cid:3) decode u t with unif { 0 , · · · , 2 r − 1 } for i = 1 , . .
. , N do assign τ it − 1 =
TraceBack
( Z t − 1 , A t − 2 ,
A it − 1 ) assign u it =
T it ( u t ) , q it =
q ( z it | x t , τ it − 1 ) assign z
it
=
F − 1 q it ( u it ) assign Z t = (cid:2)
Z t − 1 , z 1: N t (cid:3) decode j with Cat( ˜ w T ) set the ancestral lineage B T = j and B t =
A B t +1 t for t =
T − 1 , . .
, 1 for t =
T, . .
, 1 do encode u B t t with unif { u : F − 1 q Btt ( u ) =
z B t t } encode x t with g ( x t | x t − 1 , ( τ B t t − 1 , z B t t )) encode z B t t with f ( z B t t | x t − 1 , τ B t t − 1 ) if t (cid:54)
=
1 then encode v B t t − 1 with unif { v : F − 1 p t − 1 ( v ) =
A B t t − 1 } encode B t with Cat(1 /N )
return m (cid:48)
[-First is that the parent indices A it − 1 are only decoded (and then encoded back) when the resampling criteria is satisﬁed at each timestep t , otherwise each particle inherits itself, i.e., set A it − 
-]
Similarly, we can derive a coupled variant of BB-SMC which is called Bits-Back Coupled Sequential
Monte Carlo (BB-CSMC).
BB-CSMC reparameterizes the particle states z it as deterministic functions of uniform random variables u it which are coupled by a common uniform u t .
At each timestep t , instead of directly decoding z it with their approximate posterior q it , BB-CSMC ﬁrst decodes u t and then obtain u it with bijective functions T it , i.e., u it =
T it ( u t ) .
Then the particle states z it are obtained as z it
=
F − 1 q it ( u it ) , where the functions F − 1 q i t are deﬁned as equation ( ?? ).
This is not enough to sufﬁciently reduce the O ( NT ) initial bit cost since the parent indices A it − 1 are also decoded (and thus require some initial bits) for each particle.
Therefore, similarly for A it − 1 , we also need to introduce the common uniform v t − 1 and the bijective functions R it − 1 to obtain v it − 1 =
R it − 1 ( v t − 1 ) .
A it − 1 are obtained as A it − 1 =
F − 1 p t − 1 ( v it − 1 ) where p t − 1 is the categorical distribution deﬁned by normalized importance weights ˜ w t − 1 .
This reduces the initial bit cost to (2 T − 1) r − log( ˜ w T ) .
Note that r is lower bounded by log N because 2 r should be larger than N , the initial bit cost roughly scales like O ( T log N ) .
The encoding process of BB-CSMC should also be calibrated to match the modiﬁed decoding process, as with BB-CIS.
In particular, after decoding the special particle index j , BB-CSMC only encodes ( x t , B t , v B t t − 1 , u B t t , z B t t ) associated with the special particle trajectory.
The uniform v B t t − 1 and u B t t are encoded using uniform distributions over restricted sets mapped from A B t t − 1 =
B t − 1 and z B t t , respectively.
One can easily compute the net bitrate of BB-CSMC is:
[Equation]
which is comparable to BB-SMC but uses only O ( T log N ) initial bits.
C C OMPUTATIONAL C OST
The computational requirements of our McBits coders can be distilled into two components: calculation of the distributions used (which usually involves neural network computation) and the actual encoding/decoding operations themselves.
In our implementations of McBits the neural network computations tend to dominate.
For most McBits coders (e.g., BB-IS and BB-AIS), the computational complexity scales like O ( N ) , where N is the number of particles or AIS steps.
In BB-IS for example, the neural network for the approximate posterior needs to be computed once for each encoded symbol, while the neural net for the conditional likelihood is executed N times, once for each particle, in order to compute the weights for the categorical sampler.
In BB-SMC, we are required to decode/encode N times from a categorical distribution of alphabet size N at each timestep.
Naively implemented, this has computational complexity O ( TN 2 ) (in contrast to the O ( T ) computational complexity of BB-ELBO), but we believe that this can be reduced to O ( TN ) by applying the alias method (Walker, 1974).
We have not implemented the alias method, because the computational cost of decoding/encoding ancestral indices only accounted for a small fraction of the overall cost in our experiments.
The neural net computation and the encoding/decoding operations in BB-IS and BB-SMC can also be parallelized over particles, which in an ideal implementation would result in O (1) and O ( T ) time for the two methods, respectively.
However, BB-AIS is not easily parallelizable due to its sequential nature.
To test parallel performance we have written an end-to-end parallelized version of BB-IS, using the JAX framework (Bradbury et al.), results are shown in Figure 11. [Figure] well with N (number of particles) in a parallel implementation, for the BMNIST test set at batch size 1.
This demonstrates the potential practicality of the method, in some settings, with hundreds of particles. Note, the N = 1 case reduces to BB-ELBO, a simpler program which runs faster.
D E XPERIMENTAL D ETAILS D. L OSSLESS C OMPRESSION ON S YNTHETIC D ATA
We assessed the convergence properties, impact of dirty bits, and initial bit cost of our McBits coders on synthetic data.
We sometimes evaluated the ideal bitrate , which for each coder is the corresponding estimated variational bound.
Experimental details and results are presented below.
Mixture Model We used a mixture model with 1-dimensional observation and latent variables.
The alphabet sizes of the observation and the latent were 64 and 256 respectively.
The data generating distribution (prior and conditional likelihood distributions) counts were i.i.d.
sampled random integers from the range [1 , 20] and then normalized.
All the coders got access to the true data generating distribution of the model and used a uniform approximate posterior distribution.
All coders were evaluated using a message consisting of 5000 symbols i.i.d.
sampled from the model to compute the total and net bitrates.
The ideal bitrate of each coder was computed by its corresponding (empirical) negative variational bound by resampling the N particle system 100 times and averaging (except for the ELBO bound, which can be computed exactly).
Hidden Markov Model We used a hidden Markov model (HMM) with 1-dimensional observation and latent variables.
The alphabet sizes of the observation and the latent were 16 and 32 respectively.
The number of timesteps was set to 10.
The data generating distribution (prior/transition/emission distributions) counts were i.i.d.
sampled random integers from the range [1 , 20] and then normalized.
All the coders got access to the true data generating distribution of the model and used a uniform approximate posterior distribution.
All coders were evaluated using a message consisting of 5000 sequences i.i.d.
sampled from the model to compute the total and net bitrates.
The ideal bitrate of each coder was computed by its corresponding (empirical) negative variational bound of the sampled message, as with the mixture model.
The entropy of the model was also computed empirically by the negative marginal likelihood of the sampled message which was computed by the forward algorithm.
Additional Results We measured the impact of dirty bits by plotting ideal versus net bitrates on the toy mixture and the toy HMM model in Fig.
The deviation of any point to the dashed line indicates the severity of dirty bits.
Interestingly, most of our McBits coders appeared to ‘clean’ the bitstream as N increased, i.e. the net bitrate converged to the entropy, as shown in Fig. 12a.
Only BB-AIS-BitSwap did not clean the bitstream (Fig. 12a), indicating that the order of operations has a signiﬁcant impact on the cleanliness of McBits coders.
[Figure]
expect
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure] [Figure]
[Figure]
D. L OSSLESS C OMPRESSION ON I MAGES
We benchmarked the performance of BB-IS and BB-CIS on the standard train-test splits of two datasets: an alphanumeric extension of MNIST called EMNIST (Cohen et al., 2017), and CIFAR-10 (Krizhevsky, 2009).
Experimental details and results are presented below.
Datasets We used two datasets for benchmarking lossless image compression: EMNIST (Cohen et al., 2017) and CIFAR-10.
EMNIST dataset extends the MNIST dataset to handwritten digits.
There aredifferent splits provided in the dataset and we used two of them in our experiments: MNIST and Letters.
The EMNIST-MNIST split mimics the original MNIST dataset which contains 60,000 training and 10, test examples.
The EMNIST-Letters split contains 124800 training and 20800 test examples.
Both two splits were dynamically binarized following Salakhutdinov & Murray (2008).
Speciﬁcally, the observations were randomly sampled from the Bernoulli distribution with expectations equal to the real pixel values.
For the CIFAR-10 dataset, no additional preprocessing was applied.
Model For the EMNIST datasets, we used the VAE model with 1 stochastic layer as in Burda et al.
The VAE model had 50 latents with a standard Gaussian prior and factorized Gaussian approximate posterior.
The conditional likelihood distribution was modeled by the Bernoulli distribution which ﬁt with the binarized observations.
The training procudure was the same as Burda et al.
For the CIFAR10 dataset, we used the VQ-VAE model (Oord et al., 2017) with discrete latent variables.
Categorical distributions were used for both the latents and the observations.
We followed the experimental setup in Sønderby et al.
and used the VQ-VAE model with 8 latent vairables per spatial dimension.
The VQ-VAE model was trained with Gumbel-Softmax (Jang et al., 2016;
Maddison et al., 2016) relaxation with a tempreture of 0.5 and a minibatch size of 32.
The ADAM optimizer was used for training the model and the learning rate was tuned over { 1 × 10 − 4 , 5 × 10 − 4 , 1 × 10 − 3 } .
Discretization To perform bits-back coding with continuous latent variables, we need to discretize the latent space and approximate the continuous prior and the approximate posterior with their discretized variants using the same set of bins.
In MacKay (2003) and Townsend et al.
, they showed that the continuous latents could be discretized up to an arbitrary precision δ z without affecting the net bitrate as we could also get those extra bits due to discretization “back”.
The notable effect is for the total bitrate since the initial bit cost would scale with the − log δ z , which means that we prefer to use a reasonably small precision in practice.
In our EMNIST experiments, continuous latent variables were used for the VAE model and need to be discretized for compression.
We used the maximum entropy discretization in Townsend et al.
and discretized the latent space into bins that had equal mass under the standard Gaussian prior for all the coders.
Baselines We included amortized-iterative inference method (Yang et al., 2020) which also improved the compression rate by bridging the gap of the ELBO bound and the marginal likelihood as a baseline in our experiments.
When compressing a data example, it ﬁrst initializes the local variational parameters (e.g., the mean and variance for factorized Gaussian approximate posterior in our experiments) from the trained VAE inference model.
Then it optimizes the ELBO objective over the local variational parameters with stochastic gradient decent and uses the optimized variational parameters for compression.
This method introduces expensive computation at the compression stage.
In our experiments, we kept the number of optimization steps equal to the number of particles N to roughly match the computation budget in terms of the number of quries to the VAE model.
However, even so the computation cost of this method is more expensive than our BB-IS/BB-CIS coders because 1) BB-IS/BB-CIS coders with N particles introduce N queries to the VAE generative model and 1 query to the VAE inference model, while amortized-iterative inference with N optimization steps requires N quiries to the whole VAE model and N back-propagation steps; 2) the computation of the BB-IS/BB-CIS coders can be potentially parallelized over the particles, but the computation of amortized-iterative inference needs to be done in N sequential optimization steps.
In our experiments, we used ADAM optimizer for optimizing the local variational parameters and the learning rate was tuned in the range { 5 × 10 − 4 , 1 × 10 − 3 , 5 × 10 − 3 , 1 × 10 − 2 , 5 × 10 − 2 } .
Additional Results We found that the variational bounds used to train the VAE models had an impact on compression performance.
When using BB-IS with a model trained on the IWAE objective, equalizing the number of particles during compression and training resulted in better rates than BB-IS with an ELBO-trained VAE (Table 2).
Therefore, we always use our McBits coders with models trained on the corresponding variational bound.
[Table] To study the potential trade-off of applying our BB-CIS coder, we compared the compression performance of BB-IS and BB-CIS with 50 particles on EMNIST datasets in Table 3.
Although the BB-CIS coder adopts a near-random sampling strategy, its ideal bitrate and net bitrate match those of the BB-IS coder.
And it is effective for reducing the initial bit cost, as the gap between the total bitrate and the net bitrate is much smaller than that of the BB-IS coder.
These observations illustrate that we can use BB-CIS as a drop-in replacement for BB-IS nearly for free.
[Table]
Finally, we compared BB-IS and BB-CIS to other benchmark lossless compression schemes by measuring the total bitrates on EMNIST test sets (without transferring).
We also compared with amortized-iterative inference, (Yang et al., 2020) that optimizes the ELBO objective over local variational parameters for each data example at the compression stage.
To roughly match the computation budget, the number of optimization steps was set to 50 and this method is denoted as BB-ELBO-IF (50).
Both BB-IS and BB- outperformed all other baselines on both test sets, and BB-CIS was better than BB-IS in terms of total bitrate since it effectively reduces the initial bit cost.
[Table]
As the amortized-iterative inference method is orthogonal to our method, the compression performance of our McBits coder (in particular BB-IS in this experiment) can be further improved by applying amortizediterative inference at the compression stage.
Similar to the original amortized-iterative inference, one can initialize the local variational parameters from the trained VAE inference model and optimize the IWAE objective with N particles (as opposed to the ELBO objective) over them.
In Table 5, we compared the net bitrate of combining the BB-IS coder and the amortized-iterative inference method.
We observed that the amortized-iterative inference method could further boost the compression rate of our BB-IS coder. ) are shown in
[Table]
D. L OSSLESS C OMPRESSION ON S EQUENTIAL D ATA
We quantiﬁed the performance of BB-SMC on sequential data compression tasks with 4 polyphonic music datasets: Nottingham, JSB, MuseData, and Piano-midi.de (Boulanger-Lewandowski et al., 2012).
Experimental details and results are presented below.
Datasets We used 4 polyphonic music datasets to evaluate the compression performance of the BB-SMC coder on sequential datasets: Nottingham folk tunes, the JSB chorales, the MuseData library of classical piano and orchestral music, and the Piano-midi.de MIDI archive (Boulanger-Lewandowski et al., 2012).
All datasets were composed of sequences of binary 88-dimensional vectors representing active musical notes at one timestep.
For all datasets, we imitated the experimental setup presented in Maddison et al.
We used the same train/validation/test split and echoed their data preprocessing.
The only difference was that we used a chuncked version of these datasets where each sequence was chunked to sub-sequences with maximum length of 100.
We found it slightly improved the model performance and signiﬁcantly reduced the initial bit cost (as the original maximum sequence length was very large).
Models All models were based on the variational RNN architecture (Chung et al., 2015).
All distributions over latent variables were factorized Gaussians, and the output distributions were factorized Bernoullis for binary observations on 4 polyphonic music datasets.
JSB models were trained with 32 hidden units, Muse-data with 256, and all other models with 64 units.
For each aforementioned dataset, there was one model trained with the ELBO, IWAE, and FIVO objectives, respectively.
All models were trained withparticles, a batch size of 4 and the Adam optimizer with learning rate 3 × 10 − 5 .
All models were initially evaluated on the validation set, which allowed for early stopping.
In Table 6, we present our models’ performance in nats and bits to allow for easy comparison of generative modelling and compression literature.
we include the estimated data log-likelihood as model evaluation metric. We estimated the log-likelihood by computing the maximum of the ELBO, IWAE and FIVO bound with 128 particles. We include this metric for better comparison to other work. However for this work, this bound is not relevant.
Relevant metrics include the respective bounds in nats or bits per time step
[Table]
Results We compared the net bitrates of all coders for compressing each test set in Table 7.
BB-SMC clearly outperformed BB-ELBO and BB-IS with the same number of particles on all datasets.
[Table]
Then we compared our coders with benchmark lossless compression schemes in Table 8.
Our coders were comparable with those baselines and the BB-SMC coder outperformed all the baselines on the JSB dataset.
Note that the compression performance of our coders is bottlenecked by the simple VRNN architecture that we used in our experiments.
And we suppose that with more powerful and better trained VRNN models, our coders could outperform those benchmark schemes.
However, our main focus is to compare the compression performance of our BB-SMC coder and the BB-ELBO/BB-IS coders with the same VRNN architectures to show the effectiveness of BB-SMC for compressing sequential data, and this is clearly illustrated by the results.
[Table]
Discussion of Initial Bit Cost The initial bit cost of compressing sequential data using (Monte Carlo) bits-back algorithms scales linearly with both the sequence length and the number of particles, i.e., O ( NT ) .
The original four polyphonic music datasets have a special characteristic that the average and maximum sequence lengths are large but the number of sequences is small, which means that the initial bit cost is huge but cannot be sufﬁciently amortized.
Thus, if we compress the original datasets without chunking sequences, the total bitrate will be much larger than the net bitrate.
For example, there are onlysequences in the Musedata test set but the average length is 519 and the maximum length is 4273.
As a result, the total bitrates of the BB-ELBO coder and the BB-SMC coders for compressing the original dataset are 136.81 and 544.40 bits/timestep respectively and much larger than their net ones.
Therefore, we chose to chunk long sequences to short ones of a predeﬁned maximum length (100) and compress them independently, which could effectively decrease and amortize the initial bit cost.
When compressing the chunked dataset, the total bitrates of BB-ELBO and BB-SMC reduce to 12.83 and 21.39 bits/timestep respectively.
As for the initial bit cost caused by the particles, we can also use the coupled variant of BB-SMC (aka BB-CSMC) for compressing sequential data.
D. L OSSY C OMPRESSION ON I MAGES
Current state-of-the-art lossy image compressors use hierarchical latent VAEs with quantized latents that are losslessly compressed with hyperlatents (Ball´e et al., 2016; 2018; Minnen et al., 2018).
Yang et al.
(2020) observed that the marginalization gap of jointly compressing the latent and the hyperlatent can be bridged by bits-back.
Thus, our McBits coders can be used to further reduce the gap.
We experimented on a simpliﬁed setting where we used the binarized EMNIST datasets and a modiﬁcation of the VAE model with 2 stochastic layers in Burda et al. (2015).
Experimental details and results are presented below.
Lossy Compression Setup We used the binarized EMNIST datasets to benchmark the lossy compression performance.
We considered the lossy compression setup with hierarchical VAE models (Ball´e et al., 2018; Minnen et al., 2018).
Speciﬁcally, the compressing data x is transformed by trained hierarchical inference models f l and f h with parameters φ l and φ h to produce discretized latent y and hyperlatent z as y = (cid:98) µ fy (cid:101) and z = (cid:98) µ fz (cid:101) , where µ fy = f l ( x ; φ l ) and µ fz = f l ( y ; φ h ) are their continuous representations.
In our experiments, the latent is rounded to the nearest integer and the hyperlatent is discretized by the maximum entropy discretization scheme introduced in D.2 for lossless compression.
Then the latent y and the hyperlatent z are compressed with bits-back coding as in lossless compression.
On the decoder side, both y and z can be losslessly recovered and the reconstructed data ˆ x is transformed from y using the generative models g l and g h with parameters θ l and θ h .
Model We used the VAE model with 2 stochastic layers in Burda et al.
(2015) with several modiﬁcations based on Ball´e et al.
(2018) for adapting to the lossy compression setup.
Speciﬁcally, the approximate posterior distribution of the latent was a uniform distribution centered at µ y , i.e., q ( y | x ) = unif( µ fy − 12 , µ fy + 12 ) , which is a differentiable substitute for rounding during training.
The conditional likelihood distribution of the latent should also support quantization which was a factorized Gaussian distribution convolved with a standard uniform p ( y | z ) =
N ( µ gy , σ gy 2 ) ∗ unif( − 12 , 12 ) , where ( µ gy , σ gy )
=
g h
( z ; θ h ) .
The convolved distribution agrees with the discretized distribution on all integers (Ball´e et al., 2016; 2018).
This is important because the discretization of the latent would affect the compression rate which should be taken into account during training.
In contrast, the disretization of the hyperlatent does not affect the compression rate as a result of getting bits back (see the discussion in D.2), we kept the distribution over the hyperlatent unchanged as in Burda et al.
Speciﬁcally, its approximate posterior distribution q ( z | y ) was a factorized Gaussian distribution centered at µ fz and its prior distribution p ( z ) was a standard Gaussian.
The conditional likelihood distribution was kept as a factorized Bernoulli for binary observations.
Training The loss function of training the hierarchical VAE model is a relaxed rate-distortion objective:
[Equation]
The distortion is measured as the negative log likelihood of the Bernoulli observations and the rate is measure as the negative ELBO marginalized over the hyperlatent.
λ is the hyperparameter that controls the rate-distortion trade-off.
The last term is measured as 0 since q ( y | x ) is a uniform distribution.
The above rate-distortion objective is very similar to the objective function in β -VAE (Higgins et al., 2016) and can be optimized with the reparameterization trick.
Note that the ELBO rate term can be changed to the IWAE objective with multiple particles.
In our experiments, we trained the models on the binarized EMNIST-MNIST dataset and evaluated on both EMNIST-MNIST and EMNIST-Letters test sets (for evaluating the performance in the transfer setting).
we trained the model with the above loss function using the IWAE objective with
M ∈ { 1 , 5 , 50 } particles as the rate term.
For each setup, we trained models with different λ values in the range { 1 . 0 , 1 . 5 , 2 . 0 , 2 . 5 , 3 . 0 , 4 . 0 , 5 . 0 , 6 . 0 , 7 . 0 , 7 . 5 , 8 . 0 , 9 . 0 , 10 . 0 , 12 . 5 , 15 . 0 , 17 . 5 , 20 . 0 } .
For training each model, we tuned the learning rate in the range { 5 × 10 − 3 , 2 . 5 × 10 − 3 , 1 × 10 − 3 , 7 . [-
 ×−-]5 × 10 − 4 , 5 × 10 − 4 } .
Results We evaluated net bitrate savings of BB-IS compared to BB-ELBO on the EMNIST-MNIST test set with different λ values, as in Fig. 13b.
We found that BB-IS achieved more than 15% rate savings in some setups, see also the rate-distortion curves in Fig. 13a.
We also implemented these experiments for the model in Ball´e et al.
(2018), but did not observe signiﬁcant improvements.
This may be due to the speciﬁc and complex model architecture.
We also evaluated the performance of applying amortized -iterative inference (Yang et al., 2020) in lossy compression in Fig. 13c & 13d.
The main purposes were to: 1) compare the amortized-iterative inference and our BB-IS coder with similar computation budget; 2) illustrate the potential of combining amortizediterative inference with our BB-IS coder.
Speciﬁcally, we used 50 optimization steps for amortized-iterative inference to roughly match the computation budget (see discussion in D.2) with our BB-IS coder with 50 particles (denoted as BB-IS (50)).
We used a 2-stage amortized-iterative inference similar to Yang et al.
and each stage contained 25 optimization steps.
In the ﬁrst stage, both the local variational parameters of the latent and the hyperlatent were optimized with the rate-distortion objective.
In the second stage, the local variational parameters of the latent were ﬁxed while those of the hyperlatent were optimized with the rate term (i.e., negative ELBO) of the rate-distortion objective.
The learning rate was tuned in the range { 1 × 10 − 2 , 5 × 10 − 3 , 1 × 10 − 3 } .
This method is denoted as BB-ELBO-IF (50).
We observed that BB-ELBO-IF (50) improved over BB-ELBO but underperformed our BB-IS (50) in terms of rate-distortion trade-off.
We also combined BB-IS (50) with 50 optimization steps of amortized-iterative inference using negative IWAE as the rate term (denoted as BB-IS (50)-IF (50)), which we found to further improve the performance and achieved more than 20% rate savings in some setups.
We also evaluated the lossy compression performance in a transfer setting where we used models trained on the EMNIST-MNIST to compress EMNIST-Letters test set, as in Fig.
We observed that the improvement of BB-IS was not as signiﬁcant as on the EMNIST-MNIST test set.
This might be due to that although BB-IS could improve the rate term over BB-ELBO more on the EMNIST-Letters test set (as shown in Table 1 for lossless compression), the distortion term of the model trained with IWAE might not generalize well on the dataset of a slightly different distribution. [Figure] performance.
We measure the rate savings ( % ) relative to BB-ELBO for ﬁxed distortion values.
[Figure]