A Trainable Optimal Transport Embedding for Feature Aggregation and its Relationship to Attention

1 INTRODUCTION . Many scientific fields such as bioinformatics or natural language processing ( NLP ) require processing sets of features with positional information ( biological sequences , or sentences represented by a set of local features ) . 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 et al. , 2019 ; Skianis 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 labeled data . Some of these models use attention ( Bahdanau et al. , 2015 ) , a classical mechanism for aggregating features . Its typical implementation is the transformer ( Vaswani et al. , 2017 ) , which has shown to achieve state-of-the-art results for many sequence modeling tasks , e.g , in NLP ( Devlin et al. , 2019 ) or in bioinformatics ( Rives et al. , 2019 ) , when trained with self supervision on large-scale data . Beyond sequence modeling , we are interested in this paper in finding a good representation for sets of features of potentially diverse sizes , with or without positional information , when the amount of training data may be scarce . To this end , we introduce a 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 ( Peyré & Cuturi , 2019 ) and kernel methods ( Schölkopf & Smola , 2001 ) . We call this embedding OTKE ( Optimal ∗Equal contribution . †Univ . Grenoble Alpes , Inria , CNRS , Grenoble INP , LJK , 38000 Grenoble , France . ‡D.I. , UMR 8548 , École normale supérieure , Paris , France . Transport Kernel Embedding ) . Concretely , we embed feature vectors of a given set to a reproducing kernel Hilbert space ( RKHS ) and then perform a weighted pooling operation , with weights given by the transport plan between the set and a trainable reference . To gain scalability , we then obtain a finite-dimensional embedding by using kernel approximation techniques ( Williams & Seeger , 2001 ) . The motivation for using kernels is to provide a non-linear transformation of the input features before pooling , whereas optimal transport allows to align the features on a trainable reference with fast algorithms ( Cuturi , 2013 ) . Such combination provides us with a theoretically grounded , fixed-size embedding that can be learned either without any label , or with supervision . Our embedding can indeed become adaptive to 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 chromatin profiles where we achieve state-of-the-art results . We also show promising results in natural language processing tasks , where our method outperforms strong baselines . Contributions . In summary , our contribution is three-fold . We propose a new method to embed sets of features of varying sizes to fixed size representations that are well adapted to 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 . 2 RELATED WORK . Kernel methods for sets and OT-based kernels . The kernel associated with our embedding belongs to the family of match kernels ( Lyu , 2004 ; Tolias et al. , 2013 ) , which compare all pairs of features between two sets via a similarity function . 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 hyperparameters , these kernels yet can not be trained end-to-end , as opposed to our embedding that relies on a trainable reference . Efficient and trainable kernel embeddings for biological sequences have also been proposed by Chen et al . ( 2019a ; b ) . Our work can be seen as an extension of these earlier approaches by using optimal transport rather than mean pooling for aggregating local features , which performs 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 a 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-decoder architecture 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-hoc 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 a set in a model , with trainable parameters . Tsai et al . ( 2019 ) provide a view of Transformer ’ s attention via kernel methods , yet in a very different fashion where attention is cast as kernel smoothing and not as a kernel embedding . 3 PROPOSED EMBEDDING . 3.1 PRELIMINARIES . We handle sets of features in Rd and consider sets x living in X = { x | x = { x1 , . . . , xn } such that x1 , . . . , xn ∈ Rd for some n ≥ 1 } . Elements of X are typically vector representations of local data structures , such as k-mers for sequences , patches for natural images , or words for sentences . The size of x denoted by n may vary , which is not an issue since the methods we introduce may take a sequence of any size as input , while providing a 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 . Optimal transport . Our pooling mechanism will be based on the transport plan between x and x′ seen as weighted point clouds or discrete measures , which is a by-product of the optimal transport problem ( Villani , 2008 ; Peyré & Cuturi , 2019 ) . OT has indeed been widely used in alignment problems ( Grave et al. , 2019 ) . Throughout the paper , we will refer to the Kantorovich relaxation of OT with entropic regularization , detailed for example in ( Peyré & Cuturi , 2019 ) . Let a in ∆n ( probability simplex ) and b in ∆n ′ be the weights of the discrete measures ∑ i aiδxi and ∑ j bjδx′j with respective locations x and x′ , where δx is the Dirac at position x . Let C in Rn×n ′ be a matrix representing the pairwise costs for aligning the elements of x and x′ . The entropic regularized Kantorovich relaxation of OT from x to x′ is min P∈U ( a , b ) ∑ ij CijPij − εH ( P ) , ( 1 ) where H ( P ) = − ∑ ij Pij ( log ( Pij ) − 1 ) is the entropic regularization with parameter ε , which controls the sparsity of P , and U is the space of admissible couplings between a and b : U ( a , b ) = { P ∈ Rn×n ′ + : P1n = a and P > 1n′ = b } . The problem is typically solved by using a matrix scaling procedure known as Sinkhorn ’ s algorithm ( Sinkhorn & Knopp , 1967 ; Cuturi , 2013 ) . In practice , a and b are uniform measures since we consider the mass to 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 ( Schölkopf & Smola , 2001 ) map data living in a space X to a reproducing kernel Hilbert space H , associated to a positive definite kernel K through a mapping function ϕ : X → H , such that K ( x , x′ ) = 〈ϕ ( x ) , ϕ ( x′ ) 〉H . Even though ϕ ( x ) may be infinitedimensional , classical kernel approximation techniques ( Williams & Seeger , 2001 ) provide finitedimensional embeddings ψ ( x ) in Rk such that K ( x , x′ ) ≈ 〈ψ ( x ) , ψ ( x′ ) 〉 . Our embedding for sets relies in part on kernel method principles and on such a finite-dimensional approximation . 3.2 OPTIMAL TRANSPORT EMBEDDING AND ASSOCIATED KERNEL . We now present the OTKE , an embedding and pooling layer which aggregates a variable-size set or sequence of features into a fixed-size embedding . We start with an infinite-dimensional variant living in a RKHS , before introducing the finite-dimensional embedding that we use in practice . Infinite-dimensional embedding in RKHS . Given a set x and a ( learned ) reference z in X with p elements , we consider an embedding Φz ( x ) which performs the following operations : ( i ) initial embedding of the elements of x and z to a RKHS H ; ( ii ) alignment of the elements of x to the elements of z via optimal transport ; ( iii ) weighted linear pooling of the elements x into p bins , producing an embedding Φz ( x ) inHp , which is illustrated in Figure 1 . 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 interactions may benefit from pooling to 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 a pooling cell , where the elements of the input set are aggregated through a weighted sum . The weights simply reflect the similarity between the vectors of the input set and the current vector in the reference . Importantly , using a reference set enables to reduce the size of the “ attention matrix ” from quadratic to linear in the length of the input sequence . • Computing similarity weights via optimal transport . A computationally efficient similarity score between two elements is their dot-product ( Vaswani et al. , 2017 ) . In this paper , we rather consider that elements of the input set should be pooled together if they align well with the same part of the reference . Alignment scores can efficiently be obtained by computing the transport plan between the input and the reference sets : Sinkhorn ’ 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 = ( x1 , . . . , xn ) in X be an input set of feature vectors and z = ( z1 , . . . , zp ) in X be a reference set with p elements . Let κ be a positive definite kernel , e.g. , Gaussian kernel , with RKHSH and ϕ : Rd → H , its associated kernel embedding . Let κ be the n× p matrix which carries the comparisons κ ( xi , zj ) , before alignment . Then , the transport plan between x and z , denoted by the n × p matrix P ( x , z ) , is defined as the unique solution of ( 1 ) when choosing the cost C = −κ , and our embedding is defined as Φz ( x ) : = √ p× ( n∑ i=1 P ( x , z ) i1ϕ ( xi ) , . . . , n∑ i=1 P ( x , z ) ipϕ ( xi ) ) = √ p×P ( x , z ) > ϕ ( x ) , where ϕ ( x ) : = [ ϕ ( x1 ) , . . . , ϕ ( xn ) ] > . Interestingly , it is easy to show that our embedding Φz ( x ) is associated to the positive definite kernel Kz ( x , x ′ ) : = n∑ i , i′=1 Pz ( x , x ′ ) ii′κ ( xi , x ′ i′ ) = 〈Φz ( x ) , Φz ( x′ ) 〉 , ( 2 ) with Pz ( x , x′ ) : = p × P ( x , z ) P ( x′ , z ) > . This is a weighted match kernel , with weights given by optimal transport in H. The notion of pooling in the RKHS H of κ arises naturally if p ≤ 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 P ( x , z ) . This process is illustrated in Figure 1 . We acknowledge here the concurrent work by Kolouri et al . ( 2021 ) , where a 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. , Peyré & Cuturi , 2019 ) , Pz ( x , x′ ) is a valid transport plan and , empirically , a rough approximation of P ( x , x′ ) . Kz can therefore be seen as a surrogate of a well-known kernel ( Rubner et al. , 2000 ) , defined as KOT ( x , x ′ ) : = n∑ i , i′=1 P ( x , x′ ) ii′κ ( xi , x ′ i′ ) . ( 3 ) When the entropic regularization ε is equal to 0 , KOT is equivalent to the 2-Wasserstein distance W2 ( x , x ′ ) with ground metric dκ induced by kernel κ. KOT is generally not positive definite ( see Peyré & 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 Kz ) . Now , we show the relationship between this kernel and our kernel Kz , which is proved in Appendix B.1 . Lemma 3.1 ( Relation between P ( x , x′ ) and Pz ( x , x′ ) when ε = 0 ) . For any x , x′ and z in X with lengths n , n′ and p , by denoting W z2 ( x , x ′ ) : = 〈Pz ( x , x′ ) , d2κ ( x , x′ ) 〉 1/2 we have |W2 ( x , x′ ) −W z2 ( x , x′ ) | ≤ 2 min ( W2 ( x , z ) , W2 ( x′ , z ) ) . ( 4 ) This lemma shows that the distanceW z2 resulting fromKz is related to the Wasserstein distanceW2 ; yet , this relation should not be interpreted as an approximation error as our goal is not to approximate W2 , but rather to derive a trainable embedding Φz ( 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 z2 is equivalent to a distance introduced by Wang et al . ( 2013 ) , whose properties are further studied by Moosmüller & Cloninger ( 2020 ) . A major difference is that W z2 crucially relies on Sinkhorn ’ s algorithm so that the references can be learned end-to-end , as explained below .

Summary: The 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.