Original: A T RAINABLE O PTIMAL T RANSPORT E MBEDDING FOR F EATURE A GGREGATION AND ITS R ELATIONSHIP TO A TTENTION
Revised: A T RAINABLE O PTIMAL T RANSPORT E MBEDDING FOR F EATURE A GGREGATION AND ITS R ELATIONSHIP TO A TTENTION

Original: A BSTRACT
Revised: A BSTRACT

Original: We address the problem of learning on sets of features, motivated by the need of performing pooling operations in long biological sequences of varying sizes, with long-range dependencies, and possibly few labeled data.
Revised: We address the problem of learning on sets of features, motivated by the need of performing pooling operations in long biological sequences of varying sizes, with long-range dependencies, and possibly few labeled data.

Original: To address this challenging task, we introduce a parametrized representation of ﬁxed size, which embeds and then aggregates elements from a given input set according to the optimal transport plan between the set and a trainable reference.
Revised: To address this challenging task, we introduce a parametrized representation of ﬁxed size, which embeds and then aggregates elements from a given input set according to the optimal transport plan between the set and a trainable reference.

Original: Our approach scales to large datasets and allows end-to-end training of the reference, while also providing a simple unsupervised learning mechanism with small computational cost.
Revised: Our approach scales to large datasets and allows end-to-end training of the reference, while also providing a simple unsupervised learning mechanism with small computational cost.

Original: Our aggregation technique admits two useful interpretations: it may be seen as a mechanism related to attention layers in neural networks, or it may be seen as a scalable surrogate of a classical optimal transport-based kernel.
Revised: Our aggregation technique admits two useful interpretations: it may be seen as a mechanism related to attention layers in neural networks, or it may be seen as a scalable surrogate of a classical optimal transport-based kernel.

Original: We experimentally demonstrate the effectiveness of our approach on biological sequences, achieving state-of-the-art results for protein fold recognition and detection of chromatin proﬁles tasks, and, as a proof of concept, we show promising results for processing natural language sequences.
Revised: We experimentally demonstrate the effectiveness of our approach on biological sequences, achieving state-of-the-art results for protein fold recognition and detection of chromatin proﬁles tasks, and, as a proof of concept, we show promising results for processing natural language sequences.

Original: We provide an open-source implementation of our embedding that can be used alone or as a module in larger learning models at https://github.com/claying/OTK .
Revised: We provide an open-source implementation of our embedding that can be used alone or as a module in larger learning models at https://github.com/claying/OTK .

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Many scientiﬁc ﬁelds such as bioinformatics or natural language processing (NLP) require processing sets of features with positional information (biological sequences, or sentences represented by a set of local features).
Revised: Many scientiﬁc ﬁelds such as bioinformatics or natural language processing (NLP) require processing sets of features with positional information (biological sequences, or sentences represented by a set of local features).

Original: These objects are delicate to manipulate due to varying lengths and potentially long-range dependencies between their elements.
Revised: These objects are delicate to manipulate due to varying lengths and potentially long-range dependencies between their elements.

Original: For many tasks, the difﬁculty is even greater since the sets can be arbitrarily large, or only provided with few labels, or both.
Revised: For many tasks, the difﬁculty is even greater since the sets can be arbitrarily large, or only provided with few labels, or both.

Original: Deep learning architectures speciﬁcally designed for sets have recently been proposed (Lee et al., 2019; Skianis et al., 2020).
Revised: Deep learning architectures speciﬁcally designed for sets have recently been proposed (Lee et al., 2019; Skianis et al., 2020).

Original: Our experiments show that these architectures perform well for NLP tasks, but achieve mixed performance for long biological sequences of varying size with few labeled data.
Revised: Our experiments show that these architectures perform well for NLP tasks, but achieve mixed performance for long biological sequences of varying size with few labeled data.

Original: Some of these models use attention (Bahdanau et al., 2015), a classical mechanism for aggregating features.
Revised: Some of these models use attention (Bahdanau et al., 2015), a classical mechanism for aggregating features.

Original: Its typical implementation is the transformer (Vaswani et al., 2017), which has shown to achieve state-of-the-art results for many sequence modeling tasks, e.g , in NLP (Devlin et al., 2019) or in bioinformatics (Rives et al., 2019), when trained with self supervision on large-scale data.
Revised: Its typical implementation is the transformer (Vaswani et al., 2017), which has shown to achieve state-of-the-art results for many sequence modeling tasks, e.g , in NLP (Devlin et al., 2019) or in bioinformatics (Rives et al., 2019), when trained with self supervision on large-scale data.

Original: Beyond sequence modeling, we are interested in this paper in ﬁnding a good representation for sets of features of potentially diverse sizes, with or without positional information, when the amount of training data may be scarce.
Revised: Beyond sequence modeling, we are interested in this paper in ﬁnding a good representation for sets of features of potentially diverse sizes, with or without positional information, when the amount of training data may be scarce.

Original: To this end, we introduce a trainable embedding, which can operate directly on the feature set or be combined with existing deep approaches.
Revised: To this end, we introduce a trainable embedding, which can operate directly on the feature set or be combined with existing deep approaches.

Original: More precisely, our embedding marries ideas from optimal transport (OT) theory (Peyr´e & Cuturi, 2019) and kernel methods (Sch¨olkopf & Smola, 2001).
Revised: More precisely, our embedding marries ideas from optimal transport (OT) theory (Peyr´e & Cuturi, 2019) and kernel methods (Sch¨olkopf & Smola, 2001).

Original: We call this embedding OTKE (Optimal
Revised: We call this embedding OTKE (Optimal

Original: Transport Kernel Embedding).
Revised: Transport Kernel Embedding).

Original: Concretely, we embed feature vectors of a given set to a reproducing kernel Hilbert space (RKHS) and then perform a weighted pooling operation, with weights given by the transport plan between the set and a trainable reference.
Revised: Concretely, we embed feature vectors of a given set to a reproducing kernel Hilbert space (RKHS) and then perform a weighted pooling operation, with weights given by the transport plan between the set and a trainable reference.

Original: To gain scalability, we then obtain a ﬁnite-dimensional embedding by using kernel approximation techniques (Williams & Seeger, 2001).
Revised: To gain scalability, we then obtain a ﬁnite-dimensional embedding by using kernel approximation techniques (Williams & Seeger, 2001).

Original: The motivation for using kernels is to provide a non-linear transformation of the input features before pooling, whereas optimal transport allows to align the features on a trainable reference with fast algorithms (Cuturi, 2013).
Revised: The motivation for using kernels is to provide a non-linear transformation of the input features before pooling, whereas optimal transport allows to align the features on a trainable reference with fast algorithms (Cuturi, 2013).

Original: Such combination provides us with a theoretically grounded, ﬁxed-size embedding that can be learned either without any label, or with supervision.
Revised: Such combination provides us with a theoretically grounded, ﬁxed-size embedding that can be learned either without any label, or with supervision.

Original: Our embedding can indeed become adaptive to the problem at hand, by optimizing the reference with respect to a given task.
Revised: Our embedding can indeed become adaptive to the problem at hand, by optimizing the reference with respect to a given task.

Original: It can operate on large sets with varying size, model long-range dependencies when positional information is present, and scales gracefully to large datasets.
Revised: It can operate on large sets with varying size, model long-range dependencies when positional information is present, and scales gracefully to large datasets.

Original: We demonstrate its effectiveness on biological sequence classiﬁcation tasks, including protein fold recognition and detection of chromatin proﬁles where we achieve state-of-the-art results.
Revised: We demonstrate its effectiveness on biological sequence classiﬁcation tasks, including protein fold recognition and detection of chromatin proﬁles where we achieve state-of-the-art results.

Original: We also show promising results in natural language processing tasks, where our method outperforms strong baselines.
Revised: We also show promising results in natural language processing tasks, where our method outperforms strong baselines.

Original: Contributions.
Revised: Contributions.

Original: In summary, our contribution is three-fold.
Revised: In summary, our contribution is three-fold.

Original: We propose a new method to embed sets of features of varying sizes to ﬁxed size representations that are well adapted to downstream machine learning tasks, and whose parameters can be learned in either unsupervised or supervised fashion.
Revised: We propose a new method to embed sets of features of varying sizes to ﬁxed size representations that are well adapted to downstream machine learning tasks, and whose parameters can be learned in either unsupervised or supervised fashion.

Original: We demonstrate the scalability and effectiveness of our approach on biological and natural language sequences.
Revised: We demonstrate the scalability and effectiveness of our approach on biological and natural language sequences.

Original: We provide an open-source implementation of our embedding that can be used alone or as a module in larger learning models.
Revised: We provide an open-source implementation of our embedding that can be used alone or as a module in larger learning models.

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

Original: Kernel methods for sets and OT-based kernels.
Revised: Kernel methods for sets and OT-based kernels.

Original: The kernel associated with our embedding belongs to the family of match kernels (Lyu, 2004; Tolias et al., 2013), which compare all pairs of features between two sets via a similarity function.
Revised: The kernel associated with our embedding belongs to the family of match kernels (Lyu, 2004; Tolias et al., 2013), which compare all pairs of features between two sets via a similarity function.

Original: Another line of research builds kernels by matching features through the Wasserstein distance.
Revised: Another line of research builds kernels by matching features through the Wasserstein distance.

Original: A few of them are shown to be positive deﬁnite (Gardner et al., 2018) and/or fast to compute (Rabin et al., 2011; Kolouri et al., 2016).
Revised: A few of them are shown to be positive deﬁnite (Gardner et al., 2018) and/or fast to compute (Rabin et al., 2011; Kolouri et al., 2016).

Original: Except for few hyperparameters, these kernels yet cannot be trained end-to-end, as opposed to our embedding that relies on a trainable reference.
Revised: Except for few hyperparameters, these kernels yet cannot be trained end-to-end, as opposed to our embedding that relies on a trainable reference.

Original: Efﬁcient and trainable kernel embeddings for biological sequences have also been proposed by Chen et al.
Revised: Efﬁcient and trainable kernel embeddings for biological sequences have also been proposed by Chen et al.

Original: (2019a;b).
Revised: (2019a;b).

Original: Our work can be seen as an extension of these earlier approaches by using optimal transport rather than mean pooling for aggregating local features, which performs signiﬁcantly better for long sequences in practice.
Revised: Our work can be seen as an extension of these earlier approaches by using optimal transport rather than mean pooling for aggregating local features, which performs signiﬁcantly better for long sequences in practice.

Original: Deep learning for sets.
Revised: Deep learning for sets.

Original: Deep Sets (Zaheer et al., 2017) feed each element of an input set into a feed-forward neural network.
Revised: Deep Sets (Zaheer et al., 2017) feed each element of an input set into a feed-forward neural network.

Original: The outputs are aggregated following a simple pooling operation before further processing.
Revised: The outputs are aggregated following a simple pooling operation before further processing.

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

Original: (2019) propose a Transformer inspired encoder-decoder architecture for sets which also uses latent variables.
Revised: (2019) propose a Transformer inspired encoder-decoder architecture for sets which also uses latent variables.

Original: Skianis et al.
Revised: Skianis et al.

Original: (2020) compute some comparison costs between an input set and reference sets.
Revised: (2020) compute some comparison costs between an input set and reference sets.

Original: These costs are then used as features in a subsequent neural network.
Revised: These costs are then used as features in a subsequent neural network.

Original: The reference sets are learned end-to-end.
Revised: The reference sets are learned end-to-end.

Original: Unlike our approach, such models do not allow unsupervised learning.
Revised: Unlike our approach, such models do not allow unsupervised learning.

Original: We will use the last two approaches as baselines in our experiments.
Revised: We will use the last two approaches as baselines in our experiments.

Original: Interpretations of attention.
Revised: Interpretations of attention.

Original: Using the transport plan as an ad-hoc attention score was proposed by Chen et al. (2019c) in the context of network embedding to align data modalities.
Revised: Using the transport plan as an ad-hoc attention score was proposed by Chen et al. (2019c) in the context of network embedding to align data modalities.

Original: Our paper goes beyond and uses the transport plan as a principle for pooling a set in a model, with trainable parameters.
Revised: Our paper goes beyond and uses the transport plan as a principle for pooling a set in a model, with trainable parameters.

Original: Tsai et al.
Revised: Tsai et al.

Original: (2019) provide a view of Transformer’s attention via kernel methods, yet in a very different fashion where attention is cast as kernel smoothing and not as a kernel embedding.
Revised: (2019) provide a view of Transformer’s attention via kernel methods, yet in a very different fashion where attention is cast as kernel smoothing and not as a kernel embedding.

Original:  P ROPOSED E MBEDDING 3. P RELIMINARIES
Revised:  P ROPOSED E MBEDDING 3. P RELIMINARIES

Original: We handle sets of features in R d and consider sets x living in
Revised: We handle sets of features in R d and consider sets x living in

Original: X = (cid:8) x | x = { x 1 , . . .
Revised: X = (cid:8) x | x = { x 1 , . . .

Original: , x n } such that x 1 , . .
Revised: , x n } such that x 1 , . .

Original: , x n ∈ R d for some n ≥ 1 (cid:9) .
Revised: , x n ∈ R d for some n ≥ 1 (cid:9) .

Original: Elements of X are typically vector representations of local data structures, such as k -mers for sequences, patches for natural images, or words for sentences.
Revised: Elements of X are typically vector representations of local data structures, such as k -mers for sequences, patches for natural images, or words for sentences.

Original: The size of x denoted by n may vary, which is not an issue since the methods we introduce may take a sequence of any size as input, while providing a ﬁxed-size embedding.
Revised: The size of x denoted by n may vary, which is not an issue since the methods we introduce may take a sequence of any size as input, while providing a ﬁxed-size embedding.

Original: We now revisit important results on optimal transport and kernel methods, which will be useful to describe our embedding and its computation algorithms.
Revised: We now revisit important results on optimal transport and kernel methods, which will be useful to describe our embedding and its computation algorithms.

Original: Optimal transport.
Revised: Optimal transport.

Original: Our pooling mechanism will be based on the transport plan between x and x (cid:48) seen as weighted point clouds or discrete measures, which is a by-product of the optimal transport problem (Villani, 2008; Peyr´e & Cuturi, 2019).
Revised: Our pooling mechanism will be based on the transport plan between x and x (cid:48) seen as weighted point clouds or discrete measures, which is a by-product of the optimal transport problem (Villani, 2008; Peyr´e & Cuturi, 2019).

Original: OT has indeed been widely used in alignment problems (Grave et al., 2019).
Revised: OT has indeed been widely used in alignment problems (Grave et al., 2019).

Original: Throughout the paper, we will refer to the Kantorovich relaxation of OT with entropic regularization, detailed for example in (Peyr´e & Cuturi, 2019).
Revised: Throughout the paper, we will refer to the Kantorovich relaxation of OT with entropic regularization, detailed for example in (Peyr´e & Cuturi, 2019).

Original: Let a in ∆ n (probability simplex) and b in ∆ n (cid:48) be the weights of the discrete measures (cid:80) i a i δ x i and (cid:80)
Revised: Let a in ∆ n (probability simplex) and b in ∆ n (cid:48) be the weights of the discrete measures (cid:80) i a i δ x i and (cid:80)

Original: j b j δ x (cid:48)
Revised: j b j δ x (cid:48)

Original: j with respective locations x and x (cid:48) , where δ x is the Dirac at position x .
Revised: j with respective locations x and x (cid:48) , where δ x is the Dirac at position x .

Original: Let C in R n × n (cid:48) be a matrix representing the pairwise costs for aligning the elements of x and x (cid:48) .
Revised: Let C in R n × n (cid:48) be a matrix representing the pairwise costs for aligning the elements of x and x (cid:48) .

Original: The entropic regularized
Revised: The entropic regularized

Original: Kantorovich relaxation of OT from x to x (cid:48) is
Revised: Kantorovich relaxation of OT from x to x (cid:48) is

Original: [Equation]
Revised: [Equation]

Original: where H ( P ) = − (cid:80) ij P ij (log( P ij ) − 1) is the entropic regularization with parameter ε , which controls the sparsity of P , and U is the space of admissible couplings between a and b :
Revised: where H ( P ) = − (cid:80) ij P ij (log( P ij ) − 1) is the entropic regularization with parameter ε , which controls the sparsity of P , and U is the space of admissible couplings between a and b :

Original: [Equation]
Revised: [Equation]

Original: , b )
Revised: , b )

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: The problem is typically solved by using a matrix scaling procedure known as Sinkhorn’s algorithm (Sinkhorn & Knopp, 1967; Cuturi, 2013).
Revised: The problem is typically solved by using a matrix scaling procedure known as Sinkhorn’s algorithm (Sinkhorn & Knopp, 1967; Cuturi, 2013).

Original: In practice, a and b are uniform measures since we consider the mass to be evenly distributed between the points.
Revised: In practice, a and b are uniform measures since we consider the mass to be evenly distributed between the points.

Original: P is called the transport plan, which carries the information on how to distribute the mass of x in x (cid:48) with minimal cost.
Revised: P is called the transport plan, which carries the information on how to distribute the mass of x in x (cid:48) with minimal cost.

Original: Our method uses optimal transport to align features of a given set to a learned reference.
Revised: Our method uses optimal transport to align features of a given set to a learned reference.

Original: Kernel methods.
Revised: Kernel methods.

Original: Kernel methods (Sch¨olkopf & Smola, 2001) map data living in a space X to a reproducing kernel Hilbert space H , associated to a positive deﬁnite kernel K through a mapping function ϕ : X → H , such that K ( x , x (cid:48) )
Revised: Kernel methods (Sch¨olkopf & Smola, 2001) map data living in a space X to a reproducing kernel Hilbert space H , associated to a positive deﬁnite kernel K through a mapping function ϕ : X → H , such that K ( x , x (cid:48) )

Original: =
Revised: =

Original: (cid:104) ϕ
Revised: (cid:104) ϕ

Original: ( x ) , ϕ ( x (cid:48) )
Revised: ( x ) , ϕ ( x (cid:48) )

Original: (cid:105) H .
Revised: (cid:105) H .

Original: Even though ϕ ( x ) may be inﬁnitedimensional, classical kernel approximation techniques (Williams & Seeger, 2001) provide ﬁnitedimensional embeddings ψ ( x ) in R k such that K ( x , x (cid:48) )
Revised: Even though ϕ ( x ) may be inﬁnitedimensional, classical kernel approximation techniques (Williams & Seeger, 2001) provide ﬁnitedimensional embeddings ψ ( x ) in R k such that K ( x , x (cid:48) )

Original: ≈ (cid:104) ψ ( x ) , ψ ( x (cid:48) )
Revised: ≈ (cid:104) ψ ( x ) , ψ ( x (cid:48) )

Original: (cid:105) .
Revised: (cid:105) .

Original: Our embedding for sets relies in part on kernel method principles and on such a ﬁnite-dimensional approximation.
Revised: Our embedding for sets relies in part on kernel method principles and on such a ﬁnite-dimensional approximation.

Original: 3. O PTIMAL T RANSPORT E MBEDDING AND A SSOCIATED K ERNEL
Revised: 3. O PTIMAL T RANSPORT E MBEDDING AND A SSOCIATED K ERNEL

Original: We now present the OTKE, an embedding and pooling layer which aggregates a variable-size set or sequence of features into a ﬁxed-size embedding.
Revised: We now present the OTKE, an embedding and pooling layer which aggregates a variable-size set or sequence of features into a ﬁxed-size embedding.

Original: We start with an inﬁnite-dimensional variant living in a RKHS, before introducing the ﬁnite-dimensional embedding that we use in practice.
Revised: We start with an inﬁnite-dimensional variant living in a RKHS, before introducing the ﬁnite-dimensional embedding that we use in practice.

Original: Inﬁnite-dimensional embedding in RKHS.
Revised: Inﬁnite-dimensional embedding in RKHS.

Original: Given a set x and a (learned) reference z in X with p elements, we consider an embedding Φ z ( x ) which performs the following operations: (i) initial embedding of the elements of x and z to a RKHS H ; (ii) alignment of the elements of x to the elements of z via optimal transport; (iii) weighted linear pooling of the elements x into p bins, producing an embedding Φ z ( x ) in H p , which is illustrated in Figure 1.
Revised: Given a set x and a (learned) reference z in X with p elements, we consider an embedding Φ z ( x ) which performs the following operations: (i) initial embedding of the elements of x and z to a RKHS H ; (ii) alignment of the elements of x to the elements of z via optimal transport; (iii) weighted linear pooling of the elements x into p bins, producing an embedding Φ z ( x ) in H p , which is illustrated in Figure 1.

Original: Before introducing more formal details, we note that our embedding relies on two main ideas:
Revised: Before introducing more formal details, we note that our embedding relies on two main ideas:

Original: • Global similarity-based pooling using references.
Revised: • Global similarity-based pooling using references.

Original: Learning on large sets with long-range interactions may beneﬁt from pooling to reduce the number of feature vectors.
Revised: Learning on large sets with long-range interactions may beneﬁt from pooling to reduce the number of feature vectors.

Original: Our pooling rule follows an inductive bias akin to that of self-attention: elements that are relevant to each other for the task at hand should be pooled together.
Revised: Our pooling rule follows an inductive bias akin to that of self-attention: elements that are relevant to each other for the task at hand should be pooled together.

Original: To this end, each element in the reference set corresponds to a pooling cell, where the elements of the input set are aggregated through a weighted sum.
Revised: To this end, each element in the reference set corresponds to a pooling cell, where the elements of the input set are aggregated through a weighted sum.

Original: The weights simply reﬂect the similarity between the vectors of the input set and the current vector in the reference.
Revised: The weights simply reﬂect the similarity between the vectors of the input set and the current vector in the reference.

Original: Importantly, using a reference set enables to reduce the size of the “attention matrix” from quadratic to linear in the length of the input sequence.
Revised: Importantly, using a reference set enables to reduce the size of the “attention matrix” from quadratic to linear in the length of the input sequence.

Original: • Computing similarity weights via optimal transport.
Revised: • Computing similarity weights via optimal transport.

Original: A computationally efﬁcient similarity score between two elements is their dot-product (Vaswani et al., 2017).
Revised: A computationally efﬁcient similarity score between two elements is their dot-product (Vaswani et al., 2017).

Original: In this paper, we rather consider that elements of the input set should be pooled together if they align well with the same part of the reference.
Revised: In this paper, we rather consider that elements of the input set should be pooled together if they align well with the same part of the reference.

Original: Alignment scores can efﬁciently be obtained by computing the transport plan between the input and the reference sets: Sinkhorn’s algorithm indeed enjoys fast solvers (Cuturi, 2013).
Revised: Alignment scores can efﬁciently be obtained by computing the transport plan between the input and the reference sets: Sinkhorn’s algorithm indeed enjoys fast solvers (Cuturi, 2013).

Original: We are now in shape to give a formal deﬁnition.
Revised: We are now in shape to give a formal deﬁnition.

Original: Deﬁnition 3.1 ( The optimal transport kernel embedding).
Revised: Deﬁnition 3.1 ( The optimal transport kernel embedding).

Original: Let x = ( x 1 , . . .
Revised: Let x = ( x 1 , . . .

Original: , x n ) in X be an input set of feature vectors and z = ( z 1 , . .
Revised: , x n ) in X be an input set of feature vectors and z = ( z 1 , . .

Original: . , z p ) in X be a reference set with p elements.
Revised: . , z p ) in X be a reference set with p elements.

Original: Let κ be a
Revised: Let κ be a

Original: [Figure]
Revised: [Figure]

Original: ϕ ( x 1 ) ϕ
Revised: ϕ ( x 1 ) ϕ

Original: ( x 2 ) . . .
Revised: ( x 2 ) . . .

Original: ϕ ( x n )
Revised: ϕ ( x n )

Original: Φ z ( x ) 1 . . .
Revised: Φ z ( x ) 1 . . .

Original: Φ z ( x ) p
Revised: Φ z ( x ) p

Original: P
Revised: P

Original: P n
Revised: P n

Original: P 2 p positive deﬁnite kernel, e.g. , Gaussian kernel, with RKHS H and ϕ : R d → H , its associated kernel embedding.
Revised: P 2 p positive deﬁnite kernel, e.g. , Gaussian kernel, with RKHS H and ϕ : R d → H , its associated kernel embedding.

Original: Let κ be the n × p matrix which carries the comparisons κ ( x i , z j ) , before alignment.
Revised: Let κ be the n × p matrix which carries the comparisons κ ( x i , z j ) , before alignment.

Original: Then, the transport plan between x and z , denoted by the n × p matrix
Revised: Then, the transport plan between x and z , denoted by the n × p matrix

Original: P
Revised: P

Original: ( x , z ) , is deﬁned as the unique solution of (1) when choosing the cost C = − κ , and our embedding is deﬁned as 
Revised: ( x , z ) , is deﬁned as the unique solution of (1) when choosing the cost C = − κ , and our embedding is deﬁned as Φ

Original: [Equation]
Revised: [Equation]

Original: where ϕ ( x ) :
Revised: where ϕ ( x ) :

Original: =
Revised: =

Original: [ ϕ ( x 1 ) , . . . , ϕ ( x n )] (cid:62) .
Revised: [ ϕ ( x 1 ) , . . . , ϕ ( x n )] (cid:62) .

Original: Interestingly, it is easy to show that our embedding Φ z ( x ) is associated to the positive deﬁnite kernel
Revised: Interestingly, it is easy to show that our embedding Φ z ( x ) is associated to the positive deﬁnite kernel

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original:  with P z ( x , x (cid:48) ) 
Revised: , (2) with P z ( x , x (cid:48) ) :

Original: = p × P ( x , z ) P ( x (cid:48) , z ) (cid:62) .
Revised: = p × P ( x , z ) P ( x (cid:48) , z ) (cid:62) .

Original: This is a weighted match kernel, with weights given by optimal transport in H .
Revised: This is a weighted match kernel, with weights given by optimal transport in H .

Original: The notion of pooling in the RKHS H of κ arises naturally if p ≤ n .
Revised: The notion of pooling in the RKHS H of κ arises naturally if p ≤ n .

Original: The elements of x are non-linearly embedded and then aggregated in “buckets”, one for each element in the reference z , given the values of P ( x , z ) .
Revised: The elements of x are non-linearly embedded and then aggregated in “buckets”, one for each element in the reference z , given the values of P ( x , z ) .

Original: This process is illustrated in Figure 1.
Revised: This process is illustrated in Figure 1.

Revised: We acknowledge here the concurrent work by Kolouri et al.

Revised: (2021), where a similar embedding is used for graph representation.

Original: We now expose the beneﬁts of this kernel formulation, and its relation to classical non-positive deﬁnite kernel.
Revised: We now expose the beneﬁts of this kernel formulation, and its relation to classical non-positive deﬁnite kernel.

Original: Kernel interpretation.
Revised: Kernel interpretation.

Original: Thanks to the gluing lemma (see, e.g. , Peyr´e & Cuturi, 2019), P z ( x , x (cid:48) ) is a valid transport plan and, empirically, a rough approximation of P ( x , x (cid:48) ) .
Revised: Thanks to the gluing lemma (see, e.g. , Peyr ´ e & Cuturi, 2019), P z ( x , x (cid:48) ) is a valid transport plan and, empirically, a rough approximation of P ( x , x (cid:48) ) .

Original: K z can therefore be seen as a surrogate of a well-known kernel (Rubner et al., 2000), deﬁned as
Revised: K z can therefore be seen as a surrogate of a well-known kernel (Rubner et al., 2000), deﬁned as

Original: [Equation]
Revised: [Equation]

Original: When the entropic regularization ε is equal to 0 , K OT is equivalent to the 2-Wasserstein distance
Revised: When the entropic regularization ε is equal to 0 , K OT is equivalent to the 2-Wasserstein distance

Original: W 2 ( x , x (cid:48) ) with ground metric d κ induced by kernel κ .
Revised: W 2 ( x , x (cid:48) ) with ground metric d κ induced by kernel κ .

Original: K OT is generally not positive deﬁnite (see Peyr´e & Cuturi (2019), Chapter 8 . 3 ) and computationally costly (the number of transport plans to compute is quadratic in the number of sets to process whereas it is linear for K z ).
Revised: K OT is generally not positive deﬁnite (see Peyr´e & Cuturi (2019), Chapter 8 . 3 ) and computationally costly (the number of transport plans to compute is quadratic in the number of sets to process whereas it is linear for K z ).

Original: Now, we show the relationship between this kernel and our kernel K z , which is proved in Appendix B.1.
Revised: Now, we show the relationship between this kernel and our kernel K z , which is proved in Appendix B.1.

Original: Lemma 3.1 (Relation between P ( x , x (cid:48) ) and P z ( x , x (cid:48) ) when ε = 0 ) .
Revised: Lemma 3.1 (Relation between P ( x , x (cid:48) ) and P z ( x , x (cid:48) ) when ε = 0 ) .

Original: For any x , x (cid:48) and z in X with lengths n , n (cid:48) and p , by denoting W z 2 ( x , x (cid:48) ) := (cid:104)
Revised: For any x , x (cid:48) and z in X with lengths n , n (cid:48) and p , by denoting W z 2 ( x , x (cid:48) ) := (cid:104)

Original: P z ( x , x (cid:48) ) , d 2 κ ( x , x (cid:48) )
Revised: P z ( x , x (cid:48) ) , d 2 κ ( x , x (cid:48) )

Original: (cid:105) 1 / 2 we have
Revised: (cid:105) 1 / 2 we have

Original: [Equation]
Revised: [Equation]

Original: This lemma shows that the distance W z 2 resulting from K z is related to the Wasserstein distance W 2 ; yet, this relation should not be interpreted as an approximation error as our goal is not to approximate W 2 , but rather to derive a trainable embedding Φ z ( x ) with good computational properties.
Revised: This lemma shows that the distance W z 2 resulting from K z is related to the Wasserstein distance W 2 ; yet, this relation should not be interpreted as an approximation error as our goal is not to approximate W 2 , but rather to derive a trainable embedding Φ z ( x ) with good computational properties.

Original: Lemma 3.1 roots our features and to some extent self-attention in a rich optimal transport literature.
Revised: Lemma 3.1 roots our features and to some extent self-attention in a rich optimal transport literature.

Revised: In fact, W z 2 is equivalent to a distance introduced by Wang et al.

Revised: (2013), whose properties are further studied by Moosm¨uller & Cloninger (2020).

Revised: A major difference is that W z 2 crucially relies on

Revised: Sinkhorn’s algorithm so that the references can be learned end-to-end, as explained below.

Original: 3. F ROM INFINITE - DIMENSIONAL KERNEL EMBEDDING TO FINITE DIMENSION
Revised: 3. F ROM INFINITE - DIMENSIONAL KERNEL EMBEDDING TO FINITE DIMENSION

Original: In some cases, ϕ ( x ) is already ﬁnite-dimensional, which allows to compute the embedding Φ z ( x ) explicitly.
Revised: In some cases, ϕ ( x ) is already ﬁnite-dimensional, which allows to compute the embedding Φ z ( x ) explicitly.

Original: This is particularly useful when dealing with large-scale data, as it enables us to use our method for supervised learning tasks without computing the Gram matrix, which grows quadratically in size with the number of samples.
Revised: This is particularly useful when dealing with large-scale data, as it enables us to use our method for supervised learning tasks without computing the Gram matrix, which grows quadratically in size with the number of samples.

Original: When ϕ is inﬁnite or high-dimensional, it is nevertheless possible to use an approximation based on the Nystr¨om method (Williams & Seeger, 2001), which provides an embedding ψ : R d → R k such that
Revised: When ϕ is inﬁnite or high-dimensional, it is nevertheless possible to use an approximation based on the Nystr¨om method (Williams & Seeger, 2001), which provides an embedding ψ : R d → R k such that

Original: [Equation]
Revised: [Equation]

Original: Concretely, the Nystr¨om method consists in projecting points from the RKHS H onto a linear subspace F , which is parametrized by k anchor points F = Span ( ϕ ( w 1 ) , . . .
Revised: Concretely, the Nystr¨om method consists in projecting points from the RKHS H onto a linear subspace F , which is parametrized by k anchor points F = Span ( ϕ ( w 1 ) , . . .

Original: , ϕ ( w k )) .
Revised: , ϕ ( w k )) .

Original: The corresponding embedding admits an explicit form ψ ( x i ) =
Revised: The corresponding embedding admits an explicit form ψ ( x i ) =

Original: κ ( w , w ) − 1 / 2 κ ( w , x i ) , where κ ( w , w ) is the k × k Gram matrix of κ computed on the set w = { w 1 , . .
Revised: κ ( w , w ) − 1 / 2 κ ( w , x i ) , where κ ( w , w ) is the k × k Gram matrix of κ computed on the set w = { w 1 , . .

Original: . , w k } of anchor points and κ ( w , x i ) is in R k .
Revised: . , w k } of anchor points and κ ( w , x i ) is in R k .

Original: Then, there are several ways to learn the anchor points: (a) they can be chosen as random points from data; (b) they can be deﬁned as centroids obtained by K-means, see Zhang et al.
Revised: Then, there are several ways to learn the anchor points: (a) they can be chosen as random points from data; (b) they can be deﬁned as centroids obtained by K-means, see Zhang et al.

Original: ; (c) they can be learned by back-propagation for a supervised task, see Mairal (2016).
Revised: ; (c) they can be learned by back-propagation for a supervised task, see Mairal (2016).

Original: Using such an approximation within our framework can be simply achieved by (i) replacing κ by a linear kernel and (ii) replacing each element x
Revised: Using such an approximation within our framework can be simply achieved by (i) replacing κ by a linear kernel and (ii) replacing each element x

Original: i by its embedding ψ ( x i ) in R k and considering a reference set with elements in R k .
Revised: i by its embedding ψ ( x i ) in R k and considering a reference set with elements in R k .

Original: By abuse of notation, we still use z for the new parametrization.
Revised: By abuse of notation, we still use z for the new parametrization.

Original: The embedding, which we use in practice in all our experiments, becomes simply
Revised: The embedding, which we use in practice in all our experiments, becomes simply

Original: [Equation]
Revised: [Equation]

Original: where p is the number of elements in z .
Revised: where p is the number of elements in z .

Original: Next, we discuss how to learn the reference set z .
Revised: Next, we discuss how to learn the reference set z .

Original: 3. U NSUPERVISED AND S UPERVISED L EARNING OF P ARAMETER z
Revised: 3. U NSUPERVISED AND S UPERVISED L EARNING OF P ARAMETER z

Original: Unsupervised learning.
Revised: Unsupervised learning.

Original: In the fashion of the Nystr ¨ om approximation, the p elements of z can be deﬁned as the centroids obtained by K-means applied to all features from training sets in X .
Revised: In the fashion of the Nystr ¨ om approximation, the p elements of z can be deﬁned as the centroids obtained by K-means applied to all features from training sets in X .

Original: A corollary of Lemma 3.1 suggests another algorithm: a bound on the deviation term between W 2 and
Revised: A corollary of Lemma 3.1 suggests another algorithm: a bound on the deviation term between W 2 and

Original: W z 2 for m samples ( x 1 , . . .
Revised: W z 2 for m samples ( x 1 , . . .

Original: , x m ) is indeed
Revised: , x m ) is indeed

Original: [Equation]
Revised: [Equation]

Original: The right-hand term corresponds to the objective of the Wasserstein barycenter problem (Cuturi &
Revised: The right-hand term corresponds to the objective of the Wasserstein barycenter problem (Cuturi &

Original: Doucet, 2013), which yields the mean of a set of empirical measures (here the x ’s) under the OT metric.
Revised: Doucet, 2013), which yields the mean of a set of empirical measures (here the x ’s) under the OT metric.

Original: The Wasserstein barycenter is therefore an attractive candidate for choosing z .
Revised: The Wasserstein barycenter is therefore an attractive candidate for choosing z .

Original: K-means can be seen as a particular case of Wasserstein barycenter when m = 1 (Cuturi & Doucet, 2013; Ho et al., 2017) and is faster to compute.
Revised: K-means can be seen as a particular case of Wasserstein barycenter when m = 1 (Cuturi & Doucet, 2013; Ho et al., 2017) and is faster to compute.

Original: In practice, both methods yield similar results, see Appendix C, and we thus chose K-means to learn z in unsupervised settings throughout the experiments.
Revised: In practice, both methods yield similar results, see Appendix C, and we thus chose K-means to learn z in unsupervised settings throughout the experiments.

Original: The anchor points w and the references z may be then computed using similar algorithms; however, their mathematical interpretation differs as exposed above.
Revised: The anchor points w and the references z may be then computed using similar algorithms; however, their mathematical interpretation differs as exposed above.

Original: The task of representing features (learning w in R d for a speciﬁc κ ) is decoupled from the task of aggregating (learning the reference z in R k ).
Revised: The task of representing features (learning w in R d for a speciﬁc κ ) is decoupled from the task of aggregating (learning the reference z in R k ).

Original: Supervised learning.
Revised: Supervised learning.

Original: As mentioned in Section 3.1, P ( ψ ( x ) , z ) is computed using Sinkhorn’s algorithm, recalled in Appendix A, which can be easily adapted to batches of samples x , with possibly varying lengths, leading to GPU-friendly forward computations of the embedding Φ z .
Revised: As mentioned in Section 3.1, P ( ψ ( x ) , z ) is computed using Sinkhorn’s algorithm, recalled in Appendix A, which can be easily adapted to batches of samples x , with possibly varying lengths, leading to GPU-friendly forward computations of the embedding Φ z .

Original: More important, all Sinkhorn’s operations are differentiable, which enables z to be optimized with stochastic gradient descent through back-propagation (Genevay et al., 2018), e.g. , for minimizing a classiﬁcation or regression loss function when labels are available.
Revised: More important, all Sinkhorn’s operations are differentiable, which enables z to be optimized with stochastic gradient descent through back-propagation (Genevay et al., 2018), e.g. , for minimizing a classiﬁcation or regression loss function when labels are available.

Original: In practice, a small number of Sinkhorn iterations ( e.g. , 10) is sufﬁcient to compute P ( ψ ( x ) , z ) .
Revised: In practice, a small number of Sinkhorn iterations ( e.g. , 10) is sufﬁcient to compute P ( ψ ( x ) , z ) .

Original: Since the anchors w in the embedding layer below can also be learned end-to-end (Mairal, 2016), our embedding can thus be used as a module injected into any model, e.g , a deep network, as demonstrated in our experiments.
Revised: Since the anchors w in the embedding layer below can also be learned end-to-end (Mairal, 2016), our embedding can thus be used as a module injected into any model, e.g , a deep network, as demonstrated in our experiments.

Original: 3. E XTENSIONS
Revised: 3. E XTENSIONS

Original: Integrating positional information into the embedding.
Revised: Integrating positional information into the embedding.

Original: The discussed embedding and kernel do not take the position of the features into account, which may be problematic when dealing with structured data such as images or sentences.
Revised: The discussed embedding and kernel do not take the position of the features into account, which may be problematic when dealing with structured data such as images or sentences.

Original: To this end, we borrow the idea of convolutional kernel networks, or CKN (Mairal, 2016; Mairal et al., 2014), and penalize the similarities exponentially with the positional distance between a pair of elements in the input and reference sequences.
Revised: To this end, we borrow the idea of convolutional kernel networks, or CKN (Mairal, 2016; Mairal et al., 2014), and penalize the similarities exponentially with the positional distance between a pair of elements in the input and reference sequences.

Original: More precisely, we multiply P ( ψ ( x ) , z ) element-wise by a distance matrix S deﬁned as:
Revised: More precisely, we multiply P ( ψ ( x ) , z ) element-wise by a distance matrix S deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: and replace it in the embedding.
Revised: and replace it in the embedding.

Original: With similarity weights based both on content and position, the kernel associated to our embedding can be viewed as a generalization of the CKNs (whose similarity weights are based on position only), with feature alignment based on optimal transport.
Revised: With similarity weights based both on content and position, the kernel associated to our embedding can be viewed as a generalization of the CKNs (whose similarity weights are based on position only), with feature alignment based on optimal transport.

Original: When dealing with multi-dimensional objects such as images, we just replace the index scalar i with an index vector of the same spatial dimension as the object, representing the positions of each dimension.
Revised: When dealing with multi-dimensional objects such as images, we just replace the index scalar i with an index vector of the same spatial dimension as the object, representing the positions of each dimension.

Original: Using multiple references.
Revised: Using multiple references.

Original: A naive reconstruction using different references z 1 , . . .
Revised: A naive reconstruction using different references z 1 , . . .

Original: , z q in X may yield a better approximation of the transport plan.
Revised: , z q in X may yield a better approximation of the transport plan.

Original: In this case, the embedding of x becomes
Revised: In this case, the embedding of x becomes

Original: [Equation]
Revised: [Equation]

Original: with q the number of references (the factor 1 / √ q comes from the mean).
Revised: with q the number of references (the factor 1 / √ q comes from the mean).

Original: Using Eq.
Revised: Using Eq.

Original: (4), we can obtain a bound similar to (6) for a data set of m samples ( x 1 , . . .
Revised: (4), we can obtain a bound similar to (6) for a data set of m samples ( x 1 , . . .

Original: , x m ) and q references (see Appendix B. for details).
Revised: , x m ) and q references (see Appendix B. for details).

Original: To choose multiple references, we tried a K-means algorithm with 2-Wasserstein distance for assigning clusters, and we updated the centroids as in the single-reference case.
Revised: To choose multiple references, we tried a K-means algorithm with 2-Wasserstein distance for assigning clusters, and we updated the centroids as in the single-reference case.

Original: Using multiple references appears to be useful when optimizing z with supervision (see Appendix C).
Revised: Using multiple references appears to be useful when optimizing z with supervision (see Appendix C).

Original: [Equation]
Revised: [Equation]

Original: ELATION BETWEEN OUR E MBEDDING AND S ELF -A TTENTION
Revised: ELATION BETWEEN OUR E MBEDDING AND S ELF -A TTENTION

Original: Our embedding and a single layer of transformer encoder, recalled in Appendix A, share the same type of inductive bias, i.e , aggregating features relying on similarity weights.
Revised: Our embedding and a single layer of transformer encoder, recalled in Appendix A, share the same type of inductive bias, i.e , aggregating features relying on similarity weights.

Original: We now clarify their relationship.
Revised: We now clarify their relationship.

Original: Our embedding is arguably simpler (see respectively size of attention and number of parameters in Table 1), and may compete in some settings with the transformer self-attention as illustrated in Section
Revised: Our embedding is arguably simpler (see respectively size of attention and number of parameters in Table 1), and may compete in some settings with the transformer self-attention as illustrated in Section

Original: [Table]
Revised: [Table]

Original: There is a correspondence between the values, attention matrix in the transformer and ϕ , P in Deﬁnition 3.1, yet also noticeable differences.
Revised: There is a correspondence between the values, attention matrix in the transformer and ϕ , P in Deﬁnition 3.1, yet also noticeable differences.

Original: On the one hand, Φ z aligns a given sequence x with respect to a reference z , learned with or without supervision and shared across the data set.
Revised: On the one hand, Φ z aligns a given sequence x with respect to a reference z , learned with or without supervision and shared across the data set.

Original: Our weights are computed using optimal transport.
Revised: Our weights are computed using optimal transport.

Original: On the other hand, a transformer encoder performs self-alignment: for a given x i , features are aggregated depending on a similarity score between x i and the elements of x only.
Revised: On the other hand, a transformer encoder performs self-alignment: for a given x i , features are aggregated depending on a similarity score between x i and the elements of x only.

Original: The similarity score is a matrix product between queries Q and keys K matrices, learned with supervision and shared across the data set.
Revised: The similarity score is a matrix product between queries Q and keys K matrices, learned with supervision and shared across the data set.

Original: In this regard, our work complements a recent line of research questioning the dotproduct, learned self-attention (Raganato et al., 2020;
Revised: In this regard, our work complements a recent line of research questioning the dotproduct, learned self-attention (Raganato et al., 2020;

Original: Weiqiu et al., 2020).
Revised: Weiqiu et al., 2020).

Original: Selfattention-like weights can also be obtained with our embedding by computing P ( x , z i ) P ( x , z i ) (cid:62) for each reference i .
Revised: Selfattention-like weights can also be obtained with our embedding by computing P ( x , z i ) P ( x , z i ) (cid:62) for each reference i .

Original: In that sense, our work is related to recent research on efﬁcient self-attention (Wang et al., 2020; Choromanski et al., 2020), where a low-rank approximation of the self-attention matrix is computed.
Revised: In that sense, our work is related to recent research on efﬁcient self-attention (Wang et al., 2020; Choromanski et al., 2020), where a low-rank approximation of the self-attention matrix is computed.

Original: Position smoothing and relative positional encoding.
Revised: Position smoothing and relative positional encoding.

Original: Transformers can add an absolute positional encoding to the input features (Vaswani et al., 2017).
Revised: Transformers can add an absolute positional encoding to the input features (Vaswani et al., 2017).

Original: Yet, relative positional encoding (Dai et al., 2019) is a current standard for integrating positional information: the position offset between the query element and a given key can be injected in the attention score (Tsai et al., 2019), which is equivalent to our approach.
Revised: Yet, relative positional encoding (Dai et al., 2019) is a current standard for integrating positional information: the position offset between the query element and a given key can be injected in the attention score (Tsai et al., 2019), which is equivalent to our approach.

Original: The link between CKNs and our kernel, provided by this positional encoding, stands in line with recent works casting attention and convolution into a uniﬁed framework (Andreoli, 2019).
Revised: The link between CKNs and our kernel, provided by this positional encoding, stands in line with recent works casting attention and convolution into a uniﬁed framework (Andreoli, 2019).

Original: In particular, Cordonnier et al.
Revised: In particular, Cordonnier et al.

Original: (2020) show that attention learns convolution in the setting of image classiﬁcation: the kernel pattern is learned at the same time as the ﬁlters.
Revised: (2020) show that attention learns convolution in the setting of image classiﬁcation: the kernel pattern is learned at the same time as the ﬁlters.

Original: Multiple references and attention heads.
Revised: Multiple references and attention heads.

Original: In the transformer architecture, the succession of blocks composed of an attention layer followed by a fully-connected layer is called a head, with each head potentially focusing on different parts of the input.
Revised: In the transformer architecture, the succession of blocks composed of an attention layer followed by a fully-connected layer is called a head, with each head potentially focusing on different parts of the input.

Original: Successful architectures have a few heads in parallel.
Revised: Successful architectures have a few heads in parallel.

Original: The outputs of the heads are then aggregated to output a ﬁnal embedding.
Revised: The outputs of the heads are then aggregated to output a ﬁnal embedding.

Original: A layer of our embedding with non-linear kernel κ can be seen as such a block, with the references playing the role of the heads.
Revised: A layer of our embedding with non-linear kernel κ can be seen as such a block, with the references playing the role of the heads.

Original: As some recent works question the role of attention heads (Voita et al., 2019;
Revised: As some recent works question the role of attention heads (Voita et al., 2019;

Original: Michel et al., 2019), exploring the content of our learned references z may provide another perspective on this question.
Revised: Michel et al., 2019), exploring the content of our learned references z may provide another perspective on this question.

Original: More generally, visualization and interpretation of the learned references could be of interest for biological sequences.
Revised: More generally, visualization and interpretation of the learned references could be of interest for biological sequences.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We now show the effectiveness of our embedding OTKE in tasks where samples can be expressed as large sets with potentially few labels, such as in bioinformatics.
Revised: We now show the effectiveness of our embedding OTKE in tasks where samples can be expressed as large sets with potentially few labels, such as in bioinformatics.

Original: We evaluate our embedding alone in unsupervised or supervised settings, or within a model in the supervised setting.
Revised: We evaluate our embedding alone in unsupervised or supervised settings, or within a model in the supervised setting.

Original: We also consider
Revised: We also consider

Original: NLP tasks involving shorter sequences and relatively more labels.
Revised: NLP tasks involving shorter sequences and relatively more labels.

Original: 5. D ATASETS , E XPERIMENTAL S ETUP AND B ASELINES
Revised: 5. D ATASETS , E XPERIMENTAL S ETUP AND B ASELINES

Original: In unsupervised settings, we train a linear classiﬁer with the cross entropy loss between true labels and predictions on top of the features provided by our embedding (where the references z and Nystr ¨ om anchors w have been learned without supervision), or an unsupervised baseline.
Revised: In unsupervised settings, we train a linear classiﬁer with the cross entropy loss between true labels and predictions on top of the features provided by our embedding (where the references z and Nystr ¨ om anchors w have been learned without supervision), or an unsupervised baseline.

Original: In supervised settings, the same model is initialized with our unsupervised method and then trained end-to-end (including z and w ) by minimizing the same loss.
Revised: In supervised settings, the same model is initialized with our unsupervised method and then trained end-to-end (including z and w ) by minimizing the same loss.

Original: We use an alternating optimization strategy to update the parameters for both SCOP and SST datasets, as used by Chen et al.
Revised: We use an alternating optimization strategy to update the parameters for both SCOP and SST datasets, as used by Chen et al.

Original: (2019a;b).
Revised: (2019a;b).

Original: We train for 100 epochs with Adam on both data sets: the initial learning rate is 0.01, and get halved as long as there is no decrease in the validation loss for 5 epochs.
Revised: We train for 100 epochs with Adam on both data sets: the initial learning rate is 0.01, and get halved as long as there is no decrease in the validation loss for 5 epochs.

Original: The hyper-parameters we tuned include number of supports and references p, q , entropic regularization in OT ε , the bandwidth of Gaussian kernels and the regularization parameter of the linear classiﬁer.
Revised: The hyper-parameters we tuned include number of supports and references p, q , entropic regularization in OT ε , the bandwidth of Gaussian kernels and the regularization parameter of the linear classiﬁer.

Original: The best values of ε and the bandwidth were found stable across tasks, while the regularization parameter needed to be more carefully cross-validated.
Revised: The best values of ε and the bandwidth were found stable across tasks, while the regularization parameter needed to be more carefully cross-validated.

Original: Additional results and implementation details can be found in
Revised: Additional results and implementation details can be found in

Original: Appendix C.
Revised: Appendix C.

Original: Protein fold classiﬁcation on SCOP 1.75.
Revised: Protein fold classiﬁcation on SCOP 1.75.

Original: We follow the protocol described by Hou et al.
Revised: We follow the protocol described by Hou et al.

Original: for this important task in bioinformatics.
Revised: for this important task in bioinformatics.

Original: The dataset contains 19 , 245 sequences from 1 , 195 different classes of fold (hence less than 20 labels in average per class).
Revised: The dataset contains 19 , 245 sequences from 1 , 195 different classes of fold (hence less than 20 labels in average per class).

Original: The sequence lengths vary from tens to thousands.
Revised: The sequence lengths vary from tens to thousands.

Original: Each element of a sequence is a 45 -dimensional vector.
Revised: Each element of a sequence is a 45 -dimensional vector.

Original: The objective is to classify the sequences to fold classes, which corresponds to a multiclass classiﬁcation problem.
Revised: The objective is to classify the sequences to fold classes, which corresponds to a multiclass classiﬁcation problem.

Original: The features fed to the linear classiﬁer are the output of our embedding with ϕ the Gaussian kernel mapping on k -mers (subsequences of length k ) with k ﬁxed to be 10 , which is known to perform well in this task (Chen et al., 2019a).
Revised: The features fed to the linear classiﬁer are the output of our embedding with ϕ the Gaussian kernel mapping on k -mers (subsequences of length k ) with k ﬁxed to be 10 , which is known to perform well in this task (Chen et al., 2019a).

Original: The number of anchor points for Nystr¨om method is ﬁxed to 1024 and 512 respectively for unsupervised and supervised setting.
Revised: The number of anchor points for Nystr¨om method is ﬁxed to 1024 and 512 respectively for unsupervised and supervised setting.

Original: In the unsupervised setting, we compare our method to state-of-the-art unsupervised method for this task:
Revised: In the unsupervised setting, we compare our method to state-of-the-art unsupervised method for this task:

Original: CKN (Chen et al., 2019a), which performs a global mean pooling in contrast to the global adaptive pooling performed by our embedding.
Revised: CKN (Chen et al., 2019a), which performs a global mean pooling in contrast to the global adaptive pooling performed by our embedding.

Original: In the supervised setting, we compare the same model to the following supervised models:
Revised: In the supervised setting, we compare the same model to the following supervised models:

Original: CKN, Recurrent Kernel Networks (RKN) (Chen et al., 2019b), a CNN with 10 convolutional layers named DeepSF (Hou et al., 2019), Rep the Set (Skianis et al., 2020) and Set Transformer (Lee et al., 2019), using the public implementations by their authors.
Revised: CKN, Recurrent Kernel Networks (RKN) (Chen et al., 2019b), a CNN with 10 convolutional layers named DeepSF (Hou et al., 2019), Rep the Set (Skianis et al., 2020) and Set Transformer (Lee et al., 2019), using the public implementations by their authors.

Original: Rep the Set and Set Transformer are used on the top of a convolutional layer of the same ﬁlter size as CKN to extract k -mer features.
Revised: Rep the Set and Set Transformer are used on the top of a convolutional layer of the same ﬁlter size as CKN to extract k -mer features.

Original: Their model hyper-parameters, weight decay and learning rate are tuned in the same way as for our models
Revised: Their model hyper-parameters, weight decay and learning rate are tuned in the same way as for our models

Original: (see Appendix for details).
Revised: (see Appendix for details).

Original: The default architecture of Set Transformer did not perform well due to overﬁtting.
Revised: The default architecture of Set Transformer did not perform well due to overﬁtting.

Original: We thus used a shallower architecture with one Induced Set Attention Block (ISAB), one
Revised: We thus used a shallower architecture with one Induced Set Attention Block (ISAB), one

Original: Pooling by Multihead Attention (PMA) and one linear layer, similar to the one-layer architectures of CKN and our model.
Revised: Pooling by Multihead Attention (PMA) and one linear layer, similar to the one-layer architectures of CKN and our model.

Original: The results are shown in Table 2.
Revised: The results are shown in Table 2.

Original: [Table]
Revised: [Table]

Original: Detection of chromatin proﬁles.
Revised: Detection of chromatin proﬁles.

Original: Predicting the chromatin features such as transcription factor (TF) binding from raw genomic sequences has been studied extensively in recent years.
Revised: Predicting the chromatin features such as transcription factor (TF) binding from raw genomic sequences has been studied extensively in recent years.

Original: CNNs with max pooling operations have been shown effective for this task.
Revised: CNNs with max pooling operations have been shown effective for this task.

Original: Here, we consider DeepSEA dataset (Zhou & Troyanskaya, 2015) consisting in simultaneously predicting 919 chromatin proﬁles, which can be formulated as a multi-label classiﬁcation task.
Revised: Here, we consider DeepSEA dataset (Zhou & Troyanskaya, 2015) consisting in simultaneously predicting 919 chromatin proﬁles, which can be formulated as a multi-label classiﬁcation task.

Original: DeepSEA contains 4 , 935 , 024 DNA sequences of length 1000 and each of them is associated with 919 different labels (chromatin proﬁles).
Revised: DeepSEA contains 4 , 935 , 024 DNA sequences of length 1000 and each of them is associated with 919 different labels (chromatin proﬁles).

Original: Each sequence is represented as a 1000 × 4 binary matrix through one-hot encoding and the objective is to predict which proﬁles a given sequence possesses.
Revised: Each sequence is represented as a 1000 × 4 binary matrix through one-hot encoding and the objective is to predict which proﬁles a given sequence possesses.

Original: As this problem is very imbalanced for each proﬁle, learning an unsupervised model could require an extremely large number of parameters.
Revised: As this problem is very imbalanced for each proﬁle, learning an unsupervised model could require an extremely large number of parameters.

Original: We thus only consider our supervised embedding as an adaptive pooling layer and inject it into a deep neural network, between one convolutional layer and one fully connected layer, as detailed in
Revised: We thus only consider our supervised embedding as an adaptive pooling layer and inject it into a deep neural network, between one convolutional layer and one fully connected layer, as detailed in

Original: Appendix C.4.
Revised: Appendix C.4.

Original: In our embedding, ϕ is chosen to be identity and the positional encoding described in Section 3 is used.
Revised: In our embedding, ϕ is chosen to be identity and the positional encoding described in Section 3 is used.

Original: We compare our model to a state-of-the-art CNN with 3 convolutional layers and two fully-connected layers (Zhou & Troyanskaya, 2015).
Revised: We compare our model to a state-of-the-art CNN with 3 convolutional layers and two fully-connected layers (Zhou & Troyanskaya, 2015).

Original: The results are shown in Table 3.
Revised: The results are shown in Table 3.

Original: Sentiment analysis on Stanford Sentiment Treebank.
Revised: Sentiment analysis on Stanford Sentiment Treebank.

Original: SST-2 (Socher et al., 2013) belongs to the NLP GLUE benchmark (Wang et al., 2019) and consists in predicting whether a movie review is positive.
Revised: SST-2 (Socher et al., 2013) belongs to the NLP GLUE benchmark (Wang et al., 2019) and consists in predicting whether a movie review is positive.

Original: The dataset contains 70 , 042 reviews.
Revised: The dataset contains 70 , 042 reviews.

Original: The test predictions need to be submitted on the GLUE leaderboard, so that we remove a portion of the training set for validation purpose and report accuracies on the actual validation set used as a test set.
Revised: The test predictions need to be submitted on the GLUE leaderboard, so that we remove a portion of the training set for validation purpose and report accuracies on the actual validation set used as a test set.

Original: Our model is one layer of our embedding with ϕ a Gaussian kernel mapping with 64 Nystr ¨ om ﬁlters in the supervised setting, and a linear mapping in the unsupervised setting.
Revised: Our model is one layer of our embedding with ϕ a Gaussian kernel mapping with 64 Nystr ¨ om ﬁlters in the supervised setting, and a linear mapping in the unsupervised setting.

Original: The features used in our model and all baselines are word vectors with dimension 768 provided by the HuggingFace implementation (Wolf et al., 2019) of the transformer BERT (Devlin et al., 2019).
Revised: The features used in our model and all baselines are word vectors with dimension 768 provided by the HuggingFace implementation (Wolf et al., 2019) of the transformer BERT (Devlin et al., 2019).

Original: State-of-the-art accuracies are usually obtained after supervised ﬁne-tuning of pre-trained transformers.
Revised: State-of-the-art accuracies are usually obtained after supervised ﬁne-tuning of pre-trained transformers.

Original: Training a linear model on pre-trained features after simple pooling ( e.g , mean) also yields good results.
Revised: Training a linear model on pre-trained features after simple pooling ( e.g , mean) also yields good results.

Original: [CLS], which denotes the BERT embedding used for classiﬁcation, is also a common baseline.
Revised: [CLS], which denotes the BERT embedding used for classiﬁcation, is also a common baseline.

Original: The results are shown in Table 4.
Revised: The results are shown in Table 4.

Original: 5. R ESULTS AND DISCUSSION
Revised: 5. R ESULTS AND DISCUSSION

Original: In protein fold classiﬁcation, our embedding outperforms all baselines in both unsupervised and supervised settings.
Revised: In protein fold classiﬁcation, our embedding outperforms all baselines in both unsupervised and supervised settings.

Original: Surprisingly, our unsupervised model also achieves better results than most supervised baselines.
Revised: Surprisingly, our unsupervised model also achieves better results than most supervised baselines.

Original: In contrast, Set Transformer does not perform well, possibly because its implementation was not designed for sets with varying sizes, and tasks with few annotations.
Revised: In contrast, Set Transformer does not perform well, possibly because its implementation was not designed for sets with varying sizes, and tasks with few annotations.

Original: In detection of chromatin proﬁles, our model (our embedding within a deep network) has fewer layers than state-of-the-art CNNs while outperforming them, which advocates for the use of attentionbased models for such applications.
Revised: In detection of chromatin proﬁles, our model (our embedding within a deep network) has fewer layers than state-of-the-art CNNs while outperforming them, which advocates for the use of attentionbased models for such applications.

Original: Our results also suggest that positional information is important
Revised: Our results also suggest that positional information is important

Original: [Table] [Table]
Revised: [Table] [Table]

Original: (Appendix C.4;C.2), and our Gaussian position encoding outperforms the sinusoidal one introduced in Vaswani et al.
Revised: (Appendix C.4;C.2), and our Gaussian position encoding outperforms the sinusoidal one introduced in Vaswani et al.

Original: Note that in contrast to a typical transformer, which would have stored a 1000 × 1000 attention matrix, our attention score with a reference of size 64 is only 1000 × 64 , which illustrates the discussion in Section 4.
Revised: Note that in contrast to a typical transformer, which would have stored a 1000 × 1000 attention matrix, our attention score with a reference of size 64 is only 1000 × 64 , which illustrates the discussion in Section 4.

Original: In NLP, an a priori less favorable setting since sequences are shorter and there are more data, our supervised embedding gets close to a strong state-of-the-art, i.e. a fully-trained transformer.
Revised: In NLP, an a priori less favorable setting since sequences are shorter and there are more data, our supervised embedding gets close to a strong state-of-the-art, i.e. a fully-trained transformer.

Original: We observed our method to be much faster than RepSet, as fast as
Revised: We observed our method to be much faster than RepSet, as fast as

Original: Set Transformer, yet slower than ApproxRepSet (C.3).
Revised: Set Transformer, yet slower than ApproxRepSet (C.3).

Original: Using the OT plan as similarity score yields better accuracies than the dot-product between the input sets and the references (see Table 2; 4).
Revised: Using the OT plan as similarity score yields better accuracies than the dot-product between the input sets and the references (see Table 2; 4).

Original: Choice of parameters.
Revised: Choice of parameters.

Original: This paragraph sums up the impact of hyper-parameter choices.
Revised: This paragraph sums up the impact of hyper-parameter choices.

Original: Experiments justifying our claims can be found in Appendix C.
Revised: Experiments justifying our claims can be found in Appendix C.

Original: • Number of references q : for biological sequences, a single reference was found to be enough in the unsupervised case, see Table 11.
Revised: • Number of references q : for biological sequences, a single reference was found to be enough in the unsupervised case, see Table 11.

Original: In the supervised setting, Table 14 suggests that using q = 5 provides slightly better results but q = 1 remains a good baseline, and that the sensitivity to number of references is moderate.
Revised: In the supervised setting, Table 14 suggests that using q = 5 provides slightly better results but q = 1 remains a good baseline, and that the sensitivity to number of references is moderate.

Original: • Number of supports p in a reference:
Revised: • Number of supports p in a reference:

Original: Table 11 and Table 14 suggest that the sensitivity of the model to the number of supports is also moderate.
Revised: Table 11 and Table 14 suggest that the sensitivity of the model to the number of supports is also moderate.

Original: • Nystr ¨ om anchors: an anchor can be seen as a neuron in a feed-forward neural network (see expression of ψ in 3.3).
Revised: • Nystr ¨ om anchors: an anchor can be seen as a neuron in a feed-forward neural network (see expression of ψ in 3.3).

Original: In unsupervised settings, the more anchors, the better the approximation of the kernel matrix.
Revised: In unsupervised settings, the more anchors, the better the approximation of the kernel matrix.

Original: Then, the performance saturates, see Table 12.
Revised: Then, the performance saturates, see Table 12.

Original: In supervised settings, the optimal number of anchors points is much smaller, as also observed by Chen et al. (2019a), Fig 6.
Revised: In supervised settings, the optimal number of anchors points is much smaller, as also observed by Chen et al. (2019a), Fig 6.

Original: • Bandwidth σ in gaussian kernel: σ was chosen as in Chen et al.
Revised: • Bandwidth σ in gaussian kernel: σ was chosen as in Chen et al.

Original: (2019b)
Revised: (2019b)

Original: and we did not try to optimize it in this work, as it seemed to already provide good results.
Revised: and we did not try to optimize it in this work, as it seemed to already provide good results.

Original: Nevertheless, slightly better results can be obtained when tuning this parameter, for instance in SST-2.
Revised: Nevertheless, slightly better results can be obtained when tuning this parameter, for instance in SST-2.

Original: OTKE and self-supervised methods.
Revised: OTKE and self-supervised methods.

Original: Our approach should not be positioned against selfsupervision and instead brings complementary features: the OTKE may be plugged in state-ofthe-art models pre-trained on large unannotated corpus.
Revised: Our approach should not be positioned against selfsupervision and instead brings complementary features: the OTKE may be plugged in state-ofthe-art models pre-trained on large unannotated corpus.

Original: For instance, on SCOP 1.75, we use ESM1 (Rives et al., 2019), pretrained on 250 millions protein sequences, with mean pooling followed by a linear classiﬁer.
Revised: For instance, on SCOP 1.75, we use ESM1 (Rives et al., 2019), pretrained on 250 millions protein sequences, with mean pooling followed by a linear classiﬁer.

Original: As we do not have the computational ressources to ﬁne-tune ESM1-t34, we only train a linear layer on top of the extracted features.
Revised: As we do not have the computational ressources to ﬁne-tune ESM1-t34, we only train a linear layer on top of the extracted features.

Original: Using the same model, we replace the mean pooling by our (unsupervised) OTKE layer, and also only train the linear layer.
Revised: Using the same model, we replace the mean pooling by our (unsupervised) OTKE layer, and also only train the linear layer.

Original: This results in accuracy improvements as showed in Table 5.
Revised: This results in accuracy improvements as showed in Table 5.

Original: While training huge self-supervised learning models on large datasets is very effective, ESM1-t34 admits more than 2500 times more parameters than our single-layer OTKE model (260k parameters versus 670M) and our single-layer OTKE outperforms smaller versions of ESM1 (43M parameters).
Revised: While training huge self-supervised learning models on large datasets is very effective, ESM1-t34 admits more than 2500 times more parameters than our single-layer OTKE model (260k parameters versus 670M) and our single-layer OTKE outperforms smaller versions of ESM1 (43M parameters).

Original: Finally, self-supervised pre
Revised: Finally, self-supervised pre

Original: -training of a deep model including OTKE on large data sets would be interesting for fair comparison.
Revised: -training of a deep model including OTKE on large data sets would be interesting for fair comparison.

Original: [Table]
Revised: [Table]

Original: Multi-layer extension.
Revised: Multi-layer extension.

Original: Extending the OTKE to a multi-layer embedding is a natural yet not straightforward research direction: it is not clear how to ﬁnd a right deﬁnition of intermediate feature aggregation in a multi-layer OTKE model.
Revised: Extending the OTKE to a multi-layer embedding is a natural yet not straightforward research direction: it is not clear how to ﬁnd a right deﬁnition of intermediate feature aggregation in a multi-layer OTKE model.

Original: Note that for DeepSEA, our model with single-layer
Revised: Note that for DeepSEA, our model with single-layer

Original: OTKE already outperforms a multi-layer CNN, which suggests that a multi-layer OTKE is not always needed.
Revised: OTKE already outperforms a multi-layer CNN, which suggests that a multi-layer OTKE is not always needed.

Original: A CKNOWLEDGMENTS
Revised: A CKNOWLEDGMENTS

Original: JM, GM and DC were supported by the ERC grant number 714381 (SOLARIS project) and by ANR 3IA MIAI@Grenoble Alpes, (ANR-19-P3IA-0003).
Revised: JM, GM and DC were supported by the ERC grant number 714381 (SOLARIS project) and by ANR 3IA MIAI@Grenoble Alpes, (ANR-19-P3IA-0003).

Original: AA would like to acknowledge support from the ML and Optimisation joint research initiative with the fonds AXA pour la recherche and Kamet
Revised: AA would like to acknowledge support from the ML and Optimisation joint research initiative with the fonds AXA pour la recherche and Kamet

Original: Ventures, a Google focused award, as well as funding by the French government under management of Agence Nationale de la Recherche as part of the “Investissements d’avenir” program, reference
Revised: Ventures, a Google focused award, as well as funding by the French government under management of Agence Nationale de la Recherche as part of the “Investissements d’avenir” program, reference

Original: ANR-19-P3IA-0001 (PRAIRIE 3IA Institute).
Revised: ANR-19-P3IA-0001 (PRAIRIE 3IA Institute).

Original: DC and GM thank Laurent Jacob, Louis Martin,
Revised: DC and GM thank Laurent Jacob, Louis Martin,

Original: Franc¸ois-Pierre Paty and Thomas Wolf for useful discussions.
Revised: Franc¸ois-Pierre Paty and Thomas Wolf for useful discussions.

Original: Appendix
Revised: Appendix

Original: Appendix A provides some background on notions used throughout the paper; Appendix B contains the proofs skipped in the paper; Appendix C provides additional experimental results as well as details on our protocol for reproducibility.
Revised: Appendix A provides some background on notions used throughout the paper; Appendix B contains the proofs skipped in the paper; Appendix C provides additional experimental results as well as details on our protocol for reproducibility.

Original: A A DDITIONAL B ACKGROUND
Revised: A A DDITIONAL B ACKGROUND

Original: This section provides some background on attention and transformers, Sinkhorn’s algorithm and the relationship between optimal transport based kernels and positive deﬁnite histogram kernels.
Revised: This section provides some background on attention and transformers, Sinkhorn’s algorithm and the relationship between optimal transport based kernels and positive deﬁnite histogram kernels.

Original: A. S INKHORN ’ S A LGORITHM : F AST C OMPUTATION OF P κ ( x , z )
Revised: A. S INKHORN ’ S A LGORITHM : F AST C OMPUTATION OF P κ ( x , z )

Original: Without loss of generality, we consider here κ the linear kernel.
Revised: Without loss of generality, we consider here κ the linear kernel.

Original: We recall that P κ ( x , z ) is the solution of an optimal transport problem, which can be efﬁciently solved by Sinkhorn’s algorithm (Peyr ´ e & Cuturi, 2019) involving matrix multiplications only.
Revised: We recall that P κ ( x , z ) is the solution of an optimal transport problem, which can be efﬁciently solved by Sinkhorn’s algorithm (Peyr ´ e & Cuturi, 2019) involving matrix multiplications only.

Original: Speciﬁcally, Sinkhorn’s algorithm is an iterative matrix scaling method that takes the opposite of the pairwise similarity matrix K with entry
Revised: Speciﬁcally, Sinkhorn’s algorithm is an iterative matrix scaling method that takes the opposite of the pairwise similarity matrix K with entry

Original: K ij := (cid:104) x i , z j (cid:105) as input C and outputs the optimal transport plan P κ ( x , z ) =
Revised: K ij := (cid:104) x i , z j (cid:105) as input C and outputs the optimal transport plan P κ ( x , z ) =

Original: Sinkhorn ( K , ε ) .
Revised: Sinkhorn ( K , ε ) .

Original: Each iteration step (cid:96) performs the following updates
Revised: Each iteration step (cid:96) performs the following updates

Original: [Equation]
Revised: [Equation]

Original: where E = e K /ε .
Revised: where E = e K /ε .

Original: Then the matrix diag ( u ( (cid:96) ) )
Revised: Then the matrix diag ( u ( (cid:96) ) )

Original: E diag ( v ( (cid:96) ) ) converges to P κ ( x , z ) when (cid:96) tends to ∞ .
Revised: E diag ( v ( (cid:96) ) ) converges to P κ ( x , z ) when (cid:96) tends to ∞ .

Original: However when ε becomes too small, some of the elements of a matrix product Ev or E (cid:62) u become null and result in a division by 0.
Revised: However when ε becomes too small, some of the elements of a matrix product Ev or E (cid:62) u become null and result in a division by 0.

Original: To overcome this numerical stability issue, computing the multipliers u and v is preferred (see e.g. (Peyr ´ e & Cuturi, 2019, Remark 4.23)).
Revised: To overcome this numerical stability issue, computing the multipliers u and v is preferred (see e.g. (Peyr ´ e & Cuturi, 2019, Remark 4.23)).

Original: This algorithm can be easily adapted to a batch of data points x , and with possibly varying lengths via a mask vector masking on the padding positions of each data point x , leading to GPU-friendly computation.
Revised: This algorithm can be easily adapted to a batch of data points x , and with possibly varying lengths via a mask vector masking on the padding positions of each data point x , leading to GPU-friendly computation.

Original: More importantly, all the operations above at each step are differentiable, which enables z to be optimized through back-propagation.
Revised: More importantly, all the operations above at each step are differentiable, which enables z to be optimized through back-propagation.

Original: Consequently, this module can be injected into any deep networks.
Revised: Consequently, this module can be injected into any deep networks.

Original: A. A TTENTION AND TRANSFORMERS
Revised: A. A TTENTION AND TRANSFORMERS

Original: We clarify the concept of attention — a mechanism yielding a context-dependent embedding for each element of x — as a special case of non-local operations (Wang et al., 2017; Buades et al., 2011), so that it is easier to understand its relationship to the OTK.
Revised: We clarify the concept of attention — a mechanism yielding a context-dependent embedding for each element of x — as a special case of non-local operations (Wang et al., 2017; Buades et al., 2011), so that it is easier to understand its relationship to the OTK.

Original: Let us assume we are given a set x ∈ X of length n .
Revised: Let us assume we are given a set x ∈ X of length n .

Original: A non-local operation on x is a function Φ : X (cid:55)→ X such that
Revised: A non-local operation on x is a function Φ : X (cid:55)→ X such that

Original: [Equation]
Revised: [Equation]

Original: where W ( x ) i denotes the i -th column of W ( x ) , a weighting function, and V ( x ) =
Revised: where W ( x ) i denotes the i -th column of W ( x ) , a weighting function, and V ( x ) =

Original: [ v ( x 1 ) , . . .
Revised: [ v ( x 1 ) , . . .

Original: , v ( x n )] (cid:62) , an embedding.
Revised: , v ( x n )] (cid:62) , an embedding.

Original: In contrast to operations on local neighborhood such as convolutions, non-local operations theoretically account for long range dependencies between elements in the set.
Revised: In contrast to operations on local neighborhood such as convolutions, non-local operations theoretically account for long range dependencies between elements in the set.

Original: In attention and the context of neural networks, w is a learned function reﬂecting the relevance of each other elements x j with respect to the element x i being embedded and given the task at hand.
Revised: In attention and the context of neural networks, w is a learned function reﬂecting the relevance of each other elements x j with respect to the element x i being embedded and given the task at hand.

Original: In the context of the paper, we compare to a type of attention coined as dot-product selfattention , which can typically be found in the encoder part of the transformer architecture (Vaswani et al., 2017).
Revised: In the context of the paper, we compare to a type of attention coined as dot-product selfattention , which can typically be found in the encoder part of the transformer architecture (Vaswani et al., 2017).

Original: Transformers are neural network models relying mostly on a succession of an attention layer followed by a fully-connected layer.
Revised: Transformers are neural network models relying mostly on a succession of an attention layer followed by a fully-connected layer.

Original: Transformers can be used in sequence-to-sequence tasks — in this setting, they have an encoder with self-attention and a decoder part with a variant of self-attention —, or in sequence to label tasks, with only the encoder part.
Revised: Transformers can be used in sequence-to-sequence tasks — in this setting, they have an encoder with self-attention and a decoder part with a variant of self-attention —, or in sequence to label tasks, with only the encoder part.

Original: The paper deals with the latter.
Revised: The paper deals with the latter.

Original: The name self-attention means that the attention is computed using a dot-product of linear transformations of x i and x j , and x attends to itself only.
Revised: The name self-attention means that the attention is computed using a dot-product of linear transformations of x i and x j , and x attends to itself only.

Original: In its matrix formulation, dot-product self-attention is a non-local operation whose matching vector is
Revised: In its matrix formulation, dot-product self-attention is a non-local operation whose matching vector is

Original: [Equation]
Revised: [Equation]

Original: where W Q ∈ R n × d k and W K ∈ R n × d k are learned by the network.
Revised: where W Q ∈ R n × d k and W K ∈ R n × d k are learned by the network.

Original: In order to know which x j are relevant to x i , the network computes scores between a query for x i ( W Q x i ) and keys of all the elements of x ( W K x ).
Revised: In order to know which x j are relevant to x i , the network computes scores between a query for x i ( W Q x i ) and keys of all the elements of x ( W K x ).

Original: The softmax turns the scores into a weight vector in the simplex.
Revised: The softmax turns the scores into a weight vector in the simplex.

Original: Moreover, a linear mapping V ( x ) =
Revised: Moreover, a linear mapping V ( x ) =

Original: W V x , the values, is also learned.
Revised: W V x , the values, is also learned.

Original: W Q and W K are often shared (Kitaev et al., 2020).
Revised: W Q and W K are often shared (Kitaev et al., 2020).

Original: A drawback of such attention is that for a sequence of length n , the model has to store an attention matrix W with size O ( n 2 ) .
Revised: A drawback of such attention is that for a sequence of length n , the model has to store an attention matrix W with size O ( n 2 ) .

Original: More details can be found in Vaswani et al.
Revised: More details can be found in Vaswani et al.

Original: B P ROOFS B. P ROOF OF L EMMA 3. 
Revised: B P ROOFS B. P ROOF OF L EMMA 3.

Original: Proof.
Revised: Proof.

Original: First, since (cid:80) n (cid:48)
Revised: First, since (cid:80) n (cid:48)

Original: j =1 p P ( x (cid:48) , z ) jk = 1 for any k , we have
Revised: j =1 p P ( x (cid:48) , z ) jk = 1 for any k , we have

Original: W 2 ( x , z ) 2 = n (cid:88) i = p (cid:88) k
Revised: W 2 ( x , z ) 2 = n (cid:88) i = p (cid:88) k

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: P (
Revised: P (

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: (cid:88) i =
Revised: (cid:88) i =

Original: [Equation]
Revised: [Equation]

Original: (cid:88)
Revised: (cid:88)

Original: [Equation]
Revised: [Equation]

Original: n
Revised: n

Original: [Equation]
Revised: [Equation]

Original: (cid:88)
Revised: (cid:88)

Original: j =
Revised: j =

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: (cid:107) u (cid:107) 22 , with u a vector in R nn (cid:48)
Revised: (cid:107) u (cid:107) 22 , with u a vector in R nn (cid:48)

Original: p whose entries are (cid:112) p P ( x (cid:48) , z ) jk P ( x , z )
Revised: p whose entries are (cid:112) p P ( x (cid:48) , z ) jk P ( x , z )

Original: ik
Revised: ik

Original: d κ ( x i , z k ) for i = 1 , . . . , n , j = 1 , . . . , n (cid:48) and k = 1 , . . .
Revised: d κ ( x i , z k ) for i = 1 , . . . , n , j = 1 , . . . , n (cid:48) and k = 1 , . . .

Original: , p .
Revised: , p .

Original: We can also rewrite W z 2 ( x , x (cid:48) ) as an (cid:96) 2 -norm of a vector v in
Revised: We can also rewrite W z 2 ( x , x (cid:48) ) as an (cid:96) 2 -norm of a vector v in

Original: R nn (cid:48) p whose entries are (cid:112) p P ( x (cid:48) , z ) jk P ( x , z ) ik d κ ( x i , x (cid:48) j ) .
Revised: R nn (cid:48) p whose entries are (cid:112) p P ( x (cid:48) , z ) jk P ( x , z ) ik d κ ( x i , x (cid:48) j ) .

Original: Then by Minkowski inequality for the (cid:96) 2 -norm, we have
Revised: Then by Minkowski inequality for the (cid:96) 2 -norm, we have

Original: [Equation]
Revised: [Equation]

Original: where the second inequality is the triangle inequality for the distance d κ .
Revised: where the second inequality is the triangle inequality for the distance d κ .

Original: Finally, we have
Revised: Finally, we have

Original: [Equation]
Revised: [Equation]

Original: where the second inequality is the triangle inequality for the 2-Wasserstein distance.
Revised: where the second inequality is the triangle inequality for the 2-Wasserstein distance.

Original: By symmetry, we also have | W 2 ( x , x (cid:48) )
Revised: By symmetry, we also have | W 2 ( x , x (cid:48) )

Original: − W z 2 ( x , x (cid:48) )
Revised: − W z 2 ( x , x (cid:48) )

Original: | ≤ 2 W 2 ( x , z ) , which concludes the proof.
Revised: | ≤ 2 W 2 ( x , z ) , which concludes the proof.

Original: B. R ELATIONSHIP BETWEEN W 2 AND W z 2 FOR MULTIPLE REFERENCES
Revised: B. R ELATIONSHIP BETWEEN W 2 AND W z 2 FOR MULTIPLE REFERENCES

Original: Using the relation prooved in Appendix B.1, we can obtain a bound on the error term between Wand W z 2 for a data set of m samples ( x 1 , . . .
Revised: Using the relation prooved in Appendix B.1, we can obtain a bound on the error term between Wand W z 2 for a data set of m samples ( x 1 , . . .

Original: , x m ) and q references ( z 1 , . . .
Revised: , x m ) and q references ( z 1 , . . .

Original: , z q )
Revised: , z q )

Original: [Equation]
Revised: [Equation]

Original: When q = 1 , the right-hand term in the inequality is the objective to minimize in the Wasserstein barycenter problem (Cuturi & Doucet, 2013), which further explains why we considered it:
Revised: When q = 1 , the right-hand term in the inequality is the objective to minimize in the Wasserstein barycenter problem (Cuturi & Doucet, 2013), which further explains why we considered it:

Original: Once
Revised: Once

Original: W z 2 is close to the Wasserstein distance W 2 , K z will also be close to K OT .
Revised: W z 2 is close to the Wasserstein distance W 2 , K z will also be close to K OT .

Original: We extend here the bound in equation 6 in the case of one reference to the multiple-reference case.
Revised: We extend here the bound in equation 6 in the case of one reference to the multiple-reference case.

Original: The approximate 2-Wasserstein distance W z 2 ( x , x (cid:48) ) thus becomes
Revised: The approximate 2-Wasserstein distance W z 2 ( x , x (cid:48) ) thus becomes

Original: [Equation]
Revised: [Equation]

Original: Then by Minkowski inequality for the (cid:96) 2 -norm we have
Revised: Then by Minkowski inequality for the (cid:96) 2 -norm we have

Original: [Equation]
Revised: [Equation]

Original: and by equation 6 we have
Revised: and by equation 6 we have

Original: [Equation]
Revised: [Equation]

Original: Finally the approximation error in terms of Frobenius is bounded by
Revised: Finally the approximation error in terms of Frobenius is bounded by

Original: [Equation]
Revised: [Equation]

Original: In particular
Revised: In particular

Original: [Equation]
Revised: [Equation]

Original: when q = 1 that is the case of single reference
Revised: when q = 1 that is the case of single reference

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Original: where the right term equals to the objective of the Wasserstein barycenter problem
Revised: where the right term equals to the objective of the Wasserstein barycenter problem

Original: [Equation]
Revised: [Equation]

Original: which justiﬁes the choice of z when learning without supervision
Revised: which justiﬁes the choice of z when learning without supervision

Original: [Equation]
Revised: [Equation]

Original: C A DDITIONAL E XPERIMENTS AND S ETUP D ETAILS
Revised: C A DDITIONAL E XPERIMENTS AND S ETUP D ETAILS

Original: This section contains additional experiments on CIFAR-10, whose purpose is to illustrate the kernel associated with our embedding with respect to other classical or optimal transport based kernels, and test our embedding on another data modality; additional results for the experiments of the main section; details on our setup, in particular hyper-parameter tuning for our methods and the baselines.
Revised: This section contains additional experiments on CIFAR-10, whose purpose is to illustrate the kernel associated with our embedding with respect to other classical or optimal transport based kernels, and test our embedding on another data modality; additional results for the experiments of the main section; details on our setup, in particular hyper-parameter tuning for our methods and the baselines.

Original: C. E XPERIMENTS ON K ERNEL M ATRICES ( ONLY FOR SMALL DATA SETS ).
Revised: C. E XPERIMENTS ON K ERNEL M ATRICES ( ONLY FOR SMALL DATA SETS ).

Original: Here, we compare the optimal transport kernel K OT (3) and its surrogate K z (2) (with z learned without supervision) to common and other OT kernels.
Revised: Here, we compare the optimal transport kernel K OT (3) and its surrogate K z (2) (with z learned without supervision) to common and other OT kernels.

Original: Although our embedding Φ z is scalable, the exact kernel require the computation of Gram matrices.
Revised: Although our embedding Φ z is scalable, the exact kernel require the computation of Gram matrices.

Original: For this toy experiment, we therefore use
Revised: For this toy experiment, we therefore use

Original: [Table]
Revised: [Table]

Original: 5000 samples only of CIFAR-10 (images with 32 × 32 pixels), encoded without supervision using a two-layer convolutional kernel network (Mairal, 2016).
Revised: 5000 samples only of CIFAR-10 (images with 32 × 32 pixels), encoded without supervision using a two-layer convolutional kernel network (Mairal, 2016).

Original: The resulting features are 3 × 3 patches living in R d with d = 256 or 8192 .
Revised: The resulting features are 3 × 3 patches living in R d with d = 256 or 8192 .

Original: K OT and K z aggregate existing features depending on the ground cost deﬁned by − κ (Gaussian kernel) given the computed weight matrix P .
Revised: K OT and K z aggregate existing features depending on the ground cost deﬁned by − κ (Gaussian kernel) given the computed weight matrix P .

Original: In that sense, we can say that these kernels work as an adaptive pooling.
Revised: In that sense, we can say that these kernels work as an adaptive pooling.

Original: We therefore compare it to kernel matrices corresponding to mean pooling and no pooling at all (linear kernel).
Revised: We therefore compare it to kernel matrices corresponding to mean pooling and no pooling at all (linear kernel).

Original: We also compare to a recent positive deﬁnite and fast optimal transport based kernel, the Sliced Wasserstein Kernel (Kolouri et al., 2016) with 10 , 100 and 1000 projection directions.
Revised: We also compare to a recent positive deﬁnite and fast optimal transport based kernel, the Sliced Wasserstein Kernel (Kolouri et al., 2016) with 10 , 100 and 1000 projection directions.

Original: We add a positional encoding to it so as to have a fair comparison with our kernels.
Revised: We add a positional encoding to it so as to have a fair comparison with our kernels.

Original: A linear classiﬁer is trained from this matrices.
Revised: A linear classiﬁer is trained from this matrices.

Original: Although we cannot prove that K OT is positive deﬁnite, the classiﬁer trained on the kernel matrix converges when ε is not too small.
Revised: Although we cannot prove that K OT is positive deﬁnite, the classiﬁer trained on the kernel matrix converges when ε is not too small.

Original: The results can be seen on Table 6.
Revised: The results can be seen on Table 6.

Original: Without positional information, our kernels do better than Mean pooling.
Revised: Without positional information, our kernels do better than Mean pooling.

Original: When the positions are encoded, the Linear kernel is also outperformed.
Revised: When the positions are encoded, the Linear kernel is also outperformed.

Original: Note that including positions in Mean pooling and Linear kernel means interpolating between these two kernels: in the Linear kernel, only patches with same index are compared while in the Mean pooling, all patches are compared.
Revised: Note that including positions in Mean pooling and Linear kernel means interpolating between these two kernels: in the Linear kernel, only patches with same index are compared while in the Mean pooling, all patches are compared.

Original: All interpolations did worse than the Linear kernel.
Revised: All interpolations did worse than the Linear kernel.

Original: The runtimes illustrate the scalability of K z .
Revised: The runtimes illustrate the scalability of K z .

Original: C. CIFAR- 
Revised: C. CIFAR-

Original: Here, we test our embedding on the same data modality: we use CIFAR-10 features, i.e. , 60 ,images with 32 × 32 pixels and 10 classes encoded using a two-layer CKN (Mairal, 2016), one of the baseline architectures for unsupervised learning of CIFAR-10, and evaluate on the standard test set.
Revised: Here, we test our embedding on the same data modality: we use CIFAR-10 features, i.e. , 60 ,images with 32 × 32 pixels and 10 classes encoded using a two-layer CKN (Mairal, 2016), one of the baseline architectures for unsupervised learning of CIFAR-10, and evaluate on the standard test set.

Original: The very best conﬁguration of the CKN yields a small number ( 3 × 3 ) of high-dimensional ( 16 , 384 ) patches and an accuracy of 85 .
Revised: The very best conﬁguration of the CKN yields a small number ( 3 × 3 ) of high-dimensional ( 16 , 384 ) patches and an accuracy of 85 .

Original: We will illustrate our embedding on a conﬁguration which performs slightly less but provides more patches ( 16 × 16 ), a setting for which it was designed.
Revised: We will illustrate our embedding on a conﬁguration which performs slightly less but provides more patches ( 16 × 16 ), a setting for which it was designed.

Original: The input of our embedding are unsupervised features extracted from a 2-layer CKN with kernel sizes equal to 3 and 3, and Gaussian pooling size equal to 2 and 1.
Revised: The input of our embedding are unsupervised features extracted from a 2-layer CKN with kernel sizes equal to 3 and 3, and Gaussian pooling size equal to 2 and 1.

Original: We consider the following conﬁgurations of the number of ﬁlters at each layer, to simulate two different input dimensions for our embedding:
Revised: We consider the following conﬁgurations of the number of ﬁlters at each layer, to simulate two different input dimensions for our embedding:

Original: • 64 ﬁlters at ﬁrst and 256 at second layer, which yields a 16 × 16 × 256 representation for each image.
Revised: • 64 ﬁlters at ﬁrst and 256 at second layer, which yields a 16 × 16 × 256 representation for each image.

Original: • 256 ﬁlters at ﬁrst and 1024 at second layer, which yields a 16 × 16 × 1024 representation for each image.
Revised: • 256 ﬁlters at ﬁrst and 1024 at second layer, which yields a 16 × 16 × 1024 representation for each image.

Original: Since the features are the output of a Gaussian embedding, κ in our embedding will be the linear kernel.
Revised: Since the features are the output of a Gaussian embedding, κ in our embedding will be the linear kernel.

Original: The embedding is learned with one reference and various supports using K-means method described in Section 3, and compared to several classical pooling baselines, including the original
Revised: The embedding is learned with one reference and various supports using K-means method described in Section 3, and compared to several classical pooling baselines, including the original

Original: CKN’s Gaussian pooling with pooling size equal to 6.
Revised: CKN’s Gaussian pooling with pooling size equal to 6.

Original: The hyper-parameters are the entropic regularization ε and bandwidth for position encoding σ pos .
Revised: The hyper-parameters are the entropic regularization ε and bandwidth for position encoding σ pos .

Original: Their search grids are shown in Table 7 and the results in Table 8.
Revised: Their search grids are shown in Table 7 and the results in Table 8.

Original: Without supervision, the adaptive pooling of the CKN features by our embedding notably improves their performance.
Revised: Without supervision, the adaptive pooling of the CKN features by our embedding notably improves their performance.

Original: We notice that the position encoding is very important to this task, which substantially improves the performance of its counterpart without it.
Revised: We notice that the position encoding is very important to this task, which substantially improves the performance of its counterpart without it.

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: C. P ROTEIN FOLD RECOGNITION
Revised: C. P ROTEIN FOLD RECOGNITION

Original: Dataset description.
Revised: Dataset description.

Original: Our protein fold recognition experiments consider the Structural Classiﬁcation Of Proteins (SCOP) version 1.75 and 2.06.
Revised: Our protein fold recognition experiments consider the Structural Classiﬁcation Of Proteins (SCOP) version 1.75 and 2.06.

Original: We follow the data preprocessing protocols in Hou et al.
Revised: We follow the data preprocessing protocols in Hou et al.

Original: [Table]
Revised: [Table]

Original: which yields a training and validation set composed of 14699 and 2013 sequences from SCOP 1.
Revised: which yields a training and validation set composed of 14699 and 2013 sequences from SCOP 1.

Original: [Table]
Revised: [Table]

Original: and a test set of 2533 sequences from SCOP 2.06.
Revised: and a test set of 2533 sequences from SCOP 2.06.

Original: The resulting protein sequences belong to 1195 different folds, thus the problem is formulated as a multi-classiﬁcation task.
Revised: The resulting protein sequences belong to 1195 different folds, thus the problem is formulated as a multi-classiﬁcation task.

Original: The input sequence is represented as a 45-dimensional vector at each amino acid.
Revised: The input sequence is represented as a 45-dimensional vector at each amino acid.

Original: The vector consists of a 20-dimensional one-hot encoding of the sequence, a 20-dimensional position-speciﬁc scoring matrix (PSSM) representing the proﬁle of amino acids, a 3-class secondary structure represented by a one-hot vector and a 2-class solvent accessibility.
Revised: The vector consists of a 20-dimensional one-hot encoding of the sequence, a 20-dimensional position-speciﬁc scoring matrix (PSSM) representing the proﬁle of amino acids, a 3-class secondary structure represented by a one-hot vector and a 2-class solvent accessibility.

Original: The lengths of the sequences are varying from tens to thousands.
Revised: The lengths of the sequences are varying from tens to thousands.

Original: Models setting and hyperparameters.
Revised: Models setting and hyperparameters.

Original: We consider here the one-layer models followed by a global mean pooling for the baseline methods CKN (Chen et al., 2019a) and RKN (Chen et al., 2019b).
Revised: We consider here the one-layer models followed by a global mean pooling for the baseline methods CKN (Chen et al., 2019a) and RKN (Chen et al., 2019b).

Original: We build our model on top of the one-layer CKN model, where κ can be seen as a Gaussian kernel on the k-mers in sequences.
Revised: We build our model on top of the one-layer CKN model, where κ can be seen as a Gaussian kernel on the k-mers in sequences.

Original: The only difference between our model and CKN is thus the pooling operation, which is given by our embedding introduced in Section 3.
Revised: The only difference between our model and CKN is thus the pooling operation, which is given by our embedding introduced in Section 3.

Original: The bandwidth parameter of the Gaussian kernel κ on k-mers is ﬁxed to 0.6 for unsupervised models and 0.5 for supervised models, the same as used in CKN which were selected by the accuracy on the validation set.
Revised: The bandwidth parameter of the Gaussian kernel κ on k-mers is ﬁxed to 0.6 for unsupervised models and 0.5 for supervised models, the same as used in CKN which were selected by the accuracy on the validation set.

Original: The ﬁlter size k is ﬁxed to 10 and different numbers of anchor points in Nystr¨om for κ are considered in the experiments.
Revised: The ﬁlter size k is ﬁxed to 10 and different numbers of anchor points in Nystr¨om for κ are considered in the experiments.

Original: The other hyperparameters for our embedding are the entropic regularization parameter ε , the number of supports in a reference p , the number of references q , the number of iterations for Sinkhorn’s algorithm and the regularization parameter λ in the linear classiﬁer.
Revised: The other hyperparameters for our embedding are the entropic regularization parameter ε , the number of supports in a reference p , the number of references q , the number of iterations for Sinkhorn’s algorithm and the regularization parameter λ in the linear classiﬁer.

Original: The search grid for ε and λ is shown in Table 9 and they are selected by the accuracy on validation set.
Revised: The search grid for ε and λ is shown in Table 9 and they are selected by the accuracy on validation set.

Original: ε plays an important role in the performance and is observed to be stable for the same dataset.
Revised: ε plays an important role in the performance and is observed to be stable for the same dataset.

Original: For this dataset, it is selected to be 0.5 for all the unsupervised and supervised models.
Revised: For this dataset, it is selected to be 0.5 for all the unsupervised and supervised models.

Original: The effect of other hyperparameters will be discussed below.
Revised: The effect of other hyperparameters will be discussed below.

Original: For the baseline methods, the accuracies of PSI-BLAST and DeepSF are taken from Hou et al.
Revised: For the baseline methods, the accuracies of PSI-BLAST and DeepSF are taken from Hou et al.

Original: The hyperparameters for CKN and RKN can be found in Chen et al.
Revised: The hyperparameters for CKN and RKN can be found in Chen et al.

Original: (2019b).
Revised: (2019b).

Original: For Rep the
Revised: For Rep the

Original: Set (Skianis et al., 2020) and Set Transformer (Lee et al., 2019), we use the public implementations by the authors.
Revised: Set (Skianis et al., 2020) and Set Transformer (Lee et al., 2019), we use the public implementations by the authors.

Original: These two models are used on the top of a convolutional layer of the same ﬁlter size as CKN to extract k -mer features.
Revised: These two models are used on the top of a convolutional layer of the same ﬁlter size as CKN to extract k -mer features.

Original: As the exact version of Rep the Set does not provide any implementation for back-propagation to a bottom layer of it, we consider the approximate version of
Revised: As the exact version of Rep the Set does not provide any implementation for back-propagation to a bottom layer of it, we consider the approximate version of

Original: [Table] [Table]
Revised: [Table] [Table]

Original: Rep the Set only, which also scales better to our dataset.
Revised: Rep the Set only, which also scales better to our dataset.

Original: The default architecture of Set Transformer did not perform well due to overﬁtting.
Revised: The default architecture of Set Transformer did not perform well due to overﬁtting.

Original: We therefore used a shallower architecture with one ISAB, one PMA and one linear layer, similar to the one-layer architectures of CKN and our model.
Revised: We therefore used a shallower architecture with one ISAB, one PMA and one linear layer, similar to the one-layer architectures of CKN and our model.

Original: We tuned their model hyperparameters, weight decay and learning rate.
Revised: We tuned their model hyperparameters, weight decay and learning rate.

Original: The search grids for these hyperparameters are shown in Table
Revised: The search grids for these hyperparameters are shown in Table

Original: [Table]
Revised: [Table]

Original: Unsupervised embedding.
Revised: Unsupervised embedding.

Original: The kernel embedding ϕ , which is inﬁnite dimensional for the Gaussian kernel
Revised: The kernel embedding ϕ , which is inﬁnite dimensional for the Gaussian kernel

Original: [Table]
Revised: [Table]

Original: is approximated with the Nystr ¨ om method using K-means on 300000 k-mers extracted from the same training set as in Chen et al.
Revised: is approximated with the Nystr ¨ om method using K-means on 300000 k-mers extracted from the same training set as in Chen et al.

Original: (2019b).
Revised: (2019b).

Original: The reference measures are learned by using either K-means or Wasserstein to update centroids in 2-Wasserstein K-means on 3000 subsampled sequences for RAM-saving reason.
Revised: The reference measures are learned by using either K-means or Wasserstein to update centroids in 2-Wasserstein K-means on 3000 subsampled sequences for RAM-saving reason.

Original: We evaluate our model on top of features extracted from CKNs of different dimensions, representing the number of anchor points used to approximate κ .
Revised: We evaluate our model on top of features extracted from CKNs of different dimensions, representing the number of anchor points used to approximate κ .

Original: The number of iterations for Sinkhorn is ﬁxed to 100 to ensure the convergence.
Revised: The number of iterations for Sinkhorn is ﬁxed to 100 to ensure the convergence.

Original: The results for different combinations of q and p are provided in Table 11.
Revised: The results for different combinations of q and p are provided in Table 11.

Original: Increasing the number of supports p can improve the performance and then saturate it when p is too large.
Revised: Increasing the number of supports p can improve the performance and then saturate it when p is too large.

Original: On the other hand, increasing the number of references while keeping the embedding dimension ( i.e. p × q ) constant is not signiﬁcantly helpful in this unsupervised setting.
Revised: On the other hand, increasing the number of references while keeping the embedding dimension ( i.e. p × q ) constant is not signiﬁcantly helpful in this unsupervised setting.

Original: We also notice that Wasserstein Barycenter for learning the references does not outperform K-means, while the latter is faster in terms of computation.
Revised: We also notice that Wasserstein Barycenter for learning the references does not outperform K-means, while the latter is faster in terms of computation.

Original: Supervised embedding.
Revised: Supervised embedding.

Original: Our supervised embedding is initialized with the unsupervised method and then trained in an alternating fashion which was also used for CKN: we use an Adam algorithm to update anchor points in Nystr ¨ om and reference measures z , and the L-BFGS algorithm to optimize
Revised: Our supervised embedding is initialized with the unsupervised method and then trained in an alternating fashion which was also used for CKN: we use an Adam algorithm to update anchor points in Nystr ¨ om and reference measures z , and the L-BFGS algorithm to optimize

Original: We show the results for different combinations q number
Revised: We show the results for different combinations q number

Original: The reference measures z are learned with either
Revised: The reference measures z are learned with either

Original: Wasserstein
Revised: Wasserstein

Original: [Table]
Revised: [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: the classiﬁer.
Revised: the classiﬁer.

Original: The learning rate for Adam is initialized with 0.01 and halved as long as there is no decrease of the validation loss for 5 successive epochs.
Revised: The learning rate for Adam is initialized with 0.01 and halved as long as there is no decrease of the validation loss for 5 successive epochs.

Original: In practice, we notice that using a small number of Sinkhorn iterations can achieve similar performance to a large number of iteration, while being much faster to compute.
Revised: In practice, we notice that using a small number of Sinkhorn iterations can achieve similar performance to a large number of iteration, while being much faster to compute.

Original: We thus ﬁx it to 10 throughout the experiments.
Revised: We thus ﬁx it to 10 throughout the experiments.

Original: The accuracy results are obtained by averaging on 10 runs with different seeds following the setting in Chen et al.
Revised: The accuracy results are obtained by averaging on 10 runs with different seeds following the setting in Chen et al.

Original: (2019b).
Revised: (2019b).

Original: The results are shown in Table 13 with error bars.
Revised: The results are shown in Table 13 with error bars.

Original: The effect of the number of supports q is similar to the unsupervised case, while increasing the number of references can indeed improve performance.
Revised: The effect of the number of supports q is similar to the unsupervised case, while increasing the number of references can indeed improve performance.

Original: C. D ETECTION OF CHROMATIN PROFILES
Revised: C. D ETECTION OF CHROMATIN PROFILES

Original: Dataset description.
Revised: Dataset description.

Original: Predicting the functional effects of noncoding variants from only genomic sequences is a central task in human genetics.
Revised: Predicting the functional effects of noncoding variants from only genomic sequences is a central task in human genetics.

Original: A fundamental step for this task is to simultaneously predict large-scale chromatin features from DNA sequences (Zhou & Troyanskaya, 2015).
Revised: A fundamental step for this task is to simultaneously predict large-scale chromatin features from DNA sequences (Zhou & Troyanskaya, 2015).

Original: We consider here the DeepSEA dataset, which consists in simultaneously predicting 919 chromatin proﬁles including 690 transcription factor (TF) binding proﬁles for 160 different TFs, 

Revised: We consider here the DeepSEA dataset, which consists in simultaneously predicting 919 chromatin proﬁles including 690 transcription factor (TF) binding proﬁles for 160 different TFs,

Original: DNase I sensitivity proﬁles and 104 histone-mark proﬁles.
Revised: DNase I sensitivity proﬁles and 104 histone-mark proﬁles.

Original: In total, there are 4.4 million, 8000 andsamples for training, validation and test.
Revised: In total, there are 4.4 million, 8000 andsamples for training, validation and test.

Original: Each sample consists of a 1000-bp DNA sequence from the human GRCh37 reference.
Revised: Each sample consists of a 1000-bp DNA sequence from the human GRCh37 reference.

Original: Each sequence is represented as a 1000 × 4 binary matrix using onehot encoding on DNA characters.
Revised: Each sequence is represented as a 1000 × 4 binary matrix using onehot encoding on DNA characters.

Original: The dataset is available at http://deepsea.princeton.
Revised: The dataset is available at http://deepsea.princeton.

Original: edu/media/code/deepsea_train_bundle.v0.9.tar.gz .
Revised: edu/media/code/deepsea_train_bundle.v0.9.tar.gz .

Original: Note that the labels for each proﬁle are very imbalanced in this task with few positive samples.
Revised: Note that the labels for each proﬁle are very imbalanced in this task with few positive samples.

Original: For this reason, learning unsu-
Revised: For this reason, learning unsu-

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table]
Revised: [Table]

Original: pervised models could be intractable as they may require an extremely large number of parameters if junk or redundant sequences cannot be ﬁltered out
Revised: pervised models could be intractable as they may require an extremely large number of parameters if junk or redundant sequences cannot be ﬁltered out

Original: [Table]
Revised: [Table]

Original: Model architecture and hyperparameters.
Revised: Model architecture and hyperparameters.

Original: For the above reason and fair comparison, we use here our supervised embedding as a module in Deep NNs.
Revised: For the above reason and fair comparison, we use here our supervised embedding as a module in Deep NNs.

Original: The architecture of our model is shown in Table 15.
Revised: The architecture of our model is shown in Table 15.

Original: We use an Adam optimizer with initial learning rate equal to 0.01 and halved at epoch 1, 4, 8 for 15 epochs in total.
Revised: We use an Adam optimizer with initial learning rate equal to 0.01 and halved at epoch 1, 4, 8 for 15 epochs in total.

Original: The number of iterations for Sinkhorn is ﬁxed to 30.
Revised: The number of iterations for Sinkhorn is ﬁxed to 30.

Original: The whole training process takes about 30 hours on a single GTX2080TI GPU.
Revised: The whole training process takes about 30 hours on a single GTX2080TI GPU.

Original: The dropout rate is selected to be 0.
Revised: The dropout rate is selected to be 0.

Original: from the grid [0 . 1; 0 . 2; 0 . 3; 0 . 4; 0 . 5] and the weight decay is 1e-06, the same as Zhou & Troyanskaya (2015).
Revised: from the grid [0 . 1; 0 . 2; 0 . 3; 0 . 4; 0 . 5] and the weight decay is 1e-06, the same as Zhou & Troyanskaya (2015).

Original: The σ pos for position encoding is selected to be 0.1, by the validation accuracy on the grid [0 . 05; 0 . 1; 0 . 2; 0 . 3; 0 . 4; 0 . 5] .
Revised: The σ pos for position encoding is selected to be 0.1, by the validation accuracy on the grid [0 . 05; 0 . 1; 0 . 2; 0 . 3; 0 . 4; 0 . 5] .

Original: The checkpoint with the best validation accuracy is used to evaluate on the test set.
Revised: The checkpoint with the best validation accuracy is used to evaluate on the test set.

Original: Area under ROC (auROC) and area under precision curve (auPRC), averaged overchromatin proﬁles, are used to measure the performance.
Revised: Area under ROC (auROC) and area under precision curve (auPRC), averaged overchromatin proﬁles, are used to measure the performance.

Original: The hidden size d is chosen to be either 1024 or 1536.
Revised: The hidden size d is chosen to be either 1024 or 1536.

Original: Results and importance of position encoding.
Revised: Results and importance of position encoding.

Original: We compare our model to the state-of-the-art
Revised: We compare our model to the state-of-the-art

Original: CNN model DeepSEA (Zhou & Troyanskaya, 2015) with 3 convolutional layers, whose best hyperparameters can be found in the corresponding paper.
Revised: CNN model DeepSEA (Zhou & Troyanskaya, 2015) with 3 convolutional layers, whose best hyperparameters can be found in the corresponding paper.

Original: Our model outperforms DeepSEA, while requiring fewer layers.
Revised: Our model outperforms DeepSEA, while requiring fewer layers.

Original: The positional information is known to be important in this task.
Revised: The positional information is known to be important in this task.

Original: To show the efﬁcacy of our position encoding, we compare it to the sinusoidal encoding used in the original transformer (Vaswani et al., 2017).
Revised: To show the efﬁcacy of our position encoding, we compare it to the sinusoidal encoding used in the original transformer (Vaswani et al., 2017).

Original: We observe that our encoding with properly tuned σ pos requires fewer layers, while being interpretable from a kernel point of view.
Revised: We observe that our encoding with properly tuned σ pos requires fewer layers, while being interpretable from a kernel point of view.

Original: We also ﬁnd that larger hidden size d performs better, as shown in Table 16.
Revised: We also ﬁnd that larger hidden size d performs better, as shown in Table 16.

Original: ROC and PR curves for all the chromatin proﬁles and stratiﬁed by transcription factors, DNase I-hypersensitive sites and histone-marks can also be found in Figure 2.
Revised: ROC and PR curves for all the chromatin proﬁles and stratiﬁed by transcription factors, DNase I-hypersensitive sites and histone-marks can also be found in Figure 2.

Original: C. SST- 
Revised: C. SST-

Original: Dataset description.
Revised: Dataset description.

Original: The data set contains 67,349 training samples and 872 validation samples and can be found at https://gluebenchmark.com/tasks .
Revised: The data set contains 67,349 training samples and 872 validation samples and can be found at https://gluebenchmark.com/tasks .

Original: The test set contains 1,821 samples for which the predictions need to be submitted on the GLUE leaderboard, with limited number of submissions.
Revised: The test set contains 1,821 samples for which the predictions need to be submitted on the GLUE leaderboard, with limited number of submissions.

Original: As a consequence, our training and validation set are extracted from the original training set ( 80% of the original training set is used for our training set and the remaining 20% is used for our validation set), and we report accuracies on the standard validation set, used as a test set.
Revised: As a consequence, our training and validation set are extracted from the original training set ( 80% of the original training set is used for our training set and the remaining 20% is used for our validation set), and we report accuracies on the standard validation set, used as a test set.

Original: The reviews are padded with zeros when their length is shorter than the chosen sequence length (we choose 30 and 66 , the latter being the maximum review length in the data set) and the BERT
Revised: The reviews are padded with zeros when their length is shorter than the chosen sequence length (we choose 30 and 66 , the latter being the maximum review length in the data set) and the BERT

Original: [Table] [Table]
Revised: [Table] [Table]

Original: implementation requires to add special tokens [CLS] and [SEP] at the beginning and the end of each review.
Revised: implementation requires to add special tokens [CLS] and [SEP] at the beginning and the end of each review.

Original: Model architecture and hyperparameters.
Revised: Model architecture and hyperparameters.

Original: In most transformers such as BERT, the embedding associated to the token [CLS] is used for classiﬁcation and can be seen in some sense as an embedding of the review adapted to the task.
Revised: In most transformers such as BERT, the embedding associated to the token [CLS] is used for classiﬁcation and can be seen in some sense as an embedding of the review adapted to the task.

Original: The features we used are the word features provided by the BERT base-uncased version, available at https://huggingface.co/transformers/ pretrained_models.html .
Revised: The features we used are the word features provided by the BERT base-uncased version, available at https://huggingface.co/transformers/ pretrained_models.html .

Original: For this version, the dimension of the word features is 768 .
Revised: For this version, the dimension of the word features is 768 .

Original: Our model is one layer of our embedding, with ϕ the Gaussian kernel mapping with varying number of Nystr ¨ om ﬁlters in the supervised setting, and the Linear kernel in the unsupervised setting.
Revised: Our model is one layer of our embedding, with ϕ the Gaussian kernel mapping with varying number of Nystr ¨ om ﬁlters in the supervised setting, and the Linear kernel in the unsupervised setting.

Original: We do not add positonnal encoding as it is already integrated in BERT features.
Revised: We do not add positonnal encoding as it is already integrated in BERT features.

Original: In the unsupervised setting, the output features are used to train a large-scale linear classiﬁer, a Pytorch linear layer.
Revised: In the unsupervised setting, the output features are used to train a large-scale linear classiﬁer, a Pytorch linear layer.

Original: We choose the best hyper-parameters based on the accuracy of a validation set.
Revised: We choose the best hyper-parameters based on the accuracy of a validation set.

Original: In the supervised case, the parameters of the previous model, w and z , are optimized end-to-end.
Revised: In the supervised case, the parameters of the previous model, w and z , are optimized end-to-end.

Original: In this case, we tune the learning rate.
Revised: In this case, we tune the learning rate.

Original: In both case, we tune the entropic regularization parameter of optimal transport and the bandwidth of the Gaussian kernel.
Revised: In both case, we tune the entropic regularization parameter of optimal transport and the bandwidth of the Gaussian kernel.

Original: The parameters in the search grid are summed up in Table 18.
Revised: The parameters in the search grid are summed up in Table 18.

Original: The best entropic regularization and Gaussian kernel bandwidth are typically and respectively 3 .and 0 .
Revised: The best entropic regularization and Gaussian kernel bandwidth are typically and respectively 3 .and 0 .

Original: 5 for this data set.
Revised: 5 for this data set.

Original: The supervised training process takes between half an hour for smaller models (typically 128 ﬁlters in w and 3 supports in z ) and a few hours for larger models ( 256 ﬁlters and 100 supports) on a single GTX2080TI GPU.
Revised: The supervised training process takes between half an hour for smaller models (typically 128 ﬁlters in w and 3 supports in z ) and a few hours for larger models ( 256 ﬁlters and 100 supports) on a single GTX2080TI GPU.

Original: The hyper-parameters of the baselines were similarly tuned, see 19.
Revised: The hyper-parameters of the baselines were similarly tuned, see 19.

Original: Mean Pooling and [CLS] embedding did not require any tuning except for the regularization λ of the classiﬁer, which followed the same grid as in Table 18.
Revised: Mean Pooling and [CLS] embedding did not require any tuning except for the regularization λ of the classiﬁer, which followed the same grid as in Table 18.

Original: Results and discussion.
Revised: Results and discussion.

Original: As explained in Section 5, our unsupervised embedding improves the
Revised: As explained in Section 5, our unsupervised embedding improves the

Original: BERT pre-trained features while still using a simple linear model as shown in Table 17, and its supervised counterpart enables to get even closer to the state-of-the art (for the BERT base-uncased model) accuracy, which is usually obtained after ﬁne-tuning of the BERT model on the whole data set.
Revised: BERT pre-trained features while still using a simple linear model as shown in Table 17, and its supervised counterpart enables to get even closer to the state-of-the art (for the BERT base-uncased model) accuracy, which is usually obtained after ﬁne-tuning of the BERT model on the whole data set.

Original: This can be seen in Tables 20; 21.
Revised: This can be seen in Tables 20; 21.

Original: We also add a baseline consisting of one layer of classical self-attention, which did not do well hence was not reported in the main text.
Revised: We also add a baseline consisting of one layer of classical self-attention, which did not do well hence was not reported in the main text.

Original: [Table]
Revised: [Table]

Original: [Table] [Table]
Revised: [Table] [Table]