Generalization Gap in Amortized Inference
Authors: Mingtian Zhang
Peter Hayes
David Barber
Venue: NeurIPS 2022 Accept (2022)
Keywords: variational inference
amortized inference
vae
lossless compression

Abstract: The ability of likelihood-based probabilistic models to generalize to unseen data is central to many machine learning applications such as lossless compression. In this work,  we study the generalization of a popular class of probabilistic model - the Variational Auto-Encoder (VAE). We discuss the two generalization gaps that affect VAEs and show that overfitting is usually dominated by amortized inference. Based on this observation, we propose a new training objective that improves the generalization of amortized inference. We demonstrate how our method can improve performance in the context of image modeling and lossless compression.

Generalization Gap in Amortized Inference
Mingtian Zhang
Peter Hayes
David Barber
Centre for Artificial Intelligence, University College London zhang, hayes, d. barber) ucl ac uk
Abstract
The ability Akelihoodd-asee probabilistic models to generalize unseen data is central to many machine learning upplications such as lossless compression. In this work, we study generalization of a popular class of probabilistic model the Variational Auto-Encoder (VAE). We discuss the two generalization gaps that affect VAFs and show that overfitting usually dominated by amortized inference. Based on this observation, the generalization of amortized inference. We demonstrate how our method can improve performance in the context of image modeling and lossless compression.
propose new training objective that improves
Introduction
Probabilistic models have achieved great success in many machine learning applications [5 3). Given a set of training data that are sampled from an underlying data distribution Xtrain {-1, .. ZxN ~
pa(x), the goal of probabilistic modelling method to learn 9 is to minimize the Kullback-Leibe (KL) divergence
approximate pa(#) with model po A principled
(L(pa(<)) (x)) (log)a(<))
(log (logro(z))pa(e)))
where we use (.) to denote integration: (f(z)) f+))(())d The first term represents the negative entropy of the data distribution (pa) = (log pa(x)) which is constant. The second cross entropy term involves the integration the unkown data distribution Pa(x), which can be approximated by the Monte-Carlo approximation using the training dataset Xtrain
(log po po(2))
log))(#n)
(2)
==1
Therefore, estimating by minimizing the KL divergence is equivalent Maximum Likelihood Estimation (MLB) when N + co.
For finite dataset, common concern in both supervised and unsupervised learning is that the probabilistic model may overfit to the training dataset Xtrains degrading generalization perfor- mance 32 The generalization performance in the unsupervised setting can be measured by the test
likelihood 49 ΣM log pe (z')), where Nest = {z/.... ................................ z'M pa(x) is the test dataset. A model that has overfit to the training dataset Xtrain generally results in a high training likelihood but a low likelihood. Although the test likelihood is a common evaluation criterion (36), the factors that affect the generalization of unsupervised probabilistic models are less well studied in comparison to supervised learning. We posit that this because for common tasks, like sample generation or representation learning, good generalization in terms of the test likelihood is not sufficient measure of performance. For example implicit models can generate sharp samples without having likelihood function [6. 46) and representations learned by latent variable models can be arbitrarily transformed without changing the likelihood 25 However, in recent applications that use deep
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

22, 49 47 generalization in terms of the test generative models for lossless compression 38 39
likelihood directly indicates higher compression rate 49 Specifically, given a probabilistic model (o()), a lossless compressor can be constructed to compress test data point x' to bit string with length approximately equal to log, (a()). When po(x) + pa(x), the average compression length attains the entropy of the data distribution M _m=1 log, pe (2 + (pa)), which is optimal under Shannon's ' source coding theorem [4 Appendix
[ for detailed introduction. Therefore, a better test likelihood can lead to greater saving in bits and so understanding and improving
generalization of deep generative models is an important challenge.
1.1 Variational Auto-Encoder
A popular type of probabilistic model is the Variational auto-Encoder (VAE) 21 29). which assumes
a latent variable model po(x) = f po (2t a =))(2)dz.. For a nonlinear parameterization of po po(z)2) (e.g. deep neural network), the evaluation of log (2())) involves solving an intractable integration over z. In this case, the evidence lower bound (ELBO) can be used to side-step the intractability
(log po(x)) (z()) 2 (log (1 2) - log g ( ()) (= (ELBO(z, 8, o)) pa(x))
(3)
where (=(=)r) is a variational posterior parameterized by a neural network with parameter Φ. The use of an approximate posterior of the form (=(2)))) is called amortized inference. To better understand this objective, we can rewrite the expected ELBO the following
pa(a) (,))) log po (o (x) (eeo))KL(((AA (L[(qa(l pa(x)
(4)
(pa) KL(pa(=))llo(()) (KILgg(el pa(a))
(5)
const.
model learning
amortized inference
We denote the posterior family of (z(zr) as Q. which is indexed by finite dimensional 0 43
If Q flexible enough such that the true posterior po(2|r) E Q, where po(2|z) xx Po [ p))p then in the optimum of Equation 4 we have (L(g (pp
0 go (2 (z) for
x ~ Pa(x) and the ELBO will be equal to the log lig-liieddood ELBO(z, 0, 6) = log [21 6].
Many methods have been developed to increase the flexibility of Q, e.g. adding auxiliary variables
[126 26 or flow-based methods
28 to obtain a tighter ELBO.
Recent works [8 39 22 have successfully applied VAE style models to lossless compression realizing impressive performance. In this setting. the average compression length on the test data set is approximately equal to - 1 Im= ELBO(z 0, σ) (also see Appendix E Hence the better the test ELBO indicates the better the compression performance. This motivates to study the factors that affect the generalization of VAEs and find practical ways to improve the generalization of VAEs.
Mf
The contributions of our paper are summarized as follows:
We show the generalization of VAES is affected by both the generative model (decoder) and the amortized inference network (encoder); and that the overfiting of VAEs is mainly dominated by the amortized inference.
We propose a new training objective that can improve the generalization of the amortized inference without changing the model itself.
We demonstrate how the proposed method can improve the compression rate in a practical lossless compression system without scarifying any computation speed.
2 Generalization of VAEs
During training. we only have access to a finite dataset Xtrains which leads to the following Monte- Carlo approximation as our objective to train VAEs:
FE.BO O, ②)) pa(e)
ELBO (r,,, 0)
(6)
N

---

This empirical approximation will lead to the VAE overfits the training data for finite N. For example, we train VAE on the Binary MNIST dataset for I epochs and plot the Bits-per-dimension (BPD both training and testing dataset for every 100 epochs, also see Section 4 for model and training details. Figure ! visualizes the training and testing BPD, which shows the VAE model is overfitting to the training dataset.
020
Test
Train
018
The decomposition in Equation suggests that the empir- ical ELBO contains 1) a model empirical approximation:
100
200
300
4000
500
600
720
883
500
1000
Figure 1: BPD vs epochs. The training BPD decreases but the testing BPD in- creases during training, which indicates the VAE overfits to Xerain.
(L(pa(=))) Ioe (x))
> log + const., (7)
which will potentially make flexible model po( overfit the training data; and 2) an amortized inference empiri- cal approximation:
(xl=(()))))l pa(m)
Ex- (L(gs(2) (z)le))
(8)
where similarly a flexible qo(2|a) can also overfit to the training data. More specifically, we let @ be the optimal parameter of the empirical variational inference objective
arg min N E KL
(9)
and we assume for any training data point 2n E Xerain
(z(=/),)) arg min (=lgo(slaw))lpo
q" (z/xn))
where q² (2(2xn) is the realizable optimal posterior (in the Q family) for When 98 (2(2r,) overfits to Aerain 93 (2(2(='m) may not be good approximation to the true posterior pa (2) (2|=<m) for test data I'm Neet,, We refer the difference between the ELBO evaluated using (=(x)) and the ELBO evaluated using q" (2/x) (r) as the amortized inference generalization gap, which is formally defined as
KT (=]e)(())))l))(((]))))
(s)x)() Pa(x)
(10)
Equivalently, this gap can be written as the difference between two ELBOs with two different
(log Moggpolf po (2,, =) log q" (=====(())))))))>>>)>>)) (=|=)) (=|x) (logro(1: log g, (2(2)))))
patx)"
(11)
ELBO with optimal inference
ELBO with amortized inference
The inference neural network introduced by amortization the cause this inference generalization gap. It important to emphasize that this gap cannot be reduced by simply using more flexible Q. This would only make (L(gg (2(xx,) (poo))))))) smaller for the training data xr Xtrain but would not explicitly encourage better generalization performance on test data 335
To summarize, the generalization performance of VAE depends
factors:
Generative model generalization gap: defined as (p(p(())]lpo(()) and is caused by the generative model overfitting the the training data.
Amortized inference generalization gap: defined in Equation II and is caused by the amortized inference model (encoder) overfitting to the the training data.
2.1 Impact of the Generalization Gaps
The generative model generalization gap that is estimated by the test dataset (up to constant) (L(p(())lpooz)) x M Em= log Pe (2. ) + const. cannot be calculated explicitly since
In the case of VAE, the BPD is defined as the the negative ELBO (with base logarithm) normalized by the data dimension, lower BPD indicates higher ELBO.
Forr a powerful inference network we assume that there amortization gap TT], which means (=(2)))
provide the optimal g² (z|z,, (.)) for any training data In Xtrain Section for further discussion.

---

only evaluate the lower bound M we know the optimal posterior for the test data q" (=| = arg mingeo (L(((<2)r [po (=) ae )), the log-likelihood can be approximated by the lower bound logpo(<', z ELBO O, Φ) and
Em= M
ELBO(4 8, 6). Fortunately, as suggested in Equation 4
the approximation becomes an equality when pa m=) E Q. Similarly, the amortized inference generalization gap can be estimated by knowing the optimal posterior (=|xr))) for the test dataset:
MM
1 M
(1 2) - log q" (=(mm)) ( (xlor)) (log pa (1 z) log (1 (12)
We can then estimate q' (=/ar'm) by fixing (which is trained on the training dataset) and learning e" on the test dataset and assuming q" (=larm) = qø- (2/ar)n) where
ming KL (Ls pl (=(==",))) = maxş (log po (2'm 2) log qs (=(2)=)))) (=lo,,,,
(13)
This optimal inference strategy can eliminate the effect of the inference generalization gap, allowing us isolate the degree which both the generative model and amortized inference generalization gaps are contributing to the overfitting.
We take the VAE described in Section 2 and train (z(2(z) for 1k epochs with Adam (00 and Ir = 5x10 on the test data using Equation 13 to obtain the test BPD for the optimal inference strategy. In Figure 2 we plot the test ELBO (BPD) using the optimal inference strategy (green) and classic amortized inference (purple). Since for the optimal inference strategy the average likelihood
Test
Optimal
Train
amortized inference generalization gap
1 -M logpo (x((m) can be effectively approximated M by the ELBO (see Appendix for an empirical verification of the tightness of the ELBO), then the difference between the inference curves on the test set (Test and Optimal) is the amortized inference generalization gap. We observe that after eliminating the inference generalization gap. the test BPD is stable with marginal increase during training. This suggests the generative model (decoder) slightly overfit the data but that the overfitting is mainly dominated by the overfitting of the amortized inference network.
100
200
400
600
700
500
1000
Figure 2: Test BDP vs epochs. Demon- strates the amortized inference general- ization gap in a VAE trained on MNIST.
Although the optimal inference strategy can help eliminate the inference generalization gap, training qa on the test data is not practical in most applications of interest. Therefore, we now focus on improving the generalization of amortized inference without access to the test data training time.
3 Improving Generalization with Consistent Amortized Inference
We now propose an inference consistency requirement which, satisfied, would result in optimal generalization performance for amortized variational inference. Specifically when pa - pa, the amortized posterior should converge to the true posterior (=(=)²) - pa (=)::=]]] for every x ~ pa(x). Although this requirement seems natural for variational inference, the classic amortized inference training that is used for VAES 21 doesn't satisfy it. Recall the typical VAE empirical BLBO training objective
EN log (z..) - KI[gg([)/2))llo<<(le.))
(14)
When the model converges to the true distribution Pa- = Pa the training criterion for go (2(2²)
min - //2 ΣM (z(=0 |ppp (z|
(15)
can still result in the amortized posterior go (=(xr) overfitting to the training data. In principle, one could also limit the network capacity and/or add an explicit regularizer to the parameters [32) in an attempt to improve the generalization. However, this still cannot satisfy the consistency requirement in principle because it still only use the finite training dataset. Alternatively, there is another classic variational inference method that we now discuss, the wake-sleep training algorithm [3 [8] which does in fact satisfy the proposed consistency requirement.
We assume the true posterior belongs to the variational family (z(xx) C

---

3.1 Wake-Sleep Training
Defining 9a (a(zz 2) = go (z( and po(2, 2) = po (z))))(((2) the two phases of the wake-sleep
training [13 18 can be written as minimizing two different KL divergences in both x and 2 space.
Wake phase model learning: po(z]2) is trained by minimizing the KL divergence
ming (L(gs(r, =)] (2. =)) = maxe (ELBO(z, 0, o)) pa(x) + const.,
(16)
where pa))()) approximated using the training set. This is referred to as the wake phase since the model is trained on experience from the 'real environment', i.e. it uses true data samples from pa(æ:)
Sleep phase amortized inference: (2(22) is trained by minimizing the KL divergence ming (L(po(x, z) (lass, =)) = ming (=L)po())+)llao (2|())) pa(a) + const.
(17)
Leaving out the terms that are irrelevant to Φ. the objective can be estimated with Monte-Carlo log 4s (=)=r
K _k=1 log to (= where 2 p(z) and Iy pe (x) This is referred to as the sleep phase because the samples from the model used to train ge are interpreted as dreamed experience. In contrast, the training criterion for the typical VAE amortized inference (Equation 8) uses the true data samples from pa to train (o(2|zr), which we refer to as wake phase amortized inference. We notice that if perfect model po- (x) Pa(z) is used in the sleep phase amortized inference, then it is equivalent to minimizing
R 1 >
( (po (=]x))llu<( (z(=))) pa- (2) (KL(po (=<=)\[>><))=(())) pa()))
(18)
Therefore, the training of the inference network satisfies the inference consistency requirement since we can access infinite training data from Pa by sampling from pa-.
However, the wake-sleep algorithm presented lacks convergence guarantees [13] and minimizing (L (2|z)) q (z|=)) in the sleep phase doesn't 't necessarily encourage an improvement to the ELBO, which directly relates the compression rate in the lossless compression application (38). Therefore, in the next section, we propose a new variational inference scheme: reverse sleep amortized inference and demonstrate how it helps improve the generalization of the inference network in practice.
3.2 Reverse Sleep Amortized Inference
We propose use the reverse KL divergence in the sleep phase. We fix and train é using mins (KL(go<))=)lba (2]x))) (x)
pa))) = max log (2, 2) - log go go (=|²)) (f (=))))))()))
(19)
where the integration (z)) approximated by Monte-Carlo using samples from the generative
model (x)) This reverse KI. objective encourages improvements to the ELBO. When have perfect model pa- (I) = pa( the reverse sleep phase is equivalent to mins (KL (=]=) pa- (z) = ming (KL (z/=)]l pa(x)
(20)
which satisfies the inference consistency requirement.
075
Wake Inference (VAE)
The consistency requirement can also be validated empir- ically when the perfect model known pa (x) = Pa This can be achieved by using a pre-trained VAE the true data generation distribution. Therefore, we first train
Sleep Inference
R-Sleep Inference (Ours)
300
VAE fit the binary MNIST problem. The VAE has the same structure as that used in Section 2 and is trained using Adam with Ir 1x10-3 for 100 epochs. After training, we treat the pre-trained decoder pa (x|z) as the training data generator Pa(z) = pe (a =))p(())d d2) We then sample 10000 data samples from Pa to form a train- ing set Merain and 1000 samples form test set Xtest We then train new (2(2) with: 1) wake phase inference (VAE) 2) (forward) sleep inference and 3) reverse sleep inference. The network is trained using Adam with Ir = 1x10-3 3 for 100 epochs. Figure 3 shows the test BPD calculated after every training epoch. We can see the sleep phase out-performs the wake phase and the reverse sleep inference achieves the best BPD. Intuitively, this is because both the forward and reverse sleep inference use the true model to generate additional training data whereas the wake inference only has access to the finite training dataset Xtroin-
200
120
Figure 3: Test BPD vs epochs. We com- pare the consistency property between three amortized inference methods.
5

---

3.3 Reverse Half-asleep Amortized Inference with Imperfect Models
In practice our model will not be perfect pe # Pa. Empirically we find that samples from even well trained model pe may not always be sufficiently like the samples from the true data distribution. This can lead degradation in the performance of the inference network when using the reverse-sleep approach. For this reason, we propose to use a mixture distribution between the model and the empirical training data distribution as follows
(KL. (1o (z|2=) Ie (z)z))) m(x) where m(x) = (xpp(r) + (1 - --)a--
(21)
When a = O, it reduces to the standard approach used in VAE training. When a 1, we recover the reverse sleep method (Equation T9p. We find that a setting of 0.5 works well in practice. This
balances samples from the true underlying data distribution with samples from the model.
We thus refer this method reverse half-asleep since it uses both data and model samples train the amortized posterior. Intuitively, we can rewrite the Equation 2T as a sum of two positive terms <<KL( (1(]
p(w) + (1 -
(as (2/2x) (2)=)))p((((xx)) (x
(22)
Therefore, the optimal of this objective will make the first term O, which is the same requirement as the classic amortized inference (Equation 8). The second term, which is equivalent to the
reverse sleep amortized inference (Equation 19p. can
courage the inference consistency requirement: when pa = pa. the optimal of the second term will set (z(2))) (=(())) for any I ~ Pa(z). When pe is not perfect, the second term can be seen sic amortized inference objective, which can be used to pe- nalize the hypothesis space of the amortized network [32].
Sleep
Wake (a 0)
R.sloop (a 1]
regularizer added to the clas-
0.190
R-Half-asleep (a
0185
To compare with different a, we first fit a VAE (with the same structure that used in Figure 2) the Binary MNIST dataset, and then train the amortized posterior using sleep inference (Equation [7] and three different a for additional 100 epochs using Adam with learning rate 3x10- Figure 4 shows the test BPD comparison. We find the proposed reverse half-asleep method (ax 0.5) outperforms the reversed sleep method (a = 1), whereas the standard amortized inference training in VAE (a 0) leads overfitting of the inference network. We also plot the sleep inference training curve, whose BPD is less competitive since it is not directly optimizing the ELBO.
0180
Figure 4: Test BPD comparisons of Amortized inference with different a. We find the Reverse Half-asleep method (a 0.5) achieves the best BPD. The mean and std are calculated with three random seeds.
4 Generalization Experiments
We apply the reverse half-asleep to improve the generalization of VAEs on three different datasets: binary MNIST, grey MNIST 24 and CIFARIO 23). For binary and grey MNIST, we use latent dimension 16/32 and neural nets with 2 layers of 500 hidden units in both the encoder and decoder. We use Bernoulli p(æ|2) for binary MNIST and discretized logistic distribution for grey MNIST. We train the VAE with the usual amortized inference approach using Adam with Ir = 3x100 10-4 1000 epochs and save the model every 100 epochs. We then use the saved models to 1) evaluate on the test data sets, 2) conduct optimal inference by training ga(2|z0 on the test data and 3) run reverse asleep method before calculating the test BPD. For the reverse half-asleep, we train the amortized posterior for 100 epochs with Adam and ir 5x 10 To sample from pe(x), we firstly sample z' ~ p(z) and sample x' ~ p(x|2 = z²). For the optimal inference strategy, we train the amortized posterior with the same optimization scheme on the test data set for additional 500 epochs to ensure the same number of gradient steps are conducted (since training set is 5 times as big as the test set). Figure 5a and 5bb show the test BPD comparisons of binary and grey MNIST respectively and demostrate that our approach does not require further training on the test data to improve generalization performance.
for
For CI we use the convolutional ResNet 17 with 2 residual blocks and latent size 128. The observational distribution is discretized logistic distribution with linear autoregressive

---

120
Mttml
150
154
VSS
Optinal
MCC
Quitaal
H Half toleop
Sallabaseee3
816
1.0
616
1.7%
240
4.000
148
4575
130
20
900 1090
200
300
500
M00
20
100
300
Binary MNIST
(b) Grey MNIST
(c) CIFARIO
Figure 5: Test BPD comparisons among amortized inference (VAE), optimal inference strategy and the reverse half-asleep inference on three datasets. The x-axis represents the training epochs.
parameterization within channels. We train the VAE for 500 epochs with Adam and ir 5x10-1 and save the model every 100 epoch. The pre-trained VAE achieves 4.592 BPD the CIFAR1O, which is comparable with other single latent VAE models reported in 40): 4.51 BPD with VAE with latent dimension 256 and 4.67 BPD with a discrete latent VAE (VQVAE).
Ideally, when the VAE model converges to the true distribution Ps - Pas the aggregate posterior [x(2) 9 ( pa will match the prior p(=). However, for complex distribution like CIFARIO, significant mismatch between 4o and p(2) is usually observed practice [51 12). In this case, the sample x' that is generated using latent sample from the prior x' Po (o where z' ~ p(z). may be blurry or invalid. A common solution is to train another model, e.g. a PixelCNN [4 40) approximate ga(z). In case, we instead directly sample from go (z) rather than p(z) to generate samples in Equation 19 which can be done by first sampling x' ~ (a( (from the training dataset) and then sample
VAE 12
go(2|x x'). This scheme still results in consistent training objective since ga- (z) = p(z) for the optimal posterior 2-- (z|z).
We use Adam with ly = x 10-5 and train the reverse half-asleep inference for 100 epochs on the training data and train the optimal inference strategy for 500 epochs the test data, see Figure the result. We find the proposed reverse half-asleep training approach (with sampling from 4s (=)) consistently improves the generalization performance of the amortized posterior. We also apply the proposed method on a VAE trained on CIFAR 100 for 500 epochs (the rest of the experiment settings are the same as the CIFARIO case) and find our method improves the BPD from 5.288 to 5.275.
4.1 Comparisons with Regularization Methods
335 proposed to alleviate overfitting of amortized inference by optimizing linear combination between the traditional amortized inference (Fquation g and denoising objective
Recent work
a (KL(gg(>lz + e)]lpo()(=))) + [1 (t)))))((((((lll)looo
(23)
where p(e) = - ((000²1) We compare this regularizer to our method by training the amortized posterior of VAEs for an additional 100, 300 and 100 epochs on Binary, Grey MNSIT and CIFAR respectively. For the denoising regularizer, we use the same linear combination weight that used in Equation 2] and vary E {0.1, 0.2, 0.4, 0.6, 0.8. 1.0), see Table | for the comparisons. For MNIST, we find E (0.1, 0.2, 0.4) improves the generalization but larger noise levels hurts the performance. For CIFARIO, only In contrast, our method consistently achieves better generalization performance without tuning any cepee-parametedn see Figure 6 for the test BPD (evaluated every training epoch, the mean/std are calculated with 3 random seeds). Compared to the denoising approach, one limitation of our method is the requirement of model samples, which more computational expansive during training.
0.5 as
0.1 can slightly improve the generalization by 0.001 BPD.
Since the decoder is shared and fixed in all comparisons, better test ELBO indicates the predicted (2(2)))) closer to the true posterior Pa (z(2x²) under the KL divergence measure (see Equation higher ELBO with fixed 0 indicates (L( (2) (2)22) is smaller). Therefore, the proposed method can also benefit a range of tasks that require accurate prediction of the posterior on the test data. In Appendix A and B we demonstrate our method can provide better proposal distributions for the importance weighted Auto-Eneoded [7] and also improve the representation learning performance for down-stream classification tasks.
7

---

Table 1: Average test BPD comparisons with Denoising Regularizer [35
Methods
VAE
0.1
= 0.2
0.4
0.8
1.0
Ours
Binary MNIST
0.200
0.195
0.192
0.191
0.196
0.201
0.187
Grey MNIST
1.543
1.527
1.519
1.515
1.545
1.550
1.513
CIFARIO
4.592
4.591
4.598
4.614
4.651
4.667
4.572
1222
61
155
036
08
154
0145
153
470
132
400
151
(a) Binary MNIST
(b) Grey MNIST
(c) CIFAR 10
Figure 6: Test BPD evaluated after every training epoch. We find, compared to the denoising regular- izer, the proposed amortized inference training scheme consistently achieves better generalization performance in all tasks.
Application of Lossless Compression
Lossless compression is an important application of VAEs where generalization plays a key role in the compression rate. Given VAE with po(x)2), (2)x) and p(z), practical compressor can be efficiently implemented using the Bits Back algorithm [9 38 with the ANS coder [14]. See Appendix E for detailed introduction of conducting lossless compression with VAE models. In Al- gorithm [ we summarize the Bits Back procedure with amortized inference to compress/decomppess test data point x' stack that contains bit string messages. The resulting code length for data x! approximately equal to the negative ELBO
logz pe (x'12') loga p(z') + loga (=(()'²))
(24)
We have shown that qa (2(2)) may overfit the training data, degrading compression performance. To improve the compression BPD, the optimal inference strategy can also be applied in the Bits Back algorithm. In the compression stage, we can train & by
Algorithm 1 Bits Back with Amortized Inference.
Comp./decomp. stages share {po(#)=), (=()))))((2))
Compression
Draw sample 2' ~ qa(22')) from the stack. Encode x' p=(22)) onto the stack.
Encode z' p(z) onto the stack.
arg max; ELBO(2' 0)
(25)
Decompression
When the qo (2(22²) (a)) is parameterized be a Gaussian, we can just take o to be the mean and standard deviation (( d²), which only contains two training parameters. In the decompression stage, we observe that the compressed data x' recovered before the go (z|x²) is used to encode z'. Therefore, we can also train the (=(22z)) using the recovered x' to maximize the test ELBO. If the optimization procedure the same as that used the compression stage, we will the same qs- (z(x'). In practice, we need to pre-specify the number of gradient descent steps K. When K is large, we recover the optimal inference strategy and the code length approximately
Decode z' ~ p>)() from the stack.
Decode x' ~ po(x]22) 2") from the stack.
Encode z' ~ (2(2x²)) onto the stack.
- loga (o(x'2²') logg p(z') + logz 4 (2'|xx))
(26)
This observation was first proposed in 45 in the context of lossy compression and then applied
to lossless compression with Bits Back coding in [30]. Furthermore, by varying the optimization steps K in the optimal inference, can trade off between the speed and the compression rate. This
is valuable for practical applications with different speed/rate requirements. See Algorithm for summary of the Bits Back algorithm with K K-step optimal inference.

---

Although the optimal inference strategy can be used in lossless compression, it re- quires extra run-time for training at the compression stages. In contrast, our pro- posed reverse half-asleep inference scheme can improve the compression rate without scarifying any speed. Additionally, our method can also provide a better initializa- strategy to
Algorithm 2 Bits Back with K step Optimal Inference
Comp./decomp. stages share {py (2|2), go (r), p(2))
and the optimization procedure of Equation 25
Compression
Take K gradient steps Φ + oK with Equation 25
Draw sample 2' ~ -qa (2)z') from the stack.
Encode x' ~ p=(222²) onto the stack.
tion for the allow a better trade-off between compres- sion rate and speed.
Encode 2' ~ p(z) onto the stack.
Decompression
Decode x' ~ p(=) from the stack.
We implement Bits Back with ANS [4] and compare the compression among four inference methods:
Decode x' ~ p0(())z)) from the stack.
Take K gradient steps ø + 6K with Equation 25
Encode z' q,x (z|z²) onto the stack.
1. Baseline:
This is the classic VAE-based compression introduced by [8). For bi- nary and grey MNIST, both the encoder and decoder contain 2 fully connected layers with 500 hidden units and latent dimension 10. The observation distributions are Bernoulli and discretized Logistic distribution respectively. For CIFARIO. we use fully convolutional ResNets [17] with 3 residual blocks in the encoder/decoder, latent dimension 128 and discreteized Logistic distribution with channel-wise linear autoregressive(31 [31] as the observation distribution. We train both the amortized posterior and the decoder by maximizing the ELBO (Equation 3) using Adam with Ir = 3x 10-4 for 100, 100 and 500 epochs (for Binary MNIST, Grey MNIST and CIFARIO respectively), and then apply Algorithm to conduct compression.
2. Reversed Half-asleep: we do amortized inference using Equation 21 for 100 and 300 epochs with Adam optimizer (Ir = 3x 10-4) for binary and grey MNIST respectively, y, and Ir 1x10-8 for 100 epochs for CIFAR10. Other training details are the same as the baseline method.
3. Optimal Inference: take the amortized posterior (encoder) and decoder from the baseline and apply the K-step optimal inference strategy described in Algorithm 2 to do compression. We use Adam optimizer and vary the K from to 10 achieve trade-off curve between compression rate and speed. We actively choose the highest learning rate that can make the BPD consistently improve with the increment of K: Ir 5x100= for binary and grey MNIST and Ir 1 xx10- for CIFARIO.
Reversed Half-asleep + Optimal Inference: we take the encoder in method 2 and decoder from the baseline and conduct K step optimal inference. All other training details are as per method 3.
Easaine
Raswline
Ravei Inference
Riot astepp
R Holt inference
1.000
3134
K-half Optia-
wieeence
3.22
4.100
2502
4150
012
130
3374
1545
(a) Binary MNIST
(b) Grey MNIST
(c) CIFARIO
Figure 7: We plot the comparisons for different methods. The y-axis is the BPD and x-axis represents the K gradient steps in the optimal inference. The baseline and our R-Half-sleep can be seen as special cases of optimal inference with K O. We find given fixed computational budget, our method achieves lower BPD than one using traditional amortized inference training.
In Figure 7
we plot test BPD comparisons for the different methods outlined. We can see if optimization is not allowed at compression time, the use of our reverse-half-asleep method achieves better compression rate with no additional computational cost. If we allow K-step optimization during compression, for a given computational budget, the amortized posterior initialized using our
'Implementation can be found in the following repo: https: //github cwm/ com//mtomorrow/ Smaeraaliatttoooppllloootttteeee All experiments
NVIDIA V100 GPU.

---

Baseline
Ours
K=7
Baseline
Ours
K=8
BPD
0.185
0.179
0.179
BPD
4.602
4.585
4.585
Com. Time
0.006
0.006
0.013
Com. Time
0.27
0.27
0.38
Dec. Time
0.006
0.006
0,013
Dec. Time
0.26
0.26
0.38
Time Cost
0%
116.7%
Time Cost
0%
46.2%
MNIST
(b) CIFARIO
Figure 8: Compression (Com.) and decompression (Dec.) time comparison. We show that to achieve the same BPD as our method, the K-step optimal inference strategy that initializes the amortized posterior needs K 7 (binary MNIST) and K 8 (CIFARIO) steps for each test datapoint, which will cost an additional 116.7% and 46.2% of time respectively during compression.
reverse-hall waa-hhllf--seeey method also achieves lower BPD, which leads to a better trade-off between the time and compression rate. Table also reports the average time improvements of our method to compress a single MNIST and CIFARIO image respectively, which shows the effectiveness of our method.
6 Related Work
different perspective on generative models' generalization is proposed in paper [50] where the generalization is evaluated by testing if the model can generate novel combinations of features. However, the generalization defined in our work is purely measured by the test likelihood, which is different perspective and more relevant for the application of lossless compression.
Recent work [49] first studies the likelihood -ased generalization for lossless compression. They focus on the test and train data that are from different distributions whereas assume they follow the same distribution. Additionally, their model has tractable likelihood and relates to the generative model related generalization, whereas focus inference related generalization in VAEs.
Previous work [I]] studied the amortization gap in amortized inference, which caused by using q0- (2|xx) to generate posteriors for each input In rather than learning a posterior g. (2) for individually. This gap can alleviated using larger capacity encoder network. This amortization gap is fundamentally different from the inference generalization gap we discuss in this work since the latter focuses solely on test time generalization but the former problem also exists training time.
Recent work 30 proposes a compression scheme based on the IWAE 7 bound, which is tighter than the ELBO and thus improves the compression rate. However, this method has compress/decoppess multiple latent samples, which requires extra time cost. On the other hand, we focus on improving compression that only needs compress one single latent sample. Nevertheless, similar to the K K-step optimal inference strategy, our amortized training objective can also be used in the IWAE-basee method, which gives a better proposal distribution for importance sampling, see Appendix A for demonstration.
the ELBO-basee
Paper 8 considers the following data generation procedure oy pa(x), 21 po(2(x11) 22 (=(2))) and propose to enforce latent consistency between (2(x)) and (2(x)) for paired data (x1, x2) encourage the robustness of the learned representation. This procedure is close to the self-supervised contrasting learning method [10] where the augmented data is the reconstruction of the training data using the VAE model. In our paper, we want to encourage the sample from the model x' Po [z) bp(2)dz to have high ELBO under the model (Equation 19 to improve the generalization of the amortized inference and no paired data is required our procedure. Therefore. both motivations and methodologies are different from our method.
7 Conclusion
We have shown how the generalization of VAEs is largely affected by the amortized inference network and proposed a new variational inference scheme that provides better generalization as demonstrated in the application of lossless compression. Future work will study the generalization of the decoder model to further improve the performance of VAEs.
10

---

References
[1] V. Agakov and D. Barber. An auxiliary variational method. International Conference on Neural Information Processing, pages 561-566. Springer, 2004.
[2] M. Arjovsky, S. Chintala, and L., Bottou. Wasserstein generative adversarial networks. In International conference machine learning, pages 214-223. PMLR, 2017.
[3] D. Barber, A. T. Cemgil, and S. Chiappa. Bayesian time series models. Cambridge University Press, 2011.
[4] Y. Bengio, A. Courville, and P. Vincent. Representation learning: A review and new perspectives. IEEE transactions on pattern analysis and machine intelligence, 35(8):1778-118))) 2013.
[5] C. M. Bishop. Pattern recognition and machine learning. Springer, 2006.
[6] D. M. Blei, A. Kucukelbir, and J. D. McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association, (12(518):859-8777 2017.
[7] Y. Burda, R. Grosse, and R. Salakhutdinov. Importance weighted autoencoders. arXiv preprint 100::::590000000 2015.
[8] A. T. Cemgil, S. Ghaisas, K. Dvijotham, S. Gowal, and P. Kohli. Autoencoding variational autoencoder. Neural Information Processing Systems, 2020.
[9] E. Challis and D. Barber. Affine independent variational inference. Advances in Neural Information Processing Systems, 25:2186-2199 -194, 2012.
[10] Chen, S. Kornblith, M. Norouzi, and G. Hinton. simple framework for contrastive learning of visual representations. In International conference on machine learning, pages 1597-1607, PMLR, 2020.
[11] C. Cremer, X. Li, and D. Duvenaud. Inference suboptimality in variational autoencoders. In International Conference Machine Learning, pages 1078-1086. PMLR, 2018.
[12] B. Dai and Wipf. Diagnosing and enhancing vae models. arXiv preprint arXiv: 1903.0599 2019.
[13] Dayan, G. E. Hinton, R. M. Neal, and R. S. Zemel. The helmholtz machine. Neural computation, 7(5):889-99) 904, 1995.
[14] J. Duda. Asymmetric numeral systems: entropy coding combining speed of huffman coding with compression rate of arithmetic coding. arXiv preprint arxiv:1311.2555: 2013.
[15] B. J. Frey and G. E. Hinton. Free energy coding. In Proceedings of Data Compression Conference-DCCC CCCrrennc-DCC996 pages 73-81. IEEE, 1996.
[16] I. Goodfellow, J. Houget-Abadic M. Mirza, B. Xu, D. Warde-Farley, S. Orair, A. Courville, and Y. Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014.
[17] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770- 778, 2016.
[18] G. Hinton, P. Dayan, B. J. Frey, and M. Neal. The" wake-sleep" algorithm for unsupervised neural networks. Science, 268(5214): 11011118811100 1995.
[19] G. F. Hinton and D. Van Camp. Keeping the neural networks simple by minimizing the descrip- tion length of the weights. In Proceedings of the sixth annual conference on Computational learning theory, pages 5-13, 1993.
[20] D. Kingma and Ba. Adam: method for stochastic optimization. International conference on machine learning, 2015.

---

[21] D. P. Kingma and M. Welling, Auto-encoding variational bayes. International Conference on Learning Representations, 2013.
[22] F. Kingma, P. Abbeel, and J. Ho. Bit-swap: Recursive bits-back coding for lossless compression with hierarchical latent variables. In International Conference on Machine Learning, pages 3408-3417. PMLR, 2019.
[23] A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[24] Y. LeCun. The mnist database of handwritten digits. hrpp://yann. lecun. com/exdb/nnist// 1998.
[25] F. Locatello, S. Bauer, M. Lucic, G. Ractsch, S. Gelly, B. Schölkopf, and O. Bachem. Chal- lenging common assumptions in the unsupervised learning of disentangled representations. In international conference machine learning, pages 41 14 4124. PMLR, 2019.
[26] L. Maaløe, C. K. Sønderby, S. K. Sønderby, and O. Winther. Auxiliary deep generative models. In International conference on machine learning, pages 1445-1453, PMLR, 2016.
[27] D. J. MacKay. Information theory, inference and learning algorithms. Cambridge university press, 2003.
[28] D. Rezende and S. Mohamed. Variational inference with normalizing flows. In International conference on machine learning, pages 1530-1538, PMLR, 2015.
[29] D. J. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and variational inference in deep latent gaussian models. In International Conference on Machine Learning, volume 2, page 2. Citeseer, 2014.
[30] Y. Ruan, K. Ullrich, D. Severo, J. Townsend, A. Khisti, A. Doucet, A. Makhzani, and C. J. Maddison. Improving lossless compression rates via monte carlo bits-back coding, arXiv preprint arXiv:2102.11005 2021.
[31] T. Salimans, A. Karpathy, X. Chen, and D. P. Kingma. Pixelenn++: Improving the pixelenn with discretized logistic mixture likelihood and other modifications. arXiv preprint arXiv: 71170001557, 2017.
[32] S. Shalev-Shwartz and S. Ben-David. Understanding machine learning: From theory to algorithums. Cambridge university press, 2014.
[33] C. E. Shannon. A mathematical theory of communication. The Bell system technical journal, 27(3):379-423, 1948.
[34] C. E. Shannon. mathematical theory of communication. ACM SIGMOBILE mobile computing and communications review, 5(1):3-55, 2001.
[35] R. Shu, H. H. Bui. S. Zhao, M. Kochenderfer, and Ermon. Amortized inference regulariza- tion. Neural Information Processing Systems, 2018.
[36] L. Theis, A. v. d. Oord, and M. Bethge. A note on the evaluation of generative models, arXiv preprint 1r11v::511.018444 2015.
[37] J. Townsend. A tutorial the range variant of asymmetric numeral systems. arXiv preprint arXiv:2001.09 2020.
[38] J. Townsend, T. Bird, and D. Barber. Practical lossless compression with latent variables using bits back coding. International Conference on Learning Representations, 2019.
[39] J. Townsend, T. Bird. J. Kunze, and D. Barber. Hilloc: Lossless image compression with hicrarchical latent variable models. International Conference on Learning Representations, 2020.
[40] A. Van Den Oord, Vinyals, et Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
12

---

[41] A. Van Oord, N. Kalchbrenner, and K. Kavukcuoglu. Pixel recurrent neural networks. In International conference on machine learning, pages 1747-1756. PMLR, 2016.
[42] C. Wallace. Classification by minimum--messag--lngghh inference. In International Confer- ence Computing and Information, pages 72-81. Springer, 1990.
[43] Y. Wang and D. M. Blei. Frequentist connstenny of variational bayes. Journal of the American Statistical Association, 1144577::1147-1111,, 2019.
[44] I. H. Witten, R. M. Neal, and J. G. Cleary. Arithmetic coding for data compression. Communi- cations of the ACM, 30(6):520-540, 1987.
[45] Y. Yang, R. Bamler. and S. Mandt. Improving inference for neural image compression. arXiv preprint arXiv:2006. 04240, 2020.
[46] M. Zhang, P. Hayes, T. Bird, R. Habib, and D. Barber. Spread divergence. In International Conference on Muchine Learning, pages 11106-11116. PMLR, 2020.
[47] M. Zhang. J. Townsend, N. Kang, and D. Barber. Parallel neural local lossless compression. arXiv preprint arXiv:2001.05213, 2022..
[48] M. Zhang, T. Xiao, B. Paige, and D. Barber. Improving vae-based representation learning. arXiv preprint arXiv:2205. 14539, 2022.
[49] M. Zhang, A. Zhang, and S. McDonagh. On the outtof-distribution generalization of probabilis- tic image modelling. In Neural Information Processing Systems, 2021.
[50] Zhao, H. Ren, A. Yuan, J. Song, N. Goodman, and S. Ermon. Bias and generalization deep generative models: An empirical study, arXiv preprint 1r 0000229,.........111 2018.
[51] S. Zhao, J. Song, and S. Ermon. Infovae: Information maximizing variational autoencoders.
arXiv preprint arXiv: 1706.02222,, 2017.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] In Section
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as URL)? [Yes] See the footnote in page 8.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See section and 5.
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes] See Figure
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? [Yes] See the footnote in page 8.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] We cite the datasets in section 4.
(b) Did you mention the license of the assets? [N/A] We didn't find licences for MNIST
and CIFARIO.
(c) Did you include any assets either in the supplemental material as URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're 're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "ZLjlhO_pkzw", "reviewer": "Reviewer_rcEf", "summary": "This paper focuses on the generalization gap in amortized inference trained using the ELBO objective. The generalization gap can come from either the generative model p or the inference network q not generalizing from training data to test data. This paper argues that q is mainly responsible for the generalization gap, which is supported by an experiment using the VAE model and the MNIST dataset (Figures 1 and 2). The proposed fix is to train on $E_{m(x)}[ELBO]$ instead of $E_{\\hat p(x)}[ELBO]$ where $m(x) = \\alpha p_\\theta(x) + (1 - \\alpha) \\hat p(x)$, i.e. \"augment\" the empirical dataset by samples from a well-trained model. Experiments suggest that this indeed leads to better test ELBOs, which in turn leads to better compression rates.\n\nI gave this paper a 5 but am willing to increase the score if the authors address some of the comments below.", "questions": "I'd also appreciate authors' thoughts on the following points.\n\nThe paper claims that KL(q || p) is better than KL(p || q) because of compression. However, KL(p || q) also comes with its advantages, for example that it makes q spread its mass more and hence it can end up learning a better proposal for importance sampling.\n\nIn Line 107, it is claimed that $\\log p_\\theta(x)$ cannot be calculated explicitly and hence the generative model generalization gap cannot be estimated well. While it cannot be computed exactly, it can be estimated by using importance sampling, which is equivalent to the [importance weighted autoencoder (IWAE)](https://arxiv.org/abs/1509.00519) objective. In light of my previous questions about the amount of p's vs q's contribution to the genralization gap, it might be nice to see this estimation of the generative model generalization gap.\n\nThere doesn't seem to be overfitting in Figures 3 and 5c.\n\nWhat if we train using reverse half-asleep from the start instead of training using ELBO first, and only then fine tuning $\\phi$ using reverse half-asleep?\n\nIn the paragraph starting in line 228, data from $p_\\theta(x)$ is simulated by drawing $z \\sim q_\\phi(z) = \\int \\hat p(x) q_\\phi(z | x) \\mathrm dx$ instead of $z \\sim p(z)$ before drawing $x \\sim p_\\theta(x | z)$. Is this really that different from just drawing from $\\hat p(x)$ directly? Especially when looking at Figure 5c, it seems like there is underfitting going on, rather than overfitting. What would have happened if we used $p(z)$ instead of $q_\\phi(z)$?", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Originality\nAs far as I know, this is an original view on the problem and an original solution.\n\n# Quality\nI'd like to see some of the following concerns addressed.\n\nI am not entirely convinced that the generalization gap is mainly due to the amortized inference network, as claimed. I understand that in the case of this particular VAE architecture and the MNIST dataset, this is the case, as shown in Figure 2, but what about different model-dataset settings? What if the generative model family is more expressive, e.g. a PixelCNN? Surely some of p must be contributing to the generalization gap there. Another reason to believe this is that for more complex datasets like CIFAR10, there doesn't seem to be overfitting, e.g. as seen in Figure 5c.\n\nThe solution to the problem (using reverse half-asleep) only works if there is no generalization gap due to p. If there is a generalization gap due to p, the solution is ineffective. As I understand it, the proposed method is to \n1. Train using $E_{\\hat p(x)}[ELBO]$ and then\n2. Fine tune $\\phi$ (parameters of the inference network) using $E_{m(x)}[ELBO]$ where $m(x) = \\alpha p_\\theta(x) + (1 - \\alpha) \\hat p(x)$.\n\nAs suggested in lines 189-191 (\"Intuitively, this is because ...\"), this solution is effective because the inference network gets to be trained on a dataset, augmented by samples from the model, $p_\\theta(x)$. But if the model $p_\\theta(x)$ is trained on $\\hat p(x)$, how good can it be? In the extreme case of an infinitely expressive $p_\\theta(x)$, it will converge to $\\hat p(x)$ and $m(x) = \\hat p(x)$, rendering reverse half-asleep to be the same as normal ELBO.\n\n# Clarity\nOverall, the paper is clearly written. The compression section (Section 5) is difficult to understand without further context (e.g. what is the stack in Algorithms 1 and 2?), but the high level points are clear. Plot labels would have been useful. Small typos:\n- Line 153: Monte Carlo\n- Line 240: Equation number missing\n- Line 292 Sacrificing \n\n# Significance\nGeneralization of amortized inference is an important problem.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "LqPiVjLBPkg", "reviewer": "Reviewer_syyo", "summary": "The authors investigate the generalization of (amortized) variational auto-encoders in terms of their test likelihood. It is noted that \"overfitting\" in VAEs is largely contributed by the amortized inference network. To mitigate such overfitting, the authors develop reverse sleep and reverse half-sleep methods inspired by the wake-sleep algorithm, leading to better compression of the data distribution.", "questions": "My comments here are mostly a follow up to the main weakness previously mentioned.\n\n* The authors interestingly note that standard VAE training with ELBO (i.e. the wake phase) only has access to a finite training dataset. The sleep phase allows for virtually unlimited training data, and hence mitigates overfitting. One can subsequently think of how does the applied fix to the amortized network behave for the following applications other than compression to name a few:\n  * out-of-distribution detection/generalization\n  * sample generation\n  * semi-supervised learning\n\n* Another question to explore would be does this make VAEs better utilize the latent space? From the compression results, perhaps the answer is promising and tending towards yes. Is it more robust to the choice of the dimensionality of the latent space as a consequence of improving the amortization gap?\n\nI am curious if the authors have a plan for this, or could provide a small proof-of-concept for something beyond compression. ", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n* The authors develop on a likelihood-based characterization of what it means for VAEs to overfit. The demonstration on a toy example which clearly demonstrates the overfitting phenomenon is much appreciated.\n\n* The step-by-step development of the method with supporting experiments provides a relatively clear picture.\n\nWeaknesses:\n\nIt would appear that the method specifically caters to optimizing for best compression performance. Even in their own experiments, the authors note that without the half-sleep correction in the sleep phase, the generated samples may be degraded and not representative of the true data distribution. \n\nThis makes me wonder that the notion of test likelihood for generalization remains very limited in scope, and in principle only caters to compression performance (this probably would not be an issue in the case of a well-specified model, which is virtually impossible to achieve). In that sense, the insights developed in this paper may not hold across broader family of generative models.\n\nIn light of such missing broader evaluation of direct implications, I will be keeping my score to a weak accept.\n\nMinor:\n\n* Is there a typo in the summands in Eq. (8) where you need expectation w.r.t $q^*(z \\mid x^\\prime_m)$ and $q_{\\hat{\\phi}}(z \\mid x^\\prime_m)$ instead of just conditioning on $x$?\n\n* Missing closing \")\" in Eq. (5).\n\n* Line 153: Typo in \"Mon-Carlo\".", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "IZb13cxYVfM", "reviewer": "Reviewer_Ayo3", "summary": "The paper studies the generalization performance of Variational Auto-Encoders during test time, measured by the test likelihood. Generative model generalization gap and amortized inference generalization gap are introduced and it is empirically shown that the degradation in generalization is mostly dominated by the inference model. The authors then propose a new variational inference scheme named “reverse half-asleep” which helps close the gap caused by the inference model. Finally, the improvements in generalization and compression rate of using the new scheme over the traditional VAE are demonstrated on Binary MNIST, Grey MNIST, CIFAR10 datasets.", "questions": "The improvements might also be shown on CIFAR100 or ImageNet. I just wondered if the authors have performed experiments on those larger datasets. Is the trick for CIFAR10 in line 239 helpful in those settings?\n\nSome minor points:\n+ Line 140: missing the Equation number.\n+ Line 108: missing the negative sign of the KL estimation.\n+ It is clearer to include labels for axes in the figures.", "limitations": "Nothing to report.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is very well written. I can understand most of the contents in one read.\n\nThe novelty of the paper is the introduction of using reverse KL in the sleep phase during wake-sleep training. It is surprising to me that people have not done it previously.\n\nI think that this is an important paper. The technique is pretty straightforward to implement, hence would be easily adopted in practice.\n\nMy main concern, ~which is not addressed in the paper~, is on how using “reverse half-asleep” affects other metrics (like accuracy) during both training and testing phases. **Updated:** The authors pointed out in the response below that the performance on downstream classification tasks can be found in Appendix B.\n\nIt is only shown through Figure 2 that the overfitting of the VAE is mainly dominated by the overfitting of the amortized inference network. It would make the paper more technically sound if there is an analytical proof for that observation.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "I-JmU3Q-1Vz", "reviewer": "Reviewer_ofFq", "summary": "The authors examine the amortization inference generalization gap with the motivation of improving performance on lossless compresseion and image modelling. They introduce the reverse sleep amortized inference and reverse half-asleep amortized inference, components of a modified wake-sleep algorithm, and use these algorithms to train a VAE for MNIST, CIFAR10 and lossless compression experiments.", "questions": "How does a VAE perform when trained in one step (all of the parameters) using the half-asleep objective, ie. adding the additional expectation to the standard ELBO training, as oppose to the two-stage wake-sleep method? \n\nThe test error in figure 5 is still increasing with each epoch. Is this still an indication of overfitting?\n\nOn line 239 you state that you sample from q(z) instead of p(z). How is this done? Do you Monte Carlo marginalize out the training data from q(z|x)p(x)?\n\nTypos: \nLine 117: \"the VAE that is described in Figure 2\". No description is given here.\nEquation 12: The expectation is written over the joint p(x,y) but I think you want p(x).\nLine 240: Equation ?.", "limitations": "No discussion of the limitations of the method was present.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nThe investigation is well motivated by the application to lossless compression.\nI thought that the intuition presented for the aspects of the lack of generalization in section 2 were good.\n\nWeaknesses:\nI did not think the intuitions in section 2 were well supported by they toy experiments presented.\n\nI did not find the exposition clear. I often found myself confused when the authors explained background that I am already familiar with. The authors would introduce a rearrangement of the ELBO, as in equation 2, only to immediately introduce another rearrangement in equation 3 leaving me wondering why early equations were introduced. I did not find the explanation of the variational family starting on line 56 very clear. \n\nThe method used in the paper is very similar to the one used in [4] Cemgil et al. 2020, as noted in the related work. Unfortunately, there are not a lot of additional contributions to distinguish this paper from [4], other than a reinterpretation of the method.\n\nThe authors presented a very long explanation for what is a relatively simple idea: they want the inference network to converge to the correct posterior if they model converges. This does not happen in the standard ELBO because the expectation is only over the training data. They authors bootstrap their own model distribution by effectively adding an expectation over model simulated data to the ELBO. Under the very strong assumption that the model converges to the true distribution, you will me minimizing the KL between all true posteriors and q (which isn't amortization anymore), and with a flexible enough q parametrization, you will satisfy your requirement. \n\nThe model converging to the true distribution is a very strong assumption. Not much investigation is presented for whether this is a reasonable assumption (other than the fact that the inference network dominated the error in a toy example) or issues that occur when this assumption doesn't hold, which would be almost all of the time. Limitations are not discussed. A modification (half-asleep) is presented to address this issue.\n\nThe experimental section is not very strong, only applying the method to MNIST, CIFAR10 and a single compression experiment. Standard regularization techniques were not compared against. In fact, no other training methods other than those introduced in the paper were examined, including IWAE mentioned in the related work. \n\nI think it's a good idea to examine this particular aspect of VAE overfiiting but I recommend major revisions before conference acceptance.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept