Optimal Completion Distillation for Sequence Learning

1 INTRODUCTION . Recent advances in natural language processing and speech recognition hinge on the development of expressive neural network architectures for sequence to sequence ( seq2seq ) learning ( Sutskever et al. , 2014 ; Bahdanau et al. , 2015 ) . Such encoder-decoder architectures are adopted in both machine translation ( Bahdanau et al. , 2015 ; Wu et al. , 2016 ; Hassan et al. , 2018 ) and speech recognition systems ( Chan et al. , 2016 ; Bahdanau et al. , 2016a ; Chiu et al. , 2017 ) achieving impressive performance above traditional multi-stage pipelines ( Koehn et al. , 2007 ; Povey et al. , 2011 ) . Improving the building blocks of seq2seq models can fundamentally advance machine translation and speech recognition , and positively impact other domains such as image captioning ( Xu et al. , 2015 ) , parsing ( Vinyals et al. , 2015 ) , summarization ( Rush et al. , 2015 ) , and program synthesis ( Zhong et al. , 2017 ) . To improve the key components of seq2seq models , one can either design better architectures , or develop better learning algorithms . Recent architectures using convolution ( Gehring et al. , 2017 ) and self attention ( Vaswani et al. , 2017 ) have proved to be useful , especially to facilitate efficient training . On the other hand , despite many attempts to mitigate the limitations of Maximum Likelihood Estimation ( MLE ) ( Ranzato et al. , 2016 ; Wiseman and Rush , 2016 ; Norouzi et al. , 2016 ; Bahdanau et al. , 2017 ; Leblond et al. , 2018 ) , MLE is still considered the dominant approach for training seq2seq models . Current alternative approaches require pre-training or joint optimization with conditional log-likelihood . They are difficult to implement and require careful tuning of new hyper-parameters ( e.g . mixing ratios ) . In addition , alternative approaches typically do not offer a substantial performance improvement over a well tuned MLE baseline , especially when label smoothing ( Pereyra et al. , 2017 ; Edunov et al. , 2018 ) and scheduled sampling ( Bengio et al. , 2015 ) are used . In this paper , we borrow ideas from search-based structured prediction ( Daumé et al. , 2009 ; Ross et al. , 2011 ) and policy distillation ( Rusu et al. , 2016 ) and develop an efficient algorithm for optimizing seq2seq models based on edit distance1 . Our key observation is that given an arbitrary prefix ( e.g . a partial sequence generated by sampling from the model ) , we can exactly and efficiently identify all of the suffixes that result in a minimum total edit distance ( v.s . the ground truth target ) . Our training procedure , called Optimal Completion Distillation ( OCD ) , is summarized as follows : 1 . We always train on prefixes generated by sampling from the model that is being optimized . 2 . For each generated prefix , we identify all of the optimal suffixes that result in a minimum total edit distance v.s . the ground truth target using an efficient dynamic programming algorithm . 3 . We teach the model to optimally extend each generated prefix by maximizing the average log probability of the first token of each optimal suffix identified in step 2 . 1Edit distance between two sequences u and v is the minimum number of insertion , deletion , and substitution edits required to convert u to v and vice versa . The proposed OCD algorithm is efficient , straightforward to implement , and has no tunable hyperparameters of its own . Our key contributions include : • We propose OCD , a stand-alone algorithm for optimizing seq2seq models based on edit distance . OCD is scalable to real-world datasets with long sequences and large vocabularies , and consistently outperforms Maximum Likelihood Estimation ( MLE ) by a large margin . • Given a target sequence of length m and a generated sequence of length n , we present an O ( nm ) algorithm that identifies all of the optimal extensions for each prefix of the generated sequence . • We demonstrate the effectiveness of OCD on end-to-end speech recognition using attention- based seq2seq models . On the Wall Street Journal dataset , OCD achieves a Character Error Rate ( CER ) of 3.1 % and a Word Error Rate ( WER ) of 9.3 % without language model rescoring , outperforming all prior work ( Table 4 ) . On Librispeech , OCD achieves state-of-the-art WER of 4.5 % on “ test-clean ” and 13.3 % on “ test-other ” sets ( Table 5 ) . 2 BACKGROUND : SEQUENCE LEARNING WITH MLE . Given a dataset of input output pairs D ≡ { ( x , y∗ ) i } Ni=1 , we are interested in learning a mapping x→ y from an input x to a target output sequence y∗ ∈ Y . Let Y denote the set of all sequences of tokens from a finite vocabulary V with variable but finite lengths . Often learning a mapping x→ y is formulated as optimizing the parameters of a conditional distribution pθ ( y | x ) . Then , the final sequence prediction under the probabilistic model pθ is performed by exact or approximate inference ( e.g . via beam search ) as : ŷ ≈ argmax y∈Y pθ ( y | x ) . ( 1 ) Similar to the use of log loss for supervised classification , the standard approach to optimize the parameters θ of the conditional probabilistic model entails maximizing a conditional log-likelihood objective , OMLE ( θ ) = E ( x , y∗ ) ∼pD log pθ ( y∗ | x ) . This approach to learning the parameters is called Maximum Likelihood Estimation ( MLE ) and is commonly used in sequence to sequence learning . Sutskever et al . ( 2014 ) propose the use of recurrent neural networks ( RNNs ) for autoregressive seq2seq modeling to tractably optimize OMLE ( θ ) . An autoregressive model estimates the conditional probability of the target sequence given the source one token at a time , often from left-to-right . A special end-of-sequence token is appended at the end of all of target sequences to handle variable length . The conditional probability of y∗ given x is decomposed via the chain rule as , pθ ( y ∗ | x ) ≡ ∏|y∗| t=1 pθ , t ( y ∗ t | y∗ < t , x ) , ( 2 ) where y∗ < t ≡ ( y∗1 , . . . , y∗t−1 ) denotes a prefix of the sequence y∗ . To estimate the probability of a token a given a prefix y∗ < t and an input x , denoted pθ , t ( a | y∗ < t , x ) , different architectures have been proposed . Some papers ( e.g . Britz et al . ( 2017 ) ) have investigated the use of LSTM ( Hochreiter and Schmidhuber , 1997 ) and GRU ( Cho et al. , 2014 ) cells , while others proposed new architecturs based on soft attention ( Bahdanau et al. , 2015 ) , convolution ( Gehring et al. , 2017 ) , and self-attention ( Vaswani et al. , 2017 ) . Nonetheless , all of these techniques rely on MLE for learning , OMLE ( θ ) = E ( x , y∗ ) ∼pD ∑|y∗| t=1 log pθ , t ( y ∗ t | y∗ < t , x ) , ( 3 ) where pD denotes the empirical data distribution , uniform across the dataset D. We present a new objective function for optimizing autoregressive seq2seq models applicable to any neural architecture . 2.1 LIMITATIONS OF MLE FOR AUTOREGRESSIVE MODELS . In order to maximize the conditional log-likelihood ( 3 ) of an autoregressive seq2seq model ( 2 ) , one provides the model with a prefix of t− 1 tokens from the ground truth target sequence , denoted y∗ < t , and maximizes the log-probability of y∗t as the next token . This resembles a teacher walking a student through a sequence of perfect decisions , where the student learns as a passive observer . However , during inference one uses beam search ( 1 ) , wherein the student needs to generate each token ŷt by conditioning on its own previous outputs , i.e . ŷ < t instead of y ∗ < t . This creates a discrepancy between training and test known as exposure bias ( Ranzato et al. , 2016 ) . Appendix B expands this further . Concretely , we highlight two limitations with the use of MLE for autoregressive seq2seq modeling : 1 . There is a mismatch between the prefixes seen by the model during training and inference . When the distribution of ŷ < t is different from the distribution of y ∗ < t , then the student will find themselves in a novel situation that they have not been trained for . This can result in poor generalization , especially when the training set is small or the model size is large . 2 . There is a mismatch between the training loss and the task evaluation metric . During training , one optimizes the log-probability of the ground truth output sequence , which is often different from the task evaluation metric ( e.g . edit distance for speech recognition ) . There has been a recent surge of interest in understanding and mitigating the limitations of MLE for autoregressive seq2seq modeling . In Section 4 we discuss prior work in detail after presenting our approach below . 3 OPTIMAL COMPLETION DISTILLATION . To alleviate the mismatch between inference and training , we never train on ground truth target sequences . Instead , we always train on sequences generated by sampling from the current model that is being optimized . Let ỹ denote a sequence generated by sampling from the current model , and y∗ denote the ground truth target . Applying MLE to autoregressive models casts the problem of sequence learning as optimizing a mapping ( x , y∗ < t ) → y∗t from ground truth prefixes to correct next tokens . By contrast , the key question that arises when training on model samples is the choice of targets for learning a similar mapping ( x , ỹ < t ) → ? ? from generated prefixes to next tokens . Instead of using a set of pre-specified targets , OCD solves a prefix-specific problem to find optimal extensions that lead to the best completions according to the task evaluation metric . Then , OCD encourages the model to extend each prefix with the set of optimal choices for the next token . Our notion of optimal completion depends on the task evaluation metric denoted R ( · , · ) , which measures the similarity between two complete sequences , e.g . the ground truth target v.s . a generated sequence . Edit distance is a common task metric . Our goal in sequence learning is to train a model , which achieves high scores of R ( y∗ , ỹ ) . Drawing connection with the goal of reinforcement learning ( Sutton and Barto , 1998 ) , let us recall the notion of optimal Q-values . Optimal Q-values for a state-action pair ( s , a ) , denoted Q∗ ( s , a ) , represent the maximum future reward that an agent can accumulate after taking an action a at a state s by following with optimal subsequent actions . Similarly , we define Q-values for a prefix ỹ < t and the extending token a , as the maximum score attainable by concatenating [ ỹ < t , a ] with an optimal suffix y to create a full sequence [ ỹ < t , a , y ] . Formally , ∀a ∈ V , Q∗ ( ỹ < t , a ) = max y∈Y R ( y∗ , [ ỹ < t , a , y ] ) . ( 4 ) Then , the optimal extension for a prefix ỹ < t can be defined as tokens that attain the maximal Q-values , i.e . argmaxaQ ∗ ( ỹ < t , a ) . This formulation allows for a prefix ỹ < t to be sampled on-policy from the model pθ , or drawn off-policy in any way . Table 1 includes an example ground truth target from the Wall Street Journal dataset and the corresponding generated sample from a model . We illustrate that for some prefixes there exist more than a single optimal extension leading to the same edit distance . GivenQ-values for our prefix-token pairs , we use an exponential transform followed by normalization to convert Q-values to a soft optimal policy over the next token extension , π∗ ( a | ỹ < t ) = exp ( Q∗ ( ỹ < t , a ) /τ ) ∑ a′ exp ( Q ∗ ( ỹ < t , a′ ) /τ ) , ( 5 ) where τ ≥ 0 is a temperature parameter . Note the similarity of τ and the label smoothing parameter helpful within MLE . In our experiments , we used the limit of τ → 0 resulting in hard targets and no hyper-parameter tuning . Given a training example ( x , y∗ ) , we first draw a full sequence ỹ ∼ pθ ( · | x ) i.i.d . from the current model , and then minimize a per-step KL divergence between the optimal policy and the model distribution over the next token extension at each time step t. The OCD objective is expressed as , OOCD ( θ ) = E ( x , y∗ ) ∼pD Eỹ∼pθ ( ·|x ) ∑|ỹ| t=1 KL ( π∗ ( · | ỹ < t ) ‖ pθ , t ( · | ỹ < t , x ) ) . ( 6 ) For every prefix ỹ < t , we compute the optimal Q-values and use ( 5 ) to construct the optimal policy distribution π∗ . Then , we distill the knowledge of the optimal policy for each prefix ỹ < t into the parametric model using a KL loss . For the important class of sequence learning problems where edit distance is the evaluation metric , we develop a dynamic programming algorithm to calculate optimal Q-values exactly and efficiently for all prefixes of a sequence ỹ , discussed below .

Summary: The authors propose an alternative approach to training seq2seq models, which addresses concerns about exposure bias and about the typical MLE objective being different from the final evaluation metric. In particular, the authors propose to use a dynamic program to compute the optimal continuations of predicted prefixes (during training) in terms of edit distance to the true output, and then use a per-token cross entropy loss, with a target distribution that is uniform over all optimal next-tokens. The authors conduct a number of experiments, and show that this scheme allows them to attain state-of-the-art performance on end-to-end speech recognition, and that they can moreover do this without needing to pretrain the model with the MLE objective.