SlowGuess/NAACL_2022_OCR huggingface dataset:
Tobias Domhan* Eva Hasler* Ke Tran Sony Trenous Bill Byrne Felix Hieber
Amazon AI Translate
{domhant, ehasler, trnke, trenous, willbyrn, fhieber}@amazon.com
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.
Neural Machine Translation (NMT) has achieved great improvements in translation quality, largely thanks to the introduction of Transformer models (Vaswani et al., 2017). However, increasingly larger models (Aharoni et al., 2019; Arivazhagan et al., 2019) lead to prohibitively slow inference when deployed in industrial settings. Especially for real-time applications, low latency is key. A number of inference optimization speed-ups have been proposed and are used in practice:
EN: to swal low the bitter pill
DE: in den sau __ ren ap __ fel bei __ fen
GL: to bite into the sour ap ple
EN: by ho__ok or cro__ok
DE: auf biegen und brechen
GL: by bending and breaking
EN: to buy a p__ig in a po__ke
DE: die kat__ze im sack kaufen
GL: to buy the cat in the bag
EN: to swe at blood
DE: blu_t und wasser sch_wit_ zen
GL: to sweat blood and water
EN: make yourself at home!
DE: machen sie es sich bequem!
GL: make yourself comfortable!
Figure 1: Examples of subword-segmented idiomatic expressions (EN) and their German correspondences (DE) as well as an English gloss (GL) of the German expression. Alignment-based vocabulary selection: output tokens missing from the allowed set of top-k output tokens are marked in orange/bold (red/italic) for  $k = 200$  ( $k = 1000$ ).
reduced precision (Aji and Heafield, 2020), replacing self-attention with Average Attention Networks (AANs) (Zhang et al., 2018), Simpler Simple Recurrent Units (SSRUs) (Kim et al., 2019), or model pruning (Behnke and Heafield, 2020; Behnke et al., 2021).
Another technique that is very common in practice is vocabulary selection (Jean et al., 2015) which usually provides a good tradeoff between latency and automatic metric scores (BLEU) and reduced inference cost is often preferred over the loss of  $\sim 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 (Bérard et al., 2021). Despite
Human Language Technologies, pages 1861 - 1874
high parallelization in GPUs, vocabulary selection is still relevant for GPU inference for state-of-the-art 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.
- We demonstrate that alignment-based vocabulary selection is not limited by alignment model quality, but rather inherently by making target word predictions out of context (§2).
- We propose a Neural Vocabulary Selection (NVS) model based on the contextualized deep encoder representation ( $\S 3$ ).
- We show that alignment-based vocabulary selection leads to human-perceived translation quality drops not sufficiently captured by automatic metrics and that our proposed model can match an unconstrained model's quality while keeping the latency benefits of vocabulary selection (§4).
2 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 (Kasai et al., 2020).
Scope	k=200	k=1000
model	ref.	model	ref.
EN-DE	fast_align	99.6	97.5	99.9	99.7
GIZA++	99.6	97.8	100.0	99.7
MaskAlign	96.9	93.1	99.6	98.6
EN-RU	fast_align	98.7	93.8	99.9	98.9
GIZA++	98.6	94.2	99.9	99.2
MaskAlign	94.2	87.2	99.1	96.7
Table 1: Recall of unique reference and model tokens on newstest2020 with word probability tables extracted from different alignment models.
2.1 Vocabulary selection
Vocabulary selection (Jean et al., 2015), also known as lexical shortlisting or candidate selection, is a common technique for speeding up inference in sequence-to-sequence models, where the repeated computation of the softmax over the output vocabulary  $\mathcal{V}$  of size  $V$  incurs high computational cost in the next word prediction at inference time:
$$
p \left(y _ {t} \mid y _ {1: t - 1}, x; \theta\right) = \operatorname {s o f t m a x} (\boldsymbol {W} \boldsymbol {h} + \boldsymbol {b}), \tag {1}
$$
where  $W\in \mathbb{R}^{V\times d}$ $b\in \mathbb{R}^V$  and  $h\in \mathbb{R}^d$ $d$  being the hidden size of the network. Vocabulary selection chooses a subset  $\bar{\nu}\subset \nu$  , with  $\bar{V}\ll V$  , to reduce the size of matrix multiplication in Equation (1) such that
$$
\bar {p} \left(y _ {t} \mid y _ {1: t - 1}, x; \theta\right) = \operatorname {s o f t m a x} (\bar {\boldsymbol {W}} \boldsymbol {h} + \bar {\boldsymbol {b}}), \tag {2}
$$
where  $\bar{W} \in \mathbb{R}^{\bar{V} \times d}$  and  $\bar{b} \in \mathbb{R}^{\bar{V}}$ . The subset  $\bar{\nu}$  is typically chosen to be the union of the top-k target word translations for each source token, according to the word translation probabilities of a separately trained word alignment model (Jean et al., 2015; Shi and Knight, 2017). Decoding with vocabulary selection usually yields similar scores according to automatic metrics, such as BLEU (Papineni et al., 2002), compared to unrestricted decoding but at reduced latency (L'Hostis et al., 2016; Mi et al., 2016; Sankaran et al., 2017; Junczys-Dowmunt et al., 2018). In the following, we show that despite its generally solid performance, vocabulary selection based on word alignment models negatively affects translation quality, not captured by standard automatic metrics. We use models trained on WMT20 (Barrault et al., 2020a) data for all evaluations in this section, see Section 4.1 for details.
2.2 Alignment model quality
In practice, the chosen subset of allowed output words is often determined by an alignment
model, such as fast_align (Dyer et al., 2013), which provides a trade-off between the speed of alignment model training and the quality of alignments (Jean et al., 2015; Junczys-Dowmunt et al., 2018). fast_align's reparametrization of IBM model 2 (Brown et al., 1993) places a strong prior for alignments along the diagonal. We investigate whether more sophisticated alignment models can lead to better vocabulary selection, especially for language pairs with high amount of reordering. To evaluate this we compute the recall of translation model and reference tokens using GIZA++ (Och and Ney, 2003) and MaskAlign<sup>1</sup> (Chen et al., 2021) as seen in Table 1. We extract top-k word translation tables (from fast_align, GIZA++, and MaskAlign) by force-aligning the training data. Overall, GIZA++ achieves the best recall, and it is just slightly better than fast_align. MaskAlign, a state-of-the-art neural alignment model, underperforms fast_align with respect to recall. While performance of MaskAlign may be improved with careful tuning of its hyperparameters via gold alignments (Chen et al., 2021), we choose fast_align as a strong, simple baseline for vocabulary selection in the following.
2.3 Out-of-context word selection
Alignment-based vocabulary selection does not take source sentence context into account. A top-k list of translation candidates for a source word will likely cover multiple senses for common words, but may be too limited when a translation is highly dependent on the source context. Here we consider idiomatic expressions as a linguistic phenomenon that is highly context-dependent due to its semantically non-compositional nature.
Table 2 compares the recall of tokens in the reference translation when querying the translation lexicon of the alignment model for two different top-k settings. Recall is computed as the percentage of unique tokens in the reference translation that appear in the top-k lexicon, or more generally, in the set of predicted tokens according to a vocabulary selection model. We evaluate two scopes for test sets of idiomatic expressions: the full source and target sentence vs. the source and target idiomatic multi-word expressions according to metadata. The Idioms test set is an internal set of 100 English idioms in context and their human translations. ITDS is
Scope	k=200	k=1000
Idioms	ITDS	Idioms	ITDS
EN-DE	sentence	96.0	91.5	99.1	97.3
idiom	80.0	58.2	92.8	88.0
DE-EN	sentence	-	92.2	-	98.0
idiom	-	75.5	-	90.1
EN-RU	sentence	93.0	-	98.7	-
idiom	65.7	-	85.6	-
Table 2: Recall on internal test set of idioms in context and IdiomTranslationDS test set.
Scope	Wikiquote
k = 200	k = 1000
EN-DE	literal	96.5	99.2
idiomatic	74.6	91.0
EN-RU	literal	91.2	99.2
idiomatic	67.6	88.9
Table 3: Recall on English proverbs and their literal and idiomatic translations sourced from Wikiquote.
the IdiomTranslationDS $^2$  data released by Fadaee et al. (2018) with 1500 test sentences containing English and German idiomatic expressions for evaluation into and out of German, respectively. The results show that recall increases when increasing  $k$  but is consistently lower for the idiomatic expressions than for full sentences. Clearly, the idiom translations contain tokens that are on average less common than the translations of "regular" inputs. As a consequence, increasing the output vocabulary is less effective for idiom translations, with recall lagging behind by up to  $9.3\%$ . This can directly affect translation quality because the NMT model will not be able to produce idiomatic translations given an overly restrictive output vocabulary.
Table 3 shows a similar comparison but here we evaluate full literal translations vs. full idiomatic translations on a data set of English proverbs from Wikiquote<sup>3</sup>. For EN-DE, we extracted 94 triples of English sentence and two references, for ENRU we extracted 262 triples. Although in both cases recall can be improved by increasing  $k$ , it helps considerably less for idiomatic than for literal translations.
Figure 1 shows examples of idiomatic expressions from the ITDS set and the output tokens belonging to an idiomatic translation that are missing from the respective lexicon used for vocabulary
2https://github.com/marziehf/IdiomTranslationDS
<sup>3</sup>https://github.com/awslabs/sockeye/tree/naacl2022/naacl2022/wikiquote
<sup>1</sup>With default hyper-parameters from https://github.com/THUNLP-MT/Mask-Align
Scope	idioms in context
w/o adapt	w/ adapt
EN-DE	sentence	96.0	98.0
idiom	80.0	97.2
EN-RU	sentence	93.0	96.6
idiom	65.7	91.7
Table 4: Recall on internal test set of idioms in context, with and without adapting the fast_align translation lexicons  $(k = 200)$
selection. While for some of the examples, increasing the lexicon size solves the problem, for others the idiomatic translation can still not be generated because of missing output tokens.
These results demonstrate that there is room for improvement in vocabulary selection approaches when it comes to non-literal translations.
2.4 Domain mismatch in adaptation settings
Using a word alignment model to constrain the NMT output vocabulary means that this model should ideally also be adapted when adapting the NMT model to a new domain. Table 4 shows that adapting the word alignment model with relevant in-domain data (in this case, idiomatic expressions in context) yields strong recall improvements for vocabulary selection. Compared to increasing the per-source-word vocabulary as shown in Table 2, the improvement in recall for idiom tokens is larger which highlights the importance of having a vocabulary selection model which matches the domain of the NMT model. This also corroborates the finding of Bogoychev and Chen (2021) that vocabulary selection can be harmful in domain-mismatched scenarios.
We argue that integrating vocabulary prediction into the NMT model avoids the need for mitigating domain mismatch because domain adaptation will update both parts of the model. This simplifies domain adaptation since it only needs to be done once for a single model and does not require adaptation or re-training of a separate alignment model.
We use target recall as a measure for selection model quality. We see that alignment model quality only has a limited impact on target token recall with more recent models actually having lower recall overall. In domain adaptation scenarios vocabulary selection limits translation quality if the selection model is not adapted. The main challenge
for alignment-based vocabulary selection comes from its out-of-context selection of target tokens on a token-by-token basis, shown to reduce recall for translation of idiomatic, non-literal expressions. Increasing the size of the allowed set can compensate for this shortcoming at the cost of latency. However, this begs the question of whether context-sensitive selection of target tokens can achieve higher recall without increasing vocabulary size.
3 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 (Vaswani et al., 2017) 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  $\mathbf{H} \in \mathbb{R}^{d \times t}$  for  $t$  source tokens and a hidden size of  $d$ . The  $t$  source tokens are comprised of  $t - 1$  input tokens and a special  $<\mathrm{EOS}>$  token. To obtain the set of target tokens, we first project each source position to the target vocabulary size  $V$ , apply max-pooling across tokens (Shen et al., 2018), and finally use the sigmoid function,  $\sigma(\cdot)$ , to obtain
$$
\boldsymbol {z} = \sigma (\operatorname {m a x p o o l} (\boldsymbol {W H} + \boldsymbol {b})), \tag {3}
$$
where  $W \in \mathbb{R}^{V \times d}$ ,  $b \in \mathbb{R}^V$  and  $z \in \mathbb{R}^V$ . The max-pooling operation takes the per-dimension maximum across the source tokens, going from  $\mathbb{R}^{V \times t}$  to  $\mathbb{R}^V$ . Each dimension of  $z$  indicates the probability of a given target token being present in the output given the source. To obtain the target Bag-of-words (BOW), we select all tokens where  $z_i > \lambda$  as indicated in the right-hand side of Figure 2, where  $\lambda$  is a free parameter that controls the size  $\bar{V}$  of the reduced vocabulary  $\bar{\nu}$ . At inference time, the output projection and softmax at every
Figure 2: The proposed Neural Vocabulary Selection (NVS) model. A subset of the full vocabulary is selected based on encoder representation  $H$  and passed to the decoder for a reduced output projection at every decoder step.
decoder step are computed over the predicted BOW of size  $\bar{V}$  only.
We achieve goal (1) by basing predictions on the encoder representation already used by the decoder. Goal (2) is accomplished by restricting NVS to a single layer and basing the prediction on the encoder output, where we can parallelize computation across source tokens. Inference latency is dominated by non-parallelizable decoder steps (Kasai et al., 2020). By projecting to the target vocabulary per source token, each source token can "vote" on a set of target tokens. The model automatically scales to longer sentences via the max-pooling operation, acting as a union of per-token choices, fulfilling goal (3). Max-pooling does not tie the predictions across timesteps as they would be with mean-pooling which would also depend on sentence length. Additionally, we factor in a sentence-level target token prediction based on the <EOS> token. The probability of a target word being present is represented by the source position with the highest evidence, backing off to a base probability of a given word via the bias vector  $\mathbf{b}$ .
To learn the  $V\times d + V$  parameters for NVS, we use a binary cross-entropy loss with the binary ground truth vector  $\pmb {y}\in \mathbb{R}^{V}$  , where each entry indicates the presence or absence of target token  $\pmb {y}_i$  . We define the loss as
$$
L _ {\mathrm {N V S}} = \frac {1}{Z} \sum_ {i = 0} ^ {V} \boldsymbol {y} _ {i} \log (\boldsymbol {z} _ {i}) \lambda_ {p} + (1 - \boldsymbol {y} _ {i}) \log (1 - \boldsymbol {z} _ {i}),
$$
where  $\lambda_{p}$  is a weight for the positive class and
$Z = V + (\lambda_p - 1)*n_p$  is the normalizing factor, with  $n_p = \sum_i\pmb {y}_i$  being the number of unique target words. Most dimensions of the  $V$  -dimensional vector  $\pmb{y}$  will be zero as only a small number of target words are present for a given sentence. To counter the class imbalance of the negative class over the positive class the  $\lambda_{p}$  weight allows overweighting the positive class. This has the same effect as if each target word had occurred  $\lambda_{p}$  times. The NVS objective  $(L_{\mathrm{NVS}})$  is optimized jointly with the standard negative log-likelihood translation model loss  $(L_{\mathrm{MT}}):L = L_{\mathrm{NVS}} + L_{\mathrm{MT}}$
Our training setup is guided by best practices for efficient NMT to provide a strong low latency baseline: deep Transformer as encoder with a lightweight recurrent unit in the shallow decoder (Bérard et al., 2021; Kim et al., 2019), int8 quantization for CPU and half-precision GPU inference. We use the constrained data setting from WMT20 (Barrault et al., 2020b) with four language pairs English-German, German-English, English-Russian, Russian-English and apply corpus cleaning heuristics based on sentence length and language identification. We tokenize with  $\text{sacremoses}^4$  and byte-pair encode (Sennrich et al., 2016) the data with 32k merge operations.
4https://github.com/alvations/sacremoses
vocabulary selection	newstest2020	ITDS idiom test
BLEU↑	COMET↑	human eval (%)↑	CPU↓ (ms / %)	GPU↓ (ms / %)	BLEU↑	human eval (%)↑
EN-DE	34.4±0.2	0.461±0.003	100	100	999±12	208±1	25.5±0.2	100	100
34.2±0.1	0.453±0.004	98.0±1.9	96.7±1.9	51.0±0.6	78.2±0.9	25.4±0.2	96.6±1.9	97.2±1.8
34.3±0.2	0.462±0.005	100.1±2.0	100.0±1.9	66.6±0.5	82.9±0.7	25.5±0.1	98.5±1.7	100.4±1.9
34.4±0.1	0.460±0.004	-	99.5±1.7	46.6±0.5	79.8±0.9	25.5±0.2	-	99.3±1.9
34.4±0.2	0.461±0.004	100.0±1.8	-	54.1±0.7	82.2±0.9	25.5±0.2	99.1±1.9	-
34.4±0.2	0.461±0.004	-	-	65.6±0.8	86.1±0.7	25.5±0.2	-	-
DE-EN	40.7±0.2	0.645±0.001	100	100	1128±17	229±1	29.8±0.1	100	100
40.7±0.2	0.643±0.002	98.1±2.6	98.8±2.3	53.6±0.5	76.9±0.7	29.8±0.1	96.9±2.1	97.9±2.2
40.7±0.2	0.645±0.001	99.3±2.5	100.3±2.2	67.9±0.5	84.4±0.6	29.8±0.1	101.2±2.1	101.8±2.2
40.7±0.1	0.644±0.003	-	99.8±2.4	47.6±0.6	76.4±0.7	29.7±0.1	-	99.9±2.1
40.7±0.2	0.645±0.002	101.6±2.6	-	54.4±0.6	79.8±1.3	29.7±0.1	101.0±2.2	-
40.7±0.1	0.645±0.002	-	-	63.9±0.5	83.2±0.7	29.7±0.1	-	-
EN-RU	23.6±0.1	0.528±0.002	100	100	673±7	145±1
23.3±0.1	0.507±0.003	95.9±1.8	94.1±1.8	49.9±0.8	78.2±0.9
23.5±0.1	0.527±0.003	100.4±1.8	99.9±1.7	64.8±0.8	82.9±0.7
23.6±0.1	0.525±0.002	-	99.2±1.7	48.1±0.6	79.8±0.9
23.6±0.1	0.528±0.003	99.3±1.7	-	58.2±0.8	82.2±0.9
23.6±0.0	0.529±0.002	-	-	67.5±0.8	86.1±0.7
RU-EN	35.5±0.1	0.559±0.005	100	100	568±8	123±1
35.2±0.1	0.552±0.004	96.4±2.5	97.7±2.5	49.3±0.9	79.4±1.0
35.4±0.1	0.561±0.004	99.9±2.5	98.7±2.3	62.7±0.5	81.8±0.8
35.5±0.1	0.558±0.004	-	100.2±2.5	47.2±0.5	79.9±1.0
35.5±0.1	0.559±0.005	101.1±2.6	-	52.2±0.6	80.9±1.3
35.5±0.1	0.559±0.005	-	-	60.2±0.6	83.6±0.9
Table 5: Experimental results for unconstrained decoding (baseline), alignment-based VS with different  $k$ , and Neural Vocabulary Selection with varying  $\lambda$ . BLEU and COMET: mean and std of three runs with different random seeds. Human evaluation: source-based direct assessment renormalized so that the unconstrained baseline is at  $100\%$ , with  $95\%$  CI of a paired t-test. We ran two sets of human evaluations comparing 4 systems each. CPU/GPU: p90 latency in ms with  $95\%$  CI based on 30 runs with batch size 1, shown as a percentage of the baseline.
All models are Transformers (Vaswani et al., 2017) trained with the Sockeye 2 toolkit (Domhan et al., 2020). We release the NVS code as part of the Sockeye toolkit. We use a 20-layer encoder and a 2-layer decoder with self-attention replaced by SSRUs (Kim et al., 2019).
NVS and NMT objectives are optimized jointly, but gradients of the NVS objective are blocked before the encoder. This allows us to compare the different vocabulary selection techniques on the same translation model that is unaffected by the choice of vocabulary selection. All vocabulary selection methods operate at the BPE level. We use the translation dictionaries from fast_align for alignment-based vocabulary selection. We use a minimum of  $k = 200$  for alignment-based vocabulary selection which is at the upper end of what is found in previous work. Junczys-Dowmunt et al. (2018) set  $k = 100$ , Kim et al. (2019) set  $k = 75$ , and Shi and Knight (2017) set  $k = 50$ .