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 [42]: 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 [33], 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 [14, 2, 40] and representationslearned by latent variable models can be arbitrarily transformed without changing the likelihood [23].
However, in recent applications that use deep generative models for lossless compression [34, 35,20], generalization in terms of the test likelihood directly indicates higher compression rate [42].
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 [25] 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 [31].
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 [19, 27].
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) [19, 27].
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 a practical lossless compression system.
We first introduce VAE and then discuss the factors that affect its generalization.
Generalizations of VAEs
VAE [19, 27] 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 ϕ .
The useof an approximate posterior of the form q ϕ ( z | x ) is called amortized inference .
To better understandthis objective, we can rewrite the expected ELBO as following
[Equation] where the entropy H( p d ) ≡ −⟨ log p d ( x ) ⟩ p d ( x ) is a constant. We further denote the posterior familyof q ϕ ( z | x ) as Q , which is indexed by finite dimensional θ [38].
If Q is flexible enough such that thetrue posterior p θ ( z | x ) ∈ Q , where p θ ( z | x ) ∝ p θ ( x | z ) p ( z ) , then in the optimum of Equation {+2, wehave KL( q ϕ ( z | x ) || p θ ( z | x )) = 
 ⇒+} q ϕ ( z | x ) = p θ ( z | x ) for x ∼ p
d ( x ) and the ELBO becomes tight ELBO( x, θ, ϕ ) = log p θ ( x ) [19, 5].
Many methods have been developed to increase the flexibilityof Q , e.g. adding auxiliary variables [1, 24] or flow-based methods [8, 26].
In practice, we can only access a finite training dataset X train , the training objective is1 N (cid:80) Nn =1 ELBO( x n , θ, ϕ ) , which can potentially cause the VAE to over-fit to the training dataset. epochs and we plot the testing BPD every 100 epochs. Also see
Section 4 for model and training details.
For example, Figure 1 shows the training and testing curveof a VAE that is trained on the Binary MNIST dataset.
Weplot the Bits-per-dimension (BPD), which is the negativelog 2 likelihood normalized by the data dimension (lower BPD indicates higher likelihood).
We found the model haslower BPD (higher likelihood) on the training dataset buthigher BPD on the test dataset.
This indicates the modelis overfitting to the training data.
Recent work [34, 35, 20] has successfully applied VAEstyle models to lossless compression realizing impressive performance.
In this setting the average compression length on the test data set is approximately equal to− 1 M (cid:80) Mm =1 ELBO( x ′ m , θ, ϕ ) .
Hence the better the test ELBO, the better the compression performance.
This motivates finding practical ways to improve the generalizationof VAEs.
In the following sections we discuss two factors that will potentially affect the generalizations of VAEs and discuss how to alleviate this overfitting phenomenon.
Two Factors Affect the Generalization of VAEs
In practice, we only access to a finite dataset X train during training, which leads to an empirical ELBO approximation ⟨
[Equation]
The decomposition in
Equation 3 directly suggests that the empirical ELBO contains 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 train
[Equation]
where q ∗ ( z | x n ) is the realizable optimal posterior (in the Q family) for x n 1 .
For x ′ m ∈ X test , thereis no guarantee that q ˆ ϕ
( z | x ′ m ) is a good approximation to the true posterior p θ ( z | x ′ m ) .
We thus referto the difference between the ELBOs evaluated using q ˆ ϕ ( z | x ) and q ∗ ( z | x ) as the amortized inferencegeneralization gap , which is formally defined as
[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 [32].
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 8 and is caused by theoverfitting of the amortized inference (encoder).
2.2 Impact of the Generalization Gaps
The generative model generalization gap that is estimated by the test dataset (up to a constant)
[Equation]
., cannot be calculated explicitly since we canonly evaluate the lower bound − 1 M (cid:80)
Mm =1 ELBO( x ′ m , θ, ϕ ) .
Fortunately, as suggested in Equation 2,if we know the optimal posterior for the test data q ∗ ( z | x ′ m ) ≡ arg min q ∈Q KL( q ( z | x ′ m ) ||
p θ ( z | x ′ m )) ,the log-likelihood can be well approximated by the lower bound log p θ ( x ′ m )
≈ ELBO( x ′ m , θ, ϕ ) . This approximation becomes an equality when p θ ( z | x ′ m ) ∈ Q , see Appendix A for verification of thetightness of the bound.
Similarly, the amortized inference generalization gap can be estimated byknowing optimal posterior q ∗ ( z | x ′ m ) of the the test dataset:
[Equation]
We can then estimate q ∗ ( z | x ′ m ) by fixing θ (which is trained on the training dataset) and learningϕ ∗ = min ϕ KL( q ϕ ( z | x ′ m ) ||
p
θ ( z | x ′ m ))
= max ϕ (cid:10) log p θ ( x ′ m , z ) − log q ϕ ( z | x ′ m ) (cid:11) q ϕ ( z | x ′ m ) .
(10) on the test dataset and assume q ∗ ( z | x ′ m ) = q ϕ ∗ ( z | x ′ m ) .
This optimal inference strategy can helpeliminate the effect of the inference generalization gap, allowing us to isolate the degree to which boththe generative model and amortized inference generalization gap are contributing to the overfitting.
[Figure]
Therefore, we take the VAE described in Section 2.2 andtrain the q ϕ ( z | x ) for 100 epochs (with Adam [18] andlr = 5 × 10 − 4 ) on the test data using Equation 10 to obtaintest BPD for the optimal inference.
In Figure 2 we plot thetest BPD using the optimal inference and classic amortizedinference, where the green curve represents the approximation of the average likelihood 1 M (cid:80)
Mn =1 log p θ ( x ′ m ) andthe difference between two inference curves is the amortized inference generalization gap .
We observe that aftereliminating the inference generalization gap, the test BPDis stable with a marginal increase during training, whichsuggests the generative model (decoder) slightly overfitto the data but the overfitting phenomenon of the VAE ismainly dominated by the overfitting of the amortized inference network.
Although the optimal inference strategy can help eliminate the inference generalization gap, trainingq ϕ at test time is not practical in most applications of interest.
Therefore, we focus on improving thegeneralization of amortized inference at training time.
In the next section, we introduce a consistencycriterion that allows 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 [19] 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 )
[Equation]
(12) 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 [30] 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 that we now discuss, the wake-sleep training algorithm [12, 16],which satisfies the proposed consistency requirement.
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 [12, 16] 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 Monte-Carlo⟨− log q ϕ [Equation] z [Equation]
x,z [Equation] 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 5) 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 thesleep 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 [12] 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 [34].
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 ϕ using
[Equation]
) 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.
reverse sleepinference.
The network is trained using Adam with lr = 1 × 10 − 3 for 100 epochs.
Figure 3 showsthe test BPD calculated after every training epoch.
We can see the sleep phase out-performs the wakephase and the reverse sleep inference achieves the best BPD.
Intuitively, this is because both theforward and reverse sleep inference use the true model to generate additional training data whereasthe wake inference 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-sleep approach.
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 recover thereverse sleep method (Equation 16).
We find that a setting of α = 0 .
5 works well in practice.
Thisbalances samples from the true underlying data distribution with samples from the model.
We thus refer to this method as reverse half-asleep since it uses both data and model samples to trainthe amortized posterior.
Intuitively, we can rewrite the Equation 18 as a sum of two positive terms
[Equation]
Therefore, the optimal of this objective will make the first term 0, which is the same requirementas the classic amortized inference (Equation 5).
The second term, which is equivalent to the
The mean and std are calculated with three random seeds.
reverse sleep amortized inference (Equation 16), can encourage the inference consistency requirement: whenp θ = p d , the optimal of the second term will set q ϕ ( z | x ) = p θ ( z | x ) for any x ∼ p d ( x ) . When p θ is not perfect, thesecond term can be seen as a regularizer added to the classic amortized inference objective, which can be used to penalize the hypothesis space of the amortized network [30]to alleviate the over-fitting of the amortized inference.
To compare with different α , we first fit a VAE (with thesame structure as that used in Figure 2) to the Binary
MNIST dataset, and then train the amortized posteriorusing sleep inference (Equation 14) and three different αfor additional 100 epochs using Adam with learning rate3 × 10 − 4 .
Figure 4 shows the test BPD comparison.
Wefind the proposed reverse half-asleep method ( α = 0 . 5 )outperforms the reversed sleep method ( α = 1 ), whereas the standard amortized inference training in VAE ( α = 0 ) leads to overfitting of the inference network.
We also plot the sleep inference trainingcurve, whose BPD is less competitive 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 [22] and CIFAR10 [21].
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 number of gradient steps are conducted (since training set is 5 times as big as the test set).
For CIFAR10, we use the convolutional ResNet [15, 36] 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 [36]: 4.51 BPD with a VAEwith latent dimension 256 and 4.67 BPD with a discrete latent VAE (VQVAE).
Ideally, when the VAE model converges to the true distribution p θ → p d , the aggregate posteriorq ϕ ( z ) = (cid:82) q ϕ
( z | x ) p d ( x ) dx will match the prior p ( z ) .
However, for a complex distribution like CIFAR10, a significant mismatch between q ϕ ( z ) and p ( z ) is usually observed in practice [44, 11].
Inthis case, the sample x ′ that is generated using a latent sample from the prior x ′ ∼ p ( x | z ′ ) , where [Figure] z ′ ∼ p ( z ) , may be blurry or invalid.
A common solution is to train another model, e.g. a VAE[11] or a PixelCNN [37, 36] to approximate q ϕ ( x ) .
In our case, we instead directly sample fromq ϕ ( z ) instead of p ( z ) to generate samples in Equation 16, which can be done by first samplingx ′ ∼ p
d ( x ) (from the training dataset) and then sample z ′ ∼ q ϕ ( z | x = x ′ ) .
This scheme still resultsin a consistent training objective since q ϕ ∗ ( z ) =
p ( z ) for the optimal posterior q ϕ ∗ ( z | x ) .
We use Adam with lr = 1 × 10 − 5 and train the reverse half-asleep inference for 100 epochs on the trainingdata and train the optimal inference strategy for 500 epochs on the test data.
Figure 5 shows the test BPD comparison for three datasets.
We find that the reverse half-asleep training approach consistentlyimproves the generalization performance of our models as measured by the negative test ELBO.
Incontrast to the optimal inference strategy, our approach does not require further training on test datato improve generalization performance.
Recent work [32] proposed to alleviate the overfitting of the amortized inference by optimizing alinear combination between the traditional amortized inference (Equation 5) 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 an additional 100/300 and 100 epochs on Binary/Grey MNSIT and CIFARrespectively.
For the denoising regularizer, we use the same linear combination weight α = 0 .
5 asthat used in Equation 18 and vary σ ∈ { 0 . 1 , 0 . 2 , 0 . 4 , 0 . 6 , 0 . 8 , 1 . 0 } .
See Table 1 for the comparisonresults.
We find for MNIST, σ ∈ { 0 . 1 , 0 . 2 , 0 . 4 } can improve the generalization but larger level noisehurts the inference performance.
For CIFAR10, only σ = 0 .
1 can slightly improve the [-generalization by0.
-]generalizationby 0.001 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 compared 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 model generalization during training.
[Table]
[Figure]
We have shown the proposed inference method can consistently achieve better generalization.
Sincethe decoder is shared and fixed in all comparisons, better test ELBO indicates the predicted q ϕ ( z | x ′ ) is closer to the true posterior p θ ( z | x ′ ) under the KL divergence measure (see Equation 2, higher
ELBO with fixed θ indicates KL( q ϕ ( z | x )
|| p
θ ( z | x )) is smaller). Therefore, the proposed methodcan 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 distribution for theimportance weighted Auto-Encoder[6] and also improve the representation learning performancefor down-stream classification tasks. In the next section, we demonstrate an important application:lossless compression, which directly relates to the ELBO generalization of the VAE.
We show thatthe proposed method can give a free improvement on the compression rate without sacrificing anycomputation time, which is a crucial property for practical compression applications.
Application of Lossless Compression
Lossless compression is an important application of VAEs where generalization plays a key rolein the compression rate.
Given a VAE with p θ ( x | z ) , q ϕ ( z | x ) and p ( z ) , a practical compressor canbe efficiently implemented using the Bits Back algorithm [17, 34] with the ANS coder [13].
In Algorithm 1, we describe the Bits Back procedure with amortized inference to compress/decompressa test data point x ′ .
The resulting code length for data x ′ is approximately equal to the negative ELBOcode_len ( x ′
[Equation]
log [Equation] p [Equation]
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 meanand 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 [39] in the context of lossy compression and then appliedto lossless compression with Bits Back coding in [28].
Furthermore, by varying the optimizationsteps K in the optimal inference, we can trade off between the speed and the compression rate.
Thisis valuable for practical applications with different speed/rate requirements.
See Algorithm 2 for asummary of the Bits Back algorithm with K -step optimal inference.
Algorithm
Comp./decomp.
stages share { p θ ( x | z ) , q ϕ ( z | x ) } and the optimization procedure of Equation 22.
Compression
Take K gradient steps ϕ → ϕ K with Equation 22.
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 22.
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 [13] and compare the compression among four inference methods
Baseline: This is the classic VAE-based compression introduced by [34].
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 [15] with 3 residual blocks in theencoder/decoder, latent dimension 128 and discreteized Logistic distribution with channel-wise linearautoregressive[29] as the observation distribution.
We train both the amortized posterior and thedecoder by maximizing the ELBO (Equation 1) using Adam with lr = 3 × 10 − 4 for 100/100/500epochs (Binary MNIST/Grey MNIST/CIFAR10), and then apply Algorithm 1 to conduct compression.2.
Reversed Half-asleep : we do amortized inference using Equation 18 for 100/300 epochs withAdam optimizer ( lr = 3 × 10 − 4 ) for binary and grey MNIST, and lr = 1 × 10 − 5 for 100 epochs forCIFAR10.
Other training details are the same as the baseline method.3.
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.
[Figure]
and our R-Half-sleep can be seen as special cases of optimal inference with K = 0 .
We find given a fixed computational budget, our method achieves a lower BPD than one using traditional amortized inference training.
[Figure] [Table] 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.
See also Appendix ??
for a detailed discussions with related work.
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.