Revisiting Locality Sensitive Hashing for Vocabulary Selection in Fast Neural Machine Translation (2022)

Abstract: Neural machine translation models often contain large target vocabularies. The calculation of logits, softmax and beam search is computationally costly over so many classes. We investigate the use of locality sensitive hashing (LSH) to reduce the number of vocabulary items that must be evaluated and explore the relationship between the hashing algorithm, translation speed and quality. Compared to prior work, our LSH-based solution does not require additional augmentation via word-frequency lists or alignments. We propose a training procedure that produces models, which, when combined with our LSH inference algorithm increase translation speed by up to 87% over the baseline, while maintaining translation quality as measured by BLEU. Apart from just using BLEU, we focus on minimizing search errors compared to the full softmax, a much harsher quality criterion.

## introduction the computation of the output logit, softmax and beam search (the output layer) are some of the most compute-intensive tasks in current neural machine translation (nmt) models, often taking the majority of inference time for many models, on many hardware architectures, especially in deployment settings. this is mainly due to the large vocabulary size relative to other dimensions in the model. methods that reduce the effective vocabulary size can have a major impact on inference speed. vocabulary selection is one such method. however, a known downside of vocabulary selection methods is the risk of search errors if the desired output token ŷt is not a member of the reduced vocabulary v , forcing the beam search to choose a sub-optimal token. even when the impact of such search errors on bleu is minimal, search errors caused by lexical shortlisting degrade human judgements of quality @xcite . figure 1 : search errors vs. decrease in bleu for all experiments in the paper. we vary the two main hyperparameters of our lsh implementation: number of hash functions and the size k of our selected vocabulary subset. 60% of sentences have a search error before we observe a 1 bleu point degradation. the degradation in human judgement is less surprising if we inspect figure 1 which shows that translation quality, as measured by bleu, is resilient to search errors caused by lsh. only when over 60% of translations exhibit search errors is there significant bleu degradation. we examine vocabulary selection using locality sensitive hashing (lsh), and evaluate specifically in the context of neural machine translation. we introduce an lsh-based vocabulary selection algorithm and compatible models such that: foot_0 1. the models have translation quality that is better than or comparable to the baseline model; ## methods of vocabulary selection the output layer for a target vocabulary v , performs the following computations: p(y t |y 1:t-1 , x; θ) = softmax(w h + b) ŷt = argmax p(y t |y 1:t-1 , x; θ), vocabulary selection create a small subset, v ⊂ v . this will reduce the size of weight matrix, w , and bias vector, b, where equation 1 is then replaced with the more efficient equation 2 which uses w and b instead. the aim is now to find the subset, v , such that ŷt = ŷt . depending on the method, vocabulary selection (and therefore construction of v , w and b) can be static or happen dynamically per sentence (or batch), per decoder time step, or even per individual decoder hypothesis. we restrict our overview of the concept of vocabulary selection to the case where the original softmax layer remains largely unmodified except for sub-selection. methods that require complex structural reformulations of the softmax layer during training like hierarchical softmax @xcite , adaptive softmax @xcite or binary code prediction @xcite are outside the scope of this work. in-depth overviews of past and current vocabulary selection methods are provided by @xcite , @xcite , and more recently @xcite . we only repeat concepts that are either common or required to differentiate our work from previous approaches. ## our lsh-based method our work contrasts with prior research on lsh for nmt by @xcite ; @xcite in that: ## experimental setup we train a model with 6-layer transformer @xcite encoder with 6-layer ssru @xcite decoder, trained using marian @xcite and using the same toolkit for inference. this is a strong and realistic model for production mt environments which balances translation quality and efficiency. we use sentencepiece @xcite with 32,000 tokens for all models, shared between both source and target language. we use the faiss @xcite implementation of simhash hash described in section 3.1. we trained with four translation directions (german-english, english-german, spanish-english, and french-english) europarl corpus @xcite , validated on the held out development set from the same corpus ('dev2006') and tested on the out-of-domain new commentary test set ('nc-dev2007'). see we use a two stage training procedure. in the first stage, we train a translation model directly on the parallel data. we create a synthetic parallel corpus by translating the source side of the parallel corpus with the initial model. the original source is paired with these translations to form the synthetic corpus for stage two. for the second stage of training, we then consider two cases: (1) where the 2nd stage model topology is identical (i.e. the original model and the new model both have or lack the output bias) and ( 2 ) where there is some change in model topology. for case (1), we use self-training: the first-stage model are fine-tuned using the synthetic corpus. for case (2): we use sequence-level knowledge distillation: new models are distilled @xcite by training from scratch on the synthetic corpus. by abuse of terminology, in both scenarios, we call the first-stage model the teacher, and the fine-tuned or distilled model the student. translation quality was measured using sacre-bleu @xcite . we define search errors as the number of lines changed in the translation output when vocabulary selection is applied. we measure the time taken to do inference on one core of a 12 core intel xeon cpu, on a pc with 16gb ram, running ubuntu 20.04 within a wsl2 hypervisor. for short listing, we create a candidate list of target sub-word translations for each source sub-word by using word alignments obtained from fastalign @xcite . a shortlist of target subwords is created before the translation of each sentence to constrain the possible output sub-words. ## results and analysis in order to understand the contributions of different aspects of the method, we perform additional experiments for analysis. ## conclusion we demonstrate that, with the proper training procedure, using locality sensitive hashing for vocabulary selection can significantly boost translation speed while consistently producing negligible search errors. we make the following recommendations for use in practice: for existing models and greedy search, perhaps where we may not know the exact training procedure and model, we can create a model that works with lsh vocabulary selection by distilling the original model to a comparable model without output bias. using label smoothing in the distillation can improve its translation quality if the original model was not trained with it. there will be minimal search errors in using lsh while achieving significant speed improvement. to train a new model for use with greedy search, a two stage procedure should also be used where the second stage is fine-tuned on the output of the first. both stages should train models without output bias. again, the fine-tuned models can be trained with label smoothing without affecting the effectiveness of lsh. lsh vocabulary selection introduce search errors for larger beam sizes, especially when label smoothing is used during fine-tuning. therefore, if using larger beams in inference, it is recommended not to use label smoothing in the distillation or fine-tuning step. (a) de-en (b) fr-en (c) es-en figure 10: lsh k-best vs. search errors for models with output bias. (a) de-en (b) fr-en (c) es-en figure 11: lsh #hashes vs. search errors for models without output bias. (a) de-en (b) fr-en (c) es-en figure 12: lsh #hashes vs. search errors for models with output bias. (a) de-en (b) fr-en (c) es-en figure 13: search errors vs. decrease in bleu. (a) de-en (b) fr-en (c) es-en figure 15: lsh #hashes vs. search errors for teacher models with output bias using beam width of 4. (a) de-en (b) fr-en (c) es-en figure 16: comparison of translation speed (sent./sec.) vs search error between lsh and lexical shortlisting.