Improved Language Modeling by Decoding the Past

1 INTRODUCTION . Language modeling is a fundamental task in natural language processing . Given a sequence of tokens , its joint probability distribution can be modeled using the auto-regressive conditional factorization . This leads to a convenient formulation where a language model has to predict the next token given a sequence of tokens as context . Recurrent neural networks are an effective way to compute distributed representations of the context by sequentially operating on the embeddings of the tokens . These representations can then be used to predict the next token as a probability distribution over a fixed vocabulary using a linear decoder followed by Softmax . Starting from the work of Mikolov et al . ( 2010 ) , there has been a long list of works that seek to improve language modeling performance using more sophisticated recurrent neural networks ( RNNs ) ( Zaremba et al . ( 2014 ) ; Zilly et al . ( 2017 ) ; Zoph & Le ( 2016 ) ; Mujika et al . ( 2017 ) ) . However , in more recent work vanilla LSTMs ( Hochreiter & Schmidhuber ( 1997 ) ) with relatively large number of parameters have been shown to achieve state-of-the-art performance on several standard benchmark datasets both in word-level and character-level perplexity ( Merity et al . ( 2018a ; b ) ; Melis et al . ( 2018 ) ; Yang et al . ( 2017 ) ) . A key component in these models is the use of several forms of regularization e.g . variational dropout on the token embeddings ( Gal & Ghahramani ( 2016 ) ) , dropout on the hidden-to-hidden weights in the LSTM ( Wan et al . ( 2013 ) ) , norm regularization on the outputs of the LSTM and classical dropout ( Srivastava et al . ( 2014 ) ) . By carefully tuning the hyperparameters associated with these regularizers combined with optimization algorithms like NT-ASGD ( a variant of the Averaged SGD ) , it is possible to achieve very good performance . Each of these regularizations address different parts of the LSTM model and are general techniques that could be applied to any other sequence modeling problem . In this paper , we propose a regularization technique that is specific to language modeling . One unique aspect of language modeling using LSTMs ( or any RNN ) is that at each time step t , the model takes as input a particular token xt from a vocabulary W and using the hidden state of the LSTM ( which encodes the context till xt ) predicts a probability distribution wt+1 on the next token xt+1 over the same vocabulary as output . Since xt can be mapped to a trivial probability distribution over W , this operation can be interpreted as transforming distributions over W ( Inan et al . ( 2016 ) ) . Clearly , the output distribution is dependent on and is a function of xt and the context further in the past and encodes information about it . We ask the following question – How much information is it possible to decode about the input distribution ( and hence xt ) from the output distribution wt+1 ? In general , it is impossible to decode xt unambiguously . Even if the language model is perfect and correctly predicts xt+1 with probability 1 , there could be many tokens preceding it . However , in this case the number of possibilities for xt will be limited , as dictated by the bigram statistics of the corpus and the language in general . We argue that biasing the language model such that it is possible to decode more information about the past tokens from the predicted next token distribution is beneficial . We incorporate this intuition into a regularization term in the loss function of the language model . The symmetry in the inputs and outputs of the language model at each step lends itself to a simple decoding operation . It can be cast as a ( pseudo ) language modeling problem in “ reverse ” , where the future prediction wt+1 acts as the input and the last token xt acts as the target of prediction . The token embedding matrix and weights of the linear decoder of the main language model can be reused in the past decoding operation . We only need a few extra parameters to model the nonlinear transformation performed by the LSTM , which we do by using a simple stateless layer . We compute the cross-entropy loss between the decoded distribution for the past token and xt and add it to the main loss function after suitable weighting . The extra parameters used in the past decoding are discarded during inference time . We call our method Past Decode Regularization or PDR for short . We conduct extensive experiments on four benchmark datasets for word level and character level language modeling by combining PDR with existing LSTM based language models and achieve new state-of-the-art performance on three of them . 2 PAST DECODE REGULARIZATION ( PDR ) . Let X = ( x1 , x2 , · · · , xt , · · · , xT ) be a sequence of tokens . In this paper , we will experiment with both word level and character level language modeling . Therefore , tokens can be either words or characters . The joint probability P ( X ) factorizes into P ( X ) = T∏ t=1 P ( xt|x1 , x2 , · · · , xt−1 ) ( 1 ) Let ct = ( x1 , x2 , · · · , xt ) denote the context available to the language model for xt+1 . Let W denote the vocabulary of tokens , each of which is embedded into a vector of dimension d. Let E denote the token embedding matrix of dimension |W |× d and ew denote the embedding of w ∈W . An LSTM computes a distributed representation of ct in the form of its hidden state ht , which we assume has dimension d as well . The probability that the next token is w can then be calculated using a linear decoder followed by a Softmax layer as Pθ ( w|ct ) = Softmax ( htET + b ) |w = exp ( hteTw ) ∑ w′∈W exp ( hte T w′ + bw′ ) ( 2 ) where bw′ is the entry corresponding to w′ in a bias vector b of dimension |W | and |w represents projection ontow . Here we assume that the weights of the decoder are tied with the token embedding matrix E ( Inan et al . ( 2016 ) ; Press & Wolf ( 2017 ) ) . To optimize the parameters of the language model θ , the loss function to be minimized during training is set as the cross-entropy between the predicted distribution Pθ ( w|ct ) and the actual token xt+1 . LCE = ∑ t − log ( Pθ ( xt+1|ct ) ) ( 3 ) Note that Eq . ( 2 ) , when applied to all w ∈ W produces a 1 × |W | vector wt+1 , encapsulating the prediction the language model has about the next token xt+1 . Since this is dependent on and conditioned on ct , wt+1 clearly encodes information about it ; in particular about the last token xt in ct . In turn , it should be possible to infer or decode some limited information about xt from wt+1 . We argue that by biasing the model to be more accurate in recalling information about past tokens , we can help it in predicting the next token better . To this end , we define the following decoding operation to compute a probability distribution over wc ∈W as the last token in the context . Pθr ( wc|wt+1 ) = Softmax ( fθr ( wt+1E ) ET + b′θr ) ( 4 ) Here fθr is a non-linear function that maps vectors in Rd to vectors in Rd and b′θr is a bias vector of dimension |W | , together with parameters θr . In effect , we are decoding the past – the last token in the context xt . This produces a vector wrt of dimension 1 × |W | . The cross-entropy loss with respect to the actual last token xt can then be computed as LPDR = ∑ t − log ( Pθr ( xt|wt+1 ) ) ( 5 ) Here PDR stands for Past Decode Regularization . LPDR captures the extent to which the decoded distribution of tokens differs from the actual tokens xt in the context . Note the symmetry between Eqs . ( 2 ) and ( 5 ) . The “ input ” in the latter case is wt+1 and the “ context ” is provided by a nonlinear transformation of wt+1E . Different from the former , the context in Eq . ( 5 ) does not preserve any state information across time steps as we want to decode only using wt+1 . The term wt+1E can be interpreted as a “ soft ” token embedding lookup , where the token vector wt+1 is a probability distribution instead of a unit vector . We add λPDRLPDR to the loss function in Eq . ( 3 ) as a regularization term , where λPDR is a positive weighting coefficient , to construct the following new loss function for the language model . L = LCE + λPDRLPDR ( 6 ) Thus equivalently PDR can also be viewed as a method of defining an augmented loss function for language modeling . The choice of λPDR dictates the degree to which we want the language model to incorporate our inductive bias i.e . decodability of the last token in the context . If it is too large , the model will fail to predict the next token , which is its primary task . If it is zero or too small , the model will retain less information about the last token which hampers its predictive performance . In practice , we choose λPDR by a search based on validation set performance . Note that the trainable parameters θr associated with PDR are used only during training to bias the language model and are not used at inference time . This also means that it is important to control the complexity of the nonlinear function fθr so as not to overly bias the training . As a simple choice , we use a single fully connected layer of size d followed by a Tanh nonlinearity as fθr . This introduces few extra parameters and a small increase in training time as compared to a model not using PDR . 3 EXPERIMENTS . We present extensive experimental results to show the efficacy of using PDR for language modeling on four standard benchmark datasets – two each for word level and character level language modeling . For the former , we evaluate our method on the Penn Treebank ( PTB ) ( Mikolov et al . ( 2010 ) ) and the WikiText-2 ( WT2 ) ( Merity et al . ( 2016 ) ) datasets . For the latter , we use the Penn Treebank Character ( PTBC ) ( Mikolov et al . ( 2010 ) ) and the Hutter Prize Wikipedia Prize ( Hutter ( 2018 ) ) ( also known as Enwik8 ) datasets . Key statistics for these datasets is presented in Table 1 . As mentioned in the introduction , some of the best existing results on these datasets are obtained by using extensive regularization techniques on relatively large LSTMs ( Merity et al . ( 2018a ; b ) ; Yang et al . ( 2017 ) ) . We apply our regularization technique to these models , the so called AWD-LSTM . We consider two versions of the model – one with a single softmax ( AWD-LSTM ) and one with a mixture-of-softmaxes ( AWD-LSTM-MoS ) . The PDR regularization term is computed according to Eq . ( 4 ) and Eq. ( 5 ) . We call our model AWD-LSTM+PDR when using a single softmax and AWDLSTM-MoS+PDR when using a mixture-of-softmaxes . We largely follow the experimental procedure of the original models and incorporate their dropouts and regularizations in our experiments . The relative contribution of these existing regularizations and PDR will be analyzed in Section 6 . There are 7 hyperparameters associated with the regularizations used in AWD-LSTM ( and one extra with MoS ) . PDR also has an associated weighting coefficient λPDR . For our experiments , we set λPDR = 0.001 which was determined by a coarse search on the PTB and WT2 validation sets . For the remaining ones , we perform light hyperparameter search in the vicinity of those reported for AWD-LSTM in Merity et al . ( 2018a ; b ) and for AWD-LSTM-MoS in Yang et al . ( 2017 ) .

Summary: This paper proposes an additional loss term to use when training an LSTM LM.  The authors argue that, intuitively, we want the output distribution to retain some information about the context, or "past".  Given this, they use the output distribution as input to a one layer network that must predict the current token.  The loss for this network is incorporated as an additional term used when training the LM.  The authors show that by adding this loss term they can achieve SOTA (for single softmax model) perplexity on a number of LM benchmarks.