Von Mises-Fisher Loss for Training Sequence to Sequence Models with Continuous Outputs

1 INTRODUCTION . Due to the power law distribution of word frequencies , rare words are extremely common in any language ( Zipf , 1935 ) . Yet , the majority of language generation tasks—including machine translation ( Sutskever et al. , 2014 ; Bahdanau et al. , 2014 ; Luong et al. , 2015 ) , summarization ( Rush et al. , 2015 ; See et al. , 2017 ; Paulus et al. , 2018 ) , dialogue generation ( Vinyals & Le , 2015 ) , question answering ( Yin et al. , 2015 ) , speech recognition ( Graves et al. , 2013 ; Xiong et al. , 2017 ) , and others—generate words by sampling from a multinomial distribution over a closed output vocabulary . This is done by computing scores for each candidate word and normalizing them to probabilities using a softmax layer . Since softmax is computationally expensive , current systems limit their output vocabulary to a few tens of thousands of most frequent words , sacrificing linguistic diversity by replacing the long tail of rare words by the unknown word token , 〈unk〉 . Unsurprisingly , at test time this leads to an inferior performance when generating rare or out-of-vocabulary words . Despite the fixed output vocabulary , softmax is computationally the slowest layer . Moreover , its computation follows a large matrix multiplication to compute scores over the candidate words ; this makes softmax expensive in terms of memory requirements and the number of parameters to learn ( Mnih & Kavukcuoglu , 2013 ; Morin & Bengio , 2005 ; de Brébisson & Vincent , 2016 ) . Several alternatives have been proposed for alleviating these problems , including sampling-based approximations of the softmax function ( Bengio & Senecal , 2003 ; Mnih & Kavukcuoglu , 2013 ) , approaches proposing a hierarchical structure of the softmax layer ( Morin & Bengio , 2005 ; Chen et al. , 2016 ) , and changing the vocabulary to frequent subword units , thereby reducing the vocabulary size ( Sennrich et al. , 2016 ) . We propose a novel technique to generate low-dimensional continuous word representations , or word embeddings ( Mikolov et al. , 2013 ; Pennington et al. , 2014 ; Bojanowski et al. , 2017 ) instead of a probability distribution over the vocabulary at each output step . We train sequence-to-sequence 1The code is at https : //github.com/Sachin19/seq2seq-con models with continuous outputs by minimizing the distance between the output vector and the pretrained word embedding of the reference word . At test time , the model generates a vector and then searches for its nearest neighbor in the target embedding space to generate the corresponding word . This general architecture can in principle be used for any language generation ( or any recurrent regression ) task . In this work , we experiment with neural machine translation , implemented using recurrent sequence-to-sequence models ( Sutskever et al. , 2014 ) with attention ( Bahdanau et al. , 2014 ; Luong et al. , 2015 ) . To the best of our knowledge , this is the first work that uses word embeddings—rather than the softmax layer—as outputs in language generation tasks . While this idea is simple and intuitive , in practice , it does not yield competitive performance with standard regression losses like ` 2 . This is because ` 2 loss implicitly assumes a Gaussian distribution of the output space which is likely false for embeddings . In order to correctly predict the outputs corresponding to new inputs , we must model the correct probability distribution of the target vector conditioned on the input ( Bishop , 1994 ) . A major contribution of this work is a new loss function based on defining such a probability distribution over the word embedding space and minimizing its negative log likelihood ( §3 ) . We evaluate our proposed model with the new loss function on the task of machine translation , including on datasets with huge vocabulary sizes , in two language pairs , and in two data domains ( §4 ) . In §5 we show that our models can be trained up to 2.5x faster than softmax-based models while performing on par with state-of-the-art systems in terms of generation quality . Error analysis ( §6 ) reveals that the models with continuous outputs are better at correctly generating rare words and make errors that are close to the reference texts in the embedding space and are often semantically-related to the reference translation . 2 BACKGROUND . Traditionally , all sequence to sequence language generation models use one-hot representations for each word in the output vocabulary V . More formally , each word w is represented as a unique vector o ( w ) ∈ { 0 , 1 } V , where V is the size of the output vocabulary and only one entry id ( w ) ( corresponding the word ID of w in the vocabulary ) in o ( w ) is 1 and the rest are set to 0 . The models produce a distribution pt over the output vocabulary at every step t using the softmax function : pt ( w ) = esw∑ v∈V e sv ( 1 ) where , sw = Whwht + bw is the score of the word w given the hidden state h produced by the LSTM cell ( Hochreiter & Schmidhuber , 1997 ) at time step t. W ∈ RV xH and b ∈ Rv are trainable parameters . H is the size of the hidden layer h. These parameters are trained by minimizing the negative log-likelihood ( aka cross-entropy ) of this distribution by treating o ( w ) as the target distribution . The loss function is defined as follows : NLL ( pt , o ( w ) ) = − log ( pt ( w ) ) This loss computation involves a normalization proportional to the size of the output vocabulary V . This becomes a bottleneck in natural language generation tasks where the vocabulary size is typically tens of thousands of words . We propose to address this bottleneck by representing words as continuous word vectors instead of one-hot representations and introducing a novel probabilistic loss to train these models as described in §3.2 Here , we briefly summarize prior work that aimed at alleviating the sofmax bottleneck problem . 2There is prior work on predicting word embeddings , but not in conditional language generation with seq2seq . Given a word embedding dictionary , Pinter et al . ( 2017 ) train a character-level neural net that learns to approximate the embeddings . It can then be applied to infer embeddings in the same space for words that were not available in the original set . These models were trained using the ` 2 loss . 2.1 SOFTMAX ALTERNATIVES . We briefly summarize existing modifications to the sofmax layer , capitalizing on conceptually different approaches . Sampling-Based Approximations Sampling-based approaches completely do away with computing the normalization term of softmax by considering only a small subset of possible outputs . These include approximations like Importance Sampling ( Bengio & Senecal , 2003 ) , Noise Constrastive Estimation ( Mnih & Kavukcuoglu , 2013 ) , Negative Sampling ( Mikolov et al. , 2013 ) , and Blackout ( Ji et al. , 2015 ) . These alternatives significantly speed-up training time but degrade generation quality . Structural Approximations Morin & Bengio ( 2005 ) replace the flat softmax layer with a hierarchical layer in the form of a binary tree where words are at the leaves . This alleviates the problem of expensive normalization , but these gains are only obtained at training time . At test time , the hierarchical approximations lead to a drop in performance compared to softmax both in time efficiency and in accuracy . Chen et al . ( 2016 ) propose to divide the vocabulary into clusters based on their frequencies . Each word is produced by a different part of the hidden layer making the output embedding matrix much sparser . This leads to performance improvement both in training and decoding . However , it assigns fewer parameters to rare words which leads to inferior performance in predicting them ( Ruder , 2017 ) . Self Normalization Approaches Andreas et al . ( 2015 ) ; Devlin et al . ( 2014 ) add additional terms to the training loss which makes the normalization factor close to 1 , obviating the need to explicitly normalize . The evaluation of certain words can be done much faster than in softmax based models which is extremely useful for tasks like language modeling . However , for generation tasks , it is necessary to ensure that the normalization factor is exactly 1 which might not always be the case , and thus it might require explicit normalization . Subword-Based Methods Jozefowicz et al . ( 2016 ) introduce character-based methods to reduce vocabulary size . While character-based models lead to significant decrease in vocabulary size , they often differentiate poorly between similarly spelled words with different meanings . Sennrich et al . ( 2016 ) find a middle ground between characters and words based on sub-word units obtained using Byte Pair Encoding ( BPE ) . Despite its limitations ( Oda et al. , 2017 ) , BPE achieves good performance while also making the model truly open vocabulary . BPE is the state-of-the art approach currently used in machine translation . We thus use this as a baseline in our experiments . 3 LANGUAGE GENERATION WITH CONTINUOUS OUTPUTS . In our proposed model , each word type in the output vocabulary is represented by a continuous vector e ( w ) ∈ Rm where m V . This representation can be obtained by training a word embedding model on a large monolingual corpus ( Mikolov et al. , 2013 ; Pennington et al. , 2014 ; Bojanowski et al. , 2017 ) . At each generation step , the decoder of our model produces a continuous vector ê ∈ Rm . The output word is then predicted by searching for the nearest neighbor of ê in the embedding space : wpredicted = argmin w { d ( ê , e ( w ) ) |w ∈ V } where V is the output vocabulary , d is a distance function . In other words , the embedding space could be considered to be quantized into V components and the generated continuous vector is mapped to a word based on the quanta in which it lies . The mapped word is then passed to the next step of the decoder ( Gray , 1990 ) . While training this model , we know the target vector e ( w ) , and minimize its distance from the output vector ê . With this formulation , our model is directly trained to optimize towards the information encoded by the embeddings . For example , if the embeddings are primarily semantic , as in Mikolov et al . ( 2013 ) or Bojanowski et al . ( 2017 ) , the model would tend to output words in a semantic space , that is produced words would either be correct or close synonyms ( which we see in our analysis in §6 ) , or if we use synactico-semantic embeddings ( Levy & Goldberg , 2014 ; Ling et al. , 2015 ) , we might be able to also control for syntatic forms . We propose a novel probabilistic loss function—a probabilistic variant of cosine loss—which gives a theoretically grounded regression loss for sequence generation and addresses the limitations of existing empirical losses ( described in §4.2 ) . Cosine loss measures the closeness between vector directions . A natural choice for estimating directional distributions is von Mises-Fisher ( vMF ) defined over a hypersphere of unit norm . That is , a vector close to the mean direction will have high probability . VMF is considered the directional equivalent of Gaussian distribution 3 . Given a target word w , its density function is given as follows : p ( e ( w ) ; µ , κ ) = Cm ( κ ) e κµT e ( w ) , where µ and e ( w ) are vectors of dimension m with unit norm , κ is a positive scalar , also called the concentration parameter . κ = 0 defines a uniform distribution over the hypersphere and κ = ∞ defines a point distribution at µ. Cm ( κ ) is the normalization term : Cm ( κ ) = κm/2−1 ( 2π ) m/2Im/2−1 ( κ ) , where Iv is called modified Bessel function of the first kind of order v. The output of the model at each step is a vector ê of dimension m. We use κ = ‖ê‖ . Thus the density function becomes : p ( e ( w ) ; ê ) = vMF ( e ( w ) ; ê ) = Cm ( ‖ê‖ ) eê T e ( w ) ( 2 ) It is noteworthy that equation 2 is very similar to softmax computation ( except that e ( w ) is a unit vector ) , the main difference being that normalization is not done by summing over the vocabulary , which makes it much faster than the softmax computation . More details about it ’ s computation are given in the appendix . The negative log-likelihood of the vMF distribution , which at each output step is given by : NLLvMF ( ê ; e ( w ) ) = − log ( Cm ( ‖ê‖ ) ) − êTe ( w ) Regularization of NLLvMF In practice , we observe that the NLLvMF loss puts too much weight on increasing ‖ê‖ , making the second term in the loss function decrease rapidly without significant decrease in the cosine distance . To account for this , we add a regularization term . We experiment with two variants of regularization . NLLvMFreg1 : We add λ1‖ê‖ to the loss function , where λ1 is a scalar hyperparameter.4 This makes intuitive sense in that the length of the output vector should not increase too much . The regularized loss function is as follows : NLLvMFreg1 ( ê ) = − logCm ( ‖ê‖ ) − êTe ( w ) + λ1‖ê‖ NLLvMFreg2 : We modify the previous loss function as follows : NLLvMFreg2 ( ê ) = − logCm ( ‖ê‖ ) − λ2êTe ( w ) ( 3 ) − logCm ( ‖ê‖ ) decreases slowly as ‖ê‖ increases as compared the second term . Adding a λ2 < 1 the second term controls for how fast it can decrease.5

Summary: This paper describes a technique for replacing the softmax layer in sequence-to-sequence models with one that attempts to predict a continuous word embedding, which will then be mapped into a (potentially huge) pre-trained embedding vector via nearest neighbor search. The obvious choice for building a loss around such a prediction (squared error) is shown to be inappropriate empirically, and instead a von Mises-Fisher loss is proposed. Experiments conducted on small-data, small-model, greedy-search German->English, French->English and English->French scenarios demonstrate translation quality on par with BPE, and superior performance to a number of other continuous vector losses. They also provide convincing arguments that this new objective is more efficient in terms of both time and number of learned parameters.