Meta-learning richer priors for VAEs
Abstract
Variational auto-encoders have proven to capture complicated data distributions and useful latent representations, while advances in meta-learning have made it possible to extract prior knowledge from data.
We incorporate these two approaches and propose a novel flexible prior, namely the Pseudo-inputs prior, to obtain a richer latent space.
We train VAEs using the Model-Agnostic Meta-Learning (MAML) algorithm and show that it achieves comparable reconstruction performance with standard training.
However, we show that this MAML-VAE model learns richer latent representations, which we evaluate in terms of unsupervised few-shot classification as a downstream task.
Moreover, we show that our proposed Pseudo-inputs prior outperforms baseline priors, including the VampPrior, in both models, while also encouraging high-level representations through its pseudo-inputs.
1. Introduction
Variational auto-encoders (VAEs) (Kingma and Welling, 2014; Rezende et al., 2014) are deep generative models that allow to learn complicated data distributions.
However, VAEs often make use of the isotropic Gaussian prior that may over-regularize the posterior (Hoffman and Johnson, 2016), failing to capture its multi-modal nature.
For this reason, several alternative priors have been proposed, such as the Mixture of Gaussians (Dilokthanakul et al., 2016; Nalisnick et al., 2016; Kopf et al., 2021) or the VampPrior (Tomczak and
Welling, 2018).
We propose the Pseudo-inputs prior to provide more flexibility than the
VampPrior while encouraging higher-level representations.
In parallel, high-level representations of the data can also be achieved through metalearning, which aims at learning to generalise across tasks so that the model can quickly adapt to novel ones.
This requires learning specific representations that avoid to re-train in the face of new data (Santoro et al., 2016).
We combine VAEs with meta-learning with the aim of generating better latent representations.
Specifically, we employ the Model-Agnostic
Meta-Learning algorithm (MAML)
(Finn et al., 2017) to train VAEs, namely MAML-VAE.
The contributions of this work are three-fold:
• We show VAEs trained with MAML achieve comparable performances with standard training on the full dataset, and better ones on Omniglot with convolutional layers.
• We propose the Pseudo-inputs prior and show it outperforms the standard prior, the Mixture of Gaussians, and the VampPrior on VAEs and on the meta-learnt MAMLVAE model, while being more flexible and capturing higher-level representations.
• We show that meta-learning allows to generate richer latent representations, in terms of higher accuracy in unsupervised few-shot classification, evaluated as a downstream task.
Again, the Pseudo-inputs prior achieves the highest accuracy.
2. Methods 2.1. MAML-VAE
We use the standard implementation of VAEs, for which the variational posterior is assumed to be a diagonal Gaussian distribution q ϕ ( z | x ) =
N ( z ; µ ϕ ( x ) , diag( σ 2 ϕ ( x ))) modelled with a neural network parametrised by ϕ .
The decoder, on the other hand, is defined as a Bernoulli distribution modelled by a neural network p θ ( x | z ) parametrised by θ .
Lastly, we indicate the prior as p λ ( z ) parametrised by λ .
The objective optimised in VAEs is the evidence lower bound (ELBO), which is computed via a Monte Carlo estimate as
[Equation]
where z ( l ) ϕ ∼ q ϕ
( z | x ) using the reparametrization trick (Kingma and Welling, 2014).
In this work, we apply meta-learning to VAE priors by leveraging the MAML algorithm, since its model-agnostic nature allows to maintain the Bayesian structure of the VAE.
In short, the idea behind the MAML algorithm is to find the optimal parameters of the model so that in a fixed (small) number of steps over the support set D supp they can be adapted to unseen tasks over the query set D query .
The sought optimal parameters should thus be sensitive enough, allowing with small changes to obtain significant improvements in the loss function.
The training and test procedures for the resulting model, which we call
MAML-VAE, are outlined in Appendix A in Algorithm 1 and 2.
In contrast to VAEs, where parameters are updated over an arbitrary number of gradient steps, the MAML-VAE model allows only a single adaptation step over the support set to adapt to the query set, which makes the task harder.
It has also been suggested that the optimal parameters found by the MAML algorithm coincide with the optimal parameters for continual learning (Gupta et al., 2020).
Intuitively, this should lead to richer high-level latent representations.
2.2. The Pseudo-inputs prior
We propose the Pseudo-inputs prior with the intention of providing a more flexible encoding than the VampPrior, while also learning a high-level representation of the dataset.
The
Pseudo-inputs prior encodes a mixture of Gaussians by mapping K learnable h -dimensional pseudo-inputs { u k } Kk =1 to the mean vectors { µ k } Kk =1 in the d -dimensional latent space through a flexible encoding network f ψ ( · ) : u k (cid:55)→ µ k .
The k -th component of the mixture is thus simply obtained by feed-forwarding the k -th pseudo input.
We verified experimentally that learning the covariance vectors separately leads to better results.
Overall, the probability density encoded by the Pseudo-inputs prior is given by
[Equation]
where λ = { ψ, { u k } Kk =1 , { σ 2 k } Kk =1 } are the parameters of the prior.
An illustration of the proposed Pseudo-inputs prior is provided in Figure 1.
This encoding allows to freely define the complexity of the network as well as the dimension of the pseudo-inputs, proportionally to how informative they are intended to be.
We
[Figure]
could not learn this high-level representation if we instead encoded the mean vectors and covariances directly.
In particular, the Pseudo-inputs prior is meant to learn a particular parametrisation of the network that allows the pseudo-inputs to adapt to different tasks, making it suitable for meta-learning as well as transfer and continual learning.
We also suggest that pseudo-inputs could be sampled to alleviate the mode-seeking behaviour of the
Kullback-Leibler divergence (Burda et al., 2016), as outlined in Appendix A.
3. Experiments 3.1. Pseudo-inputs prior: rate-distortion analysis
In order to understand the behavior of the Pseudo-inputs prior, we perform the ratedistortion analysis (Alemi et al., 2018) and compare it against the standard prior,
the
Mixture of Gaussians prior (MoG)—encoding mean and covariance vectors directly—and the VampPrior.
In order to do so, we leverage the β -VAE objective proposed in Higgins et al.
(2017), which is defined as L := RE + β KL with the distortion being the reconstruction error RE := − (cid:82) q ( x ) dx (cid:82) q ϕ
( z | x ) log p θ ( x | z ) dz and the rate being the Kullback-Leibler divergence KL :
= (cid:82) q ( x ) dx (cid:82) q ϕ ( z | x ) log q ϕ ( x | z ) p λ ( z ) dz .
In practice, we trained the VAE model with the β -VAE objective for β ∈ { 0 . 01 , 0 . 1 , 0 .
We compare the priors in terms of their ELBO (RE + KL) on Omniglot (Lake et al., 2015) and Quickdraw (Jongejan et al., 2016), which we report in Figure 2 and in Appendix B in Figure 4, respectively.
It is apparent that the least expressive prior is almost everywhere the standard Gaussian, as it achieves the highest ELBOs for most β values.
In contrast, the proposed Pseudo-inputs prior achieves the lowest ELBO values for β = 0 .
5 , 1 , 2 consistently for Omniglot and for
Quickdraw.
This β range is actually the most significant as it is the closest to the actual
ELBO bound for β
= 1.
In fact, only β = 1 constitutes a lower bound on the marginal log-likelihood while the more extreme values β = 0 .
1 , 5 , 10 allow only to test the autoencoding and auto-decoding behavior.
Furthermore, when β > 1, the over-pruning biases of variational inference are enhanced (St¨uhmer et al., 2019).
In particular, for β = 1, the proposed Pseudo-inputs prior significantly outperforms the other priors.
Therefore, we can state that based on the rate-distortion analysis, the proposed Pseudo-inputs prior achieves better performances, consistently for Omniglot and Quickdraw, suggesting it is more expressive than the VampPrior, the MoG prior, and the standard prior.
= 0.
R E +
K L = 0.
= 0.
=     =      =     =170.
3.2. MAML-VAE and VAE comparison
In this section, we compare performances between the VAE model and the proposed MAMLVAE model.
Furthermore, we show the proposed Pseudo-inputs prior outperforms baseline priors in both models on both Omniglot and Quickdraw.
We provide details about how Omniglot and Quickdraw are organised into tasks for meta-learning and about the models in Appendix A. We also compare the models implemented with convolutional layers, namely the MAML-CVAE and CVAE models, to see whether convolutions help generate a latent space suitable for unsupervised few-shot classification and, if so, in what proportion to the
MAML-VAE and VAE models.
Whenever the “MAML-” prefix is used, all parameters (including the prior, if applicable) are learnt through the MAML algorithm, while standard training is used otherwise.
The results obtained are reported in Table 1.
On the one hand, results clearly show the proposed Pseudo-inputs prior outperforms baseline priors.
This is true on Omniglot and Quickdraw and consistently across all four models.
This suggests that results are robust and that the Pseudo-inputs prior is suitable both for standard and meta-learning, since it is able to efficiently learn a dataset-specific representation and a higher-level one.
On the other hand, results show that the VAE and the MAML-VAE models, and their convolutional counterparts CVAE and MAML-CVAE, achieve comparable performances across different priors on Omniglot and Quickdraw.
The difference is more evident when convolutional layers are employed as the MAML-CVAE model outperforms the CVAE model over all priors on Omniglot while it achieves worse results on Quickdraw.
This is notable, as meta-learning is in general a harder task.
3.3. Unsupervised few-shot classification in the latent space
In order to quantify whether the latent space learnt in the MAML-VAE model is richer than the one learnt with VAEs, we measure its suitability to unsupervised few-shot classification, which we evaluate as a downstream task of prototypical networks (Snell et al., 2017).
Specifically, we exploit prototypical networks to evaluate the embedding generated
[Table]
by variational posterior q ϕ ∗ ( z | x ) depending on whether it was learnt through the VAE or the
MAML-VAE model.
Note that the variational posterior is trained through either standard training or the MAML algorithm and that the learnt parameters are left unchanged.
This makes the task unsupervised since the models, and hence the variational posteriors, are trained without any label information.
Instead of averaging over samples of the variational posterior, we evaluate the latent space generated by its mean vector µ ϕ ∗ ( x ) directly.
Let us now consider a few-shot task as an N -way classification with S support and Q query shots, taken respectively from D supp and D query , both drawn from the test set D test .
At test time, the support set allows to obtain { c suppk } Nk =1 prototypical representations of the N classes as:
[Equation]
where D k is the set of support points labelled with class k .
The accuracy of the unsupervised few-shot classification task is then computed on the query set based on the prototypical class representations obtained on the support set.
This is computed as softmax over the
Euclidean distance between the class prototypes { c suppk } Nk =1 and the encoded query points.
We evaluated the models on Omniglot over various N -way S -shot classification tasks: results for the MAML-VAE and VAE model are reported in Table 2 while for the MAMLCVAE and CVAE models in Appendix B in Table 4.
On the one hand, we see that the
MAML-VAE model achieves significantly higher accuracy than the VAE model across all priors, which indicates that meta-learning encourages a significantly richer latent space.
When convolutional layers are employed, the difference shrinks, which indicates that increased capacity of the convolutional encoder and decoder allows for a less expressive prior distribution.
However, when either the VampPrior or the Pseudo-inputs prior are employed, it is still beneficial to exploit both contributions.
On the other hand, we showed that the
[Table]
CACTUs (Hsu et al., 2019), as detailed in Appendix B in Table 5.
4. Related Work
In recent years, many different priors have been proposed for VAEs (e.g., Fortuin et al., 2019a, 2020;
Manduchi et al., 2019, 2021; Jazbec et al., 2020, 2021;
Ashman et al., 2020).
However, to the best of our knowledge, none of these approaches used meta-learning to improve the priors using related tasks.
The prior most related to our approach is the
VampPrior (Tomczak and Welling, 2018), which we also use as a baseline in our experiments.
While meta-learning (Thrun and Pratt, 1998; Baxter, 2000) has also been recently connected to Bayesian models (e.g., Grant et al., 2018; Finn et al., 2018; Yoon et al., 2018;
Fortuin et al., 2019b; Rothfuss et al., 2021), it has not yet been used for VAEs.
5. Conclusion
We have shown that meta-learning, and specifically the MAML algorithm, can be beneficial for VAEs, since it enables richer latent representations.
In particular, we have shown that the MAML-VAE model achieves comparable performances with VAEs in terms of test marginal log-likelihood, despite meta-learning being a much harder task, while it yields better performances in downstream few-shot classification tasks.
Moreover, we have proposed the Pseudo-inputs prior, which allows to flexibly encode a mixture of Gaussians while also learning a high-level representation of the data through its pseudo-inputs.
We have shown that this prior outperforms the standard prior, the Mixture of Gaussians, and the
VampPrior, in standard supervised and meta-learning settings, and on different data sets.
Appendix A. Dataset and models
Omniglot
Omniglot (Lake et al., 2015) contains images of handwritten characters from 50 alphabets: 30 are intended for training ( background set) and 20 for test ( evaluation set).
Each character is drawn 20 times from 20 different people.
For both the VAE and
MAML-VAE models we refer to the background-evaluation split.
However, the MAML algorithm requires the dataset to be organised into tasks, which we do in N-way K-shot fashion, where training is performed over K support points for a task that consists of N classes.
Specifically, we identify a task as learning a particular character, so each task is defined by the 20 handwritten images of each character, which are split into support and query images.
Overall, the dataset so defined is composed of 1020 tasks for training,for validation and 420 for test, for a total of 20400, 3400 and 8400 data points, respectively.
For the Omniglot meta-learning dataset we leverage the Torchmeta implementation (Deleu et al., 2019), which allows to easily build N-way K-shot tasks.
In order to ensure images were normalised identically, we used the Torchmeta package for the VAE model as well.
Quickdraw
Quickdraw (Jongejan et al., 2016) is a dataset containing 50 million drawings belonging to 345 classes obtained by asking different users to draw the object represented by each class.
Since the objects can be drawn in several ways, examples belonging to the same class may look significantly different.
The meta-learning task on Quickdraw is thus to capture the “idea” of the represented object, which is a much more challenging task than learning a specific character as in Omniglot.
To our knowledge there is no easy-to-use meta-learning implementation of the Quickdraw data set, so we decided to build our own.
We selected 200 drawings for each class and we split training, validation and test so that each contains disjoint sets of classes.
We picked 50 classes for training, 10 for validation and 25 for test, for a total of 10000, 2000 and 5000 data points, respectively.
We report the classes used in each split in Table 3.
When choosing the classes for training, validation and test we paid attention to keep a consistent variety of shapes
[Table]
[Table]
Both Quickdraw and Omniglot consist of square grey images composed of 28 × 28 pixels
[Table]
In this work we decided to binarise each image by applying a filter, namely a mask for which pixels greater than 0.5 are set to 1 while those that are smaller than 0.5 are set to 0.
This way a sharper and more continuous image is obtained.
In contrast, it is also common to binarise the images through a Bernoulli sampling (Tomczak and Welling, 2018).
However, this technique does not guarantee that the resulting shape is continuous, which is a crucial aspect both for the Omniglot and the Quickdraw dataset.
In fact, in both cases the images are created through a (mostly) continuous stroke of pen.
Examples of how the Bernoulli sampling can distort the resulting images and a comparison with our filtering approach are reported across Figure 3( a )-
( h ).
[Figure]
Sampling the pseudo-inputs
The encoder f ψ is a mapping from the pseudo-inputs space to the latent space.
Instead of learning the pseudo-inputs, they may also be defined over some simple distribution so that its samples are then drawn and mapped in the latent space.
One way to do so is to sample the pseudo-inputs from the standard Gaussian distribution u k ∼ N ( c k , σ 2 1 d ), where σ 2 is an hyper-parameter and c k a fixed (or trainable) mean vector, and to do so at each forward pass.
Intuitively, the parameter σ 2 controls the variance of the encoded samples ˆ µ k , which are obtained as ˆ µ k = f ψ ( c k + σϵ ) with ϵ ∼ N (0 , 1 d ).
Assuming the optimisation converges, samples of the pseudo-inputs are then mapped into samples of the encoded prior distribution.
Sampling the pseudo-inputs at each forward pass during optimisation does not guarantee convergence though.
However, if the analogous problem where pseudo-inputs are trainable converges, then for sufficiently small ϵ it also converges for trainable c k values.
We show in our experiments that c k need not to be trainable for convergence and we suggest that sampling pseudo-inputs allows to alleviate the mode-seeking behaviour typical of the Kullback-Leibler divergence (Burda et al., 2016).
Input: D train
Output: Φ
Initialise Φ := { θ, ϕ, λ } ; while not converged do
Sample a mini-batch of task(s) τ from D supp , D query ∼ D train ; for each task τ ∈ D supp do Φ ′ ← Φ − η in ∇ Φ L τ (Φ); end Φ ← Φ − η out ∇ Φ (cid:80) τ
∼D query L τ (cid:0) Φ ′ (Φ) (cid:1) ; end
Input: Φ ∗ , D test
Output: L
L ← 0; for D supp , D query in D test do Φ ′ ← Φ ∗ − η in ∇ Φ L τ (Φ ∗ );
L ← L + (cid:80) τ
∼D query L τ (Φ ′ (Φ ∗ )); end
L ← L/
|D test |
Details about the VAE and MAML-VAE models
The VAE model is trained with a learning rate of 0 .
0005 whereas the MAML-VAE model with an inner learning rate η in = 0 .
05 and an outer learning rate η out = 0 .
The MAML-VAE model is trained in a 1-way fashion with 5 support and 15 query shots (20 images per task).
We employ ADAM (Kingma and Ba, 2017) as optimizer for the VAE model and for the MAML-VAE model in the outerloop, while in the inner loop we use simple stochastic gradient descent.
The computational bottleneck of the MAML-VAE model is in the outer-loop, which requires back-propagating through the gradients of the inner-loop.
Therefore, in MAML-VAE we take a single gradient update in the inner loop.
In each model we perform 100 warm-up steps (Bowman et al., 2016) to ensure that the reconstruction error and the Kullback-Leibler term in the ELBO are properly optimised.
Furthermore, the stop criterion used for training is in all settings a 50 look-ahead steps early stopping.
As suggested in (Tomczak and Welling, 2018), we use a 40-dimensional latent space for both the VAE and MAML-VAE models.
Both models are implemented on top of their code, which relies on PyTorch (Paszke et al., 2019).
We also provide an implementation of the models that exploits convolutional layers, which we call the CVAE and MAML-CVAE models.
In order to make a fair comparison we maintain the same architecture, only replacing fully connected layers with convolutional ones.
The trained models are evaluated in terms of the marginal log-likelihood on the test set, which is computed using the importance weighting estimate proposed in the IWAE model (Burda et al., 2016) with 5000 samples.
Concerning the Pseudo-inputs prior, its encoder f ψ is implemented with three fully connected layers of dimension h × 100, 100 × 100, and 100 × d with h = 50 and d = 40.
We applied the Exponential Linear Unit (ELU) (Clevert et al., 2016) as activation function on each layer.
In order to make a fair comparison, the MoG prior, the VampPrior and the Pseudo-inputs prior are all implemented with 100 components.
Appendix B. Further results
= 0.
R E +
K L = 0.
= 0.
=      =      =     =250.
( N ; S, Q ) (5; 1 , 19) (5; 5 , 15) (20; 1 , 19) (20; 5 , 15)
[Table] [Table]
[Table] [Table]
21 ±
[Table]
47 ±
[Table] [Table]
[Table] [Table] [Table] [Table]
[Table] [Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
Vamp
[Table] [Table]
81 ±
[Table]
46 ±
[Table] [Table]
[Table] [Table]
Comparison with other unsupervised approaches
We compare our results to a common baseline for unsupervised few-shot classification that exploits meta-learning: the CACTUs model (Hsu et al., 2019).
However, note that our model is not intended for unsupervised few-shot classification, but is rather used to evaluate the generated latent space.
The
CACTUs model automatically builds tasks from unlabelled data and then exploits metalearning to find a learning procedure that allows to solve the constructed tasks.
In order to construct tasks from unlabelled data it learns an appropriate embedding function mapping from the data space and then generates a partition over the embedded data-points leveraging k-means (MacQueen, 1967).
The meta-learner is then trained over the so-defined labelled data through either the MAML algorithm
(“CACTUs-MAML”) or prototypical networks (“CACTUs-ProtoNets”).
However, the learnt embedding may be used directly for downstream supervised learning.
In particular, we consider the implementations with a linear classifier (“Emb. linear classifier”) and a MLP with one hidden layer of 128 units (“Emb.
MLP classifier”), which include the most performing ones.
Two unsupervised embedding learning algorithms are used: ACAI (Berthelot et al., 2018) and BiGAN (Donahue et al., 2017; Dumoulin et al., 2017).
In Table 5 we compare their results and their baselines against our most performing models in terms of few-shot classification accuracy.
[Table]
The comparison shows that the MAML-CVAE model implemented with Pseudo-inputs prior outperforms all BiGAN implementations.
Furthermore, it achieves comparable accuracy to the ACAI implementation of CACTUs-ProtoNets but lower accuracy with respect to the ACAI implementation of CACTUs-MAML.
We regard this as a significant result as our method was not meant for unsupervised few-shot classification.