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: The paper provides an optimal transport (OT) based algorithm for improving existing summary networks for learning from set-structured data. The proposed approach views each set as a distribution over a set of global prototypes. To learn the distribution over the global prototypes, the proposed approach minimizes the OT distance to the set's empirical distribution over data points. Empirical results demonstrate that the proposed framework improves upon the existing summary network approaches as well as metric-based few-shot classification and generative modeling applications.