SlowGuess/NAACL_2022_OCR huggingface dataset:
Felix Stahlberg and Shankar Kumar
{fstahlberg, shankarkumar} @google.com
The softmax layer in neural machine translation is designed to model the distribution over mutually exclusive tokens. Machine translation, however, is intrinsically uncertain: the same source sentence can have multiple semantically equivalent translations. Therefore, we propose to replace the softmax activation with a multi-label classification layer that can model ambiguity more effectively. We call our loss function Single-label Contrastive Objective for Non-Exclusive Sequences (SCONES). We show that the multi-label output layer can still be trained on single reference training data using the SCONES loss function. SCONES yields consistent BLEU score gains across six translation directions, particularly for medium-resource language pairs and small beam sizes. By using smaller beam sizes we can speed up inference by a factor of  $3.9\mathrm{x}$  and still match or improve the BLEU score obtained using softmax. Furthermore, we demonstrate that SCONES can be used to train NMT models that assign the highest probability to adequate translations, thus mitigating the "beam search curse". Additional experiments on synthetic language pairs with varying levels of uncertainty suggest that the improvements from SCONES can be attributed to better handling of ambiguity.
Conventional neural machine translation (NMT) models learn the probability  $P(\mathbf{y}|\mathbf{x})$  of the target sentence  $\mathbf{y}$  given the source sentence  $\mathbf{x}$  (Kalchbrenner and Blunsom, 2013; Sutskever et al., 2014). This framework implies that there is a single best translation for a given source sentence: if there were multiple valid translations  $\mathbf{y}_1$  and  $\mathbf{y}_2$  they would need to share probability mass (e.g.  $P(\mathbf{y}_1|\mathbf{x}) = 0.5$  and  $P(\mathbf{y}_2|\mathbf{x}) = 0.5$ ), but such
a distribution could also represent model uncertainty, i.e. the case when either  $\mathbf{y}_1$  or  $\mathbf{y}_2$  are correct translations. Therefore, learning a single distribution over all target language sentences does not allow the model to naturally express intrinsic uncertainty $^2$  (Padó et al., 2009; Dreyer and Marcu, 2012; Ott et al., 2018; Stahlberg et al., 2022), the nature of the translation task to allow multiple semantically equivalent translations for a given source sentence. A single distribution over all sequences represents uncertainty by assigning probabilities, but it cannot distinguish between different kinds of uncertainty (e.g. model uncertainty versus intrinsic uncertainty).
Therefore, in this work we frame machine translation as a multi-label classification task (Tsoumakas and Katakis, 2007; Zhang and Zhou, 2014). Rather than learning a single distribution  $P(\mathbf{y}|\mathbf{x})$  over all target sentences  $\mathbf{y}$  for a source sentence  $\mathbf{x}$ , we learn binary classifiers for each sentence pair  $(\mathbf{x},\mathbf{y})$  that indicate whether or not  $\mathbf{y}$  is a valid translation of  $\mathbf{x}$ . In this framework, intrinsic uncertainty can be represented by setting the probabilities of two (or more) correct translations  $\mathbf{y}_1$  and  $\mathbf{y}_2$  to 1 simultaneously. The probabilities for each translation are computed using separate binary classifiers, and thus there is no requirement that the probabilities sum to one over all translations. In practice, the probability of a complete translation is decomposed into a product of the token-level probabilities. Thus we replace the softmax output layer in Transformer models (Vaswani et al., 2017) with sigmoid activations that assign a probability between 0 and 1 to each token in the vocabulary at each time step. We propose a loss function, Single-label Contrastive Objective for Non-Exclusive Sequences (SCONES) that allows us to train our models on single reference training data. Our work is inspired by noise-contrastive es
2This is sometimes referred to as aleatoric uncertainty in the literature (Der Kiureghian and Ditlevsen, 2009).
1https://en.wikipedia.org/wiki/Cream_ tea#Variations
Human Language Technologies, pages 4950 - 4961
timation (NCE) (Gutmann and Hyvarinen, 2010; Mnih and Teh, 2012). Unlike NCE, whose primary goal was to efficiently train models over large vocabularies, our motivation for SCONES is to model non-exclusive outputs.
We demonstrate multiple benefits of training NMT models using SCONES when compared to standard cross-entropy with regular softmax. We report consistent BLEU score gains between  $1\% - 9\%$  across six different translation directions. SCONES with greedy search typically outperforms softmax with beam search, resulting in inference speed-ups of up to  $3.9\mathrm{x}$  compared to softmax without any degradation in BLEU score.
SCONES can be tuned to mitigate some of the pathologies of traditional NMT models. Softmax-based models have been shown to assign the highest probability to either empty or inadequate translations (modes) (Stahlberg and Byrne, 2019; Eikema and Aziz, 2020). This behavior manifests itself as the "beam search curse" (Koehn and Knowles, 2017): increasing the beam size may lead to worse translation quality. We show that SCONES can be used to train models that a) assign the highest probability to adequate translations and b) do not suffer from the beam search curse.
Finally, we use SCONES to train models on synthetic translation pairs that we generate by sampling from the IBM Model 3 (Brown et al., 1993). By varying the sampling temperature, we control the level of ambiguity in the language pair. We show that SCONES is effective in improving the adequacy of the highest probability translation for highly ambiguous translation pairs, confirming our intuition that SCONES can handle intrinsic uncertainty well.
2 Training NMT models with SCONES
We denote the (subword) vocabulary as  $\mathcal{V} = \{w_1, \ldots, w_{|\mathcal{V}|}\}$ , the special end-of-sentence symbol as  $w_1 = \langle \mathrm{s} \rangle$ , the source sentence as  $\mathbf{x} = \langle x_1, \ldots, x_{|\mathbf{x}|} \rangle \in \mathcal{V}^*$ , a translation as  $\mathbf{y} = \langle y_1, \ldots, y_{|\mathbf{y}|} \rangle \in \mathcal{V}^*$ , and a translation prefix as  $\mathbf{y}_{\leq i} = \langle y_1, \ldots, y_i \rangle$ . We use a center dot “.” for string concatenations. Unlike conventional NMT that models a single distribution  $P(\mathbf{y} | \mathbf{x})$  over all target language sentences, SCONES learns a separate binary classifier for each sentence pair  $(\mathbf{x}, \mathbf{y})$ . We define a Boolean function  $t(\cdot, \cdot)$  that indicates
whether  $\mathbf{y}$  is a valid translation of  $\mathbf{x}$ :
$$
t (\mathbf {x}, \mathbf {y}) := \left\{ \begin{array}{l l} \text {t r u e} & \text {i f y i s a t r a n s l a t i o n o f x} \\ \text {f a l s e} & \text {o t h e r w i s e} \end{array} . \right.
$$
We do not model  $t(\cdot, \cdot)$  directly. To guide decoding, we learn variables  $z_{\mathbf{x},\mathbf{y}}$  which generalize  $t(\cdot, \cdot)$  to translation prefixes:
$$
z _ {\mathbf {x}, \mathbf {y}} := \left\{ \begin{array}{l l} 1 & \exists \mathbf {y} ^ {\prime} \in \mathcal {V} ^ {*}: t (\mathbf {x}, \mathbf {y} \cdot \mathbf {y} ^ {\prime}) = \text {t r u e} \\ 0 & \text {o t h e r w i s e} \end{array} , \right. \tag {2}
$$
i.e.  $z_{\mathbf{x},\mathbf{y}}$  is a binary label for the pair  $(\mathbf{x},\mathbf{y})$  consisting of source sentence  $\mathbf{x}$  and the translation prefix  $\mathbf{y}$ :  $z_{\mathbf{x},\mathbf{y}} = 1$  iff.  $\mathbf{y}$  is a prefix of a valid translation of  $\mathbf{x}$ . We decompose its probability as a product of conditionals to facilitate left-to-right beam decoding:
$$
\begin{array}{l} P \left(z _ {\mathbf {x}, \mathbf {y}} = 1 | \mathbf {x}\right) := \prod_ {i = 1} ^ {| \mathbf {y} |} P \left(z _ {\mathbf {x}, \mathbf {y} _ {\leq i}} = 1 \mid z _ {\mathbf {x}, \mathbf {y} _ {<   i}} = 1, \mathbf {x}\right) \\ = \quad \prod_ {i = 1} ^ {| \mathbf {y} |} P \left(z _ {\mathbf {x}, \mathbf {y} _ {\leq i}} = 1 | \mathbf {x}, \mathbf {y} _ {<   i}\right). \tag {3} \\ \end{array}
$$
We assign the conditional probabilities by applying the sigmoid activation function  $\sigma (\cdot)$  to the logits:
$$
P \left(z _ {\mathbf {x}, \mathbf {y} <   i \cdot w} = 1 | \mathbf {x}, \mathbf {y} <   i\right) = \sigma \left(f (\mathbf {x}, \mathbf {y} <   i) _ {w}\right), \tag {4}
$$
where  $w \in \mathcal{V}$  is a single token,  $f(\mathbf{x},\mathbf{y}_{<i}) \in \mathbb{R}^{|\mathcal{V}|}$  are the logits at time step  $i$ , and  $f(\mathbf{x},\mathbf{y}_{<i})_w$  is the logit corresponding to token  $w$ . The only architectural difference to a standard NMT model is the output activation: instead of the softmax function that yields a single distribution over the full vocabulary, we use multiple sigmoid activations in each logit component to define separate Bernoulli distributions for each item in the vocabulary (Fig. 1). However, using such a multi-label classification view requires a different training loss function because, unlike the probabilities from a softmax, the probabilities in Eq. 4 do not provide a normalized distribution over the vocabulary. An additional challenge is that existing MT training datasets typically do not provide more than one reference translation. Our SCONES loss function aims to balance two token-level objectives using a scaling factor  $\alpha \in \mathbb{R}^+$ :
$$
\mathcal {L} (\mathbf {x}, \mathbf {y}) = \frac {1}{| \mathbf {y} |} \sum_ {i = 1} ^ {| \mathbf {y} |} \mathcal {L} _ {\text {S C O N E S}} (\mathbf {x}, \mathbf {y}, i), \tag {5}
$$
3As a base case we define  $P(z_{\mathbf{x},\epsilon} = 1|\mathbf{x}) = 1$  for the empty translation prefix.
Figure 1: Multi-way NMT Transformer architecture for non-exclusive target sequences.
$$
\mathcal {L} _ {\text {S C O N E S}} (\mathbf {x}, \mathbf {y}, i) = \mathcal {L} _ {+} (\mathbf {x}, \mathbf {y}, i) + \alpha \mathcal {L} _ {-} (\mathbf {x}, \mathbf {y}, i). \tag {6}
$$
$\mathcal{L}_{+}(\cdot)$  aims to increase the probability  $P(z_{\mathbf{x},\mathbf{y}_{\leq i}} = 1|\mathbf{x},\mathbf{y}_{< i})$  of the gold label  $y_{i}$  since it is a valid extension of the translation prefix  $\mathbf{y}_{< i}$ :
$$
\begin{array}{l} \mathcal {L} _ {+} (\mathbf {x}, \mathbf {y}, i) = - \log P \left(z _ {\mathbf {x}, \mathbf {y} _ {\leq i}} = 1 \mid \mathbf {x}, \mathbf {y} _ {<   i}\right) \tag {7} \\ = - \log \sigma (f (\mathbf {x}, \mathbf {y} _ {<   i}) _ {y _ {i}}). \\ \end{array}
$$
$\mathcal{L}_{-}(\cdot)$  is designed to reduce the probability  $P(z_{\mathbf{x},\mathbf{y}_{< i}\cdot w} = 1|\mathbf{x},\mathbf{y}_{< i})$  for all labels  $w$  except for the gold label  $y_{i}$ :
$$
\begin{array}{l} \mathcal {L} _ {-} (\mathbf {x}, \mathbf {y}, i) = - \sum_ {w \in \mathcal {V} \backslash \{y _ {i} \}} \log P (z _ {\mathbf {x}, \mathbf {y} <   i \cdot w} = 0 | \mathbf {x}, \mathbf {y} <   i) \\ = - \sum_ {w \in \mathcal {V} \backslash \{y _ {i} \}} \log \left(1 - \sigma \left(f \left(\mathbf {x}, \mathbf {y} _ {<   i}\right) _ {w}\right)\right). \tag {8} \\ \end{array}
$$
Appendix C provides an implementation of SCONES in JAX (Bradbury et al., 2018).
During inference we search for the translation  $\mathbf{y}^*$  that ends with  $
$  and has the highest probability of being a translation of  $\mathbf{x}$ :
$$
\begin{array}{l} \mathbf{y}^{*} = \operatorname *{arg  max}_{\mathbf{y}\in \{\mathbf{w}:\langle \text{$s$} \rangle |\mathbf{w}\in \mathcal{V}^{*}\}}P(z_{\mathbf{x},\mathbf{y}} = 1|\mathbf{x}) \\ \stackrel {\text {E q s .} 3, 4} {=} \underset {\mathbf {y} \in \{\mathbf {w} \cdot <   / s > | \mathbf {w} \in \mathcal {V} ^ {*} \}} {\arg \max } \sum_ {i = 1} ^ {| \mathbf {y} |} \log \sigma \left(f \left(\mathbf {x}, \mathbf {y} _ {<   i}\right) _ {y _ {i}}\right). \tag {9} \\ \end{array}
$$
We approximate this decision rule with vanilla beam search. The same inference code is used for both our softmax baselines and the SCONES-trained models. The only difference is that the
Parameter	Value
Attention dropout rate	0.1
Attention layer size	512
Dropout rate	0.1
Embedding size	512
MLP dimension	2,048
Number of attention heads	8
Number of layers	6
Training batch size	256
Total number of parameters	121M
Table 1: Transformer hyper-parameters.
Language pair	#Training sentence pairs
Unfiltered	Filtered
German-English	39M	33M
Finnish-English	6.6M	5.5M
Lithuanian-English	2.3M	2.0M
Table 2: MT training set sizes.
logits from SCONES models are transformed by a sigmoid instead of a softmax activation, i.e. no summation over the full vocabulary is necessary.
Relation to noise-contrastive estimation Our SCONES loss function is related to noise-contrastive estimation (NCE) (Gutmann and Hyvarinen, 2010; Mnih and Teh, 2012) because both methods reformulate next word prediction as a multi-label classification problem, and both losses have a "positive" component for the gold label, and a "negative" component for other labels. Unlike NCE, the negative loss component  $(\mathcal{L}_{-}(\cdot))$  in SCONES does not require sampling from a noise distribution as it makes use of all tokens in the
Technically, SCONES could be written as an instance of NCE with a scaling factor  $\alpha$  and an exhaustive enumeration of negative NCE samples.
Greedy search	Beam search (beam size = 4)
de-en	en-de	fi-en	en-fi	lt-en	en-lt	de-en	en-de	fi-en	en-fi	lt-en	en-lt
Softmax	38.8	38.7	26.9	18.5	26.3	11.5	39.6	39.4	27.7	19.0	26.9	12.0
SCONES	39.9	39.1	27.6	19.5	27.7	12.5	40.3	39.8	28.4	20.0	28.9	12.6
Rel. improvement	+2.7‡	+1.2	+2.8†	+5.4‡	+5.3‡	+8.5‡	+1.7†	+0.9	+2.7†	+5.5‡	+7.4‡	+5.7
Table 3: BLEU score gains from SCONES over our NMT softmax baselines with tuned  $\alpha$  -values (Table 5). Using a paired bootstrap method (Koehn, 2004), we highlight improvements that are statistically significant either at a .05 level  $(\dagger)$  or a .01 level  $(\ddagger)$
Greedy search	Beam search (beam size = 4)
de-en	en-de	fi-en	en-fi	lt-en	en-lt	de-en	en-de	fi-en	en-fi	lt-en	en-lt
Softmax	70.44	68.08	68.93	66.16	68.52	56.68	70.78	68.48	69.56	66.44	69.20	57.61
SCONES	70.69	67.55	69.28	67.32	68.96	58.68	70.88	67.99	69.72	67.91	69.95	59.48
Table 4: BLEURT (Sellam et al., 2020) scores (BLEURT-20 checkpoint) for SCONES and our NMT softmax baselines with tuned  $\alpha$ -values (Table 5).
Language pair	α
de-en	0.5
en-de	0.5
fi-en	0.7
en-fi	1.0
lt-en	0.7
en-lt	0.9
Table 5: Values of  $\alpha$  that yield the best greedy BLEU scores on the respective development sets.
vocabulary besides the gold token. This is possible because we operate on a limited 32K subword vocabulary whereas NCE is typically used to efficiently train language models with much larger word-level vocabularies (Mnih and Teh, 2012). NCE has a "self-normalization" property (Gutmann and Hyvarinen, 2010; Pihlaja et al., 2010; Mnih and Teh, 2012; Goldberger and Melamud, 2018) which can reduce computation by avoiding the expensive partition function for distributions over the full vocabulary. To do so, NCE uses the multi-label classification task as a proxy problem. By contrast, in SCONES, the multi-label classification perspective is used to express the intrinsic uncertainty in MT and is not simply a proxy for the full softmax. Thus the primary motivation for SCONES is not self-normalization over the full vocabulary.
In this work our focus is to compare NMT models trained with SCONES with well-trained standard softmax-based models. Thus we keep our setup simple, reproducible, and computationally economical. We trained Transformer models (Table 1) in six translation directions - German-English (deen), Finnish-English (en-fi), Lithuanian-English (lt-en), and the reverse directions - on the WMT19
(Barrault et al., 2019) training sets as provided by TensorFlow Datasets. We selected these language pairs to experiment with different training set sizes (Table 2). The training sets were filtered using language ID and simple length-based heuristics, and split into subwords using joint 32K SentencePiece (Kudo and Richardson, 2018) models. All our models were trained until convergence on the development set (between 100K and 700K training steps) using the LAMB (You et al., 2020) optimizer in JAX (Bradbury et al., 2018). Our softmax baselines are trained by minimizing cross-entropy without label smoothing. Our multi-way NMT models are trained by minimizing the SCONES loss function from Sec. 2, also without label smoothing. We evaluate our models on the WMT19 test sets (Barrault et al., 2019) with SacreBLEU (Post, 2018), using the WMT18 test sets as development sets to tune  $\alpha$ .
4.1 Translation quality
Table 3 compares our SCONES-based NMT systems with the softmax baselines when  $\alpha$  is tuned based on the BLEU score on the development set (Table 5). SCONES yields consistent improvements across the board. For four of six language pairs (all except en-de and fi-en), SCONES with greedy search is even able to outperform the softmax models with beam search. The language pairs with fewer resources (fi  $\leftrightarrow$  en, lt  $\leftrightarrow$  en) benefit from SCONES training much more than the high-resource language pairs (de  $\leftrightarrow$  en). SCONES still yields gains for all language directions except
$^{5}$ https://www.tensorflow.org/datasets/catalog/wmt19Translate
6Comparable to http://wmt.ufal.cz/.
Figure 2: BLEU scores as a function of GPU decoding speeds (median over five runs) for softmax and SCONES with greedy search and beam search with beam sizes 2, 4, and 8 (annotated with  $\times$ ).
English-German when we use BLEURT (Sellam et al., 2020) instead of BLEU as the evaluation measure (Table 4).
4.2 Decoding speed
Our softmax-based models reach their (near) optimum BLEU score with a beam size of around 4. Most of our SCONES models can achieve similar or better BLEU scores with greedy search. Replacing beam-4 search with greedy search corresponds to a  $3.9\mathrm{x}$  speed-up  $(2.76\rightarrow 10.64$  sentences per second) on an entry-level NVIDIA Quadro P1000 GPU with a batch size of 4.7 Fig. 2 shows the BLEU scores for all six translation directions as a function of decoding speed. Most of the speed-ups are due to choosing a smaller beam size and not due to SCONES avoiding the normalization over the full vocabulary. We expect further speed-ups when comparing models with larger vocabularies.
4.3 Mitigating the beam search curse
One of the most irksome pathologies of traditional softmax-based NMT models is the "beam search curse" (Koehn and Knowles, 2017): larger beam sizes improve the log-probability of the translations, but the translation quality gets worse. This happens because with large beam sizes, the model prefers translations that are too short. This phenomenon has been linked to the local normalization in sequence models (Sountsov and Sarawagi, 2016; Murray and Chiang, 2018) and poor model calibration (Kumar and Sarawagi, 2019). Stahlberg and Byrne (2019) showed that modes are often empty and suggested that the inherent bias of the model towards short translations is often obscured by beam search errors. Stahlberg et al. (2022) provided strong evidence that this length deficiency is due to the intrinsic uncertainty of the MT task. Given that models trained with SCONES explicitly take into account inherent uncertainty, we ran an experiment to determine whether these models are
<sup>7</sup>As an additional optimization, our greedy search implementation operates directly on the logits without applying the output activations.
Figure 3: German-English BLEU score as a function of beam size.
Figure 4: German-English length ratio (hypothesis length / reference length) as a function of beam size.
more robust to the beam search curse compared to softmax trained models.
Fig. 3 plots the BLEU score as a function of the beam size. The sharp decline of the green curve for large beam sizes reflects the beam search curse for the softmax baseline. SCONES seems to be less affected at larger beam sizes, particularly for small  $\alpha$ -values: the BLEU score for SCONES with  $\alpha = 0.2$  (solid purple curve) is stable for beam sizes greater than 100. Fig. 4, which displays the length ratio (the hypothesis length divided by the reference length) versus beam size, suggests that the differences in BLEU trajectories are partly due to translation lengths. Translations obtained using softmax become shorter at higher beam sizes whereas for SCONES with  $\alpha = 0.2$ , there is no such steep decrease in length. To study the impact of  $\alpha$  in the absence of beam search errors we ran the exact depth-first search algorithm of Stahlberg and Byrne (2019) to find the translation with global highest probability. $^{8}$  The adequacy of the transla
Figure 5: German-English BLEU scores and length ratios (hypothesis length / reference length) for exact search.
Figure 6: Mean and standard deviation (error bars) of log-probabilities of the global highest probability translations (found using exact search) and the empty translations for German-English.
tions found by exact search depends heavily on  $\alpha$  (Fig. 5). With exact search, small  $\alpha$ -values yield adequate translations, but  $\alpha \approx 1.0$  performs similar to the softmax baseline: the BLEU score drops because hypotheses are too short. Table 6 shows that SCONES with  $\alpha = 0.2$  consistently outperforms the softmax baselines by a large margin with exact search. Fig. 6 sheds some light on why SCONES with small  $\alpha$  does not prefer empty translations. A small  $\alpha$  leads to a larger gap between the log-probabilities of the exact search translation and the empty translation that arises from higher log-probabilities for the exact-search translation along with smaller variances. Intuitively, a small  $\alpha$  reduces the importance of the negative loss component  $\mathcal{L}_{-}(\cdot)$  in Eq. 6, and thus biases each binary classifier towards predicting the true label.
4.4 Reducing the number of beam search errors
Fig. 7 displays the percentage of beam search errors, the fraction of sentences for which beam
The maximum number of explored states per sentence was set to 1M. This threshold was reached for less than  $1.45\%$  of the German-English sentences. See Appendix A for other language directions.
Beam search (beam size = 4)	Exact search
de-en	en-de	fi-en	en-fi	lt-en	en-lt	de-en	en-de	fi-en	en-fi	lt-en	en-lt
Softmax	39.6	39.4	27.7	19.0	26.9	12.0	23.7	15.6	16.7	10.1	14.2	7.1
SCONES (α = 0.2)	39.3	38.9	27.7	19.6	27.9	12.7	39.1	37.2	26.7	18.7	25.6	12.1
Table 6: BLEU scores of beam search and exact search for all six translation directions.
Figure 7: Number of beam search errors for German-English as a function of the beam size.
search did not find the global best translation, as a function of beam size. We confirm the findings of Stahlberg and Byrne (2019) for softmax models: the percentage of search errors remains at a relatively high level of around  $20\%$  even for very large beam sizes. Increasing the beam size is most effective in reducing the number of search errors for SCONES with a small value of  $\alpha$ . However, a small  $\alpha$  does not always yield the best overall BLEU score (Fig. 3). Taken together, these observations provide an insight into model errors in NMT: If we describe the "model error" as the mismatch between the global most likely translation and an adequate translation (following Stahlberg and Byrne (2019)), a small  $\alpha$  would simultaneously lead to both fewer search errors (Fig. 7) and fewer model errors (Tab. 6). Counter-intuitively, however, BLEU scores peak at slightly higher  $\alpha$ -values (Tab. 5). A more sophisticated notion of model errors and search errors is needed to understand the complex inherent biases of beam search for neural sequence-to-sequence models.
5 Experiments with synthetic language pairs
Our main motivation for SCONES is to equip the model to naturally represent intrinsic uncertainty, i.e. the existence of multiple correct target sentences for the same source sentence. To examine the characteristics of SCONES as a function of uncertainty, we generated synthetic language pairs that differ by the level of ambiguity. For this pur