Learning to Use Future Information in Simultaneous Translation

1 INTRODUCTION . Neural machine translation ( NMT ) is an important task for the machine learning community and many advanced models have been designed ( Sutskever et al. , 2014 ; Bahdanau et al. , 2014 ; Vaswani et al. , 2017 ) . In this work , we work on a more challenging task in NMT , simultaneous translation ( also known as simultaneous interpretation ) , which is widely used in international conferences , summits and business . Different from standard NMT , simultaneous NMT has a stricter requirement for latency . We can not wait to the end of a source sentence but have to start the translation right after reading the first few words . That is , the translator is required to provide instant translation based on a partial source sentence . Simultaneous NMT is formulated as a prefix-to-prefix problem ( Ma et al. , 2019 ; 2020 ; Xiong et al. , 2019 ) , where a prefix refers to a subsequence starting from the beginning of the sentence to be translated . In simultaneous NMT , we face more uncertainty than conventional NMT , since the translation starts with a partial source sentence rather than the complete one . Wait-k inference ( Ma et al. , 2019 ) is a simple yet effective strategy in simultaneous NMT where the translation is k words behind the source input . Rather than instant translation of each word , wait-k inference actually leverages k more future words during inference phase . Obviously , a larger k can bring more future information , and therefore results in better translation quality but at the cost of larger latency . Thus , when used in real-world applications , we should have a relatively small k for simultaneous NMT . While only small k values are allowed in inference , we observe that wait-m training with m > k will lead to better accuracy for wait-k inference . Figure 1 shows the results of training with wait-m but test with wait-3 on IWSLT ’ 14 English→German translation dataset . If training with m = 3 , we will obtain a 22.79 BLEU score . If we set m to larger values such as 7 , 13 or 21 and test with wait-3 , we can get better BLEU scores . That is , the model can benefit from the availability of more future information in training . This is consistent with the observation in ( Ma et al. , 2019 ) . The challenge is how much future information we should use in training . As shown in Figure 1 , using more future information does not monotonically improve the translation accuracy of wait-k inference , mainly because that more future information results in a larger mismatch between training and inference ( i.e. , m − k more words are used in training than inference ) . Besides , due to the diversity of the natural language , intuitively , using different m ’ s for different sentences will lead to better performance . Even for the same sentence pair , the optimal m for training might vary in different training stages . In this work , we propose an algorithm that can automatically determine how much future information to use in training for simultaneous NMT . Given a pre-defined k , we want to maximize the performance of wait-k inference . We have a set of M training strategies wait-m with different waiting thresholds m ( m ∈ { 1 , 2 , · · · , M } ) . We introduce a controller such that given a training sample , the controller dynamically selects one of these training strategies so as to maximize the validation performance on waitk inference . Which wait-m training strategy to select is based on the data itself and the network status of the current translation model . The controller and the translation model are jointly trained , and the learning process is formulated as a bi-level optimization problem ( Sinha et al. , 2018 ) , where one optimization problem is nested within another . Our contribution is summarized as follows : ( 1 ) We propose a new method for simultaneous NMT , where a controller is introduced to adaptively determine how much future information to use for training . The controller and the translation model are jointly learned through bi-level optimization . ( 2 ) Experiments on four datasets show that our method improves the wait-k baseline by 1 to 3 BLEU scores , and also consistently outperforms several heuristic baselines leveraging future information . 2 RELATED WORK . Previous work on simultaneous translation can be categorized by whether using a fixed decoding scheduler or an adaptive one . Fixed policies usually use pre-defined rules to determine when to read or to write a new token ( Dalvi et al. , 2018 ; Ma et al. , 2019 ) . Wait-k is the representative method for fixed scheduler ( Ma et al. , 2019 ) , where the decoding is always k words behind the source input . Wait-k achieves good results in terms of translation quality and controllable latency , and has been used in speech-related simultaneous translation ( Zhang et al. , 2019 ; Ren et al. , 2020 ) . For methods that use adaptive schedulers , Cho & Esipova ( 2016 ) proposed wait-if-worse ( WIW ) and wait-if-diff ( WID ) methods which generate a new target word if its probability does not decrease ( for WIW ) or the generated word is unchanged ( for WID ) after reading a new source token . Grissom II et al . ( 2014 ) and Gu et al . ( 2017 ) used reinforcement learning to train the read/write controller , while Zheng et al . ( 2019a ) obtained it in a supervised way . Alinejad et al . ( 2018 ) added a “ predict ” operator to the controller so that it can anticipate future source inputs . Zheng et al . ( 2019b ) introduced a “ delay ” token into the target vocabulary indicating that the model should read a new word instead of generating a new one . Arivazhagan et al . ( 2019 ) proposed monotonic infinite lookback attention ( MILk ) , which first used a hard attention model to determine when to read new source tokens , and then a soft attention model to perform translation . Ma et al . ( 2020 ) extended MILk into a multi-head version and proposed monotonic multihead attention ( MMA ) with two variants : MMA-IL ( Infinite Lookback ) which has higher translation quality by looking back at all available source tokens , and MMA-H ( ard ) which is more computational efficient by limiting the attention span . Besides , Zheng et al . ( 2020a ) extended wait-k to an adaptive strategy by training multiple wait-m models with different m ’ s and adaptively selecting a decoding strategy during inference . Zheng et al . ( 2020b ) explored a new setting , where at each timestep , the translation model over generates the target words and corrects them in a timely fashion . 3 PROBLEM FORMULATION AND BACKGROUND . In this section , we first introduce the notations used in this work , followed by the formulation of wait-k strategy , and then we introduce our network architecture adapted from ( Ma et al. , 2019 ) . 3.1 NOTATIONS AND FORMULATION . Let X and Y denote the source language domain and target language domain . For any x ∈ X and y ∈ Y , let xi and yi denote the i-th token in x and y respectively . Lx and Ly are the numbers of tokens in x and y. x≤t represents a prefix of x , which is the subsequence x1 , x2 , · · · , xt , and similarly for y≤t . Let Dtr and Dva denote the training and validation sets , both of which are collections of bilingual sentence pairs . The wait-k strategy ( Ma et al. , 2019 ) is defined as follows : given an input x ∈ X , the generation of the translation y is always k tokens behind reading x . That is , at the t-th decoding step , we generate token yt based on x≤t+k−1 ( more strictly , x≤min { t+k−1 , Lx } ) . Our goal is to obtain a model f : X 7→ Y with parameter θ that can achieve better results with wait-k inference . 3.2 MODEL ARCHITECTURE . Our model for simultaneous NMT is based on Transformer model ( Vaswani et al. , 2017 ) . The model includes an encoder and a decoder , which are used for incrementally processing the source and target sentences respectively . Both the encoder and decoder are stacked of L blocks . We mainly introduce the differences compared with the standard Transformer . ( 1 ) Incremental encoding : Let hlt denote the output of the t-th position from block l. For ease of reference , let H l≤t denote { hl1 , hl2 , · · · , hlt } , and let h0t denote the embeddings of the t-th token . An attention model attn ( q , K , V ) , takes a query q ∈ Rd , a set of keys K and values V as inputs . K and V are of equal size , q ∈ Rd where d ∈ Z+ is the dimension , ki ∈ Rd and vi ∈ Rd are the i-th key and value . attn is defined as follows : attn ( q , K , V ) = |K|∑ i=1 αiWvvi , αi = exp ( ( Wqq ) > ( Wkki ) ) Z , Z = |K|∑ i=1 exp ( ( Wqq ) > ( Wkki ) ) , ( 1 ) where W ’ s are the parameters to be optimized . In the encoder side , hlt is obtained in a unidirectional way : hlt = attn ( h l−1 t , H l−1 ≤t , H l−1 ≤t ) . That is , the model can only attend to the previously generated hidden representations , and the computation complexity is O ( L2x ) . In comparison , ( Ma et al. , 2019 ) still leverages bidirectional attention , whose computation complexity is O ( L3x ) . We find that unidirectional attention is much more efficient than bidirectional attention without much accuracy drop ( see Appendix D.1 for details ) . ( 2 ) Incremental decoding : Since we use wait-k strategy , the decoding starts before reading all inputs . At the t-th decoding step , the decoder can only read x≤t+k−1 . When t ≤ Lx−k , the decoder greedily generates one token at each step , i.e. , the token is yt = argmaxw∈V P ( w|y≤t−1 ; HL≤t+k−1 ) , where V is the vocabulary of the target language . When t > Lx − k , the model has read the full input sentence and can generate words using beam search ( Ma et al. , 2019 ) . 4 OUR METHOD . We first introduce our algorithm in Section 4.1 , and then we discuss its relationship with several other heuristic algorithms that leverage future information in Section 4.2 . 4.1 ALGORITHM . Let f ( · · · ; θ ) denote a translation model parameterized by θ , and let ϕ denote the controller parameterized by ω to guide the training process of f . f ( · · · ; θ∗ ( ω ) ) is the translation model obtained under the guidance of the controller ϕ ( · · · ; ω ) , where θ∗ ( ω ) is the corresponding parameter . For each training data ( x , y ) , the controller ϕ adaptively assigns a training task wait-m , where m ∈ { 1 , 2 , · · · , M } , andM ∈ Z+ is a pre-defined hyperparameter . The input of ϕ consists of two parts : ( i ) the information of the training data ( x , y ) ; ( ii ) the network status of the translation model f . For ease of reference , denote these input features as Ix , y , f . We will discuss how to design Ix , y , f in Section 5.1 . LetMk ( Dva ; θ∗ ( ω ) ) denote the validation metric , which is evaluated on the validation set Dva with model f ( · · · ; θ∗ ( ω ) ) and wait-k inference . We formulate the training process of f and ϕ as a bi-level optimization , where two optimization problems are nested together . In the inner-optimization , given a ω , we want to obtain the model f ( · · · , θ∗ ( ω ) ) that can minimize the loss function ` on the training set Dtr under the guidance of the controller ϕ ( · · · , ω ) . In the outer-optimization , given a translation model θ∗ ( ω ) , we optimize ω to maximize the validation performance Mk . The mathematical formulation is shown as follows : max ω Mk ( Dva ; θ∗ ( ω ) ) ; s.t . θ∗ ( ω ) = argmin θ 1 |Dtr| ∑ ( x , y ) ∼Dtr Em∼ϕ ( Ix , y , f ; ω ) ` ( x , y , m ; θ ) ; where ` ( x , y , m ; θ ) = ∑ ( x , y ) logP ( y|x ; θ , m ) = ∑ ( x , y ) |y|∑ t=1 logP ( yt|y≤t−1 , x≤t+m−1 ) . ( 2 ) We optimize Equation 2 in an alternative way , where we first optimize θ with a given ω , and then update ω using the REINFORCE algorithm . We repeat the above process until convergence . Details are in Algorithm 1 : Algorithm 1 : The optimization algorithm . 1 Input : Training episode E ; internal update iterations T ; learning rate ηθ of the translation model ; learning rate ηω of the controller ; batch size B ; initial parameters ω , θ ; 2 for e← 1 : E do 3 Init a buffer to store states and actions : B = { } ; 4 for t← 1 : T do 5 Randomly sample a mini-batch of data De , t from Dtr with batch size B ; 6 Assign a wait-m task to each data : D̃ = { ( x , y , m ) | ( x , y ) ∈ De , t , m ∼ ϕ ( Ix , y , f ; ω ) } , where the batch size is B , and m is sampled from to the output distribution of ϕ ; 7 Update the buffer : B ← B ∪ { ( Ix , y , f , m ) | ( x , y , m ) ∈ D̃ } ; 8 Update the translation model : θ ← θ − ( ηθ/B ) ∇θ ∑ ( x , y , m ) ∈D̃ ` ( x , y , m ; θ ) ; 9 Calculate the validation performance as the reward : Re =Mk ( Dva ; θe , T ) ; 10 Update the controller : ω ← ω + ηωRe ∑ ( I , m ) ∈B∇ω logP ( ϕ ( I ; ω ) = m ) . 11 Return θ. Algorithm 1 consists of E episodes ( i.e. , the outer loop ) , and each episode consists of T update iterations ( i.e. , the inner loop ) . The inner loop ( from step 4 to step 8 ) aims to optimize the θ , where we can update the parameter with any gradient based optimizer like momentum SGD , Adam ( Kingma & Ba , 2015 ) , etc . The outer loop ( from step 2 to step 10 ) aims to optimize ω. ϕ ( Ix , y , f ; ω ) can be regarded as a policy network , where the state is Ix , y , f , the action is the choice of the task wait-m , m ∈ { 1 , 2 , · · · , M } , and the reward is the validation performance Re ( step 9 ) . At the end of each episode , we update ω using REINFORCE algorithm ( step 10 ) .

Summary: This paper proposes a new training method for wait-k simultaneous translation. Rather than training on prefix pairs where the target prefix lags the source by k tokens, it uses an RL controller to determine an optimal lag for each sentence pair. The controller uses a small set of features intended to capture training progress, and is trained with REINFORCE to minimize wait-k loss on a validation set, in alternation with main training steps. This method shows consistent gains over various wait-k training heuristics, and some gains over other approaches that adapt the lag at inference time.