Neural Embeddings for Nearest Neighbor Search Under Edit Distance

The edit distance between two sequences is an important metric with many applications . The drawback , however , is the high computational cost of many basic problems involving this notion , such as the nearest neighbor search . A natural approach to overcoming this issue is to embed the sequences into a vector space such that the geometric distance in the target space approximates the edit distance in the original space . However , the known edit distance embedding algorithms , such as Chakraborty et al . ( 2016 ) , construct embeddings that are data-independent , i.e. , do not exploit any structure of embedded sets of strings . In this paper we propose an alternative approach , which learns the embedding function according to the data distribution . Our experiments show that the new algorithm has much better empirical performance than prior data-independent methods . 1 Introduction . The edit distance is a popular metric for computing the dissimilarity between two strings . Given strings x and y , the edit distance dE ( x , y ) between them is equal to the minimum number of character insertions , deletions and substitutions needed to transform x into y . The metric and its many variants have been defined and studied since the 1960s ( Levenshtein , 1966 ; Wagner and Fischer , 1974 ) . It has found many applications in computational biology ( to measure the similarity between DNA/RNA sequences ) , natural language processing , information theory and other fields . From the algorithmic perspective , the main drawback of edit distance is its computational complexity . The best known algorithm for computing the distance between two strings of length n takes time that is roughly quadratic in n ( Masek and Paterson , 1980 ) ; there is also evidence that this runtime is essentially optimal ( Backurs and Indyk , 2015 ) . Furthermore , many applications require finding the closest string in a large collection of ( say ) N strings . Thus , even if we could estimate the edit distance in time O ( n ) , computing the distances to all strings in the collection would take O ( Nn ) time , which is infeasible for large data sets . A natural approach to resolving both issues relies on low-distortion embeddings . The idea is to construct a mapping f that maps strings into a d-dimensional space < d , such that the edit distance dE ( x , y ) is approximated by the distance ‖f ( x ) − f ( y ) ‖ between the images of the strings , up to some factor D called distortion . If the mapping f can be evaluated efficiently , this immediately implies an efficient way to compute the edit distance . Furthermore , embedding strings into space equipped with a norm ‖ · ‖ makes it possible to use nearest neighbor search algorithms designed for that norm . For many normed spaces , e.g . Minkowski norms ` dp , many efficient approximate nearest neighbor search algorithms are available . This dual goal has motivated a long line of research focused on designing efficient low-distortion embeddings of variants of edit distance into normed spaces ( Muthukrishnan and Sahinalp , 2000 ; Cormode et al. , 2001 ; Ostrovsky and Rabani , 2005 ; Batu et al. , 2006 ; Jowhari , 2012 ; Andoni and Onak , 2012 ; Chakraborty et al. , 2016 ) . The method given in the last paper ( Chakraborty et al. , 2016 ) was shown to yield good results in practice , in the context of large scale similarity joins , i.e. , identifying all pairs of similar sequences ( Zhang and Zhang , 2017 ) . A weakness of the aforementioned algorithms is that they are data-independent , i.e. , they do not exploit any structure of embedded sets of strings . In contrast , over the last decade , a large body of research on “ learning to hash ” showed that data-dependent methods often yield ( binary ) embeddings with a significantly lower distortion than their data-independent counterparts ( Wang et al. , 2018 ) . Such methods , however , are typically tailored to obtaining binary representations of vectors that already live in a geometric ( typically Hilbert ) space . It is thus natural to investigate whether similar benefits of data-dependent methods can be obtained for edit distance embeddings as well . In this paper we investigate whether the “ learning to hash ” approach can yield lower distortion embeddings than those obtained using data-independent methods . Since our goal is to obtain distance-preserving mappings of sequences , we use Recurrent Neural Networks as a model implementing those mappings . Specifically , we leverage a two-layer GRU network , depicted in Figure 2 and described in more details in Section 4 . The network inputs a sequence of characters ( e.g. , A , G , T or C for DNA sequences ) , and outputs a vector of symbols , such that the distance between the vectors approximates the edit distance between the original sequences . The standard loss function for optimizing the network is the triplet loss ( see Figure 1 ) , which has two problems in our scenario . The first problem is distance scaling : triplet loss only controls the relative distance ordering among data points after embedding , rather than the absolute embedding distance between any two data points . This prevents us from getting accurate edit distance estimation from the embedding . To alleviate this issue , we designed a warm-start phase ( Phase 1 ) before optimizing triplet loss ( Phase 2 ) , which initializes the embedding so that the discrepancy between the original edit distance and the embedding distance is minimized . The second problem is worst case performance of our model . Due to the nature of learning , the network we learned fits the training data better , and may not do well on unseen hard examples . To alleviate this issue , we append a back up phase ( Phase 3 ) , which deals with hard cases . We emphasize that all the three phases are optimizing the same network . We run Adam algorithm for each phase until convergence , and then use its final state as the initial state for the next phase . Although in principle one may use a single phase with combination of different loss functions to train our model , this requires additional hyperparameter adjustment for the combining ratios . Our current three phase training does not need additional hyperparameters , and therefore can easily fit different scenarios . The Phase 3 design is non-trivial . Our starting observation is that Recurrent Neural Networks that we employ are not “ compatible ” with the data-independent mappings of Chakraborty et al . ( 2016 ) ( which we call CGK from now on ) . This is because CGK can generate different number of copies for every input character , which is a discrete decision that is hard to be learned by neural networks . We show , however , that this “ extra power ” of the algorithm is not necessary . In particular , we introduce a “ truncated ” version of the CGK algorithm ( referred to as CGK ’ ) , that outputs only 1 or 2 symbols per round . We then show that CGK ’ has the same guarantees as the original CGK algorithm . Based on this observation , we design our Phase 3 to be an optimized version of CGK ’ . Specifically , since each round there will be either 1 or 2 symbols in CGK ’ , we can directly compare the two cases and use the difference as approximation to the gradients to train our model . Empirically , Phase 3 does not affect the recall on ordinary datasets much , but greatly improves performance on hard instances . To evaluate the embeddings , we employ them for nearest neighbor search in a collection of sequences , and compare their performance to CGK . As in Zhang and Zhang ( 2017 ) we use three data sets : Gen50ks ( containing sequences from the human genome project ) , UniRef ( protein sequence dataset from the UniProt project ) and Trec ( containing textual information from Medline , an online medical information database ) . After learning the embeddings , we generate a set of queries , and for each query we identify C sequences in the data set ( called candidates ) that are closest to the query in the embedded space . We then measure “ top-K recall ” , i.e. , the fraction of the actual top K nearest neighbors ( computed using the exact edit distance ) returned as one of the C candidates . Our results demonstrate that , over a wide range of parameters K and C , our learned embedding consistently achieves higher recalls than the algorithm of CGK , often by a factor of 3 or more . Similarly , our method often obtains identical recall to that of CGK while using an order of magnitude fewer candidates . This means that in a typical nearest neighbor search “ filtering ” pipeline , where the candidates are quickly identified in the embedded space and then the exact edit distance is computed between the candidates and queries to identify the final answers , our method reduces the number of exact edit distance computations by up to one order of magnitude . Moreover , our learned embedding also has much more organized structure than CGK , as shown in Appendix A.4 . 2 Related work . 2.1 Edit Distance Computing . Exact and approximate algorithms for edit distance computation have been broadly studied , see Navarro ( 2001 ) for an overview . For exact algorithms , Wagner and Fischer ( 1974 ) gave O ( n2 ) -time algorithm using dynamic programming . This was later improved to O ( n2/ log n ) by Masek and Paterson ( 1980 ) . Faster approximation algorithms were studied as well , see e.g. , Landau et al . ( 1998 ) ; Batu et al . ( 2003 ) ; Bar-Yossef et al . ( 2004 ) ; Batu et al . ( 2006 ) ; Andoni et al . ( 2010 ) ; Andoni and Onak ( 2012 ) . The current state of the art consists of a near-linear algorithm of Andoni et al . ( 2010 ) guaranteeing ( log n ) c-approximation ( with a trade-off between the running time and the exponent c ) , and a constant factor approximation O ( n2−2/7 ) -time algorithm of Chakraborty et al . ( 2018 ) . To the best of our knowledge , neither algorithm has been implemented . 2.2 Edit Distance Embedding . As mentioned in the introduction , there has been a long line of research focused on designing embeddings of various variants of the edit distance into geometric space . For the edit distance as defined in the introduction , the best known algorithm is due to Ostrovsky and Rabani ( 2005 ) , who designed an embedding into the Hamming space with distortion 2O ( √ logn log logn ) . This algorithm was quite complex , and to the best of our knowledge has never been implemented . Recently , Chakraborty et al . ( 2016 ) gave a substantially different embedding with distortion O ( k ) where k is the edit distance between the original strings . The empirical evaluation of the algorithm , given in Zhang and Zhang ( 2017 ) , demonstrated that the embedding yields good results in practice , and that the empirical distortion of this embedding is typically much lower than the worst case bound . Zhang and Zhang ( 2019 ) also proposed a novel string partition based algorithm for edit similarity join , without generating embeddings . Gomez-Bigorda et al . ( 2017 ) presented a novel approach for word embedding in Levenshtein space to improve downstream handwritten word spotting . Their method is similar to our Phase 1 , but our model deals with longer sequences instead of words , and also adopts triplet loss training to capture sequence similarity ( Phase 2 ) and CGK based algorithm to learn sequence alignment ( Phase 3 ) .

Summary: In this paper, the authors propose an approach for learning embeddings for strings using a three-phase approach. Each phase uses a different neural network to learn embeddings for strings such that the distance between strings in embedding space approximates the edit distance between the strings. A modest set of empirical results suggests that the proposed approach outperforms hand-crafted embeddings.