Generalization Gap in Amortized Inference
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 modelthe 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
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 X train = { x 1 , · · · , x N } ∼ p d ( x ) , the goal of probabilistic modelling is to approximate p d ( x ) with a model p θ ( x ) .
A principled method to learn θ is to minimize the Kullback-Leibler (KL) divergence
[Equation]
where we use ⟨·⟩ to denote integration: ⟨ f ( x ) ⟩ p ( x ) ≡ (cid:82) f ( x ) p ( x ) dx .
The first term represents the negative entropy of the data distribution − H ( p d ) ≡ ⟨ log p d ( x ) ⟩ p d ( x ) , which is a constant.
The second cross entropy term involves the integration over the unkown data distribution p d ( x ) , which can be approximated by the Monte-Carlo approximation using the training dataset X train
[Equation]
Therefore, estimating θ by minimizing the KL divergence is equivalent to Maximum Likelihood
Estimation (MLE) when N → ∞ .
For a finite dataset, a common concern in both supervised and unsupervised learning is that the probabilistic model may overfit to the training dataset X train , degrading generalization performance [32].
The generalization performance in the unsupervised setting can be measured by the test likelihood [49]: 1 M (cid:80)
Mn =1 log p θ ( x ′ m ) , where X test = { x ′ 1 , . . .
, x ′ M } ∼ p
d ( x ) is the test dataset.
A model that has overfit to the training dataset X train generally results in a high training likelihood but a low test 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 is because for common tasks, like sample generation or representation learning, good generalization in terms of the test likelihood is not a sufficient measure of performance.
For example implicit models can generate sharp samples without having a likelihood function [16, 2, 46] and representations learned by latent variable models can be arbitrarily transformed without changing the likelihood [25].
However, in recent applications that use deep generative models for lossless compression [38, 39, 22, 49, 47], generalization in terms of the test likelihood directly indicates higher compression rate [49].
Specifically, given a probabilistic model p θ ( x ) , a lossless compressor can be constructed to compress a test data point x ′ to a bit string with length approximately equal to − log 2 p θ ( x ′ ) .
When p θ ( x ) → p d ( x ) , the average compression length attains the entropy of the data distribution − 1 M (cid:80) Mm =1 log 2 p θ ( x ′ m )
→ H ( p d ) , which is optimal under Shannon’s source coding theorem [34], see Appendix E for a detailed introduction.
Therefore, a better test likelihood can lead to a 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 p θ ( x ) =
(cid:82) p θ ( x | z ) p ( z ) dz .
For a nonlinear parameterization of p θ ( x | z ) (e.g. a deep neural network), the evaluation of log p θ ( x ) involves solving an intractable integration over z .
In this case, the evidence lower bound (ELBO) can be used to side-step the intractability
[Equation] [Equation]
where q ϕ ( z | x ) is a variational posterior parameterized by a neural network with parameter ϕ .
The use of an approximate posterior of the form q ϕ ( z | x ) is called amortized inference .
To better understand this objective, we can rewrite the expected ELBO as the following
[Equation]
We denote the posterior family of q ϕ ( z | x ) as Q , which is indexed by a finite dimensional θ [43].
If Q is flexible enough such that the true posterior p θ ( z | x ) ∈ Q , where p θ ( z | x ) ∝ p θ ( x | z ) p ( z ) , then in the optimum of Equation 4, we have KL( q ϕ ( z | x ) || p
θ ( z | x ))
=
⇒ q ϕ ( z | x ) =
p θ ( z | x ) for x ∼ p
d ( x ) and the ELBO will be equal to the log-likelihood ELBO( x, θ, ϕ ) = log p θ ( x ) [21, 6].
Many methods have been developed to increase the flexibility of Q , e.g. adding auxiliary variables [1, 26] or flow-based methods [9, 28], to obtain a tighter ELBO.
Recent works [38, 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 M (cid:80)
Mm =1 ELBO( x ′ m , θ, ϕ ) (also see Appendix E).
Hence the better the test ELBO indicates the better the compression performance.
This motivates us to study the factors that affect the generalization of VAEs and find practical ways to improve the generalization of VAEs.
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.
Generalization of VAEs
During training, we only have access to a finite dataset X train , which leads to the following MonteCarlo approximation as our objective to train VAEs:
[Equation]
This empirical approximation will lead to the VAE overfits to the training data for finite N .
For example, we train a
VAE on the Binary MNIST dataset for 1k epochs and plot the Bits-per-dimension (BPD) 1 of both training and testing dataset for every 100 epochs, also see Section 4 for model and training details.
Figure 1 visualizes the training and testing BPD, which shows the VAE model is overfitting to the training dataset.
The decomposition in Equation 5 suggests that the empirical ELBO contains 1) a model empirical approximation :
[Equation]
which will potentially make a flexible model p θ ( x ) overfit to the training data; and 2) an amortized inference empirical approximation :
[Equation]
where similarly a flexible q ϕ ( z | x ) can also overfit to the training data.
More specifically, we let ˆ ϕ
be the 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 2 .
When q ˆ ϕ ( z | x n ) overfits to X train , q ˆ ϕ ( z | x ′ m ) may not be a good approximation to the true posterior p θ ( z | x ′ m ) for test data x ′ m ∈ X test , We refer to the difference between the ELBO evaluated using q ˆ ϕ ( z | x ) and the ELBO evaluated using q ∗ ( z | x ) as the amortized inference generalization 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 by amortization is the cause of this inference generalization gap.
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 would not explicitly encourage better generalization performance on test data [35].
To summarize, the generalization performance of a VAE depends on two factors:
• Generative model generalization gap: defined as KL( p d ( x ) || p
θ ( x ))
and is caused by the generative model overfitting to the the training data.
• Amortized inference generalization gap: defined in Equation 11 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 a constant)
KL( p d ( x ) || p
θ ( x ))
≈ − 1 M (cid:80)
M m =1 log p θ ( x ′ m )
+ const.
cannot be calculated explicitly since we can only evaluate the lower bound − 1 M (cid:80)
Mm =1 ELBO( x ′ m , θ, ϕ ) .
Fortunately, as suggested in Equation 4, if we know the optimal posterior for the test data q ∗ ( z | x ′ m ) ≡ arg min q
[Equation]
q
[Equation]
| x
[Equation]
|| p θ
[Equation]
z | x ′ m )) , the log-likelihood can be approximated by the lower bound log p θ ( x ′ m )
≈ ELBO( x ′ m , θ, ϕ ) and the approximation becomes an equality when p θ ( z | x ′ m ) ∈ Q .
Similarly, the amortized inference generalization gap can be estimated by knowing the optimal posterior q ∗ ( z | x ′ m ) for the test dataset:
M
M (cid:88)
[Equation]
We can then estimate q ∗ ( z | x ′ m ) by fixing θ (which is trained on the training dataset) and learning ϕ ∗ on the test dataset and assuming q ∗ ( z | x
[Equation]
m )
=
q ϕ ∗ ( z | x ′ m ) , where
[Equation]
[Equation]
This optimal inference strategy can eliminate the effect of the inference generalization gap, allowing us to isolate the degree to which both the generative model and amortized inference generalization gaps are contributing to the overfitting.
[Figure]
We take the VAE described in Section 2 and train q ϕ ( z | x ) for 1k epochs with Adam [20] and lr = 5 × 10 − 4 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
M (cid:80)
Mn =1 log p θ ( x ′ ( m ) ) can be effectively approximated by the ELBO (see Appendix A for an empirical verification of the tightness of the ELBO), then the difference between the two 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 a marginal increase during training.
This suggests the generative model (decoder) slightly overfit to the data but that the overfitting is mainly dominated by the overfitting of the amortized inference network.
Although the optimal inference strategy can help eliminate the inference generalization gap, training q ϕ 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 at training time.
Improving Generalization with Consistent Amortized Inference
We now propose an inference consistency requirement which, if satisfied, would result in optimal generalization performance for amortized variational inference.
Specifically when p θ → p d , the amortized posterior should converge to the true posterior q ϕ ( z | x ) →
p θ ( z | x ) 3 for every x ∼ p
d ( 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 ELBO training objective
[Equation]
When the model converges to the true distribution p θ ∗ =
p d the training criterion for q ϕ ( z | x )
[Equation]
[Equation]
can still result in the amortized posterior q ϕ ( z | x ) 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 [13, 18], which does in fact satisfy 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-sleep training [13, 18] 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 the model 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 ϕ ( z | x ′ m ) ⟩ 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 is referred to as the sleep phase because the samples from the model used to train q ϕ are interpreted as dreamed experience.
In contrast, the training criterion for the typical VAE amortized inference (Equation 8) uses the true data samples from p d to train q ϕ ( z | x ) , which we refer to as wake phase amortized inference .
We notice that if a perfect model p θ ∗ ( x ) =
p d ( x ) is used in the sleep phase amortized inference, then it is equivalent to minimizing
[Equation]
Therefore, the training of the inference network satisfies the inference consistency requirement since we can access infinite training data from p d by sampling from p θ ∗ .
However, the wake-sleep algorithm presented lacks convergence guarantees [13] 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 [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 to use the reverse KL divergence in the sleep phase.
We fix θ and train ϕ using
[Equation]
where the integration ⟨·⟩ p θ ( x ) is approximated by Monte-Carlo using samples from the generative model p θ ( x ) .
This reverse KL objective encourages improvements to the ELBO.
When we have a perfect 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 the true data generation distribution.
Therefore, we first train a VAE to fit the binary MNIST problem.
The VAE has the same structure as that used in Section 2 and is trained using Adam with lr = 1 × 10 − 3 for 100 epochs.
After training, we treat the pre-trained decoder p θ ′
( x | z ) as the training data generator p d ( x ) ≡ (cid:82) p θ ′ ( x | z ) p ( z ) dz .
We then 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 sleep inference.
The network is trained using Adam with lr = 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 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 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 well trained model p θ may not always be sufficiently like the samples from the true data distribution.
This can lead to 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
[Equation]
where m
[Equation]
When α = 0 , it reduces to the standard approach used in VAE training.
When α = 1 , we recover the reverse sleep method (Equation 19).
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 to this method as reverse half-asleep since it uses both data and model samples to train the amortized posterior.
Intuitively, we can rewrite the Equation 21 as a sum of two positive terms
[Equation] [Equation]
Therefore, the optimal of this objective will make the first term 0, which is the same requirement as the classic amortized inference (Equation 8).
The second term, which is equivalent to the
We find the Reverse Half-asleep method ( α = 0 . 5 ) achieves the best BPD.
The mean and std are calculated with three random seeds.
reverse sleep amortized inference (Equation 19), can encourage the inference consistency requirement: when p
[Equation]
p
[Equation]
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, the second 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 [32].
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 17) and three different α for additional 100 epochs using Adam with learning rate 3 × 10 − 4 .
Figure 4 shows the test BPD comparison.
We find 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 training curve, 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 [24] and CIFAR10 [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 ( 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 for 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 q ϕ ( z | x ) on the test data and 3) run reverse half-asleep method before calculating the test BPD.
For the reverse half-asleep, we train the amortized posterior for 100 epochs with Adam and lr = 5 × 10 − 4 .
To sample from p θ ( x ) , we firstly sample z ′ ∼ p ( z ) and sample x ′ ∼ p ( x | z = z ′ ) .
For the optimal inference strategy, we train the amortized posterior with the same optimization scheme on the test data set for additionalepochs 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 5b 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 CIFAR10, we use the convolutional ResNet [17, 40] with 2 residual blocks and latent size 128.
The observational distribution is a discretized logistic distribution with linear autoregressive
[Figure]
parameterization within channels.
We train the VAE for 500 epochs with Adam and lr = 5 × 10 −and 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 [40]: 4.51 BPD with a 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 p θ → p d , the aggregate posterior q ϕ ( 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 [51, 12].
In this case, the sample x ′ that is generated 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 [12] or a PixelCNN [41, 40] to approximate q ϕ ( z ) .
In our case, we instead directly sample from q ϕ ( z ) rather than p ( z ) to generate samples in Equation 19, which can be done by first sampling x ′ ∼ p
d ( x ) (from the training dataset) and then sample z ′ ∼ q ϕ
( z | x = x ′ ) .
This scheme still results in 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 training data and train the optimal inference strategy for 500 epochs on the test data, see Figure 5c for the result.
We find the proposed reverse half-asleep training approach (with sampling from q ϕ ( z ) ) consistently improves the generalization performance of the amortized posterior.
We also apply the proposed method on a VAE trained on CIFAR100 for 500 epochs (the rest of the experiment settings are the same as the CIFAR10 case) and find our method improves the BPD from 5.288 to 5.275.
4.1 Comparisons with Regularization Methods
Recent work [35] proposed to alleviate overfitting of amortized inference by optimizing a linear combination between the traditional amortized inference (Equation 8) and a denoising objective
[Equation]
where p ( ϵ ) =
N (0 , σ 2 I ) .
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 α = 0 .
5 as that used in Equation 21 and vary σ ∈ { 0 . 1 , 0 . 2 , 0 . 4 , 0 . 6 , 0 . 8 , 1 . 0 } , see Table 1 for the comparisons.
For MNIST, we find σ ∈ { 0 . 1 , 0 . 2 , 0 . 4 } improves the generalization but larger noise levels hurts the performance.
For CIFAR10, only σ = 0 .
1 can slightly improve the generalization by 0.001 BPD.
In contrast, our method consistently achieves better generalization performance without tuning any hyper-parameters, 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 is more computational expansive during training.
Since the 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 4, higher ELBO with fixed θ indicates KL( q ϕ ( z | x ) ||
p
θ ( z | x )) 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-Encoder [7] and also improve the representation learning performance for down-stream classification tasks.
[Table]
[Figure]
[Table]
[Figure]
[Table]
Application of Lossless Compression
Lossless compression is an important application of VAEs where generalization plays a key role in the compression rate.
Given a VAE with p θ ( x | z ) , q ϕ ( z | x ) and p ( z ) , a practical compressor can be efficiently implemented using the Bits Back algorithm [19, 38] with the ANS coder [14].
See
Appendix E for a detailed introduction of conducting lossless compression with VAE models.
In Algorithm 1, we summarize the Bits Back procedure with amortized inference to compress/decompress a test data point x ′ to a stack that contains bit string messages.
The resulting code length for data x ′ is approximately equal to the negative ELBO
[Equation]
Comp./decomp.
stages share { p θ ( x | z ) , q ϕ ( z | x ) , p ( z ) } .
Compression
Draw sample z ′ ∼ q ϕ
( z | x ′ ) from the stack.
Encode x ′ ∼ p θ ( x | z
′ ) onto the stack.
Encode z ′ ∼ p
( z ) onto the stack.
Decompression Decode z ′ ∼ p ( z ) from the stack.
Decode x ′ ∼ p θ ( x | z ′ ) from the stack.
Encode z ′ ∼ q ϕ
( z | x ′ ) onto the stack.
We have shown that q ϕ ( z | x ) may overfit to 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
[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 decompression stage, we observe that the compressed data x ′ is recovered before the q ϕ ( z | x ′ ) is used to encode z ′ .
Therefore, we can also train the q ϕ ( z | x ′ ) using the recovered x ′ to maximize the test ELBO.
If the optimization procedure is the same as that used in the compression stage, we will get the same q ϕ ∗ ( 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 is approximately
[Equation]
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, we can trade off between the speed and the compression rate.
This is valuable for practical applications with different speed/rate requirements.
See Algorithm 2 for a summary of the Bits Back algorithm with K -step optimal inference.
Comp./decomp.
stages share { p θ ( x | z ) , q ϕ ( z | x ) , p ( z ) } and the optimization procedure of Equation 25.
Compression
Take K gradient steps ϕ → ϕ K with Equation 25.
Draw sample z ′ ∼ q ϕ K ( z | x ′ ) from the stack.
Encode x ′ ∼ p θ ( x | z ′ ) onto the stack.
Encode z ′ ∼ p ( z ) onto the stack.
Decompression
Decode z ′ ∼ p ( z ) from the stack.
Decode x ′ ∼ p θ ( x | z ′ ) from the stack.
Take K gradient steps ϕ → ϕ K with Equation 25.
Encode z ′ ∼ q ϕ K ( z | x ′ ) onto the stack.
Although the optimal inference strategy can be used in lossless compression, it requires extra run-time for training at the compression stages.
In contrast, our proposed reverse half-asleep inference scheme can improve the compression rate without scarifying any speed.
Additionally, our method can also provide a better initialization for the optimal-inference strategy to allow a better trade-off between compression rate and speed.
We implement 4 Bits Back with ANS [14] and compare the compression among four inference methods:
Baseline: This is the classic VAE-based compression
introduced by
For
binary 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 CIFAR10, 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] as the observation distribution.
We train both the amortized posterior and the decoder by maximizing the ELBO (Equation 3) using Adam with lr = 3 × 10 − 4 for 100, 100 and 500 epochs (for Binary MNIST, Grey MNIST and CIFAR10 respectively), and then apply Algorithm 1 to conduct
compression.
Reversed Half-asleep : we do amortized inference using Equation 21 for 100 and 300 epochs with Adam optimizer ( lr = 3 × 10 − 4 ) for binary and grey MNIST respectively, and lr = 1 ×
10 − 5 forepochs for CIFAR10.
Other training details are the same as the baseline method.
Optimal Inference:
we
take the amortized posterior (encoder) and decoder from the baseline and
apply
the K -step optimal inference strategy described in Algorithm
to
do compression
We use Adam
optimizer
and vary the K from
to 10 to
achieve a
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 : lr = 5 × 10 − 3 for binary and grey MNIST and lr = 1 × 10 − 3 for CIFAR10.
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.
[Figure]
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
[Figure]
, which will cost an additional reverse-half-asleep method also achieves lower BPD, which leads to a better trade-off between the time and compression rate.
Table 8 also reports the average time improvements of our method to compress a single MNIST and CIFAR10 image respectively, which shows the effectiveness of our method.
Related Work
A 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 a different perspective and more relevant for the application of lossless compression.
Recent work [49] first studies the likelihood-based generalization for lossless compression.
They focus on the test and train data that are from different distributions whereas we assume they follow the same distribution.
Additionally, their model has a tractable likelihood and relates to the generative model related generalization , whereas we focus on inference related generalization in VAEs.
Previous work [11] studied the amortization gap in amortized inference, which is caused by using q ϕ ∗ ( z | x n ) to generate posteriors for each input x n rather than learning a posterior q ∗ n ( z ) for x n individually.
This gap can be alleviated using a 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 at 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 to compress/decompress multiple latent samples, which requires extra time cost.
On the other hand, we focus on improving the ELBO-based compression that only needs to compress one single latent sample.
Nevertheless, similar to the K -step optimal inference strategy, our amortized training objective can also be used in the IWAE-based method, which gives a better proposal distribution for importance sampling, see
Appendix A for a demonstration.
Paper [8] considers the following data generation procedure x 1 ∼ p
d ( x ) , z 1 ∼ p θ ( z | x 1 ) , x 2 ∼ p θ ( x | z 1 ) and propose to enforce latent consistency between q ϕ ( z | x 1 ) and q ϕ ( z | x 2 ) for paired data ( x 1 , x 2 ) to 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 ′ ∼ (cid:82) p θ ( x | z ) p ( z ) dz to have high ELBO under the model (Equation 19) to improve the generalization of the amortized inference and no paired data is required in our procedure.
Therefore, both motivations and methodologies are different from our method.
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.