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: This paper studies the problem of information retrieval that searches similar items to a given query. The paper introduces a new tokenizer and claims that Seq2Tok can convert a continuous sequence (e.g. audio) to discrete tokens. The purpose of the conversion is to enable easier retrieval. The paper conducts extensive experiments and claims that Seq2Tok achieves consistent performance on audio retrieval tasks.