B REAKING THE S OFTMAX B OTTLENECK : A H IGH -R ANK RNN L ANGUAGE M ODEL
A BSTRACT
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
I NTRODUCTION
As a fundamental task in natural language processing, statistical language modeling has gone through signiﬁcant development from traditional Ngram language models to neural language models in the last decade (Bengio et al., 2003; Mnih & Hinton, 2007; Mikolov et al., 2010).
Despite the huge variety of models, as a density estimation problem, language modeling mostly relies on a universal auto-regressive factorization of the joint probability and then models each conditional factor using different approaches.
Speciﬁcally, given a corpus of tokens X =
( X 1 , . . .
, X T ) , the joint probability P ( X ) factorizes as P ( X ) =
(cid:81) t P ( X t | X <t ) =
(cid:81) t P ( X t | C t ) , where C t =
X <t is referred to as the context of the conditional probability hereafter.
Based on the factorization, recurrent neural networks (RNN) based language models achieve stateof-the-art results on various benchmarks (Merity et al., 2017; Melis et al., 2017; Krause et al., 2017).
A standard approach is to use a recurrent network to encode the context into a ﬁxed size vector, which is then multiplied by the word embeddings (Inan et al., 2016; Press & Wolf, 2017) using dot product to obtain the logits.
The logits are consumed by the Softmax function to give a categorical probability distribution over the next token.
In spite of the expressiveness of RNNs as universal approximators (Schäfer & Zimmermann, 2006), an unclear question is whether the combination of dot product and Softmax is capable of modeling the conditional probability, which can vary dramatically with the change of the context.
In this work, we study the expressiveness of the aforementioned Softmax-based recurrent language models from a perspective of matrix factorization.
We show that learning a Softmax-based recurrent language model with the standard formulation is essentially equivalent to solving a matrix factorization problem.
More importantly, due to the fact that natural language is highly context-dependent, the matrix to be factorized can be high-rank.
This further implies that standard Softmax-based language models with distributed (output) word embeddings do not have enough capacity to model natural language.
We call this the Softmax bottleneck .
We propose a simple and effective method to address the Softmax bottleneck.
Speciﬁcally, we introduce discrete latent variables into a recurrent language model, and formulate the next-token probability distribution as a Mixture of Softmaxes (MoS).
Mixture of Softmaxes is more expressive than Softmax and other surrogates considered in prior work.
Moreover, we show that MoS learns matrices that have much larger normalized singular values and thus much higher rank than Softmax and other baselines on real-world datasets.
We evaluate our proposed approach on standard language modeling benchmarks.
MoS substantially improves over the current state-of-the-art results on benchmarks, by up to 3.6 points in terms of perplexity, reaching perplexities 47.69 on Penn Treebank and 40.68 on WikiText-2.
We further apply MoS to a dialog dataset and show improved performance over Softmax and other baselines.
Our contribution is two-fold.
First, we identify the Softmax bottleneck by formulating language modeling as a matrix factorization problem.
Second, we propose a simple and effective method that substantially improves over the current state-of-the-art results.
L ANGUAGE M ODELING AS M ATRIX F ACTORIZATION
As discussed in Section 1, with the autoregressive factorization, language modeling can be reduced to modeling the conditional distribution of the next token x
given the context c .
Though one might argue that a natural language allows an inﬁnite number of contexts due to its compositionality (Pinker, 1994), we proceed with our analysis by considering a ﬁnite set of possible contexts.
The unboundedness of natural language does not affect our conclusions, which will be discussed later.
We consider a natural language as a ﬁnite set of pairs of a context and its conditional next-token distribution 2 L = { ( c 1 , P ∗ ( X | c 1 )) , · · · , ( c N , P ∗ ( X | c N )) } , where N is the number of possible contexts.
We assume P ∗ > 0 everywhere to account for errors and ﬂexibility in natural language.
Let { x 1 , x 2 , · · · , x M } denote a set of M possible tokens in the language L .
The objective of a language model is to learn a model distribution P θ ( X | C ) parameterized by θ to match the true data distribution P ∗ ( X | C ) .
In this work, we study the expressiveness of the parametric model class P θ ( X | C ) .
In other words, we are asking the following question: given a natural language L , does there exist a parameter θ such that P θ ( X | c ) =
P ∗ ( X | c ) for all c in L ?
We start by looking at a Softmax-based model class since it is widely used.
2. S OFTMAX
The majority of parametric language models use a Softmax function operating on a context vector (or hidden state)
h c and a word embedding w x to deﬁne the conditional distribution P θ ( x | c ) .
More speciﬁcally, the model distribution is usually written as
[Equation]
where h c is a function of c , and w x is a function of x .
Both functions are parameterized by θ .
Both the context vector h c and the word embedding w x have the same dimension d .
The dot product h (cid:62) c w x is called a logit .
To help discuss the expressiveness of Softmax, we deﬁne three matrices:
[Equation]
  where H θ ∈ R N × d , W θ ∈ R M × d , A ∈ R N × M , and the rows of H θ , W θ , and A correspond to context vectors, word embeddings, and log probabilities of the true data distribution respectively.
We use the subscript θ because ( H θ , W θ ) is effectively a function indexed by the parameter θ , from the joint function family U .
Concretely, H θ is implemented as deep neural networks, such as a recurrent network, while W θ is instantiated as an embedding lookup.
We further specify a set of matrices formed by applying row-wise shift to A
F ( A ) =
{ A + ΛJ N,M | Λ is diagonal and Λ ∈ R N × N } , where J N,M is an all-ones matrix with size N × M .
Essentially, the row-wise shift operation adds an arbitrary real number to each row of A .
Thus, F ( A ) is an inﬁnite set.
Notably, the set F ( A ) has two important properties (see Appendix A for the proof), which are key to our analysis.
Property 1.
For any matrix A (cid:48) , A (cid:48) ∈ F ( A ) if and only if Softmax ( A (cid:48) )
=
P ∗ .
In other words,
F ( A ) deﬁnes the set of all possible logits that correspond to the true data distribution.
Property 2.
For any A 1 (cid:54) =
A 2 ∈ F ( A ) , | rank ( A 1 ) − rank (
A 2 ) | ≤ 1 .
In other words, all matrices in F ( A ) have similar ranks, with the maximum rank difference being 1.
Based on the Property 1 of F ( A ) , we immediately have the following Lemma.
Lemma 1.
Given a model parameter θ , H θ W (cid:62) θ ∈ F ( A ) if and only if P θ ( X | c ) =
P ∗ ( X | c ) for all c in L .
Now the expressiveness question becomes: does there exist a parameter θ and A (cid:48) ∈ F ( A ) such that
[Equation]
This is essentially a matrix factorization problem.
We want the model to learn matrices H θ and W θ that are able to factorize some matrix A (cid:48)
∈ F ( A ) .
First, note that for a valid factorization to exist, the rank of H θ W (cid:62) θ has to be at least as large as the rank of A (cid:48) .
Further, since H θ ∈ R N × d and
W θ ∈ R M × d , the rank of H θ W (cid:62) θ is strictly upper bounded by the embedding size d .
As a result, if d ≥ rank ( A (cid:48) ) , a universal approximator can theoretically recover A (cid:48) .
However, if d < rank ( A (cid:48) ) , no matter how expressive the function family U is, no ( H θ , W θ ) can even theoretically recover A (cid:48) .
We summarize the reasoning above as follows (see Appendix A for the proof).
Proposition 1.
Given that the function family U is a universal approximator, there exists a parameter θ such that P θ ( X | c ) =
P ∗ ( X | c ) for all c in L if and only if d ≥ min
A (cid:48) ∈ F ( A ) rank ( A (cid:48) ) .
Combining Proposition 1 with the Property 2 of F ( A ) , we are now able to state the Softmax Bottleneck problem formally.
Corollary 1.
(Softmax Bottleneck)
If d < rank ( A ) − 1 , for any function family U and any model parameter θ , there exists a context c in L such that P θ ( X
| c ) (cid:54)
=
P ∗ ( X | c ) .
The above corollary indicates that when the dimension d is too small, Softmax does not have the capacity to express the true data distribution.
Clearly, this conclusion is not restricted to a ﬁnite language L .
When L is inﬁnite, one can always take a ﬁnite subset and the Softmax bottleneck still exists.
Next, we discuss why the Softmax bottleneck is an issue by presenting our hypothesis that
A is high-rank for natural language.
2. H YPOTHESIS : N ATURAL L ANGUAGE IS H IGH -R ANK
We hypothesize that for a natural language L , the log probability matrix A is a high-rank matrix.
It is difﬁcult (if possible) to rigorously prove this hypothesis since we do not have access to the true data distribution of a natural language.
However, it is suggested by the following intuitive reasoning and empirical observations:
• Natural language is highly context-dependent (Mikolov & Zweig, 2012).
For example, the token “north” is likely to be followed by “korea” or “korean” in a news article on international politics, which however is unlikely in a textbook on U.S. domestic history.
We hypothesize that such subtle context dependency should result in a high-rank matrix A .
• If A is low-rank, it means humans only need a limited number (e.g. a few hundred) of bases, and all semantic meanings can be created by (potentially) negating and (weighted) averaging these bases.
However, it is hard to ﬁnd a natural concept in linguistics and cognitive science that corresponds to such bases, which questions the existence of such bases.
For example, semantic meanings might not be those bases since a few hundred meanings may not be enough to cover everyday meanings, not to mention niche meanings in specialized domains.
• Empirically, our high-rank language model outperforms conventional low-rank language models on several benchmarks, as shown in Section 3.
We also provide evidences in Section 3.3 to support our hypothesis that learning a high-rank language model is important.
Given the hypothesis that natural language is high-rank, it is clear that the Softmax bottleneck limits the expressiveness of the models.
In practice, the embedding dimension d is usually set at the scale of 10 2 , while the rank of A can possibly be as high as M (at the scale of 10 5 ), which is orders of magnitude larger than d .
Softmax is effectively learning a low-rank approximation to A , and our experiments suggest that such approximation loses the ability to model context dependency, both qualitatively and quantitatively (Cf. Section 3).
2. E ASY F IXES ?
Identifying the Softmax bottleneck immediately suggests some possible “easy ﬁxes”.
First, as considered by a lot of prior work, one can employ a non-parametric model, namely an Ngram model (Kneser & Ney, 1995).
Ngram models are not constrained by any parametric forms so it can universally approximate any natural language, given enough parameters.
Second, it is possible to increase the dimension d (e.g., to match M ) so that the model can express a high-rank matrix A .
However, these two methods increase the number of parameters dramatically, compared to using a low-dimensional Softmax.
More speciﬁcally, an Ngram needs ( N × M ) parameters in order to express A , where N is potentially unbounded.
Similarly, a high-dimensional Softmax requires ( M ×
M ) parameters for the word embeddings.
Increasing the number of model parameters easily leads to overﬁtting.
In past work, Kneser & Ney (1995) used back-off to alleviate overﬁtting.
Moreover, as deep learning models were tuned by extensive hyper-parameter search, increasing the dimension d beyond several hundred is not helpful
(Merity et al., 2017; Melis et al., 2017; Krause et al., 2017).
Clearly there is a tradeoff between expressiveness and generalization on language modeling.
Naively increasing the expressiveness hurts generalization.
Below, we introduce an alternative approach that increases the expressiveness without exploding the parametric space.
2. M IXTURE OF S OFTMAXES : A H IGH -R ANK L ANGUAGE M ODEL
We propose a high-rank language model called Mixture of Softmaxes (MoS) to alleviate the Softmax bottleneck issue.
MoS formulates the conditional distribution as
[Equation]
where π c,k is the prior or mixture weight of the k -th component, and h c,k is the k -th context vector associated with context c .
In other words, MoS computes K Softmax distributions and uses a weighted average of them as the next-token probability distribution.
Similar to prior work on recurrent language modeling (Merity et al., 2017; Melis et al., 2017; Krause et al., 2017), we ﬁrst apply a stack of recurrent layers on top of X to obtain a sequence of hidden states ( g 1 , · · · , g T ) .
The prior and the context vector for context c t are parameterized as π c t ,k =
[Equation]
h c t ,k = tanh ( W h,k g t ) where w π,k and W h,k are model parameters.
Our method is simple and easy to implement, and has the following advantages:
• Improved expressiveness (compared to Softmax).
MoS is theoretically more (or at least equally) expressive compared to Softmax given the same dimension d .
This can be seen by the fact that MoS with K = 1 is reduced to Softmax.
More importantly, MoS effectively approximates A by
[Equation]
where Π k is an ( N × N ) diagonal matrix with elements being
the prior π c,k .
Because ˆ A MoS is a nonlinear function ( log_sum_exp ) of the context vectors and the word embeddings, ˆ A MoS can be arbitrarily high-rank.
As a result, MoS does not suffer from the rank limitation, compared to
Softmax.
• Improved generalization (compared to Ngram).
Ngram
models
and high-dimensional Softmax
(Cf. Section 2.3) improve the expressiveness but do not generalize well.
In contrast, MoS does not have a generalization issue due to the following reasons.
First, MoS deﬁnes the following generative process: a discrete latent variable k is ﬁrst sampled from { 1 , · · · , K } , and then the next token is sampled based on the k -th Softmax component.
By doing so we introduce an inductive bias that the next token is generated based on a latent discrete decision (e.g., a topic), which is often safe in language modeling (Blei et al., 2003).
Second, since ˆ A MoS is deﬁned by a nonlinear function and not restricted by the rank bottleneck, in practice it is possible to reduce d to compensate for the increase of model parameters brought by the mixture structure.
As a result, MoS has a similar model size compared to Softmax and thus is not prone to overﬁtting.
2. M IXTURE OF C ONTEXTS : A L OW -R ANK B ASELINE
Another possible approach is to directly mix the context vectors (or logits) before taking the Softmax, rather than mixing the probabilities afterwards as in MoS. Speciﬁcally, the conditional distribution is parameterized as
[Equation]
(cid:48) (cid:17) , (2) where h c,k and π c,k share the same parameterization as in MoS. Despite its superﬁcial similarity to
MoS, this model, which we refer to as mixture of contexts (MoC), actually suffers from the same rank limitation problem as Softmax.
This can be easily seen by deﬁning h (cid:48)
c
= (cid:80)
Kk =1 π c,k h c,k , which turns the MoC parameterization (2) into P θ ( x | c ) = exp h
(cid:48)(cid:62) c w x (cid:80) x (cid:48)
exp h (cid:48)(cid:62)
c
w x (cid:48) .
Note that this is equivalent to the Softmax parameterization (1).
Thus, performing mixture in the feature space can only make the function family U more expressive, but does not change the fact that the rank of H θ W (cid:62) θ is upper bounded by the embedding dimension d .
In our experiments, we implement MoC as a baseline and compare it experimentally to MoS.
E XPERIMENTS 3. M AIN R ESULTS
We conduct a series of experiments with the following settings:
• Following previous work (Krause et al., 2017; Merity et al., 2017; Melis et al., 2017), we evaluate the proposed MoS model on two widely used language modeling datasets, namely Penn Treebank (PTB) (Mikolov et al., 2010) and WikiText-2 (WT2) (Merity et al., 2016) based on perplexity.
For fair comparison, we closely follow the regularization and optimization techniques introduced by Merity et al.
We heuristically and manually search hyper-parameters for MoS based on the validation performance while limiting the model size (see Appendix B.1 for our hyper-parameters).
• To investigate whether the effectiveness of MoS can be extended to even larger datasets, we conduct an additional language modeling experiment on the 1B Word dataset (Chelba et al., 2013).
Speciﬁcally, we lower-case the text and choose the top 100K tokens as the vocabulary.
A standard neural language model with 2 layers of LSTMs followed by a Softmax output layer is used as the baseline.
Again, the network size of MoS is adjusted to ensure a comparable number of parameters.
Notably, dropout was not used, since we found it not helpful to either model (see Appendix B.2 for more details).
• To show that the MoS is a generic structure that can be used to model other context-dependent distributions, we additionally conduct experiments in the dialog domain.
We use the Switchboard dataset (Godfrey & Holliman, 1997) preprocessed by Zhao et al.
(2017) 4 to train a Seq2Seq (Sutskever et al., 2014) model with MoS added to the decoder RNN.
Then, a Seq2Seq model using Softmax and another one augmented by MoC with comparable parameter sizes
[Table] [Table]
[Table] [Table]
[Table]
are used as baselines.
For evaluation, we include both the perplexity and the precision/recall of Smoothed Sentence-level BLEU, as suggested by Zhao et al.
When generating responses, we use beam search with beam size 10, restrict the maximum length to
[Table]
and retain the top-5 responses
[Table]
The language modeling results on PTB and WT2 are presented in Table 1 and Table 2 respectively.
With a comparable number of parameters, MoS outperforms all baselines with or without dynamic evaluation, and substantially improves over the current state of the art, by up to 3.6 points in perplexity.
[Table]
The improvement on the large-scale dataset is even more signiﬁcant.
As shown in Table 3, MoS outperforms Softmax by over 5.6 points in perplexity.
It suggests the effectiveness of MoS is not limited to small datasets where many regularization techniques are used.
Note that with limited computational resources, we didn’t tune the hyper-parameters for MoS.
[Table]
Further, the experimental results on Switchboard are summarized in Table 4 5 .
Clearly, on all metrics,
MoS outperforms MoC and Softmax, showing its general effectiveness.
3. A BLATION S TUDY
To further verify the improvement shown above does come from the MoS structure rather than adding another hidden layer or ﬁnding a particular set of hyper-parameters, we conduct an ablation study on both PTB and WT2.
Firstly, we compare MoS with an MoC architecture with the same number of layers, hidden sizes, and embedding sizes, which thus has the same number of parameters.
In addition, we adopt the hyper-parameters used to obtain the best MoS model (denoted as MoS hyper-parameters), and train a baseline AWD-LSTM.
To avoid distractive factors and save computational resources, all ablative experiments excluded the use of ﬁnetuing and dynamic evaluation.
The results are shown in Table 5.
Compared to the vanilla AWD-LSTM, though being more expressive, MoC performs only better on PTB, but worse on WT2.
It suggests that simply adding another hidden layer or employing a mixture structure in the feature space does not guarantee a better performance.
On the other hand, training AWD-LSTM using MoS hyper-parameters severely hurts the performance, which rules out hyper-parameters as the main source of improvement.
[Table]
3. V ERIFY THE R OLE OF R ANK
While the study above veriﬁes that MoS is the key to achieving the state-of-the-art performance, it is still not clear whether the superiority of MoS comes from its potential high rank, as suggested by our theoretical analysis in Section 2.
In the sequel, we take steps to verify this hypothesis.
• Firstly, we verify that MoS does induce a high-rank log-probability matrix empirically, while MoC and Softmax fail.
On the validation or test set of PTB with tokens X = { X 1 , . . .
, X T } , we compute the log probabilities
log P
X i | X <i ) ∈
R
M } Tt =1 for each token
using
all three
models
Then, for each model, we
stack all T log probability vectors
into a
T × M matrix resulting in
ˆ
A MoS
, ˆ A MoC and ˆ A
Softmax
Theoretically, the number of
non
-zero
singular values
of
a matrix is equal
to its rank.
However, performing
singular value
decomposition
of real valued matrices using numerical approaches
often encounter roundoff errors
Hence
, we
adopt the expected roundoff error suggested by Press (2007) when estimating the ranks of ˆ A MoS , ˆ A MoC and ˆ A Softmax .
The
estimated ranks are shown in Table 6.
As predicted by our theoretical analysis, the matrix ranks induced by Softmax and MoC are both limited by the corresponding embedding sizes.
By contrast, the matrix rank obtained from MoS does not suffer from this constraint, almost reaching full rank ( M = 10000 ).
In appendix C.1, we give additional evidences for the higher rank of
MoS.
[Table] [Table]
• Secondly, we show that, before reaching full rank, increasing the number of mixture components in MoS also increases the rank of the log-probability matrix, which in turn leads to improved performance (lower perplexity).
Speciﬁcally, on PTB, with other hyper-parameters ﬁxed as used in section 3.1, we vary the number of mixtures used in MoS and compare the corresponding empirical rank and test perplexity without ﬁnetuning.
Table summarizes the results.
This clear
positive correlation between rank and performance strongly supports the our theoretical
analysis in section Moreover, note that after reaching almost full rank (i.e., using 15 mixture components), further increasing the number of components degrades the performance due to overﬁtting (as we inspected the training and test perplexities).
• In addition, as performance improvement can often come from better regularization, we investigate whether MoS has a better, though unexpected, regularization effect compared to Softmax.
We consider the 1B word dataset where overﬁtting is unlikely and no explicit regularization technique (e.g., dropout) is employed.
As we can see from the left part of Table 3, MoS and Softmax achieve a similar generalization gap, i.e., the performance gap between the test set and the training set.
It suggests both models have similar regularization effects.
Meanwhile, MoS has a lower training perplexity compared to Softmax, indicating that the improvement of MoS results from improved expressiveness.
• The last evidence we provide is based on an inverse experiment.
Empirically, we ﬁnd that when Softmax does not suffer from a rank limitation, e.g., in character-level language modeling, using MoS will not improve the performance.
Due to lack of space, we refer readers to Appendix C. for details.
3. A DDITIONAL ANALYSIS
MoS computational time The expressiveness of MoS does come with a computational cost— computing a K -times larger Softmax.
To give readers a concrete idea of the inﬂuence on training time, we perform detailed analysis in Appendix C.3.
As we will see, computational wall time of
MoS is actually sub-linear w.r.t.
the number of Softmaxes K .
In most settings, we observe a two to three times slowdown when using MoS with up to 15 mixture components.
Qualitative analysis Finally, we conduct a case study on PTB to see how MoS improves the next-token prediction in detail.
Due to lack of space, we refer readers to Appendix C.4 for details.
The key insight from the case study is that MoS is better at making context-dependent predictions.
Speciﬁcally, given the same immediate preceding word, MoS will produce distinct next-step prediction based on long-term context in history.
By contrast, the baseline often yields similar next-step prediction, independent of the long-term context.
R ELATED WORK
In language modeling, Hutchinson et al.
(2011; 2012) have previously considered the problem from a matrix rank perspective.
However, their focus was to improve the generalization of Ngram language models via a sparse plus low-rank approximation.
By contrast, as neural language models already generalize well, we focus on a high-rank neural language model that improves expressiveness without sacriﬁcing generalization.
Neubig & Dyer (2016) proposed to mix Ngram and neural language models to unify and beneﬁt from both.
However, this mixture might not generalize well since an Ngram model, which has poor generalization, is included.
Moreover, the fact that the two components are separately trained can limit its expressiveness.
Levy & Goldberg (2014) also considered the matrix factorization perspective, but in the context of learning word embeddings.
In a general sense, Mixture of Softmaxes proposed in this work can be seen as a particular instantiation of the long-existing idea called Mixture of Experts (MoE) (Jacobs et al., 1991).
However, there are two core differences.
Firstly, MoE has usually been instantiated as mixture of Gaussians to model data in continuous domains (Jacobs et al., 1991; Graves, 2013; Bazzani et al., 2016).
More importantly, the motivation of using the mixture structure is distinct.
For Gaussian mixture models, the mixture structure is employed to allow for a parameterized multi-modal distribution.
By contrast, Softmax by itself can parameterize a multi-modal distribution, and MoS is introduced to break the Softmax bottleneck as discussed in Section 2.
There has been previous work (Eigen et al., 2013; Shazeer et al., 2017) proposing architectures that can be categorized as instantiations of MoC, since the mixture structure is employed in the feature space.
6 The target of Eigen et al.
(2013) is to create a more expressive feed-forward layer through the mixture structure.
In comparison, Shazeer et al.
(2017) focuses on a sparse gating mechanism also on the feature level, which enables efﬁcient conditional computation and allows the training of a very large neural architecture.
In addition to having different motivations from our work, all these
MoC variants suffer from the same rank limitation problem as discussed in Section 2.
Finally, several previous works have tried to introduce latent variables into sequence modeling (Bayer & Osendorfer, 2014; Gregor et al., 2015; Chung et al., 2015; Gan et al., 2015; Fraccaro et al., 2016; Chung et al., 2016).
Except for (Chung et al., 2016), these structures all deﬁne a continuous latent variable for each step of the RNN computation, and rely on the SGVB estimator (Kingma & Welling, 2013) to optimize a variational lower bound of the log-likelihood.
Since exact integration is infeasible, these models cannot estimate the likelihood (perplexity) exactly at test time.
Moreover, for discrete data, the variational lower bound is usually too loose to yield a competitive approximation compared to standard auto-regressive models.
As an exception, Chung et al.
(2016) utilizes Bernoulli latent variables to model the hierarchical structure in language,
where the
Bernoulli sampling is replaced by a thresholding operation at test time to give perplexity estimation.
C ONCLUSIONS
Under the matrix factorization framework, the expressiveness of Softmax-based language models is limited by the dimension of the word embeddings, which is termed as the Softmax bottleneck.
Our proposed MoS model improves the expressiveness over Softmax, and at the same time avoids overﬁtting compared to non-parametric models and naively increasing the word embedding dimensions.
Our method improves the current state-of-the-art results on standard benchmarks by a large margin, which in turn justiﬁes our theoretical reasoning: it is important to have a high-rank model for natural language.
A CKNOWLEDGMENTS
This work was supported by the DARPA award D17AP00001, the Google focused award, and the
Nvidia NVAIL award.
A P ROOFS Proof of Property
Proof.
For any A (cid:48) ∈ F ( A ) , let P A (cid:48)
( X | C ) denote the distribution deﬁned by applying Softmax on the logits given by A (cid:48) .
Consider row i column j , by deﬁnition any entry in A (cid:48) can be expressed as
A (cid:48) ij =
A ij + Λ ii .
It follows
[Equation]
For any A (cid:48)(cid:48) ∈ { A (cid:48)(cid:48) | Softmax ( A (cid:48)(cid:48) )
=
P ∗ } , for any i and j , we have
[Equation]
It follows that for any i , j , and k ,
[Equation]
c i
[Equation]
As a result,
[Equation]
This means each row in A (cid:48)(cid:48) can be obtained by adding a real number to the corresponding row in
A .
Therefore, there exists a diagonal matrix Λ ∈ R N × N such that
[Equation]
It follows that A (cid:48)(cid:48) ∈ F ( A ) .
Proof of Property
Proof.
For any A 1 and A 2 in F ( A ) , by deﬁnition we have A 1 =
A + Λ 1 J N,M , and A 2 =
A + Λ 2 J N,M where Λ 1 and Λ 2 are two diagonal matrices.
It can be rewritten as
[Equation]
Let S be a maximum set of linearly independent rows in A 2 .
Let e N be an all-ones vector with dimension N .
The i -th row vector a 1 ,i in A 1 can be written as
[Equation]
Because a 2 ,i is a linear combination of vectors in S , a 1 ,i is a linear combination of vectors in
S ∪ { e N } .
It follows that
[Equation]
Similarly, we can derive
[Equation]
Therefore,
[Equation]
Proof of Proposition
Proof.
If there exists a parameter θ such that P θ ( X | c ) =
P ∗ ( X | c ) for all c in L , by Lemma 1, we have H θ W (cid:62) θ ∈ F ( A ) .
As a result, there exists a matrix A (cid:48) ∈ F ( A ) such that H θ W (cid:62) θ
= A (cid:48) .
Because H θ and W θ are of dimensions ( N × d ) and ( M × d ) respectively, we have
[Equation]
If d ≥ min A (cid:48)(cid:48) ∈
F ( A ) rank ( A (cid:48)(cid:48) ) , there exist matrices A (cid:48) ∈ F ( A ) , H (cid:48)
∈ R N × d and W (cid:
) ∈ R M × d , such that A (cid:48) can be factorized as A (cid:48)
=
H (cid:48) W (cid:48)(cid:62) .
Because U is a universal approximator, there exists θ such that H θ
= H (cid:48) and W θ = W (cid:48) .
By Lemma 1, P θ ( X | c ) =
P ∗ ( X | c ) for all c in L .
B E XPERIMENT SETTING AND H YPER - PARAMETERS B. PTB AND WT
The hyper-parameters used for MoS in language modeling experiment is summarized below.
[Figure]
The hyper-parameters used for dynamic evaluation of MoS is summarized below.
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
B. 1B W ORD D ATASET
For training, we use all of the 100 training shards.
For validation, we use two shards from the heldout set, namely [heldout-00, heldout-10] .
For test, we use another three shards from the heldout set, namely [heldout-20, heldout-30, heldout-40] .
The hyper-parameters are listed below.
[Figure]
C A DDITIONAL EXPERIMENTS C. H IGHER EMPIRICAL RANK OF M O S COMPARED TO M O C AND S OFTMAX
In section 3, we compute the rank of different models based on the non-zero singular values of the empirical log-likelihood matrix.
Since there can be roundoff mistakes, a less error-prone approach is to directly study the distribution of singular values.
Speciﬁcally, if more singular values have relatively larger magnitude, the rank of the matrix tends to be higher.
Motivated from this intuition,
[Figure]
we visualize the distribution of the singular values.
To account for the different magnitudes of singular values from different models, we ﬁrst normalize all singular values to [0 , 1] .
Then, we plot the cumulative percentage of normalized singular values, i.e., percentage of normalized singular values below a threshold, in Figure 1.
As we can see, most of the singular values of Softmax and
MoC concentrate on an area with very low values.
In comparison, the concentration area of the MoS singular values is not only several orders larger, but also spans a much wider region.
Intuitively, MoS utilizes the corresponding singular vectors to capture a larger and more diverse set of contexts.
[Figure]
[Table]
[Figure]
What’s more, another indicator of high rank is that the model can precisely capture the nuance of difference contexts.
If a model can better capture the distinctions among contexts, we expect the nextstep conditional distributions to be less similar to each on average.
Based on this intuition, we use the expected pairwise Kullback–Leibler divergence (KLD), i.e., E c,c (cid:48) ∼C
[KLD( P ( X | c ) (cid:107) P ( X | c (cid:48) ))] where C denotes all possible contexts, as another metric to evaluate the ranks of the three models (MoS, MoC and Softmax).
Practically, we sample c, c (cid:48) from validation or test data of PTB to get the empirical estimations for the three models, which are shown in the right half of Table 11.
As we expected, MoS achieves higher expected pairwise KLD, indicating its superiority in covering more contexts of the next-token distribution.
C. A N INVERSE EXPERIMENT ON CHARACTER - LEVEL LANGUAGE MODELING
[Table]
Here, we detail the inverse experiment, which shows that when Softmax does not suffer from a rank limitation, using MoS will not improve the performance.
Notice that character-level language modeling (CharLM) is exactly such a problem, because the rank of the log-likelihood matrix is upper bounded by the vocabulary size, and CharLM usually has a very limited vocabulary (tens of characters).
In this case, with the embedding size being hundreds in practice, Softmax is no longer a bottleneck in this task.
Hence, we expect MoS to yield similar performance to Softmax on CharLM.
We conduct experiments of CharLM using the text8 dataset (Mahoney, 2011), which consists of 100M characters including only alphabetical characters and spaces derived from Wikipedia.
We follow Mikolov et al.
and use the ﬁrst 90M characters for training, the next 5M for validation and the ﬁnal 5M for testing.
The standard evaluation metric bit-per-character (BPC) is employed.
We employ a 1-layer 1024-unit LSTM followed by Softmax as the baseline.
For MoS, we consider 7 or 10 mixtures and reduce the hidden and/or embedding size to match the baseline capacity.
When decreasing the hidden and/or embedding size, we either keep both the same, or make the hidden size relatively larger.
The results are summarized in Table 12.
Clearly, the Softmax and MoS obtain the same BPC on the test set and comparable BPC on the validation set, which well match our hypothesis.
Since the only difference in word-level language modeling is the existence of the
Softmax bottleneck, the distinct behavior of MoS again supports our hypothesis that it is solving the
Softmax bottleneck problem.
[Table]
We evaluate the additional computational cost introduced by MoS. We consider two sets of controlled experiments.
In the ﬁrst set, we compare the training time of MoS and Softmax using the same batch sizes.
In the second set, we compare the training time of two methods using the hyperparameter settings that achieve the best performance for each model (i.e., the settings in Tables 1, 2, and 3).
In both sets, we control two models to have comparable model sizes.
The results on the three datasets are shown in Table 13.
Thanks to the efﬁciency of matrix multiplication on GPU, the computational wall time of MoS is actually sub-linear w.r.t.
the number of
Softmaxes K .
In most settings, we observe a two to three times slowdown when using MoS. Specifically, the “bs” setting measures the computational cost introduced by MoS given enough memory, which is 1.9x, 2.5x, and 3.8x slowdown on PTB, WT2, and 1B respectively.
The “best-1” setting is usually slower compared to “bs”, because a single batch does not ﬁt into the memory of a single
GPU using MoS, in which case we have to split one batch into multiple small ones, resulting in further slowdown.
In this sense, the gap between “best-1” and “bs” measures the computational cost introduced due to the increase of memory consumed by MoS.
The “best-3” alleviates this issue by using three GPUs, which allows larger-batch training for MoS.
In this case, we reduce the computational cost to 2.9x on WT2 and 2.1x on 1B with our best performing model.
Note that the computational cost is closely related to the batch size, which is interleaved with optimization.
Though how batch sizes affect optimization remains an open question and might be task dependent, we believe the “best-1” and “best-3” settings well reﬂect the actual computational cost brought by MoS on language modeling tasks.
C. Q UALITATIVE A NALYSIS
Since MoC shows a stronger performance than Softmax on PTB, the qualitative study focuses on the comparison between MoC and MoS. Concretely, given the same context (previous tokens), we search for prediction steps where MoS achieves lower negative log loss than MoC by a margin.
We show some representative cases in Table 14 with the following observations:
• Comparing the ﬁrst two cases, given the same preceding word “N”, MoS ﬂexibly adjusts its top predictions based on the different topic quantities being discussed in the context.
In comparison, MoC emits quite similar top choices regardless of the context, suggesting its inferiority in make context-dependent predictions.
• In the 3rd case, the context is about international politics, where country/region names are likely to appear.
MoS captures this nuance well, and yields top choices that can be used to complete a country name given the immediate preceding word “south”.
Similarly, in the 4th case, MoS is able to include “ual”, a core entity of discussion in the context, in its top predictions.
In contrast, MoC gives rather generic predictions irrieselevant to the context in both cases.
• For the 5th and the 6th example, we see MoS is able to exploit less common words accurately according to the context, while MoC fails to yield such choices.
This well matches our analysis that MoS has the capacity of modeling context-dependent language.
[Table]
managed
[Table]
investment
[Table]
quality prop
[Table]
as other <unk> governments particularly poland and the soviet union have recently discovered initial steps to open up society can create a momentum for radical change that becomes difﬁcult
[Table]
to
[Table]
unk
[Table]
as the days go by the south african government will be ever more hard pressed to justify the continued <
[Table]
the government is watching closely to see if their presence in the <unk
[Table]
leads to increased
[Table]
unk
[Table]
protests and violence if it does pretoria will use this as a reason to keep mr
[Table]
<
unk> behind bars
[Table]
ﬁrst
[Table]
unk
[Table]
ﬁrst
[Table]
china ’s <unk> <unk> program has achieved some successes in <unk> runaway economic growth and stabilizing prices but has failed to eliminate serious defects in state planning and an <unk> drain on state budgets <eos> the ofﬁcial china daily said retail prices of <unk> foods have n’t risen since last december but acknowledged that huge government __?__
[Table]