The Devil is in the Details: On the Pitfalls of Vocabulary Selection in Neural Machine Translation (2022)

Abstract: Vocabulary selection, or lexical shortlisting, is a well-known technique to improve latency of Neural Machine Translation models by constraining the set of allowed output words during inference. The chosen set is typically determined by separately trained alignment model parameters, independent of the source-sentence context at inference time. While vocabulary selection appears competitive with respect to automatic quality metrics in prior work, we show that it can fail to select the right set of output words, particularly for semantically non-compositional linguistic phenomena such as idiomatic expressions, leading to reduced translation quality as perceived by humans. Trading off latency for quality by increasing the size of the allowed set is often not an option in real-world scenarios. We propose a model of vocabulary selection, integrated into the neural translation model, that predicts the set of allowed output words from contextualized encoder representations. This restores translation quality of an unconstrained system, as measured by human evaluations on WMT newstest2020 and idiomatic expressions, at an inference latency competitive with alignment-based selection using aggressive thresholds, thereby removing the dependency on separately trained alignment models.

## introduction neural machine translation (nmt) has achieved great improvements in translation quality, largely thanks to the introduction of transformer models @xcite . however, increasingly larger models @xcite lead to prohibitively slow inference when deployed in industrial settings. especially for real-time applications, low latency is key. a number of inference optimization speedups have been proposed and are used in practice: reduced precision (aji and heafield, 2020), replacing self-attention with average attention networks (aans) @xcite , simpler simple recurrent units (ssrus) @xcite , or model pruning @xcite . another technique that is very common in practice is vocabulary selection @xcite which usually provides a good tradeoff between latency and automatic metric scores (bleu) and reduced inference cost is often preferred over the loss of ∼0.1 bleu. vocabulary selection is effective because latency is dominated by expensive, repeated decoder steps, where the final projection to the output vocabulary size contributes to a large portion of time spent @xcite . despite high parallelization in gpus, vocabulary selection is still relevant for gpu inference for state-of-theart models. however, we show that standard methods of vocabulary selection based on alignment model dictionaries lead to quality degradations not sufficiently captured by automatic metrics such as bleu. we demonstrate that this is particularly true for semantically non-compositional linguistic phenomena such as idiomatic expressions, and aggressive thresholds for vocabulary selection. for example, see figure 1 for alignment-model based vocabulary selection failing to include tokens crucial for translating idiomatic expressions in the set of allowed output words. while less aggressive thresholds can reduce the observed quality issues, it also reduces the desired latency benefit. in this paper we propose a neural vocabulary selection model that is jointly trained with the translation model and achieves translation quality at the level of an unconstrained baseline with latency at the level of an aggressively thresholded alignment-based vocabulary selection model. our contributions are as follows: ## pitfalls of vocabulary selection we first describe vocabulary selection and then analyze its shortcomings. throughout the paper, we use the recall of unique target sentence tokens as a proxy for measuring vocabulary selection quality, i.e. the reachability of the optimal translation. we use the average vocabulary size in inference decoder steps across sentences as a proxy for translation latency since it directly impacts decoding speed @xcite . ## neural vocabulary selection (nvs) we incorporate vocabulary selection directly into the neural translation model, instead of relying on a separate statistical model based on token translation probabilities. this enables predictions based on contextualized representations of the full source sentence. it further simplifies the training procedure and domain adaptation, as we do not require a separate training procedure for an alignment model. the goal of our approach is three-fold. we aim to ( 1 ) keep the general transformer @xcite translation model architecture, (2) incur only a minimal latency overhead that amortizes by cheaper decoder steps due to smaller output vocabularies, and (3) scale well to sentences of different lengths. figure 2 shows the neural vocabulary selection (nvs) model. we base the prediction of output tokens on the contextualized hidden representation produced by the transformer encoder @xmath0. the t source tokens are comprised of t -1 input tokens and a @xmath1. to obtain the set of target tokens, we first project each source position to the target vocabulary size v , apply max-pooling across tokens @xcite , and finally use the sigmoid function, σ(•), to obtain to learn the v × d + v parameters for nvs, we use a binary cross-entropy loss with the binary ground truth vector @xmath2, where each entry indicates the presence or absence of target token y i . we define the loss as ## related work there are two dominant approaches to generate a restricted set of target word candidates (i) using an external model and (ii) using the nmt system itself. in the first approach, a short-list of translation candidates is generated from word-alignments @xcite , phrase table, and the most common target words @xcite . @xcite propose an additional method using support vector machines to predict target candidates from a sparse representation of the source sentence. in the second approach, @xcite build alignment probability table from the softattention layer from decoder to encoder. however, applying their method to multi-head attention in transformer is non-trivial as attention may not capture word-alignments in multiple attention layers @xcite . @xcite use local sensitive hashing to shrink the target vocabu-source idiomatic target i thought i would be nervous , but i was cool as a cu cum ber . wunder wa ffe they said he was ' a little bit under the weather ' . sich nicht wohl fühlen i still can 't wra p my head around it . closely related to our work is @xcite , who predict all words in a target sentence from the initial hidden state of the decoder. our nvs model differs from theirs in that we make a prediction for each source token and aggregate the results via max-pooling to scale with sentence length. recent work of @xcite illustrates the risk associated with reducing latency via vocabulary selection in domain-mismatched settings. our work takes this a step further by providing a detailed analysis on the shortcomings of vocabulary selection and proposing a model to mitigate them. related to our findings on non-compositional expressions, @xcite evaluate the effect of methods used to speed up decoding in transformer models on gender bias and find minimal bleu degradations but reduced gendered noun translation performance on a targeted test set. ## conclusions alignment-based vocabulary selection is a common method to heavily constrain the set of allowed output words in decoding for reduced latency with only minor bleu degradations. we showed with human evaluations and a targeted qualitative analysis that such translations are perceivably worse. even recent automatic metrics based on pre-trained neural networks, such as comet, are only able to capture the observed quality degradations in two out of four language pairs. human-perceived quality is negatively affected both for generic translations, represented by newstest2020, as well as for idiomatic translations. increasing the vocabulary selection threshold can alleviate the quality issues at an increased single sentence translation latency. to preserve both translation latency and quality we proposed a neural vocabulary selection model that is directly integrated into the translation model. such a joint model further simplifies the training pipeline, removing the dependency on a separate alignment model. our model has higher reference token recall at similar vocabulary sizes, translating into higher quality at similar latency.