Neural Distance Embeddings for Biological Sequences

1 Introduction . Over the course of evolution , biological sequences constantly mutate and a large part of biological research is based on the analysis of these mutations . Biologists have developed accurate statistical models to estimate the evolutionary distance between pairs of sequences based on their edit distance D ( s1 , s2 ) : the minimum number of ( weighted ) insertions , deletions or substitutions required to transform a string s1 into another string s2 . However , the computation of this edit distance kernel D with traditional methods is bound to a quadratic complexity and hardly parallelizable , making its computation a bottleneck in large scale analyses , such as microbiome studies [ 1 , 2 , 3 ] . Furthermore , the accurate computation of similarities among multiple sequences , at the foundation of critical tasks such as hierarchical clustering and multiple sequence alignment , is computationally intractable even for relatively small numbers of sequences . Problems that in other spaces are relatively simple become combinatorially hard in the space of sequences defined by the edit distance . For example , finding the Steiner string , a classical problem in bioinformatics that can be thought of as computing the geometric median in the space of sequences , is NP-complete . Classical algorithms and heuristics [ 4 , 5 , 6 , 7 ] widely used in bioinformatics for these tasks are data-independent and , therefore , can not exploit the low-dimensional manifold assumption that characterises real-world data [ 8 , 9 , 10 ] . Leveraging the available data to produce efficient and data- ∗Correspondence to gcorso @ mit.edu 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . dependent heuristics and representations would greatly accelerate large-scale analyses that are critical to biological research . While the number of available biological sequences has grown exponentially over the past decades , machine learning approaches to problems related to string matching [ 11 , 12 ] have not been adopted widely in bioinformatics due to their limitation in accuracy and speed . In contrast to most tasks in computer vision and NLP , string matching problems are typically formalised as combinatorial optimisation problems . These discrete formulations do not fit well with the current deep learning approaches . Moreover , representation learning methods based on Euclidean spaces struggle to capture the hierarchical structure that characterises real-world biological datasets due to evolution . Finally , common self-supervised learning approaches , very successful in NLP , are less effective in the biological context where relations tend to be between sequences rather than between bases [ 13 ] . In this work , we present Neural Distance Embeddings ( NeuroSEED ) , a general framework to produce representations for biological sequences where the distance in the embedding space is correlated with the evolutionary distance D between sequences . NeuroSEED provides fast approximations of the distance kernel D , low-dimensional representations for biological sequences , tractable analysis of the relationship between multiple sequences in the embedding geometry and a way to decode novel sequences . Firstly , we reformulate several existing approaches into NeuroSEED highlighting their contributions and limitations . Then , we examine the task of embedding sequences to preserve the edit distance that is the basis of the framework . This analysis reveals the importance of data-dependent approaches and of using a geometry that matches the underlying data distribution well . The hyperbolic space is able to capture the implicit hierarchical structure given by biological evolution and provides an average 22 % reduction in embedding RMSE against the best competing geometry . We show the potential of the framework and its wide applicability by analysing two fundamental tasks in bioinformatics involving the relations between multiple sequences : hierarchical clustering and multiple sequence alignment . For both tasks , unsupervised approaches using NeuroSEED encoders are able to match the accuracy of common heuristics while being orders of magnitude faster . For hierarchical clustering , we also explore a method based on the continuous relaxation of Dasgupta ’ s cost in the hyperbolic space which provides a 15x runtime reduction at similar quality levels . Finally , for multiple sequence alignment , we devise an original approach based on variational autoencoders that matches the performance of competitive baselines while significantly reducing the runtime complexity . As a summary our contributions are : ( i ) We introduce NeuroSEED , a general framework to map sequences in geometric vector spaces , and reformulate existing approaches into it . ( ii ) We show how the hyperbolic space can bring significant improvements to the data-dependent analysis of biological sequences . ( iii ) We propose several heuristic approaches to classical bioinformatics problems that can be constructed on top of NeuroSEED embeddings and provide significant running time reduction against classical baselines . 2 Bioinformatics tasks . The field of bioinformatics has developed a wide range of algorithms to tackle the classical problems that we explore . Here we present the tasks and briefly mention their motivation and some of the baselines we test . More details are provided in Appendix B. Edit distance approximation In this work , we always deal with the classical edit distance where the same weight is given to every string operation , but all the approaches developed can be applied to any distance function of choice ( which is given as an oracle ) . For example , when using amino acid sequences , one of the different metric variants of the classical substitution matrices such as mPAM250 [ 15 ] would be a good choice . As baseline approximation heuristics , we take k-mer [ 5 ] , which is the most commonly used alignment-free method and represents sequences by the frequency vector of subsequences of a certain length , and FFP [ 16 ] , another alignment-free method which looks at the Jensen-Shannon divergence between distributions of k-mers . Hierarchical clustering ( HC ) Discovering the intrinsic hierarchical structure given by evolutionary history is a critical step of many biological analyses . Hierarchical clustering ( HC ) consists of , given a pairwise distance function , defining a tree with internal points corresponding to clusters and leaves to datapoints . Dasgupta ’ s cost [ 17 ] measures how well the tree generated respects the similarities between datapoints . As baselines we consider classical agglomerative clustering algorithms ( Single [ 18 ] , Complete [ 19 ] and Average Linkage [ 6 ] ) and the recent technique [ 20 ] that uses a continuous relaxation of Dasgupta ’ s cost in the hyperbolic space . Multiple sequence alignment ( MSA ) Aligning three or more sequences is used for the identification of active and binding sites as well as conserved protein structures , but finding its optimal solution is NP-complete . A related task to MSA is the approximation of the Steiner string which minimises the sum of the distances ( consensus error ) to the sequences in a set . Datasets To evaluate the heuristics we chose three datasets containing different portions of the 16S rRNA gene , crucial in microbiome analysis [ 21 ] , one of the most promising applications of our approach . The first , Qiita [ 21 ] , contains more than 6M sequences of up to 152 bp that cover the V4 hyper-variable region . The second , RT988 [ 11 ] , has only 6.7k publicly available sequences of length up to 465 bp covering the V3-V4 regions . Both datasets were generated by Illumina MiSeq [ 22 ] and contain sequences of approximately the same length . Qiita was collected from skin , saliva and faeces samples , while RT988 was from oral plaques . The third dataset is the Greengenes full-length 16S rRNA database [ 23 ] , which contains more than 1M sequences of length between 1,111 to 2,368 . Moreover , we used a dataset of synthetically generated sequences to test the importance of data-dependent approaches . A full description of the data splits for each of the tasks is provided in Appendix B.4 . 3 Neural Distance Embeddings . The underlying idea behind the NeuroSEED framework , represented in Figure 1 , is to map sequences in a continuous space so that the distance between embedded points is correlated to the one between sequences . Given a distribution of sequences and a distance function D between them , any NeuroSEED approach is formed by four main components : an embedding geometry , an encoder model , a decoder model , and a loss function . Embedding geometry The distance function d between the embedded points defines the geometry of the embedding space . While this factor has been mostly ignored by previous work [ 11 , 24 , 25 , 26 , 27 ] , we show that it is critical for this geometry to reflect the relations between the sequences in the domain . In our experiments , we provide a comparison between Euclidean , Manhattan , cosine , squared Euclidean ( referred to as Square ) and hyperbolic distances ( details in Appendix D ) . Encoder model The encoder model fθ maps sequences to points in the embedding space . In this work we test a variety of models as encoder functions : linear layer , MLP , CNN , GRU [ 28 ] and transformer [ 29 ] with local and global attention . The details on how the models are adapted to the sequences are provided in Appendix C. Chen et al . [ 24 ] proposed CSM , an encoder architecture based on the convolution of subsequences . However , as also noted by Koide et al . [ 12 ] , this model does not perform well when various layers are stacked and , due to the interdependence of cells in the dynamic programming routine , it can not be efficiently parallelised on GPU . Decoder model For some tasks it is also useful to decode sequences from the embedding space . This idea , employed in Section 7.2 and novel among the works related to NeuroSEED , enables to apply the framework to a wider set of problems . Loss function The simplest way to train a NeuroSEED encoder is to directly minimise the MSE between the sequences ’ distance and its approximation as the distance between the embeddings : L ( θ , S ) = ∑ s1 , s2∈S ( D ( s1 , s2 ) − α d ( fθ ( s1 ) , fθ ( s2 ) ) ) 2 ( 1 ) where α is a constant or learnable scalar . Depending on the application that the learned embeddings are used for , the MSE loss may be combined or substituted with other loss functions such as the triplet loss for closest string retrieval ( Appendix F ) , the relaxation of Dasgupta ’ s cost for hierarchical clustering ( Section 7.1 ) or the sequence reconstruction loss for multiple sequence alignment ( Section 7.2 ) . There are at least five previous works [ 11 , 24 , 25 , 26 , 27 ] that have used approaches that can be described using the NeuroSEED framework . These methods , summarised in Table 1 , show the potential of approaches based on the idea of NeuroSEED , but share two significant limitations . The first is the lack of analysis of the geometry of the embedding space , which we show to be critical . The second is that the range of tasks is limited to edit distance approximation ( EDA ) and closest string retrieval ( CSR ) . We highlight how this framework has the flexibility to be adapted to significantly more complex tasks involving relations between multiple sequences such as hierarchical clustering and multiple sequence alignment . 4 Related work . Hyperbolic embeddings Hyperbolic geometry is a non-Euclidean geometry with constant negative sectional curvature and is often referred to as a continuous version of a tree for its ability to embed trees with arbitrarily low distortion . The advantages of mapping objects with implicit or explicit hierarchical structure in the hyperbolic space have also been shown in other domains [ 30 , 31 , 32 , 10 ] . In comparison , this work deals with a very different space defined by the edit distance in biological sequences and , unlike most of the previous works , we do not only derive embeddings for a particular set of datapoints , but train an encoder to map arbitrary sequences from the domain in the space . Sequence Distance Embeddings The clear advantage of working in more computationally tractable spaces has motivated significant research in Sequence Distance Embeddings [ 33 ] ( also known as low-distortion embeddings ) approaches to variants of the edit distance [ 34 , 35 , 36 , 37 , 38 , 39 , 40 ] . However , they are all data-independent and have shown weak performance on the ‘ unconstrained ’ edit distance . Hashing and metric learning NeuroSEED also fits well into the wider research on learning to hash [ 41 ] , where the goal is typically to map a high dimensional vector space into a smaller one preserving distances . Another field related to NeuroSEED is metric learning [ 42 , 43 ] , where models are trained to learn embeddings from examples of similar and dissimilar pairs . Locality-sensitive hashing Finally , the work presented is complementary to the line of research in locality-sensitive hashing methods . Researchers have developed a series of heuristics especially for the tasks of sequence clustering [ 44 , 45 ] and local alignment [ 46 ] . These use as subroutines embeddings/features based on alignment-based or alignment-free methods , and therefore , fall into the limitations we highlight in the paper . Future work could analyse how to improve these heuristics with NeuroSEED embeddings .

Summary: This paper presents an encoder-based framework  (called NeuroSEED) to learn the embeddings of DNA sequences in geometric vector spaces that preserve edit distance between them. Learning an efficient representation that captures the underlying hierarchical relationships allows faster applications of downstream applications like Edit distance Approximation (EDA), Closest String Retrieval (CSR), Multiple Sequence Alignment (MSA), and Hierarchical clustering (HA). The paper explores a variety of neural network models for encoding the sequences as well as different geometric vector spaces. The tasks are formulated as unsupervised (with general MSE loss) and supervised (with tailored loss functions).  It reports hyperbolic space to give the state-of-the-art performance for all the downstream tasks and claims a faster compute time (in terms of time complexity) when using these embeddings. The results are reported on two relevant microbiome analysis datasets - Qiita and RT988.