Adaptive Input Representations for Neural Language Modeling

1 INTRODUCTION . Language modeling is a basic task in natural language processing , with many applications such as speech recognition ( Arisoy et al. , 2012 ) and statistical machine translation ( Schwenk et al. , 2012 ; Vaswani et al. , 2013 ; Baltescu & Blunsom , 2015 ) . Recently , much progress has been made by neural methods ( Bengio et al. , 2003 ; Mikolov et al. , 2010 ) based on LSTMs ( Józefowicz et al. , 2016 ) , gated convolutional networks ( Dauphin et al. , 2017 ) and self-attentional networks ( Al-Rfou et al. , 2018 ) . There are different choices for the basic unit we wish to model , including full words ( Bengio et al. , 2003 ) , characters for the input ( Kim et al. , 2016 ) , or also the output ( Merity et al. , 2018 ) as well as sub-words ( Buckman & Neubig , 2018 ; Mielke & Eisner , 2018 ) . Word-based models are particularly challenging since computing probabilities for all 800K words of the BILLION WORD benchmark is still a substantial part of the overall computation ( Chen et al. , 2016 ) . A popular approach to lower the computational burden is to structure the output vocabulary so that not all probabilities need to be computed . The hierarchical softmax does this by introducing latent variables or clusters to simplify normalization ( Goodman , 2001 ; Morin & Bengio , 2005 ; Mikolov et al. , 2011 ) . This has been further improved by the adaptive softmax which introduces a variable capacity scheme for output word embeddings , assigning more parameters to frequent words and fewer parameters to rare words ( Grave et al. , 2017 ) . In this paper , we introduce adaptive input embeddings which extend the adaptive softmax to input word representations . This factorization assigns more capacity to frequent words and reduces the capacity for less frequent words with the benefit of reducing overfitting to rare words . For a competitive setup on the BILLION WORD benchmark , adaptive input embeddings reduce the number of parameters in the input and output layers by 23 % while achieving higher accuracy over fixed size embeddings . When the adaptive input representations are tied with an adaptive softmax in the output , then the number of parameters is reduced by a total of 61 % . Our experiments compare models based on word inputs , character inputs , as well as sub-word units using a self-attention architecture ( Vaswani et al. , 2017 ) . We show that models with adaptive word representations can outperform very strong character-based models while training more than twice as fast . We also substantially improve adaptive softmax by introducing additional dropout regularization in the tail projection . On the WIKITEXT-103 benchmark we achieve a perplexity of 18.7 , a 1Code and pre-trained models available at http : //github.com/pytorch/fairseq reduction of 10.5 perplexity over the previously best published result . On the larger BILLION WORD benchmark our best model with adaptive input embeddings achieves 23.02 perplexity , a reduction of nearly 5 perplexity over the next best previously published result . 2 RELATED WORK . Adaptive word representations are inspired by the adaptive softmax work Grave et al . ( 2017 ) which first described a GPU friendly way to construct a hierarchical softmax and showed that it performs very competitively compared to a full softmax , while offering significantly faster speed and a lower memory footprint . Merity et al . ( 2018 ) use a modified version of adaptive softmax which does not reduce the dimensionality of less frequent words in order to be able to share output embeddings with the input . This setup is akin to a hierarchical softmax with tied weights . We show that variable-sized input embeddings can perform better than fixed sized embeddings . Furthermore , this also enables weight sharing with an adaptive softmax output layer . Merity et al . ( 2018 ) evaluates both character-based and word-based factorizations but does not directly compare them to each other . We perform a direct comparison of word-based and characterbased input vocabularies and also compare to a sub-word factorization for both the input and output . Recently , Al-Rfou et al . ( 2018 ) demonstrated that self-attentional models can perform very well on language modeling tasks where the input and output is both characters . We also consider word-based benchmarks . 3 ADAPTIVE INPUT REPRESENTATIONS . The adaptive softmax exploits the fact that the distribution of word types in natural language follows a Zipfian distribution in order to improve the computation of the output probabilities . We apply the same intuition for input word embeddings with the motivation to reduce the number of parameters which frees up capacity for other parts of the model . We define a number of clusters that partitions the frequency ordered vocabulary V = V1 ∪ V2 , . . . , Vn−1 ∪ Vn such that Vi ∩ Vj = ∅ for ∀i , j , and i 6= j , where V1 contains the most frequent words and Vn the least frequent words . We will refer to V1 as the head and to any subsequent clusters loosely as tail . We reduce the capacity for each cluster by a factor of k. That is , if words in V1 have dimension d , then words in Vn have dimension dkn−1 . We typically set k = 4 following Grave et al . ( 2017 ) . Next , we add linear projections W1 ∈ Rd×d , . . . , Wn ∈ Rd/k n−1×d to map the embeddings of each cluster to dimension d so that the concatenated output of the adaptive input embedding layer can be easily used by the subsequent model ( Figure 1 ) . We also project V1 which already has dimension d. When presented with a number of input words , the adaptive input embedding layer partitions the words into the various clusters , performs separate lookups in the embedding tables and then projects to dimension d , followed by concatenating the embeddings in the original order . Weight sharing . When the output layer is an adaptive softmax with the same partition of V , d , and k as the adaptive input layer , then we can tie the weights ( Inan et al. , 2016 ; Press & Wolf , 2017 ) . This further reduces the number of parameters and can simultaneously improve performance ( §5 ) . We can share both the parameters for the actual words as well as the projections W1 , . . . , Wn Sharing the word embeddings is straightforward except for the head where the adaptive softmax has n− 1 additional embeddings for the remaining clusters which are not shared with the input . We share all projections , except for the head projection which is not available in the adaptive softmax since the model output is directly multiplied with the output word embeddings for the head band . Performance decreased when we added a head projection to the adaptive softmax in the output , regardless of when it was shared or not . Sharing both the word embeddings as well as the projections performed very well on WIKITEXT-103 but on BILLION WORD we only share the word embeddings as we found that this performed better on the validation set . 4 EXPERIMENTAL SETUP . 4.1 MODEL . We follow most of the architectural choices described in Vaswani et al . ( 2017 ) but use only a decoder network . We add sinusoidal position embeddings to the input layer and stack N = 16 blocks for both BILLION WORD and WIKITEXT-103 . Each block contains two sub-blocks : the first is a multihead self-attention module with H = 16 heads . The second sub-block is a feed-forward module ( FFN ) of the form ReLU ( W1X+ b1 ) W2+ b2 where W1 ∈ Re×eff , W1 ∈ Reff×e and e = 1024 , eff = 4096 unless otherwise stated . Different to Vaswani et al . ( 2017 ) we apply layer normalization before the self-attention and FFN blocks instead of after , as we find it leads to more effective training . Sub-blocks are surrounded by a residual connection ( He et al. , 2015 ) . We use a dropout rate of 0.1 and attention dropout of 0.1 for BILLION WORD models , and increase regularization for WIKITEXT-103 by using dropout 0.3 , and 0.1 ReLU dropout as well as attention dropout 0.1 . We use the same hyperparameters for all models trained on the same dataset in order to enable a like for like comparison . When the dimensionality of the input or output layer differs from e , then we add a simple linear projection with no bias . 4.2 DATASETS . We experiment on the BILLION WORD benchmark and WIKITEXT-103 . BILLION WORD contains 768M word tokens and has a vocabulary of about 800K word types , which corresponds to words with more than 3 occurrences in the training set ( Chelba et al. , 2013 ) . The training data of WIKITEXT-103 comprises about 100M tokens and a vocabulary of around 260K , corresponding to types with more than 3 occurrences in the training data ( Merity et al. , 2016 ) . The dataset is composed of shuffled Wikipedia articles where the context carries across sentences . 4.3 BATCHING . For BILLION WORD we batch individual sentences since the corpus does not contain document structure . For WIKITEXT-103 we partition the training data into blocks of 512 contiguous tokens ignoring document boundaries . Evaluation is the same except that we require blocks to contain complete sentences totaling up to 512 tokens.2 2 Respecting document boundaries may lead to better results and we leave this to future work . We limit the number of tokens per GPU to a maximum threshold B per GPU . That is , we add examples of similar length until we reach this threshold . When we train on multiple GPUs , each GPU processes B tokens using the same model parameters . This increases the effective batch size to the product of the number of GPUs and B . For BILLION WORD models we use B = 2048 and typically train on 32 GPUs , giving an effective batch size of 65K tokens . The smaller vocabulary of WIKITEXT-103 enables increasing B to 4096 and we train on 8 GPUs . We found that large batch training is beneficial for this dataset and we therefore accumulate gradient updates over two batches before committing a parameter update ( Ott et al. , 2018a ) . This gives an effective batch size of 65K tokens for WIKITEXT-103 . 4.4 INPUT AND OUTPUT LAYER HYPERPARAMETERS . Embedding sizes . For fixed size word input layers and softmax output layers we generally use embeddings of size 512 for WIKITEXT-103 . When we use an adaptive softmax in the output and fixed size word embeddings for the input , then we use dimension 256 for the input embeddings for BILLION WORD and 64 for WIKITEXT-103 . We tuned this choice on the validation set ( Appendix A ) . BPE inputs and outputs have embeddings of size 1024 . Character CNN . We model character inputs by convolving the representations of all characters in a word following Kim et al . ( 2015 ) which applies several filters , then max pooling , a number of highway layers and a projection . Character embeddings have size 128 and we apply seven filters of size 1x128 , 2x256 , 3x384 , 4x512 , 5x512 , 6x512 , 7x512 , where 3x128 indicates a filter processing three characters that outputs 128 features . We use a single highway layer for WIKITEXT-103 , and two for BILLION WORD . We do not add start of word and end of word markers as they did not improve validation accuracy . We train on the same pre-processed data as the other models , with unknown tokens in both the inputs and outputs . Adaptive input representations and adaptive softmax . We use an adaptive softmax output layer to train models with large word-based vocabularies . For adaptive word inputs and adaptive softmax , we use embeddings of size d = 1024 for the head and reduce the size of subsequent clusters by a factor of k = 4 . For WIKITEXT-103 , we have three bands of size 20K ( d=1024 ) , 40K ( d=256 ) and 200K ( d=64 ) . For BILLION WORD the bands are 60K ( d=1024 ) , 100K ( d=256 ) , and 640K ( d=64 ) . Sub-word models . We learn a byte-pair encoding ( BPE ) of 32K codes on the training data of each benchmark ( Sennrich et al. , 2016 ) . After applying the code to the training data we obtain a vocabulary of 33,337 tokens for WIKITEXT-103 and 32,347 tokens for BILLION WORD . BPE input/output embeddings have size 1024 . The final evaluation is in terms word-level perplexity to be comparable to other models . The probability of a word is the product of the sub-word units .

Summary: This paper introduced a new architecture for input embeddings of neural language models: adaptive input representation (ADP). ADP allowed a model builder to define a set of bands of input words with different frequency where frequent words have larger embedding size than the others. The embeddings of each band are then projected into the same size. This resulted in lowering the number of parameters.