O PTIMAL C OMPLETION D ISTILLATION FOR S EQUENCE L EARNING
Google Brain
A BSTRACT
We present Optimal Completion Distillation (OCD), a training procedure for optimizing sequence to sequence models based on edit distance.
OCD is efﬁcient, 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 ﬁrst identify the set of optimal sufﬁxes that minimize the total edit distance, using an efﬁcient dynamic programming algorithm.
Then, for each position of the generated sequence, we deﬁne a target distribution that puts an equal probability on the ﬁrst token of each optimal sufﬁx.
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.
I NTRODUCTION
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 efﬁcient 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 difﬁcult 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 efﬁcient algorithm for optimizing seq2seq models based on edit distance 1 .
Our key observation is that given an arbitrary preﬁx ( e.g. a partial sequence generated by sampling from the model), we can exactly and efﬁciently identify all of the sufﬁxes 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:
We always train on preﬁxes generated by sampling from the model that is being optimized.
For each generated preﬁx, we identify all of the optimal sufﬁxes that result in a minimum total edit distance v.s. the ground truth target using an efﬁcient dynamic programming algorithm.
We teach the model to optimally extend each generated preﬁx by maximizing the average log probability of the ﬁrst token of each optimal sufﬁx identiﬁed in step 2.
The proposed OCD algorithm is efﬁcient, 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 identiﬁes all of the optimal extensions for each preﬁx 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).
B ACKGROUND : S EQUENCE L EARNING 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 ﬁnite vocabulary V with variable but ﬁnite lengths.
Often learning a mapping x → y is formulated as optimizing the parameters of a conditional distribution p θ ( y | x ) .
Then, the ﬁnal 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 classiﬁcation, 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 preﬁx of the sequence y ∗ .
To estimate the probability of a token a given a preﬁx 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. L IMITATIONS OF MLE FOR A UTOREGRESSIVE M ODELS
In order to maximize the conditional log-likelihood (3) of an autoregressive seq2seq model (2), one provides the model with a preﬁx 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:
There is a mismatch between the preﬁxes 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 ﬁnd 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.
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.
O PTIMAL C OMPLETION D ISTILLATION
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 (cid:101)
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 preﬁxes 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 , (cid:101)
y <t )
→ ??
from generated preﬁxes to next tokens.
Instead of using a set of pre-speciﬁed targets, OCD solves a preﬁx-speciﬁc problem to ﬁnd optimal extensions that lead to the best completions according to the task evaluation metric.
Then, OCD encourages the model to extend each preﬁx 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 ∗ , (cid:101) 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 deﬁne Q-values for a preﬁx (cid:101)
y <t and the extending token a , as the maximum score attainable by concatenating [ (cid:101)
y <t , a ] with an optimal sufﬁx y to create a full sequence [ (cid:101)
y <t , a, y ] .
Formally,
[Equation]
Then, the optimal extension for a preﬁx (cid:101) y <t can be deﬁned as tokens that attain the maximal Q-values, i.e. argmax a Q ∗ ( (cid:101)
y <t , a ) .
This formulation allows for a preﬁx (cid:101)
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 preﬁxes there exist more than a single optimal extension leading to the same edit distance.
Given Q -values for our preﬁx-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 ﬁrst draw a full sequence (cid:101)
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 preﬁx (cid:101) 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 preﬁx (cid:101)
y <t into the parametric model using a KL loss.
For the important class of sequence learning problems where edit
[Table]
distance is the evaluation metric, we develop a dynamic programming algorithm to calculate optimal
Q -values exactly and efﬁciently for all preﬁxes of a sequence (cid:101) y , discussed below.
3. O PTIMAL Q- VALUES FOR E DIT D ISTANCE
We propose a dynamic programming algorithm to calculate optimal Q-values exactly and efﬁciently for the reward metric of negative edit distance, i.e. R ( y ∗ , (cid:101) y ) =
− D edit ( y ∗ , (cid:101) y ) .
Given two sequences y ∗ and (cid:101) y , we compute the Q-values for every preﬁx (cid:101)
y <t and any extending token a ∈ V with an asymptotic complexity of O ( | y ∗ | .
| (cid:101) y | + |V| .
| (cid:101) y | ) .
Assuming that | y ∗ | ≈ | (cid:101) 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| ) .
When this assumption does not hold, e.g. genetic applications,
OCD is less efﬁcient 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 efﬁciency 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 (cid:101) y and y ∗ to each other based on,  
[Equation]
Table 2 shows an example edit distance table for sequences “Satrapy” and “Sunday”.
Our goal is to identify the set of all optimal sufﬁxes y ∈ Y that result in a full sequences [ (cid:101)
y <
i , y ] with a minimum edit distance v.s. y ∗ .
Lemma 1.
The edit distance resulting from any potential sufﬁx y ∈ Y is lower bounded by m i ,
[Equation]
Proof.
Let’s consider the path P that traces D edit ([ (cid:101)
y <
i , y ] , y ∗ ) back to D edit ( (cid:101) 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]
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 D edit ([ (cid:101)
y <
i , y ] , y ∗ ) ≥ D edit ( (cid:101)
y <
i , y ∗ <k ) ≥
m i .
Then, consider any k such that D edit ( (cid:101)
y <
i , y ∗ <k )
= m i .
Let y ∗≥ k ≡ ( y ∗ k , . . .
, y ∗| y ∗ | ) denote a sufﬁx of y ∗ .
We conclude that D edit ([ (cid:101)
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 sufﬁx for (cid:101)
y <
i .
Further, it is straightforward to prove by contradiction that the set of optimal sufﬁxes is limited to sufﬁxes y ∗≥ k corresponding to D edit ( (cid:101)
y <
i , y ∗ <k )
= m i .
Since the set of optimal completions for (cid:101) y <i is limited to y ∗≥ k , the only extensions that can lead to maximum reward are the starting token of such sufﬁxes ( y ∗ k ).
Since D edit ( (cid:101) y <
i , y ∗ <k )
=
m i as well, we can identify the optimal extensions by calculating the edit distances between all preﬁxes of (cid:101)
y and all preﬁxes of y ∗ which can be efﬁciently calculated by dynamic programming in O ( | (cid:101) y
| .
| y ∗ | ) .
For a preﬁx (cid:101) y <i after we calculate the minimum edit distance m i among all preﬁxes of y ∗ , we set the Q ∗ ( (cid:101)
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 modiﬁed Levenshtein algorithm to efﬁciently compute the Q ∗ ( (cid:101)
y <
i , a ) for all i and a in Appendix A.
R ELATED W ORK
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 preﬁxes, 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 ﬁnding 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 preﬁxes 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 preﬁx 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 difﬁcult 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 ﬁnd exact Q-values efﬁciently in O ( V + T ) per step.
Unlike L2S and SeaRNN, which require ground truth preﬁxes 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 ﬁrst 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 difﬁcult 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.
E XPERIMENTS
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 signiﬁcantly outperforms our own strong baselines including MLE (Maximum Likelihood Estimation with label smoothing) and SS (scheduled sampling with a well-tuned schedule).
Moreover, OCD signiﬁcantly outperforms all prior work, achieving a new state-of-the-art on two competitive benchmarks.
5. W ALL S TREET J OURNAL
The WSJ dataset is readings of three separate years of the Wall Street Journal.
We use the standard conﬁguration of si 284 for training, dev 93 for validation and report both test Character Error Rate (CER) and Word Error Rate (WER) on eval 92 .
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.
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 ﬁrst analyze some key characteristics of the OCD model separately, and then compare our results with other baselines and state-of-the-art methods.
Training preﬁxes and generalization .
We emphasize that during training, the generated preﬁxes sampled from the model do not match the ground truth sequence, even at the end of training.
We
[Table] [Table]
[Table]
et
[Table]
+ LSD
[Table]
CTC (Kim et al.; 2017) 7.
- seq2seq + TwinNet
(Serdyuk et al.; 2018) 6.
- seq2seq + MLE + REINFORCE
(Tjandra et al.; 2018) 6.
Our Implementation seq2seq + MLE 3.
seq2seq + SS 3. seq2seq + OCD 3. deﬁne
OCD preﬁx mismatch as the fraction of OCD training tokens that do not match corresponding ground truth training tokens at each position. Assuming that the generated preﬁx 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.
Figure 1 shows that OCD preﬁxes mismatch is more than 25% for the most of the training.
This suggests that OCD and MLE are training on very different input preﬁx trajectories.
Further, Figure 2 depicts 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.
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 ﬁxed probability schedule of (1 .
→ 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 .
→ 1 . 0) model is optimizing the log likelihood of ground truth ( a.k.a.
Hamming distance).
The signiﬁcant drop in CER of SS (1 .
→ 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 preﬁxes are simply always sampled from the model from the start of training.
We note that even ﬁne tuning a pre-trained SS model which achieves 3 .
6% CER with 100% sampling increases the CER to 3 .
This emphasizes the importance of making the loss a function of the model input preﬁxes, as opposed to the ground truth preﬁxes.
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 ﬁnal 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 signiﬁcantly 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. L IBRISPEECH
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 [Table] [Table] [Table]
Cold Fusion (Sriram et al., 2018)
[Table]
Invariant
[Table]
[Table]
et al
[Table]
[Table]
Pretraining+seq2seq+CTC (
[Table]
et al., 2018)
[Table]
Our Implementationseq2seq+MLE 2.
[Table]
seq2seq + OCD
[Table]
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.
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.
Table 5 compares our results with other recent works and the MLE baseline on
Librispeech.
C ONCLUSION
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 preﬁx, 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. A CKNOWLEDGEMENTS
We thank Geoffrey Hinton for his valuable feedback and reviews.
We thank Samy Bengio, Navdeep
Jaitly and Jamie Kiros for their help in reviewing the manuscript as well.
We thank Zhifeng Chen and Yonghui Wu for their generous technical help
A PPENDIX A OCD A LGORITHM
Complexity.
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.
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.
Run through.
As an example of how this algorithm works, consider the sequence “SUNDAY” as reference and “SATURDAY” as hypothesis.
Table A.1 ﬁrst shows how to extract optimal targets and their respective Q ∗ -values from the table of edit distances between all preﬁxes of reference and all preﬁxes of hypothesis.
At each row highlighted cells indicate the preﬁxes 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 preﬁx “SA” of the hypothesis can lead to the minimum total edit distance.
Concatenating with both reference sufﬁxes, “UNDAY” and “NDAY” will result in an edit distance of
Therefore, predicting “U” or “N” at step 2 can lead to the maximum attainable reward of ( − 1) [Table] [Table] Y
A PPENDIX B E XPOSURE 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. Ranzato et al.
; Bahdanau et al.
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 difﬁculty reducing the variance of the gradients to perform proper credit assignment.
Accordingly, most policy gradient approaches Ranzato et al.
; Bahdanau et al.
; Wu et al.
pre-train the model using teacher forcing.
By contrast, the OCD method proposed in this paper deﬁnes an optimal completion policy π ∗ t for any off-policy preﬁx 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 classiﬁcation problem, where the input to the classiﬁer 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 preﬁxes 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 preﬁx
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 preﬁxes.
During inference, when one conducts beam search with a large beam size then one is more likely to discover wrong generalization of p θ (ˆ y t | ˆ y <t , x ) , because the sequence is optimized globally.
A natural strategy to remedy this issue
[Figure]
Figure B.2 illustrates how increasing the beam size for MLE and SS during inference decreases their performance on WSJ datasets to above 11% WER.
OCD suffers a degradation in the performance too
but it never gets above 10% WER.
A PPENDIX C A RCHITECTURE
WSJ.
The input audio signal is converted into 80 -dimensional ﬁlterbank features computed every 10 ms with delta and delta-delta acceleration, normalized with per-speaker mean and variance generated by Kaldi (Povey et al., 2011).
Our encoder uses 2-layers of convolutions with 3 × 3 ﬁlters, stride 2 × 2 and 32 channels, followed by a convolutional LSTM with 1D-convolution of ﬁlter width 3 , 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 .
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 experiments 2 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 .
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 conﬁgs are the same as the WSJ setup.
A PPENDIX D H AMMING DISTANCE VS E DIT D ISTANCE DURING TRAINING
Figure D.3 plots the edit distance on training data of OCD and MLE for ﬁxed hamming distances during training.
The plot shows that for a ﬁxed 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.