A Trainable Optimal Transport Embedding for Feature Aggregation and its Relationship to Attention
Authors: Grégoire Mialon
Dexiong Chen
Alexandre d'Aspremont
Julien Mairal
Venue: ICLR 2021 Poster (2021)
Keywords: bioinformatics
optimal transport
kernel methods
attention
transformers

Abstract: 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. To address this challenging task, we introduce a parametrized representation of fixed 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. 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. 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. 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 profiles tasks, and, as a proof of concept, we show promising results for processing natural language sequences. 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.

Published as a conference paper at ICLR 2021
A TRAINABLE OPTIMAL TRANSPORT EMBEDDING FOR FEATURE AGGREGATION AND ITS RELATIONSHIP TO ATTENTION
Grégoire Mialon" 1 Dexiong Chen" Alexandre d'Aspremont Julien Mairal
i Inria, 3 CNRS-ENS
(gregoi re. mia lon, dexi ong chen, julien....... ma mairalllaa wtw.mmaaaaaiiiaaaa....
jeinria.fr fr, aspremon
ABSTRACT
We address the problem of learning on sets 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. To address this challeng- ing task, we introduce parametrized representation of fixed size, which embeds and then aggregates elements from given input set according the optimal trans- port plan between the set and trainable reference. 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. Our ag- gregation technique admits two useful interpretations: it may be seen as a mech- scal-
anism related to attention layers in neural networks, it may be seen able surrogate of a classical optimal transport-basee kernel. We experimentally demonstrate the effectiveness of our approach on biological sequences, achieving state-of-the-an results for protein fold recognition and detection of chromatin pro- files tasks, and, proof of concept, natural language sequences. We provide an open-source implementation of our embedding that can be used alone or as a module in larger learning models https: //tgt//gtttb//////// com c cllyyng/OTK
show promising results for processing
INTRODUCTION
Many scientific fields such as bioinformatics or natural language processing (NLP) require process- ing sets of features with positional information (biological sequences, sentences represented by set of local features). These objects are delicate to manipulate due to varying lengths and potentially long-range dependencies between their elements. For many tasks, the difficulty is even greater since the sets can be arbitrarily large, or only provided with few labels, or both.
Deep learning architectures specifically designed for sets have recently been proposed (Lee al. 2019 Skianns et al.] 2020). Our experiments show that these architectures perform well for NLP tasks, but achieve mixed performance for long biological sequences of varying size with few la- beled data. Some of these models use attention (Bahdanau et al.] 2015),, classical mechanism for aggregating features. Its typical implementation is the transformer (Vaswani et al., 2017). which has shown to achieve state-of-the-ar results for many sequence modeling tasks, e.g. in NLP (De- vlin et al. 2019) or in bioinformaties (Rives al., 2019), when trained with self supervision on large-scale data. Beyond sequence modeling, we are interested in this paper in finding good rep- resentation for sets of features of potentially diverse sizes, with or without positional information, when the amount of training data may be scarce. To this end, we introduce trainable embedding, which can operate directly on the feature set or be combined with existing deep approaches.
More precisely, our embedding marries ideas from optimal transport (OT) theory (Peyré & Cuturi 2019) and kernel methods (Schölkopf Smola 200I). We call this embedding OTKE (Optimal
"Equal contribution.
1 Univ. Grenoble Alpes, Inria, CNRS, Grenoble INP, LJK, 38000 Grenoble, France. D.I., UMR 8548, Fcole normale supérieure, Paris, France.

---

Published as a conference paper at ICLR 2021
Transport Kernel Embedding). Concretely, we embed feature vectors of given set to reproducing kernel Hilbert space (RKHS) and then perform a weighted pooling operation, with weights given by the transport plan between the set and trainable reference. To gain scalability, we then obtain finite-dimensional embedding by using kernel approximation techniques (Williams & Seeger] 2001). The motivation for using kernels is to provide non-linear transformation of the input features be- fore pooling, whereas optimal transport allows to align the features on a trainable reference with fast algorithms (Cuturi 2013). Such combination provides us with theoretically grounded, fixed-size embedding that can be learned either without any label, or with supervision. Our embedding can indeed become adaptive the problem at hand, by optimizing the reference with respect to a given task. It can operate on large sets with varying size, model long-range dependencies when positional information is present, and scales gracefully to large datasets. We demonstrate its effectiveness on biological sequence classification tasks, including protein fold recognition and detection of chro- matin profiles where we achieve statc-of-the-ar results. We also show promising results in natural language processing tasks, where method outperforms strong basclines.
Contributions. In summary, our contribution is three-fold. We propose a new method to embed sets of features of varying sizes fixed size representations that are well adapted downstream machine learning tasks, and whose parameters can be learned in either unsupervised or supervised fashion. We demonstrate the scalability and effectiveness of our approach on biological and natural language sequences. We provide an open-source implementation of our embedding that can be used alone or as a module in larger learning models.
RELATED WORK
Kernel methods for sets and OT-based kernels. The kernel associated with our embedding be- longs to the family of match kernels (Lyu 2004 Tolias et al. 2013), which compare all pairs of fea- tures between two sets via similarity function. Another line of research builds kernels by matching features through the Wasserstein distance. A few of them are shown to be positive definite (Gardner et al. 2018) and/or fast to compute (Rabin et al., 2011 Kolouri et al. 2016). Except for few hyper- parameters, these kernels yet cannot be trained end-to-end, as opposed to our embedding that relies on trainable reference. Efficient and trainable kernel embeddings for biological sequences have also been proposed by CCenn et al. (2019a b). Our work can be seen as an extension of these car- lier approaches by using optimal transport rather than mean pooling for aggregating local features, which performs significantly better for long sequences in practice.
Deep learning for sets. Deep Sets (Zaheer et al. 2017) feed each element of an input set into feed-forward neural network. The outputs are aggregated following a simple pooling operation before further processing. Lee et al. (2019) propose a Transformer inspired encoder-decods archi- tecture for sets which also uses latent variables. [Skianis et al. (2020) compute some comparison costs between an input set and reference sets. These costs are then used as features in a subsequent neural network. The reference sets are learned end-to-end. Unlike our approach, such models do not allow unsupervised learning. We will use the last two approaches as baselines in our experiments.
Interpretations of attention. Using the transport plan as an ad-hoe attention score was proposed by Chen et al.) (2019c) in the context of network embedding to align data modalities. Our paper goes beyond and uses the transport plan as a principle for pooling set in a model, with trainable parameters. Tsai et al. (2019) provide view of Transformer's attention via kernel methods, yet in very different fashion where attention is cast as kernel smoothing and not as a kernel embedding.
PROPOSED EMBEDDING
3.1 PRELIMINARIES
We handle sets of features in Rdd and consider sets x living in
X x= {x x= (x)..., { such that X1,..., Xn E IRd for some n 2 1}.
Elements of X' are typically vector representations of local data structures, such k-mers for quences, patches for natural images, or words for sentences. The size of x denoted by n may vary, which is not an issue since the methods we introduce may take sequence of any size as input, while providing a fixed-size embedding. We now revisit important results on optimal transport and kernel methods, which will be useful to describe our embedding and its computation algorithms.
2

---

Published as a conference paper at ICLR 2021
Optimal transport. Our pooling mechanism will be based the transport plan between and x' seen as weighted point clouds or discrete measures, which is a by-product of the optimal transport problem (Villami 2008; Peyré & Cuturi, 2019). OT has indeed been widely used in alignment problems (Grave et al, 2019). TTrouuhouu the paper, we will refer to the Kantorovich relaxation of OT with entropic regularization, detailed for example in (Peyré & Cuturi 2019). Let in A (probability simplex) and b in A"' be the weights of the discrete measures D; a,0x, and Σ3 by0x;
with respective locations x and x', where 8x is the Dirac position x. Let C in Rax''' be matrix representing the pairwise costs for aligning the elements of x and x'. The entropic regularized Kantorovich relaxation of OT from x'
Ec.Po-eH(P)) min (aU(a.
(1)
where HPP) = - P (log(()))))) 1) is the entropic regularization with parameter E, which controls the sparsity of P. and the space of admissible couplings between a and b:
b) = (P E Raxn', Pl, a and P 1 p-
The problem is typically solved by using a matrix scaling procedure known as Sinkhorn's 's algo- rithm (Sinkhorn & Knopp 1967; Cuturi 2013). In practice, and b are uniform measures since we consider the mass be evenly distributed between the points. P is called the transport plan, which carries the information on how to distribute the mass of x in x' with minimal cost. Our method uses optimal transport to align features of a given set to a learned reference.
Kernel methods. Kernel methods (Schölkopf & Smola, 2001) map data living in a space X' to a reproducing kernel Hilbert space H. associated to positive definite kernel K through mapping function @ : X - H. such that K(x,'' x') = (p(x), (x))) Even though s(x) may be infinite- dimensional, classical kernel approximation techniques (Williams & Seeger 2001) provide finite- dimensional embeddings (x) in & such that ((x, x') z (1((x, (x))). Our embedding for sets relies in part kernel method principles and on such mite--dimensional approximation.
3.2 OPTIMAL TRANSPORT EMBEDDING AND ASSOCIATED KERNEL.
We now present the OTKE, an embedding and pooling layer which aggregates variable-size set or sequence of features into fixed-size embedding. We start with an anfnite-dimensional variant living in a RKHS, before introducing the minee-dimeniionl embedding that we use in practice.
Infinite-dimensionall embedding in RKHS. Given set x and a (learned) reference z in X' with p elements, we consider an embedding + (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 (x) in HP, which is illustrated in Figure |
Before introducing more formal details, we note that our embedding relies on two main ideas:
Global similarity-based pooling using references. Learning on large sets with long-range interac- tions may benefit from pooling reduce the number of feature vectors. 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. To this end, each element in the reference set corresponds to pooling cell, where the elements of the input set are aggregated through weighted sum. The weights simply reflect the similarity between the vectors of the input set and the current vector in the reference. Importantly, using reference set enables to reduce the size of the "attention matrix" from quadratic to linear in the length of the input sequence.
Computing similarity weights via optimal transport. A computationally efficient similarity score between two elements is their dot- dot-prrduct (Vaswani et al. 2017). In this paper, we rather consider that elements of the input should be pooled together if they align well with the same part of the reference. Alignment scores can efficiently be obtained by computing the transport plan between the input and the reference sets: Sinkhorn's 's algorithm indeed enjoys fast solvers (Cuturi 2013).
We are now in shape to give a formal definition.
Definition 3.1 (The optimal transport kernel embedding). Let x = (x)..... .,, x,) in X' be an input set of feature vectors and (Z1,..., 2p)) in X' be reference set with p elements. Let be
3

---

Published as a conference paper at ICLR 2021
p(x1)
p(xx)]
((x..)
X
X2
P1
P2p
Pn
Pmt
P2p
P
Zp
2
4x((x);
<2((x),
Figure I: The input point cloud x is transported onto the reference z = (Z)...... zp) (left), yielding the optimal transport plan P, Pxx(x,2) used to aggregate the embedded features and form $2(x) (right).
positive definite kernel, e.g. Gaussian kernel, with RKHS H. and y Rd > H. its associated kernel embedding. Let R. be the n x p matrix which carries the comparisons x(x,, Zz). before alignment.
Then, the transport plan between and denoted by the n x p matrix P(x, z). is defined as the unique solution
(I)
when choosing the cost = -k, and our embedding is defined as
(x) = vp Σ
P(xx, ()xx(x))((xs),
P(x,
Σ P(x, 2) ()pp()))))
vp P(x,2) (x()),
where p(xx) == (x(x)), (x(x)))...,,$(xx ((xx))
Interestingly, it easy to show that our embedding (x) is associated to the positive definite kernel
Ka(x,x)) = Σ
x4) = (x), (x())),
(2)
6,==1
with (x,x''))= X 2(x,z)P(x',z)) This is weighted match kernel, with weights given by optimal transport in H. The notion pooling the RKHS H of arises naturally if S n. 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 (xx, 2). This process is illustrated in Figure [ We acknowledge here the concurrent work by Kolouri al. (2021), 1), where similar embedding is used for graph representation. We now expose the benefits of this kernel formulation, and its relation to classical non-positive definite kernel.
Kernel interpretation. Thanks to the gluing lemma (see, e.g., Peyré & Cuturi 2019), (x valid transport plan and, empirically, rough approximation P(x,x') K2 can therefore be seen as surrogate of well-known kernel (Rubner et al. 2000), defined as
Kor Kx (x, x)x,
Σ P(x, () 2,1=1
When the entropic regularization equal to O, Kor is equivalent the 2- 2-sssserstein distance (2(x,x) x') with ground metric dy induced by kernel к. Kor is generally not positive definite (see Peyré & 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 K2). Now, we show the relationship between this kernel and our kernel K2, which is proved Appendix B.I
Lemma 3.1 (Relation between P(x, x') and Pa Px(x,x' when & = 0). For any x, x' and z in X with lengths n, n' and p. by denoting W (x, x') = (Pz(x, x'), de (x) // we have W (x, x') < 2 min [2 z), (x'x)
(2(xx x')
This lemma shows that the distance W resulting from K2 is related to the Wasserstein distance W2; yet, this relation should not be interpreted as an approximation error as our goal is not to approx- imate W2, but rather to derive trainable embedding (x(x) with good computational properties. Lemma 3.1 roots our features and to some extent self-attention in a rich optimal transport literature. In fact, W is equivalent to distance introduced by Wang al. 2013), whose properties fur- ther studied by MMoosmüller Cloninger (2020). A major difference that W crucially relies on Sinkhorn's 'ss algorithm so that the references can be learned end-to-end, explained below.
4

---

Published as a conference paper at ICLR 2021
3.3 FROM INFINITE- DIMANSENAIM KERNEL EMBEDDING TO FINITE DIMENSION
In some cases, (x() is already minte-dimessional which allows to compute the embedding (x(x) explicitly. This is particularly useful when dealing with large-scale data, as enables us to use our method for supervised learning tasks without computing the Gram matrix, which grows quadrati- cally in size with the number of samples. When o is infinite or high-dimensional. it is nevertheless possible to an approximation based on the Nyström method (Williams & Seeger 2001). which provides an embedding v Rd > R2 such that
(1+(),), ((x)))=+ R k(x;, x,).
Concretely, the Nyström method consists in projecting points from the RKHS H onto a linear sub- :p( (((w))): The corre- space F, which parametrized by & anchor points F = -1)1//2 xw),,,), where k(w, w) is the sponding embedding admits an explicit form (xx)) K(w, w)
k X k Gram matrix of к computed on the set w [Wi,..., W&J of anchor points and K( (w, x;) is in RR Then, there are several ways to learn the anchor points: (a) they can be chosen as random points from data; (b) they can be defined as centroids obtained by K-means, see [Zhang et al. (2008); (c) they can be learned by back- propagation for a supervised task, see MMiral (2016).
Using such an approximation within our framework can be simply achieved by (i) replacing k by linear kernel and (ii) replacing each element x; by its embedding (x,) in Rh and considering reference set with elements in RR. By abuse of notation, we still use z for the new parametrization. The embedding, which we use in practice in all our experiments, becomes simply
$ xx)) = vp Σ (()()), (x))
(()(x), 2) (xx))((x))
= VD x (()()), 2) (x)) E R
(5)
where p is the number of elements in z. Next, we discuss how to learn the reference set z.
3.4 UNSUPERVISED AND SUPERVISED LEARNING OF PARAMETER z
Unsupervised learning. In the fashion of the Nyström approximation, the p elements of z can be defined as the centroids obtained by K-means applied to all features from training sets in X. A corollary of Lemma 3.1 suggests another algorithm: a bound on the deviation term between W2 and W for m samples
x") is indeed
m W2 [ z)-
E2
W2(x', x²) WZ x))
771
+1=1
The right-hand term corresponds to the objective of the Wasserstein barycenter problem (Cuturi & Doucet 20133), which yields the mean of a set of empirical measures (here the x's) under the OT K-means
metric. The Wasserstein barycenter is therefore an attractive candidate for choosing can be seen as a particular case of Wasserstein barycenter when m = (Cuturi & Doucet 2013 Ho let al. 2017) and faster to compute. 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. The anchor points w and the references may be then computed using similar algorithms; however, their mathematical interpretation differs as exposed above. The task of representing features (learning w in I for specific x) decoupled from the task of aggregating (learning the reference in R²*).
Supervised learning. As mentioned in Section 3.1 (x) z) is computed using Sinkhorn's 's al- gorithm, 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 Φx More impor- tant, all Sinkhorn's operations are differentiable, which enables to optimized with stochastic gradient descent through back-propagation (Genevay et al.) 2018), e.g., for minimizing a classifica- tion or regression loss function when labels are available. In practice, a small number of Sinkhorn iterations (e.g., 10) is sufficient to compute P P( ((i(x), z). Since the anchors w in the embedding layer below can also be learned mndto--nn (Mairal 2016). our embedding can thus be used as a module injected into any model, e.g, deep network, demonstrated in our experiments.
5

---

Published as a conference paper at ICLR 2021
3.5 EXTENSIONS
Integrating positional information into the embedding. 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. To this end, we borrow the idea of convolutional kernel networks, CKN (Mairal 2016 Mairal al., 2014). and penalize the similarities exponentially with the positional distance between a pair of elements in the input and reference sequences. More precisely, we multiply (x)( z) element-wise by a distance matrix S defined as:
S = e (1/=11/))
and replace it in the embedding. With similarity weights based both on content and position, the ker- nel associated to our embedding can be viewed as generalization of the CKNs (whose similarity weights are based on position only), with feature aligmment based on optimal transport. When deal- ing with multi dimensioonl objects such as images, we just replace the index scalar with an index vector of the same spatial dimension as the object, representing the positions of each dimension.
Using multiple references. naive reconstruction using different references z ,..., yield better approximation of the transport plan. In this case, the embedding of x becomes
22 in X may
P xe (x) = 1/v ( (x), (x)
with the number of references (the factor 1/v9 comes from the mean). Using Eq. 4 we can obtain a bound similar (6) for data set of m samples (x* ...., x"") and q references (see Appendix B.2 for details). To choose multiple references, we tried a K-means algorithm with 2- Wasserstein dis- tance for assigning clusters, and we updated the centroids as in the single-reference case. Using
multiple references appears to be useful when optimizing z with supervision (see Appendix C
4 RELATION BETWEEN OUR EMBEDDING AND SELF-ATTENTION
Our embedding and a single layer of transformer encoder, recalled in Appendix A share the same type of inductive bias, iee, aggregating features relying on similarity weights. We now clarify their relationship. Our embedding is arguably simpler (see respectively size of attention and number of parameters in Table T), and may compete in some settings with the transformer self-attention as illustrated in Section
Shared reference versus self-attention. There is a correspondence between the values, attention matrix in the transformer and sp. P in Definition B.I. yet also no- ticeable differences. On one hand, & aligns given sequence x with respect to a reference z, learned with or without super- vision, and shared across the data set. Our weights are computed using optimal trans- port. On the other hand, transformer coder performs self-alignment: for a given xi, features are aggregated depending on a similarity score between X; and the el- ements of only. The similarity score is a matrix product between queries Q and keys K matrices, learned with supervi- sion and shared across the data set. In this regard, our work complements re- cent line of research questioning the dot- product, learned self-attention (Raganato et al.] 2020; |Weiqiu et al. 2020). Self- attention-like weights can also be obtained
Table 1: Relationship between o, and transformer self- attention. k: function describing how the transformer integrates positional information; n. sequence length; q: number of references attention heads; d: di- mension of the embeddings; p. number of supports in Typically, p < d. In recent transformer architec- tures, positional encoding requires learning additional parameters (~ gd3).
Self-Attention
Attention score Size of score Alignment w.r.t: Learned + Shared Nonlinear mapping
W = -WT
P (mp)
O(n?)
x itself
z z
W and Q
Feed -oowwfor
or +b
Position encoding Nb, parameters Supervision
k(te, qd Needed
qpd Not needed
with embedding by computing P(x, 2, P TT)TT for each reference i. In that sense, our work related to recent research on efficient self-attention (Wang et al.) 2020) [hhromanski et al. 2020). where low-rank approximation of the self-attention matrix computed.

---

Published as a conference paper at ICLR 2021
Position smoothing and relative positional encoding. Transformers can add absolute posi- tional encoding to the imput features (Vaswani et al. 2017). Yet, relative positional encoding (Dai et al. 2019) is 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. The link between CKNs and our kernel, provided by this positional encoding, stands in line with recent works casting attention and convolution into a unified frame- work (Andreoli] 2019). In particular, Cordonnier et al. (2020) show that attention learns convolution in the setting of image classification: the kernel pattern is learned at the same time as the filters.
Multiple references and attention heads. 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. Successful architectures have a few heads in parallel. The outputs of the heads are then aggregated to output a final embedding. A layer of our embedding with non-linean kernel K can be seen as such a block, with the references playing the role of the heads. As some recent works question the role of attention heads (Voita et al. [2019] Michel al. 2 exploring the content of our learned references z may provide another perspec- tive on this question. More generally, visualization and interpretation of the learned references could be of interest for biological sequences.
5 EXPERIMENTS
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. We evaluate our embedding alone in unsupervised or supervised settings, or within model in the supervised setting. We also consider NLP tasks involving shorter sequences and relatively more labels.
5.1 DATASETS, EXPERIMENTAL SETUP AND BASELINES
In unsupervised settings, we train a linear classifier with the cross entropy loss between true la- bels and predictions on top of the features provided by our embedding (where the references and Nyström anchors w have been learned without supervision), or an unsupervised baseline. In supervised settings, the same model is initialized with our unsupervised method and then trained end-to-end (including and w) by minimizing the same loss. We use an alternating optimization strategy to update the parameters for both SCOP and SST datasets, as used by Chen et al. (2019a b). 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. The hyper-parameter tuned include number of supports and references p. entropic regularization in OT e, the band- width of Gaussian kernels and the regularization parameter of the linear classifier. The best values of e and the bandwidth were found stable across tasks, while the regularization parameter needed to be more carefully cross-validated. Additional results and implementation details can be found in Appendix C
Protein fold classification on SCOP 1.75. We follow the protocol described by Hou et al. (2019) for this important task in bioinformatics. The dataset contains 19, 245 sequences from 1, 195 dif- ferent classes of fold (hence less than 20 labels in average per class). The sequence lengths vary from tens to thousands. Each element of sequence is 45-dimensional vector. The objective is to classify the sequences to fold classes, which corresponds to a multiclass classification problem. The features fed the linear classifier are the output of embedding with yo the Gaussian kernel map- ping on k-mers (subsequences of length k) with fixed to be 10, which is known to perform well in this task (Chen al. 2019a). The number of anchor points for Nyström method fixed 1024 and 512 respectively for unsupervised and supervised setting. In the unsupervised setting, we compare our method to state-of-the-arr unsupervised method for this task: CKN (Chen et al. 2019a), which performs global mean pooling in contrast to the global adaptive pooling performed by our embed- ding. In the supervised setting, we compare the same model to the following supervised models: CKN, Recurrent Kernel Networks (RKN) (Chen et al. 2019b), 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. Rep the Set and Set Transformer are used on the top of a convolutional layer of the same filter size as CKN to extract k-mer features. Their model hyper- parameeer,, weight decay and learning rate are tuned in the same way for our models
7

---

Published as a conference paper at ICLR 2021
(see Appendix for details). The default architecture of Set Transformer did not perform well due to overfitting. We thus used a shallower architecture with one Induced Set Attention Block (ISAB), one Pooling by Multihead Attention (PMA) and one linear layer, similar to the one-layer architectures of CKN and our model. The results are shown in T
Table 2: Classification accuracy (top 1/5/10) on test set for SCOP 1.75 for different unsupervised and supervised baselines, averaged from 10 different runs (q references x p supports).
Method
Unsupervised
Supervised
DeepSF (Hou et al.
2019
Not available.
73.0/90.3/94.5
CKN Chen et al. 201 9a
RKN Chen ct al.
2019b
Not available.
Set Transtormer Lec cl al_
2019
Not available.
7922==4.6091.551444444..3333
Approximate Rep the Set (Skiams ct al. 2020
Not available.
$4.5=00.669.00004495555777
Ours (dot-product instead of OT)
7822H1..99911--0796.000 55.8:0.299533--.111.60-00 18.7--0.3965--00227773:-
$755-10.799.55--222999990
Ours (Unsup.: 1 100 Sup.: x 10)
Detection of chromatin profiles. Predicting the chromatin features such as transcription factor (TF) binding from raw genomic sequences has been studied extensively in recent years. CNNs with max pooling operations have been shown effective for this task. Here, we consider DeepSEA dataset (Zhou & Troyanskayal 2015) consisting in simultancously predicting 919 chromatin profiles, which can be formulated as multi- label classification task. DeepSEA contains 4, 935, 024 DNA se- quences of length 1000 and each of them is associated with 919 different labels (chromatin profiles). Each sequence is represented
1000 x 4 binary matrix through one-hot encoding and the objec- tive is to predict which profiles a given sequence possesses. As this problem is very imbalanced for each profile, learning unsupervised model could require an extremely large number of parame- ters. 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 Appendix C.4 In our embedding, is chosen to be identity and the positional encoding described in Section is used. We compare our model to saute-of-the-as art CNN with convolutional layers and two fully aally-connectdd layers (Zhou & Troyanskaya 2015). The results are shown in Table 3
Sentiment analysis on Stanford Sentiment Treebank. 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. The dataset contains 70, 042 reviews. 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 test set. Our model is one layer of our embedding with yp a Gaussian kernel mapping with 64 Nyström filters in the supervised setting, and linear mapping in the unsupervised setting. 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 al., 2019). State-of-the-ar accuracies after supervised fine-tuning of pre-trained transformers. Training a linear model on pre-trained features after simple pooling (e.g, mean) also yields good results. [CLS], which denotes the BERT embedding used for classification, is also a common baseline. The results are shown in Table 4
usually obtained
5.2 RESULTS AND DISCUSSION
In protein fold classification, our embedding outperforms all baselines in both unsupervised and supervised settings. Surprisingly, our unsupervised model also achieves better results than most supervised baselines. 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. In detection of chromatin profiles, our model (our embedding within a deep network) has fewer layers than Maateoffthe-aa CNNs while outperforming them, which advocates for the use of attention- based models for such applications. Our results also suggest that positional information is important
Table 3: Results for prediction of chromatin profiles on the DeepSEA dataset. The metrics are area under ROC (auROC) and area under PR curve (auPRC), averaged over 919 chromatin profiles. Due to the huge size of the dataset, we only provide results based on a single run.
Method
auROC
auPRC
DeepSEA (hoou& & Troyanskaya 2015
0.933
0.342
Ours with position encoding (Simusodaal Vaswani et al. 2017 /Our))
0.917/0.936
0.31 1/0.360
8

---

Published as a conference paper at ICLR 2021
Table 4: Classification accuracies for SST-2 reported on standard validation set, averaged from 10 different runs (q references supports).
Method
Unsupervised
Supervised
[CLS]] embedding (Devlin et al.
2019
84.6=0.3
90.3++.11
Mcan Pooling of BERIicatures features Devvlin ct al.
2019
85.3=0.4
90.8++011
Approximate Rep the Set (Skianis et al. 2020
Not available.
86,8+0.9
Rep the Set (Skianis al.
2020
Not available.
87.1
Set Transformer (Lee al. 2019
Not available.
87.910.8
Ours (dot-product instead of OT)
85.7
86.911.11
Ours (Unsupervised:
x 300. Supervised: x 30)
86.8-0.3
88.1-10.8
(Appendix C.4 C.2 and our Gaussian position encoding outperforms the sinusoidal one introduced in Vaswani et al. (2017). Note that in contrast to typical transformer, which would have stored 1000 1000 attention matrix, our attention score with a reference of size 64 is only 1000 x 64, which illustrates the discussion in Section 4 In NLP, priori less favorable setting since sequences shorter and there are more data, our supervised embedding gets close to a strong state-of-the-aa Lee. fully-trained transformer. We observed our method to be much faster than RepSet, as fast as Set Transformer, yet slower than ApproxRepSet (C.3) Using the OT plan as similarity score yields better accuracies than the dot-product between the input sets and the references (see Table 2
Choice of parameters. This paragraph sums up the impact of ments justifying our claims can be found in Appendix
choices. Experi-
Number of references q: for biological sequences, single reference was found to be enough in the unsupervised case, see Table I1 In the supervised setting, Table 14 suggests that using remains good baseline, and that the sensitivity
q 5 provides slightly better results but to number of references is moderate.
Number of supports p in reference: Table 11 and Table 14 suggest that the sensitivity of the
model to the number of supports is also moderate.
Nyström anchors: an anchor can be seen as neuron in a feed-forward neural network (see expression of &b in 3.3). In unsupervised settings, the more anchors, the better the approximation of the kernel matrix. Then, the performance saturates, see Table 12 In supervised settings, the optimal number of anchors points is much smaller, as also observed by Chen et al. (2019a), Fig 6. Bandwidth a in gaussian kernel: a was chosen as in [Cen et al. (2019b) and we did not try to optimize it in this work, as it seemed to already provide good results. Nevertheless, slightly better results can be obtained when tuning this parameter, for instance in SST-2.
OTKE and self-supervised methods. Our approach should not be positioned against self- supervision and instead brings complementary features: the OТКE may be plugged in state-of- the-art models pre-trained on large unannotated corpus. For instance, on SCOP 1.75, we use ESM- 1 (Rives et al., 2019), pretrained on 250 millions protein sequences, with mean pooling followed by a linear classifier. As we do not have the computational ressources fine-tune ESM1-134, we only train a linear layer top of the extracted features. Using the same model, we replace the mean pooling by our (unsupervised) OTKE layer, and also only train the linear layer. This results in ac- curacy improvements as showed in Table While training huge self-supervised learning models large datasets is very effective, ESMI-134 admits more than 2500 times more parameters than our Langlellyer OTKE model (260k parameters versus 670M) and our single-layer OTKE outperforms smaller versions of ESMI (43M parameters). Finally, self-supervised pre-training of a deep model including OTKE large data sets would be interesting for fair comparison.
Table 5: Classification accuracy (top 1/5/10) results of our unsupervised embedding for SCOP 1.75 with pre-trained ESM models (Rives et al. 2019).
Model
parameters
Mean Pooling
Unsupervised OTKE
155M--6-43- ESMI-134-670M-URSN
43M
44.01/93.17/95.0 94.95/97.32/97.91
85.91/93.2/995.30 95.22/97.32/98.03
670M
Multi-layer extension. Extending the OTKE to multi-layer embedding is natural yet not straightforward research direction: it is not clear how to find a right definition of intermediate fea- ture aggregation in a multi-layer OTKE model. Note that for DeepSEA, our model with single-layer OTKE already outperforms a multi-layer CNN, which suggests that a multi-layer OTKE is not al- ways needed.

---

Published as a conference paper at ICLR 2021
ACKNOWLEDGMENTS
JM. GM and DC were supported by the ERC grant number 714381 (SOLARIS project) and by ANR 3IA MIAI@Grenoble Alpes, (ANR-19-P3A--0003), AA would like acknowledge support from the ML and Optimisation joint research initiative with the fonds AXA pour la recherche and Kamet 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' a'vveeii"" program, reference ANR-19-P31A-000 (PRAIRIE 31A Institute). DC and GM thank Laurent Jacob, Louis Martin, François-Pierre Paty and Thomas Wolf for useful discussions.
REFERENCES
Jean-Marc Andreoli. Convolution, attention and structure embedding. In arXiv preprint arXiv: 1905.01289v5, 2019.
Dzimitry Bahdanau, Kyunghyun Cho, and Joshua Bengio. Neural machine translation by jointly learning to align and translate. In International Conference on Learning Representations (ICLR), 2015.
Antoni Buades, Bartomeu Coll, and Jean-Michel Morel. Non-Local Means Denoising. Image Processing On Line, 1:208-212, 2011.
Dexiong Chen, Laurent Jacob, and Julien Mairal. Biological sequence modeling with convolutional kernel networks. Bioinformatics, pp. 35(18):324433002 2019a.
Dexiong Chen, Laurent Jacob, and Julien Mairal. Recurrent kernel networks. In Advances in Neural Information Processing Systems (NeurIPS), 2019b.
Liqun Chen, Guoyin Wang. Chenyang Tao, Dinghan Shen, Pengyu Cheng, Xinyuan Zhang, Wenlin Wang, Yizhe Zhang, and Lawrence Carin. Improving textual network embedding with global attention via optimal transport. Proceedings of the Annual Meeting of the Association for Computational Linguistics (ACL), 2019c.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song. Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, David Belanger, Lucy Colwell, and Adrian Weller. Masked language modeling for proteins via linearly scalable long-context transformers. In arXiv:2006603555v2 2020.
Jean-Baptiste Cordonnier, Andreas Loukas, and Martin Jaggi. On the relationship between self- attention and convolutional layers. In International Conference on Learning Representations (ICLR), 2020.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems (NeurIPS). 2013.
Marco Cuturi and Amnad Doucet. Fast computation of wasserstein barycenters. In International Conference on Machine Learning (ICML), 2013.
Zihang Dai, Yiming Yang, Jaime Carbonell, Quoc V. Le, and Ruslan Salakhutdinov. Transformer-xi Attentive language models beyond a fixed-length context. In Proceedings of the Annual Meeting of the Association for Computational Linguistics (ACL), 2019.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the North American Chapter of the Association for Computational Linguistics (NAACL), 2019.
Andrew Gardner, Christian A. Duncan, Jinko Kanno, and Rastko R. Selmic. On the definiteness of earth mover's 's distance and its relation to set intersection. IEEE Transactions Cybernetics, 48 (11):3184-3166, 2018.
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn diver- gences. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2018.
10

---

Published as a conference paper at ICLR 2021
Edouard Grave, Armand Joulin, and Quentin Berthet. Unsupervised alignment of embeddings with wasserstein procrustes. In International Conference on Artificial Intelligence and Statistics (AIS- TATS), 2019.
Nhat Ho, XoanLong Nguyen, Mikhail Yurochkin, Hung Hai Bui, Viet Huynh, and Dinh Phung. Multilevel clustering via wasserstein means. In International Conference on Machine Learning (ICML), 2017.
Jie Hou, Badri Adhikari, and Jianlin Cheng. Deepsf: deep convolutional neural network for map- pingprotein sequences folds. Bioinformatics, pp. 34(5):1295-1300) 2019.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations (ICLR), 2020.
Soheil Kolouri, Yang Zou, and Gustavo K. Rohde. Sliced wasserstein kernels for probability distri- butions. In Proceedings of the Conference Computer Vision and Pullern Recognition (CVPR), 2016.
Soheil Kolouri, Navid Naderializadeh, Gustavo K. Rohde, and Heiko Hoffmann. Wasserstein em- bedding for graph learning. In International Conference on Learning Representations (ICLR), 2021.
Juho Lee, Yoonho Lee, Jungtack Kim, Adam R. Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-based permutation invariant neural networks. In International Conference on Machine Learning (ICML). 2019.
Siwei Lyu. Mercer kernels for object recognition with local features. In Conference on Computer Vision and Puttern Recognition (CVPR), 2004.
Julien Mairal. End-to-end kernel learning with supervised convolutional kernel networks. In Ad- vances Neural Information Processing Systems (NeurIPS), 2016.
Julien Mairal, Piotr Koniusz, Zaid Harchaoui, and Cordelia Schmid. Convolutional kernel networks. In Advances Neural Information Processing Systems (NeurIPS), 2014.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Caroline Moosmüller and Alexander Cloninger. Linear optimal transport embedding: Provable fast wasserstein distance computation and classification for nonlinear problems, 2020.
Gabriel Peyré and Marco Cuturi. Computational optimal transport. Foundations and Trends in Machine Learning, 11(56)):55--06, 2019.
Julien Rabin, Gabriel Peyré, Julie Delon, and Mare Bernot. Wasserstein barycenter and its applica- tion to texture mixing. In International Conference on Scale Space and Variational Methods in Computer Vision (SSVM), 2011.
Alessandro Raganato, Yves Scherrer, and Tiedemann Jörg. Fixed encoder self-attention patterns in transformer-based machine translation. In arXiv preprint arXiv: 2002. 0002./0260. /0260, 2020.
Alexander Rives, Siddharth Goyal, Joshua Meier, Demi Guo, Myle Out, C. Lawrence Zitnick, Jerry Ma, and Rob Fergus. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. In bioRxiv 622803, 2019.
Yossi Rubner, Carlo Tomasi, and Leonidas J. Guibad. The earth mover's distance metric for image retrieval. International Journal of Computer Vision, 121, 2000.
Bernhard Schölkopf and Alexander J. Smola. Learning with kernels: support vector machines, regularization, optimization, and beyond. MIT press, 2001.
Richard Sinkhorn and Paul Knopp. Concerning nonnegative matrices and doubly stochastic matri- ces. Pacific Journal of Mathematics, (1),), 1967.

---

Published as a conference paper at ICLR 2021
Konstantinos Skianis, Giannis Nikolentzos, Stratis Limnios, and Michalis Vazirgiannis. Rep the set: Neural networks for learning set representations. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2020.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2013.
Giorgos Tolias, Yannis Avrithis, and Hervé Jégou. To aggregate or not to aggregate: Selective match kernels for image search. In Proceedings of the International Conference Computer Vision (ICCV), 2013.
Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Transformer dissection: A unified understanding of transformer's attention via the lens of kernel. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2019.
Ashish Vaswani, Noam Shazeer, Nikki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
Cédric Villani. Optimal Transport: Old and New. Springer, 2008.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Proceedings of the Annual Meeting of the Association for Computational Linguistics (ACL), 2019.
Alex Wang. Amanpreet Singh. Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowmnan. Glue: multi-task benchmark and analysis platform for natural language understanding. In Inter- national Conference on Learning Representations (ICI.R) 2019.
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. In arXiv preprint arXiv: 2006.04768:13 2020.
Wei Wang, Dejan Slepcev, Saurav Basu, John A. Ozolek, and Gustavo K. Rohde. A linear optimal transportation framework for quantifying and visualizing variations in sets of images. Interna- tional Journal of Computer Vision, 101(2):254-200 2013.
Xiaolong Wang, Ross B. Girshick, Abhinav Gupta, and Kaiming He. Non-local neural networks. In Proceedings of the Conference on Computer Vision and Pattern Recognition (CVPR), 2017.
You Weiqiu, Simeng Sun, and Mohit lyyer. Hard-coded gaussian attention for neural machine trans- lation. In Proceedings of the Annual Meeting of the Association for Computational Linguistics (ACL), 2020.
Christopher K.I. Williams and Matthias Seeger. Using the nyström method to speed up kernel machines. In Advances in Neural Information Processing Systems (NeurIPS), 2001.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowic7, and Jamie Brew. Huggingface's trans- formers: State-of-the-at natural language processing. In arXiv preprint arXiv: 1910.03771, 2019.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbhakhsh, Barnaáás Póczos, Ruslan Salakhutdinov, and Alexander J. Smola. Deep sets. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
Kai Zhang, Ivor W. Tsang, and James T. Kwok. Improved nyström low-rank approximation and error analysis. In International Conference on Machine Learning (ICML), 2008
Jian Zhou and Olga G Troyanskaya. Predicting effects of noncoding variants with deep learning- based sequence model. Nature methods, 22(10)931-9344 2015.
12

---

Published as a conference paper at ICLR 2021
Appendix
Appendix provides some background on notions used throughout the paper; Appendix(Bcontaiinn the proof' skipped in the paper; Appendix provides additional experimental results as well as details on our protocol for reproducibility.
A ADDITIONAL BACKGROUND
This section provides some background on attention and transformers, Sinkhorn's algorithm and the relationship between optimal transport based kernels and positive definite histogram kernels.
A.1 SINKHORN'S S ALGORITHM: FAST COMPUTATION OF P. (x,
Without loss of generality, we consider here k the linear kernel. We recall that (,, z) is the solu- tion of an optimal transport problem, which can be efficiently solved by Sinkhorn's 's algorithm (Peyné & Cuturi 2019) involving matrix multiplications only. Specifically, Sinkhorn's 's algorithm is an it- erative matrix scaling method that takes the opposite of the pairwise similarity matrix K with entry Kij := (x, 2j) as input C and outputs the optimal transport plan PA P^ z) = Sinkhorn Each iteration step l performs the following updates
1/n and ((+1) 1/p Ev(e)
(e+1)
E (())))
where E = eK/e, Then the matrix diag( (u Ediag( converges to P Px(x,, z) when tends to 00. However when e becomes too small, some of the elements of a matrix product Ev or E become null and result in division by 0. To overcome this numerical stability issue, computing the multipliers u and v is preferred (see e.g. (Peyré & Cuturi 2019 Remark 4.23)). This algorithm can be easily adapted to batch of data points x, and with possibly varying lengths via mask vector masking on the padding positions of each data point x, leading to GPU-friendly computation. More importantly, all the operations above at each step differentiable, which enables to be optimized through back-propagation. Consequently, this module can be injected into any deep networks.
A.2 ATTENTION AND TRANSFORMERS
We clarify the concept of attention - mechanism yielding a context-dependent embedding for each element of x as special case of non-local operations (Wang et al. 2017 Buades et al.) 201T), so that easier to understand its relationship to the OTK. Let us assume we are given a set X of length A non-local operation x is function & X + X such that
(xx);; = Σ w(x, x,)xx)))())) = W(x), Vxx),
where W(x); denotes the i-th column of W(x), weighting function, and V(x) (o(x)).....) ux,)) "T an embedding. In contrast to operations on local neighborhood such as con- volutions, lon-local operations theoretically account for long range dependencies between elements in the set. In attention and the context of neural networks, w is learned function reflecting the rel- evance of each other elements x, with respect to the element x, being embedded and given the task at hand. In the context of the paper, we compare to type of attention coined dot-product self- attention, which can typically be found the encoder part of the transformer architecture (Vaswani et al. 2017). Transformers are neural network models relying mostly a succession of atten- tion layer followed by a fully-comnected layer. Transformers can be used in cewtlenc-tt-eequeeee tasks in this setting, they have encoder with self-attention and decoder part with variant of self-attention -, or in sequence to label tasks, with only the encoder part. The paper deals with the latter. The name self-attention means that the attention is computed using dot-product of linear
13

---

Published as a conference paper at ICLR 2021
transformations of x, and X,, and attends to itself only. In its matrix formulation, dot- product self-attention is a non-local operation whose matching vector is
[ W vds
W(x), Softmax
where Wi E guxd, and Wx E Rnxd are learned by the network. In order to know which x are relevant to x, the network computes scores between a query for x; (Wqx;) and keys of all the elements of x (WK x)).). The sofumax turns the scores into weight vector in the simplex. Moreover, a linear mapping Vxx) = Wvx, the values, is also learned. We and Wx are often shared (Kitaev) al. 2020). A drawback of such attention is that for sequence of length n, the model has to store an attention matrix W with size O(n²). More details can be found in Vaswani et al. (2017).
B PROOFS
B.1 PROOF OF LEMMA 3.1
Proof. First, since Σラ pP(x', 2)zx = for any k, we have
n
p
W 22
ΣΣ
P (x, 2) (x,,
k=1
pP 2) 2) 2) 2x)
with vector in Ren'r whose entries are p 2) P(x, Z)u,d, (x;, Zz) for 1,..., j I....., n' and 1.....P. We can also rewrite W (x, x') as an f2-norm of a vector in Ran't P whose entries are /pP(x', 2)jx P(x, f--norm, have
Then by Minkowski inequality for the
122(x,2) (x(xxx))
|u| l|vl|al
I 11
v|12
1/2
pP 2) P(x, 2), () 2τ) d, (x,, x') ,)2
1/2
pP z) 2)
1/2
P 2)
(2(x'z).
where the second inequality is the triangle inequality for the distance d,. Finally, have
[2'xx,') Wz(xxx') x')
(x, W2(x, z)) + [W2(x,2) Wz(xx
W2(X', 2) + W2(x', =2 z),
where the second inequality is the triangle inequality for the 2- Wasserstein distance. By symmetry, also have [x, x') W7 (x, x') 2W2(x,2) z), which concludes the proof. n1
14

---

Published as a conference paper at ICLR 2021
B.2 RELATIONSHIP BETWEEN W2 AND W FOR MULTIPLE REFERENCES
Using the relation prooved in Appendix B.I we can obtain a bound on the error term between W22 and W for a data of m samples (x+ x") and q references (z² ;...,.99 x²)
E
12(x', := m 1J=1
W2 ......" (x(,x))
4 mq
(x),))2)
(9)
When 9 1, the right-hand term in the inequality is the objective minimize in the Wasserstein barycenter problem (Cuturi & Doucet 2013). which further explains why we considered it: Once W is close to the Wasserstein distance W2. Ka will also be close to Kor- We extend here the bound in equation O in the case of one reference to the multiple-reference case. The approximate 2- Wasserstein distance (x x') thus becomes
W2 W (x,x')=
1
Σ
? P x'), dx
4/2
(1
1
Σ i=1
W (x,x')²
1/2
Then by Minkowski inequality for the fe-norm we have
3/2
-
4/2l
4
W2((x, x') W22 (x,x')1 =
xxx
Σ
W (x, x)2
-
1/2
W (2((x,x)) "V)(x,')-W5"" W3''5xxxx00, (x,x))²
9
and by equation 6 we have
4/2
(x)xxx)) - W22 (x, (xxx'|5 S
min (xinnWxx,,2)) W2(x',
Finally the approximation error in terms of Frobenius is bounded by
4
82 m2 1,j=1
W2 x²) - W.2 ....."" (x' x2)2
W3 (x),z))
mq
i=1 j=1
In particular, when q 1 that is the case of single reference, we have
m
W 2),
where the right term equals the objective of the Wasserstein barycenter problem, which justifies the choice of z when learning without supervision.
C ADDITIONAL EXPERIMENTS AND SETUP DETAILS
This section contains additional experiments on CIFAR-10. 10, whose purpose is to illustrate the kernel associated with our embedding with respect 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.
C.I EXPERIMENTS ON KERNEL MATRICES (ONLY FOR SMALL DATA SETS).
Here, compare the optimal transport kernel Kor ( and its surrogate Kg (2) (with z learned
without supervision) to common and other OT kernels. Although our embedding 2, is scalable, the exact kernel require the computation of Gram matrices. For this toy experiment, we therefore use
15

---

Published as a conference paper at ICLR 2021
Table 6: Classification accuracies for 5000 samples of CIFAR- using CKN features (Mairal 2016 and forming Gram matrix. A random baseline would yield 10%.
Dataset
(3 x 3), 256
Kernel
Accuracy
Runtime
Mean Pooling
58.5
30 sec
Flatten
67.6
30 sec
Sliced Wasserstein Kolouri
2016
63.8
min
Sliced- Wasserstein Kolourt et al. 2016 + sin. pos enc. Devlin et al. 2019
66.0
~ min
Kor
64.5
~ 20 min
Kor + our pos enc.
67.1
20 min
K₂
67.9
30 sec
Kx our pos enc.
70.2
30
5000 samples only of CIFAR-I0 (images with 32 x 32 pixels), encoded without supervision using a two-layer convolutional kernel network (Mairal 2016). The resulting features are 3 x 3 patches living in RR with d 256 or 8192. Kor and K₂ aggregate existing features depending ground cost defined by -x (Gaussian kernel) given the computed weight matrix P. In that sense, can say that these kernels work as adaptive pooling. We therefore compare it to kernel matrices corresponding to mean pooling and no pooling at all (linear kernel). We also compare to a recent positive definite and fast optimal transport based kernel, the Sliced Wasserstein Kernel (Kolouri et al. 2016) with 10, 100 and 1000 projection directions. We add positional encoding to it so as to have a fair comparison with our kernels. A linear classifier is trained from this matrices. Although we cannot prove that Kor is positive definite, the classifier trained on the kernel matrix converges when 8 is not too small. The results can be seen on Table [ Without positional information, our kernels do better than Mean pooling. When the positions are encoded, the Linear kernel is also outperformed. 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. All interpolations did worse than the Linear kernel. The runtimes illustrate the scalability of Kz-
the
C.2 CIFAR-10
Here, we test our embedding on the same data modality: we use CIFAR-10 features, ie., 60, 000 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. The very best configuration of the CKN yields small number x 3) of nigh-dimensional (16, 384) patches and an accuracy of 85.8%. We will illustrate our embedding on configuration which performs slightly less but provides more patches (16 16), setting for which it was designed.
The input of our embedding are unsupervised features extracted from a 2-layer CKN with kernel sizes equal to and 3, and Gaussian pooling size equal to 2 and 1. We consider the following configurations of the number of filters at each layer, to simulate two different input dimensions for our embedding:
64 filters first and 256 at second layer, which yields 16 16 x 256 representation for each image.
256 filters at first and 1024 at second layer, which yields 16 16 x 1024 representation for each image.
Since the features are the output of Gaussian embedding, k in our embedding will be the linear kernel. 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 CKN's Gaussian pooling with pooling size equal to 6. The hyper-parameters are the entropic regu- larization E and bandwidth for position encoding pos- Their search grids are shown in Table 7 and the results in Table & Without supervision, the adaptive pooling of the CKN features by our embed- ding notably improves their performance. We notice that the position encoding is very important to this task, which substantially improves the performance of its counterpart without it.
16

---

Published as a conference paper at ICLR 2021
Table 7: Hyperparameter search range for CIFAR-10
Hyperparameter
Search range
Entropic regularization
0.0:0.1;0001:000000 0.5: 0.,,
Position encoding bandwidth pps
Table 8: Classification results using unsupervised representations for CIFAR- 10 for two feature configurations (extracted from a 2-layer unsipervised CKN with different number of filters). We consider here our embedding with one reference and different number of supports, learned with K-means, with without position encoding (PE).
Method
Nb. supports
16 x 16 x 256
16 x 16 x 1024
Flatten
73.1
76.1
Mean pooling
64.9
73.4
Gaussian pooling MMiral 2016
77.5
82.0
Ours
75.6
79,3
Ours (with PF)
78.0
82.2
Ours
77.9
80.1
64
Ours (with PF)
81.4
83.2
Ours
78.4
80.7
144
Ours (with PE)
81.8
83.4
C.3 PROTEIN FOLD RECOGNITION
Dataset description. Our protein fold recognition experiments consider the Structural Classifica- tion Of Proteins (SCOP) version 1.75 and 2.06. We follow the data preprocessing protocols in Hou et al. 2219 which yields training and validation set composed of 14699 and 2013 sequences from SCOP 1.75, and test set of 2533 sequences from SCOP 2.06. The resulting protein sequences belong 1195 different folds, thus the problem is formulated as multi-classification task. The input sequence is represented as a 45-dimensional vector at each amino acid. The vector consists of a 20-dimensional one-hot encoding of the sequence, a 20-dimensional position-specific scoring matrix (PSSM) representing the profile of amino acids, 3-class secondary structure represented by a one-hot vector and 2-class solvent accessibility. The lengths of the sequences are varying from tens to thousands.
Models setting and hyperparameters. We consider here the one-layer models followed by global mean pooling for the baseline methods CKN (Chen et al. (019a) and RKN (Chen et al., 2019b). We build our model on top of the one-layer CKN model, where k can seen as a Gaus- sian kernel the k-mers in sequences. The only difference between our model and CKN is thus the pooling operation, which is given by our embedding introduced in Section 3 The bandwidth parameter of the Gaussian kernel supervised models, the same as used in CKN which were selected by the accuracy on the valida- tion The filter size k is fixed 10 and different numbers of anchor points in Nyström for are considered in the experiments. The other hyperparameters for our embedding are the entropic regularization parameter the number of supports in reference p. the number of references q, the number of iterations for Sinkhorn's 's algorithm and the regularization parameter À in the linear classifier. The search grid for E and À is shown Table and they are selected by the accuracy on validation set. e plays important role in the performance and is observed to be stable for the same dataset. For this dataset, is selected be 0.5 for all the unsupervised and supervised models. The effect of other hyperparameters will be discussed below.
k-mers is fixed 0.6 for unsupervised models and 0.5 for
For the baseline methods, the accuracies of PSI-BLAST and DeepSF are taken from Hou et al. (2019). The hyperparameters for CKN and RKN can be found in Chen et al. (2019b). For Rep the Set (Skianis et al. 2020) and Set Transformer (Lee et al. 2019), we use the public implementations by the authors. These two models are used the top of convolutional layer of the same filter size as CKN to extract k-mer features. As the exact version of Rep the Set does not provide any implementation for back-propagation bottom layer of it, we consider the approximate version of
17

---

Published as a conference paper at ICLR 2021
Table 9: Hyperparameter search grid for SCOP 1.75
Hyperparameter
Search range
E for Sinkhorn
[1.0; 0.5;; 0.1;0.1 0.01] 1/2
À for classifier (unsupervised setting)
classifier (supervised setting)
[Ie-6:le-5:le4:lee-3]
Table 10: Hyperparameter search grid for SCOP 1.75 baselines.
Model and Hyperparameter
Search range
Sppro:
Hidden Sets x Cardinality
[20;30; 50; 100) x [10; 20; 50]
ApproxRepSet: Learning Rate
[0.0001; 0.0005; 0.001]
ApproxRepSet: Weight Decay
[Ie-S:le-4:1e3:3le-2]]
Transformer: Heads x Dim Hidden
[1:4:8] [64; 128; 256]
Set Transformer: Learning Rate
0.000;
Set Transformer: Weight Decay
[lc-5:le-41e-3:lee2]]
Rep the Set only, which also scales better to dataset. The default architecture of Set Transformer did not perform well due to overfitting. We therefore used a shallower architecture with one ISAB, one PMA and one linear layer, similar the one-layer architectures of CKN and model. We tuned their model hyperparameters, weight decay and learning rate. The search grids for these hyperparameters are shown in Table 10
Unsupervised embedding. The kernel embedding ş, which is infinite dimensional for the Gaus- sian kernel, is approximated with the Nyström method using K-means on 300000 k-mers extracted from the same training set as in Chen al. (2019b). The reference measures are learned by using either K-means or Wasserstein to update centroids in Wasserstein K-means on 3000 subsampled model top of features extracted from CKNs of different dimensions, representing the number of anchor points used to approximate k. The num- ber of iterations for Sinkhorn is fixed to 100 ensure the convergence. The results for different combinations of g and are provided in Table 11 Increasing the number of supports p can improve large. On the other hand, increasing the number of
sequences for RAM-savvng reason. We evaluate
the performance and then saturate it when p references while keeping the embedding dimension (i.e. q) constant is not significantly helpful in this unsupervised setting. We also notice that Wasserstein Barycenter for learning the references does not outperform K-means, while the latter is faster in terms of computation.
Supervised embedding. Our supervised embedding is initialized with the unsupervised method and then trained in alternating fashion which was also used for CKN: to update anchor points in Nyström and reference measures z, and the L-BFGS algorithm to optimize
an Adam algorithm
Table 11: Classification accuracy (top 1/5/10) results of our unsupervised embedding for SCOP 1.75. We show the results for different combinations (number references x number of supports p)- The reference measures are learned with either K-means Wasserstein barycenter for updating centroids.
Embedding size (g X p)
Nb. filters
Method
9
10
50
100
200
76.5/91.5/94.4
77.5/91.7/94.55
79.4/92.4/94
78.7/92.1/95..
K-means
5
9/77
77.8/91.7/94..6
78.6/91.9/94
18.1//2.1/94.7
62.7/85.8/91.1 10
76.5/91.0/94.2
78.1/92.2/944.9
78. 6/..292//
128
1
64.0/85.9/91
71.6/88.9/93.2
77.2/91.4/944
77.5/91.9/9 19999
Wass. bary.
5
70.5/89. 1/93.0
76.6/91.3/94.4
78.4/91.7/94.3
9/94.7
10
63.0/85.7/91.0
75.9//1.4/944.3
77.5/91.//4.6
77.9/92.0/94.7
1
84 84.4/550//96
04.6995.0/97.0
85.7/95.3/96.7
85 4/95 2/96.7
1024
K-means
5
81. 11.11/94.000 0/9662
84.9/94.8/9 8/96.8
84.7/94
$5.2/95 0/96.7
10
79.8/93.5/95.9
6/.19
84.4/94.7/96.7
84.8/94.9/96.7
18

---

Published as a conference paper at ICLR 2021
Table 12: Classification accuracy (top 1/5/10) results of unsupervised embedding for SCOP 1.75. We show the results for different number of Nyström anchors. The number of references and supports are fixed and 100.
Number of anchors
Accuracies
1024
55.8/95.3/96.8
2048
86.6/95. 05.6/9959/97..2
3072
87.8/96. 1/97.4
Table 13: Classification accuracy (top 1/5/10) of supervised models for SCOP 1.75. The accuracies obtained by averaging 10 different runs. We show the results of using either one reference with 50 supports or references with 10 supports. Here DeepSF is a 10-layer CNN model.
Method
Runtime
Top 1/5/10 accuracy on SCOP 2.06
PSI-BLAST Hoo et al. 2019
84, 84.33/86.4888 48/87.34
DeepSF Hoo ot al 2019
73.00/90.25/944 40/94.33-.9.
Set Transtormrr LLee
2019
3.3h
ApproxRepSet (Skianis et 0. 2020
2h
Number of filters
128
512
CKN CCen al. 2019
1.5h
N4.11:H011194422HH022202 16.11.
RKN Len et
65.29H02779195550.586...8
Ours
2. (1 50)
3,5h
88 011.22229576-.133777 10
@x (5 10)
84.68 0.588:00500/4..6 1-0.44..68--0..00000 1..29999.999999911... 0.18
88.66:+0.25/95.99 -.1.5990000.11/7.333333333
the classifier. The learning rate for Adam is initialized with 0.01 and halved as long there is no decrease of the validation loss for 5 successive epochs. practice, we notice that using a small number Sinkhorn iterations can achieve similar performance being much faster to compute. We thus fix it to 10 throughout the experiments. The accuracy results are obtained by averaging on 10 runs with different seeds following the setting Chen al. (2019b). The results are shown in Table 13 with error bars. The effect of the number of supports similar the unsupervised case, while increasing the number references can indeed improve performance.
large number of iteration, while
C.4 DETECTION OF CHROMATIN PROFILES
Dataset description. Predicting the functional effects of noncoding variants from only genomic sequences is central task in human genetics. A fundamental step for this task is simultane- ously predict large-scale chromatin features from DNA sequences (Zhou & Troyanskaya 2015). We consider here the DeepSEA dataset, which consists simultaneously predicting 919 chromatin profiles including 690 transcription factor (TF) binding profiles for 160 different TFs, 125 DNase I sensitivity profiles and 104 histone-mark profiles. total, there are 4.4 million. 8000 and 455024 samples for training, validation and test. Each sample consists of a 1000-bp DNA sequence from the human GRCh37 reference. Each sequence represented as a 1000 x binary matrix using one- hot encoding on DNA characters. The dataset is available at ht : /eeessea.eeeeeeaeeeeeeeee .rr.ppri wwwaapprninoooo.oo edu/medde attde//////////////ee/eee/eeeeee /a//////e/eeeseaaaaaeaaaae/aaeee train bundle. v0. 9. tar gz Note that the labels for each profile are very imbalanced in this task with few positive samples. For this reason, learning unsu-
Table 14: Classification accuracy (top 1/5/10) results of supervised embedding for SCOP 1.75. We show the results for different combinations of (number of references x number of supports p)- The reference measures are learned with K-means.
Embedding size (g p)
10
50
100
200
q
88.3/95.5/97,
88 98.4/95.8/97.2
87.1/94.9/6..7
$7.7/94.9/96.2
q 2
87.8/95.8/9 8/97.0 0
89 2/99/
86.5/94.9/96.6
87. 87.6/94.9/96.3
q 5
$7.0/95.1 1/96.7
87.4/95.4/97.0
87. 17.4/94.7/96.
g 10
84.5/93.6/95
89.8 89.8/96.0/97.2
88.0/95.7/97.0
85.6/94.4966.1
19

---

Published as a conference paper at ICLR 2021
Table 15: Model architecture for DeepSEA dataset.
Model architecture
Convid(in channels=4, out channels=d, kernel sizc=16) + ReLU
(Ours) EmbeddingLayerrin channels=d, supports=64, references=l, c 1.0, PE=True, Ope = 0.1) Linear(in channels=d, out channels=d) ReL.U Dropout(0.4)
Linear(in channelsed 64, out channels=919) + ReLU
Lin channels=919, out channels=919)
Results for prediction of chromatin profiles on the DeepSEA dataset. The metrics are area under ROC (auROC) and area under PR curve (auPRC), averaged over 919 chromatin profiles. The accuracies are averaged from 10 different runs. Armed with the positional encoding (PE) de- scribed in Section S} our embedding outperforms the state-of-the-art model and another model of our embedding with the PE proposed in Vaswani et al. (2017).
Table 16:
Method
DeepSEA
Ours
Ours (d 1024)
Ours (d 1536)
Position encoding
Sinusoidal Vaswani et al. 2017
Ours
Ours
auROC
0.933
0.917 0.311
0.935
0.936
auPRC
0.342
0.354
0.360
pervised models could be intractable as they may require an extremely large number of parameters junk or redundant sequences cannot be filtered out.
Model architecture and hyperparameters. For the above reason and fair comparison, we use here our supervised embedding as module in Deep NNs. The architecture of our model is shown Table [5] 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. The number of iterations for Sinkhorn is fixed to 30. The whole training process takes about 30 hours on a single GTX20S0TI GPU. The dropout rate selected to be 0.4 from the grid 10 0.2; 0.3; 0.4; 0.5) and the weight decay is le-06, the same as Zhou & Troyanskaya (2015). The pos for position encoding is selected to be 0.1. by the validation accuracy on the grid (0.05; 0.2; 0.3; 0.4; 0.5). The checkpoint with the best validation accuracy is used to evaluate on the test set. Area under ROC (auROC) and area under precision curve (auPRC), averaged over 919 chromatin profiles, are used to measure the performance. The hidden size d is chosen to be either 1024 or 1536.
Results and importance of position encoding. We compare our model to the saae-of-thhean CNN model DeepSEA (Zhou & Troyanskaya] 2015) with 3 convolutional layers, whose best hyper- parameters can be found in the corresponding paper. Our model outperforms DeepSEA, while requiring fewer layers. The positional information is known to be important in this task. To show the efficacy of our position encoding, we compare it to the sinusoidal encoding used in the original transformer (Vaswani al. 2017). We observe that our encoding with properly tuned pos requires fewer layers, while being interpretable from a kernel point of view. We also find that larger hidden size d performs better, as shown in Table 16 ROC and PR curves for all the chromatin profiles and stratified by transcription factors, DNase l-hypersensitive sites and histone-mate marks can also be found Figure
C.5 SST-2
Dataset description. The data set contains 67,349 training samples and 872 validation samples and can be found at / weg.uuuecchaaaaammm ples for which the predictions need to be submitted on the GLUE leaderboard, with limited number of submissions. As 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 validation set), and we report accuracies set. 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
com/ task/ The test set contains 1,821 sam-
the standard validation set, used as test
20

---

Published as a conference paper at ICLR 2021
ROC cerves jag 1000-1 137)
Piceoss Dog kuR 202)
02
06
00
64
52
C4
lav
Figure 2: ROC and PR curves for all the chromatin profiles (first row) and stratified by transcrip- tion factors (left column), DNase I-hypersensitive sites (middle column) and histone-marks (right
column). The profiles with positive samples fewer than 50 on the test set are not taken into account.
21

---

Published as a conference paper at ICLR 2021
Table 17: Accuracies on standard validation set for SST-2 with unsupervised features depending on the number of references and supports, The references were computed using K-means on samples for multiple references and K-means on patches for multiple supports. The size of the imput BERT features is (length x dimension). The accuracies are averaged from 10 different runs. (g references x p supports)
BERT Input Feature Size
(30 x 768)
(66 768)
Features
Pre-trained
Finc-tuned
Pre-trained
Fine-tuned
[CLS] Flatten
84.6++..3
90.3++0.1
86.0+10.2
92.8=0.1
84.9=+044
91.0++.11
85.2+0.3
92.5=0.1
Mcan pooling
85.3+00.3
90.8+0.1
85.4+003
92.6=0.2
d₂ (1 x 3)
85,5+0011
90.9+0.1
86.5+0.1
92.6=0.1
@₂ (1 x 10)
85.1=0.4
90.9+0.1
85.9
92.6=0.1
4, x 30)
86.3+0.3
90.8+0.1
86.6#0.55
92.6=0.1
42 x 100)
85.7+10.7
90.911011
86.6:10.1
92.7-0.1
4₂ (1 x 300)
86.
90.9
87.2
92.7-0.1
Table 18: Hyperparameter search grid for SST-2.
Hyperparameter
Search range
Entropic regularization E
[0.5; 1.0; 3.0; 10.0] 10 n 10.3)
X for classifier (unsupervised setting)
Gaussian kernel bandwidth
[0.4; 5: 0.6; 0.7; 0.8] 00.1:001;000001
Learning rate (supervised setting)
implementation requires to add special tokens [CLS] and [SEP] at the beginning and the end each review.
Model architecture and hyperparameters. In most transformers such as BERT, the embedding associated to the token [CLS] is used for classification and can be seen in some sense as em- bedding of the review adapted to the task. The features we used are the word features provided by wwtpttatttooooooocoocoaco
the BERT base-uncased version, available at htttp::/hhugingeee pretrained mmddels html, For this version, the dimension of the word features is 768. Our model is one layer of our embedding, with o the Gaussian kernel mapping with varying number of Nyström filters in the supervised setting, and the Linear kernel in the unsupervised setting. We do not add positonnal encoding as is already integrated BERT features. In the unsupervised setting, the output features are used to train large-scale linear classifier, Pytorch linear layer. We choose the best hhper--arameee parameters based on the accuracy of a validation set. In the supervised case, optimized end-to-end In this case, tune the
the parameters of the previous model, and learning rate. In both case, we tune the entropic regularization parameter of optimal transport and the bandwidth of the Gaussian kernel. The parameters in the search grid are summed up in Table [8] The best entropic regularization and Gaussian kernel bandwidth are typically and respectively 3.0 and 0.5 for this data set. The supervised training process takes between half hour for smaller models (typically 128 filters in w and 3 supports in z) and few hours for larger models (256 filters and 100 supports) on single GTX20S0TI GPU. The hyper-parameters of the baselines were simi- larly tuned, see 19 Mean Pooling and [CLS] embedding did not require any tuning except for the regularization À of the classifier, which followed the same grid as in Table 18
Results and discussion. As explained in Section BERT pre- pre-trained features while still using a simple linear model as shown in Table 17 and its supervised counterpart enables to get even eloser the state-of-the art (for the BERT unca-dd model) accuracy, which is usually obtained after fine-tuning of the BERT model on the whole data set. This can be seen in Tables 20 21] We also add baseline consisting of one layer of classical self-attention, which did not do well hence was not reported in the main text.
unsupervised embedding improves the
22

---

Published as a conference paper at ICLR 2021
Table 19: Hyperparameter search grid for SST-2 baselines.
Model and Hyperparameter
Search range
RepSet and ApproxRepSet: Hidden Sets x Cardinality
[4; 20; 30; 50; 100] x [3; 10; 20; 30; 50] [0.0001] 0.001; 0.01]
ApproxRepSet: Learning Rate
Set Transformer: Heads Dim Hidden
[1: [8: 16; 64; 128]
Transformer: Learning Rate
[0.001; 0.01]
Table 20: Classification accuracy on standard validation set of supervised models for SST-2, with pre-trained BERT (30 768) features. The accuracies of our embedding were averaged from 3 dif- ferent runs before being run 10 times for the best results for comparison with baselines, cf. Section 5 10 Sinkhorn iterations were used. (q references p supports). Accuracy on SST-2
Method
Number of Nyström filters
32
64
128
2₂ (1
88.38
88.38
88.18
2z (1 10)
88.11
88.15
87.61
@z (1 30)
88.30
88.30
88.26
@ (4 x 3)
88.07
88.26
88.30
@, (4 10)
87.6
87.84
88.11
4, (4 30)
88.18
88.68 68.68
88.07
Table 21: Classification accuracy standard validation set of all baselines for SST-2, with pre- trained BERT (30 x 768) features, averaged from 10 different runs.
Method
Accuracy SST-2
[CLS] embedding DDvllin al.
2019 Devlin al.
90.3
Mean Pooling of BERT Teatures
2019
90.8 + 0.1
One Self-Attention Layer
Vaswann al. 2017
83.7+0.1
Approximate Rep the Set
Skianis et al. 2020
86.8+0.99
Rep the Set (Skianis et al.
2020)
87.1
Set Transformer (Lee al.
2019
87.9
dz (1 x 30) (dot-product instead of OT)
16.
23


Reviews:
[{"review_id": "d7rq4cqrlo", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a kernel embedding for a set/sequence of features, based on the optimal transport distance to a set of references, leading to a fixed-dimensional embedding of variable length sequences.\nThe set of references can be obtained as cluster centers over the full dataset (unsupervised), or learned based on a downstream objective (supervised).\nThe method is somewhat related to a single layer of self-attention with an optimal transport map instead of dot product attention.\nExperiments on classifying variable-length sequences are presented between protein, chromatin and NLP sentiment classification.\n\nI recommend weak accept because (+) a well-motivated novel embedding/pooling architecture with (-) somewhat weak, initial proof-of-concept style results and (-) somewhat strange positioning against self-supervised methods and self-attention.\n\nStrenghts:\n* Well-motivated novel layer for embedding a variable-size set or sequence to a *fixed dimensional* embedding space (a point that's not directly apparent though, consider emphasizing this point)\n* Relatively elegant formulation grounded in OT and kernel methods.\n* The same framework allows an unsupervised and supervised variant.\n\nWeaknesses:\n* Experimental results are not very strong. For protein fold classification, I expect a comparison to transformer-based methods on proteins (eg ESM-1, https://github.com/facebookresearch/esm), simply average-pooled over the full sequence. For sentiment analysis, it is cool to see the unsupervised/linear probe on fixed features, beating the frozen BERT - however not competitive to BERT+finetuning which is how it is usually done.\n* Given weak results, the impact of this method is not very clear. It is only shown on shallow single-layer OTKE, without clear discussion how this could be extended to deeper architectures.\n* Given the shallow architecture, the paper is somewhat mispositioned, given the (a) introduction calling out limitations of self-supervised transformers and (b) comparison to self-attention in Sec 4: this suggests inserting the OTKE as a drop-in replacement for the dot product in a transformer model, trained with a self-supervised objective like MLM. Given the framing it seems such a logical step, it surprises me to see at least a discussion of it missing. This would be akin to recent work [1,2] of approximating self-attention by mapping to a fixed-dimensional space.\nAlternatively, it seems like the method could be applied on top of strong/SotA MLM-transformer embeddings to enable better pooling than just [CLS] token or mean pooling (both protein fold prediction and sentiment classification).\n\nMinor comments / questions:\n* The labeling \"unsupervised\" in Table 2 and 4 is somewhat confusing/misleading, rather I'd name it \"linear probe\" or \"supervised finetuning\". Maybe re-clarify in intro of 5.1 \"supervised == learning the z also\"?\n* What are the input features for protein fold classification? it is just mentioned they are 45-dimensional, not what is their nature (amino acid identity + pssm?)\n* On end-to-end learning by back-propping through unrolled sinkhorn (Sec 3.4): this seems to be the method introduced in [3]? I think this paper should be cited in this context.\n* For clarity: In both abstract, end of introduction, start of Sec 3.1 and Sec3.3, it is simply not clear to the reader what is the input/output of your method: a phrase like this would be much needed: \"an embedding+pooling layer which aggregates a variable-size set or sequence of features to a fixed-size embedding\".\n\n\n- [1] arXiv:2006.03555 Choromanski et al. Masked Language Modeling for Proteins via Linearly Scalable Long-Context Transformers\n- [2] arXiv:2006.04768 Wang et al. Linformer: Self-Attention with Linear Complexity\n- [3] Genevay, A., Peyr´e, G., Cuturi, M., et al. (2018). Learning Generative Models with Sinkhorn Divergences. In AISTATS.\n\n\n\n----\n### EDIT: reply to authors' response\nI thank the authors for their in-depth response and revision. The additional results comparing to pre-trained features definitely adds perspective, and it shows OTKE giving a very modest boost indeed, mostly over the weaker input features. The revision of the paper also improves the clarity, but my comment around Sec4 remains.\nI stand by my original score for the 3 reasons I had originally listed.", "extra_scores": {}, "extra_text": {}}, {"review_id": "5vM8iwoHn5D", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Summary**:\nThe authors propose a new way to aggregate the embeddings of elements in a set (or sequence) by comparing it with respect to (trainable) reference set(s) via Optimal Transport (OT). The motivation to build such a pooling operation is derived from self-attention and the authors suggest an OT spin to that (e.g., the different reference sets/measures can be thought of as different heads in attention). This is, however, done in a principled way with the help of kernel embeddings and not just ad-hoc using the transport plan as the attention matrix.\n \n**Pros**:\n-They properly bridge the gap of OT with attention via their non-local pooling perspective. This allows them to have a similar feature aggregator, but which requires linear memory as compared to quadratic for the vanilla self-attention (although there have been several recent linear memory variants to it, any comments?). \n-The method results in an improvement over other baselines for several biology-based sequencing applications (also they demonstrate early results for NLP). \n-The paper is well written and clear.\n\n**Cons**: \n-The empirical results are okay, but could have been better or tested more extensively. Also, there are some peculiarities about their empirical analysis. \n-This kind of technique has promise for use in NLP, but I am slightly worried that right now there are just too many \"knobs\" whose correct configuration will have to be worked out (anchor points, reference measures, nystrom, etc). However, I think this can be left to future work. \n\n*-> It would be great if the authors can answer the questions below:* \n\nEmpirical analysis:\n-In the experiments for Table 2, 3, 4, do you use the same number of heads for other baselines as the number of references in your case?\n\n-It seems quite funny that, for their unsupervised results, the optimal number of supports p equal 1, which would basically imply that optimal transport wasn’t much of use and the features of x_i were just pooled based on an average. Do you have any comments on this and why this is happening?\n\n-The NLP experiments are relatively weak and could still have been strengthened by testing on more tasks from the GLUE benchmark. Another surprising thing is that even though scores of your method unsupervised are higher than mean pooling of BERT features, the latter significantly outperforms in the supervised case. Do you have any explanation for this?\n\n-Besides, why do you write in the text that NLP is “an a priori less favourable setting”?\n\n-Can you give some precise estimates of how the runtime of this OTKE compares to dot product attention? Also, how expensive is the Nystrom procedure?\n\nMiscellaneous: \n-This is a bit of nitpicking, but the authors are a bit lazy while citing, and in several places just cite the textbooks for OT and Kernels (which is not a problem). However, it is useful for the reader to additionally have the accurate references as well, e.g. for Sinkhorn’s algorithm perhaps also cite Sinkhorn & Knopp (1967) and Cuturi (2013); for OT theory, Villani 2008 etc.  \n-Typos: \"a set x and a reference z in X and a reference z”\n-ISAB and PMA (unexplained acronyms)\n\n**Conclusion**:  Overall, I think the idea is quite interesting, and carries the potential for OT-based building blocks in NLP, or even analysing the benefit provided by attention. Hence, I am in favour of accepting this paper.  \n\n(PS: I might be slightly biased because I like OT and have also been thinking about ideas on similar lines connecting attention and OT.)", "extra_scores": {}, "extra_text": {}}, {"review_id": "38vfg9_DABA", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes a transport-based feature representation for the input of a set of vectors. The feature is defined through the transport to reference sets. Vectors in reference sets can be learned by supervised or unsupervised approaches. The relation with self-attention is also discussed. Experiments empirically show a good performance of the proposed method.\n\nTransporting a set of inputs to trainable references is seemingly a more sophisticated extension of the classical basis function regression with trainable basis (basis corresponds to references), and thus, I think the basic idea would be reasonable and easy to have an intuition. My major concern is on about the settings of the number of references and supports, which would have a large number of possible settings. Except for that, I currently only have minor comments.\n\nDiscussing the selection of the number of references and supports in more detail would have been informative. In my understanding, these are most important hyper-parameters. Currently, no practical recipes and analysis of sensitivity to these settings are provided. How were they chosen in the experiments?\n\nThe number of supports can be different for each one of references. Therefore, there exist a huge number of possible settings even when only considering moderate sizes.\n\nThe optimized references seemingly have some interpretation about discovery of an important set of vectors. I'm not fully sure, but it seems to provide a different view point from self-attention, and thus, it is perhaps worth mentioning if possible. \n\nA naive extension could be to make multiple transportation layers (transport z to another references). Is it meaningless? (I guess, at least, some activation function is required) If it has some meaning, providing comments may be informative.\n\nIn eq.(5), 'R^k x p' should be 'R^p x k'?", "extra_scores": {}, "extra_text": {}}, {"review_id": "0z-mF1x74-k", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "##########################################################################\n\nSummary:\n\nThe paper proposes a parametrized embedding based on optimal transport in the kernel space. The method is scalable and the parameters of the embedding can be learned in either unsupervised or supervised fashion. The paper demonstrates the effectiveness of the purposed approach using real data in bioinformatics and natural language processing. \n\nOverall, the proposed embedding seems to be novel and well-motivated. The empirical results are impressive. Also, the paper is well written and technically sound. I don't have major concerns about the paper. \n\n\n##########################################################################\n\nMinor comments: \n\n- Contributions paragraph: \"in either unsupervised and supervised fashion\" --> \"in either unsupervised or supervised fashion\"\n- The line below Eq. (1): the entropy should be - \\sum P ( log P - 1 ) ? \n- Section 3.2 \"Infinite-dimensional embedding in RKHS\" line 1: \"a reference z in X\" repeated twice.\n- Third line below Eq. (4): \\Phi_z(x) instead of \\Phi_z(z)\n- First line in Section 5.2: \"either unsupervised or supervised settings\" --> \"both unsupervised and supervised settings\"", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers agreed that the paper proposes some interesting and novel ideas on the use of OT for pooling. It also provides some nice insights and strong experimental results. As suggested by one of the reviewer, a discussion about the impact of the number of references may be of interest though.


Decision: Accept (Poster)