Title: Breaking the Softmax Bottleneck: A High-Rank RNN Language Model
Authors: Zhilin Yang
Zihang Dai
Ruslan Salakhutdinov
William W. Cohen
Venue: ICLR 2018

Abstract: We formulate language modeling as a matrix factorization problem, and show that the expressiveness of Softmax-based models (including the majority of neural language models) is limited by a Softmax bottleneck. Given that natural language is highly context-dependent, this further implies that in practice Softmax with distributed word embeddings does not have enough capacity to model natural language. We propose a simple and effective method to address this issue, and improve the state-of-the-art perplexities on Penn Treebank and WikiText-2 to 47.69 and 40.68 respectively. The proposed method also excels on the large-scale 1B Word dataset, outperforming the baseline by over 5.6 points in perplexity.


Review discussion:
The authors argue in this paper that due to the limited rank of the  context-to-vocabulary logit matrix in the currently used version of the softmax output layer, it is not able to capture the full complexity of language.
As a result, they propose to use a mixture of softmax output layers instead where the mixing probabilities are context-dependent, which allows to obtain a full rank logit matrix in complexity linear in the number of mixture components (here 15). This leads to improvements in the word-level perplexities of the PTB and wikitext2 data sets, and Switchboard BLEU scores.
The question of the expressiveness of the softmax layer, as well as its suitability for word-level prediction, is indeed an important one which has received too little attention. This makes a lot of the questions asked in this paper extremely relevant to the field.
However, it is unclear that the rank of the logit matrix is the right quantity to consider.
For example, it is easy to describe a rank D NxM matrix where up to 2^D lines have max values at different indices.
Further, the first two "observations" in Section 2.2 would be more accurately described as "intuitions" of the authors. As they write themselves "there is no evidence showing that semantic meanings are fully linearly correlated." Why then try to link "meanings" to basis vectors for the rows of A?
To be clear, the proposed model is undoubtedly more expressive than a regular softmax, and although it does come at a substantial computational cost (a back-of-the envelope calculation tells us that computing 15 components of 280d MoS takes the same number of operations as one with dimension 1084 = sqrt (280*280*15)), it apparently manages not to drastically increase overfitting, which is significant.
Unfortunately, this is only tested on relatively small data sets, up to 2M tokens and a vocabulary of size 30K for language modeling. They do constitute a good starting place to test a model, but given the importance of regularization on those specific tasks, it is difficult to predict how the MoS would behave if more training data were available, and if one could e.g. simply try a 1084 dimension embedding for the softmax without having to worry about overfitting.
Another important missing experiment would consist in varying the number of mixture components (this could very well be done on WikiText2). This could help validate the hypothesis: how does the estimated rank vary with the number of components? How about the performance and pairwise KL divergence?
This paper offers a promising direction for language modeling research, but would require more justification, or at least a more developed experimental section.
Pros:
- Important starting question
- Thought-provoking approach
- Experimental gains on small data sets
Cons:
- The link between the intuition and reality of the gains is not obvious
- Experiments limited to small data sets, some obvious questions remain
Thank you for the valuable feedback.
[[Rank and meanings]] It could be possible that A is low-rank for a natural language as it is hard to rule out this possibility rigorously, but we hypothesize that A is high-rank. Our hypothesis is supported by our intuitive reasoning and empirical experiments. Empirically, we give three more pieces of evidences supporting our hypothesis that the rank is the key bottleneck of Softmax and MoS improves the performance by solving the rank bottleneck (section 3.2):
- Before the rank saturates to the full rank, using more mixture components in MoS continues to increase the rank of the log-probability matrix.
Further, when the rank increases, the perplexity also decreases.
- MoS has a similar generalization gap compared to Softmax, which rules out the concern that the improvement actually comes from some unexpected regularization effects of MoS.
- In character-level language modeling, since the largest possible rank is upper bounded by the limited vocabulary size, Softmax does not suffer from the rank bottleneck. In this case, Softmax and MoS have almost the same performance, which matches our analysis.
We agree that linking semantic meanings to bases lacks rigor and this would be better described as intuitions. We have made corresponding changes in the paper.
[[Computation vs. Capacity]]: It is true that MoS involves a larger amount of computation compared to the standard Softmax.
However, [Collins et al] suggests the capacity of neural language models is mostly related to the number of parameters, rather than computation.
Moreover, powerful models often require a larger amount of computation.
For example, the attention based seq2seq model involves much more computation compared to the vanilla seq2seq.
[[Large-scale experiment]]: We have added a “large-scale language modeling experiment” using the 1B Word Dataset (section 3.1), where MoS significantly outperforms the baseline model with a large margin. This indicates that MoS consistently outperforms Softmax, regardless of the scale of the dataset.
Also, note that PTB and WT2 are two de-facto benchmarks widely used in previous work on language modeling. None of the following papers had experiments on datasets larger than WT2: Zoph & Le ICLR 2017, Zilly et al ICML 2017, Inan et al ICLR 2017, Grave et al ICLR 2017, Merity et al ICLR 2017.
[[Varying the number of mixtures]]: Thanks for the suggestion. We performed this experiment, whose result is summarized in the second bullet point of section 3.2 (updated version). As expected, the number of mixture components is positively correlated with the empirical rank. More importantly, before the rank saturates to the full rank, MoS with a higher rank leads to a better performance (lower perplexity).
[Collins et al] Capacity and Trainability in Recurrent Neural Networks
The authors have added some important experiments whose results support their claim, and I do believe that the current version of the paper makes a stronger case. I am still not satisfied that the rank explanation fully captures what is going on there, although it certainly correlates with better results, but this paper will provide an important data point for future research and I am raising my score from 5 to 7.
The authors has addressed my concerns, so I raised my rating.
The paper is grounded on a solid theoretical motivation and the analysis is sound and quite interesting.
There are no results on large corpora such as 1 billion tokens benchmark corpus, or at least medium level corpus with 50 million tokens. The corpora the authors choose are quite small, the variance of the estimates are high, and similar conclusions might not be valid on a large corpus.
[1] provides the results of character level language models on Enwik8 dataset, which shows regularization doesn't have much effect and needs less tuning. Results on this data might be more convincing.
The results of MOS is very good, but the computation complexity is much higher than other baselines. In the experiments, the embedding dimension of MOS is slightly smaller, but the number of mixture is 15. This will make it less usable, I think it's necessary to provide the training time comparison.
Finally experiments on machine translation or speech recognition should be done and to see what improvements the proposed method could bring for BLEU or WER.
[1] Melis, Gábor, Chris Dyer, and Phil Blunsom. "On the state of the art of evaluation in neural language models." arXiv preprint arXiv:1707.05589 (2017).
[2] Joris Pelemans, Noam Shazeer, Ciprian Chelba, Sparse Non-negative Matrix Language Modeling,  Transactions of the Association for Computational Linguistics, vol. 4 (2016), pp. 329-342
[3] Shazeer et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR 2017
It's true that regularization doesn't play as important a role with large datasets as it does with small datasets, but enwik8 is character based and it's unclear whether this paper's arguments would apply. Sticking to word level corpora, I'd much sooner recommend Wikitext-103 than the Billion Word corpus which has issues.
Furthermore, language modelling improvements are interesting in their own right without having to validate them via MT or speech recognition.
Thanks for your valuable comments.
[[Large-scale experiment]]: We’ve added a “large-scale language modeling experiment” using the 1B Word Dataset (section 3.1), where MoS significantly outperforms the baseline model by a large margin. This indicates that MoS consistently outperforms Softmax, regardless of the scale of the dataset.
Also, note that PTB and WT2 are two de-facto benchmarks widely used in previous work on language modeling. None of the following papers had experiments on datasets larger than WT2: Zoph & Le ICLR 2017, Zilly et al ICML 2017, Inan et al ICLR 2017, Grave et al ICLR 2017, Merity et al ICLR 2017.
[[Character-level LM]]: Firstly, note that the largest possible rank of the log-probability matrix is upper bounded by the vocabulary size. In character-level LM, the vocabulary size is usually much smaller than the embedding size. In this case, Softmax does not suffer from the rank bottleneck problem, and we expect MoS and Softmax to achieve similar performance in practice. To verify our expectation, we perform character-level LM experiment on the text8 dataset, where MoS and Softmax indeed achieve almost the same performance (section 3.2 & appendix C.2 in the updated version).
[[Training time]]: We have added the training time analysis for MoS and provided empirical numbers in the updated versions of the paper (section 3.3 & Appendix C.3). In general, computational wall time of MoS is actually sub-linear w.r.t. the number of mixture components. In most settings, we observe a two to three times slowdown compared to Softmax when using up to 15 components for MoS. We believe such additional computational cost is acceptable for the following reasons:
- MoS is highly parallelizable, meaning that using more machines can always speed up the computation almost linearly.
- The field of deep learning systems (both hardware and software) is making rapid progress. It might be possible to further optimize MoS on GPUs for fast computation. More developed hardware systems would also further reduce the computational cost.
- Historically, important techniques sometimes come with an additional computational cost, e.g., LSTM, attention, deep ResNets. We believe that with MoS, the extra cost is reasonable and the gain is substantial.
[[Application to MT/ASR]]: We believe this is best left to future research, as performing rigorous experiments and careful comparison for such a real-world applications is non-trivial. And we believe language modeling is of its own importance already.
Language models are important components to many NLP tasks. The current state-of-the-art language models are based on recurrent neural networks which compute the probability of a word given all previous words using a softmax function over a linear function of the RNN's hidden state. This paper argues the softmax is not expressive enough and proposes to use a more flexible mixture of softmaxes. The use of a mixture of softmaxes is motivated from a theoretical point of view by translating language modeling into matrix factorization.
Pros:
--The paper is very well written and easy to follow. The ideas build up on each other in an intuitive way.
--The idea behind the paper is novel: translating language modeling into a matrix factorization problem is new as far as I know.
--The maths is very rigorous.
--The experiment section is thorough.
Cons:
--To claim SOTA all models need to be given the same capacity (same number of parameters). In Table 2 the baselines have a lower capacity. This is an unfair comparison
--I suspect the proposed approach is slower than the baselines. There is no mention of computational cost. Reporting that would help interpret the numbers.
The SOTA claim might not hold if baselines are given the same capacity.
But regardless of this, the paper has very strong contributions and deserves acceptance at ICLR.
Thank you for the valuable comments.
[[Claim of SOTA]]: We believe the 2M difference in the number of parameters is negligible compared to the number of parameters we use (i.e., 35M).
In fact, we ran MoS in another setting with 31M parameters and got 63.59 on WT2 without finetuning, compared to 63.33 obtained by our best-performing model.
[[Training time]]: Thanks for the suggestion and we have added the training time analysis for MoS and provided empirical numbers in the updated versions of the paper (section 3.3 & Appendix C.3). In general, computational wall time of MoS is actually sub-linear w.r.t. the number of mixture components. In most settings, we observe a two to three times slowdown compared to Softmax when using up to 15 components for MoS. We believe such additional computational cost is acceptable for the following reasons:
- MoS is highly parallelizable, meaning that using more machines can always speed up the computation almost linearly.
- The field of deep learning systems (both hardware and software) is making rapid progress. It might be possible to further optimize MoS on GPUs for fast computation. More developed hardware systems would also further reduce the computational cost.
- Historically, important techniques sometimes come with an additional computational cost, e.g., LSTM, attention, deep ResNets. We believe that with MoS, the extra cost is reasonable and the gain is substantial.
To provide faithful answers to the reviews and comments, we have conducted a series of additional experiments and updated the paper accordingly. The core changes are summarized as follows:
(1) we add a “large-scale language modeling experiment” using the 1B Word Dataset (section 3.1)
(2) we give three more pieces of “evidence supporting our theory” that the rank is the key bottleneck of Softmax and MoS improves the performance by breaking the rank bottleneck (section 3.2):
- Empirically, before the rank saturates to the full rank, using more mixture components in MoS continues to increase the rank of the log-probability matrix.
Further, the higher the rank is, the lower the perplexity that can be achieved.
- MoS has a similar generalization gap compared to Softmax, which rules out the concern that the improvement actually comes from some unexpected regularization effects of MoS.
- In character-level language modeling, since the largest possible rank is upper bounded by the limited vocabulary size, Softmax does not suffer from the rank bottleneck. In this case, Softmax and MoS have almost the same performance, which matches our analysis.
(3) we perform “training time analysis” for MoS and provide empirical training time comparison (section 3.3 & Appendix C.3)

Meta-review:
role: PC
data: comment: Viewing language modeling as a matrix factorization problem, the authors argue that the low rank of word embeddings used by such models limits their expressivity and show that replacing the softmax in such models with a mixture of softmaxes provides an effective way of overcoming this bottleneck. This is an interesting and well-executed paper that provides potentially important insight. It would be good to at least mention prior work related to the language modeling as matrix factorization perspective (e.g. Levy & Goldberg, 2014).

Decision: Accept (Oral)