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 paper considers a shortcoming of sequence to sequence models trained using maximum likelihood estimation. In particular, a model trained in this way can be biased in the sense that training sequences typically have different sets of prefixes compared to test sequences. As a result, at the prediction time the model does not generalize well and for a given input sequence the decoder constructs a label sequence which reflects the training label sequences rather than the actual target label.