Title: OPTIMAL COMPLETION DISTILLATION FOR SEQUENCE LEARNING

Abstract: We present Optimal Completion Distillation (OCD), a training procedure for optimizing sequence to sequence models based on edit distance. OCD is efficient, has no hyper-parameters of its own, and does not require pretraining or joint optimization with conditional log-likelihood. Given a partial sequence generated by the model, we first identify the set of optimal suffixes that minimize the total edit distance, using an efficient dynamic programming algorithm. Then, for each position of the generated sequence, we define a target distribution that puts an equal probability on the first token of each optimal suffix. OCD achieves the state-of-theart performance on end-to-end speech recognition, on both Wall Street Journal and Librispeech datasets, achieving 9.3% and 4.5% word error rates, respectively.

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.

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 attentionbased 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 } N i=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:

EQUATION

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,

O MLE (θ) = E (x,y * )∼p D 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 O MLE (θ). 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,

EQUATION

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,

EQUATION

where p D 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 ŷt by conditioning on its own previous outputs, i.e. ŷ<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 ŷ<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].

EQUATION

Then, the optimal extension for a prefix y <t can be defined as tokens that attain the maximal Q-values, i.e. argmax a Q * ( 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.

Given Q-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,

EQUATION

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,

EQUATION

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

Target sequence y * a s _ h e _ t a l k s _ h i s _ w i f e Generated sequence y a s _ e e _ t a l k s _ w h o s e _ w i f e Optimal extensions for a s _ h e _ t a l k s _ h i i s _ _ w i f e edit distance (OCD h h i w targets) _

Table 1 : A sample sequence y * from the Wall Street Journal dataset, where the model's prediction y is not perfect. The optimal next characters for each prefix of y based on edit distance are shown in blue. For example, for the prefix "as_e" there are 3 optimal next characters of "e", "h", and "_". All of these 3 characters when combined with proper suffixes will result in a total edit distance of 1.

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.

3.1. OPTIMAL Q-VALUES FOR EDIT DISTANCE:

We propose a dynamic programming algorithm to calculate optimal Q-values exactly and efficiently for the reward metric of negative edit distance, i. When this assumption does not hold, e.g. genetic applications, OCD is less efficient than MLE. However, in practice, the wall clock time is dominated by the forward and backward passes of a neural networks, and the OCD cost is often negligible. We discuss the efficiency of OCD further in Appendix A.

Recall the Levenshtein algorithm (Levenshtein, 1966) for calculating the minimum number of edits (insertion, deletion and substitution) required to convert sequences y and y * to each other based on,

   D edit ( y <-1 , :) = ∞ D edit (:, y * <-1 ) = ∞ D edit ( y <0 , y * <0 ) = 0 , D edit ( y <i , y * <j ) = min    D edit ( y <i-1 , y * <j ) + 1 D edit ( y <i , y * <j-1 ) + 1 D edit ( y <i-1 , y * <j-1 ) + 1[ỹ i = y * j ] .

(7) Table 2 shows an example edit distance table for sequences "Satrapy" and "Sunday". Our goal is to identify the set of all optimal suffixes y ∈ Y that result in a full sequences [ y <i , y] with a minimum edit distance v.s. y * . Lemma 1. The edit distance resulting from any potential suffix y ∈ Y is lower bounded by m i ,

EQUATION

Proof. Let's consider the path P that traces D edit ([ y <i , y], y * ) back to D edit ( y <0 , y * <0 ) connecting each cell to an adjacent parent cell, which provides the minimum value among the three options in (7). Such a path for tracing edit distance between "Satrapy" and "Sunday" is shown in Table 2 .

Table 2 : Each row corresponds to a prefix of "SATRAPY" and shows edit distances with all prefixes of "SUNDAY". We also show OCD targets (optimal extensions) for each prefix, and minimum value along each row, denoted m i (see ( 8)). We highlight the trace path for D edit ("Satrapy", "Sunday"). Suppose the path P crosses row i at a cell (i, k). Since the operations in ( 7) are non-decreasing, the edit distance along the path cannot decrease, so

Edit Distance Table OCD Targets mi S U N D A Y 0 1 2 3 4 5 6 S 0 S 1 0 1 2 3 4 5 U 0 A 2 1 1 2 3 3 4 U, N 1 T 3 2 2 2 3 4 4 U, N, D 2 R 4 3 3 3 3 4 5 U, N, D, A 3 A 5

D edit ([ y <i , y], y * ) ≥ D edit ( y <i , y * <k ) ≥ m i .

Then, consider any k such that D edit ( y <i , y * <k ) = m i . Let y * ≥k ≡ (y * k , . . . , y * |y * | ) denote a suffix of y * . We conclude that D edit ([ y <i , y * ≥k ], y * ) = m i , because on the one hand there is a particular edit path that results in m i edits, and on the other hand m i is a lower bound according to Lemma 1. Hence any such y * ≥k is an optimal suffix for y <i . Further, it is straightforward to prove by contradiction that the set of optimal suffixes is limited to suffixes y * ≥k corresponding to D edit ( y <i , y * <k ) = m i . Since the set of optimal completions for y <i is limited to y * ≥k , the only extensions that can lead to maximum reward are the starting token of such suffixes (y * k ). Since D edit ( y <i , y * <k ) = m i as well, we can identify the optimal extensions by calculating the edit distances between all prefixes of y and all prefixes of y * which can be efficiently calculated by dynamic programming in O(| y|.|y * |). For a prefix y <i after we calculate the minimum edit distance m i among all prefixes of y * , we set the Q * ( y <i , y * k ) = -m i for all k where y * <k has edit distance equal to m i . We set the Q * for any other token to -m i -1. We provide the details of our modified Levenshtein algorithm to efficiently compute the Q * ( y <i , a) for all i and a in Appendix A.

4. RELATED WORK:

Our work builds upon Learning to Search (Daumé III and Marcu, 2005) and Imitation Learning techniques (Ross et al., 2011; Ross and Bagnell, 2014; Sun et al., 2018) , where a student policy is optimized to imitate an expert teacher. DAgger (Ross et al., 2011) in particular is closely related, where a dataset of trajectories from an expert teacher is aggregated with samples from past student models, and a policy is optimized to mimic a given expert policy π * at various states. Similarly, OCD aims to mimic an optimal policy π * at all prefixes, but in OCD, the behavior policy is directly obtained from an online student. Further, the oracle policy is not provided during training, and we obtain the optimal policy by finding optimal Q-values. AggreVaTeD (Sun et al., 2017) assumes access to an unbiased estimate of Q-values and relies on variance reduction techniques and conjugate gradients to address a policy optimization problem. OCD calculates exact Q-values and uses regular SGD for optimization. Importantly, our roll-in prefixes are drawn only from the student model, and we do not require mixing in ground truth (a.k.a. expert) samples. Cheng and Boots (2018) showed that mixing in ground truth samples is an essential regularizer for value aggregation convergence in imitation learning.

Our work is closely related to Policy Distillation (Rusu et al., 2016) , where a Deep Q-Network (DQN) agent (Mnih et al., 2015) that is previously optimized is used as the expert teacher. Then, action sequences are sampled from the teacher and the learned Q-value estimates are distilled (Hinton et al., 2014) into a smaller student network using a KL loss. OCD adopts a similar loss function, but rather than estimating Q-values using bootstrapping, we estimate exact Q-values using dynamic programming. Moreover, we draw samples from the student rather than the teacher.

Similar to OCD, the learning to search (L2S) techniques such as LOLS (Chang et al., 2015) and Goodman et al. (2016) also attempt to estimate the Q-values for each state-action pair. Such techniques examine multiple roll-outs of a generated prefix and aggregate the return values. SeaRNN (Leblond et al., 2018) approximates the cost-to-go for each token by computing the task loss for as many roll-outs as the vocabulary size at each time step with a per step complexity of O(V T ). It is often difficult to scale approaches based on multiple roll-outs to real world datasets, where either the sequences are long or the vocabulary is large. OCD exploits the special structure in edit distance and find exact Q-values efficiently in O(V + T ) per step. Unlike L2S and SeaRNN, which require ground truth prefixes to stabilize training, we solely train on model samples.

Approaches based on Reinforcement Learning (RL) have also been applied to sequence prediction problems, including REINFORCE (Ranzato et al., 2016) , Actor-Critic (Bahdanau et al., 2017) and Self-critical Sequence Training (Rennie et al., 2017) . These methods sample sequences from the model's distribution and backpropagate a sequence-level task objective (e.g. edit distance). Beam Search Optimization (Wiseman and Rush, 2016) and Edit-based Minimum Bayes Risk (EMBR) (Prabhavalkar et al., 2018) is similar, but the sampling procedure is replaced with beam search. These training methods suffer from high variances and credit assignment problems. By contrast, OCD takes advantage of the decomposition of the sequence-level objective into token level optimal completion targets. This reduces the variance of the gradient and stabilizes the model. Crucially, unlike most RL-based approaches, we neither need MLE pretraining or joint optimization with log-likelihood. Bahdanau et al. (2016b) also noticed some of the nice structure of edit distance, but they optimize the model by regressing its outputs to edit distance values leading to suboptimal performance. Rather, we first construct the optimal policy and then use knowledge distillation for training. Independently, Karita et al. (2018) also decomposed edit distance into the contribution of individual tokens and used this decomposition within the EMBR framework. That said, Karita et al. (2018) do not theoretically justify this particular choice of decomposition and report high variance in their gradient estimates.

Reward Augmented Maximum Likelihood (RAML) (Norouzi et al., 2016) and its variants (Ma et al., 2017; Elbayad et al., 2018; Wang et al., 2018) are also similiar to RL-based approaches. Instead of sampling from the model's distribution, RAML samples sequences from the true exponentiated reward distribution. However, sampling from the true distribution is often difficult and intractable. RAML suffers from the same problems as RL-based methods in credit assignment. SPG (Ding and Soricut, 2017) changes the policy gradient formulation to sample from a reward shaped model distribution. Therefore, its samples are closer than RAML to the model's samples. In order to facilitate sampling from their proposed distribution SPG provides a heuristic to decompose ROUGE score. Although SPG has a lower variance due to their biased samples, it suffers from the same problems as RAML and RL-based methods in credit assignment.

Generally, OCD excels at training from scratch, which makes it an ideal substitution for MLE. Hence, OCD is orthogonal to methods which require MLE pretraining or joint optimization.

5. EXPERIMENTS:

We conduct our experiments on speech recogntion on the Wall Street Journal (WSJ) (Paul and Baker, 1992) and Librispeech (Panayotov et al., 2015) benchmarks. We only compare end-to-end speech recognition approaches that do not incorporate language model rescoring. On both WSJ and Librispeech, our proposed OCD (Optimal Completion Distillation) algorithm significantly outperforms our own strong baselines including MLE (Maximum Likelihood Estimation with label smoothing) and SS (scheduled sampling with a well-tuned schedule). Moreover, OCD significantly outperforms all prior work, achieving a new state-of-the-art on two competitive benchmarks.

5.1. WALL STREET JOURNAL:

The WSJ dataset is readings of three separate years of the Wall Street Journal. We use the standard configuration of si284 for training, dev93 for validation and report both test Character Error Rate (CER) and Word Error Rate (WER) on eval92. We tokenize the dataset to English characters and punctuation. Our model is an attention-based seq2seq network with a deep convolutional frontend as used in Zhang et al. (2017) . During inference, we use beam search with a beam size of 16 for all of our models. We describe the architecture and hyperparameter details in Appendix C. We first analyze some key characteristics of the OCD model separately, and then compare our results with other baselines and state-of-the-art methods.

Training prefixes and generalization. We emphasize that during training, the generated prefixes sampled from the model do not match the ground truth sequence, even at the end of training. We (Graves and Jaitly; 2014) 8.4 27.3 Gram-CTC (Liu et al.; 2017) -16.7 seq2seq (Bahdanau et al.; 2016a) 6.4 18.6 seq2seq + TLE (Bahdanau et al.; 2016b) 5.9 18.0 seq2seq + LS (Chorowski and Jaitly; 2017) -10.6 seq2seq + CNN (Zhang et al.; 2017) -10.5 seq2seq + LSD (Chan et al.; 2017) -9.6 seq2seq + CTC (Kim et al.; 2017) 7.4 -seq2seq + TwinNet (Serdyuk et al.; 2018) 6.2 -seq2seq + MLE + REINFORCE (Tjandra et al.; 2018) 6 Impact of edit distance. We further investigate the role of the optimizer by experimenting with different losses. Table 3 compares the test CER and WER of the schedule sampling with a fixed probability schedule of (1.0 → 1.0) and OCD model. Both of the models are trained only on sampled trajectories. The main difference is their optimizers, where the SS(1.0 → 1.0) model is optimizing the log likelihood of ground truth (a.k.a. Hamming distance). The significant drop in CER of SS(1.0 → 1.0) emphasizes the necessity of pretraining or joint training with MLE for models such as SS. OCD is trained from random initialization and does not require MLE pretraining, nor does it require joint optimization with MLE. We also emphasize that unlike SS, we do not need to tune an exploration schedule, OCD prefixes are simply always sampled from the model from the start of training. We note that even fine tuning a pre-trained SS model which achieves 3.6% CER with 100% sampling increases the CER to 3.8%. This emphasizes the importance of making the loss a function of the model input prefixes, as opposed to the ground truth prefixes. Appendix D covers another aspect of optimizing Edit distance rather than Hamming distance.

Target distribution. Another baseline which is closer to MLE framework is selecting only one correct target. Table 3 compares OCD with several Optimal Completion Target (OCT) models. In OCT, we optimize the log-likelihood of one target, which at each step we pick dynamically based on the minimum edit distance completion similar to OCD. We experiment with several different strategies when there is more than one character that can lead to minimum CER. In the OCT (Shortest), we select the token that would minimize the CER and the final length of the sequence. In the OCT (Same #Words), we select the token that in addition to minimum CER, would lead to the closest number of words to the target sequence. We show that OCD achieves significantly better CER and WER over the other optimization strategies compared in Table 3 . This highlights the importance of optimizing for the entire set of optimal completion targets, as opposed to a single target.

State-of-the-art. Our model trained with OCD optimizes for CER; we achieve 3.1% CER and 9.3% WER, substantially outperforming our baseline by 14% relatively on CER and 12% relatively on WER. In terms of CER, our work substantially outperforms prior work as compared in Table 4 , with the closest being Tjandra et al. (2018) trained with policy gradients on CER. In terms of WER, our work is also outperforming Chan et al. (2017) , which uses subword units while our model emits characters.

5.2. LIBRISPEECH:

For the Librispeech dataset, we train on the full training set (960h audio data) and validate our results on the dev-other set. We report the results both on the "clean" and "other" test set. We use Byte Pair Encoding (BPE) (Sennrich et al., 2016) for the output token segmentation. BPE token set is an open Wav2letter (Collobert et al., 2016) 6.9 7.2 --Gated ConvNet (Liptchinsky et al., 2017) -6.7 -20.8 Cold Fusion (Sriram et al., 2018) 3.9 7.5 9.3 17.0 Invariant Representation Learning (Liang et al., 2018) 3.3 -11.0 -Pretraining+seq2seq+CTC (Zeyer et al., 2018) -4.9 -15.4

Our Implementation seq2seq + MLE 2.9 5.7 8.4 15.4 seq2seq + OCD 1.7 4.5 6.4 13.3

vocabulary set since it includes the characters as well as common words and n-grams. We use 10k BPE tokens and report both CER and WER as the evaluation metric. We describe the architecture and hyperparameter details in Appendix C.

Fig. 3 shows the validation and training WER curves for MLE and OCD. OCD starts outperforming MLE on training decodings after training for 13 epochs and on validation decodings after 9 epochs. Our MLE baseline achieves 5.7% WER, while OCD achieves 4.5% WER on test-clean (21% improvement) and improves the state-of-the-art results over Zeyer et al. (2018) . test-other is the more challenging test split ranked by the WER of a model trained on WSJ (Panayotov et al., 2015) mainly because readers accents deviate more from US-English accents. On test-other our MLE baseline achieves 15.4%, while our OCD model achieves 13.3% WER, outperforming the 15.4% WER of Zeyer et al. (2018) . Table 5 compares our results with other recent works and the MLE baseline on Librispeech.

6. CONCLUSION:

This paper presents Optimal Completion Distillation (OCD), a training procedure for optimizing autoregressive sequence models base on edit distance. OCD is applicable to on-policy or off-policy trajectories, and in this paper, we demonstrate its effectiveness on samples drawn from the model in an online fashion. Given any prefix, OCD creates an optimal extension policy by computing the exact optimal Q-values via dynamic programming. The optimal extension policy is distilled by minimizing a KL divergence between the optimal policy and the model. OCD does not require MLE initialization or joint optimization with conditional log-likelihood. OCD achieves 3.1% CER and 9.3% WER on the competitive WSJ speech recognition task, and 4.5% WER on Librispeech without any language model. OCD outperforms all published work on end-to-end speech recognition, including our own well-tuned MLE and scheduled sampling baselines without introducing new hyper-parameters.

Procedure 1 EditDistanceQ op returns Q-values of the tokens at each time step based on the minimum edit distance between a reference sequence r and a hypothesis sequence h of length t.

1: for j in (0..t) do 2:

d j ← j + 1 3: for i in (1..t) do 4: minDist ← i 5: subCost ← i -1 6: insCost ← i + 1 7:

for j in (0..t -1) do 8: 

if h i-1 =

Q i,r1 ← 1 20: for j in (1..t) do 21: if d j = minDist then 22: Q i,rj+1 ← 1 23:

for all tokens k do 24: Although the loss calculation has the same complexity as MLE, online sampling from the model to generate the input of next RNN cell (as in OCD and SS) is generally slower than reading the ground truth (as in MLE). Therefore, overall a naive implementation of OCD is ≤ 20% slower than our baseline MLE in terms of number of step time. However, since OCD is stand alone and can be trained off-policy, we can also train on stale samples and untie the input generation worker from the training workers. In this case it is as fast as the MLE baseline.

Q i,k ← Q i,k -1-minDist return Q APPENDIX A OCD ALGORITHM

Run through. As an example of how this algorithm works, consider the sequence "SUNDAY" as reference and "SATURDAY" as hypothesis. Table A .1 first shows how to extract optimal targets and their respective Q * -values from the table of edit distances between all prefixes of reference and all prefixes of hypothesis. At each row highlighted cells indicate the prefixes which has minimum edit distance in the row. The next character at these indices are the Optimal targets for that row. At each step the Q * -value for the optimal targets is negative of the minimum edit distance and for the non-optimal characters it is one smaller.

Table A.1 also illustrates how appending the optimal completions for the prefix "SA" of the hypothesis can lead to the minimum total edit distance. Concatenating with both reference suffixes, "UNDAY" and "NDAY" will result in an edit distance of 1. Therefore, predicting "U" or "N" at step 2 can lead to the maximum attainable reward of (-1).

Table A .1: Top: Each row corresponds to a prefix of "SATURDAY" and shows edit distances with all prefixes of "SUNDAY", along with the optimal targets and their Q * -value at that step. The highlighted cells indicate cells with minimum edit distance at each row. Bottom: An example of appending suffixes of "SUNDAY" with minimum edit distance to the prefix "SA".

Edit Distance:

OCD Targets Q-values S U N D A Y 0 1 2 3 4 5 6 S 0 S 1 0 1 2 3 4 5 U 0 A 2 1 1 2 3 3 4 U, N -1 T 3 2 2 2 3 4 4 U, N, D -2 U 4 3 2 3 3 4 5 N -2 R 5 4 3 3 4 4 5 N, D -3 D 6 5 4 4 3 4 5 A -3 A 7 6 5 5 4 3 4 Y -3 Y 8 7 6 6 5 4 3 
 -3 S U N D A Y 0 1 2 3 4 5 6 S 1 0 1 2 3 4 5 A 2 1 1 2 3 3 4 U 1 N 1 D 1 A 1 Y 1 S U N D A Y 0 1 2 3 4 5 6 S 1 0 1 2 3 4 5 A 2 1 1 2 3 3 4 N 1 D 1 A 1 Y 1

APPENDIX B EXPOSURE BIAS:

A key limitation of teacher forcing for sequence learning stems from the discrepancy between the training and test objectives. One trains the model using conditional log-likelihood O CLL , but evaluates the quality of the model using empirical reward O ER .

Unlike teacher forcing and Scheduled Sampling (SS), policy gradient approaches (e.g. Wu et al. (2016) pre-train the model using teacher forcing. By contrast, the OCD method proposed in this paper defines an optimal completion policy π * t for any off-policy prefix by incorporating the ground truth information. Then, OCD optimizes a token level log-loss and alleviates the credit assignment problem. Finally, training is much more stable, and we do not require initialization nor joint optimization with MLE.

There is an intuitive notion of exposure bias Ranzato et al. ( 2016) discussed in the literature as a limitation of teacher forcing. We formalize this notion as follows. One can think of the optimization of the log loss (3) in an autoregressive models as a classification problem, where the input to the classifier is a tuple (s, y * <t ) and the correct output is y * i , where y * <t ≡ (y * 1 , . . . , y * t-1 ). Then the training dataset comprises different examples and different prefixes of the ground truth sequence. The key challenge is that once the model is trained, one should not expect the model to generalize to a new prefix y <t that does not come from the training distribution of P (y * <t ). This problem can become severe as y <t becomes more dissimilar to correct prefixes. During inference, when one conducts beam search with a large beam size then one is more likely to discover wrong generalization of p θ (ŷ t |ŷ <t , x), because the sequence is optimized globally. A natural strategy to remedy this issue followed by 3 LSTM layers with 256 cell size. We also apply batch-normalization between each layer in the encoder. The attention-based decoder is a 1-layer LSTM with 256 cell size with content-based attention. We use Xavier initializer (Glorot and Bengio, 2010) and train our models for 300 epochs of batch size 8 with 8 async workers. We separately tune the learning rate for our baseline and OCD model, 0.0007 for OCD vs 0.001 for baseline. We apply a single 0.01 drop of learning rate when validation CER plateaus, the same as for our baseline. Both happen around 225 epoch. We implemented our experiments2 in TensorFlow (Abadi et al., 2016) .

Librispeech. Since the dataset is larger than WSJ, we use a larger batch size of 16, smaller learning rate of 0.0005 for baseline and 0.0003 for OCD. Models are trained for 70 epochs. We remove the convolutional LSTM layers of the encoder, increase the number of LSTM layers in the encoder to 6, and increase the LSTM cell size to 384. All other configs are the same as the WSJ setup. 

:

Edit 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.

We are in the process of releasing the code for OCD.

Figures and Tables:

e. R(y * , y) = -D edit (y * , y). Given two sequences y * and y, we compute the Q-values for every prefix y <t and any extending token a ∈ V with an asymptotic complexity of O(|y * |.| y| + |V|.| y|). Assuming that |y * | ≈ | y| ≤ |V|, our algorithm does not increase the time complexity over MLE, since computing the cross-entropy losses in MLE also requires a complexity of O(|y * |.|V|).

Figure 1: Fraction of OCD training prefix tokens on WSJ which does not match ground truth.

Figure 2: WSJ validation Character Error Rate (CER) per training CER for MLE and OCD.

Figure 3: Librispeech training and validation WER per training epoch for OCD and MLE.

The total time complexity for calculating the sequence loss using OCD is O(T 2 + |V |T ) where V is the vocabulary size and T is the sequence length. MLE loss has a time complexity of O(|V |T ) for calculating the softmax loss at each step. Therefore, assuming that O(T ) ≤ O(|V |) OCD does not change the time complexity compared to the baseline seq2seq+MLE. The memory cost of the OCD algorithm is O(T + |V |T ) = O(|V |T ), O(T ) for the dynamic programming in line 4 -line 13 of Proc. 1 and O(|V |T ) for storing the stepwise Q values. MLE also stores the one-hot encoding of targets at each step with a cost of O(|V |T ). Therefore, the memory complexity does not change compared to the MLE baseline either.

Figure D.3 plots the edit distance on training data of OCD and MLE for fixed hamming distances during training.The plot shows that for a fixed Hamming distance (which is the metric that MLE correlates with more), OCD achieves a lower edit distance compared to MLE. This gives evidence that OCD is indeed optimizing for edit distance as intended.

Figure D.3: WSJ training Character Error Rate (CER) of MLE and OCD over Character Accuracy at different checkpoints during training.

WSJ Character Error Rate (CER) and Word Error Rate (WER) of different baselines. Schedule Sampling optimizes for Hamming distance and mixes samples from the model and ground truth with a probability schedule (start-of-training → end-of-training). OCD always samples from the model and optimizes for all characters which minimize Edit distance. Optimal Completion Target optimizes for one character which minimizes edit distance and another criteria (shortest or same #words).

Table content: <table><tr><td>Training Strategy</td><td colspan="2">CER WER</td></tr><tr><td>Schedule Sampling (1.0 → 1.0)</td><td>12.1</td><td>35.6</td></tr><tr><td>Schedule Sampling (0.0 → 1.0)</td><td>3.8</td><td>11.7</td></tr><tr><td>Schedule Sampling (0.0 → 0.55)</td><td>3.6</td><td>10.2</td></tr><tr><td>Optimal Completion Target (Shortest)</td><td>3.8</td><td>12.7</td></tr><tr><td>Optimal Completion Target (Same #Words)</td><td>3.3</td><td>10.2</td></tr><tr><td>Optimal Completion Distillation</td><td>3.1</td><td>9.3</td></tr></table>

Character Error Rate (CER) and Word Error Rate (WER) results on the end-to-end speech recognition WSJ task. We report results of our Optimal Completion Distillation (OCD) model, and well-tuned implementations of maximum likelihood estimation (MLE) and Scheduled Sampling (SS).

Table content: <table><tr><td>Model</td></tr></table>

define OCD prefix mismatch as the fraction of OCD training tokens that do not match corresponding ground truth training tokens at each position. Assuming that the generated prefix sequence is perfectly matched with the ground truth sequence, then the OCD targets would simply be the following tokens of the ground truth sequence. Hence, OCD becomes equivalent to MLE. Figure1shows that OCD prefixes mismatch is more than 25% for the most of the training. This suggests that OCD and MLE are training on very different input prefix trajectories. Further, Figure2depicts validation CER as a function of training CER for different model checkpoints during training, where we use beam search on both training and validation sets to obtain CER values. Even at the same training CER, we observe better validation error for OCD, which suggests that OCD improves generalization of MLE, possibly because OCD alleviates the mismatch between training and inference.

Table content: <table><tr><td/><td>.1</td><td>-</td></tr><tr><td>Our Implementation</td><td/><td/></tr><tr><td>seq2seq + MLE</td><td>3.6</td><td>10.6</td></tr><tr><td>seq2seq + SS</td><td>3.6</td><td>10.2</td></tr><tr><td>seq2seq + OCD</td><td>3.1</td><td>9.3</td></tr></table>

Character Error Rate (CER) and Word Error Rate (WER) on LibriSpeech test sets.

Table content: <table><tr><td>Model</td></tr></table>

r j then

Table content: <table><tr><td>9:</td><td>repCost ← 0</td></tr><tr><td>10:</td><td>else</td></tr><tr><td>11:</td><td>repCost ← 1</td></tr><tr><td>12:</td><td>cheapest ← min(subCost + repCost, d j + 1, insCost)</td></tr><tr><td>13:</td><td>subCost ← d j</td></tr><tr><td>14:</td><td>insCost ← cheapest + 1</td></tr><tr><td>15:</td><td>d j ← cheapest</td></tr><tr><td>16:</td><td>if j &lt; minDist then</td></tr><tr><td>17:</td><td>minDist ← d j</td></tr><tr><td>18:</td><td>if minDist = i then</td></tr><tr><td>19:</td><td/></tr></table>

Ranzato et al. (2016); Bahdanau et al. (2017)) and OCD aim to optimize the empirical reward objective (4) on the training set. We illustrate four different training strategies of MLE, SS, Policy Gradient and OCD in Figure B.1. The drawback of policy gradient techniques is twofold: 1) they cannot easily incorporate ground truth sequence information except through the reward function, and 2) they have difficulty reducing the variance of the gradients to perform proper credit assignment. Accordingly, most policy gradient approaches Ranzato et al. (2016); Bahdanau et al. (2017);

Table content: <table/>

Peer review (iclr2019):
 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. This is a very nice paper; it is generally well written, it gets excellent results, and it contains a comprehensive set of experiments and analysis. It is also quite exciting to see an approach to alleviating exposure bias that does not require pretraining with MLE. Accordingly, my suggestions mainly relate to the presentation and related work: - It seems a bit strange to argue that the proposed approach doesn*t increase the time complexity over MLE. While technically true (as the authors note) if the vocabulary is bigger than the sequence lengths, the difference in (on policy) training time will presumably be felt when dealing with very long sequences, or with cases where the number of labels per time-step is small, like in character-level generation or in seq2seq style sequence labeling. - I think it*s difficult to argue that the proposed approach isn*t essentially a modification of imitation learning/learning-to-search algorithms like, say, AggreVaTe or LOLS (Chang et al., ICML 2015). As far as I can tell, the only differences are that cross entropy is used rather than a cost-sensitive classifier, and, perhaps, that the training is done in minibatches (with no aggregation). - Relatedly, while it is interesting that the loss uses all the optimal completion tokens, it should be noted that there is much work in transition-based parsing that adopts a learning-to-search approach and uses losses that incorporate multiple optimal next-predictions as given by a "dynamic oracle"; see Goldberg and Nivre (COLING, 2012) and others. - I think it*s also worth noting that training approaches like MIXER and others can target arbitrary rewards (and not just those where we can efficiently compute optimal next-steps), and so the proposed approach is a compelling competitor to MIXER-like approaches on problems such as machine translation or image captioning only to the extent that training with edit-distance is useful for such problems. Do you have a sense of whether training with edit-distance does indeed improve performance on such tasks? Pros: - well written and interesting - good experiments, results, and analysis Cons: - perhaps slightly more similar to previous work than is argued Update after author response: thanks for your response; I think the revised paper largely addresses my comments and those of the other reviewers, and I continue to hope it is accepted. Here are two small notes on the related work section of the revised paper: - In distinguishing OCD from DAgger, you note that the optimal policy is computed rather than provided at training time. In fact, structured prediction applications of SEARN (Daume III et al., 2009, which should also be cited) and DAgger often have this flavor too, such as when using them for sequence labeling (where optimal continuations are calculated based on Hamming distance). - Please include a reference to Goldberg and Nivre*s (2012) dynamic oracle work.