 version of the adaptive softmax (Grave et al., 2016b), extended to allow for tied weights (Inan et al., 2016; Press & Wolf, 2016). While other softmax approximation strategies exist (Morin & Bengio, 2005; Gutmann & Hyvärinen, 2010), the adaptive softmax has been shown to achieve results close to that of the full softmax whilst maintaining high GPU efficiency.
The adaptive softmax uses a hierarchy determined by word frequency to reduce computation time. Words are split into two levels: the first level (the short-list) and the second level, which forms clusters of rare words. Each cluster has a representative token in the short-list which determines the overall probability assigned across all words within that cluster. Due to Zipf’s law, most words will only require a softmax over the short-list during training, reducing computation and memory usage. Clusters on the second level are also constructed such that the matrix multiplications required for a standard batch are near optimal for GPU efficiency.
In the original adaptive softmax implementation, words within clusters on the second level feature a reduced embedding size, reducing from embedding dimensionality d to d4. This minimizes both total parameter count and the size of the softmax matrix multiplications. The paper justifies this by noting that rare words are unlikely to need full embedding size fidelity due to how infrequently they occur.
Weight tying cannot be used with this memory optimization however. As weight tying re-uses the word vectors from the embedding as the targets in the softmax, the embedding dimensionality must be equal for all words. We discard the adaptive softmax memory optimization in order to utilize tied weights. Counter-intuitively, weight tying actually reduces the memory usage even more than the adaptive softmax memory optimization, as weight tying allows us to halve the memory usage by re-using the word vectors from the embedding layer.
When a dataset is processed using a small vocabulary, all the words can be placed into the short-list, thus reducing the adaptive softmax to a standard softmax.
Table 1: Statistics of the character-level Penn Treebank, character-level enwik8 dataset, and WikiText-103. The out of vocabulary (OoV) rate notes what percentage of tokens have been replaced by an ⟨unk⟩ token, not applicable to character-level datasets.
Our work is over three datasets, two character-level datasets (Penn Treebank, enwik8) and a large-scale word-level (WikiText-103) dataset. A set of overview statistics for these datasets are presented in Table 1.
Table 2: Bits Per Character (BPC) on character-level Penn Treebank.
Table 3: Bits Per Character (BPC) on enwik8.
In Mikolov et al. (2010), the Penn Treebank dataset (Marcus et al., 1994) was processed to generate a character-level language modeling dataset. The dataset exists in both a word- and character-level form.
While the dataset was originally composed of Wall Street Journal articles, the preprocessed dataset removes many features considered important for capturing language modeling. Oddly, the vocabulary of the words in the character-level dataset is limited to 10,000 - the same vocabulary as used in the word level dataset. This vastly simplifies the task of character-level language modeling as character transitions will be limited to those found within the limited word level vocabulary.
In addition to the limited vocabulary, the character-level dataset is all lower case, has all punctuation removed (other than as part of certain words such as u.s. or mr.), and replaces all numbers with N. All of these would be considered important subtasks for a character-level language model to cover.
When rare words are encountered in both the character and word level datasets, they’re replaced with the token <unk>. This makes little sense for the character-level dataset, which doesn’t suffer from out of vocabulary issues as the word level dataset would. As the <unk> token is the only token to use the < and > characters, a sufficiently advanced model should always output unk> upon seeing < as this is the only time angle brackets are used.
We later explore these issues and their impact by comparing models trained on character-level Penn Treebank with a more realistic character-level dataset, enwik8.
We report our results, using the bits per character (BPC) metric, in Table 2. The model uses a three-layered LSTM with a BPTT length of 150, embedding sizes of 200 and hidden layers of size 1000. We regularize the model using LSTM dropout, weight dropout (Merity et al., 2018), and weight decay. These values are tuned coarsely. For full hyper-parameter values, refer to Table 6. We train the model using the Adam (Kingma & Ba, 2014) optimizer with a learning rate of 2.7×10−3 for 500 epochs and reduce the learning rate by 10 on epochs 300 and 400. While the model was not optimized for convergence speed, it takes only 8 hours to train on an NVIDIA Volta GPU. Both our LSTM and QRNN model beat the current state-of-the-art though using more parameters. Of note is that the QRNN model uses 6 layers to achieve the same result as the LSTM with only 3 layers, suggesting that the limited recurrent computation capacity of the QRNN may be an issue.
The Hutter Prize Wikipedia dataset (Hutter, 2018), also known as enwik8, is a byte-level dataset consisting of the first 100 million bytes of a Wikipedia XML dump. For simplicity we shall refer to it as a character-level dataset. Within these 100 million bytes are 205 unique tokens. The Hutter Prize was launched in 2006 and is focused around compressing the enwik8 dataset as efficiently as possible.
The XML dump contains a wide array of content, including English articles, XML data, hyperlinks, and special characters. As the data has not been processed, it features many of the intricacies of language modeling, both natural and artificial, that we would like our models to capture. For our experiments, we follow a standard setup where the train, validation and test sets consist of the first 90M, 5M, and 5M characters, respectively.
We report our results in Table 3. The model uses a three-layered LSTM with a BPTT length of 200, embedding sizes of 400 and hidden layers of size 1850. The only explicit regularizer we employ is the weight dropped LSTM (Merity et al., 2018) of magnitude 0.2. For full hyper-parameter values, refer to Table 6. We train the model using the Adam (Kingma & Ba, 2014) optimizer with default hyperparameter for 50 epochs and reduce the learning rate by 10 on epochs 25 and 35.
This dataset is far more challenging than character-level Penn Treebank for multiple reasons. The dataset has 18 times more training tokens and the data has not been processed at all, maintaining far more complex character to character transitions than that of the character-level Penn Treebank. Potentially as a result of this, the QRNN based model underperforms models with comparable numbers of parameters. The limited recurrent computation capacity of the QRNN appears to become a major issue when moving toward realistic character-level datasets.
The WikiText-2 (WT2) and WikiText-103 (WT103) datasets introduced in Merity et al. (2017b) contain lightly preprocessed Wikipedia articles, retaining the majority of punctuation and numbers. The WT2 data set contains approximately 2 million words in the training set and 0.2 million in validation and test sets. The WT103 data set contains a larger training set of 103 million words and the same validation and testing set as WT2. As the Wikipedia articles are relatively long and are focused on a single topic, capturing and utilizing long term dependencies can be key to models obtaining strong performance.
The underlying model used in this paper achieved state-of-the-art on language modeling for the word-level PTB and WikiText-2 datasets (Merity et al., 2018) We show that, in conjunction with the tied adaptive softmax, we achieve state-of-the-art perplexity on the WikiText-103 data set using the AWD-QRNN framework; see Table 4. We opt to use the QRNN as the LSTM is 3 times slower, as reported in Table 5, and a QRNN based model’s performance on word-level datasets has been found to equal that of an LSTM based model Merity et al. (2018).
By utilizing the QRNN and tied adaptive softmax, we were able to train to a state-of-the-art result with a NVIDIA Volta GPU in 12 hours. The model used consisted of a 4-layered QRNN model with an embedding size of 400 and 2500 nodes in each hidden layer. We trained the model using a batch size of 60 and a BPTT length of 140 using the Adam optimizer (Kingma & Ba, 2014) for 14 epochs, reducing the learning rate by 10 on epoch 12. To avoid over-fitting, we employ the regularization strategies proposed in (Merity et al., 2018) including variational dropout, random sequence lengths, and L2-norm decay. The values for the model hyperparameters were tuned only coarsely; for full hyper-parameter values, refer to Table 6.
Table 4: Perplexity on the word-level WikiText-103 dataset. The model was trained for 12 hours (14 epochs) on an NVIDIA Volta.
Table 5: Mini-batch timings during training on WikiText-103. This is a 3.1 times speed-