Original: L OSSLESS C OMPRESSION USING C ONTINUOUSLYI NDEXED N ORMALIZING F LOWS
Revised: L OSSLESS C OMPRESSION USING C ONTINUOUSLYI NDEXED N ORMALIZING F LOWS

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Recently, a class of deep generative models known as continuously-indexed ﬂows (CIFs) have expanding the modelling capacity of normalizing ﬂows (NFs) in the context of both density estimation and variational inference.
Revised: Recently, a class of deep generative models known as continuously-indexed ﬂows (CIFs) has expanded the modelling capacity of normalizing ﬂows (NFs) in the context of both density estimation and variational inference.

Original: CIFs are provably more general and expressive than NFs, but do not induce a closed-form density model and thus require additional considerations when applying in the same contexts that NFs have shown promise.
Revised: CIFs are provably more general and expressive than NFs, but do not induce a closed-form density model and thus require additional considerations when applying in the same contexts that NFs have shown promise.

Original: One such area is lossless compression, where NFs have been used as the density model to develop a compression scheme, known as local bits-back, with expected codelength approximately equal to the average negative log-likelihood of the NF density model.
Revised: One such area is lossless compression, where NFs have been used as the density model to develop a compression scheme, known as local bits-back, with expected codelength approximately equal to the average negative log-likelihood of the NF density model.

Original: Here, we propose to extend the local bits-back scheme to CIF-based density models, as the improved expressiveness inherent in CIFs stands to reduce the expected codelength of compressed data.
Revised: Here, we propose to extend the local bits-back scheme to CIF-based density models, as the improved expressiveness inherent in CIFs stands to reduce the expected codelength of compressed data.

Original: We also leverage recent work on compression schemes built with hierarchical variational auto-encoders – as hierarchical CIFs can themselves be seen as interpolating between these and NFs – gaining further expressiveness in our density models and effectiveness in our compression scheme.
Revised: We also leverage recent work on compression schemes built with hierarchical variational auto-encoders – as hierarchical CIFs can themselves be seen as interpolating between these and NFs – gaining further expressiveness in our density models and effectiveness in our compression scheme.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: As datasets become larger, so too do the challenges in both modelling the data-generating process and providing efﬁcient compression schemes.
Revised: As datasets become larger, so too do the challenges in both modelling the data-generating process and providing efﬁcient compression schemes.

Original: These two problems are linked through the framework of entropy encoding, which describes a family of lossless compression techniques making use of a model of the data probability distribution.
Revised: These two problems are linked through the framework of entropy encoding, which describes a family of lossless compression techniques making use of a model of the data probability distribution.

Original: A combination of recent advancements on both fronts have led to the development of practical lossless compression algorithms for high-dimensional data.
Revised: A combination of recent advancements on both fronts have led to the development of practical lossless compression algorithms for high-dimensional data.

Original: Perhaps most notably, Townsend et al.
Revised: Perhaps most notably, Townsend et al.

Original: (2019) has demonstrated how the entropy encoding scheme known as asymmetric numeral systems (ANS) (Duda, 2009) can be integrated with bits-back coding (Frey & Hinton, 1996) to build a powerful compression method backed by variational auto-encoders (VAEs) (Kingma & Welling, 2014).
Revised: (2019) has demonstrated how the entropy encoding scheme known as asymmetric numeral systems (ANS) (Duda, 2009) can be integrated into bits-back coding (Frey & Hinton, 1996) to build a powerful compression method backed by variational auto-encoders (VAEs) (Kingma & Welling, 2014).

Original: Other approaches also combining ANS with bits-back coding, but making use of other deep generative models such as hierarchical VAEs (Kingma et al., 2019) and normalizing ﬂows (NFs)
Revised: Other approaches also combining ANS with bits-back coding, but making use of other deep generative models such as hierarchical VAEs (Kingma et al., 2019) and normalizing ﬂows (NFs)

Original: (Ho et al., 2019), have emerged and shown promising empirical results, achieving codelengths close to theoretical optimum.
Revised: (Ho et al., 2019), have emerged and shown promising empirical results, achieving codelengths close to theoretical optimum.

Original: In all cases, the expected codelength of the compression scheme is negatively correlated with the log-likelihood (or evidence lower bound) of the generative model, implying that a more expressive generative model will correspond to a better algorithm for lossless compression.
Revised: In all cases, the expected codelength of the compression scheme is negatively correlated with the log-likelihood (or evidence lower bound) of the generative model, implying that a more expressive generative model will correspond to a better algorithm for lossless compression.

Original: One way to directly and provably improve the expressiveness of an NF-based model is to use continuously-indexed ﬂows (CIFs) (Cornish et al., 2020).
Revised: One way to directly and provably improve the expressiveness of an NF-based model is to use continuously-indexed ﬂows (CIFs) (Cornish et al., 2020).

Original: This recently-proposed framework augments the generative process of any standard NF with additional indexing variables, relaxing the restrictive bijectivity constraint imposed on the baseline ﬂow.
Revised: This recently-proposed framework augments the generative process of any standard NF with additional indexing variables, relaxing the restrictive bijectivity constraint imposed on the baseline ﬂow.

Original: The cost of increased modelling capacity is a now intractable density model, but this improved expressiveness has been able to overcome any obstacles stemming from this intractability in the context of both density estimation (Cornish et al., 2020) and variational inference (Caterini et al., 2020).
Revised: The cost of increased modelling capacity is a now intractable density model, but this improved expressiveness has been able to overcome any obstacles stemming from this intractability in the context of both density estimation (Cornish et al., 2020) and variational inference (Caterini et al., 2020).

Original: In this work we propose an efﬁcient lossless compression scheme for using hierarchical CIFs as the probability model for bits-back coding.
Revised: In this work we propose an efﬁcient lossless compression scheme for using hierarchical CIFs as the probability model for bits-back coding.

Original: The scheme we propose builds on the ideas of BB-ANS (Townsend et al., 2019) and local bits-back coding (Ho et al., 2019) to allow for compression using a single layer CIF model, as well as Bit-Swap (Kingma et al., 2019) that allows us to decrease the number of auxiliary bits required to utilize a hierarchical CIF model which signiﬁcantly decreases the codelengths in practice, especially for short data sequences.
Revised: The scheme we propose builds on the ideas of BB-ANS (Townsend et al., 2019) and local bits-back coding (Ho et al., 2019) to allow for compression using a single layer CIF model, as well as Bit-Swap (Kingma et al., 2019) that allows us to decrease the number of auxiliary bits required to utilize a hierarchical CIF model which signiﬁcantly decreases the codelengths in practice, especially for short data sequences.

Original:  B ACKGROUND 2. C ONTINUOUSLY -I NDEXED F LOWS
Revised:  B ACKGROUND 2. C ONTINUOUSLY -I NDEXED F LOWS

Original: CIFs (Cornish et al., 2020) propose to model continuous data deﬁned over some space X as the
Revised: CIFs (Cornish et al., 2020) propose to model continuous data deﬁned over some space X as the

Original: X -marginal of
Revised: X -marginal of

Original: [Equation]
Revised: [Equation]

Original: where p Z is a density deﬁned over some latent space Z , p U | Z is a parametrized conditional density deﬁned over an indexing space U , and F : Z × U
Revised: where p Z is a density deﬁned over some latent space Z , p U | Z is a parametrized conditional density deﬁned over an indexing space U , and F : Z × U

Original: → X is a function such that F ( · ; u ) is a bijection for each u ∈ U .
Revised: → X is a function such that F ( · ; u ) is a bijection for each u ∈ U .

Original: If we parametrize F such that, e.g. F ( · ; 0)
Revised: If we parametrize F such that, e.g. F ( · ; 0)

Original: =
Revised: =

Original: f ( · ) for a standard normalizing ﬂow transformation f : Z → X , then it is easy to see that equation 1 strictly generalizes the typical normalizing ﬂow density model Z ∼ p
Revised: f ( · ) for a standard normalizing ﬂow transformation f : Z → X , then it is easy to see that equation 1 strictly generalizes the typical normalizing ﬂow density model Z ∼ p

Original: Z , X
Revised: Z , X

Original: = f ( Z ) .
Revised: = f ( Z ) .

Original: The cost for this improvement is an intractable density model p X ( x )
Revised: The cost for this improvement is an intractable density model p X ( x )

Original: =
Revised: =

Original: (cid:82) p X,U ( x , u ) d u , although tractable p X,U emerges from equation 1.
Revised: (cid:82) p X,U ( x , u ) d u , although tractable p X,U emerges from equation 1.

Original: We can also stack the ﬁnal two steps of equation 1 to gain further expressiveness as per Cornish et al.
Revised: We can also stack the ﬁnal two steps of equation 1 to gain further expressiveness as per Cornish et al.

Original: Speciﬁcally, model X now as the Z L -marginal of the following L -layered model:
Revised: Speciﬁcally, model X now as the Z L -marginal of the following L -layered model:

Original: [Equation]
Revised: [Equation]

Original: where the ﬁnal two steps are repeated for (cid:96) ∈ { 1 , . .
Revised: where the ﬁnal two steps are repeated for (cid:96) ∈ { 1 , . .

Original: , L } .
Revised: , L } .

Original: We can write the joint density over ( X, U 1: L ) recursively, but p X remains intractable, cf.
Revised: We can write the joint density over ( X, U 1: L ) recursively, but p X remains intractable, cf.

Original: equation 12 of Cornish et al.
Revised: equation 12 of Cornish et al.

Original: However, we can learn the parameters of p X using variational inference, constructing an inference model
Revised: However, we can learn the parameters of p X using variational inference, constructing an inference model

Original: [Equation]
Revised: [Equation]

Original: where z L := x , z (cid:96) :
Revised: where z L := x , z [Equation]

Original: = F − 1 (cid:96) +

Revised: := F [Equation] 1 (cid:96) +

Original: ( z (cid:96) +1 ; u (cid:96) +1 ) recursively for (cid:96) ∈ { 1 , . . .
Revised: ( z (cid:96) +1 ; u (cid:96) +1 ) recursively for (cid:96) ∈ { 1 , . . .

Original: , L } , and each q U (cid:96) | Z (cid:96) is a parametrized conditional density.
Revised: , L } , and each q U (cid:96) | Z (cid:96) is a parametrized conditional density.

Original: We can then train a CIF density model by maximizing the evidence lower bound (ELBO) objective, given for a single point x ∈ X as:
Revised: We can then train a CIF density model by maximizing the evidence lower bound (ELBO) objective, given for a single point x ∈ X as:

Original: [Equation]
Revised: [Equation]

Original: It is important to note that the factorization of the true posterior over indexing variables p U 1: L | X matches that of equation 3 (Cornish et al., 2020, Appendix B.6), which helps to increase the value of
Revised: It is important to note that the factorization of the true posterior over indexing variables p U 1: L | X matches that of equation 3 (Cornish et al., 2020, Appendix B.6), which helps to increase the value of

Original: L ( x ) as equation 4 is maximized in q when q U 1: L | X
Revised: L ( x ) as equation 4 is maximized in q when q U 1: L | X

Original: = p U 1: L | X .
Revised: = p U 1: L | X .

Original: 2. B ITS - BACK CODING
Revised: 2. B ITS - BACK CODING

Original: Bits-back coding (Frey & Hinton, 1996) is a method that allows for lossless compression using density models with latent variables, which Townsend et al.
Revised: Bits-back coding (Frey & Hinton, 1996) is a method that allows for lossless compression using density models with latent variables, which Townsend et al.

Original: (2019) recently combined with asymmetric numeral systems (ANS)
Revised: (2019) recently combined with asymmetric numeral systems (ANS)

Original: (Duda, 2009) to devise a practical compression scheme based on variational auto-encoders (VAEs).
Revised: (Duda, 2009) to devise a practical compression scheme based on variational auto-encoders (VAEs).

Original: Speciﬁcally, suppose we have access to a VAE with generative model p X,Z ( x , z ) :
Revised: Speciﬁcally, suppose we have access to a VAE with generative model p X,Z ( x , z ) :

Original: = p X | Z ( x | z ) · p Z ( z ) and encoder q Z | X ( z | x ) trained over a dataset D :=
Revised: = p X | Z ( x | z ) · p Z ( z ) and encoder q Z | X ( z | x ) trained over a dataset D :=

Original: { x i } i .
Revised: { x i } i .

Original: Given a stack-like entropy encoding scheme such as ANS, we can iteratively encode each point x i in the dataset onto a pre-existing bit stack – here denoted m – by performing the following steps: (i) decode z i from m using q Z | X ( · | x i ) ; (ii) encode x i onto m using p X | Z ( · | z i ) ; (iii) encode z i onto m using p Z .
Revised: Given a stack-like entropy encoding scheme such as ANS, we can iteratively encode each point x i in the dataset onto a pre-existing bit stack – here denoted m – by performing the following steps: (i) decode z i from m using q Z | X ( · | x i ) ; (ii) encode x i onto m using p X | Z ( · | z i ) ; (iii) encode z i onto m using p Z .

Original: We can reverse this procedure to then decode the dataset from m once encoding is complete.
Revised: We can reverse this procedure to then decode the dataset from m once encoding is complete.

Original: The expected codelength for a point x is approximately equal to the negative evidence lower bound
Revised: The expected codelength for a point x is approximately equal to the negative evidence lower bound

Original: E z ∼ q Z | X ( ·| x ) [log q Z | X
Revised: E z ∼ q Z | X ( ·| x ) [log q Z | X

Original: ( z | x ) − log p X,Z ( x , z )] , which conveniently also serves as the
Revised: ( z | x ) − log p X,Z ( x , z )] , which conveniently also serves as the

Original: VAE objective function, underscoring the relationship between performant generative models and efﬁcient compression schemes.
Revised: VAE objective function, underscoring the relationship between performant generative models and efﬁcient compression schemes.

Original: 2. L OCAL BITS - BACK CODING
Revised: 2. L OCAL BITS - BACK CODING

Original: Local bits-back coding is a method allowing for lossless compression using normalizing ﬂows (Ho et al., 2019).
Revised: Local bits-back coding is a method allowing for lossless compression using normalizing ﬂows (Ho et al., 2019).

Original: The key idea behind this scheme is reinterpreting z and x , which in the context of standard NFs are related by a bijection x = f ( z ) , as random variables which have instead a fuzzy relationship centred on a sharply peaked normal distribution.
Revised: The key idea behind this scheme is reinterpreting z and x , which in the context of standard NFs are related by a bijection x = f ( z ) , as random variables which have instead a fuzzy relationship centred on a sharply peaked normal distribution.

Original: In the context of CIFs, the ﬂow F is a now a bijection between z and x conditioned on u as noted above.
Revised: In the context of CIFs, the ﬂow F is a now a bijection between z and x conditioned on u as noted above.

Original: Thus, given u ∈ U , the fuzzy relationship between z and x is given by
Revised: Thus, given u ∈ U , the fuzzy relationship between z and x is given by

Original: [Equation]
Revised: [Equation]

Original: where, as per Ho et al.
Revised: where, as per Ho et al.

Original: (2019), σ > 0 is a small scalar parameter and ¯ x are the centres of hypercube bins of volume δ x for the discretization of continuous data x .
Revised: (2019), σ > 0 is a small scalar parameter and ¯ x are the centres of hypercube bins of volume δ x for the discretization of continuous data x .

Original: Furthermore, J ( x , u ) denotes the
Revised: Furthermore, J ( x , u ) denotes the

Original: Jacobian of F − 1 , with respect to its ﬁrst argument (keeping the second ﬁxed), evaluated at ( x , u ) .
Revised: Jacobian of F − 1 , with respect to its ﬁrst argument (keeping the second ﬁxed), evaluated at ( x , u ) .

Original:  M ETHOD
Revised:  M ETHOD

Original: We ﬁrst consider using just a single layer CIF as the density model for bits back coding, and then extend it to thehierarchical CIF casein the second part of this section.
Revised: We ﬁrst consider using just a single layer CIF as the density model for bits back coding, and then extend it to the case of hierarchical CIFs.

Original: Note that throughout this section the notation u is intended as per CIF model notation of Cornish et al.
Revised: Note that throughout this section the notation u is intended as per CIF model notation of Cornish et al.

Original: (2020), rather than the dequantization notation of Ho et al.
Revised: (2020), rather than the dequantization notation of Ho et al.

Original: Weassume we discretize the continuous space X into bins of volume δ x , and we denote the centres of the bins the data was quantized to as ¯ x .
Revised: We discretize the continuous space X into bins of volume δ x , and we denote the centres of the bins the data was quantized to as ¯ x .

Original: Similarly, U and Z are also discretized into bins of volume δ u and δ z , respectively, and discretized values are denoted as ¯ u and ¯ z .
Revised: Similarly, U and Z are also discretized into bins of volume δ u and δ z , respectively, and discretized values are denoted as ¯ u and ¯ z .

Original: NB: As is standard, each distribution that we are decoding/encoding with respect to is deﬁned over continuous data; these can be approximately converted into discrete distributions by taking the previous continuous density value at the respective bin centre and multiplying by the discretization volume, e.g. Q (¯ u 
Revised: NB: As is standard, each distribution that we are decoding/encoding with respect to is deﬁned over continuous data; these can be approximately converted into discrete distributions by taking the previous continuous density value at the respective bin centre and multiplying by the discretization volume, e.g. Q (¯ u | ¯ x ) =

Original: [Equation]
Revised: q (¯ u | ¯ x ) δ u .

Original: The CIF-based lossless coding scheme requires one additional constraint on the CIF architecture: the neural networks parameterizing the distribution p U | Z ( u | z ) must be constrained to be Lipschitz continuous.
Revised: The CIF-based lossless coding scheme requires one additional constraint on the CIF architecture: the neural networks parameterizing the distribution p U | Z ( u | z ) must be constrained to be Lipschitz continuous.

Original: The reasons for that are given in Appendix A.2.
Revised: The reasons for that are given in Appendix A.2.

Original: 3. S INGLE LAYER
Revised: 3. S INGLE LAYER

Original: The encoding procedure is outlined in Algorithm 1 below, while the decoding procedure is presented in Algorithm 3 in Appendix A.1. 
Revised: The encoding procedure is outlined in Algorithm 1 below, while the decoding procedure is presented in Algorithm 3 in Appendix A.1. [Table] auxiliarybitsused

Original: [Table]
Revised: [Table]

Original: 1: Decode ¯ u ∼ q U | X ( · | ¯ x ) δ u 2: Decode ¯ z ∼ ˜ p Z | X,U ( · | ¯ x , ¯ u ) δ z 3: Encode ¯ x using ˜ p 
Revised: ¯ u ∼ q U | X ( · | ¯ x ) δ u 2: Decode ¯ z ∼ ˜ p Z | X,U ( · | ¯ x , ¯ u ) δ z [Table] ¯ x using ˜ p X | Z,U ( · | ¯ z , ¯ u ) δ x

Revised: [Table]

Original: X | Z,U ( · | ¯ z , ¯ u ) δ x 4: Encode ¯ u using p U | Z ( · | ¯ z ) δ u 5: Encode ¯ z using p Z ( · ) δ z
Revised: Encode ¯ u using p U | Z ( · | ¯ z ) δ u [Table] Encode ¯ z using p Z ( · ) δ z

Original: The resulting expected asymptotic codelength of such a scheme, i.e. ignoring the impact of the initial auxiliary random bits used, is approximately equal to the sum of the CIF negative ELBO, a constant depending on the discretization precision, and second order term in σ stemming from the use of Local
Revised: The resulting expected asymptotic codelength of such a scheme, i.e. ignoring the impact of the initial auxiliary random bits used, is approximately equal to the sum of the CIF negative ELBO, a constant depending on the discretization precision, and second order term in σ stemming from the use of Local

Original: Bits-Back coding (derivation in the Appendix A.2):
Revised: Bits-Back coding (derivation in the Appendix A.2):

Original: [Equation]
Revised: [Equation]

Original: The expected number of initial auxiliary bits required for this scheme to operate is
Revised: The expected number of initial auxiliary bits required for this scheme to operate is

Original: [Equation]
Revised: [Equation]

Original: where p D is the true marginal data distribution over x ∈ X .
Revised: where p D is the true marginal data distribution over x ∈ X .

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Revised: [Table]

Original: 3. H IERARCHICAL M ODEL
Revised: 3. H IERARCHICAL M ODEL

Original: Practically
Revised: Practically

Original: [Table]
Revised: [Table]

Original: the only way CIF models can be expressive enough to achieve competitive compression rates is via multi-layer stacking.
Revised: the only way CIF models can be expressive enough to achieve competitive compression rates is via multi-layer stacking.

Original: However, a naïve application of bits-back coding to multiple layers of CIF would lead to a linear growth in the required number of auxiliary random bits with respect to the number of ﬂow layers L .
Revised: However [Table] a naïve application of bits-back coding to multiple layers of CIF would lead to a linear growth in the required number of auxiliary random bits with respect to the number of ﬂow layers L .

Original: We overcome this limitation and reduce the required number of auxiliary bits by applying the Bit-Swap algorithm (Kingma et al., 2019)
Revised: We overcome this limitation and reduce the required number of auxiliary bits by applying the Bit-Swap algorithm (Kingma et al., 2019)

Original: [Table]
Revised: [Table]

Original: as the stacking of CIF layers is much like the hierarchical VAEs for which this was originally designed
Revised: as the stacking of CIF layers is much like the hierarchical VAEs for which this was originally designed

Original: [Table] [Table]
Revised: [Table] [Table]

Original: In the hierarchical case, the resulting expected codelength is analogous to the single layer case:
Revised: In the hierarchical case, the resulting expected codelength is analogous to the single layer case:

Original: [Equation]
Revised: [Equation]

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: Previous works (Townsend et al., 2019; Ho et al., 2019; Kingma et al., 2019) have empirically shown that the codelengths they achieve are close to the theoretical optimum for the generative models they have investigated.
Revised: Previous works (Townsend et al., 2019; Ho et al., 2019; Kingma et al., 2019) have empirically shown that the codelengths they achieve are close to the theoretical optimum for the generative models they have investigated.

Original: Hence we anticipate our compression scheme to achieve codelengths as predicted by the (negative) log-likelihood performance of CIFs (Cornish et al., 2020).
Revised: Hence we anticipate our compression scheme to achieve codelengths as predicted by the (negative) log-likelihood performance of CIFs (Cornish et al., 2020).

Original: The results reported by
Revised: The results reported by

Original: Cornish et al.
Revised: Cornish et al.

Original: (2020) are obtained using an importance sampling estimator with K = 100 samples, which implies that in order to achieve codelengths equal to those we would have to combine our method with the recently-introduced McBits method, which supports the application of a range of
Revised: (2020) are obtained using an importance sampling estimator with K = 100 samples, which implies that in order to achieve codelengths equal to those we would have to combine our method with the recently-introduced McBits method, which supports the application of a range of

Original: Monte Carlo inference algorithms as part of the bits-back coding scheme (Ruan et al., 2021).
Revised: Monte Carlo inference algorithms as part of the bits-back coding scheme (Ruan et al., 2021).

Original: In future work we intend to extend and formalize our lossless compression scheme for the entire family of models covered by the SurVAE framework (Nielsen et al., 2020), along with providing a practical implementation of our proposed method.
Revised: In future work we intend to extend and formalize our lossless compression scheme for the entire family of models covered by the SurVAE framework (Nielsen et al., 2020), along with providing a practical implementation of our proposed method.

Original: A A PPENDIX A. D ECODING ALGORITHMS
Revised: A A PPENDIX A. D ECODING ALGORITHMS

Original: [Table] [Table]
Revised: [Table] [Table]

Original: | Z

Original: [Equation]

Original: |

Original: [Equation]

Original: z
Revised: U

Original: [Equation]
Revised: [Equation]

Original: u 7: ¯ x ← ¯ z L
Revised: 7: ¯ x ← ¯ z L

Original: A. E XPECTED CODELENGTH
Revised: A. E XPECTED CODELENGTH

Original: To compute the expected codelength let us sum the expected codelengths of the individual terms coded in Algorithm 1 in their coding order:
Revised: To compute the expected codelength let us sum the expected codelengths of the individual terms coded in Algorithm 1 in their coding order:

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: Now
Revised: Now

Original: [Equation]
Revised: [Equation]

Original: we notice that the second, third, and fourth terms are very close to the CIF joint distribution log p X,U ( x , u ) = log | det J ( x , u ) |
Revised: we notice that the second, third, and fourth terms are very close to the CIF joint distribution log p X,U ( x , u ) = log | det J ( x , u ) |

Original: + log p U | Z (
Revised: + log p U [Equation]

Original: [Equation]
Revised: [Equation]

Original: but now z is not exactly equal to F − 1 ( x ; u ) , as the two are related according to the local bits-back model
Revised: but now z is not exactly equal to F − 1 ( x ; u ) , as the two are related according to the local bits-back model

Original: [Equation]
Revised: [Equation]

Original: for some (cid:15) ∼ N ( 0 , I ) (with ¯ z representing the quantized version of this relationship).
Revised: for some (cid:15) ∼ N ( 0 , I ) (with ¯ z representing the quantized version of this relationship).

Original: However, since σ (cid:28) 1 , we expect equation 11 to approximately hold and thus (like in Ho et al.
Revised: However, since σ (cid:28) 1 , we expect equation 10 to approximately hold and thus (like in Ho et al.

Original: we consider the asymptotics of equation 9 in σ .
Revised: we consider the asymptotics of equation 9 in σ .

Original: Consider ﬁrst log p Z , which will typically be a standard normal distribution.
Revised: Firstly, let us consider log p Z , which will typically be a standard normal distribution.

Original: If we assume this to be the case, then for any ξ ∈ R d , with (cid:15) ∼ N ( 0 , I ) :
Revised: If we assume this to be the case, then for any ξ ∈ R d , with (cid:15) ∼ N ( 0 , I ) :

Original: [Equation]
Revised: [Equation]

Original: where the O ( σ ) term disappears because  (cid:15) has mean 0 .
Revised: where the O ( σ ) term disappears because E (cid:15) [ (cid:15) ] = 0 .

Original: Therefore,
Revised: Therefore,

Original: [Equation]
Revised: [Equation]

Original: with the approximation error stemming from quantization.
Revised: with the approximation error stemming from quantization.

Original: As for the log p U | Z term, we have a bit more work to do.
Revised: As for the log p U | Z term, we have a bit more work to do.

Original: We willoften assume this is a conditional 
Revised: We will assume this is a conditional diagonal

Revised: Gaussian with a covariance matrix diag ( s ( z )) where diag ( · ) returns a matrix with consecutive values of the vector s on the diagonal.

Original: Gaussian deﬁned as
Revised: log p U | Z is hence deﬁned as

Original: [Equation]
Revised: [Equation]

Original: where µ and Σ are two different outputs of the same neural network, and Σ is diagonal.
Revised: where µ and s are two different outputs of the same neural network.

Original: We see that log p U | Z dependsquadratically on µ ( z ) and 
Revised: We see that log p U | Z depends on µ ( z ) and s ( z ) .

Original: (cid:112) Σ( z ) − 1 (as we typically parametrize the standard deviations and not the variances).

Original: Thus, if µ and Σ are parametrized such that small changes in their inputs can only induce small changes in their outputs – say if they are 1-Lipschitz functions – then we can say that (refer to discussion starting from equation 12)
Revised: Thus, if µ and s are parametrized such that small changes in their inputs can only induce small changes in their outputs—say if they are 1-Lipschitz functions—then, as we show below (discussion starting from equation 14), we can say that

Original: [Equation]
Revised: [Equation]

Original: 2 ) , where again the error originates from the discretization.
Revised: where again the error originates from the discretization.

Original: Thus [Equation] we can ﬁnally rewrite equation 9 as:
Revised: Thus, we can ﬁnally rewrite equation 9 as:

Original: [Equation]
Revised: [Equation]

Original: Now, assuming δ u is small enough that the probability mass function q U | X (¯ u | ¯ x ) δ u is approximately equivalent to the probability distribution function q U |
Revised: and now, assuming δ u is small enough that the probability mass function q U | X (¯ u | ¯ x ) δ u is approximately equivalent to the probability distribution function q U |

Original: X ( u | ¯ x ) [Equation] then
Revised: X ( u | ¯ x ) , then

Original: [Equation]
Revised: [Equation]

Original: Now [Equation] we give further details on the analysis of p U | Z (¯ u | ¯ z ) : First recall that (up to discretization error)
Revised: Now, we give further details on the analysis of p U | Z (¯ u | ¯ z ) : First recall that (up to discretization error)

Original: [Equation]
Revised: [Equation]

Original: for (cid:15) ∼ N ( 0 , I ) .
Revised: for (cid:15) ∼ N ( 0 , I ) .

Original: Now let’s write out p U | Z :
Revised: Now let’s write out p U | Z :

Original: [Equation]
Revised: [Equation]

Revised: For term (b),

Original: [Equation]
Revised: [Equation]

Original: Taking Taylor expansion of ¯ σ (¯ z ) − 1 and µ (¯ z ) around y =
Revised: taking Taylor expansion of s ( y + σ J (cid:15) ) − 1 and µ ( y + σ J (cid:15) ) around y =

Original: F − 1 (¯ x ; ¯ u ) 
Revised: F − 1 (¯ x ; ¯ u ) =

Original: [Equation]
Revised: [Equation]

Original: Since E (cid:15) [ (cid:15) ] ==
Revised: rearranging, for some term C independent of σ and (cid:15) we get =

Original: [Equation]
Revised: [Equation]

Original: Now, term (a) E (cid:15) [ (a) ] = E (cid:15) (cid:34) −
Revised: since (cid:15) is independent of all other terms

Original: [Equation]
Revised: [Equation]

Original: (cid:88) i log ¯ σ ( ¯ z ) i (cid:35) (16)
Revised: and since E (cid:15) [ (cid:15) ] =

Original: Taking Taylor expansion of log ¯ σ (¯ z ) − 1 around y =
Revised: [Equation]

Original: F − 1 (¯ x ; ¯ u ) =
Revised: O ( σ 2 ) .

Original: E (cid:15) (cid:34) [Equation] i
Revised: Now, for term (a),

Original: [Equation]
Revised: [Equation]

Original:  log ¯ [Equation] [Equation] [Equation]
Revised: taking Taylor expansion of log s ( y + σ J (cid:15) ) around y

Original: +
Revised: =

Original: [Equation] ¯ 
Revised: F − 1 (¯ x ; ¯ u )

Original: [Equation]
Revised: [Equation]

Original: ∇ y ¯
Revised: and again, since (cid:15) is independent of all other terms and E (cid:15) [ (cid:15) ] =

Original: [Equation]
Revised: [Equation]

Original: y
Revised: Hence

Original: [Equation]
Revised: [Equation]

Original: σ(cid:15) + O ( σ 2 ) (cid:1)(cid:35) (17)

Original: Again, by the fact that E (cid:15) [ (cid:15) ] = = − 1 2 (cid:88) i log ¯ σ ( y ) i (18)
Revised: (16) leading to equation 13.