Language Model Based Target Token Importance Rescaling for Simultaneous Neural Machine Translation
Year: 2023

Abstract:
The decoder in simultaneous neural machine translation receives limited information from the source while having to balance the opposing requirements of latency versus translation quality. In this paper, we use an auxiliary targetside language model to augment the training of the decoder model. Under this notion of target adaptive training, generating rare or difficult tokens is rewarded which improves the translation quality while reducing latency. The predictions made by a language model in the decoder are combined with the traditional cross entropy loss which frees up the focus on the source side context. Our experimental results over multiple language pairs show that compared to previous state of the art methods in simultaneous translation, we can use an augmented target side context to improve BLEU scores significantly. We show improvements over the state of the art in the low latency range with lower average lagging values (faster output). 1

heading: 1 Introduction
text: Simultaneous Machine Translation (SiMT; Grissom II et al. (2014); Cho and Esipova (2016)) is a special case of neural machine translation (NMT; Vaswani et al. (2017)) that aims to produce real time-translations in the target language from a streaming input in the source language. The cornerstone of this task, as well as a key challenge, is the trade-off between the translation quality and the latency in producing the translations. This balance is ensured by a fixed (Ma et al., 2019; Elbayad et al., 2020) or adaptive (Arivazhagan et al., 2019; Ma et al., 2020; Zhang and Feng, 2022b) read/write policy that determines whether to wait for the next source token (a READ action) or to generate a translation (a WRITE action). Adaptive policies dynamically predict the action based on
1github.com/sfu-natlang/target_rescale_siMT ∗Equal contribution. Listing order is random. Work
performed while AJ was visiting SFU Natlang Lab.
the current source and target contexts (Zheng et al., 2020). Although adaptive policies achieve a better latency/BLEU trade-off, they often fail to account for the varying importance of different tokens when deciding a READ/WRITE action.
In Figure 1 (top), there is a negative correlation between the normalized negative log likelihoods of output tokens as measured by MMA (a SiMT model with an adaptive policy; Ma et al. 2020) versus a left-to-right language model (LM). This reflects a translation which the SiMT model is confident about, but which the LM regards as poor English (possibly due to the semantic mismatch
341
evident in “warm in winter”). Since a simultaneous policy can only access partial source context, its outputs are likely to reflect imperfect guesses such as these, particularly when translating in realtime between language pairs with different word orderings (Subject-Object-Verb) and very long compounds. As a result, training objectives which treat all translated tokens with equal importance are suboptimal.
In the context of translation, content words are generally considered more informative than function words (Chen et al., 2020). This is because content words carry the main semantic and lexical meaning of a sentence, while function words provide grammatical context and help to convey the syntactic structure of a sentence. Similarly, highfrequency words that are easier for the translation model to generate may sometimes carry less information than the desirable low-frequency (rare) words that the model struggles to generate (Chen et al., 2017). To this end, Zhang et al. (2022b) proposed to leverage conditional mutual information (MI) to estimate the weight-coefficients between the source and target to reweigh the importance of each target token. However, such an approach hasn’t been explored to address simultaneous or streaming MT to the best of our knowledge, and the lack of a complete source context makes the adaptation of this method to SiMT non-trivial. To improve simultaneous MT, Alinejad et al. (2018) proposed a prediction mechanism on the source side to get future information and aid the lack of information on target-side for translation. Instead of directly predicting a source token, Zhang and Feng (2022a) predict its aligned future-position for a given target token to guide its policy. On the other hand, Zhang and Feng (2022b) and Zhang et al. (2022a) explored policies that assign varying importance to source/target tokens based on their level of information, with more informative tokens having a greater influence on the model.
In this paper, we propose a technique to alleviate this problem in SiMT using an information theoretic approach and an adaptive training paradigm. Inspired by the recent work in using pointwise mutual information for guiding the decoder in fullsentence (non-simultaneous) translation (Lee et al., 2022), we differentiate the importance of various target tokens by their dependence on the source sentence. As shown in Figure 1 (bottom), to guide our simultaneous translation model, we incorporate
a language model that provides an additional signal indicating the importance of each target token or sentence. This target-context aware estimation leverages the relative probabilities of the translation model and language model to guide the generation process by explicitly re-weighting the training loss of each target token in the translation. Experiments show the strength of our simple method, outperforming several strong baselines in terms of both latency and BLEU scores. We perform exhaustive analysis to show that our model performs particularly well on translating low frequency words and longer sentences.

heading: 2 Background
text: Target adaptive training (Lin et al., 2017) in NMT addresses the token imbalance problem (Gu et al., 2020). While a translation model is conventionally trained with conditional maximum likelihood estimation or cross-entropy:
LCE(x,y) = − N∑
j=1
log p (yj | y<j ,x) (1)
adaptive training rescales this objective by assigning static or dynamic weights to further guide the translation model:
Ladapt(x,y) = − N∑
j=1
wj log p (yj | y<j ,x) (2)
Frequency based approaches (Gu et al., 2020; Xu et al., 2021) to assign these weights are promising but maintaining a frequency count can be an expensive overhead and would not be directly transferable to a simultaneous setting. More recently, Zhang et al. (2022b) proposed to leverage pointwise mutual information (MI) to estimate the weight-coefficients between the source x and target y as :
MI(x,y) = log
( p(x,y)
p(x).p(y)
) (3)
which can reflect the importance of target tokens for translation models.
Monotonic Infinite-Lookback Attention Arivazhagan et al. (2019) models a Bernoulli variable to make the READ or WRITE decision at every time step, while processing the input sequence incrementally. Ma et al. (2020) present monotonic
multihead attention to extend this policy to the multihead attention of transformers. For each encoder state in MMA, every head in the cross-attention of each decoder layer produces a probability pi,j that dictates if it should write target token yj while having read till the source token xi, or wait for more inputs. This is computed using the softmax energy:
energyi,j =
( mjW K ( si−1WQ )T √ dk )
i,j
pi,j = Sigmoid ( energyi,j
) (4)
where m signifies the encoder states, W the input projection matrix for query Q and key K, and dk is the dimension of the attention head. The probability pi,j is then used to parameterize the Bernoulli random variable:
bi,j ∼ Bernoulli (pi,j) (5)
If bi,j = 1 then the model performs a WRITE action on yj based on previous source tokens, otherwise it performs a READ.
Our method is based on MMA and we use it as our main simultaneous policy. To mitigate the negative impact of outlier heads2 on the read/write path, we have made slight modifications to MMA to ensure more stable performance. Instead of allowing the heads in each decoder layer to independently determine the READ/WRITE action, we now share the READ/WRITE action between the decoder layers. This adjustment helps to avoid outlier heads that could potentially disrupt the system performance and stability (Indurthi et al., 2022).

heading: 3 Approach

heading: 3.1 Target-context Aware Information Quotient
text: Inspired by Lee et al. (2022), we leverage the pointwise mutual information (MI) between each target token and its source context under the condition of previous target context. For a target token yj and the streaming source context x ≤ i, factoring in the partially constructed target prefix y < j gives the target information quotient (TIQ) is calculated
2In MMA, every head in the transformer multihead attention independently decides its read/write action and has access to all previous encoder states. The write action only takes place when the slowest head has arrived to a write decision.
as:
TIQ (yj) = log
( p (yj ,x≤i | y<j)
p (yj | y<j) · p (x≤i | y<j)
)
= log
( p (yj | x≤i,y<j) · p (x≤i | y<j)
p (yj | y<j) · p (x≤i | y<j)
)
= log
( p (yj | x≤i,y<j)
p (yj | y<j)
)
= log
( pSiMT (yj)
pLM (yj)
)
(6)
where pSiMT(.) is the simultaneous translation model probability and pLM(.) is the auxiliary targetside language model of the same size as the translation decoder. By decomposing the conditional joint distribution, this can be formalized as the log quotient of the streaming translation model probability and target language model probability. This captures the information of a target token conditioned on the target context and uses it to rescale the loss, thereby making the model pay more attention to more “informative" words.
To incorporate weights into the adaptive training objective (equation 2), two separate weights are used:
Token-level weight is used to determine weights of loss from each target token yj and streaming source context, conditioned on the obtained partial translation at the current timestep. We use a token TIQ measure and normalise it to reduce variance:
TIQtok = ( TIQ(yj)− µtok ) /σtok (7)
where µtok, and σtok are the mean and standard deviation of TIQ(yj) respectively, for every sentence.
Sentence-level weight on the other hand, is token-level TIQ is aggregated and averaged over the target sentence length |y|:
TIQsen =
  1
|y|
|y|∑
j=1
TIQ(yj)− µsen   /σsen (8)
where µsen, and σsen are the mean and standard deviation of TIQ(yj) respectively, over a batch.
The final rescaling factor to assign weights in equation 2 is calculated as:
wj = (λtokTIQtok + 1) · (λsenTIQsen + 1) (9)
The rescaling allows the model to learn the source side information for a particular target token yj , while also factoring in the target context so far.
Given that information (from source) is constrained in the nature of this task, this additional signal of the target context acts as reinforcement for translation. The likelihood score from the LM should serve to strengthen the predictive capability of the decoder. Frequent words would have a higher LM score and therefore a smaller weight wj . On the other hand, rare words would be scored lower by the LM, and thus have a higher rescaling weight wj , allowing the model to focus on them more.

heading: 3.2 Final Training Objective with Adaptive Weights and Latency Constraints
text: In MMA models, following Ma et al. (2020), we use the weighted average of differentiable lagging metric C (Arivazhagan et al., 2019) over all the attentions heads as the weighted average latency Lavg constraint3.
The MMA model uses both these loss terms in its final loss, with the hyperparameters λavg and λvar respectively. Combining the latency average loss and the target-context aware information quotient, the final training objective for our model is:
LMMA+TC = Ladapt(TIQ) + λavgLavg (10)
where Ladapt is adaptive cross-entropy loss from equation 2 with TIQ (equation 9) as its rescalingweight, and λavg is a hyperparameter to control the latency constraint.
3Early experiments with other policies such as GMA and Wait-info showed the approach to be ineffective. The explicit latency loss in MMA is crucial for the working of target adaptive training for simultaneous MT.

heading: 4 Experiments

heading: 4.1 Data
text: IWSLT’15 English ↔ Vietnamese (133K pairs) with TED tst2012 (1553 pairs) as validation set and TED tst2013 (1268 pairs) as test set. The vocabulary sizes of English and Vietnamese are 17K and 7.7K respectively.
IWSLT’14 English ↔ German (160K pairs) with validation set and test set of 7283 and 6750 pairs respectively. The vocabulary size of German is 13.5K and 9.98K for English.

heading: 4.2 Baselines and Model Settings
text: The following are the main baselines we compare our method against:
Offline Transformer (Vaswani et al., 2017) model for full-sentence translation.
Wait-k policy (Ma et al., 2019) which is a fixedpolicy that reads k source tokens initially, and then alternates between reading and writing.
Efficient Wait-k (Elbayad et al., 2020) uses multiple k’s to train a Wait-k model and relieves the constraint of test k being equal to train k.
Monotonic Multihead Attention (MMA; Ma et al. (2020)) extends infinite lookback attention (Arivazhagan et al., 2019) to all the Transformer heads.
Wait-Info (Zhang et al., 2022a) quantifies source and target token info to decide R/W action.
We also juxtapose our method against several other baselines on the En → Vi direction:
Gaussian Multihead Attention (GMA; Zhang and Feng (2022a)) that predicts the aligned source position for a target token and rescales attention with a gaussian distribution centred at this position.
ITST (Zhang and Feng, 2022b) finds the optimal information transport between source and target.
Adaptive Wait-k(Zheng et al., 2020) dynamically chooses an optimal k in the wait-k policy at every step.
MoE Wait-k (Zhang and Feng, 2021b) uses attention heads as experts trained with different k with the wait-k policy.
MMA+TC (ours) is the proposed MMA model with target context aware adaptive training objective. We use an auxiliary target-side LM decoder of the same configuration as the MT decoder. Note that the LM is only used during training and discarded at test time. We do not use extra data.
The implementation of our method is based on fairseq (Ott et al., 2019). Following MMA, we use transformer (Vaswani et al., 2017) with 6 encoder and decoder layers and 4 monotonic attention heads for the IWSLT datasets En↔Vi, De↔En. All baselines are trained with same configurations and are trained with 16k tokens. Our auxiliary language model follows the decoder settings in the model.

heading: 4.3 Evaluation
text: We evaluate using BLEU (Papineni et al., 2002) for translation quality and Average Lagging (AL) (Ma et al., 2019) for latency. AL denotes the lagging behind the ideal policy (Wait-0). Other metrics used are Average Proportion (AP) (Cho and Esipova, 2016) and Differentiable Average Lagging (DAL) (Arivazhagan et al., 2019). Given a read/write policy gi, AL is :
AL = 1
τ
τ∑
i=1
gi − i− 1 |y|/|x| (11)
where τ = argmaxi (gi = |x|), |x| and |y| are source sentence and target sentence lengths respectively.

heading: 5 Results
text: Figure 2 shows the comparison of BLEU vs. Latency (in terms of Average Lagging) of our method against previous methods on the IWSLT’15 Vi → En and IWSLT’14 En ↔ De directions. For Vi → En, we observe a significant improvement in the BLEU scores at the same latencies, compared to the baselines. We also reach the offline translation quality in low AL on this dataset. In the En → De, De → En directions too, there is a boost in the translation quality, more noticeably for lower latencies. The plots show that our method boosts translation quality in the earlier latencies and the effect of reweighing is more pronounced in these regions, where the source context is more limited. In higher latency regions, when the source information window increases, the other baselines start to reach our BLEU score in the English-German directions.
In Figure 3, we compare against several stateof-the-art methods on the En → Vi. Our method gets better translation quality compared all others, in the low-latency zone, matching the offline score at 3.86 AL. We show the BLEU vs. AL plot in a low latency range to compare performance in the more challenging area of this task, the low latency points.

heading: 6 Analysis

heading: 6.1 Token-level vs. Sentence-Level Weight
text: Ablation Study The two hyperparameters in our method are Sentence-Level Weight and TokenLevel Weight, which determine the sentence and token-level effect of rescaling with LM. In Fig. 5
we report the BLEU scores with different hyperparameter settings on Vi-En. (AL across the table are similar as experiments are done with the same λ). We set the values of these hyperparameters to 0.2 in all our experiments.

heading: 6.2 Effect on Low-frequency Words
text: With reweighing loss using the Language Model likelihood, we aim to reduce the effect of frequency imbalance in the corpus on training. We compare our translations against MMA on rare and frequent words. In addition to an overall BLEU improvement, we also see an improvement in the F-measure of rare words. As shown in Figure 4, our method does better on extremely rare words (freq ≤ 10). Table 1 shows that while the baseline overfits to the most frequent words, our method captures rare words, from the bottom two frequency bins (50- 70% and 70-100%), better. The results show that our method makes the model train better on rare words and remedy the effect of token imbalance.
Content word occurrences. Zhang et al. (2022a) show that focusing on the right content words in the target is crucial to getting the necessary target information in a subcutaneous translation setting. Following Moradi et al. (2019) we inspect the content words generated by our model using spacy to get POS tags over the translations. As evident from Table 2, our model recovers more content words in the translations wrt the reference.

heading: 6.3 Effect on Translation Length
text: Following the rationale of Lakew et al. (2019) in NMT and Zhang and Feng (2022d) in simultaneous translation, we inspect the translation quality of our model on varying target sentence lengths in Figure 8 and observe that our method shows a big improvement in BLEU on the longer sentences. Our method prevents the model from over-producing words (as seen in the top figure in Figure 8). We hypothesize that this is because the model does not generate as many words (and overuse them) from the most frequent word bin (see Table 1, top 10% bin) as MMA. Our target sentence lengths are consistently less than MMA’s and are closer to the ground truth sentence lengths (as shown in the bin 0, Fig. 8 (top)).

heading: 6.4 Effect on Translation Paths
text: Attention Heatmaps and READ/WRITE Sequences. Figure 6 compares attention heatmaps from MMA and MMA+TC (our method) on the Vi → En direction. As evident, our method performs READ actions in smaller intervals between predicting consecutive WRITE actions.
Consider the READ/WRITE actions generated by MMA and MMA+TC for the given source sen-
tence are:
Src: Chúng tôi còn vt mu đây . Nó còn khá m .
MMA: RRR W RRR WWW RRR WW RRR W RR WWWWW
Ours: RRR W RRR WW RR W R WW RRR W R W R WWWW
In this example, MMA reads more than required for a write in certain places. It shows that at a similar lag, our model gets a higher probability of a WRITE action, compared to MMA, after having read the same number of source words.
Sufficiency of the READ actions. Zhang and Feng (2022c) introduce a metric of sufficiency ASuf in Read/Write paths with the notion that too
many but not all necessary READs would result in high latency while few but not sufficient READ actions would exclude needed information and could cause poor translation quality. When the ground truth aligned source position of the jth target word is denoted by aj , and the number of source words read when writing target jth word is denoted by rj :
ASuf = 1
|y|
|y|∑
j=1
1aj≤rj (12)
We compare our method against MMA and WaitInfo on AL=4 with the sufficiency metric. Using equation(12) across sentences of varying lengths, we evaluate the read-write paths of each model, against reference alignments from Eflomal (Östling and Tiedemann, 2016)4. In Figure 7, we can see a clearly increasing and higher score on sufficiency as compared to the baselines - Wait-Info and MMA. This signifies that our target-context augmented training helps the model read sufficient source tokens required for producing a translation, while maintaining the same latency as others, showing that the model learns and correctly gauges the information it requires to translate a target token, and
4We use the Eflomal library to get alignment priors from IWSLT’15 Vi-En train set, and use them to generate alignments for the test set. https://github.com/ robertostling/eflomal
makes READ actions accordingly.
Ratio of Aligned READ actions. We compare MMA and our Read-Write policy against the reference source-target alignments by computing the overlap between the hard alignments and the translation path for all output translations :
IoUa,r = ∑
i=1
( intersection(ai, ri)
union(ai, ri)
) (13)
where ai is the reference alignment matrix for the ith sentence, made by setting all aligned source positions to 1 and ri is the upper triangular matrix set to 1 using reads from the policy.5 The IoU scores for our policy and for MMA are shown in Figure 8 (bottom) with varying sentence lengths. Our policy shows a stronger adherence to the source-target monotonic alignment path.

heading: 7 Related Work
text: Simultaneous Translation. Fixed Policy methods (Ma et al., 2019; Elbayad et al., 2020) follow the fixed rule of waiting for the first k source tokens before generating a target token, and alternate thereafter. Adaptive Wait-k (Zheng et al., 2020) dynamically chooses the best k at every step. Han et al. (2020) applied meta learning in wait-k. Zhang and Feng (2021b) use each attention head as an expert of wait-k policy whereas Zhang and Feng (2021a)
5We choose this metric to show the extent to which the policy follows the source-target alignments. In an ideal setting, IoU = 1.
introduce a character level wait-k policy. But fixed policy methods aren’t feasible for complex inputs and cannot adapt to them. Full-sentence MT has also been leveraged to augment the policy with future information (Zhang et al., 2020; Alinejad et al., 2021). But using such oracle or gold (Zheng et al., 2019; Arthur et al., 2021) READ/WRITE actions does not optimize policy with translation quality. Alinejad et al. (2018) proposes providing future-information on the source side using prediction. Grissom II et al. (2014) predict unseen verbs and uses reinforcement learning to learn when to trust these predictions and when to wait for more input. In contrast, we leverage target side context to strengthen the simultaneous translations.
Zhang and Feng (2022c) train two models on either language directions and make their policies converge. Wilken et al. (2020) propose external ground-truth alignments to train the policy. Papi et al. (2023) use cross attention scores to guide policy. Infinite-lookback (Arivazhagan et al., 2019) and chunkwise (Chiu* and Raffel*, 2018) attention propose to use a soft monotonic attention over previous encoder states. We use a variant of the policy proposed by Ma et al. (2020) that adapts monotonic attention to the multihead architecture of the Transformer. GMA (Zhang and Feng, 2022a) predicts the aligned source position of the current target token and rescales attention based on it. But these methods treat all words equally during training whereas our method improves upon MMA via adaptive training.
Some recent work explores capturing and quantifying information from the source tokens and use it to model READ/WRITE actions (Zhang et al., 2022a; Zhang and Feng, 2022b). But these works do not use the target context in their information. Unlike their quantization method, we present a simple scoring by using an auxiliary target-side LM.
Adaptive Training for MT. Target adaptive objectives have been explored by (Lin et al., 2017) which uses probability of a class to scale, but actually only scale down high frequency classes; (Jiang et al., 2019) which directly uses normalized frequency count but have high variance. (Gu et al., 2020) use a chi-square and an exponential distribution function with frequency. However these use only static word frequency. BMI (Xu et al., 2021) attempt to capture mutual information between each source and target token. CBMI (Zhang et al., 2022b) incorporate target context as well, in
mutual information. However, these adaptive methods are not directly transferable to the streaming nature of our task.

heading: 8 Conclusion
text: We have presented a simple technique for rescaling target-token importance in simultaneous translation using an information theoretic approach and an adaptive training paradigm. We differentiate the importance of various target tokens by their dependence on the source sentence. To guide our simultaneous translation model, we incorporate a targetside language model that provides an additional signal indicating the importance of each target token or sentence under the condition of the previous target context. Our model shows strong performance on several datasets and outperforms several state-ofthe-art techniques in the low latency range (AL<5). Further analysis shows that our technique is better able to translate long sentences and those with rare words. We also showed that the translation path (read/write action sequence) has a stronger correlation to the source-target alignment.
Limitations and Future Work
Since our auxiliary target-side LM decoder is spawned with the same configuration as the MT decoder, this significantly adds to the model size at training time. This makes it difficult to scale/slower to train with translation models of large size. While this problem can be easily mitigated by using a GPU of larger memory, we would like to explore more efficient ways of incorporating the target context which we leave for future work. Secondly, even though our method gives a significant boost to translation quality in the early latencies, it relies on the MMA (Ma et al., 2020) policy that has some limitations in terms of latency because of a suboptimal decision making using multiple heads (Indurthi et al., 2022). While our policy shows improvement, it could be further optimized, for instance, in following reference alignments more closely which would have a positive effect on latency. Finally, using additional monolingual data is also a viable direction for future work to strengthen the language model used in the approach.

heading: Acknowledgements
text: We would like to thank the anonymous reviewers for their helpful comments. We would also like to thank Logan Born for the numerous discussions
that helped shape this paper and Dr. Angel Chang for lending us the GPU resources. The research was partially supported by the Natural Sciences and Engineering Research Council of Canada grants NSERC RGPIN-2018-06437 and RGPAS-2018522574 and a Department of National Defence (DND) and NSERC grant DGDND-2018-00025 to the third author.

heading: A Hyperparameters
text: All models were trained on 2 x Titan RTX with 24 GB memory each. An entire training run finishes within 2.5 hours with fp32 completing about 40 epochs.
B Detailed Results

heading: Future Work
text: Limitations and Future Work
Since our auxiliary target-side LM decoder is spawned with the same configuration as the MT decoder, this significantly adds to the model size at training time. This makes it difficult to scale/slower to train with translation models of large size. While this problem can be easily mitigated by using a GPU of larger memory, we would like to explore more efficient ways of incorporating the target context which we leave for future work. Secondly, even though our method gives a significant boost to translation quality in the early latencies, it relies on the MMA (Ma et al. , 2020) policy that has some limitations in terms of latency because of a suboptimal decision making using multiple heads (Indurthi et al. , 2022). While our policy shows improvement, it could be further optimized, for instance, in following reference alignments more closely which would have a positive effect on latency. Finally, using additional monolingual data is also a viable direction for future work to strengthen the language model used in the approach