Seq2Tok: Deep Sequence Tokenizer for Retrieval

1 INTRODUCTION . Sequence Retrieval aims at retrieving sequences similar to a query sequence , with the constraint that an ordered alignment exists between the query and the target sequence . The examples include speech search , where the order of constituent units , such as phonemes , syllables or words , remains same ; and music search – query by humming or query by example – where the order of constituent units , such as relative notes or phrases , remains same . Apart from audio , the problem also extends to tasks such as handwritten word search and gesture search . One can define similarity metrics over the sequences using Dynamic Time Warping ( DTW ) . But those can be very inefficient if the sequences are continuous valued and have high sampling rates , as in audio . Also , the high variability of query sequences makes the hand-made metrics cumbersome and ineffective . A more effective way of sequence retrieval is by mapping them to discrete tokens . For instance , for linguistic content based speech retrieval , audio is mapped to discrete linguistic tokens , such as phonemes , syllables or words , and search is performed over these tokens . For query-by-humming based music search , audio is mapped to discrete melody-related tokens , such as notes . For query-by-example based music search , audio is mapped to rule-based landmarks in spectrograms [ Wang ( 2003 ) ] . Hence , each problem uses a domain-specific hand-made mapping to tokens defined by a domain expert . In this paper , we propose an audio tokenizer which maps the raw audio feature sequence to a sequence of tokens . The tokenizer is trained from pairs of query and target sequences , without any expert-defined tokens . The proposed algorithm makes use of the fact that relevant landmarks appear in the same order in both the query and the target sequences . This reduces the dimensionality of data and preserves the information relevant for effective retrieval . Existing text search methods , namely , locally sensitive hashing , edit distance , etc. , can then be used to query over these tokens . 2 RELATED WORK . Sequence Labelling . With expert-defined tokens , various methods are popularly used for mapping sequences to tokens . In conventional methods , Hidden Markov Models [ Rabiner & Juang ( 1986 ) ] and Conditional Random Fields [ Lafferty et al . ( 2001 ) ] have been popularly used for sequence labelling . These methods involve a significant amount of domain knowledge and many assumptions to make tractable models , which are avoided by End-to-End learning models such as Recurrent Neural Networks ( RNNs ) using Connectionist Temporal Classification framework [ Graves et al . ( 2006 ) ] . Sequence labeling can be used for sequence retrieval by converting the sequences to tokens , which are easy to search over . But this approach inevitably depends upon expert-defined tokens . Unsupervised Speech Recognition . Modern Representation learning techniques such as Contrastive Predictive Coding ( CPC ) [ van den Oord et al . ( 2018 ) ] and Autoregressive Predictive Coding ( APC ) [ Chung & Glass ( 2020 ) ] are able to generate state-of-the-art continuous representations which encode the slowly varying phoneme features in raw speech . Wav2Vec [ Schneider et al . ( 2019 ) ] generated continuous representations with an encoder pretrained to distinguish future timesteps from a set of distractors sampled from a proposal distribution by optimizing a CPC-based contrastive loss . After pretraining the models , the representations are mapped to phoneme tokens via minimization of the Connectionist Temporal Classification ( CTC ) Loss [ Graves et al . ( 2006 ) ] . With inspiration from VQ-VAE [ van den Oord et al . ( 2017 ) ] which marked the beginning of generation of discrete representations , VQ-Wav2Vec [ Baevski et al . ( 2019 ) ] discretized the representations generated by Wav2Vec [ Schneider et al . ( 2019 ) ] with either a K-Means Vector Quantizer [ Baevski et al . ( 2019 ) ] or a Gumble-Softmax based Vector Quantizer [ Baevski et al . ( 2019 ) ] to generate quantized speech representation used to pretrain a BERT [ Devlin et al . ( 2018 ) ] . Wav2Vec 2.0 [ Baevski et al . ( 2020 ) ] discretized and masked the representations to solve a CPC-based contrastive task over the quantized representations which are learned jointly . VQ-APC [ Chung et al . ( 2020 ) ] discretized the representations to solve APC task over the quantized representations which are learned jointly . Note , the main motivation to discretize raw audio in a latent space for all aforementioned works are to explicitly control information content encoded in the representations rather than achieve sequence tokenization . The codes or tokens generated by the models aren ’ t constrained to be interpretable and hence initialised in large numbers . Performing CTC fine-tuning in essence groups codes in the vector quantizer module of each model together and achieves codebook compression to the number of phonemes or linguistic units . The discretization of representations do result in performance increase on various downstream tasks , with a pretrained Wav2Vec 2.0 [ Baevski et al . ( 2020 ) ] achieving state of the art results on speech recognition by being fine-tuned on only 10 minutes of transcribed audio . SeqRQ-AE [ Liu et al . ( 2019 ) ] is the closest to our work which uses a VQ-VAE [ van den Oord et al . ( 2017 ) ] to estimate representations for each phoneme in speech . The VQ-VAE has a codebook containing the same number of members as the number of phonemes . The codebook is trained via phonetic clustering , temporal segmentation and CTC finetuning on a small amount of transcribed audio . Neural Audio Fingerprinting . Now Playing [ Arcas et al . ( 2017 ) ] and [ Chang et al . ( 2020 ) ] use a Neural Network Fingerprinter ( NNFP ) module outputting representations which are efficient for search in query-by-example tasks where the difference between query and actual song is pretty minute in comparison to humming where only the melody is sung . Now Playing [ Arcas et al . ( 2017 ) ] trains representations by optimizing the Triplet loss [ Schroff et al . ( 2015 ) ] and [ Chang et al . ( 2020 ) ] trains representations by simulating the Maximum Inner Product Search on mini-batches of representations . Cross Domain Alignment . Given a pair of semantically similar inputs for training , tasks such as visual question answering ( text and image ) and machine translation ( text ) involve learning of an alignment . The alignment here is not ordered and the inputs may be from different modalities . Attention models have been used to find alignment between output entities and input regions [ Yao et al . ( 2018 ) ] . [ Chen et al . ( 2020 ) ] use Gromov-Wasserstein distance between output and input entities to match them . However , there is no notion of tokens there , rather the salient entities in the input are represented as vectors in a graph . Graph Matching . Generation of representations for graphs can be done by Graph Neural Networks [ Gori et al . ( 2005 ) ] . Graph matching is used to find similarity of structured graphs [ Li et al . ( 2019 ) ] . However , they perform the matching jointly on the pair of inputs , rather than representing each input independently . This makes them unsuitable for the search problem at hand due to large runtime complexity . The distance metrics used for graph matching are based on edit distance [ Li et al . ( 2019 ) ] and Wasserstein distance [ Chen et al . ( 2020 ) ] . 3 METHODOLOGY . 3.1 PROBLEM STATEMENT . Given a pair of semantically similar feature sequences ( Xi , Xj ) of length Ti and Tj respectively and feature sequence Xk of length Tk with no similarity to either of the feature sequences Xi or Xj , we aim to map sequence Xi to sequence of tokens Ti of length T ′i ≤ Ti , sequence Xj to sequence of tokens Tj of length T ′j ≤ Tj , and sequence Xk to sequence of tokens Tk of length T ′k ≤ Tk where all the tokens belong to a finite alphabet A such that , the edit distance between Ti and Tj is minimum while that between Tk and Ti as well as Tk and Tj is maximum . 3.2 MODEL ARCHITECTURE . Wav2Tok is comprised of an encoder , which , given input x ∈ Rn , outputs a representation z ∈ Rm , and a tokenizer , which , given the encoder output z , outputs a token τ which belongs to a finite K-element alphabet A = { 1 , ... , K } . Given , a temporal sequence X = [ x1 , ... , xT ] of length T where xt is the feature vector at time t. Our model performs frame-wise labelling of X to output a sequence of tokens T = [ τ1 , ... , τT ] where τt ∈ A is the token assigned to input xt . The encoder f : X 7→ Z encodes xt in a hyper-spherical space Z using a BiLSTM network [ Schuster & Paliwal ( 1997 ) ] followed by an L-2 normalization layer . The BiLSTM allows zt to summarise information in both directions , and hence , encode surrounding context . The tokenizer g : Z 7→ T is a K-means vector quantizer which vector quantizes encoder output zt with K quantizers or cluster centroids wtokk ∈ Z ; k = 1 , ... , K and outputs token τt as , τt = argmaxk { zt · wtokk } . The centroids are estimated prior to training via spherical K-means clustering in the spherical space Z generated by the freshly initialised encoder . The centroids are further tuned with training iterations . 3.3 MODEL TRAINING . Given a batch of tuples of similar sequences ( Xi , Xi′ ) where Xi = [ xit ∈ Rn ; t = 1 , ... , Ti ] and Xi′ = [ xi ′ t′ ∈ Rn ; t′ = 1 , ... , Ti′ ] . The encoder f maps these sequences to representations Zi = f ( Xi ) and Zi′ = f ( Xi′ ) . The tokenizer g then labels Zi framewise to token sequence Ti and Zi′ framewise to token sequence Ti′ . For training , the tuple ( Zi , Zi′ ) are concatenated to form the composed sequence Zii′ = [ Zi ; Zi′ ] of length Tii′ = Ti + Ti′ . The tuple ( Ti , Ti′ ) are also concatenated to form the token sequence Tii′ = [ Ti ; Ti′ ] of length Tii′ = Ti+Ti′ corresponding to sequenceZii′ . Given Tii′ = [ τ ii ′ 1 , ... , τ ii′ Tii′ ] and Zii′ = [ zii ′ 1 , ... , zii ′ Tii′ ] , we form aK ′ ≤ K length list of prototypes Pii′ = { pτ1ii′ , ... , p τK′ ii′ } where { τ1 , ... , τK′ } ∈ A are allK ′ unique tokens that occur in Tii′ and pτk′ii′ is a prototype for token τk′ ∈ A generated from all vectors in Zii′ mapping to τk′ as , p̃τk′ii′ = 1∑Tii′ t=1 1τ ii′t =τk′ Tii′∑ t=1 zii ′ t 1τ ii′t =τk′ ( 1 ) pτk′ii′ = p̃τk′ii′ ||p̃τk′ii′ || ( 2 ) Let Pii′ and Pjj′ be lists of prototypes generated given tuples ( Xi , Xi′ ) and ( Xj , Xj′ ) , respectively . We sample tuples of prototypes ( pτii′ , pτ ′ jj′ ) where pτii′ is sampled from Pii′ and pτ ′ jj′ is sampled from Pjj′ . Approach A . Given ( pτii′ , pτ ′ jj′ ) , we optimize the following loss function L11 to increase the cosine similarity between pτii′ and pτ ′ jj′ to 1 if τ = τ ′ and to decrease the cosine similarity to 0 if τ 6= τ ′ , L11 ( pτii′ , pτ ′ jj′ , t1 ) = BinXEnt ( p τ ii′ · pτ ′ jj′ , t1 ) ( 3 ) BinXEnt ( pτii′ · pτ ′ jj′ , t1 ) = t1 · log ( pτii′ · pτ ′ jj′ ) + ( 1− t1 ) · log ( 1− pτii′ · pτ ′ jj′ ) ( 4 ) where BinXEnt ( . ) is the binary cross-entropy function , pτii′ · pτ ′ jj′ gives a cosine similarity score as both the vectors are L-2 normalized , and target t1 is 0 if τ 6= τ ′ else 1 . We also optimize the following loss function L12 to increase the cosine similarity of pτii′ or pτ ′ jj′ with respect to the k th quantizer wtokk in our tokenizer g : Z 7→ T to 1 if τ or τ ′ = k repectively and to decrease the same to 0 if τ or τ ′ 6= k , L12 ( pτii′ , pτ ′ jj′ ) = K∑ k=1 ( BinXEnt ( pτii′ · wtokk , t1k ) +BinXEnt ( pτ ′ jj′ · wtokk , t2k ) ) ( 5 ) where target t1k is 0 if τ 6= k else 1 and target t2k is 0 if τ ′ 6= k else 1 . The constrastive loss function L1 for Approach A thus formed with the union of the above two loss functions is given as , L1 ( pτii′ , pτ ′ jj′ , t1 ) = L11 ( pτii′ , pτ ′ jj′ , t1 ) + L12 ( pτii′ , pτ ′ jj′ ) ( 6 ) Approach B . Given ( pτii′ , pτ ′ jj′ ) , a relation network R : Rm 7→ R generates relation score R ( pτii′ , pτ ′ jj′ ) in the range [ 0 , 1 ] which is closer to 0 if τ 6= τ ′ and closer to 1 if τ = τ ′ . The relation score acts as a non-linear similarity metric and is given as , R ( pτii′ , pτ ′ jj′ ) = σ ( WTr |pτii′ − pτ ′ jj′ | ) ( 7 ) where σ ( . ) is the sigmoid function , |.| is the element-wise absolute value function , andWr ∈ Rm is the weight matrix of the relation network R : Rm 7→ R. Given R ( pτii′ , pτ ′ jj′ ) , we optimize the following loss function L21 to increase the relation score between pτii′ and pτ ′ jj′ to 1 if τ = τ ′ and to decrease the relation score to 0 if τ 6= τ ′ , L21 ( pτii′ , pτ ′ jj′ , t2 ) = BinXEnt ( R ( pτii′ , pτ ′ jj′ ) , t2 ) ( 8 ) where target t2 is 0 if τ 6= τ ′ else 1 . We also optimize the following loss function L22 to increase the relation score of pτii′ or pτ ′ jj′ with respect to the k th quantizer wtokk in our tokenizer g : Z 7→ T to 1 if τ or τ ′ = k respectively and to decrease the same to 0 if τ or τ ′ 6= k , L22 ( pτii′ , pτ ′ jj′ ) = K∑ k=1 ( BinXEnt ( R ( pτii′ , wtokk ) , t1k ) +BinXEnt ( R ( pτ ′ jj′ , w tok k ) , t2k ) ) ( 9 ) where target t3k is 0 if τ 6= k else 1 and target t4k is 0 if τ ′ 6= k else 1 . The contrastive loss function L2 for Approach B thus formed with the union of the above two loss functions is given as , L2 ( pτii′ , pτ ′ jj′ , t2 ) = L22 ( pτii′ , pτ ′ jj′ , t2 ) + L22 ( pτii′ , pτ ′ jj′ ) ( 10 ) A spherical K-means vector quantization loss as inspired from the K-means vector quantization loss in VQ-Wav2Vec [ Baevski et al . ( 2019 ) ] may be iterated to add more movement to the centroids and the representations . Given prototype pτii′ for token τ ∈ A generated from sequence of representations Zii′ = [ zii ′ 1 , ... , zii ′ Tii′ ] and sequence of tokens Tii′ = [ τ ii ′ 1 , ... , τ ii′ Tii′ ] , our vector quantization loss is given as , LSK ( pτii′ ) = 1∑Tii′ t=1 1τ ii′t =τ Tii′∑ t=1 ( 2− ( sg ( zii ′ t ) · wtokτ + γ · zii ′ t · sg ( wtokτ ) ) ) 1τ ii′t =τ ( 11 ) where γ is a positive constant , and sg ( . ) is the stop-gradient operator . The stop-gradient operator converts the parameter input to a constant during gradient computation . Increasing sg ( zii ′ t ) · wtokτ pushes wtokτ towards zii ′ t keeping zii ′ t static and increasing zii ′ t · sg ( wtokτ ) pushes zii ′ t towards wtokτ keeping wtokτ static . With γ set to 0 , iterating the loss only affects wtokτ . Given ( pτii′ , pτjj′ ) , the final loss function L which we minimize to optimize our model parameters is given as , L ( pτii′ , pτ ′ jj′ , ta ) = La ( pτii′ , pτ ′ jj′ , ta ) + α · ( LSK ( pτii′ ) + LSK ( pτ ′ jj′ ) ) ( 12 ) where α is a positive constant , ta is the target , and a = 1 corresponds to Approach A while a = 2 corresponds to Approach B .

Summary: The paper presents an algorithm to do sequence matching using a tokenizations of the original sequences, and evaluates it in three different settings. The test environments are non-standard setups of public databases, with a closed vocabulary of sequences, i.e. training, validation, and testing sets contain the same sequences.