Parallel Scheduled Sampling

1 INTRODUCTION . Auto-regressive models are a popular choice for generating sequences of any kind including audio ( van den Oord et al. , 2016b ) , images ( van den Oord et al. , 2016a ) , and text ( Sutskever et al. , 2014 ; Cho et al. , 2014 ) . Here , the joint probability of the sequence is factorized in a pre-determined order during train and test time . For example , auto-regressive models for text generation factorize the joint probability left-to-right . The text sequence is generated by a decoder network left-to-right , one token ( word or word-piece or character ) at a time and are widely used in text generation tasks such as summarization ( Liu et al. , 2018 ) , machine translation ( Sutskever et al. , 2014 ) and dialog response generation ( Budzianowski et al. , 2018 ) in the encoder-decoder ( Cho et al. , 2014 ; Sutskever et al. , 2014 ) setting . Such models are typically trained by teacher-forcing ( Williams and Zipser , 1989 ) where ground-truth history is fed to the model as input , which at test time is replaced by the model prediction . Auto-regressive models applied in any domain suffer from this train-test time discrepancy . Scheduled Sampling ( Bengio et al. , 2015 ) aims to mitigate the discrepancy between train and test time in teacher-forcing by randomly replacing some tokens in the history with the model ’ s prediction . More concretely , at a given time step in generating the output sequence , the model is conditioned either on ground-truth or model prediction from the previous time-step with some probability . The probability of selecting model predicted token is gradually increased as training progresses . This procedure potentially allows the model to recover from its own errors , and Bengio et al . ( 2015 ) observe better empirical performance in natural language parsing , image captioning , and speech recognition compared to teacher-forced training . Scheduled Sampling has also been used to get better performance in other tasks such as video prediction ( Finn et al. , 2016 ) , knowledge base completion ( Miwa and Bansal , 2016 ) and piano music transcription ( Sigtia et al. , 2016 ) . A key bottleneck in training models with Scheduled Sampling is its inherently sequential nature . Unlike teacher-forcing , tokens must be processed one time-step at a time . The sequential procedure makes Scheduled Sampling impractical for training neural networks , particularly on problems involving long sequence generation . In this work , we describe a simple technique to parallelize Scheduled Sampling . Given an input example , we first generate an entire model prediction sequence in parallel by conditioning on ground-truth history ( equivalent to forward-pass of teacher-forcing ) . Then , we employ a ( parallel ) mixing step where we generate a new sequence whose token at every time step is either from the model prediction or the ground-truth . Finally , we perform training as in teacher-forcing by conditioning on the sequence obtained from mixing . In Section 2.3 we show that by performing multiple passes of parallel prediction and mixing , we obtain a conditioning sequence that converges to a sample decode from the model . Our key contributions are , • A novel , fully parallelizable approach to reducing train-test discrepancy of auto-regressive models without modification . We believe this to be the first approach that does not require full-sequence decoding within the training loop . We find the method produces the same empirical benefits of Scheduled Sampling while using as little as 0.3 % of the training time . • We prove equivalence of the proposed approach to Scheduled Sampling under certain choices of hyperparameters . This recovers the clear interpolation between training-time teacher-forcing and test-time decoding described in Bengio et al . ( 2015 ) . • We extensively evaluate our approach on four auto-regressive tasks in text and image domains . We find that Parallel Scheduled Sampling matches Scheduled Sampling ’ s benefits , takes massively less compute time , and significantly improves performance compared to teacher-forcing on most tasks . In dialog response generation task , Parallel Scheduled Sampling achieves 1.6 BLEU score ( 11.5 % ) improvement over teacher-forcing while in image generation it achieves 20 % and 13.8 % improvement in Frechet Inception Distance ( FID ) and Inception Score ( IS ) respectively . 2 METHOD . Our proposed technique can be applied to both conditional and unconditional auto-regressive generative models . For notational simplicity , we consider the task of conditional sequence generation . The training set is given in terms of N input-output sequences { ( xi , yi ) } ni=1 , where xi is the input and target yi is the desired output . The target yi is a variable-length sequence of Ti tokens ( or pixels ) , ( yi1 , y i 2 , . . . , y i Ti ) , whereas xi may be variable-length ( as in translation ) or fixed-length ( as in image captioning ) . The goal is to learn a model that accurately predicts yi given xi . We use y1 : t to denote the sequence of tokens ( y1 , y2 , . . . , yt ) . 2.1 TEACHER-FORCING AND DECODING . Given an input x and a target y , the log-probability of the target can be decomposed autoregressively : P ( y|x ) = T∏ t=1 P ( yt|y1 : t−1 , x ) Auto-regressive sequence generation models learn to assign high likelihood to token yt given previous target tokens y1 : t−1 = ( y1 , . . . , yt−1 ) and inputs x via a learned likelihood model pθ . Neural language models such as RNNs ( Mikolov et al. , 2010 ) and Transformer ( Vaswani et al. , 2017 ) adopt left-to-right decomposition while image generation models ( van den Oord et al. , 2016a ; Parmar et al. , 2018 ) adopt a raster-scan order . Such models are typically trained with teacher-forcing ( Williams and Zipser , 1989 ) . In teacherforcing , the log likelihood of the training set is directly maximized , θ∗ = argmax θ Ltf ( θ ) = argmax θ N∑ i=1 Ti∑ t=1 log pθ ( y i t|yi1 : t−1 , xi ) ( 1 ) Importantly , teacher-forcing conditions on gold target prefixes y1 : t−1 , enabling backpropagation through all timesteps with a single pass of inference . At inference time , beam or sample decoding is often used to generate a candidate target ŷi . In this regime , target tokens are generated one at a time while conditioning on previously-generated tokens . ŷt ∼ pθ ( ŷt|ŷ1 : t−1 , x ) or ŷt = argmax y pθ ( y|ŷ1 : t−1 , x ) A potential failure mode for teacher-forcing-trained models is in conditioning on previously unobserved target prefixes ŷ1 : t−1 . As the model has not conditioned on these prefixes at training time , it may generate bland , repetitive , or nonsensical candidate targets ( Holtzman et al. , 2019 ) . 2.2 SCHEDULED SAMPLING . Scheduled Sampling ( Bengio et al. , 2015 ) , hereafter Sequential Scheduled Sampling is a training technique designed to bridge the gap between teacher-forcing and sample decoding . In its simplest form , Sequential Scheduled Sampling generates tokens ỹ1 : t and conditions on these target prefixes during training . Sequential Scheduled Sampling uses the same objective function as teacher-forcing ( Equation 1 ) except the conditioning tokens ỹ1 : t are a random mixture of gold tokens y1 : t and sampled tokens ŷ1 : t instead of gold tokens y1 : t. See Algorithm 1 for implementation . Algorithm 1 Sequential Scheduled Sampling ( single example ) for all timesteps t = 1 , . . . , Ti do Sample ŷt ∼ pθ ( ŷt|ỹ1 : t−1 , x ) Choose next conditioning token , ỹt = { yt with probability 1− p ŷt with probability p end for return Accumulate loss ∑ t log pθ ( yt|ỹ1 : t−1 , x ) As p→ 0 , we condition on y1 : t−1 as in teacher-forcing , and as p→ 1 , we condition on ŷt : t−1 as in sample decoding . Typically a schedule will be used to gradually increase p over the course of training . As illustrated in Bengio et al . ( 2015 ) , Scheduled Sampling leads to a performance improvement in a variety of language generation tasks . In spite of its benefits , Sequential Scheduled Sampling is inherently a sequential algorithm : choosing conditioning token ỹt requires conditioning autoregressively on tokens ỹ1 : t−1 . While this is natural for sequential architectures such as RNNs and LSTMs , it is poorly suited to self-attending feed-forward models such as Transformer where inference for multiple timesteps can be carried out simultaneously . 2.3 PARALLEL SCHEDULED SAMPLING . We propose a natural extension to Sequential Scheduled Sampling called Parallel Scheduled Sampling . Whereas Sequential Scheduled Sampling selects conditioning tokens one after another , we propose generating conditioning tokens for all timesteps in parallel over the course of one or more passes . While this technique requires strictly more operations than Sequential Scheduled Sampling , it is better suited to hardware accelerators such as GPUs and TPUs ( Jouppi et al. , 2017 ) . Moreover , we find in our experiments that only a modest number of passes is necessary for improving model performance . Parallel Scheduled Sampling generates conditioning tokens for all timesteps simultaneously . The procedure consists of multiple passes , each pass consisting of parallel sampling and mixing steps ( Figure 1 ) . In the first pass , the algorithm conditions on gold tokens y1 : t , generating tokens ŷ1 : t i.i.d . according to pθ ( ŷt|y1 : t−1 , x ) . Sampling tokens in the first pass is equivalent to the forward-pass of teacher-forcing . The sampled tokens , ŷ1 : t , are mixed ( in parallel ) with gold tokens , y1 : t , to produce conditioning tokens for the next pass , ỹ1 : t. We now describe the multiple-pass procedure . Let ŷk,1 : t , and ỹk,1 : t denote sampled and mixed tokens respectively on pass k. The mixed tokens from pass k , ỹk,1 : t , are used for conditioning on pass k + 1 in place of gold tokens y1 : t. Finally , the loss is calculated as before , conditioning on the final mixture of gold and sampled tokens ỹK,1 : t. See Algorithm 2 for implementation . Algorithm 2 Parallel Scheduled Sampling ( single example ) Set ỹ0 , t = yt . for all passes k = 1 , . . . , K do for all timesteps t in parallel do Sample ŷk , t ∼ pθ ( ŷk , t|ỹk−1,1 : t−1 , x ) . if t < k then Copy ỹk , t = ỹk−1 , t else Sample ỹk , t = { yt with probability 1− p ŷk , t with probability p end if end for end for return Accumulate loss ∑ t log pθ ( yt|ỹK,1 : t−1 , x ) Finally , we prove that by running the sampling and mixing steps for multiple passes as described in Algorithm 2 , the final sample from Parallel Scheduled Sampling converges to a random sample decode from the model when p = 1 and K ≥ T . Theorem 2.1 . Consider a sequence of tokens z = ( z1 , z2 , . . . , zT ) of length T . Let p = 1 andK ≥ T be fixed . Then the likelihood of z1 : T under Parallel Scheduled Sampling ’ s proposal distribution1 over conditioning tokens on pass K , qKθ ( z1 : T ) , is identical to random sample decoding ’ s , pθ ( z1 : T ) , qKθ ( z1 : T ) = pθ ( z1 : T ) . Proof . We begin by establishing notation . Let pθ ( z1 : T ) be the likelihood of a sequence z1 : T according to random sample decoding . Let qKθ ( z1 : t ) be the likelihood of the same according to Parallel Scheduled Sampling ’ s proposal distribution on pass K. The proof proceeds by induction . First we show that the proposal distribution for the first token matches random sampling ’ s on the first pass , q1θ ( z1 ) = pθ ( z1 ) . Then we show that if q K θ ( z1 : t ) = pθ ( z1 : t ) holds for some K , it also holds for all K ′ > K. Finally , we show that if the previous 1We drop conditioning on x in the following for conciseness statement holds , it also holds for tokens z1 : t+1 on pass K + 1 . Thus , it follows that the proposal distribution matches random sampling ’ s for all T tokens so long as K ≥ T . Base Case : Consider the proposal distribution for the first token on the first pass , z1 . As p = 1 , the first token is sampled from pθ ( z1 ) by construction . Thus , q1θ ( z1 ) = pθ ( z1 ) Induction over K : Suppose that the proposal distribution for tokens qKθ ( z1 : t ) = pθ ( z1 : t ) some K ≥ t. Then the equality also hold for the proposal distribution on pass K + 1 . This follows trivially as tokens z1 : t are “ copied ” from pass K to K + 1 and thus their likelihood is unchanged , qK+1θ ( z1 : t ) = q K θ ( z1 : t ) = pθ ( z1 : t ) Induction over t : Suppose that the proposal distribution matches random sample decoding ’ s for the first t tokens for t = K ; that is , qKθ ( z1 : t ) = pθ ( z1 : t ) . We show that the statement holds for pass K + 1 for tokens z1 : t+1 . First , recall that by construction the proposal distribution for token zt+1 given previous tokens z1 : t is the same as random sampling ’ s when t ≥ K , qK+1θ ( zt+1|z1 : t ) = pθ ( zt+1|z1 : t ) Note that this only holds when p = 1 . Then , qK+1θ ( z1 : t+1 ) = q K+1 θ ( zt+1|z1 : t ) q K+1 θ ( z1 : t ) = qK+1θ ( zt+1|z1 : t ) q K θ ( z1 : t ) = pθ ( zt+1|z1 : t ) pθ ( z1 : t ) = pθ ( z1 : t+1 ) Where we use the chain rule , induction over K for z1 : t , the inductive assumption for qKθ ( z1 : t ) , and the definition of qK+1θ ( zt+1|z1 : t ) when t ≥ K .

Summary: This paper proposes a parallelized version of Scheduled Sampling, which is proposed initially to mitigate Exposure Bias. Exposure Bias occurs when sequence generation models are trained by teacher-forcing, where the t-th token, $y_t$, is forced to be correctly estimated given ground truths, $y_1, ..., y_{t-1}$. Scheduled Sampling randomly selects the previous sequence from generated one or ground truth. Since sequential Scheduled Sampling requires a long time to train, the authors propose a way to parallelize it without degrading the performance. Experimental results for dialog response generation, summarization, and image generation (as a grammar model on images) demonstrate that the proposed method is faster than the original Scheduled Sampling. Additionally, the performance gain is shown in comparison to baselines trained by teacher-forcing.