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: The submission proposes an algorithm for the semi-supervised meta-learning (unsupervised meta-training + supervised meta-testing) setting of [1], which adapts the few-shot learning + evaluation setting of [2, 3] by omitting classification labels at meta-training time. The algorithm makes use of a variational auto-encoder (VAE) formulation defined over a hierarchical model that describes the decomposition of a dataset into tasks of datapoint-target pairs (i.e., the meta-learning setup). The prior distribution of the hierarchical VAE is taken to be a mixture of Gaussians to facilitate the construction of pseudo-labels at meta-training time. The algorithm is evaluated on the Omniglot and miniImageNet few-shot classification tasks (with labels unused at meta-training time).