Year: 2021

Abstract:
Learning from set-structured data is a fundamental problem that has recently attracted increasing attention, where a series of summary networks are introduced to deal with the set input. In fact, many meta-learning problems can be treated as set-input tasks. Most existing summary networks aim to design different architectures for the input set in order to enforce permutation invariance. However, scant attention has been paid to the common cases where different sets in a metadistribution are closely related and share certain statistical properties. Viewing each set as a distribution over a set of global prototypes, this paper provides a novel optimal transport (OT) based way to improve existing summary networks. To learn the distribution over the global prototypes, we minimize its OT distance to the set empirical distribution over data points, providing a natural unsupervised way to improve the summary network. Since our plug-and-play framework can be applied to many meta-learning problems, we further instantiate it to the cases of few-shot classification and implicit meta generative modeling. Extensive experiments demonstrate that our framework significantly improves the existing summary networks on learning more powerful summary statistics from sets and can be successfully integrated into metric-based few-shot classification and generative modeling applications, providing a promising tool for addressing set-input and meta-learning problems.

heading: 1 INTRODUCTION
text: 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.

heading: 2 BACKGROUND

heading: 2.1 SUMMARY NETWORKS FOR SET-STRUCTURED INPUT
text: 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).

heading: 2.2 OPTIMAL TRANSPORT
text: 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.

heading: 3 PROPOSED FRAMEWORK
text: 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.

heading: 3.1 LEARNING GLOBAL PROTOTYPES AND SET REPRESENTATION VIA OPTIMAL TRANSPORT
text: 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.

heading: 3.2 APPLICATION TO METRIC-BASED FEW-SHOT CLASSIFICATION
text: As a challenging meta-learning problem, few-shot classification has recently attracted increasing attention, where one representative method is metric-based few-shot classification algorithms. Taking the ProtoNet (Snell et al., 2017) as the example, we provide a simple but effective method to improve its classification performance with the help of OT distance and summary network, where we refer the reader to ProtoNet for more details. In a specific, ProtoNet represents each class by computing an M -dimensional representation cj ∈ RM , with an embedding function fφ1 : Rd → RM , where we adopt the same φ1 as the learnable parameters, following the feature extractor in summary network in Equation 1 for simplicity. Formally, ProtoNet adopts the average pooling to aggregate the embedded features of support points belonging to its class into vector cj = 1|Sj | ∑ (xji,yji)∈Sj fφ1 (xji). ProtoNet then compares the distance between a query point fφ1(x) to the cj in the same embedding space. Motivated by the summary network, we further introduce a feed-forward network gφ2 to map the cj into the hj used to defineQj distribution over B. Therefore, the functions gφ2 and fφ1 can be jointly optimized by minimizing the OT distance between Pj and Qj and the original classification loss in ProtoNet, where cj is still used to perform classification since hj is a normalized weight vector. Only introducing matrix B and gφ2 , whose parameters are usually negligible compared to the original feature extractor fφ1 , our proposed method provides a simple but effective way to enforce
Algorithm 1 Workflow about minimizing the OT distance between Pj and Qj . Require: Datasets D1:J , batch size m, learning rate α, initial summary network parameters φ, initial global prototype matrix B, cost function C and hyper-parameter . while B, φ has not converged do
Randomly choose j from 1, 2, ..,J Sample the real data {xji}mi=1 from set j, which is denoted as empirical distribution Pj in Equation 3; Compute the set representation hj=Softmax(Sφ({xji}mi=1)) with summary network in Equation 1; Represent the Qj with global prototype matrix B and statistics hj in Equation 4; Compute the loss OT (Pj , Qj) between Pj and Qj with Sinkhorn algorithm in Equation F; B← B+ αgB, where gB ← ∇B [OT (Pj , Qj)]; φ← φ+ αgφ, where gφ ← ∇φ [OT (Pj , Qj)];
end while
the fφ1 to learn more powerful representation cj of each class and feature fφ1(x) of query sample, benefiting the metric-based few-shot classification.

heading: 3.3 APPLICATION TO IMPLICIT META GENERATIVE MODELS
text: Considering implicit meta generative modeling is still a challenging but important task in metalearning, we further present how to construct the model by introducing set representation hj as summary statistics, where we consider GAN-based implicit models. Specifically, given pj ∼ pM, we aim to construct a parametrized pushforward (i.e., generator) of reference Gaussian distribution ρ, denoted as Tθ(·, pj)]ρ, to approximate the marginal distribution pj , where θ summarizes the parameters of pushforward. Since it is unaccessible to the distribution of interest pj , we replace pj with Pj and use the summary network to encode set Dj into hj as discussed above, which is further fed into the generator serving as the conditional information, denoted as Tθ(z;hj), z ∼ ρ. To enforce the pushforward Tθ(z;hj) to fit the real distribution Pj as well as possible, we introduce a discriminator fw following the standard GAN (Goodfellow et al., 2014). Generally, our GAN-based model consists of three components. Summary network Sφ(·) focuses on learning the summary statistics hj by minimizing the OT distance between Pj and Qj . The pushforward aims to push the combination of a random noise vector z and statistics hj to generate samples that resemble the ones from Pj , where we simply adopt a concatenation for hj and z although other choices are also available. Besides, fw tries to distinguish the “fake” samples from the “real” samples in set Dj . Therefore, we optimize the implicit meta generative model by defining the objective function as:
min B,φ,θ max w
V (B, φ, θ, w) = LOT + 1
J J∑ j=1 Ex∼Pj [log fw(x)] + Ez∼ρ[log(1− fw(Tθ(z;hj))] (8)
In addition to the standard GAN loss, we can also adopt the Wasserstein GAN (WGAN) of Arjovsky et al. (2017) to approximate Pj . It is also flexibly to decide the input fed into fw. For example, following the conditional GAN (CGAN) of Mirza & Osindero (2014), we can combine hj and data points (generated or real), where the critic can be denoted as fw(x,hj) and fw(Tθ(z;hj),hj), respectively. Since we focus on fitting the meta-distribution with the help of the summary network and OT loss, we leave the problem-specific design of generator, critic, and summary network as future work for considerable flexibility in architectures. We defer to Appendix A the illustration of our proposed model in Fig. 3 and detailed algorithm in Algorithm 2.

heading: 4 RELATED WORK
text: Learning Summary Representation of Set-input. There are two lines for learning the set representation. The first line aims to design more powerful summary networks, which are reviewed in Introduction and Section 2.1 and omitted here due to the limited space. The another line assumes a/some to-be-learned reference set(s), and optimizes the distance between original sets (or features of the observed data points) and the reference set(s) with OT theory or other distance measures, to learn set representation. For example, RepSet (Skianis et al., 2020) computes some comparison costs between the input sets and some to-be-learned reference sets with a network flow algorithm, such as bipartite matching. These costs are then used as set representation in a subsequent neural network. However, unlike our framework, RepSet does not allow unsupervised learning and mainly focuses on classification tasks. The optimal Transport Kernel Embedding (OTKE) (Mialon et al.,
2021) marries ideas from OT theory and kernel methods (Schölkopf et al., 2002), and aligns features of a given set to a trainable reference distribution. Wasserstein Embedding for Graph Learning (WEGL) (Kolouri et al., 2021) also uses a similar idea to the linear Wasserstein embedding as a pooling operator for learning from sets of features. To our best knowledge, both of them view the reference distribution as the barycenter and compute the set-specific representation by aggregating the features (embedded with kernel methods) in a given set with adaptive weight, defined by the transport plan between the given set and the reference. Different from them, we assume J probability distributions (rather than one reference distribution with an uniform measure) over these shared prototypes by taking set representations h1:J as the measures, to approximate the corresponding J empirical distributions, respectively. Then we naturally use the summary network as the encoder to compute h1:J , which can be jointly optimized with the shared prototypes by minimizing the general OT loss in an unsupervised way. For a given set, we can directly compute its representation with summary network, avoiding iteratively optimizing the transport plan between given set and learned reference like OTKE and WEGL. These differences between them and ours lead to different views of set representation learning and different frameworks as well. To learn compact representations for sequential data, Cherian & Aeron (2020) blend contrastive learning, adversarial learning, OT, and Riemannian geometry into one framework. However, our work directly minimises the OT distance between empirical distribution and the to-be-learned distribution, without contrastive learning and explicit generative processes from original set to the adversarial noise set.
Metric-based few-shot classification methods. Our method has a close connection with metricbased few-shot classification algorithms. For example, MatchingNet (Vinyals et al., 2016) and ProtoNet (Snell et al., 2017) learned to classify samples by computing distances to representatives of each class. Using an attention mechanism over a learned embedding of the support set to predict classes for the query set, MatchingNet (Vinyals et al., 2016) can be viewed as a weighted nearestneighbor classifier applied within an embedding space. ProtoNet (Snell et al., 2017) takes a class’s prototype to be the mean of its support set in the learned embedding space, which further performs classification for an embedded query point by finding the nearest class prototype. Importantly, the global prototypes in our paper are shared among all sets, which is different from the specific prototype for each class in ProtoNet but suitable for our case. Due to the flexibility of our method, we can project the average aggregated feature vector derived from the encoder in ProtoNet or MatchingNet, into hj by introducing a simple neural network, which can be jointly optimized with the encoder by minimizing the classification loss and OT loss. Our novelty is that OT loss can be naturally used to improve the learning of encoder without changing the existing model architectures or algorithms a lot. Another recent work for learning multiple centres is infinite mixture prototypes (IMP) (Allen et al., 2019), which represents each class by a set of clusters and infers the number of clusters with Bayesian nonparametrics. However, in our work, the centers are shared for all classes and set-specific feature extracted from the summary network serves as the proportion of centers, where the centers and summary network can be jointly learned with OT loss.
Meta GAN-based Models. As discussed by Hong et al. (2020a), meta GAN-based models can be roughly divided into optimization-based, fusion-based, and transformation-based methods. Clouâtre & Demers (2019) and Liang et al. (2020) integrated GANs with meta-learning algorithms to realize the optimization-based methods, including model-agnostic meta-learning (MAML) (Finn et al., 2017) and Reptile (Nichol et al., 2018). Hong et al. (2020b;c) fused multiple conditional images by combining matching procedure with GANs, providing fusion-based methods. For transformation based methods, Antoniou et al. (2017) and Hong et al. (2020a) combine only one image and the random noise into the generator to produce a slightly different image from the same category, without using the multiple images from same category. Besides, a recent work that connects existing summary network with GAN is MetaGAN (Zhang et al. 2018), which feeds the output of the summary network into the generator and focuses on few-shot classification using MAML. The key difference of these models from ours is that we adopt the OT distance to optimize the summary network for capturing each sets’ summary statistics, which are assimilated into the existing GAN-based models as conditional information. In this general framework, the summary network, generator, and discriminator can be flexibly chosen for specific tasks.

heading: 5 EXPERIMENTS
text: We conduct extensive experiments to evaluate the performance of our proposed prototype-oriented framework for meta-learning. We show that our method outperforms state-of-art baselines, including
improving summary networks, few-shot generation, and few-shot classification. Unless specified otherwise, we set the hyper-parameter of entropic constraint as = 0.1, the maximum iteration number in Sinkhorn algorithm as Itermax = 200 for all experiments and adopt the Adam optimizer (Kingma & Ba, 2015) with learning rate 0.001. We repeat all experiments 5 times and report the mean and standard deviation on corresponding test datasets.

heading: 5.1 EXPERIMENTS ABOUT OT DISTANCE IN SUMMARY NETWORK
text: To evaluate the effectiveness of OT loss in improving the summary network designed for setstructured input, we conduct three tasks on two classical architectures: DeepSets (Zaheer et al. 2017) and Set Transformer (Lee et al. 2019), where the former uses standard feed-forward neural networks and the latter adopts the attention-based network architecture. For Set Transformer and DeepSets, the summary network is defined in Equation 1 and optimized by the task-specific loss; for Set Transformer(+OT) and DeepSets(+OT), the summary network, defined as in Equations 1 and 7, is optimized with matrix B by both the OT loss and task-specific loss. We provide more details about toy datasets, network architectures, and training schemes in Appendix B.
Amortized Clustering with Mixture of Gaussians (MoGs): We consider the task of maximum likelihood of MoGs with C components, denoted as P (x;θ) = ∑C c=1 πcN ( x | µc,diag ( σ2c ))
. Given the dataset X={x1:n} generated from the MoG, the goal is to train a neural network, which takes X as input set and outputs parameters θ. Each dataset contains n ∈ [100, 500] points on a 2D plane, each of which is sampled from one ofC Gaussians. Table 1 reports the test average likelihood of different models with varying C ∈ {4, 8}, where we set K = 50 for all C. We observe that Set Transformer outperforms DeepSets largely, validating the effectiveness of attention mechanisms in this task. Besides, both Set Transformer(+OT) and DeepSets(+OT) improve their baselines, showing that the OT loss can encourage the summary networks to learn more efficient summary statistics.
Point Cloud Classification: Here, we evaluate our method on the task of point cloud classification using the ModelNet40 (Chang et al., 2015) dataset 1, which consists of 3D objects from 40 different categories. By treating each object as a point cloud, we represent it as a set of N vectors in R3 (x; y; z-coordinates). Table 1 reports the classification accuracy, where we perform the experiments with varying N ∈ {64, 1024}, set K = 40. Clearly, both of the DeepSets and Set Transformer can be improved by adding the OT loss. Notably, fewer points would lead to lower performance, where the OT loss plays a more important role. Taking this task as the example, we further study our model’s sensitivity to the hyper-parameter in Fig. 4 of the Appendix B.5.
Sum of Digits: Following Zaheer et al. (2017), we aim to compute the sum of a given set of digits, where we consider MNIST8m (Loosli et al., 2007) containing 8 million instances of 28×28 grey-scale stamps of digits in {0, ..., 9}. By randomly sampling a subset of maximum M = 10 images from MNIST8m, we build N = 100k “sets” of training, where we denote the sum of digits in that set as the set-label. We construct 100k sets of test MNIST digits, where we vary the M starting from 10 all the way up to 100. The output of the summary network is a scalar, predicting the sum ofM digits. In this case, we adopt L1 as the task-specific loss and set K = 10. We show the accuracy of digit summation for different algorithms in Fig. 2 and find that the OT loss can enhance the summary networks to
achieve better generalization. In this task, we also explore the convergence rate and the learned transport plan matrix of Sinkhorn algorithm in Fig. 5 of the Appendix B.6.

heading: 5.2 EXPERIMENTS ON FEW-SHOT CLASSIFICATION
text: To explore whether our proposed method can improve the metric-based few-shot classification, we consider two commonly-used algorithms as the baselines, including ProtoNet (Snell et al., 2017) and MatchNet (Vinyals et al., 2016). Denoting the feature extractor in each algorithm as fφ1 , we
1We adopt the point-cloud dataset directly from the authors of Zaheer et al. 2017
consider several popular backbones, including ResNet10 and ResNet34 (He et al., 2016). Recalling the discussions in section 3.2, to enforce fφ1 to learn more powerful image features, we additionally introduce matrix B and net gφ2 and learn the model by minimizing the OT distance and classification errors. We perform the experiments on the CUB (Welinder et al., 2010) and miniImageNet (Ravi & Larochelle, 2016). As a fine-grained few-shot classification benchmark, CUB contains 200 different classes of birds with a total of 11, 788 images of size 84×84×3, where we split the dataset into 100 base classes, 50 validation classes, and 50 novel classes following Chen et al. (2019). miniImageNet is derived from ILSVRC-12 dataset (Russakovsky et al., 2015), consisting of 84 × 84 × 3 images from 100 classes with 600 random samples in each class. We follow the splits used in previous work (Ravi & Larochelle, 2016), which splits the dataset into 64 base classes, 16 validation classes, and 20 novel classes. Table 2 reports the 5way5shot and 5way10shot classification results of different methods on miniImageNet and CUB. We see that introducing the OT distance and summary network can consistently improve over baseline classifiers, and the performance gain gradually increases with the development of number of network layers. This suggests that our proposed plug-and-play framework can be flexibly used to enhance the metric-based few-shot classification, without the requirement of designing complicated models on purpose.

heading: 5.3 EXPERIMENTS ON FEW-SHOT GENERATION
text: Here, we consider few-shot generation task to investigate the effectiveness of our proposed implicit meta generative framework, where we consider CGAN (Mirza & Osindero, 2014) and DAGAN (Antoniou et al., 2017) as baselines for their ability to generate conditional samples. For CGANbased models, we adopt summary network as the encoder to extract the feature vector hj from a given setDj , which is further fed into the generator and discriminator as the conditional information. Since the original DAGAN only assimilates one image into the generator, in our framework, we replace the encoder in DAGAN with summary network to learn the set representation. For DAGANbased models, we adopt the same way with the original DAGAN to construct the real/samples for the critic and explain here for clarity: we sample two sets (D1j , D2j ) from the same distribution or category; then we represent the real samples as the combination of D1j and D2j and the fake ones as the combination of D1j and D̂1j from the generator (conditioned on D1j). Different from our framework that separately optimizes the summary network using the OT loss, all other models for comparison in this paper jointly optimize the encoder (e.g., summary network) with the generator by the generator loss. Besides, for a fair comparison, we also consider introducing the additional reconstruction loss (mean square error, MSE) to optimize the generator and encoder in baselines. We consider the DeepSets as the summary network for its simple architecture. For 3D natural images,
we adopt the pretrained densenet (Iandola et al., 2014) to extract features from each data pints, and take the features as the input to summary network. To evaluate the quality of generated samples, we adopt commonly used metric Fréchet Inception Distance (FID) (Heusel et al., 2017), where we only report the FID score (Heusel et al., 2017) considering the notable performance gap between our model and the compared ones. We provide several examples on toy datasets to show the efficiency of our proposed model in Appendix D.
Rotated MNIST: Following Wu et al. (2020), we artificially transform each image in MNIST dataset (LeCun, 1998) with 18 rotations (-180 to 180 by 20 degrees), leading to 18 distributions characterized by angle A. We choose 9 interleaved distributions for training and the rest as unseen distributions for testing. We consider the CGAN-based and DAGAN-based models, respectively. During the test stage, for each unseen distribution, we randomly sample 1000 real images and generate 20 fake images based on every 20 real images and repeat this process 50 times (i.e., 1000/20), resulting in 1000 generated samples for each method. We summarize the test performance in Table 3 with varying A. We can find that our proposed framework allows for better generalization to related but unseen distributions at test time, indicating the OT loss can enforce the summary network to capture more salient characteristics.
Natural Images: We further consider few-shot image generation on Flowers (Nilsback & Zisserman, 2008) and Animal Faces (Deng et al., 2009), where we follow seen/unseen split provided in Liu et al. (2019). Flowers dataset contains 8189 images of 102 categories, which are divided into 85 training seen and 17 testing unseen categories; Animal Faces dataset contains 117, 574 animal faces collected from 149 carnivorous animal categories, which are split into 119 training seen and 30 testing unseen categories. We present the example images generated by DAGAN and DAGAN(+OT) and network architectures in Appendix E for the limited space, where we also compute the FID scores with the similar way in Rotated MNIST experiment. We find that our method achieves the lowest FID and has the ability to generate more realistic natural images compared with baselines. This indicates the summary network in our proposed framework can successfully capture the important summary statistics within the set, beneficial for the few-shot image generation.

heading: 6 CONCLUSION
text: In this paper, we present a novel method to improve existing summary networks designed for setstructured input based on optimal transport, where a set is endowed with two distributions: one is the empirical distribution over the data points, and another is the distribution over the learnable global prototypes. Moreover, we use the summary network to encode input set as the prototype proportion (i.e., set representation) for global centers in corresponding set. To learn the distribution over global prototypes and summary network, we minimize the OT distance between two distributions in terms of the defined cost function. Only additionally introducing the acceptable parameters, our proposed model provides a natural and unsupervised way to improve the summary network. In addition to the set-input problems, our plug and play framework has shown appealing properties that can be applied to many meta-learning tasks, where we consider the cases of metric-based few-shot classification and implicit meta generative modeling. Extensive experiments have been conducted, showing that our proposed framework achieves state-of-the-art performance on both improving existing summary networks and meta-learning models for set-input problems. Due to the flexibility and simplicity of our proposed framework, there are still some exciting extensions. For example, an interesting future work would be to apply our method into approximate Bayesian computation for posterior inference.

heading: A ALGORITHMS AND ILLUSTRATION OF OUR PROPOSED MODEL
text: The pseudo code for the implicit meta generative modeling is provided in Algorithm 2.

heading: B EXPERIMENTAL SETTINGS ABOUT INTRODUCING OT DISTANCE INTO THE SUMMARY NETWORKS
text: B.1 DETAILS FOR AMORTIZED CLUSTERING WITH MIXTURES OF GAUSSIANS
We generate the 2D toy datasets following the Lee et al. (2019), where we additionally vary the C (the number of components) from 4 to 8. Below, we present the detailed generation process about the toy datasets:
1. Specify the number of components C for 2D toy dataset. 2. Generate the number of data points, n ∼ Uniform(100, 500). 3. Sample the mean vector for C components.
µc,d ∼ Uniform(−4, 4), c = 1, . . . , C, d = 1, 2
4. Sample the cluster labels. π ∼ Dir ( 1> ) , zi ∼ Categorical(π), i = 1, . . . , n, zi = 1, . . . , C
5. Generate data from spherical Gaussian. xi ∼ N ( µzi , (0.3) 2I ) , i = 1, . . . , n
B.2 DETAILS ABOUT SET-TRANSFORMER-BASED AND DEEPSETS-BASED ARCHITECTURES USED IN MOGS EXPERIMENTS
DeepSets In terms of the DeepSets, the fφ1 in summary network contains 3 permutation-equivariant layers with 256 channels followed by mean-pooling over the set structure. Then the resulting vector representation zj of the set is then fed to a fully connected layer with 512 units followed by a linear layer 512 × C(1 + 2 ∗ 2), where C denotes the number of components. We use ELU activation at all layers. To introduce the OT loss into the DeepSets, we further feed the zj into a fully connected layer with 512 units followed by a 50-way softmax unit and also introduce the global matrix B ∈ R2×50. Set Transformer To perform the MoGs experiments, we adopt the same architecture for Set Transformer following Lee et al. (2019), whose parameters are reported in Table 4. To introduce the OT loss, we also add the two fully connected layers with 256 units on the resulting vector zj followed by a 50-way softmax unit, and a global prototype matrix B ∈ R2×50.
B.3 DETAILS ABOUT SET-TRANSFORMER-BASED AND DEEPSETS-BASED ARCHITECTURES USED IN SUM OF DIGITS
DeepSets Following the official code in Zaheer et al. (2017), we adopt the default architecture to implement the DeepSets, where we first project the image into a 128-dimensional vector with three convolutional layers and apply summary network on the 128-dimensional vectors. To build DeepSets(+OT), we take the set representation zj after sum-pooling in summary network as the input and introduce fully connected layer with 128 units followed by a 10-way softmax unit, and a global prototype matrix B ∈ R128×10. Set Transformer For Set Transformer, we follow the similar structure used in MoGs experiments, where we also project the image with three convolutional layers and output a scalar. To build Set Transformer(+OT), we take the representation zj after sum-pooling in summary network as the input and introduce fully connected layer with 128 units followed by a 10-way softmax unit, and a center matrix B ∈ R128×10.
B.4 DETAILS ABOUT SET-TRANSFORMER-BASED AND DEEPSETS-BASED ARCHITECTURES USED IN POINT CLOUD CLASSIFICATION
DeepSets For original DeepSets, we adopt the same architecture with Zaheer et al. (2017). In a specific, the fφ1 in summary network contains 3 permutation-equivariant layers with 256 channels followed by max-pooling over the set structure. Then the resulting vector representation zj of the set is then fed to a fully connected layer with 256 units followed by a 40-way softmax unit. We
use Tanh activation at all layers and dropout on the layers after set-max-pooling (i.e. two dropout operations) with 50% dropout rate. To introduce the OT loss into the DeepSets, we further feed the zj into a fully connected layer with 256 units followed by a 40-way softmax unit, with 70% dropout rate. Besides, we additionally introduce the center matrix B ∈ R3×40. Set Transformer We also adopt the same architecture to implement the Set Transformer, where we summarize the parameters in Table 5, following Lee et al. (2019). To improve the Set Transformer with OT distance, we also introduce a fully connected layer with 256 units followed by a 40-way softmax unit, with 90% dropout rate, and a center matrix B ∈ R3×40.
B.5 PARAMETER SENSITIVITY
In the previous experiments, we fix the value of as 0.1, controlling the weight of the entropic regularisation in the Sinkhorn algorithm. Notably, unless specified otherwise, we specify the construction of C as Cik = 1− cos (xji,βk). Therefore, the cost function provides an upper-bounded positive similarity metric, making the has the corresponding reasonable range as a prior knowledge. Here, we study our DeepSets(+OT)’s sensitivity to . We consider the point cloud classification task and each object is represented as a set of N = 20 vectors. As shown in Fig. 4, we report the performance of DeepSets(+OT) on point cloud classification task with varying , where DeepSets serves as the baseline. It can be seen that our model is robust to the . Besides, all the results of DeeepSets(+OT) with different are superior than that of DeeepSets, indicating the effectiveness of our method. By fine-tuning for each dataset in each task, we might obtain better results than those reported in our experiments. However, we aim to validate our method instead of exhaustively tuning this hyper-parameter and thus we set = 0.1, which can achieve the acceptable result.
B.6 CONVERGENCE RATE OF SINKHORN ALGORITHM
In this paper, we set the maximum iteration number as Itermax = 200 in Sinkhorn algorithm for all experiments. As shown in Fig. 5, we visualize the convergence rate of Sinkhorn algorithm, where we consider the task about “sum of digits” (DeepSets+OT). The upper figure shows the convergence rate of Sinkhorn algorithm. The bottom figure visualizes the transport plan matrix with varying iterations. We find that the 200 iterations are typically enough for Sinkhorn algorithm and we can learn a sparse transport plan matrix T when the algorithm converge. Notably, the transport plan matrix needs to satisfy two marginal constraints, defined by the probability measures of two distributions, respectively. Recall that the empirical distribution has an unchanged uniform probability measure, so the learned transport plan matrix is dense for the Nj observed samples. In terms of another dis-
matrix T = diag
(
u(`)
)
Kdiag
(
v(`)
)
computed at iteration of Sinkhorn’s iterations.
tribution, its probability measure is the set-specific representation, weighting the importance of K shared centers for corresponding set. Therefore, it is reasonable that transport plan matrix is sparse for K centers.

heading: C EXPERIMENTAL SETTINGS ABOUT FEW-SHOT CLASSIFICATION
text: Denote the prototype for set j (computed by fφ1 ) in few-shot classification as cj . We consider two backbones for fφ1 , including ResNet10 and ResNet34, which produce the 512-dimensional cj . To improve the metric-based few-shot classification with our framework, taking the cj as input, we further construct the gφ2 . Specifically, we introduce a fully connected network with architecture as 512 → 256 units with ReLU function followed by a X-way (X=64 for CUB, and X=128 for miniImageNet) Softmax function and a center matrix B ∈ R512×X. We conduct 10000 tasks of the training set Dtr to train the model while 1000 tasks of the test set Dte to evaluate the learned model. And Dtr ∩ Dte = ∅. We run 60 epochs to train the model on CUB and miniImageNet. The model is trained using Adam optimizer with default settings (learning rate 1e − 3, β = (0.9, 0.999), and = 1e− 8) on one Nvidia Geforce RTX3090 GPU.

heading: D ADDITIONAL EXPERIMENTAL RESULTS ON FEW-SHOT GENERATION ABOUT TOY DATASETS
text: We test our algorithm through a series of synthetic data sets and realistic data sets. For synthetic datasets, we set Tθ , fw and Sφ as fully connected neural networks, where Tθ , fw have 4 hidden layers and fφ1 and gφ2 (we adopt DeepSets) have 3 hidden layers. Each layer has 200 nodes, and the activation function is chosen as RELU, where we adopt the softmax in the final layer.
Normal distribution on 2D toy data: We first consider the 2D normal case, where the training data contains 10K sets and each set contains 100 data points from N(µ,Σ). We sample the mean, variance, and covariance from U [−5, 5], U [1, 2], and U [−0.5, 0.5], respectively. Fig. 6 shows the real (gray points) and generated samples (red points) by different models given unseen test sets, where we only consider CGAN-based methods for the simple toy data. We find that our model (third column) can improve the resistance to mode collapse compared with CGAN+MSE (second column) and better fit the unseen test distributions than CGAN (first column). This result indicates the OT loss can spur the summary network to capture more desired statistics for unseen distributions.
One-dimensional Gaussian distributions: In this case, we generate another collection of synthetic 1 − D datasets based on Gaussian parametric family, where the means and variances are sampled from U[−1, 1] and U[0.5, 2] respectively. The training data contains 10K sets each containing 50 samples. We also visualize the pdfs of gaussian distributions with randomly means and variance in Fig 7, which are used to sample test data, and show the 500 data points generated by the pushforward. For this experiment, we set the dimension of z and summary vector s as 2.
Multi-family distribution on 1D toy data: To validate if our proposed model can capture many types of distributional families simultaneously, we construct a collection of synthetic 1-D datasets each containing 100 samples from either an Exponential, Gaussian or Laplacian distribution with equal probability. For Gaussian and Laplacian distributions, means and variances are sampled from U [−1, 1] and U [0.5, 2] respectively; for Exponential distributions, rates are sampled from U [0.5, 2]. Fig. 8 visualizes the pdfs of six one-dimensional test distributions with different means and variances and the generated data points. It is interesting to observe that the generated data points can fit the corresponding pdf well, indicating our model can generalize to different distributions with varying parameters. Besides, our model performs slightly worse on the Exponential distributions, perhaps attributing to the fact that it is the only non-symmetric distribution.

heading: E DETAILS ABOUT NATURAL IMAGE GENERATIONS AND THE RESULTS
text: We use denseNet proposed by Huang et al. (2018) as the backbone of summary network, then a pooling operation is conducted as Zaheer et al. (2017) does. And a 2-layer fully connected network with ReLU activation function is finally employed to embed the 4096-D visual features
into the corresponding 512-D set representations hj . As for conditional generator (conditioned on set representations as well as Gaussian noise), we introduce a 2-layer embedding network [100→600→100] with LeakyReLU activation function to embed the input noise n. Besides, we use a 5-layer deconvolution network [ConvTranspose2d(X + 100, 512, 4, 1, 0)→ ConvTranspose2d(512, 256, 4, 2, 1)→ConvTranspose2d(512, 128, 4, 2, 1)→ConvTranspose2d(128, 64, 4,
2, 1)→ConvTranspose2d(64, 3, 4, 2, 1)], where X = 64, 128 for oxford and animal face datasets respectively with BatchNorm along channels and ReLU activation function to deconvolute the concatenated noise embeddings and set representations cat([n,hj ]) as fake output images xfake : RN×3×64×64. Finally, a 5-layer discriminator network [Conv2d(2*3, 64, 4, 2, 1)→Conv2d(64, 128, 4, 2, 1))→Conv2d(128, 256, 4, 2, 1)→Conv2d(256, 512, 4, 2, 1)→Conv2d(512, 1, 4, 1, 0)] with BatchNorm as well as LeakyReLU activation function at the first fourth deconvolutional layers and the last layer without BatchNorm while with sigmoid activation funtion to distinguish the true or fake generated images. We present the generated results in Figure 9.

heading: F THE DIFFERENCE BETWEEN OUR MODEL AND BARYCENTER PROBLEM
text: In this section, we clarify the difference between Wasserstein barycenter and our method. Specifically, for j-th distribution, we denoted Pj as its empirical distribution consisting of Nj samples, expressed as Pj = ∑Nj i=1 1 Nj δxji ,xji ∈ Rd. Notably, aj = [ 1Nj ] ∈ Σ
Nj represents the probability measure for distribution Pj . For another thing, we can represent Pj with another to-be-learned distribution Qj , defined as Qj = ∑K k=1 hjkδβk ,βk ∈ Rd. Here bj = [hjk] ∈ ΣK is the probability measure for distribution Qj , which can be computed using summary network Sφ(xi,1:Nj ) and serves as the representation for set j. And βk is the k-th prototype in the same space of the observed data points, which is the k-th column of B ∈ RK×d, a learnable global prototype matrix. To optimize the β1:K and the summary network Sφ for computing hj , we minimize the average OT loss (between Qj and Pj) for all training sets. We rewrite the Equation (6) here for convenience:
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)) .
Usually, this equation can also be represented:
min B,φ
1
J J∑ j=1 (OT (aj , bj)) . (9)
In terms of Wasserstein barycenter, we adopt the same notations for consistency. Following (Altschuler & Boix-Adserà, 2021), given J empirical distributions P1:J and their respective probability measures [a1, . . . ,aJ ] supported on Rd and a vector λ ∈ ΣJ , their corresponding Wasserstein
barycenter can be viewed as another distribution Q, i.e., Q = ∑K k=1mkδβk ,βk ∈ Rd, where βk is the k-th column of B, m = [mk] ∈ ΣK is the probability measure for distribution Q. Then we can learn the barycenter (i.e., B andm) by minimizing
min B,m
1
J J∑ j=1 λjW (Pj , Q) = min B,m 1 J J∑ j=1 λjW (aj ,m) (10)
where aboveW(·, ·) denotes the squared 2-Wasserstein distance. By comparing the equation 10 and equation 9, we can find that our model learns a Qj to approximate Pj for each distribution j but “barycenter” problem learns a shared Q as the barycenter for all P1:J . Therefore, after minimizing the average loss for all training sets with Equation (6), we can use the probability measure bj = [hjk] (i.e., hj) to represent the empirical distribution Pj . Especially, we can directly map the test set to its representation by using the summary network. However, since the probability measure m in “barycenter” is shared by all distributions, it can not represent a specific set. Therefore, to achieve the set representation, it might need to first compute the transport plan between test set and the barycenter Q and then aggregate the data points (or features) within the test set by taking the transport plan as the weight. Therefore, our model produces a more intuitive solution to learn the set representation, which can take full advantage of the existing summary networks and provides a promising tool for addressing set-input and meta-learning problems.

Review summaries:
 This work proposes a method to improve set representation, and applies it in the context of meta-learning. More precisely, the method consists in jointly learning a summary network and prototypes using an optimal transport loss: the prototypes and summary network output should minimize the sum of the optimal transport costs for all the sets of the meta dataset. Experiments are conducted in the context of few-shot learning and other tasks used to evaluate summary networks.

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. 

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.