Locality-Sensitive Hashing for Long Context Neural Machine Translation (2022)

Abstract: After its introduction the Transformer architecture quickly became the gold standard for the task of neural machine translation. A major advantage of the Transformer compared to previous architectures is the faster training speed achieved by complete parallelization across timesteps due to the use of attention over recurrent layers. However, this also leads to one of the biggest problems of the Transformer, namely the quadratic time and memory complexity with respect to the input length. In this work we adapt the locality-sensitive hashing approach of Kitaev et al. (2020) to self-attention in the Transformer, we extended it to cross-attention and apply this memory efficient framework to sentence- and document-level machine translation. Our experiments show that the LSH attention scheme for sentence-level comes at the cost of slightly reduced translation quality. For document-level NMT we are able to include much bigger context sizes than what is possible with the baseline Transformer. However, more context does neither improve translation quality nor improve scores on targeted test suites.

## introduction after its introduction in 2017, the transformer architecture @xcite quickly became the gold standard for the task of neural machine translation (nmt) @xcite . furthermore, variants of the transformer have since been used very successfully for a variety of other tasks such as language modeling (lm) @xcite , natural language understanding (nlu) @xcite , speech translation (st) @xcite , automatic speech recognition (asr) @xcite and image processing @xcite . a major advantage of the transformer compared to previous architectures is the faster training speed achieved by complete parallelization across timesteps. however, this also leads to one of the biggest problems of the transformer, namely the quadratic time and memory complexity of attention layers with respect to the sequence length. for sentence-level nmt this is not a big issue as most of the time the length of sequences is relatively short and can be handled efficently, even if subword segmentation is applied @xcite . however, this drastically changes when moving towards character-level @xcite or document-level (tiedemann and scherrer, 2017) nmt. especially for the latter, speed and memory issues are one of the biggest roadblocks towards 'true' document level systems @xcite . this leads to the situation where most works make do with including just a few sentences as a form of 'local' context information @xcite @xcite or heavily compressing the document information @xcite @xcite . more recently research focus has been shifting towards more efficient attention calculation for longer input sentences in several lm and nlu tasks @xcite . among these works is the approach by @xcite , in which the authors propose to make the attention matrix sparse by pre-selecting the relevant positions. they report good results on the lm objective while at the same time drastically reducing computational complexity. in this work we take the approach of @xcite as a starting point to improve the efficiency of (document-level) nmt systems. our contribution is three-fold: the problem of quadratic time and memory complexity of the attention framework has received increasing attention since the success of the transformer architecture @xcite . for asr, st and image processing the complexity can be reduced with relative ease by reducing the size of the time dimension with convolutional @xcite or pooling layers @xcite . furthermore, it is possible to restrict the attention to a few neighboring positions @xcite . however, this is not optimal for text input, as neighboring input words do not necessarily have the same strong correlation as neighboring audio frames or image pixels. existing work on improving the text processing complexity of the transformer mainly focuses on the case where all attention inputs come from the same embedding space, e.g. language modeling: dai et al. (2019) and rae et al. (2019) utilize a segment-level recurrence mechanism similar to what has been used in recurrent architectures. wang et al. (2020) project the time dimension of key and value down to a smaller, fixed-size dimension while leaving the queries untouched. directly altering the attention computation, child et al. (2019), sukhbaatar et al. (2019) and qiu et al. (2020) limit the attention to a local neighborhood or a fixed stride while zaheer et al. (2020) and beltagy et al. ( 2020 ) combine multiple sparse attention masks. in a more flexible approach, matching positions can be pre-selected using a locality-sensitive hashing function @xcite or clustering @xcite . in the present work, we pick one of the most efficient and best performing approaches up to date, namely the approach by @xcite and apply it to the task of machine translation. we confirm that the concepts can work for the self-attention in nmt systems and expand the framework for the case of cross-attention. most work related to document-level nmt limit the inter sentence context to few neighboring sentences. the simplest approach which we also follow in the present work, is to concatenate consecutive sentences using a special sentence separator token @xcite . there exist more sophisticated approaches which utilize separate encoders for the context information @xcite but later work seems to suggest that these approaches do not significantly outperform the simpler concatenation approach @xcite . in the realm of nmt, not so much work exists regarding improving the efficiency of the system and the work that exists mainly focuses on documentlevel nmt. @xcite propose to compress the context into a single vector which then can be attended to as an additional token embedding. tu et al. (2018) and kuang et al. (2018 ) utilize a cache that holds context information. zhang et al. (2020) and bao et al. (2021) mask out the attention energies between tokens from different sentences, showing that the full context is not necessary to achieve good translation performance. raganato et al. (2020) and you et al. (2020) replace most attention heads with fixed patterns but only for sentence-level nmt and only for self-attention as they report a severe degradation when doing the same for the cross-attention. there exist several different ways to implement lsh @xcite . the lsh scheme used by @xcite and consecutively in this work was proposed by @xcite . lsh has also been successfully applied to efficiently calculate pairwise embedding similarity for information retrieval @xcite . @xcite use lsh to pre-select embeddings in the softmax operation of an nmt system to speed up the decoding process. ## locality-sensitive hashing attention at the core of the transformer architecture is the attention mechanism that compares a sequence of queries q 1 , . . . q i to a sequence of key-value pairs (k 1 , v 1 ), . . . (k j , v j ) via a soft-lookup α(j|i) = α(q i , j, k j 1 ) and maps them to context vectors to compute the full sequence of context vectors, o(ij) operations are required. in the special case attention range of self-attention, i.e. @xmath0, the amount of operations grows quadratically with the sequence length i. since this can be problematic for long sequences, @xcite proposed to use locality-sensitive hashing (lsh) attention. in the following, we first describe the concept of lsh for self-attention, here we omit the leftto-right masking originally used @xcite and describe the concept for bidirectional self-attention instead. afterwards, we describe our extension of lsh to cross-attention. and α is normalized over p i instead of {1, . . . , j}. the hashing function h maps to a small number of classes {1, . . . , n hash } and is locality-sensitive, i.e. if two vectors are close-by they are likely to get assigned the same hash value. @xcite consider the case of self-attention and approximate the set p i to keep computation efficient. first the original sequence of keys is sorted by their hash value as primary criterion and original sequence order as secondary criterion. the resulting sequence is cut into chunks c i of fixed size and how to find an adequate key chunk for each query chunk? hashing and chunking is done for both the key and the query sequences, resulting in two different chunk sequences. we propose to calculate an alignment from the query chunks to the key chunks. for each query chunk c we find an aligned key chunk k(c) that contains queries with similar hash classes. to do this, the range of hash classes (h min , h max ) of the query chunk c is determined. next, we enumerate all key chunks k 1 , . . . , k n and search for the first key chunk k j 1 that contains an entry hashed to h min and the last key chunk k j 2 that corresponds to h max . then the middle chunk align each query chunk what happens if a query belongs to a hash class that is not represented in the aligned key chunk? since no keys are found that are close to the current query q i , we use the average value of the aligned query chunk. that is, we set pi := k(c i ) and obtain c (lsh throughout our experiments both key and query chunks are of equal size ℓ chunk . the lsh crossattention is shown in figure 2 . to reduce the impact of the chunking we compute attention not only within the aligned chunk but also one chunk to the left and right, similar to @xcite . this is applied both in selfand cross-attention. for unidirectional attention components, only the left context is considered. @xcite show that multi-round hashing can help to improve the performance of lsh attention systems. for multi-round hashing different hash functions h r are used to determine the corresponding (chunked) hash classes p r i and the context vector is calculated over the union ## experimental setup we evaluate our extensions to the attention by training transformer @xcite models with varying attention mechanisms on four mt tasks: the wmt 2016 news translation romanian to english data with 612k parallel sentences (europarl v8 & se times), the wmt 2019 english to german data with 329k parallel sentences (news commentary v14), as well as the iwslt 2017 english to german and english to italian data consisting of 232k and 206k parallel sentences (ted talks). the data is pre-processed by applying 20k spm merge operations (15k for both iwslt tasks) @xcite . the average sentence length for both wmt tasks is 30 subwords and 24 subwords for the iwslt tasks. the wmt @xmath1. for document-level systems we utilize this information in a pre-processing step by simply concatenating the k preceding sentences on source and target side to each sentence pair like tiedemann and scherrer (2017) do, but experiment with larger context sizes @xmath2, 3, 9, 12}. in particular @xmath3. in between the concatenated sentences we add a special separator token. we do not utilize right side context to ensure source and target have roughly the same length. the general system architecture follows the 'base' configuration of @xcite with 6 encoder/decoder layers of feature dimension d @xmath4, 8 attention heads and key/value dimension d @xmath5. we share the source/target embeddings as well as the transposed projections and employ training dropout of 30 % (20 % for @xmath6. all models are implemented in re-turnn @xcite . we use the adam optimizer (kingma and ba, 2015) with initial learning rate of 10 -3 . after training the systems for 200 checkpoints ( 1 /4 of all data for wmt @xmath7, 1 /2 for wmt @xmath8, we select the best checkpoint based on the dev perplexity on which we report bleu using sacrebleu @xcite and ter using tercom @xcite ) on an unseen test set. as systems with larger documentcontext see more frames in each epoch, we already stop training after 100 checkpoints for @xmath9. we find that the converged document-level systems are able to predict the correct number of target sentences with almost perfect accuracy. we extract the last predicted sentence for each sample and then calculate bleu and ter on the sentence-level data. when deploying lsh in the cross-attention, we found it crucial for training stability to first shuffle the key and query sequences as secondary criterion before sorting by hash classes. this helps during training in cases where the amount of queries/keys with the same hash class exceeds the window size. ## extended analysis 6.1 hash quality to evaluate the impact of approximating the full attention lsh we train systems with varying number of hash classes n hash in the cross attention. as described in section 3, queries may only attend to keys of the same hash class. the results for this are shown in in a first step we want to answer the question whether lsh attention actually makes use of different hash classes. otherwise, if one hash class is over-or underrepresented, the chunk size used by the system will not be large enough to actually attend to all relevant keys. to verify this, we extract the distribution of all key and query vectors the system generated on the development set and count the sizes of all hash classes. we find that indeed the hash classes are approximately equally distributed, i.e. all have a size close to 1 n hash . increasing the number of hash classes decreases the number of keys each query can attend to. this also decreases translation performance in terms of bleu and ter, but only minorly: the system using 8 hash classes, i.e. only attending to one eighth of all keys per query, only performs 1.1 % bleu worse than the baseline when also using lsh during inference. the previous results all also use lsh during inference. alternatively, we also experiment with full attention during inference after training the system with lsh. in this case, performance is almost equal to the lsh-restricted attention, even when using many hash classes. for each sentence pair, we extract the attention weights using full attention and sum over the key positions the lsh system attends to. this is the share of full attention covered by the lsh approximation, which however in the lsh system is renormalized to have a sum of 1 for each query. the average of this over all dev sentences and attention heads is shown in the last column of ## conclusion we present a method to make the transformer nmt architecture more memory-efficient when handling long input sequences. this is achieved by pre-selecting the most relevant candidates in selfattention and cross-attention using an lsh scheme that has been successfully applied for language modeling in previous work. we modify the existing lsh scheme to work in the nmt framework and conduct experiments on both sentence-level and document-level nmt tasks. our experiments show that the lsh attention scheme can be used for sentence-level nmt, although the approximation comes at the cost of slightly reduced translation quality. for documentlevel nmt we are able to include much bigger context sizes than what is possible with the baseline transformer. however, more context does neither improve translation quality nor improve scores on targeted test suites. in the future, we plan to use this approach for speech translation where long input sequences are a more pressing issue.