Learning Prototype-oriented Set Representations for Meta-Learning 

1 INTRODUCTION . Machine learning models , such as convolutional neural networks for images ( He et al . 2016 ) and recurrent neural networks for sequential data ( Sutskever et al . 2014 ) , have achieved great success in taking advantage of the structure in the input space ( Maron et al . 2020 ) . However , extending them to handle unstructured input in the form of sets , where a set can be defined as an unordered collections of elements , is not trivial and has recently attracted increasing attention ( Jurewicz & StrømbergDerczynski , 2021 ) . Set-input is relevant to a range of problems , such as understanding a scene formed of a set of objects ( Eslami et al . 2016 ) , classifying an object composed of a set of 3D points ( Qi et al . 2017 ) , and estimating summary statistics from a set of data points for implicit generative models ( Chen et al . 2021 ) . Moreover , many meta-learning problems , which process different but related tasks , may also be viewed as set-input tasks ( Lee et al. , 2019 ) , where an input set corresponds to the training dataset of a single task . Therefore , we broaden the scope of set related applications by including traditional set-structured input problems and most meta-learning problems . Both of them aim to improve the quick adaptation ability for unseen sets , even though the latter is more difficult because of limited samples or the occurrence of new categories for classification problems . For a set-input , the output of the model must not change if the elements of the input set are reordered , which entails permutation invariance of the model . To enforce this property , multiple researchers have recently focused on designing different network architectures , which can be referred to as a summary network for compressing the set-structured data into a fixed-size output . For example , the prominent works of Zaheer et al . ( 2017 ) and Edwards & Storkey ( 2017 ) combined the standard feed-forward neural networks with a set-pooling layer , which have been proven to be universal approximators of continuous permutation invariant functions ( Zaheer et al . 2017 ) . Lee et al . ( 2019 ) further introduced Set Transformer to encode and aggregate the features within the set using the multi-head attention mechanism . Maron et al . ( 2020 ) designed deep models and presented a principled approach to learn sets of symmetric elements . Despite the effectiveness and recent popularity of these works in set-input problems , there are several shortcomings for existing summary networks , which could hinder their applicability and further extensions : 1 ) The parameters of the summary network are typically optimized by a task-specific loss function , which could limit the models ’ flexibility . 2 ) A desideratum of a summary network is to extract set features , which have enough ability to represent the summary statistics of the input set and thus benefit the corresponding set-specific task ; but for many existing summary networks , there is no clear evidence or constraint that the outputs of the summary network could describe the set ’ s summary statistics well . These limits still remain even using the recent more carefully designed summary networks , while sets with limited samples further exacerbate the problem . To address the above shortcomings , we present a novel and generic approach to improve the summary networks for set-structured data and adapt them to meta-learning problems . Motivated by meta-learning that aims to extract transferable patterns useful for all related tasks , we assume that there are K global prototypes ( i.e. , centers ) among the collection of related sets , and each prototype or center is encouraged to capture the statistical information shared by those sets , similar to the “ topic ” in topic modeling or “ dictionary atom ” in dictionary learning . Specifically , for the jth set , we consider it as a discrete distribution Pj over all the samples within the set ( in data or feature space ) . At the same time , we also represent this set with another distribution Qj ( in the same space with Pj ) , supported on K global prototypes with a K-dimensional set representation hj . Since hj measures the importance of global prototypes for set j , it can be treated as the prototype proportion for summarizing the salient characteristics of set j . Moreover , the existing summary networks can be adopted to encode set j as hj for their desired property of permutation invariance . In this way , we can formulate the learning of summary networks as the process of learning a Pj to be as close to Qj as possible , a process facilitated by leveraging the optimal transport ( OT ) distance ( Peyré & Cuturi 2019 ) . Therefore , the global prototypes and summary network can be learned by jointly optimizing the task-specific loss and OT distance between Pj and Qj in an end-to-end manner . We can refer to this method as prototype-oriented framework for meta-learning , which is applicable to a range of unsupervised and supervised tasks , such as set-input problems solved by summary networks , meta generation ( Hong et al. , 2020c ; Antoniou et al. , 2017 ) , metric-based few-shot classification ( Snell et al. , 2017 ) , and learning statistics for approximate Bayesian computation ( Chen et al. , 2021 ) . Since our plug-and-play framework can be applied to many meta-learning problems , this paper further instantiates it to the cases of metric-based few-shot classification and implicit meta generative modeling . We summarize our contributions as follows : ( 1 ) We formulate the learning of summary network as the distribution approximation problem by minimizing the distance between distribution over data points and another distribution over global prototypes . ( 2 ) We leverage optimal transport distance to measure the distance between the distributions for use in a joint learning algorithm . ( 3 ) We apply our method to metric-based few-shot classification and construct implicit meta generative model , where summary network is used to extract the summary statistics from set and optimized by the OT loss . Experiments on several meta-learning tasks demonstrate that introducing the OT loss into existing summary networks can extract more effective set representations for the corresponding tasks , which can be also integrated into existing few-shot classification and GAN frameworks , producing a new way to learn the set ’ summary statistics applicable to many applications . 2 BACKGROUND . 2.1 SUMMARY NETWORKS FOR SET-STRUCTURED INPUT . To deal with the set-structured input Dj = { xj,1 : Nj } and satisfy the permutation invariance in set , a remarkably simple but effective summary network is to perform pooling over embedding vectors extracted from the elements of a set . More formally , Sφ ( Dj ) = gφ2 ( pool ( { fφ1 ( xj1 ) , . . . , fφ1 ( xjNj ) } ) ) , ( 1 ) where fφ1 ( · ) acts on each element of a set and gφ2 ( pool ( · ) ) aggregates these encoded features and produces desired output , and φ = { φ1 , φ2 } denotes the parameters of the summary network . Most network architectures for set-structured data follow this structure ; see more details from the previous works ( Lee et al. , 2019 ; Zaheer et al. , 2017 ; Edwards & Storkey , 2017 ; Maron et al. , 2020 ) . 2.2 OPTIMAL TRANSPORT . Although OT has a rich theory , we limit our discussion to OT for discrete distributions and refer the reader to Peyré & Cuturi ( 2019 ) for more details . Let us consider p and q as two discrete probability distributions on the arbitrary space X ⊆ Rd , which can be formulated as p = ∑n i=1 aiδxi and q = ∑m j=1 bjδyj . In this case , a ∈ Σn and b ∈ Σm , where Σm denotes the probability simplex of Rm . The OT distance between a and b is defined as OT ( a , b ) = min T∈U ( a , b ) 〈T , C〉 , ( 2 ) where 〈· , ·〉means the Frobenius dot-product ; C ∈ Rn×m≥0 is the transport cost function with element Cij = C ( xi , yj ) ; T ∈ Rn×m > 0 denotes the doubly stochastic transport probability matrix such that U ( a , b ) : = { T | ∑n i Tij = bj , ∑m j Tij = ai } . To relax the time-consuming problem when opti- mising the OT distance , Cuturi ( 2013 ) introduced the entropic regularization , H = − ∑ ij Tij lnTij , leading to the widely-used Sinkhorn algorithm for discrete OT problems . 3 PROPOSED FRAMEWORK . In meta-learning , given a meta-distribution pM of tasks , the marginal distribution pj of task j is sampled from pM for j ∈ J , where J denotes a finite set of indices . For example , we can sample pj from pM with probability 1J when pM is uniform over a finite number of marginals . During meta-training , direct access to the distribution of interest pj is usually not available . Instead , we will observe a set of data points Dj = { xji } Nji=1 , which consists of Nj i.i.d samples from pj over R d. We can roughly treat the meta-learning problems as the set-input tasks , where dataset Dj from pj corresponds to an input set . To learn more representative features from related but unseen sets in meta-learning problems , we adopt the summary network as the encoder to extract set representations and improve it by introducing the OT loss and global prototypes , providing many applications . Besides , we also provide the applications to metric-based few-shot classification and implicit generative framework by assimilating the summary statistics . Below we describe our model in detail . 3.1 LEARNING GLOBAL PROTOTYPES AND SET REPRESENTATION VIA OPTIMAL TRANSPORT . Given J sets from meta-distribution pM , we can represent each set Dj from meta-distribution pM as an empirical distribution over Nj samples on the original data space , formulated as Pj = ∑Nj i=1 1 Nj δxji , xji ∈ Rd . ( 3 ) Since all sets ( distributions ) draw from meta-distribution pM are closely related , it is reasonable to assume that these sets share some statistical information . Motivated by dictionary learning and topic modeling , we define the shared information as the learnable global prototype matrix B = { βk } ∈ Rd×K , where K means the number of global prototypes and βk denotes the distributed representation of the k-th prototype in the same space of the observed data points ( i.e. , “ topic ” in topic modeling ) . Given the global prototype matrix B , each set j can be represented with a K-dimensional weight vector hj ∈ Σk ( e.g. , “ topic proportion ” in topic modeling ) , where hjk indicates the weight of the prototype βk for set j . Therefore , we can represent set Dj with another distribution Qj on global prototypes β1 : K , defined as Qj = ∑K k=1 hjkδβk , βk ∈ Rd , ( 4 ) where hj can be viewed as set representation for describing set j . Since set j can be represented as Qj and Pj , we propose to learn set-specific representation hj andK global prototypes B by pushing Qj towards Pj : OT ( Pj , Qj ) = min B , hj 〈T , C〉 def.= Nj∑ i K∑ k CikTik , ( 5 ) where C ∈ RNj×K≥0 is the transport cost matrix . In this paper , to measure the distance between data point xji in set j and prototype βk , unless specified otherwise , we specify the construction of C as Cik = 1 − cos ( xji , βk ) , for providing an upper-bounded positive similarity metric . Besides , the transport probability matrix T ∈ RNj×K > 0 should satisfy Π ( a , b ) : = { T | T1K = a , T > 1Nj = b } with Tik = T ( xji , βk ) , where a = [ 1Nj ] ∈ Σ Nj and b = [ hjk ] ∈ ΣK denote the respective probability vectors for distribution Pj in Equation 3 and Qj Equation 4 . Since hj should be invariant to the permutations of the samples in set j , we adopt summary network to encode the set of Nj points . For unsupervised tasks , taking the summary network in Equation 1 as the example , we can directly add a Softmax activation function into Sφ to enforce the simplex constraint in set representation hj , denoted as hj = Softmax ( Sφ ( Dj ) ) . As shown in Fig . 1 , given J sets , to learn the global prototype matrix B and the summary network parameterized by φ , we adopt the entropic constraint ( Cuturi , 2013 ) and define the average OT loss for all training sets as LOT = min B , φ 1 J J∑ j=1 Nj∑ i K∑ k CikTik − Nj∑ i K∑ k −TiklnTik = min B , φ 1 J J∑ j=1 ( OT ( Pj , Qj ) ) , ( 6 ) where is a hyper-parameter for entropic constraint . Algorithm 1 describes the workflow of OT distance for improving summary network under unsupervised tasks . For supervised tasks , set j is denoted as Dj = { xj,1 : Nj , yj } , where yj is the ground-truth output determined by specific task . As hj is a normalized weight vector , directly using it to realize the corresponding task may be undesired . Denoting zj=pool ( { fφ1 ( xj1 ) , . . . , fφ1 ( xjNj ) } ) , we project it to the following vectors : hj = fe ( zj ) , ŷj = fλ ( zj ) , ( 7 ) where hj and ŷj are responsible for the OT and task-specific losses , respectively . Now the summary network parameters φ̃= { e , λ , φ1 } and global prototypes B can be learned by jointly optimizing the task-specific loss ( computed by ŷj and yj ) and OT loss in Equation F. In summary , minimizing the OT distance between the prototype distribution Qj and empirical distribution Pj provides a principled and unsupervised way to encourage the summary network to capture the set ’ s summary statistics , where a suite of summary networks can be integrated into this plug-and-play framework . Therefore , it realizes efficient learning from new sets for both unsupervised and supervised tasks .

Summary: This work introduces a straightforward development for set representation learning in the meta-learning context based on the intuition that the sets encountered in real-world meta-learning tasks tend to have common attributes, as illustrated in Figure 1. The idea is to jointly learn these common attributes $\beta_{[1:K]}$, referred to as 'global centres' or 'global prototypes', and the parameters of a summary network using an Optimal Transport derived loss function that compares the empirical set distribution $P_j$ with a set summary distribution over the global centres/prototypes, $Q_j=\sum_{k=1}^K h_{jk}\delta_{\beta_k}$. The idea is elegant in its simplicity and effectiveness.