Original: L EARNING P ROTOTYPE - ORIENTED S ET R EPRESENTATIONS FOR M ETA -L EARNING
Revised: L EARNING P ROTOTYPE - ORIENTED S ET R EPRESENTATIONS FOR M ETA -L EARNING

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: 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.
Revised: 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.

Original: In fact, many meta-learning problems can be treated as set-input tasks.
Revised: In fact, many meta-learning problems can be treated as set-input tasks.

Original: Most existing summary networks aim to design different architectures for the input set in order to enforce permutation invariance.
Revised: Most existing summary networks aim to design different architectures for the input set in order to enforce permutation invariance.

Original: However, scant attention has been paid to the common cases where different sets in a metadistribution are closely related and share certain statistical properties.
Revised: However, scant attention has been paid to the common cases where different sets in a metadistribution are closely related and share certain statistical properties.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Since our plug-and-play framework can be applied to many meta-learning problems, we further instantiate it to the cases of few-shot classiﬁcation and implicit meta generative modeling.
Revised: Since our plug-and-play framework can be applied to many meta-learning problems, we further instantiate it to the cases of few-shot classiﬁcation and implicit meta generative modeling.

Original: Extensive experiments demonstrate that our framework signiﬁcantly improves the existing summary networks on learning more powerful summary statistics from sets and can be successfully integrated into metric-based few-shot classiﬁcation and generative modeling applications, providing a promising tool for addressing set-input and meta-learning problems.
Revised: Extensive experiments demonstrate that our framework signiﬁcantly improves the existing summary networks on learning more powerful summary statistics from sets and can be successfully integrated into metric-based few-shot classiﬁcation and generative modeling applications, providing a promising tool for addressing set-input and meta-learning problems.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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).
Revised: 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).

Original: However, extending them to handle unstructured input in the form of sets, where a set can be deﬁned as an unordered collections of elements, is not trivial and has recently attracted increasing attention (Jurewicz & StrømbergDerczynski, 2021).
Revised: However, extending them to handle unstructured input in the form of sets, where a set can be deﬁned as an unordered collections of elements, is not trivial and has recently attracted increasing attention (Jurewicz & StrømbergDerczynski, 2021).

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: Therefore, we broaden the scope of set related applications by including traditional set-structured input problems and most meta-learning problems.
Revised: Therefore, we broaden the scope of set related applications by including traditional set-structured input problems and most meta-learning problems.

Original: Both of them aim to improve the quick adaptation ability for unseen sets, even though the latter is more difﬁcult because of limited samples or the occurrence of new categories for classiﬁcation problems.
Revised: Both of them aim to improve the quick adaptation ability for unseen sets, even though the latter is more difﬁcult because of limited samples or the occurrence of new categories for classiﬁcation problems.

Original: 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.
Revised: 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.

Original: 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 ﬁxed-size output.
Revised: 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 ﬁxed-size output.

Original: For example, the prominent works of Zaheer et al.
Revised: For example, the prominent works of Zaheer et al.

Original: 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).
Revised: 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).

Original: Lee et al.
Revised: Lee et al.

Original: (2019) further introduced Set Transformer to encode and aggregate the features within the set using the multi-head attention mechanism.
Revised: (2019) further introduced Set Transformer to encode and aggregate the features within the set using the multi-head attention mechanism.

Original: Maron et al.
Revised: Maron et al.

Original: (2020) designed deep models and presented a principled approach to learn sets of symmetric elements.
Revised: (2020) designed deep models and presented a principled approach to learn sets of symmetric elements.

Original: 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-speciﬁc loss function, which could limit the models’ ﬂexibility.
Revised: 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-speciﬁc loss function, which could limit the models’ ﬂexibility.

Original: 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 beneﬁt the corresponding set-speciﬁc 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.
Revised: 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 beneﬁt the corresponding set-speciﬁc 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.

Original: These limits still remain even using the recent more carefully designed summary networks, while sets with limited samples further exacerbate the problem.
Revised: These limits still remain even using the recent more carefully designed summary networks, while sets with limited samples further exacerbate the problem.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Speciﬁcally, for the j th set, we consider it as a discrete distribution P j over all the samples within the set (in data or feature space).
Revised: Speciﬁcally, for the j th set, we consider it as a discrete distribution P j over all the samples within the set (in data or feature space).

Original: At the same time, we also represent this set with another distribution Q j (in the same space with P j ), supported on K global prototypes with a K -dimensional set representation h j .
Revised: At the same time, we also represent this set with another distribution Q j (in the same space with P j ), supported on K global prototypes with a K -dimensional set representation h j .

Original: Since h j 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 .
Revised: Since h j 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 .

Original: Moreover, the existing summary networks can be adopted to encode set j as h j for their desired property of permutation invariance.
Revised: Moreover, the existing summary networks can be adopted to encode set j as h j for their desired property of permutation invariance.

Original: In this way, we can formulate the learning of summary networks as the process of learning a P j to be as close to Q j as possible, a process facilitated by leveraging the optimal transport (OT) distance (Peyr ´ e & Cuturi 2019).
Revised: In this way, we can formulate the learning of summary networks as the process of learning a P j to be as close to Q j as possible, a process facilitated by leveraging the optimal transport (OT) distance (Peyr ´ e & Cuturi 2019).

Original: Therefore, the global prototypes and summary network can be learned by jointly optimizing the corresponding task loss and OT distance between P j and Q j in an end-to-end manner.
Revised: Therefore, the global prototypes and summary network can be learned by jointly optimizing the task-speciﬁc loss and OT distance between P j and Q j in an end-to-end manner.

Original: We can refer to this approach 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 classiﬁcation (Snell et al., 2017), and learningsummary statistics for approximate Bayesian computation (Chen et al., 2021).
Revised: 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 classiﬁcation (Snell et al., 2017), and learning statistics for approximate Bayesian computation (Chen et al., 2021).

Original: 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 classiﬁcation and implicit meta generative modeling.
Revised: 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 classiﬁcation and implicit meta generative modeling.

Original: 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.
Revised: 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.

Original: (2) We leverage optimal transport distance to measure the distance between the distributions for use in a joint learning algorithm.
Revised: (2) We leverage optimal transport distance to measure the distance between the distributions for use in a joint learning algorithm.

Original: We apply our method to metric-based few-shot classiﬁcation and construct implicit meta generative model, where summary network is used to extract the summary statistics from set and optimized by the OT loss.
Revised: We apply our method to metric-based few-shot classiﬁcation and construct implicit meta generative model, where summary network is used to extract the summary statistics from set and optimized by the OT loss.

Original: 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 classiﬁcation and GAN frameworks, producing a new way to learn the set’ summary statistics applicable to many applications.
Revised: 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 classiﬁcation and GAN frameworks, producing a new way to learn the set’ summary statistics applicable to many applications.

Original:  B ACKGROUND 2. S UMMARY NETWORKS FOR SET - STRUCTURED INPUT
Revised:  B ACKGROUND 2. S UMMARY NETWORKS FOR SET - STRUCTURED INPUT

Original: To deal with the set-structured input D j = { x j, 1: N j } 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.
Revised: To deal with the set-structured input D j = { x j, 1: N j } 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.

Original: More formally,
Revised: More formally,

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: Despite their effectiveness, there is no clear evidence that the output of these summary networks could describe a set’s summary statistics well, which have been proven crucial for the set-structured inference problems (Chen et al., 2021) and may be also important to other set-speciﬁc tasks.

Original: 2. O PTIMAL T RANSPORT
Revised: 2. O PTIMAL T RANSPORT

Original: Although OT has a rich theory, we limit our discussion to OT for discrete distributions and refer the reader to Peyr´e & Cuturi (2019) for more details.
Revised: Although OT has a rich theory, we limit our discussion to OT for discrete distributions and refer the reader to Peyr´e & Cuturi (2019) for more details.

Original: Let us consider p and q as two discrete probability distributions on the arbitrary space X ⊆ R d , which can be formulated as p = (cid:80) ni
Revised: Let us consider p and q as two discrete probability distributions on the arbitrary space X ⊆ R d , which can be formulated as p = (cid:80) ni

Original: =1 a i δ x i and q = (cid:80) mj
Revised: =1 a i δ x i and q = (cid:80) mj

Original: =1 b j δ y j .
Revised: =1 b j δ y j .

Original: In this case, a ∈ Σ n and b ∈ Σ m , where Σ m denotes the probability simplex of R m .
Revised: In this case, a ∈ Σ n and b ∈ Σ m , where Σ m denotes the probability simplex of R m .

Original: The OT distance between a and b is deﬁned as
Revised: The OT distance between a and b is deﬁned as

Original: [Equation]
Revised: [Equation]

Original: where (cid:104)· , ·(cid:105) means the Frobenius dot-product; C ∈ R n × m ≥is the transport cost function with element
Revised: where (cid:104)· , ·(cid:105) means the Frobenius dot-product; C ∈ R n × m ≥is the transport cost function with element

Original: C ij = C
Revised: C ij = C

Original: ( x i , y j ) ; T ∈ R n × m >denotes the doubly stochastic transport probability matrix such that
Revised: ( x i , y j ) ; T ∈ R n × m >denotes the doubly stochastic transport probability matrix such that

Original: U ( a , b ) :=
Revised: U ( a , b ) :=

Original: { T | (cid:80) ni T ij = b j , (cid:80) mj T ij = a i } .
Revised: { T | (cid:80) ni T ij = b j , (cid:80) mj T ij = a i } .

Original: To relax the time-consuming problem when optimising the OT distance, Cuturi (2013) introduced the entropic regularization, H = − (cid:80) ij T ij ln T ij , leading to the widely-used Sinkhorn algorithm for discrete OT problems.
Revised: To relax the time-consuming problem when optimising the OT distance, Cuturi (2013) introduced the entropic regularization, H = − (cid:80) ij T ij ln T ij , leading to the widely-used Sinkhorn algorithm for discrete OT problems.

Original:  P ROPOSED FRAMEWORK
Revised:  P ROPOSED FRAMEWORK

Original: In meta-learning, given a meta-distribution p M of tasks, the marginal distribution p j of task j is sampled from p M for j ∈ J , where J denotes a ﬁnite set of indices.
Revised: In meta-learning, given a meta-distribution p M of tasks, the marginal distribution p j of task j is sampled from p M for j ∈ J , where J denotes a ﬁnite set of indices.

Original: For example, we can sample p j from p M with probability 1 J when p M is uniform over a ﬁnite number of marginals.
Revised: For example, we can sample p j from p M with probability 1 J when p M is uniform over a ﬁnite number of marginals.

Original: During meta-training, direct access to the distribution of interest p j is usually not available.
Revised: During meta-training, direct access to the distribution of interest p j is usually not available.

Original: Instead, we will observe a set of data points D j = { x ji } N j i =1 , which consists of N j i.i.d samples from p j over R d .
Revised: Instead, we will observe a set of data points D j = { x ji } N j i =1 , which consists of N j i.i.d samples from p j over R d .

Original: We can roughly treat the meta-learning problems as the set-input tasks, where dataset D j from p j corresponds to an input set.
Revised: We can roughly treat the meta-learning problems as the set-input tasks, where dataset D j from p j corresponds to an input set.

Original: In this paper, 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.
Revised: 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.

Original: Besides, we also provide the applications to metric-based few-shot classiﬁcation and implicit generative framework by assimilating the summary statistics.
Revised: Besides, we also provide the applications to metric-based few-shot classiﬁcation and implicit generative framework by assimilating the summary statistics.

Original: Below we describe our model in detail.
Revised: Below we describe our model in detail.

Original: 3. L EARNING GLOBAL PROTOTYPES AND SET REPRESENTATION VIA OPTIMAL TRANSPORT
Revised: 3. L EARNING GLOBAL PROTOTYPES AND SET REPRESENTATION VIA OPTIMAL TRANSPORT

Original: Given J sets from meta-distribution p M , we can represent each set D j from meta-distribution p M as an empirical distribution over N j samples on the original data space, formulated as
Revised: Given J sets from meta-distribution p M , we can represent each set D j from meta-distribution p M as an empirical distribution over N j samples on the original data space, formulated as

Original: [Equation]
Revised: [Equation]

Original: Since all sets (distributions) draw from meta-distribution p M are closely related, it is reasonable to assume that these sets share some statistical information.
Revised: Since all sets (distributions) draw from meta-distribution p M are closely related, it is reasonable to assume that these sets share some statistical information.

Original: Motivated by dictionary learning and topic modeling, we deﬁne the shared information as the learnable global prototype matrix B = { β k } ∈ R d × 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).
Revised: Motivated by dictionary learning and topic modeling, we deﬁne the shared information as the learnable global prototype matrix B = { β k } ∈ R d × 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).

Original: Given the global prototype matrix B , each set j can be represented with a
Revised: Given the global prototype matrix B , each set j can be represented with a

Original: K -dimensional weight vector h
Revised: K -dimensional weight vector h

Original: j ∈ Σ k
Revised: j ∈ Σ k

Original: ( e.g. , “topic proportion” in topic modeling), where h jk indicates the weight of the prototype β k for set j .
Revised: ( e.g. , “topic proportion” in topic modeling), where h jk indicates the weight of the prototype β k for set j .

Original: Therefore, we can represent set D j with another distribution Q j on global prototypes β 1: K , deﬁned as
Revised: Therefore, we can represent set D j with another distribution Q j on global prototypes β 1: K , deﬁned as

Original: [Equation]
Revised: [Equation]

Original: where h j can be viewed as set representation for describing set j .
Revised: where h j can be viewed as set representation for describing set j .

Original: Since set j can be represented as
Revised: Since set j can be represented as

Original: Q j and P j , we propose to learn set-speciﬁc representation
Revised: Q j and P j , we propose to learn set-speciﬁc representation

Original: h j and K global prototypes B by pushing
Revised: h j and K global prototypes B by pushing

Original: [Figure]
Revised: Q j towards P j :

Original: [Equation]
Revised: [Equation]

Original:  where C ∈ R N j × K ≥is the transport cost matrix.
Revised: 5) where C ∈ R N j × K ≥is the transport cost matrix.

Original: In this paper, to measure the distance between data point x ji in set j and prototype β k , unless speciﬁed otherwise, we specify the construction of C as
Revised: In this paper, to measure the distance between data point x ji in set j and prototype β k , unless speciﬁed otherwise, we specify the construction of C as

Original: C ik = 1 − cos ( x ji , β k ) , for providing an upper-bounded positive similarity metric.
Revised: C ik = 1 − cos ( x ji , β k ) , for providing an upper-bounded positive similarity metric.

Original: Besides, the transport probability matrix T ∈ R N j × K >should satisfy Π( a , b ) :
Revised: Besides, the transport probability matrix T ∈ R N j × K >should satisfy Π( a , b ) :

Original: = (cid:8) T | T1 K
Revised: = (cid:8) T | T1 K

Original: = a , T (cid:62) 1 N j = b (cid:9) with T ik = T ( x ji , β k ) , where a = [ 1 N j ] ∈ Σ N j and b = [ h jk ] ∈ Σ K denote the respective probability vectors for distribution P j in Equation 3 and Q j Equation 4. 
Revised: = a , T (cid:62) 1 N j = b (cid:9) with T ik = T ( x ji , β k ) , where a = [ 1 N j ] ∈ Σ N j and b = [ h jk ] ∈ Σ K denote the respective probability vectors for distribution P j in Equation 3 and Q j Equation 4. [Figure]

Original:  Since h j should be invariant to the permutations of the samples in set j , we adopt summary network to encode the set of N j points.
Revised: [Figure] Since h j should be invariant to the permutations of the samples in set j , we adopt summary network to encode the set of N j points.

Original: For unsupervised tasks, taking the summary network in Equationas the example, we can directly add a Softmax activation function into S φ to enforce the simplex constraint in set representation h j , denoted as h j = Softmax( S φ ( D j )) .
Revised: For unsupervised tasks, taking the summary network in Equationas the example, we can directly add a Softmax activation function into S φ to enforce the simplex constraint in set representation h j , denoted as h j = Softmax( S φ ( D j )) .

Original: As shown in Fig.
Revised: As shown in Fig.

Original: 1, given
Revised: 1, given

Original: J sets, to learn the global prototype matrix B and the summary network parameterized by φ , we adopt the entropic constraint (Cuturi, 2013) and deﬁne the average OT loss for all training sets as
Revised: J sets, to learn the global prototype matrix B and the summary network parameterized by φ , we adopt the entropic constraint (Cuturi, 2013) and deﬁne the average OT loss for all training sets as

Original: L OT = min B ,φ
Revised: L OT = min B ,φ

Original: [Equation]
Revised: [Equation]

Original: where (cid:15) is a hyper-parameter for entropic constraint.
Revised: where (cid:15) is a hyper-parameter for entropic constraint.

Original: Algorithm 1 describes the workﬂow of OT distance for improving summary network under unsupervised tasks.
Revised: Algorithm 1 describes the workﬂow of OT distance for improving summary network under unsupervised tasks.

Original: For supervised tasks, set j is denoted as D j = { x j, 1: N j , y j } , where y j is the ground-truth output determined by speciﬁc task.
Revised: For supervised tasks, set j is denoted as D j = { x j, 1: N j , y j } , where y j is the ground-truth output determined by speciﬁc task.

Original: As h j is a normalized weight vector, directly using it to realize the corresponding task may be undesired.
Revised: As h j is a normalized weight vector, directly using it to realize the corresponding task may be undesired.

Original: Denoting z
Revised: Denoting z

Original: j =pool (cid:0)(cid:8) f
Revised: j =pool (cid:0)(cid:8) f

Original: φ 1 ( x j 1 ) , . . .
Revised: φ 1 ( x j 1 ) , . . .

Original: , f φ 1 (cid:0) x jN j (cid:1)(cid:9)(cid:1) , we project it to the following vectors:
Revised: , f φ 1 (cid:0) x jN j (cid:1)(cid:9)(cid:1) , we project it to the following vectors:

Original: [Equation]
Revised: [Equation]

Original: where h j and ˆ y j are responsible for the OT and task-speciﬁc losses, respectively.
Revised: where h j and ˆ y j are responsible for the OT and task-speciﬁc losses, respectively.

Original: Now the summary network parameters ˜
Revised: Now the summary network parameters ˜

Original: φ = { e, λ, φ 1 } andthe K global prototypes B can be learned by jointly optimizing the task-speciﬁc loss (computed by ˆ y j and y j ) and OT loss in Equation 6.
Revised: φ = { e, λ, φ 1 } and global prototypes B can be learned by jointly optimizing the task-speciﬁc loss (computed by ˆ y j and y j ) and OT loss in Equation F.

Original: In summary, minimizing the OT distance between the prototype distribution Q j and empirical distribution P j 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-andplay framework.
Revised: In summary, minimizing the OT distance between the prototype distribution Q j and empirical distribution P j 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.

Original: Therefore, it realizes efﬁcient learning from new sets for both unsupervised and supervised tasks.
Revised: Therefore, it realizes efﬁcient learning from new sets for both unsupervised and supervised tasks.

Original: 3. A PPLICATION TO M ETRIC - BASED F EW - SHOT C LASSIFICATION
Revised: 3. A PPLICATION TO M ETRIC - BASED F EW - SHOT C LASSIFICATION

Original: As a challenging meta-learning problem, few-shot classiﬁcation has recently attracted increasing attention, where one representative method is metric-based few-shot classiﬁcation algorithms.
Revised: As a challenging meta-learning problem, few-shot classiﬁcation has recently attracted increasing attention, where one representative method is metric-based few-shot classiﬁcation algorithms.

Original: Taking the ProtoNet (Snell et al., 2017) as the example, we provide a simple but effective method to improve its classiﬁcation performance with the help of OT distance and summary network, where we refer the reader to ProtoNet for more details.
Revised: Taking the ProtoNet (Snell et al., 2017) as the example, we provide a simple but effective method to improve its classiﬁcation performance with the help of OT distance and summary network, where we refer the reader to ProtoNet for more details.

Original: In a speciﬁc, ProtoNet represents each class by computing an
Revised: In a speciﬁc, ProtoNet represents each class by computing an

Original: M -dimensional representation c j ∈ R M , with an embedding function f φ 1 : R d → R M , where we adopt the same φ 1 as the learnable parameters, following the feature extractor in summary network in Equation 1 for simplicity.
Revised: M -dimensional representation c j ∈ R M , with an embedding function f φ 1 : R d → R M , where we adopt the same φ 1 as the learnable parameters, following the feature extractor in summary network in Equation 1 for simplicity.

Original: Formally, ProtoNet adopts the average pooling to aggregate the embedded features of support points belonging to its class into vector c j = 1 | S j | (cid:80) ( x ji ,
Revised: Formally, ProtoNet adopts the average pooling to aggregate the embedded features of support points belonging to its class into vector c j = 1 | S j | (cid:80) ( x ji ,

Original: y ji ) ∈ S
Revised: y ji ) ∈ S

Original: j f φ 1 ( x ji ) .
Revised: j f φ 1 ( x ji ) .

Original: ProtoNet then compares the distance between a query point f φ 1 ( x ) to the c j in the same embedding space.
Revised: ProtoNet then compares the distance between a query point f φ 1 ( x ) to the c j in the same embedding space.

Original: Motivated by the summary network, we further introduce a feed-forward network g φ 2 to map the c j into the h j used to deﬁne Q j distribution over B .
Revised: Motivated by the summary network, we further introduce a feed-forward network g φ 2 to map the c j into the h j used to deﬁne Q j distribution over B .

Original: Therefore, the functions g φ 2 and f φ 1 can be jointly optimized by minimizing the OT distance between P j and Q j and the original classiﬁcation loss in ProtoNet, where c j is still used to perform classiﬁcation since h j is a normalized weight vector.
Revised: Therefore, the functions g φ 2 and f φ 1 can be jointly optimized by minimizing the OT distance between P j and Q j and the original classiﬁcation loss in ProtoNet, where c j is still used to perform classiﬁcation since h j is a normalized weight vector.

Original: 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 the f φ 1 to learn more powerful representation
Revised: 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 the f φ 1 to learn more powerful representation

Original: c j of each class and feature f φ 1 ( x ) of query sample, beneﬁting the metric-based few-shot classiﬁcation.
Revised: c j of each class and feature f φ 1 ( x ) of query sample, beneﬁting the metric-based few-shot classiﬁcation.

Original: 3. A PPLICATION TO I MPLICIT M ETA G ENERATIVE MODELS
Revised: 3. A PPLICATION TO I MPLICIT M ETA G ENERATIVE MODELS

Original: 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
Revised: 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

Original: h j as summary statistics, where we consider GAN-based implicit models.
Revised: h j as summary statistics, where we consider GAN-based implicit models.

Original: Speciﬁcally, given p j ∼ p M , we aim to construct a parametrized pushforward ( i.e. , generator) of reference Gaussian distribution ρ , denoted as T θ ( · , p j ) (cid:93)ρ , to approximate the marginal distribution p j , where θ summarizes the parameters of pushforward.
Revised: Speciﬁcally, given p j ∼ p M , we aim to construct a parametrized pushforward ( i.e. , generator) of reference Gaussian distribution ρ , denoted as T θ ( · , p j ) (cid:93)ρ , to approximate the marginal distribution p j , where θ summarizes the parameters of pushforward.

Original: Since it is unaccessible to the distribution of interest p j , we replace p j with P j and use the summary network to encode set D j into h j as discussed above, which is further fed into the generator serving as the conditional information, denoted as T θ ( z ; h j ) , z ∼ ρ .
Revised: Since it is unaccessible to the distribution of interest p j , we replace p j with P j and use the summary network to encode set D j into h j as discussed above, which is further fed into the generator serving as the conditional information, denoted as T θ ( z ; h j ) , z ∼ ρ .

Original: To enforce the pushforward T θ ( z ; h j ) to ﬁt the real distribution P j as good as possible, we introduce a discriminator f w following the standard GAN (Goodfellow et al., 2014).
Revised: To enforce the pushforward T θ ( z ; h j ) to ﬁt the real distribution P j as well as possible, we introduce a discriminator f w following the standard GAN (Goodfellow et al., 2014).

Original: Generally, our GAN-based model consists of three components.
Revised: Generally, our GAN-based model consists of three components.

Original: Summary network S φ ( · ) focuses on learning the summary statistics h j by minimizing the OT distance between P j and Q j .
Revised: Summary network S φ ( · ) focuses on learning the summary statistics h j by minimizing the OT distance between P j and Q j .

Original: The pushforward aims to push the combination of a random noise vector z and statistics h j to generate samples that resemble the ones from P j , where we simply adopt a concatenation for h j and z although other choices are also available.
Revised: The pushforward aims to push the combination of a random noise vector z and statistics h j to generate samples that resemble the ones from P j , where we simply adopt a concatenation for h j and z although other choices are also available.

Original: Besides, f w tries to distinguish the “fake” samples from the “real” samples in set D j .
Revised: Besides, f w tries to distinguish the “fake” samples from the “real” samples in set D j .

Original: Therefore, we optimize the implicit meta generative model by deﬁning the objective function as:
Revised: Therefore, we optimize the implicit meta generative model by deﬁning the objective function as:

Original: [Equation]
Revised: [Equation]

Original: In addition to the standard GAN loss, we can also adopt the Wasserstein GAN (WGAN) of Arjovsky et al.
Revised: In addition to the standard GAN loss, we can also adopt the Wasserstein GAN (WGAN) of Arjovsky et al.

Original: (2017) to approximate P j .
Revised: (2017) to approximate P j .

Original: It is also ﬂexibly to decide the input fed into f w .
Revised: It is also ﬂexibly to decide the input fed into f w .

Original: For example, following the conditional GAN (CGAN) of Mirza & Osindero (2014), we can combine h j and data points (generated or real), where the critic can be denoted as f w ( x , h j ) and f w ( T θ ( z ; h j ) , h j ) , respectively.
Revised: For example, following the conditional GAN (CGAN) of Mirza & Osindero (2014), we can combine h j and data points (generated or real), where the critic can be denoted as f w ( x , h j ) and f w ( T θ ( z ; h j ) , h j ) , respectively.

Original: Since we focus on ﬁtting the meta-distribution with the help of the summary network and OT loss, we leave the problem-speciﬁc design of generator, critic, and summary network as future work for considerable ﬂexibility in architectures.
Revised: Since we focus on ﬁtting the meta-distribution with the help of the summary network and OT loss, we leave the problem-speciﬁc design of generator, critic, and summary network as future work for considerable ﬂexibility in architectures.

Original: We defer to Appendix A the illustration of our proposed model in Fig.
Revised: We defer to Appendix A the illustration of our proposed model in Fig.

Original: 4 and detailed algorithm in Algorithm 2.
Revised: 3 and detailed algorithm in Algorithm 2.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Learning Summary Representation of Set-input.
Revised: Learning Summary Representation of Set-input.

Original: We ﬁrst consider summary networks reviewed in  Section 2.1 as the closest line of related works toours. Recently, models for unstructured inputs in the form of sets have rapidly gained attention.
Revised: There are two lines for learning the set representation. The ﬁrst 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.

Original:  For example, Zaheer et al.
Revised: 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 ﬂow 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 classiﬁcation tasks.

Original:  and Edwards & Storkey (2017) introduced the leading approaches by adopting standard feed-forward neural networks and  a set-pooling layer to constrain the permutation invariance.
Revised: The optimal Transport Kernel Embedding (OTKE) (Mialon et al., 2021) marries ideas from OT theory and kernel methods (Sch¨olkopf et al., 2002), and aligns features of a given set to a trainable reference distribution.

Original: Although remarkably simple, these architectures have been proved to be a universal approximator for any set function (Zaheer et al. 2017). Considering attention-based models perform a weighted summation over a set of features, Lee et al.
Revised: 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-speciﬁc representation by aggregating the features (embedded with kernel methods) in a given set with adaptive weight, deﬁned 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 h 1: J as the measures, to approximate the corresponding J empirical distributions, respectively.

Original: (2019) designed Set Transformer to process every element in an input set, which can explicitly model higher-order interactions among the elements in a set.
Revised: Then we naturally use the summary network as the encoder to compute h 1: J , which can be jointly optimized with the shared prototypes by minimizing the general

Original: Wagstaff et al.
Revised: OT loss in an unsupervised way.

Original: (2019) provided a theoretical study on the limitations of representing functions on sets with such networks.
Revised: 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.

Original: Maron et al. (2020) further presented a principled approach to learning sets of symmetric elements.
Revised: To learn compact representations for sequential data, Cherian & Aeron (2020) blend contrastive learning, adversarial learning, OT, and Riemannian geometry into one framework.

Original: Different from them that focus on designing more stronger models, we naturally introduce a general framework to improve existing summary networks based on the OT distance  and global prototypes.
Revised: 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.

Original: Metric-based few-shot classiﬁcation methods.
Revised: Metric-based few-shot classiﬁcation methods.

Original: Recently, researchers proposed simple but effective few-shot classiﬁcation algorithms, where most related work to ours is metric-based few-shot classiﬁcation algorithms.
Revised: Our method has a close connection with metricbased few-shot classiﬁcation algorithms.

Original: 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.
Revised: 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.

Original: 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 nearest-neighbor classiﬁer applied within an embedding space.
Revised: 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 classiﬁer applied within an embedding space.

Original: 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 classiﬁcation for an embedded query point by ﬁnding the nearest class prototype.
Revised: 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 classiﬁcation for an embedded query point by ﬁnding the nearest class prototype.

Original: Importantly, the global prototypes in our paper are shared among all sets, which is different from the speciﬁc prototype for each class in ProtoNet but suitable for our case.
Revised: Importantly, the global prototypes in our paper are shared among all sets, which is different from the speciﬁc prototype for each class in ProtoNet but suitable for our case.

Original: Different from these studies usually compute each class’ prototype by averaging the encoded features of data points within the class and learn the encoder based on the classiﬁcation loss, we further project the average aggregated feature vector to h j and learn thewhole encodernetwork byjointly minimizing the classiﬁcation loss and OT distance.
Revised: Due to the ﬂexibility of our method, we can project the average aggregated feature vector derived from the encoder in ProtoNet or MatchingNet, into h j by introducing a simple neural network, which can be jointly optimized with the encoder by minimizing the classiﬁcation 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 inﬁnite 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-speciﬁc 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.

Original: Meta GAN-based Models.
Revised: Meta GAN-based Models.

Original: As discussed by Hong et al. (2020a), meta GAN-based models can be roughly divided into optimization-based, fusion-based, and transformation-based methods.
Revised: As discussed by Hong et al. (2020a), meta GAN-based models can be roughly divided into optimization-based, fusion-based, and transformation-based methods.

Original: Clou ˆ atre & Demers (2019) and Liang et al.
Revised: Clou ˆ atre & Demers (2019) and Liang et al.

Original: (2020) integrated GANs with meta-learning algorithms to realize the optimization-based methods, including model-agnostic meta-learning (MAML)
Revised: (2020) integrated GANs with meta-learning algorithms to realize the optimization-based methods, including model-agnostic meta-learning (MAML)

Original: (Finn et al., 2017) and Reptile (Nichol et al., 2018).
Revised: (Finn et al., 2017) and Reptile (Nichol et al., 2018).

Original: Hong et al.
Revised: Hong et al.

Original: (2020b;c) fused multiple conditional images by combining matching procedure with GANs, providing fusion-based methods.
Revised: (2020b;c) fused multiple conditional images by combining matching procedure with GANs, providing fusion-based methods.

Original: For transformation based methods, Antoniou et al.
Revised: For transformation based methods, Antoniou et al.

Original: (2017) and Hong et al.
Revised: (2017) and Hong et al.

Original: (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.
Revised: (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.

Original: 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 classiﬁcation using MAML.
Revised: 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 classiﬁcation using MAML.

Original: 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.
Revised: 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.

Original: In this general framework, the summary network, generator, and discriminator can be ﬂexibly chosen for speciﬁc tasks.
Revised: In this general framework, the summary network, generator, and discriminator can be ﬂexibly chosen for speciﬁc tasks.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We conduct extensive experiments to evaluate the performance of our proposed prototype-oriented framework for meta-learning.
Revised: We conduct extensive experiments to evaluate the performance of our proposed prototype-oriented framework for meta-learning.

Original: We show that our method outperforms state-of-art baselines, including improving summary networks, few-shot generation, and few-shot classiﬁcation.
Revised: We show that our method outperforms state-of-art baselines, including improving summary networks, few-shot generation, and few-shot classiﬁcation.

Original: Unless speciﬁed otherwise, we set the hyper-parameter of entropic constraint as (cid:15) = 0 .
Revised: Unless speciﬁed otherwise, we set the hyper-parameter of entropic constraint as (cid:15) = 0 .

Original: 1  for all experiments and adopt the Adam optimizer (Kingma & Ba, 2015) with learning rate 0 .
Revised: 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 .

Revised: We repeat all experiments 5 times and report the mean and standard deviation on corresponding test datasets.

Original: 5. E XPERIMENTS ABOUT OT DISTANCE IN S UMMARY N ETWORK
Revised: 5. E XPERIMENTS ABOUT OT DISTANCE IN S UMMARY N ETWORK

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: For Set Transformer and
Revised: For Set Transformer and

Original: DeepSets, the summary network is deﬁned in Equation 1 and optimized by the task-speciﬁc loss; for Set Transformer(+OT) and DeepSets(+OT), the summary network, deﬁned as in Equations 1 and 7, is optimized with matrix B by both the OT loss and task-speciﬁc loss.
Revised: DeepSets, the summary network is deﬁned in Equation 1 and optimized by the task-speciﬁc loss; for Set Transformer(+OT) and DeepSets(+OT), the summary network, deﬁned as in Equations 1 and 7, is optimized with matrix B by both the OT loss and task-speciﬁc loss.

Original: We provide more details about toy datasets, network architectures, and training schemes in Appendix B.
Revised: We provide more details about toy datasets, network architectures, and training schemes in Appendix B.

Original: Amortized Clustering with Mixture of Gaussians (MoGs): We consider the task of maximum likelihood of MoGs with C components, denoted as P ( x ; θ )
Revised: Amortized Clustering with Mixture of Gaussians (MoGs): We consider the task of maximum likelihood of MoGs with C components, denoted as P ( x ; θ )

Original: = (cid:80) Cc =
Revised: = (cid:80) Cc =

Original: π c N (cid:0) x | µ c , diag (cid:0) σ 2 c (cid:1)(cid:1) .
Revised: π c N (cid:0) x | µ c , diag (cid:0) σ 2 c (cid:1)(cid:1) .

Original: Given the dataset X = { x 1: n } generated from the MoG, the goal is to train a neural network, which [Table] [Table] [Table]
Revised: Given the dataset X = { x 1: n } generated from the MoG, the goal is to train a neural network, which takes X as input set and outputs parameters θ .

Original: [Table] [Table] 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 of C Gaussians.
Revised: Each dataset contains n ∈ [100 , 500] points on a 2D plane, each of which is sampled from one of C Gaussians.

Original: Table 1 summarizes the test average likelihood of different models with varying C ∈ { 4 , 8 } [Table] where we set K = 50 for all C .
Revised: Table 1 reports the test average likelihood of different models with varying C ∈ { 4 , 8 } , where we set K = 50 for all C .

Original: We observe that Set Transformer outperforms DeepSets largely, validating the effectiveness of attention mechanisms in this task.
Revised: We observe that Set Transformer outperforms DeepSets largely, validating the effectiveness of attention mechanisms in this task.

Original: Besides, both Set Transformer(+OT) and DeepSets(+OT) improve their baselines, showing that the OT loss can encourage the summary networks to learn more efﬁcient summary statistics.
Revised: Besides, both Set Transformer(+OT) and DeepSets(+OT) improve their baselines, showing that the OT loss can encourage the summary networks to learn more efﬁcient summary statistics.

Original: Point Cloud Classiﬁcation: Here, we evaluate our method on the task of point cloud classiﬁcation using the ModelNet40 (Chang et al., 2015) dataset 1 , which consists of 3D objects from 40 different categories.
Revised: Point Cloud Classiﬁcation: Here, we evaluate our method on the task of point cloud classiﬁcation using the ModelNet40 (Chang et al., 2015) dataset 1 , which consists of 3D objects from 40 different categories.

Original: By treating each object as a point cloud, we represent it as a set of N vectors in R 3 (x; y; z-coordinates).
Revised: By treating each object as a point cloud, we represent it as a set of N vectors in R 3 (x; y; z-coordinates).

Original: Table 1 reports the classiﬁcation accuracy, where we perform the experiments with varying N ∈ { 64 , 1024 } , set K = 40 , and use the Euclidean distance to deﬁne the cost function.
Revised: Table 1 reports the classiﬁcation accuracy, where we perform the experiments with varying N ∈ { 64 , 1024 } , set K = 40 .

Original: Clearly, both of the DeepSets and Set Transformer can be improved by adding the OT loss.
Revised: Clearly, both of the DeepSets and Set Transformer can be improved by adding the OT loss.

Original: Notably, fewer points would lead to lower classiﬁcation results, where the OT loss plays a more important role.
Revised: Notably, fewer points would lead to lower performance, where the OT loss plays a more important role.

Revised: Taking this task as the example, we further study our model’s sensitivity to the hyper-parameter (cid:15) in Fig.

Revised: 4 of the Appendix B.5.

Original: [Figure]
Revised: [Figure]

Original: Sum of Digits:
Revised: Sum of Digits:

Original: Following Zaheer et al.
Revised: Following Zaheer et al.

Original: (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 } .
Revised: (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 } .

Original: By randomly sampling a subset of maximum M = 10 images from MNIST8m, we build N = 100 k “sets” of training, where we denote the sum of digits in that set as the set-label.
Revised: By randomly sampling a subset of maximum M = 10 images from MNIST8m, we build N = 100 k “sets” of training, where we denote the sum of digits in that set as the set-label.

Original: We construct 100 k sets of test MNIST digits, where we vary the M starting from 10 all the way up to 100 .
Revised: We construct 100 k sets of test MNIST digits, where we vary the M starting from 10 all the way up to 100 .

Original: The output of the summary network is a scalar, predicting the sum of M digits.
Revised: The output of the summary network is a scalar, predicting the sum of M digits.

Original: In this case, we adopt L1 as the task-speciﬁc loss and set K = 10 .
Revised: In this case, we adopt L1 as the task-speciﬁc loss and set K = 10 .

Original: We show the accuracy of digit summation for different algorithms in Fig.
Revised: We show the accuracy of digit summation for different algorithms in Fig.

Original: 2 and ﬁnd
Revised: 2 and ﬁnd

Original: that the OT loss can enhance the summary networks to achieve better generalization.
Revised: that the OT loss can enhance the summary networks to achieve better generalization.

Revised: In this task, we also explore the convergence rate and the learned transport plan matrix of Sinkhorn algorithm in Fig.

Revised: 5 of the Appendix B.6.

Original: 5. E XPERIMENTS ON F EW - SHOT C LASSIFICATION
Revised: 5. E XPERIMENTS ON F EW - SHOT C LASSIFICATION

Original: To explore whether our proposed method can improve the metric-based few-shot classiﬁcation, we consider two commonly-used algorithms as the baselines, including ProtoNet (Snell et al., 2017) and MatchNet (Vinyals et al., 2016).
Revised: To explore whether our proposed method can improve the metric-based few-shot classiﬁcation, we consider two commonly-used algorithms as the baselines, including ProtoNet (Snell et al., 2017) and MatchNet (Vinyals et al., 2016).

Original: Denoting the feature extractor in each algorithm as f φ 1 , we  consider several popular backbones, including ResNet10 and ResNet34 (He et al., 2016).
Revised: Denoting the feature extractor in each algorithm as f φ 1 , we [Table] consider several popular backbones, including ResNet10 and ResNet34 (He et al., 2016).

Original: 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 classiﬁcation errors.
Revised: 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 classiﬁcation errors.

Original: We perform the experiments on the CUB (Welinder et al., 2010) and miniImageNet (Ravi &
Revised: We perform the experiments on the CUB (Welinder et al., 2010) and miniImageNet (Ravi &

Original: Larochelle, 2016).
Revised: Larochelle, 2016).

Original: As a ﬁne-grained few-shot classiﬁcation 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 intobase classes, 50 validation classes, and 50 novel classes following Chen et al.
Revised: As a ﬁne-grained few-shot classiﬁcation 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 intobase classes, 50 validation classes, and 50 novel classes following Chen et al.

Original: 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.
Revised: 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.

Original: We follow the splits used in previous work (Ravi & Larochelle, 2016), which splits the dataset into 64 base classes, 16 validation classes [Table] and 20 novel classes.
Revised: 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.

Original: Table 2 reports the 5way5shot and 5way10shot classiﬁcation results of different methods on miniImageNet and CUB.
Revised: Table 2 reports the 5way5shot and 5way10shot classiﬁcation results of different methods on miniImageNet and CUB.

Original: We see that introducing the OT distance and summary network can consistently improve over baseline classiﬁers, and the performance gain gradually increases with the development of number of network layers.
Revised: We see that introducing the OT distance and summary network can consistently improve over baseline classiﬁers, and the performance gain gradually increases with the development of number of network layers.

Original: This suggests that our proposed plug-and-play framework can be ﬂexibly used to enhance the metric-based few-shot classiﬁcation, without the requirement of designing complicated models on purpose. 
Revised: This suggests that our proposed plug-and-play framework can be ﬂexibly used to enhance the metric-based few-shot classiﬁcation, without the requirement of designing complicated models on purpose. [Table]

Original: [Table]
Revised: [Table]

Original: 5. E XPERIMENTS ON F EW - SHOT G ENERATION
Revised: 5. E XPERIMENTS ON F EW - SHOT G ENERATION

Original: 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.
Revised: 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.

Original: For CGANbased models, we adopt summary network as the encoder to extract the feature vector h j from a given set D j , which is further fed into the generator and discriminator as the conditional information.
Revised: For CGANbased models, we adopt summary network as the encoder to extract the feature vector h j from a given set D j , which is further fed into the generator and discriminator as the conditional information.

Original: 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.
Revised: 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.

Original: 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 ( D 1 j , D 2 j ) from the same distribution or category; then we represent the real samples as the combination of D 1 j and D 2 j and the fake ones as the combination of D 1 j and ˆ D 1 j from the generator (conditioned on D 1 j ).
Revised: 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 ( D 1 j , D 2 j ) from the same distribution or category; then we represent the real samples as the combination of D 1 j and D 2 j and the fake ones as the combination of D 1 j and ˆ D 1 j from the generator (conditioned on D 1 j ).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We consider the DeepSets as the summary network for its simple architecture.
Revised: We consider the DeepSets as the summary network for its simple architecture.

Original: 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.
Revised: 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.

Original: To evaluate the quality of generated samples, we adopt commonly used metric Fr ´ echet 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.
Revised: To evaluate the quality of generated samples, we adopt commonly used metric Fr´echet 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.

Original: We provide several examples on toy datasets to show the efﬁciency of our proposed model in Appendix D.
Revised: We provide several examples on toy datasets to show the efﬁciency of our proposed model in Appendix D.

Original: [Table]
Revised: [Table]

Original: Rotated MNIST:
Revised: Rotated MNIST:

Original: Following Wu et al.
Revised: Following Wu et al.

Original: (2020), we artiﬁcially 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 .
Revised: (2020), we artiﬁcially 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 .

Original: We choose 9 interleaved distributions for training and the rest as unseen distributions for testing.
Revised: We choose 9 interleaved distributions for training and the rest as unseen distributions for testing.

Original: We consider the CGAN-based and DAGAN-based models, respectively.
Revised: We consider the CGAN-based and DAGAN-based models, respectively.

Original: 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.
Revised: 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.

Original: We summarize the test performance in Tablewith varying A .
Revised: We summarize the test performance in Tablewith varying A .

Original: We can ﬁnd 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.
Revised: We can ﬁnd 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.

Original: 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
Revised: 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

Original: Liu et al.
Revised: Liu et al.

Original: Flowers dataset contains 8189 images of 102 categories, which are divided intotraining 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.
Revised: Flowers dataset contains 8189 images of 102 categories, which are divided intotraining 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.

Original: We present the example images generated by DAGAN and DAGAN(+OT)in Fig. 3, and observe that our method achieves the lowest FID, where we  compute the FID scores with the similar way in Rotated MNIST experiment and present the network architecture in Appendix E.
Revised: 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.

Original: We ﬁnd that our proposed model has the ability to generate more realistic natural images compared with baselines.
Revised: We ﬁnd that our method achieves the lowest FID and has the ability to generate more realistic natural images compared with baselines.

Original: This indicates the summary network in our proposed framework can successfully capture the important summary statistics within the set, beneﬁcial for the few-shot image generation.[Figure]
Revised: This indicates the summary network in our proposed framework can successfully capture the important summary statistics within the set, beneﬁcial for the few-shot image generation.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: 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.
Revised: 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.

Original: Moreover, we use the summary network to encode input set as the prototype proportion (i.e., set representation) for global centers in corresponding set.
Revised: Moreover, we use the summary network to encode input set as the prototype proportion (i.e., set representation) for global centers in corresponding set.

Original: To learn the distribution over global prototypes and summary network, we minimize the OT distance between two distributions in terms of the deﬁned cost function.
Revised: To learn the distribution over global prototypes and summary network, we minimize the OT distance between two distributions in terms of the deﬁned cost function.

Original: Only additionally introducing the acceptable parameters, our proposed model provides a natural and unsupervised way to improve the summary network.
Revised: Only additionally introducing the acceptable parameters, our proposed model provides a natural and unsupervised way to improve the summary network.

Original: 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 classiﬁcation and implicit meta generative modeling.
Revised: 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 classiﬁcation and implicit meta generative modeling.

Original: 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.
Revised: 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.

Original: Due to the ﬂexibility and simplicity ofthe our proposed framework, there are still some exciting extensions of our work.
Revised: Due to the ﬂexibility and simplicity of our proposed framework, there are still some exciting extensions.

Original: For example, an interesting future work would be to apply our proposed framework into approximate Bayesian computation for posterior inference.
Revised: For example, an interesting future work would be to apply our method into approximate Bayesian computation for posterior inference.

Original: A A LGORITHMS AND ILLUSTRATION OF OUR PROPOSED MODELS
Revised: A A LGORITHMS AND ILLUSTRATION OF OUR PROPOSED MODEL

Original: The pseudo code for introducing the OT distance to improve the summary network under the unsupervised tasks is presented in Algorithm 1.

Original: Besides, the pseudo code for the implicit meta generative modeling is provided in Algorithm 2.[-Algorithm 
-]
Revised: The pseudo code for the implicit meta generative modeling is provided in Algorithm 2.[-Algorithm 
-]

Original: The workﬂow of our proposed implicit meta generative framework.

Original: Require: Datasets D 1: J , initial discriminator parameters w , initial generator parameters θ , initial summary network parameters φ , initial matrix B , the cost function C , the number of critic iterations per generator iteration η critic , the batch size m , learning rate α and the hyper-parameter (cid:15) .

Original: while w , θ, B , φ has not converged do

Original: Randomly choose j from 1 , 2 , .., J for t = 1 , · · · , η critic do

Original: Sample real data set D j = { x ji } mi =1 from j -th empirical distribution P j ;

Original: Embed observed batch data into the statistics h j = Softmax( S φ ( D j )) ;

Original: Sample prior samples end for ;

Original: Represent the Q j with global prototype matrix B and set representation/statistics h j

Original: Compute loss

Original: (cid:15) ( P j , Q j ) between P j and Q j with Sinkhorn algorithm in Equation details;

Original: ; end while

Original: B E XPERIMENTAL SETTINGS ABOUT INTRODUCING OT DISTANCE INTO THE S UMMARY N ETWORKS B. D ETAILS FOR AMORTIZED CLUSTERING WITH MIXTURES OF G AUSSIANS
Revised: B E XPERIMENTAL SETTINGS ABOUT INTRODUCING OT DISTANCE INTO THE S UMMARY N ETWORKS B. D ETAILS FOR AMORTIZED CLUSTERING WITH MIXTURES OF G AUSSIANS

Original: We generate the 2D toy datasets following the Lee et al.
Revised: We generate the 2D toy datasets following the Lee et al.

Original: (2019), where we additionally vary the C (the number of components) from 4 to 8 .
Revised: (2019), where we additionally vary the C (the number of components) from 4 to 8 .

Original: Below, we present the detailed generation process about the toy datasets:
Revised: Below, we present the detailed generation process about the toy datasets:

Original: Specify the number of components C for 2D toy dataset. 2.
Revised: Specify the number of components C for 2D toy dataset. 2.

Original: Generate the number of data points, n ∼ Uniform(100 , 500) .
Revised: Generate the number of data points, n ∼ Uniform(100 , 500) .

Original: Sample the mean vector for C components.
Revised: Sample the mean vector for C components.

Original: µ c [Equation] d ∼ Uniform( − 4 , 4) , 
Revised: µ c,d ∼ Uniform( − 4 , 4) , c

Original: [Equation]
Revised: [Equation]

Revised: , C, d = 1 ,4.

Original: Sample the cluster labels.
Revised: Sample the cluster labels.

Original: [Equation]
Revised: [Equation]

Revised: (cid:62) (cid:17) , z i ∼ Categorical( π ) , i = 1 , . . .

Revised: , n, z i = 1 , . . .

Revised: C 5.

Original: Generate data from spherical Gaussian.
Revised: Generate data from spherical Gaussian.

Original: [Equation]
Revised: [Equation]

Original: B. D ETAILS ABOUT S ET -T RANSFORMER - BASED AND D EEP S ETS - BASED ARCHITECTURES USED IN M O G S EXPERIMENTS
Revised: B. D ETAILS ABOUT S ET -T RANSFORMER - BASED AND D EEP S ETS - BASED ARCHITECTURES USED IN M O G S EXPERIMENTS

Original: 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.
Revised: 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.

Original: Then the resulting vector representation z j 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.
Revised: Then the resulting vector representation z j 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.

Original: We use ELU activation at all layers.
Revised: We use ELU activation at all layers.

Original: To introduce the OT loss into the DeepSets, we further feed the z j into a fully connected layer with 512 units followed by a 50-way softmax unit and also introduce the global matrix
Revised: To introduce the OT loss into the DeepSets, we further feed the z j into a fully connected layer with 512 units followed by a 50-way softmax unit and also introduce the global matrix

Original: B ∈ R 2 × 50 .
Revised: B ∈ R 2 × 50 .

Original: [Table] [Table] Set Transformer To perform the MoGs experiments, we adopt the same architecture for Set Transformer following Lee et al.
Revised: Set Transformer To perform the MoGs experiments, we adopt the same architecture for Set Transformer following Lee et al.

Original: (2019), whose parameters are reported in Table 4.
Revised: (2019), whose parameters are reported in Table 4.

Original: To introduce the OT loss, we also add the two fully connected layers with 256 units on the resulting vector z j followed by a 50-way softmax unit [Table] and a global prototype matrix B ∈ R 2 × 50 .
Revised: To introduce the OT loss, we also add the two fully connected layers with 256 units on the resulting vector z j followed by a 50-way softmax unit, and a global prototype matrix B ∈ R 2 × 50 .

Original: B. D ETAILS ABOUT S ET -T RANSFORMER - BASED AND D EEP S ETS - BASED ARCHITECTURES USED IN SUM OF DIGITSDeepSets
Revised: B. D ETAILS ABOUT S ET -T RANSFORMER - BASED AND D EEP S ETS - BASED ARCHITECTURES USED IN SUM OF DIGITS

Original:  Following the ofﬁcial code in Zaheer et al.
Revised: DeepSets Following the ofﬁcial code in Zaheer et al.

Original: (2017), we adopt the default architecture to implement the DeepSets, where we ﬁrst project the image into a 128 -dimensional vector with three convolutional layers and apply summary network on the 128 -dimensional vectors.
Revised: (2017), we adopt the default architecture to implement the DeepSets, where we ﬁrst project the image into a 128 -dimensional vector with three convolutional layers and apply summary network on the 128 -dimensional vectors.

Original: To build
Revised: To build

Original: DeepSets(+OT), we take the set representation
Revised: DeepSets(+OT), we take the set representation

Original: z
Revised: z

Original: j 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 ∈ R 128 × 10 .
Revised: j 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 ∈ R 128 × 10 .

Original: 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.
Revised: 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.

Original: To build Set
Revised: To build Set

Original: Transformer(+OT), we take the representation z
Revised: Transformer(+OT), we take the representation z

Original: j 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 ∈ R 128 × 10 .
Revised: j 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 ∈ R 128 × 10 .

Original: B. D ETAILS ABOUT S ET -T RANSFORMER - BASED AND D EEP S ETS - BASED ARCHITECTURES USED IN POINT CLOUD CLASSIFICATION
Revised: B. D ETAILS ABOUT S ET -T RANSFORMER - BASED AND D EEP S ETS - BASED ARCHITECTURES USED IN POINT CLOUD CLASSIFICATION

Original: DeepSets For original DeepSets, we adopt the same architecture with Zaheer et al.
Revised: DeepSets For original DeepSets, we adopt the same architecture with Zaheer et al.

Original: In a speciﬁc, the f φ 1 in summary network contains 3 permutation-equivariant layers with 256 channels followed by max-pooling over the set structure.
Revised: In a speciﬁc, the f φ 1 in summary network contains 3 permutation-equivariant layers with 256 channels followed by max-pooling over the set structure.

Original: Then the resulting vector representation z j of the set is then fed to a fully connected layer with 256 units followed by a 40-way softmax unit. 
Revised: Then the resulting vector representation z j of the set is then fed to a fully connected layer with 256 units followed by a 40-way softmax unit. We [Table]

Original: 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.
Revised: [Table] use Tanh activation at all layers and dropout on the layers after set-max-pooling (i [Table] e. two dropout operations) with 50% dropout rate.

Original: To introduce the OT loss into the DeepSets, we further feed the z j into a fully connected layer with 256 units followed by a 40-way softmax unit, with 70% dropout rate.
Revised: To introduce the OT loss into the DeepSets, we further feed the z j into a fully connected layer with 256 units followed by a 40-way softmax unit, with 70% dropout rate.

Original: Besides, we additionally introduce the center matrix B ∈ R 3 × 40 .
Revised: Besides, we additionally introduce the center matrix B ∈ R 3 × 40 .

Original: Set Transformer
Revised: Set Transformer

Original: We also adopt the same architecture to implement the Set Transformer, where we summarize the parameters in Table 5, following Lee et al.
Revised: We also adopt the same architecture to implement the Set Transformer, where we summarize the parameters in Table 5, following Lee et al.

Original: 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 ∈ R 3 × 40 . 
Revised: 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 ∈ R 3 × 40 . B. P ARAMETER SENSITIVITY

Revised: In the previous experiments, we ﬁx the value of (cid:15) as 0 . 1 , controlling the weight of the entropic regularisation in the Sinkhorn algorithm.

Original:  C E XPERIMENTAL SETTINGS ABOUT FEW - SHOT CLASSIFICATION
Revised: Notably, unless speciﬁed otherwise, we specify the construction of C as C ik = 1 − cos ( x ji , β k ) . Therefore, the cost function provides an upper-bounded positive similarity metric, making the (cid:15) has the corresponding reasonable range as a prior knowledge. Here, we study our DeepSets(+OT)’s sensitivity to (cid:15) . We consider the point cloud classiﬁcation task and each object is represented as a set of N = 20 vectors.

Revised: As shown in Fig.

Revised: 4, we report the performance of DeepSets(+OT) on point cloud classiﬁcation task with varying (cid:15) , where DeepSets serves as the baseline.

Revised: It can be seen that our model is robust to the (cid:15) .

Revised: Besides, all the results of

Revised: DeeepSets(+OT) with different (cid:15) are superior than that of DeeepSets, indicating the effectiveness of our method.

Revised: By ﬁne-tuning (cid:15) for each dataset in each task, we might obtain better results than those reported in our experiments.

Revised: However, we aim to validate our method instead of exhaustively tuning this hyper-parameter and thus we set (cid:15) = 0 .

Revised: 1 , which can achieve the acceptable result.

Revised: B. C ONVERGENCE RATE OF S INKHORN ALGORITHM

Revised: In this paper, we set the maximum iteration number as Itermax = 200 in Sinkhorn algorithm for all experiments.

Revised: As shown in Fig.

Revised: 5, we visualize the convergence rate of Sinkhorn algorithm, where we consider the task about “sum of digits” (DeepSets+OT).

Revised: The upper ﬁgure shows the convergence rate of Sinkhorn algorithm.

Revised: The bottom ﬁgure visualizes the transport plan matrix with varying iterations.

Revised: We ﬁnd that the 200 iterations are typically enough for Sinkhorn algorithm

Revised: and we can learn a sparse transport plan matrix T when the algorithm converge.

Revised: Notably, the transport plan matrix needs to satisfy two marginal constraints, deﬁned by the probability measures of two distributions, respectively.

Revised: Recall that the empirical distribution has an unchanged uniform probability measure, so the learned transport plan matrix is dense for the N j observed samples.

Revised: In terms of another dis- tribution, its probability measure is the set-speciﬁc representation, weighting the importance of K shared centers for corresponding set.

Revised: Therefore, it is reasonable that transport plan matrix is sparse for K centers.

Revised: C E XPERIMENTAL SETTINGS ABOUT FEW - SHOT CLASSIFICATION

Original: Denote the prototype for set j (computed by f φ 1 ) in few-shot classiﬁcation as c j .
Revised: Denote the prototype for set j (computed by f φ 1 ) in few-shot classiﬁcation as c j .

Original: We consider two backbones for f φ 1 , including ResNet10 and ResNet34, which produce the 512 -dimensional c j .
Revised: We consider two backbones for f φ 1 , including ResNet10 and ResNet34, which produce the 512 -dimensional c j .

Original: To improve the metric-based few-shot classiﬁcation with our framework, taking the c j as input, we further construct the g φ 2 .
Revised: To improve the metric-based few-shot classiﬁcation with our framework, taking the c j as input, we further construct the g φ 2 .

Original: Speciﬁcally, we introduce a fully connected network with architecture as 512 → 256 units with ReLU function followed by a X -way
Revised: Speciﬁcally, we introduce a fully connected network with architecture as 512 → 256 units with ReLU function followed by a X -way

Original: ( X =64 for CUB, and X =128 for miniImageNet) Softmax function and a center matrix B ∈ R 512 × X .
Revised: ( X =64 for CUB, and X =128 for miniImageNet) Softmax function and a center matrix B ∈ R 512 × X .

Original: We conduct 10000 tasks of the training set D tr to train the model while 1000 tasks of the test set D te to evaluate the learned model.
Revised: We conduct 10000 tasks of the training set D tr to train the model while 1000 tasks of the test set D te to evaluate the learned model.

Original: And D tr ∩ D te = ∅ .
Revised: And D tr ∩ D te = ∅ .

Original: We run 60 epochs to train the model on CUB and miniImageNet.
Revised: We run 60 epochs to train the model on CUB and miniImageNet.

Original: The model is trained using Adam optimizerwith default settings (learning rate 1 e − 3 , β = (0 . 9 , 0 . 999) , and (cid:15) = 1 e − 8 ) on one Nvidia Geforce RTX3090 GPU.
Revised: The model is trained using Adam optimizer

Original: D A DDITIONAL EXPERIMENTAL RESULTS ON FEW - SHOT GENERATION ABOUT TOY DATASETS
Revised: D A DDITIONAL EXPERIMENTAL RESULTS ON FEW - SHOT GENERATION ABOUT TOY DATASETS

Original: We test our algorithm through a series of synthetic data sets and realistic data sets.
Revised: We test our algorithm through a series of synthetic data sets and realistic data sets.

Original: For synthetic datasets, we set T θ , f w and S φ as fully connected neural networks, where T θ , f w have 4 hidden layers and f φ 1 and g φ 2 (we adopt DeepSets) have 3 hidden layers.
Revised: For synthetic datasets, we set T θ , f w and S φ as fully connected neural networks, where T θ , f w have 4 hidden layers and f φ 1 and g φ 2 (we adopt DeepSets) have 3 hidden layers.

Original: Each layer has 200 nodes, and the activation function is chosen as RELU, where we adopt the softmax in the ﬁnal layer.
Revised: Each layer has 200 nodes, and the activation function is chosen as RELU, where we adopt the softmax in the ﬁnal layer.

Original: Normal distribution on 2D toy data: We ﬁrst consider the 2D normal case, where the training data contains 10 K sets and each set contains 100 data points from N ( µ , Σ ) .
Revised: Normal distribution on 2D toy data: We ﬁrst consider the 2D normal case, where the training data contains 10 K sets and each set contains 100 data points from N ( µ , Σ ) .

Original: We sample the mean, variance, and covariance from U [ − 5 , 5] , U [1 , 2] , and U [ − 0 .
Revised: We sample the mean, variance, and covariance from U [ − 5 , 5] , U [1 , 2] , and U [ − 0 .

Original: 5] , respectively.
Revised: 5] , respectively.

Original: Fig.
Revised: Fig.

Original: 5 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.
Revised: 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.

Original: We ﬁnd that our model (third column) can improve the resistance to mode collapse compared with CGAN+MSE (second column) and better ﬁt the unseen test distributions than CGAN (ﬁrst column).
Revised: We ﬁnd that our model (third column) can improve the resistance to mode collapse compared with CGAN+MSE (second column) and better ﬁt the unseen test distributions than CGAN (ﬁrst column).

Original: This result indicates the OT loss can spur the summary network to capture more desired statistics for unseen distributions.
Revised: This result indicates the OT loss can spur the summary network to capture more desired statistics for unseen distributions.

Original: 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.
Revised: 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.

Original: The training data contains 10 K sets each containingsamples.
Revised: The training data contains 10 K sets each containingsamples.

Original: We also visualize the pdfs of gaussian distributions with randomly means and variance in Fig 6, which are used to sample test data, and show the 500 data points generated by the pushforward.
Revised: 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.

Original: For this experiment, we set the dimension of z and summary vector s as 2 .
Revised: For this experiment, we set the dimension of z and summary vector s as 2 .

Original: 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.
Revised: 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.

Original: For Gaussian and Laplacian distributions, means and variances are sampled from
Revised: For Gaussian and Laplacian distributions, means and variances are sampled from

Original: U [ − 1 , 1] and U [0 . 5 , 2] respectively; for Exponential distributions, rates are sampled from U [0 . 5 , 2] .
Revised: U [ − 1 , 1] and U [0 . 5 , 2] respectively; for Exponential distributions, rates are sampled from U [0 . 5 , 2] .

Original: Fig.
Revised: Fig.

Original: 7 visualizes the pdfs of six one-dimensional test distributions with different means and variances and the generated data points.
Revised: 8 visualizes the pdfs of six one-dimensional test distributions with different means and variances and the generated data points.

Original: It is interesting to observe that the generated data points can ﬁt the corresponding pdf well, indicating our model can generalize to different distributions with varying parameters.
Revised: It is interesting to observe that the generated data points can ﬁt the corresponding pdf well, indicating our model can generalize to different distributions with varying parameters.

Original: Besides, our model performs slightly worse on the Exponential distributions, perhaps attributing to the fact that it is the only non-symmetric distribution.
Revised: Besides, our model performs slightly worse on the Exponential distributions, perhaps attributing to the fact that it is the only non-symmetric distribution.

Original: E D ETAILS ABOUT NATURAL IMAGE GENERATIONS 
Revised: E D ETAILS ABOUT NATURAL IMAGE GENERATIONS AND THE RESULTS

Original: We use denseNet proposed by Huang et al.
Revised: We use denseNet proposed by Huang et al.

Original: as the backbone of summary network, then a pooling operation is conducted as Zaheer et al.
Revised: as the backbone of summary network, then a pooling operation is conducted as Zaheer et al.

Original: (2017) does.
Revised: (2017) does.

Original: And a 2-layer fully connected network with ReLU activation function is ﬁnally employed to embed the 4096-D visual features into the corresponding 512-D set representations h j .
Revised: And a 2-layer fully connected network with ReLU activation function is ﬁnally employed to embed the 4096-D visual features into the corresponding 512-D set representations h j .

Original: 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 .
Revised: 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 .

Original: Besides, we use a 5-layer deconvolution network [ConvTranspose2d( X + 100, 512, 4, 1, 0) → ConvTranspose2d(512, 256, 4, 2, 1)
Revised: Besides, we use a 5-layer deconvolution network [ConvTranspose2d( X + 100, 512, 4, 1, 0) → ConvTranspose2d(512, 256, 4, 2, 1)

Original: → ConvTranspose2d(512, 128, 4, 2, 1) →
Revised: → ConvTranspose2d(512, 128, 4, 2, 1) →

Original: 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 , h j ]) as fake output images x fake : R N × 3 × 64 × 64 .
Revised: ConvTranspose2d(128, 64, 4, → 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 , h j ]) as fake output images x fake : R N × 3 × 64 × 64 .

Original: Finally, a 5-layer discriminator network [Conv2d(2*3, 64, 4, 2, 1) → Conv2d(64, 128, 4, 2, 1))
Revised: Finally, a 5-layer discriminator network [Conv2d(2*3, 64, 4, 2, 1) → Conv2d(64, 128, 4, 2, 1))

Original: → 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 ﬁrst fourth deconvolutional layers and the last layer without BatchNorm while with sigmoid activation funtion to distinguish the true or fake generated images. 
Revised: → 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 ﬁrst 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. [Figure] F T HE DIFFERENCE BETWEEN OUR MODEL AND BARYCENTER PROBLEM

Revised: In this section, we clarify the difference between Wasserstein barycenter and our method.

Revised: Specifically, for j -th distribution, we denoted P j as its empirical distribution consisting of N j samples, expressed as P j = (cid:80) N j i =1 1 N j δ x ji , x ji ∈ R d .

Revised: Notably, a j = [ 1 N j ] ∈ Σ N j represents the probability measure for distribution P j .

Revised: For another thing, we can represent P j with another to-be-learned distribution Q j , deﬁned as

Revised: Q j = (cid:80) K k =1 h jk δ β k , β k ∈ R d .

Revised: Here b j = [ h jk ] ∈ Σ K is the probability measure for distribution

Revised: Q j , which can be computed using summary network S φ ( x i, 1: N j ) and serves as the representation for set j .

Revised: And β k is the k -th prototype in the same space of the observed data points, which is the k -th column of B ∈ R K × d , a learnable global prototype matrix.

Revised: To optimize the β 1: K and the summary network S φ for computing h j , we minimize the average OT loss (between Q j and P j ) for all training sets.

Revised: We rewrite the Equation (6) here for convenience:

Revised: [Equation]

Revised: Usually, this equation can also be represented:

Revised: [Equation]

Revised: In terms of Wasserstein barycenter, we adopt the same notations for consistency.

Revised: Following (Altschuler & Boix-Adser`a, 2021), given J empirical distributions P 1: J and their respective probability measures [ a 1 , . . .

Revised: , a J ] supported on R d and a vector λ ∈ Σ J , their corresponding Wasserstein barycenter can be viewed as another distribution Q , i.e., Q = (cid:80) Kk

Revised: =1 m k δ β k , β k ∈ R d , where β k is the k -th column of B , m =

Revised: [ m k ] ∈ Σ K is the probability measure for distribution Q .

Revised: Then we can learn the barycenter (i.e., B and m ) by minimizing

Revised: [Equation]

Revised: where above W ( · , · ) denotes the squared 2-Wasserstein distance.

Revised: By comparing the equation 10 and equation 9, we can ﬁnd that our model learns a Q j to approximate P j for each distribution j but “barycenter” problem learns a shared Q as the barycenter for all P 1: J .

Revised: Therefore, after minimizing the average loss for all training sets with Equation (6), we can use the probability measure b j = [ h jk ] (i.e., h j ) to represent the empirical distribution P j .

Revised: Especially, we can directly map the test set to its representation by using the summary network.

Revised: However, since the probability measure m in “barycenter” is shared by all distributions, it can not represent a speciﬁc set.

Revised: Therefore, to achieve the set representation, it might need to ﬁrst 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.

Revised: 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.

Revised: We will clarify the difference between our model and “barycenter” problem in the appendix of the revised version if necessary.