Meta-GMVAE: Mixture of Gaussian VAE for Unsupervised Meta-Learning

1 INTRODUCTION . Unsupervised learning is one of the most fundamental and challenging problems in machine learning , due to the absence of target labels to guide the learning process . Thanks to the enormous research efforts , there now exist many unsupervised learning methods that have shown promising results on real-world domains , including image recognition ( Le , 2013 ) and natural language understanding ( Ramachandran et al. , 2017 ) . The essential goal of unsupervised learning is obtaining meaningful feature representations that best characterize the data , which can be later utilized to improve the performance of the downstream tasks , by training a supervised task-specific model on the top of the learned representations ( Reed et al. , 2014 ; Cheung et al. , 2015 ; Chen et al. , 2016 ) or fine-tuning the entire pre-trained models ( Erhan et al. , 2010 ) . Meta-learning , whose objective is to learn general knowledge across diverse tasks , such that the learned model can rapidly adapt to novel tasks , shares the spirit of unsupervised learning in that both seek more efficient and effective learning procedure over learning from scratch . However , the essential difference between the two is that most meta-learning approaches have been built on the supervised learning scheme , and require human-crafted task distributions to be applied in fewshot classification . Acquiring labeled dataset for meta-training may require a massive amount of human efforts , and more importantly , meta-learning limits its applications to the pre-defined task distributions ( e.g . classification of specific set of classes ) . Two recent works have proposed unsupervised meta-learning that can bridge the gap between unsupervised learning and meta-learning by focusing on constructing supervised tasks with pseudo-labels from the unlabeled data . To do so , CACTUs ( Hsu et al. , 2019 ) clusters data in the embedding space learned with several unsupervised learning methods , while UMTRA ( Khodadadeh et al. , 2019 ) assumed that each randomly drawn sample represents a different class and augmented each pseudoclass with data augmentation ( Cubuk et al. , 2018 ) . After constructing the meta-training dataset with such heuristics , they simply apply supervised meta-learning algorithms as usual . Despite the success of the existing unsupervised meta-learning methods , they are fundamentally limited , since 1 ) they only consider unsupervised learning for heuristic pseudo-labeling of unlabeled data , and 2 ) the two-stage approach makes it impossible to recover from incorrect pseudo-class assignment when learning the unsupervised representation space . In this paper , we propose a principled unsupervised meta-learning model based on Variational Autoencoder ( VAE ) ( Kingma & Welling , 2014 ) and set-level variational inference using self-attention ( Vaswani et al. , 2017 ) . Moreover , we introduce multi-modal prior distributions , a mixture of Gaussians ( GMM ) , assuming that each modality represents each class-concept in any given tasks . Then the parameter of GMM is optimized by running Expectation-Maximization ( EM ) on the observations sampled from the set-dependent variational posterior . In this framework , however , there is no guarantee that each modality obtained from EM algorithm corresponds to a label . To realize modality as label , we deploy semi-supervised EM at meta-test time , considering the support set and query set as labeled and unlabeled observations , respectively . We refer to our method as Meta-Gaussian Mixture Variational Autoencoders ( Meta-GMVAE ) ( See Figure 1 for high-level concept ) . While our method can be used as a full generative model for generating the samples ( images ) , the ability to generalize to generate samples may not be necessary for capturing the meta-knowledge for non-generative downstream tasks . Thus , we propose another version of Meta-GMVAE that reconstructs high-level features learned by unsupervised representation learning approaches ( e.g . Chen et al . ( 2020 ) ) . To investigate the effectiveness of our framework , we run experiments on two benchmark fewshot image classification datasets , namely Omiglot ( Lake et al. , 2011 ) and Mini-Imagenet ( Ravi & Larochelle , 2017 ) . The experimental results show that our Meta-GMVAE obtains impressive performance gains over the relevant unsupervised meta-learning baselines on both datasets , obtaining even better accuracy than fully supervised MAML ( Finn et al. , 2017 ) while utilizing as small as 0.1 % of the labeled data on one-shot settings in Omniglot dataset . Moreover , our model can generalize to classification tasks with different number of ways ( classes ) without loss of accuracy . Our contribution is threefold : • We propose a novel unsupervised meta-learning model , namely Meta-GMVAE , which metalearns the set-conditioned prior and posterior network for a VAE . Our Meta-GMVAE is a principled unsupervised meta-learning method , unlike existing methods on unsupervised meta-learning that combines heuristic pseudo-labeling with supervised meta-learning . • We propose to learn the multi-modal structure of a given dataset with the Gaussian mixture prior , such that it can adapt to a novel dataset via the EM algorithm . This flexible adaptation to a new task , is not possible with existing methods that propose VAEs with Gaussian mixture priors for single task learning . • We show that Meta-GMVAE largely outperforms relevant unsupervised meta-learning baselines on two benchmark datasets , while obtaining even better performance than a supervised metalearning model under a specific setting . We further show that Meta-GMVAE can generalize to classification tasks with different number of ways ( classes ) . 2 RELATED WORK . Unsupervised learning Many prior unsupervised learning methods have developed proxy objectives which is either based on reconstruction ( Vincent et al. , 2010 ; Higgins et al. , 2017 ) , adversarially obtained image fidelity ( Radford et al. , 2016 ; Salimans et al. , 2016 ; Donahue et al. , 2017 ; Dumoulin et al. , 2017 ) , disentanglement ( Bengio et al. , 2013 ; Reed et al. , 2014 ; Cheung et al. , 2015 ; Chen et al. , 2016 ; Mathieu et al. , 2016 ; Denton & Birodkar , 2017 ; Kim & Mnih , 2018 ; Ding et al. , 2020 ) , clustering ( Coates & Ng , 2012 ; Krähenbühl et al. , 2016 ; Bojanowski & Joulin , 2017 ; Caron et al. , 2018 ) , or contrastive learning ( Chen et al. , 2020 ) . In the unsupervised learning literature , the most relevant work to ours are methods that use Gaussian Mixture priors for variational autoencoders . Dilokthanakul et al . ( 2016 ) ; Jiang et al . ( 2017 ) consider single task learning and therefore , the learned prior parameter is fixed after training , and thus can not adapt to new tasks . CURL ( Rao et al. , 2019 ) learns a network that outputs Gaussian mixture priors over a sequence of tasks for unsupervised continual learning . However CURL can not adapt to a new task without training on it , while our framework can generalize to a new task without any training , via amortized inference with a dataset ( task ) encoder . Also , our model does not learn Gaussian mixture priors but rather obtain them on the fly using the expectation-maximization algorithm . Meta-learning Meta-learning ( Thrun & Pratt , 1998 ) shares the intuition of unsupervised learning in that it aims to improve the model performance on an unseen task by leveraging prior knowledge , rather than learning from scratch . While the literature on meta-learning is vast , we only discuss relevant existing works for few-shot image classification . Metric-based meta-learning ( Koch et al. , 2015 ; Vinyals et al. , 2016 ; Snell et al. , 2017 ; Oreshkin et al. , 2018 ; Mishra et al. , 2018 ) is one of the most popular approaches , where it learns to embed the data instances of the same class to be closer in the shared embedding space . One can measure the distance in the embedding space by cosine similarity ( Vinyals et al. , 2016 ) , or Euclidean distance ( Snell et al. , 2017 ) . On the other hand , gradient-based meta-learning ( Finn et al. , 2017 ; 2018 ; Li et al. , 2017 ; Lee & Choi , 2018 ; Ravi & Beatson , 2019 ; Flennerhag et al. , 2020 ) aims at learning a global initialization of parameters , which can rapidly adapt to a novel task with only a few gradient steps . Moreover , some previous works ( Hewitt et al. , 2018 ; Edwards & Storkey , 2017 ; Garnelo et al. , 2018 ) tackle meta-learning by modeling the set-dependent variational posterior with a single global latent variable , however , we model the variational posterior conditioned on each data instances . Moreover , while all of these works assume supervised learning scenarios where one has access to full labels in meta-training stage , we focus on unsupervised setting in this paper . Unsupervised meta-learning One of the main limitations of conventional meta-learning methods is that their application is strictly limited to the tasks from a pre-defined task distribution . A few works ( Hsu et al. , 2019 ; Khodadadeh et al. , 2019 ) have been proposed to resolve this issue by combining unsupervised learning with meta-learning . The main idea is to construct meta-training dataset in an unsupervised manner by leveraging existing supervised meta-learning models . CACTUs ( Hsu et al. , 2019 ) deploy several deep metric learning ( Berthelot et al. , 2019 ; Donahue et al. , 2017 ; Caron et al. , 2018 ; Chen et al. , 2016 ) to episodically cluster the unlabeled dataset , and then train MAML ( Finn et al. , 2017 ) and Prototypical Networks ( Snell et al. , 2017 ) on the constructed data . UMTRA ( Khodadadeh et al. , 2019 ) assumes that each randomly drawn sample is from a different class from others , and use data augmentation ( Cubuk et al. , 2018 ) to construct synthetic task distribution for meta-training . Instead of only deploying unsupervised learning for constructing meta-training task distributions , we propose an unsupervised meta-learning model that meta-learns set-level variational posterior by matching the multi-modal prior distribution representing latent classes . 3 UNSUPERVISED META-LEARNING WITH META-GMVAES . In this section , we describe our problem setting with respect to unsupervised meta-learning , and demonstrate our approach . The graphical illustration of our model for unsupervised meta-training and supervised meta-test is depicted in Figure 2 . 3.1 PROBLEM STATEMENT . Our goal is to learn unsupervised feature representations which can be transferred to wide range of downstream few-shot classification tasks . As suggested by Hsu et al . ( 2019 ) ; Khodadadeh et al . ( 2019 ) , we only assume an unlabeled dataset Du = { xu } Uu=1 in the meta-training stage . We aim toward applying the knowledge learned during unsupervised meta-training stage to novel tasks in meta-test stage , which comes with a modest amount of labeled data ( or as few as a single example per class ) for each task . As with most meta-learning methods , we further assume that the labeled data are drawn from the same distribution as that of the unlabeled data , with a different set of classes . Specifically , the goal of a K-way S-shot classification task T is to correctly predict the labels of query data points Q = { xq } Qq=1 , using S support data points and labels S = { ( xs , ys ) } Ss=1 per class , where S is relatively small ( i.e . between 1 and 50 ) . 3.2 META-LEVEL GAUSSIAN MIXTURE VAE . Unsupervised meta-training We now describe the meta-learning framework for learning unsupervised latent representations that can be transferred to human-designed few-shot image-classification tasks . In particular , we aim toward learning multi-modal latent spaces for Variational Autoencoder ( VAE ) in an episodic manner . We use the Gaussian mixture for the prior distribution pψ ( z ) = ∑K k=1 pψ ( y = k ) pψ ( z|y = k ) , where ψ is the parameter of the prior network . Then the generative process can be described as follows : • y ∼ pψ ( y ) , where y corresponds to the categorical L.V . for a single mode . • z ∼ pψ ( z|y ) , where z corresponds to the Gaussian L.V . responsible for data generation . • x ∼ pθ ( x|z ) , where θ is the parameter of the generative model . The above generative process is similar to those from the previous works ( Dilokthanakul et al. , 2016 ; Jiang et al. , 2017 ) on modeling the VAE prior with Gaussian mixtures . However , they target single-task learning and the parameter of the prior network is fixed after training such as equation 1c in Dilokthanakul et al . ( 2016 ) and equation 5 in Jiang et al . ( 2017 ) , which is suboptimal since a meta-learning model should be able to adapt and generalize to a novel task . To learn the set-dependent multi-modalities , we further assume that there exists a parameter ψi for each episodic datasetDi = { xj } Mj=1 , which is randomly drawn from the unlabeled datasetDu . Then we derive the variational lower bound for the marginal log-likelihood of Di as follows : log pθ ( Di ) = M∑ j=1 log pθ ( xj ) = M∑ j=1 log ∫ pθ ( xj |zj ) pψi ( zj ) qφ ( zj |xj , Di ) qφ ( zj |xj , Di ) dzj ( 1 ) ≥ M∑ j=1 [ Ezj∼qφ ( zj |xj , Di ) [ log pθ ( xj |zj ) + log pψi ( zj ) − log qφ ( zj |xj , Di ) ) ] ] ( 2 ) ≈ M∑ j=1 1 N N∑ n=1 [ log pθ ( xj |z ( n ) j ) + log pψi ( z ( n ) j ) − log qφ ( z ( n ) j |xj , Di ) ] ( 3 ) = : L ( θ , φ , ψi , Di ) , z ( n ) j i.i.d∼ qφ ( zj |xj , Di ) . ( 4 ) Here the lower bound for each datapoint is approximated by Monte Carlo estimation with the sample size N . Following the convention of the VAE literature , we assume that the variational posterior qφ ( zj |xj , Di ) follows an isotropic Gaussian distribution . Algorithm 1 Meta-training Require : An unlabeled dataset Du 1 : Initialize parameters θ , φ 2 : while not done do 3 : Sample B episode datasets { Di } Bi=1 from Du 4 : for all i ∈ [ 1 , B ] do 5 : Draw n MC samples from qφ ( zj |xj , Di ) 6 : Initialize πk as 1/K and randomly choose K different points for µk . 7 : Compute optimal parameter ψ∗i using Eq 7 8 : end for 9 : Update θ , φ using L ( θ , φ , { Di } Bi=1 ) in Eq 9 . 10 : end while Algorithm 2 Meta-test for an episode Require : A test task T = S ∪ Q 1 : Set D = { xs } Ss=1 ∪ { xq } Qq=1 2 : Draw n MC samples from qφ ( zj |xj , D ) 3 : Initialize µk = ∑S , N s , n=1 1y ( n ) s =k z ( n ) s∑S , N s , n=1 1y ( n ) s =k and σ2k = I 4 : Compute optimal parameter ψ∗ using Eq 10 5 : Compute p ( yq|xq , D ) using Eq 11 6 : Infer the label yq = argmax k p ( yq = k|xq , D ) 7 : 8 : Set-dependent variational posterior Our derivation of the evidence lower bound in Eq 4 is similar to that of the hierarchical VAE framework , such as equation 3 in Edwards & Storkey ( 2017 ) and equation 4 in Hewitt et al . ( 2018 ) , in that we use the i.i.d assumption that the log likelihood of a dataset equals the sum over the log-likelihoods of each individual data point . Yet , previous works assume that each input set consists of data instances from a single concept ( e.g . a class ) , therefore , they encode the dataset into a single global latent variable ( e.g . qφ ( z|D ) ) . This is not appropriate for unsupervised meta-learning where labels are unavailable . Thus we learn a set-conditioned variational posterior qφ ( zj |xj , Di ) , which models a latent variable to encode each data xj within the given datasetDi into the latent space . Specifically , we model the variational posterior qφ ( zj |xj , Di ) using the self-attention mechanism ( Vaswani et al. , 2017 ) as follows : H = TransformerEncoder ( f ( Di ) ) µj =WµHj + bµ , σ 2 j = exp ( Wσ2Hj + bσ2 ) qφ ( zj |xj , Di ) = N ( zj ; µj , σ2j ) ( 5 ) Here we deploy TransformerEncoder ( · ) , a neural network based on the multi-head self-attention mechanism proposed by Vaswani et al . ( 2017 ) , to model the dependency between data instances , and f is a convolutional neural network ( or an identity function for the Mini-ImageNet ) which takes each data in Di as an input . Moreover , we use the reparameterization trick ( Kingma & Welling , 2014 ) to train the model with backpropagation since the stochastic sampling process z ( n ) j i.i.d∼ qφ ( zj |xj , Di ) is non-differentiable . Expectation Maximization As discussed before , we assume that the parameter ψi of the prior Gaussian Mixture is task-specific and characterizes the given dataset Di . To obtain the task-specific parameter that optimally explain the given dataset , we propose to locally maximize the lower bound in Eq 4 with respect to the prior parameter ψi . We can obtain the optimal parameter ψ∗i by solving the following optimization problem : ψ∗i = argmax ψi L ( θ , φ , ψi , Di ) = argmax ψi M , N∑ j , n=1 log pψ ( z ( n ) j ) , z ( n ) j i.i.d∼ qφ ( zj |xj , Di ) , ( 6 ) where we only consider the term related to the task-specific parameter ψi , and eliminate the normalization term 1N since it does not change the solution of the optimization problem . The above formula implies that the optimal parameter maximizes the log-likelihood of observations which can be drawn from the variational posterior distribution . However , we do not have an analytic solution for Maximum Likelihood Estimation ( MLE ) of a GMM . The most prevalent approach for estimating the parameters for the mixture of Gaussian is solving it with Expectation Maximization ( EM ) algorithm . To this end , we propose to optimize the taskspecific parameter of GMM prior distribution using EM algorithm as follows : ( E-step ) Qj , n ( k ) : = p ( y ( n ) j = k|z ( n ) j ) = πkN ( z ( n ) j ; µk , I ) ∑ k πkN ( z ( n ) j ; µk , I ) ( M-step ) µk : = ∑M , N j , n=1Qj , n ( k ) z ( n ) j∑M , N j , n=1Qj , n ( k ) , πk : = ∑M , N j , n=1Qj , n ( k ) ∑K k=1 ∑M , N j , n=1Qj , n ( k ) ψi : = { ( µk , I , πk ) } Kk=1 , ( 7 ) where πk , µk , and N ( · ) denote the mixing probability of k-th component , mean parameter , and normal distribution , respectively . We assume that the covariance matrix of Gaussian distribution is fixed with the identity matrix I , following the assumption of original VAE on the prior distribution . We initialize { πk } Kk=1 and { µk } Kk=1 as 1K and randomly drawn K different points , respectively . We can obtain MLE solution for the parameters of GMM , by iteratively performing E-step and M-step until the log-likelihood converges . We found that using a fixed number of iterations for the EM algorithm does not degrade the performance , and consider it as a hyperparameter of our framework . Training objective Note that we want to maximize the variational lower bound of the marginal loglikelihood over all the episode datasets Di that can be sampled from Du . We use stochastic gradient ascent with respect to the variational parameter φ and the generative parameter θ , to maximize the following objective : L ( θ , φ , { Di } Bi=1 ) : = 1 B B∑ i=1 [ max ψi L ( θ , φ , ψi , Di ) ] ( 8 ) = 1 B B∑ i=1 M∑ j=1 1 N N∑ n=1 [ log pθ ( xj |z ( n ) j ) + log pψ∗i ( z ( n ) j ) − log qφ ( z ( n ) j |xj , Di ) ] . ( 9 ) Here we use B mini-batch of episode datasets , where each dataset consists of M datapoints . The task-specific parameter ψ∗i for each episode dataset Di is obtained by EM algorithm in Eq 7 . Supervised meta-test By introducing the multi-modal prior distribution into a generative learning framework , our model learns pseudo-class concepts by clustering latent features with EM algorithm . However , there is no guarantee that each modality obtained by EM algorithm corresponds to the label we are interested in at the meta-test stage . To realize modality as label in downstream fewshot image classification tasks , we deploy semi-supervised EM algorithm instead . Given a task T consisting of support set S = { ( xs , ys ) } Ss=1 and query set Q = { xq } Q q=1 , we use both the support set and query set as an episode dataset D = { xs } Ss=1 ∪ { xq } Q q=1 and draw latent variables from the variational posterior qφ ( zj |xj , D ) . Note that we abbreviate the index i since we consider a single task for now . We then perform semi-supervised EM algorithm as follows : ( E-step ) Qq , n ( k ) : = p ( y ( n ) q = k|z ( n ) q ) = N ( z ( n ) q ; µk , σ2k ) ∑ kN ( z ( n ) q ; µk , σ2k ) ( M-step ) µk : = ∑S , N s , n=1 1y ( n ) s =k z ( n ) s + ∑Q , N q , n=1Qq , n ( k ) z ( n ) q∑S , N s , n=1 1y ( n ) s =k + ∑Q , N q , n=1Qq , n ( k ) , σ2k : = ∑S , N s , n=1 1y ( n ) s =k ( z ( n ) s − µk ) 2 + ∑Q , N q , n=1Qq , n ( k ) ( z ( n ) q − µk ) 2∑S , N s , n=1 1y ( n ) s =k + ∑Q , N q , n=1Qq , n ( k ) ψ : = { ( µk , σ2k , 1 K ) } Kk=1 , ( 10 ) where 1 denotes an indicator function . We fix the mixing probability as 1K since the labels in each task T are uniformly distributed . Moreover , we utilize diagonal covariance σ2k to obtain more accurate statistics for the inference . We initialize µk and σ2k as the average value of support latent representations and the identity matrix I , respectively . Similar to the meta-training stage , we obtain the MLE solution for the parameters of GMM , by performing E-step and M-step for a fixed number of iterations . Finally , we compute the conditional probability of p ( yq|xq , D ) using the obtained parameters ψ∗ as follows : p ( yq|xq , D ) = Eqφ ( zq|xq , D ) [ pψ∗ ( yq|zq ) ] ≈ 1 N N∑ n=1 pψ∗ ( yq|z ( n ) q ) , z ( n ) q i.i.d∼ qφ ( zq|xq , D ) . ( 11 ) Here we compute pψ∗ ( yq|z ( n ) q ) with Bayes rule , and we reuseN different Monte Carlo samples that is drawn for Eq 10 , where the prediction of query ŷq = argmax k p ( yq = k|xq , D ) . We present the pseudo-code of the algorithm for training and inference of Meta-GMVAE in the Algorithm 1 and 2 . Visual feature reconstruction While our method is a generative model that can generate samples from output distribution , the ability to generate samples may not be necessary for discriminative downstream tasks ( Chen et al. , 2020 ) . Moreover , we found that VAEs almost fail to learn in MiniImageNet dataset with the architecturally limited constraints of the meta-learning literature . Thus , we propose a high-level feature reconstruction objective instead for Mini-ImageNet dataset . We experimentally find that the recently proposed constrastive learning framework , namely SimCLR ( Chen et al. , 2020 ) , is the most effective for our settings . Specifically , SimCLR learns high-level representation by performing a constrastive prediction task on pairs of augmented examples derived from a minibatch . We train SimCLR on the unsupervised datasetDu = { xu } Uu=1 , and use high-level features extracted by SimCLR as an input for our framework .

Summary: This paper proposes a method for unsupervised meta-learning based on using a variational autoencoder (VAE). The variational autoencoder model they use differs from the typical one in that it considers episode-specific datasets, where the approximate posterior can be computed as a function of the set (using transformer architecture) rather than using an individual example. Additionally, they use a mixture of Gaussian distribution as a prior, whose parameters are learned per-episode using the EM algorithm. For the supervised evaluation phase, in order to adapt the learned prior to the few-shot dataset setting, semi-supervised EM is run using both support and query sets to adapt the mixture of Gaussian distribution to the evaluation dataset. Then, the query set predictions are obtained using the learned prior and posterior from the VAE model. Experimental evaluation is conducted on the Omniglot and Mini-ImageNet benchmarks and the proposed method is compared against other unsupervised meta-learning methods, mainly CACTUs and UMTRA. An interesting aspect about the Mini-ImageNet experiments are that because learning the VAE directly for this high-dimensional data may be difficult, the authors use features from a SimCLR-trained model as input for their VAE model. The proposed method seems to perform favorably across both of the benchmarks when varying the number of "shots".