Generalization Gap in Amortized Inference
Anonymous Author(s)
AffiliationAddress email
Abstract
The ability of likelihood-based probabilistic models to generalize to unseen datais central to many machine learning applications such as lossless compression.
Inthis work, we study the generalizations of a popular class of probabilistic models-
the Variational Auto-Encoder (VAE).
We point out the two generalization gapsthat can affect the generalization ability of VAEs and show that the over-fittingphenomenon is usually dominated by the amortized inference network.
Based onthis observation we propose a new training objective, inspired by the classic wakesleep algorithm, to improve the generalizations properties of amortized inference.
We also demonstrate how it can improve generalization performance in the contextof image modeling and lossless compression
Introduction
Probabilistic models have achieved great success in modelling real-life data.
Given a set of trainingdata that are sampled from an underlying data distribution X train = { x 1 , . . .
, x N } ∼ p
d ( x ) , thegoal of probabilistic modelling is to approximate p d ( x ) with a model p θ ( x ) .
The parameter θ isusually learned by Maximum Likelihood Estimation (MLE): θ ∗ = arg max θ 1 N (cid:80) Nn =1 log p θ ( x n ) .
As is a common concern in supervised learning, a flexible probabilistic model may also overfit to afinite training dataset X train degrading generalization performance.
The generalization performanceof probabilistic models can be measured by the test likelihood [36]: 1 M (cid:80)
Mn =1 log p θ ( x ′ m ) , where
X test = { x ′ 1 , . . .
, x ′ M } ∼ p
d ( x ) is the test dataset.
Although the test likelihood is a common evaluation criterion for probabilistic models [29], thefactors that affect the generalization of deep generative models are less studied.
We posit that this isbecause for common tasks, like sample generation or representation learning, good generalization interms of the test likelihood is not a sufficient measure of performance.
For example implicit modelscan generate sharp samples without having a likelihood function [10, 2, 35] and representationslearned by latent variable models can be arbitrarily transformed without changing the likelihood [19].
However, in recent applications that use deep generative models for lossless compression [30, 31,16], generalization in terms of the test likelihood directly indicates higher compression rate [36].
Specifically, given a probabilistic model p θ ( x ) , a lossless compressor can be constructed to compressa test data x ′ to a bit string with length approximately equal to − log 2 p θ ( x ′ )
(see [21] for a detailedintroduction).
For a set of test data X test = { x ′ 1 , . . .
, x ′ M } ∼ p
d ( x ) with M → ∞ , when p θ ( x ) →p d ( x ) , the average compression length attains data entropy − 1 M (cid:80)
Mm =1 log 2 p θ ( x ′ m )
→ H( p d ( x )) ,which is optimal under Shannon’s source coding theorem [27].
Therefore, a better test likelihood canlead to a greater saving in bits and so understanding and improving generalization of deep generativemodels is an important challenge.
A popular type of probabilistic model is the latent variable model p θ ( x ) =
(cid:82) p θ ( x | z ) p ( z ) dz , where zis an unobserved latent variable [15, 23].
For a non-linear parameterization of p θ ( x | z ) , the evaluationof log p θ ( x ) is usually intractable and variational inference technique are used to form the evidence lower bound (ELBO) to the log-likelihood log p θ ( x ) .
The non-linear latent variable model withamortized inference is also referred to as the Variational Auto-Encoder (VAE) [15, 23].
In this case,the generalization performance of a VAE not only depends on the generative model p θ ( x | z ) but isalso affected by the variational inference procedure.
We are thus interested in studying the factorsthat affect the generalization of VAEs and how to improve the generalization in practice.
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 the overfiting of the VAE is mainlydominated by the overfiting of the amortized inference network.
• We propose a new plug-in amortized inference scheme which can improve the generalizationwithout changing the network structure.
We further demonstrate how the proposed methodcan improve the compression rate in the practical lossless compression system.
In the following sections, we first introduce the VAE model and then discuss the factors that affect itsgeneralization.
Generalizations of VAEs
VAE [15, 23] assumes a latent variable model p θ ( x ) =
(cid:82) p θ ( x | z ) p ( z ) dz , where the integration overz is intractable for non-linear p θ ( x
| z ) .
The evidence lower bound (ELBO) can be used
[Equation]
(1) where q ϕ ( z | x ) is a variational posterior parameterized by a neural network with parameter ϕ .
be found in Section 4.
The use of an approximate posterior of the form q ϕ ( z | x )is called amortized inference .
We denote the posteriorfamily of q ϕ ( z | x ) as Q : Q = { q : q ϕ ( z | x ) } for a givenx .
If Q is flexible enough such that p θ ( z | x ) ∈ Q thenat the optimum of Equation [-2, wehave KL( q ϕ ( z | x ) || p θ ( z | x )) = 
 ⇒-]1 q ϕ ( z | x ) = p θ ( z | x ) for every x ∈ p
d ( x ) and the inequality in Equation 1 becomesan equality.
Many methods have been developed to increase the flexibility of Q . For example, adding auxiliaryvariables [1, 20], or flow-based methods [5, 22].
In practice we only access a finite training dataset X train ,the training objective is 1 N (cid:80) Nn =1 ELBO( x n , θ, ϕ ) , whichcan potentially cause the VAE to over-fit to the trainingdataset.
For example, Figure 1 shows the training andtesting curve of a VAE that is trained on the Binary MNISTdataset.
We plot the Bits-per-dimension (BPD), whichis the negative log 2 likelihood normalized by the datadimension (lower BPD indicates higher likelihood).
We found the model has lower BPD (higherlikelihood) on the training dataset but higher BPD on the test dataset.
This indicates the model isoverfitting to the training data.
Recent work [30, 31, 16] has successfully applied VAE style models to lossless compression realizingimpressive performance.
In this setting the average compression length on the test data set isapproximately equal to − 1 M (cid:80) Mm =1 ELBO( x ′ m , θ, ϕ ) .
Hence the better the test ELBO, the better thecompression performance.
This motivates finding practical ways to improve the generalization of VAEs.
In the following sections we discuss two factors that will potentially affect the generalizationsof VAEs and discuss how to alleviate this overfitting phenomenon.
2.1 Two Factors Affect the Generalization of VAEs
To better understand what affects the generalization of VAEs we re-write the expected ELBO as
[Equation]
where p θ ( z | x ) ∝ p θ ( x | z ) p ( z ) is the true model posterior and the data entropy H( p d ) = −⟨ log p d ( x ) ⟩ p d ( x ) is a constant.
Given finite training dataset X train , the training objective 2 willcontain 1) a model empirical approximation :
[Equation]
which will potentially make a flexible model p θ ( x ) over-fit to the training data; and 2) an amortizedinference empirical approximation :
[Equation]
where similarly a flexible q ϕ ( z | x ) can also over-fit to the training data.
More specifically, we let ˆ ϕ bethe optimal parameter of the empirical variational inference objective
[Equation]
and we assume for any training data point x n ∈ X
[Equation]
[Equation]
where q ∗ ( z | x n ) is the realizable optimal posterior (in the Q family) for x n 1 .
However, for x ′ m ∈ X test ,there is no guarantee that q ˆ ϕ
( z | x ′ m ) is a good approximation to the true posterior p θ ( z | x ′ m ) .
We thusrefer to the difference between the ELBOs evaluated using q ˆ ϕ ( z | x ) and q ∗ ( z | x ) as the amortizedinference generalization gap , which is formally defined as(cid:68)
[Equation]
Equivalently, this gap can be written as the difference between two ELBOs with two different q
[Equation]
The inference neural network introduced with amortization is the cause of this inference generalizationgap.
It is important to emphasize that this gap cannot be reduced by simply using a more flexible Q .
This would only make KL( q ϕ ( z | x n ) ||
p
θ ( z | x n )) smaller for the training data x n ∈ X train but wouldnot explicitly encourage better generalization performance on test data [28].
To summarize, the generalization performance of VAE models depends on two factors:
• Generative model generalization gap: defined as KL( p d ( x ) || p
θ ( x ))
and is caused by the
overfitting of the generative model (decoder).
• Amortized inference generalization gap: defined in Equation 7 and is caused by the overfitting of the amortized inference (encoder).
We study the contributions of both gaps to the overfitting phenomenon of VAE models.
2.2 Impact of the Generalization Gaps
The generative model generalization gap that is estimated by the test dataset KL( p d ( x ) || p θ ( x ))
≈
Mn = log p θ ( x ′ m ) + const. , cannot be calculated explicitly since log p θ ( x ′ m ) is intractable in the
VAE case.
However, the amortized inference generalization gap (Equation 7) can be estimated byknowing the test dataset X test .
Specifically, we approximate Equation 7 using test dataset
[Equation]
where q ∗ ( z | x ′ m ) ≡ arg min q ∈Q KL( q ( z | x ′ m ) ||
p
θ ( z | x ′ m )) .
Therefore, we can fix θ and minimize KL( q ϕ ( z | x ′ m ) || p
θ ( z | x ′ m )) with respect to ϕ on test data x ′ m , which is equivalent to [Equation]
Updating ϕ using this objective would result in an optimal inference strategy.
We can use this strategyto eliminate the effect of the inference generalization gap, allowing us to isolate the degree to which both the generative model and amortized inference generalization gap are contributing to overfitting
[Figure]
To visualize the amortized generalization gap, we takethe VAE that is described in Figure 2 and only train theq ϕ ( z | x ) for 100 additional epochs (with Adam [14] andlr = 5 × 10 − 4 ) on the test data using Equation 9 to obtaintest BPD for the optimal inference.
In Figure 2 we plot thetest BPD using the optimal inference and classic amortizedinference.
The difference between two inference curvesis the amortized inference generalization gap .
We cansee that after eliminating the inference generalization gap,the test BPD is stable during training, which indicates theoverfitting phenomenon of this VAE is dominated by theoverfitting of the amortized inference network.
Althoughthe optimal inference strategy can help eliminate the inference generalization gap, training q ϕ at test time is notpossible in most applications of interest.
Therefore, we focus on how to improve the generalizationof amortized inference at training time.
In the next section, we introduce a consistency criterion thatallows us to further understand the impact of amortized inference on generalization.
Improving Generalization with Consistent Amortized Inference
We now propose an inference consistency requirement , which if satisfied, would result in optimalgeneralization for amortized variational inference, i.e. when p θ → p d , the amortized posterior shouldconverge to the true posterior q ϕ ( z | x ) → p θ ( z | x ) 2 for every x ∼ p
d ( x ) .
Although this requirementseems natural for variational inference, the classic amortized inference training that is widely usedfor VAEs [15] doesn’t satisfy it.
Recall the typical VAE empirical ELBO training objective
[Equation]
When the model converges to the true distribution p θ ∗ =
p d the training criterion for q ϕ ( z | x )min
[Equation]
11) can still result in the amortized posterior q ϕ ( z | x ) over-fitting to the training data.
In principle, onecould also limit the network capacity and/or add an explicit regularizer to the parameters [26] in anattempt to improve the generalization.
However, this still cannot satisfy the consistency requirementin principle since it doesn’t change the underlying training objective.
Alternatively, there is anotherclassic variational inference method, the wake-sleep training algorithm [8, 12], which satisfies theproposed consistency requirement, that we now discuss.
3.1 Wake-Sleep Training
Defining q ϕ ( x, z )
=
q ϕ ( z | x ) p d ( x ) and p θ ( x, z ) =
p θ ( x | z ) p ( z ) , the two phases of the wake-sleeptraining [8, 12] can be written as minimizing two different KL divergences in both x and z space.
Wake phase model learning: p θ ( x | z ) is trained by minimizing the KL divergence
[Equation]
where ⟨·⟩ p d ( x ) is approximated using the training set.
This is referred to as the wake phase since themodel is trained on experience from the ‘real environment’, i.e. it uses true data samples from p d ( x ) .
Sleep phase amortized inference: q ϕ ( z | x ) is trained by minimizing the KL divergence
[Equation]
Leaving out the terms that are irrelevant to ϕ , the objective can be estimated with Mon-Carlo⟨− log q ϕ ( z | x ) ⟩ p θ ( x,z ) ≈ 1 K (cid:80)
Kk =1 − log q ϕ
( z k | x k ) , where z k ∼ p ( z ) and x k ∼ p θ ( x | z k ) .
This isreferred to as the sleep phase because the samples from the model used to train q ϕ are interpretedas dreamed experience.
In contrast, the training criterion for the typical VAE amortized inference (Equation 4) uses the true data samples to train q ϕ ( z | x ) .
So following the wake-sleep nomenclaturewe can refer to this as wake phase amortized inference .
We notice that a perfect model used in the sleep phase is equivalent to minimizing
[Equation]
Therefore, the training of the inference network satisfies the inference consistency requirement sincewe can access infinite training data from p d by sampling from p θ ∗ .
However, the wake-sleep algorithm presented lacks convergence guarantees [8] and minimizing
KL( p θ ( z | x ) || q ϕ
( z | x )) in the sleep phase doesn’t necessarily encourage an improvement to the ELBO,which directly relates to the compression rate in the lossless compression application [30].
Therefore,in the next section, we propose a new variational inference scheme: reverse sleep amortized inferenceand demonstrate how it helps improve the generalization of the inference network in practice.
3.2 Reverse Sleep Amortized Inference
We propose to use the reverse KL divergence in the sleep phase.
We thus fix θ and train ϕ usingmin ϕ (cid:10)
[Equation]
(15) where the integration ⟨·⟩ p θ ( x ) is approximated by Monte-Carlo using samples from the generativemodel p θ ( x ) .
This reverse KL objective encourages improvements to the ELBO.
When we have aperfect model p θ ∗ ( x ) = p d ( x ) the reverse sleep phase is equivalent to
[Equation]
[Equation]
which satisfies the inference consistency requirement .
The consistency requirement can also be validated empirically when the perfect model is known p θ ∗ ( x )
= p d ( x ) .
This can be achieved by using a pre-trained VAE as thetrue data generation distribution.
Therefore, we first traina VAE to fit the binary MNIST problem.
The VAE hasthe same structure as that used in Section 2 and is trainedusing Adam with lr = 1 × 10 − 3 for 100 epochs.
Aftertraining, we treat the pre-trained decoder p θ ′
( x | z ) as thetraining data generator p d ( x ) ≡ (cid:82) p θ ′ ( x | z ) p ( z ) dz .
Wethen sample 10000 data samples from p d to form a training set X train and 1000 samples to form a test set X test .
We then train a new q ϕ ( z | x ) with: 1. wake phase inference (VAE) 2.
(forward) sleep inference and 3.
reversesleep inference.
The network is trained using Adam withlr = 1 × 10 − 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 phaseand the reverse sleep inference achieves the best BPD.
Intuitively, this is because both the forwardand reverse sleep inference use the true model to generate additional training data whereas the wakeinference only has access to the finite training dataset X train .
3.3 Reverse Half-asleep Amortized Inference with Imperfect Models
In practice our model will not be perfect p θ ̸ = p d .
Empirically we find that samples from even a welltrained model p θ may not always be sufficiently like the samples from the true data distribution.
Thiscan lead to a degradation in the performance of the inference network when using the reverse-sleepapproach.
For this reason, we propose to use a mixture distribution between the model and theempirical training data distribution as follows
[Equation]
where m ( x ) ≡ αp θ ( x ) + (1 − α )
ˆ p d .
When α = 0 , it reduces to the standard approach used in VAE training.
When α = 1 , we recoverthe reverse sleep method (Equation 4).
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.
The mean and std are calculated with three random seeds.
We thus refer to this method as reverse half-asleep sinceit uses both data and model samples to train the amortizedposterior.
To compare with different α , we first fit a VAE(with the same structure as that used in Figure 2) to the
Binary MNIST dataset, and then train the amortized posterior using sleep inference (Equation 13) and three differentα for additional 100 epochs using Adam with learning rate3 × 10 − 4 .
Figure 4 shows the test BPD comparison, wherewe find the proposed reverse half-asleep method ( α = 0 .
5 )outperforms the reversed sleep method ( α = 1 ), whereasthe standard amortized inference training in VAE ( α = 0 )leads to overfitting of the inference network.
We also plotthe sleep inference as a comparison, whose BPD is lesscompetitive since it is not directly optimizing the ELBO.
Generalization Experiments
We apply the reverse half-asleep to improve the generalization of VAEs on three different datasets:binary MNIST, grey MNIST [18] and CIFAR10 [17].
For binary and grey MNIST, we use latentdimension 16/32 and neural nets with 2 layers of 500 hidden units in both the encoder and decoder.
We use Bernoulli p ( x | z ) for binary MNIST and discretized logistic distribution for grey MNIST.
We train the VAE with the usual amortized inference approach using Adam with lr = 3 × 10 − 4 for1000 epochs and save the model every 100 epochs.
We then use the saved models to 1) evaluateon the test data sets, 2) conduct optimal inference by training q ϕ ( z | x ) on the test data and 3) runreverse half-asleep method before calculating the test BPD.
For the reverse half-asleep, we trainthe amortized posterior for 100 epochs with Adam and lr = 5 × 10 − 4 .
To sample from p θ ( x ) , wefirstly sample z ′ ∼ p ( z ) and sample x ′ ∼ p ( x | z = z ′ ) .
For optimal inference, we train the amortizedposterior with the same optimization scheme on the test data set for additional 500 epochs to ensurethe same gradient steps are conducted (since training dataset is 5 times as big as the test set).
For CIFAR10, we use the convolutional ResNet [11, 32] with 2 residual blocks and latent size128.
The observational distribution is a discretized logistic distribution with linear autoregressiveparameterization within channels.
We train the VAE for 500 epochs with Adam and lr = 5 × 10 − 4and save the model every 100 epoch.
The pre-trained VAE achieves 4.592 BPD on the CIFAR10,which is comparable with other single latent VAE models reported in [32]: 4.51 BPD with a VAEwith latent dimension 256 and 4.67 BPD with a discrete latent VAE (VQVAE).
Ideally, when the VAEmodel converges to the true distribution p θ → p d , the aggregate posterior q ϕ ( z ) = (cid:82) q ϕ
( z | x ) p d ( x ) dxwill match the prior p ( z ) .
However, for a complex distribution like CIFAR10, a significant mismatchbetween q ϕ ( z ) and p ( z ) is usually observed in practice [38, 7].
In this case, the sample x ′ that isgenerated using a latent sample from the prior x ′ ∼ p ( x | z ′ ) where z ′ ∼ p ( z ) may be blurry or invalid.
A common solution is to train another model, e.g. a VAE [7] or a PixelCNN [33, 32] to approximateq ϕ ( x ) .
In our case, we instead directly sample from q ϕ ( z ) instead of p ( z ) to generate samples in
Equation .
This scheme still results in a consistent training objective since q ϕ ∗ ( z ) =
p ( z ) for optimalposterior q ϕ ∗ ( z | x ) .
We use Adam with lr = 1 × 10 − 5 and train the reverse half-asleep inference for100 epochs on the training data and train the optimal inference strategy for 500 epochs on the testdata.
Figure 5 show the test BPD comparison for three datasets.
We find that the reverse half-asleeptraining approach consistently improves the generalization performance of our models as measuredby the negative test ELBO.
In contrast to the optimal inference strategy, our approach does not requirefurther training on test data to improve generalization performance.
Paper [28] proposed to alleviate the overfitting of the amortized inference by optimizing a linearcombination between the traditional amortized inference (Equation 4) and a denoising objective
[Equation]
where p ( ϵ ) = N(0 , σ 2 I ) .
We compare this regularizer to our method by training the amortizedposterior of VAEs for additional 100/300 and 100 epochs on Binary/Grey MNSIT and CIFARrespectively.
For denoising regularizer, we use the same linear combination weight α = 0 .
5 as thatused in Equation 17 and vary σ ∈ { 0 . 1 , 0 . 2 , 0 .
See Table 1 for the comparison results.
We find for MNIST, σ ∈ { 0 . 1 , 0 . 2 , 0 . 4 } can improve the generalization but larger level noise hurts [Figure] the inference performance.
For CIFAR10, only σ = 0 .
1 can slightly improve the {+generalization by0.
+} BPD.
In contrast, our method can achieve better generalization performance without tuningany hyper-parameters.
Our method requires sampling from the model, which does require additionalcomputations, but is negligible comparing to the whole training procedure.
Figure 6 reports the test
BPD curves evaluated after every training epoch (the mean and std are calculated over 3 randomseeds).
We find our method consistently improves the model generalization during training.
[Table]
[Figure]
Application of Lossless Compression
Lossless compression is an important application of VAEs where generalization plays a key role inthe compression rate.
Given a VAE with decoder p θ ( x | z ) , amortized posterior q ϕ ( z | x ) and prior p ( z ) ,a practical compressor can be efficiently implemented using the Bits Back algorithm [13, 30] withthe ANS coder [9].
In Algorithm 1, we describe the Bits Back procedure with amortized inferenceto compress/decompress a test data point x ′ .
The resulting code length for data x ′ is approximatelyequal to the negative one-data ELBO
[Equation]
Algorithm
Comp. and decomp.
stages share { p θ ( x | z ) , q ϕ ( z | x ) } .
Compression
Draw sample z ′ ∼ q ϕ
( z | x ′ ) from the stack.
Encode x ′ ∼ p θ ( x | z ′ ) , z ′ ∼ p ( z ) onto the stack.
Decompression
Decode z ′ ∼ p ( z ) , x ′ ∼ p θ ( x | z ′ ) from the stack.
Encode z ′ ∼ q ϕ
( z | x ′ ) onto the stack.
We have shown that q ϕ ( z | x ) may over-fitto the training data, degrading compressionperformance.
To improve the compression
BPD, the optimal inference method canalso be applied in the Bits Back algorithm.
In the compression stage, we can train ϕ by
[Equation]
When the q ϕ ( z | x ′ ) is parameterized to be a
Gaussian, we can just take ϕ to be the mean and standard deviation N( ϕ µ , ϕ 2 σ ) , which only contains two training parameters.
In the decompressionstage, we observe that the compressed data x ′ is recovered before the q ϕ ( z | x ′ ) is used to encodez ′ .
Therefore, we can also train the q ϕ ( z | x ′ ) using the recovered x ′ to maximize the test ELBO.
Ifthe optimization procedure is the same as that used in the compression stage, we will get the sameq ϕ
∗
z | x ′ ) .
In practice, we need to pre-specify the number of gradient descent steps K .
When K islarge, we recover the optimal inference scheme and the code length is approximately
[Equation]
This observation was first proposed in [34] in the context of lossy compression and then applied tolossless compression with Bits Back coding in [24].
Furthermore, by varying the optimization steps K in the optimal inference, we can also control a trade-off between the speed and the compressionrate.
This is valuable for practical applications with different speed/rate requirements.
See Algorithm2 for a summary of the Bits Back algorithm with K -step optimal inference.
Algorithm 2 Bits
Comp./decomp.
stages share { p θ ( x | z ) , q ϕ ( z | x ) } and the optimization procedure of Equation 20.
Compression
Take K gradient steps ϕ → ϕ K with Equation 20.
Draw sample z ′ ∼ q ϕ K ( z | x ′ ) from the stack.
Encode x ′ ∼ p θ ( x | z ′ ) , z ′ ∼ p ( z ) onto the stack.
Decompression
Decode z ′ ∼ p ( z ) , x ′ ∼ p θ ( x | z ′ ) from the stack.
Take K gradient steps ϕ → ϕ K with Equation 20.
Encode z ′ ∼ q ϕ K ( z | x ′ ) onto the stack.
Although the optimal inference can be usedin the lossless compression, it requires extra run-time for training at the compression stages.
On the other hand, the proposed reverse half-asleep inference schemecan improve the compression rate withoutscarifying any speed.
Additionally, ourmethod can also provide a better initialization for the optimal-inference to allow abetter trade-off between compression rateand speed.
Therefore, we implement 3 Bits
Back with ANS [9] and compare the compression among four inference methods:
Baseline: This is the classic VAE-based compression introduced by [30].
For binary and greyMNIST, both the encoder and decoder contain 2 fully connected layers with 500 hidden units andlatent dimension 10.
The observation distributions are Bernoulli and discretized Logistic distributionrespectively.
For CIFAR10, we use fully convolutional ResNets [11] with 3 residual blocks inboth the encoder and decoder, latent dimension 128 and discreteized Logistic distribution withchannel-wise linear autoregressive[25] as the observation distribution.
We train both the amortizedposterior and the decoder by maximizing the ELBO (Equation 1) using Adam with lr = 3 × 10 − 4for 100/100/500 epochs (Binary MNIST/Grey MNIST/CIFAR10), and then apply Algorithm 1 toconduct compression.
Reversed Half-asleep : we train the amortized posterior using the proposed training criterion (Equation 17) for 100/300 epochs using the Adam optimizer with lr = 3 × 10 − 4 for binary and greyMNIST, and lr = 1 × 10 − 5 for 100 epochs for CIFAR10.
All other training details are the same asthe baseline method.
Optimal Inference: we take the amortized posterior (encoder) and decoder from the baselineand apply the K -step optimal inference described in Algorithm 2 to do compression.
We use Adamoptimizer and vary the K from 1 to 10 to achieve a trade-off curve between compression rate andspeed.
We actively choose the highest learning rate that can make the BPD consistently improve withthe increment of K : lr = 5 × 10 − 3 for binary and grey MNIST and lr = 1 × 10 − 3 for CIFAR10.4.
Reversed Half-asleep + Optimal Inference: we take the encoder in method 2 and decoder fromthe baseline and conduct K -step optimal inference.
All other training details are as per method 3.
In Figure 7, we plot test BPD comparisons for different methods.
We can see if optimization is notallowed at compression time, the use of our reverse-half-asleep method achieves better compressionrate with no additional computational cost.
If we allow K -step optimization during compression,for a given computational budget, the amortized posterior initialized using our reverse-half-asleepmethod also achieves lower BPD, which leads to a better trade-off between the time and compressionrate.
See Table 8 for the average time improvements of our method to compress a single MNIST and
CIFAR10 image respectively.
[Figure] [Figure]  Related Work A different perspective on generative models’ generalization is proposed in paper [37] where thegeneralization is evaluated by testing if the model can generate novel combinations of features.
However, the generalization defined in our paper is purely measured by the test likelihood, which is adifferent perspective and more relevant for the application of lossless compression.
Recent work [36] first studies the likelihood-based generalization for lossless compression.
Theyfocus on the test and train data that are from different distributions whereas we assume they followthe same distribution.
Additionally, their model has a tractable likelihood and relates to the generativemodel related generalization , whereas we focus on the inference related generalization in VAEs.
Previous work [6] studied the amortization gap in amortized inference, which is caused by usingq ϕ ∗ ( z | x n ) to generate posteriors for each input x n rather than learning a posterior q ∗ n ( z ) for x nindividually. This gap can be alleviated using a larger capacity encoder network. This amortizationgap is fundamentally different from the inference generalization gap we discuss in this work since thelatter focuses solely on test time generalization but the former problem also exists at training time. Paper [4] also uses samples from the model to train the encoder, which aims to make the learnedrepresentations more robust to adversarial attacks, whereas the motivation of our method is to improvegeneralization performance in terms of the test likelihood.
Although the task in their paper is different,our work can provide a different view to explain their method: good generalization of q ϕ ( z | x ) canbenefit downstream tasks like representation learning performance.
Recent work [24] proposes a compression scheme based on the IWAE [3] bound, which is tighter thanthe ELBO and thus improves the compression rate.
However, this method has to compress/decompressmultiple latent samples, which requires extra time cost.
On the other hand, we focus on improvingthe ELBO-based compression that only needs to compress one single latent sample.
Nevertheless,similar to the K -step optimal inference method, our amortized training objective can also be used inthe IWAE-based method, which gives a better proposal distribution for importance sampling.
Conclusion
We have shown how the generalization of VAEs is largely affected by the amortized inference networkand proposed a new variational inference scheme that provides better generalization as demonstratedin the application of lossless compression.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
In the end of section 4.
(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 tothem?
[Yes]
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]
If you ran 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 a 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 4 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 6.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See the footnote in page 8.
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 CIFAR10.
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to 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]