Attention Forcing for Sequence-to-sequence Model Training

1 INTRODUCTION . Auto-regressive sequence-to-sequence ( seq2seq ) models with attention mechanism are widely used in a variety of areas including Neural Machine Translation ( NMT ) ( Neubig , 2017 ; Huang et al. , 2016 ) and speech synthesis ( Shen et al. , 2018 ; Wang et al. , 2018 ) , also known as Text-To-Speech ( TTS ) . These models excel at connecting sequences of different length , but can be difficult to train . A standard approach is teacher forcing , which guides a model with reference output history during training . This makes the model unlikely to recover from its mistakes during inference , where the reference output is replaced by generated output . One alternative is to train the model in free running mode , where the model is guided by generated output history . This approach often struggles to converge , especially for attention-based models , which need to infer the correct output and align it with the input at the same time . Several approaches are introduced to tackle the above problem , namely scheduled sampling ( Bengio et al. , 2015 ) and professor forcing ( Lamb et al. , 2016 ) . Scheduled sampling randomly decides , for each time step , whether the reference or generated output token is added to the output history . The probability of choosing the reference output token decays from 1 to 0 with a heuristic schedule . A natural extension is sequence-level scheduled sampling , where the decision is made for each sequence instead of token . Professor forcing views the seq2seq model as a generator . During training , the generator operates in both teacher forcing mode and free running mode . In teacher forcing mode , it tries to maximize the standard likelihood . In free running mode , it tries to fool a discriminator , which is trained to tell if the model is running in teacher forcing mode . To make training stable , the above approaches require either a well tuned schedule , or a well trained discriminator . This paper introduces attention forcing , which guides the model with generated output history and reference attention . This approach makes training stable by decoupling the learning of the output and that of the alignment . There is no need for a schedule or a discriminator . Furthermore , for cascaded systems like many TTS systems , attention forcing can be particularly useful . A model trained with attention forcing can generate ( in attention forcing mode ) output sequences aligned with the references . These output sequences can be used to train a downstream model , enabling it to fix some upstream errors . The TTS experiments show that attention forcing yields significant gain in speech quality . The NMT experiments show that for tasks where various re-orderings of the output are valid , guiding the model with generated output history can be problematic , while guiding the model with reference attention yields slight but consistent gain in BLEU score ( Papineni et al. , 2002 ) . 2 SEQUENCE-TO-SEQUENCE GENERATION . Sequence-to-sequence generation can be defined as the problem of mapping an input sequence x1 : L to an output sequence y1 : T . From a probabilistic perspective , a model θ estimates the distribution of y1 : T given x1 : L , typically as a product of distributions conditioned on output history : p ( y1 : T |x1 : L ; θ ) = ∏T t=1 p ( yt|y1 : t−1 , x1 : L ; θ ) ( 1 ) Ideally , the model is trained through minimizing the KL-divergence between the true distribution p ( y1 : T |x1 : L ) and the estimated distribution : θ̂ = argmin θ Ex1 : L∼p ( x1 : L ) KL ( p ( y1 : T |x1 : L ) ||p ( y1 : T |x1 : L ; θ ) ) = argmin θ Ex1 : L∼p ( x1 : L ) Ey1 : T∼p ( y1 : T |x1 : L ) log ( p ( y1 : T |x1 : L ) /p ( y1 : T |x1 : L ; θ ) ) ( 2 ) In practice , this is approximated by minimizing the Negative Log-Likelihood ( NLL ) of some training data { y ( n ) 1 : T , x ( n ) 1 : L } N1 , sampled from the true distribution : θ̂ = argmin θ − ∑N n=1 log p ( y ( n ) 1 : T |x ( n ) 1 : L ; θ ) ( 3 ) While L and T are functions of n , the subscripts are omitted to simplify notations , i.e . Ln and Tn are written as L and T . At inference stage , given an input x∗1 : L , the output ŷ1 : T can be obtained through searching for the most probable sequence from the estimated distribution : ŷ1 : T = argmax y1 : T p ( y1 : T |x∗1 : L ; θ̂ ) ( 4 ) The exact search is computationally expensive , and is often approximated by greedy search if the output space is continuous , or beam search if the output space is discrete ( Bengio et al. , 2015 ) . 2.1 ATTENTION-BASED SEQ2SEQ MODEL . Attention mechanisms ( Bahdanau et al. , 2014 ; Chorowski et al. , 2015 ) are commonly used to connect sequences of different length . This paper focuses on attention-based encoder-decoder models . For these models , the probability p ( yt|y1 : t−1 , x1 : L ; θ ) is estimated as : p ( yt|y1 : t−1 , x1 : L ; θ ) ≈ p ( yt|y1 : t−1 , αt , x1 : L ; θ ) ≈ p ( yt|st , ct ; θy ) ( 5 ) st = f ( y1 : t−1 ; θs ) ( 6 ) ct = f ( αt , x1 : L ; θc ) ( 7 ) θ = { θy , θs , θc } . αt is an alignment vector ( a set of attention weights ) . st is a state vector representing the output history y1 : t−1 , and ct is a context vector summarizing x1 : L for the prediction of yt . The following equations , as well as figure 1 , give a more detailed illustration of how αt , st and ct can be computed : h1 : L = f ( x1 : L ; θh ) ( 8 ) st = f ( st−1 , yt−1 ; θs ) ( 9 ) αt = f ( st , h1 : L ; θα ) ( 10 ) ct = ∑L l=1 αt , lhl ( 11 ) ŷt ∼ p ( yt|st , ct ; θy ) ( 12 ) First the encoder maps x1 : L to an encoding sequence h1 : L. For each decoder time step , st is updated with yt−1 . Based on h1 : L and st , the attention mechanism computes αt , and then ct as the weighted sum of h1 : L. Finally , the decoder estimates a distribution based on st and ct , and optionally generates an output token ŷt by either sampling or taking the most probable token . Note that the output history y1 : t−1 plays an important role , as it impacts p ( yt|st , ct ; θy ) through both st and ct. Also note that there are many forms of attention-based encoder-decoder models . While attention forcing is illustrated with this particular form , it is not limited to it . 2.2 TRAINING APPROACHES . As shown in equations 2 and 3 , minimizing the KL-divergence between the true distribution and the model distribution can be approximated by minimizing the NLL . This motivates the approach to train the model in teacher forcing mode , where p ( yt|y1 : t−1 , x1 : L ; θ ) is computed with the correct output history y1 : t−1 , as shown in equations 5 and 6 . In this case , the loss can be written as : L ( T ) y ( θ ) = − ∑N n=1 log p ( y ( n ) 1 : T |x ( n ) 1 : L ; θ ) = − ∑N n=1 ∑T t=1 log p ( y ( n ) t |y ( n ) 1 : t−1 , x ( n ) 1 : L ; θ ) ( 13 ) This approach yields the correct model ( zero KL-divergence ) if the following assumptions hold : 1 ) the model is powerful enough ; 2 ) the model is optimized correctly ; 3 ) there is enough training data to approximate the expectation shown in equation 2 . In practice , these assumptions are often not true , hence the model is prone to make mistakes . To illustrate the problem , suppose there is a reference output y∗1 : T for the test input x ∗ 1 : L. Due to data sparsity in high-dimensional space , x ∗ 1 : L is likely to be unseen during training . If the probability p ( y∗t |y∗1 : t−1 , x∗1 : L ; θ ) is wrongly estimated to be small at time step t , the probability of the reference output sequence p ( y∗1 : T |x∗1 : L ; θ ) will also be small , i.e . it will be unlikely for the model to generate y∗1 : T . In practice , the model can be assessed by some loss D ( y∗1 : T , ŷ1 : T ) between the reference output y∗1 : T and the generated output ŷ1 : T . Taking the expected value yields the Bayes risk : Eŷ1 : T∼p ( y1 : T |x∗1 : L ; θ ) D ( y ∗ 1 : T , ŷ1 : T ) . This motivates training the model with the following loss : L ( B ) y ( θ ) = ∑N n=1 Eŷ1 : T∼p ( y1 : T |x ( n ) 1 : L ; θ ) D ( y ( n ) 1 : T , ŷ1 : T ) ≈ ∑N n=1 ∑M m=1 p ( ŷ ( n , m ) 1 : T |x ( n ) 1 : L ; θ ) D ( y ( n ) 1 : T , ŷ ( n , m ) 1 : T ) ( 14 ) ŷ ( n , m ) is sampled from the estimated distribution p ( y1 : T |x ( n ) 1 : L ; θ ) . D is minimal when the two sequences are equal . So the model is trained to not only assign high probability to the reference sequences in the training data , but also assign low probability to other sequences . This makes minimum Bayes risk training prone to overfitting . Very often , D is computed at sub-sequence level . Examples include BLEU score for NMT , word error rate for speech recognition and root mean square error for TTS . So if an approach trains the model to predict the reference output , based on erroneous output history , it will indirectly reduce the Bayes risk . One example is to train the model in free running mode , where p ( yt|y1 : t−1 , x1 : L ; θ ) is estimated with the generated output history : p ( yt|y1 : t−1 , x1 : L ; θ ) ≈ p ( yt|ŷ1 : t−1 , x1 : L ; θ ) ≈ p ( yt|st , ct ; θy ) ( 15 ) st = f ( ŷ1 : t−1 ; θs ) ( 16 ) ŷt is obtained from the estimated distribution p ( yt|st , ct ; θy ) , as shown in equation 12 . ( The approaches discussed in this section are designed for all auto-regressive models , with or without attention mechanism . So the realization ct is not shown . ) The corresponding loss function is : L ( F ) y ( θ ) = − ∑N n=1 ∑T t=1 log p ( y ( n ) t |ŷ ( n ) 1 : t−1 , x ( n ) 1 : L ; θ ) ( 17 ) Note that if there is enough data and modeling power , and the model is optimized correctly , the distribution ∏T t=1 p ( yt|ŷ1 : t−1 , x1 : L ; θ ) can be the same as the true distribution p ( y1 : T |x1 : L ) . The problem with this approach is that training often struggles to converge . One concern is that the model needs to learn to infer the correct output and align that with the input at the same time . Therefore , several approaches , namely scheduled sampling and professor forcing , are proposed to train the model in a mode between teacher forcing and free running . Scheduled sampling ( Bengio et al. , 2015 ) randomly decides , for each time step , whether the reference or generated output token is added to the output history ỹ1 : t−1 . For this approach , p ( yt|y1 : t−1 , x1 : L ; θ ) is estimated as : p ( yt|y1 : t−1 , x1 : L ; θ ) ≈ p ( yt|ỹ1 : t−1 , x1 : L ; θ ) ≈ p ( yt|st , ct ; θy ) ( 18 ) st = f ( ỹ1 : t−1 ; θs ) ( 19 ) ỹt = { yt with probability ŷt with probability 1− ( 20 ) gradually decays from 1 to 0 with a heuristic schedule . Considering that during training , ỹ1 : t−1 is mostly an inconsistent mixture of the reference output and the generated output , a natural extension is sequence-level scheduled sampling ( Bengio et al. , 2015 ) , where the decision is made for each sequence instead of token : ỹ1 : t−1 = { y1 : t−1 with probability ŷ1 : t−1 with probability 1− ( 21 ) This type of training improves the results of many experiments , but sometimes leads to worse results ( Wang et al. , 2017 ; Bengio et al. , 2015 ) . One concern is that the decay schedule does not fit the learning pace of the model . Professor forcing ( Lamb et al. , 2016 ) is an alternative trade-off . During training , the model θ is viewed as a generator , which generates two output sequences for each input sequence , respectively in teacher forcing mode and free running mode1 . For the training example { y ( n ) 1 : T , x ( n ) 1 : L } , let y ′ ( n ) 1 : T denote the output generated in teacher forcing mode , and ŷ ( n ) 1 : T the output generated in free running forcing mode , this can be expressed as : ∀t y′ ( n ) t ∼ p ( yt|y ( n ) 1 : t−1 , x ( n ) 1 : L ; θ ) ( 22 ) ∀t ŷ ( n ) t ∼ p ( yt|ŷ ( n ) 1 : t−1 , x ( n ) 1 : L ; θ ) ( 23 ) In addition to the final output , some intermediate output sequences are saved . Let β′ ( n ) 1 : T and β̂ ( n ) 1 : T denote the intermediate output sequences generated respectively in teacher forcing and free running mode . These generated sequences form a dataset { y′ ( n ) 1 : T , β ′ ( n ) 1 : T , ŷ ( n ) 1 : T , β̂ ( n ) 1 : T } N1 that is used to train a discriminator ψ. ψ is trained to predict the probability that a group of sequences is generated in teacher forcing mode , and the loss function is : Lψ ( ψ|θ ) = − ∑N n=1 ( log ( f ( y ′ ( n ) 1 : T , β ′ ( n ) 1 : T ; ψ ) ) + log ( 1− f ( ŷ ( n ) 1 : T , β̂ ( n ) 1 : T ; ψ ) ) ) ( 24 ) While this loss function is optimized w.r.t . ψ , it depends on θ , hence the notation ψ|θ . For the generator θ , there are three training objectives . The first one is the standard likelihood shown in equation 13 . The second one is to fool the discriminator in free running mode : L ( F ) β ( θ|ψ ) = − ∑N n=1 log ( f ( ŷ ( n ) 1 : T , β̂ ( n ) 1 : T ; ψ ) ) ( 25 ) The third one , which is optional , is to fool the discriminator in teacher forcing mode : L ( T ) β ( θ|ψ ) = − ∑N n=1 log ( 1− f ( y′ ( n ) 1 : T , β ′ ( n ) 1 : T ; ψ ) ) ( 26 ) This approach makes the distribution p ( yt|ŷ1 : t−1 , x1 : L ; θ ) estimated in free running mode similar to the corresponding distribution p ( yt|y1 : t−1 , x1 : L ; θ ) estimated in teacher forcing mode . In addition , it regularizes some hidden layers , encouraging them to behave as if in teacher forcing mode . The disadvantage is that it requires designing and training the discriminator .

Summary: This paper proposes an alternative mechanism of training the attention values of a sequence to sequence learning model as applied to tasks like speech synthesis and translation.  During training they compute two forms of attention: (1) the standard soft-attention from a decoder fed with teacher forced output, and (2) the inference-time attention from a decoder fed with predicted outputs.  Their training objective consists of two terms: The first is the token-wise cross entropy loss but by conditioning on the predicted output  but with teacher-forced attention.  The second is a KL distance between the above two types of attention distributions.   Experiments with mechanical  turks indicate that their attention forcing mechanism is strongly preferred over the existing teacher forced output and attention model.  On translation their method provides little or no improvement.