Softmax Bottleneck Makes Language Models Unable to Represent Multi-mode Word Distributions (2022)

Abstract: Neural language models (LMs) such as GPT-2 estimate the probability distribution over the next word by a softmax over the vocabulary. The softmax layer produces the distribution based on the dot products of a single hidden state and the embeddings of words in the vocabulary. However, we discover that this single hidden state cannot produce all probability distributions regardless of the LM size or training data size because the single hidden state embedding cannot be close to the embeddings of all the possible next words simultaneously when there are other interfering word embeddings between them. In this work, we demonstrate the importance of this limitation both theoretically and practically. Our work not only deepens our understanding of softmax bottleneck and mixture of softmax (MoS) but also inspires us to propose multi-facet softmax (MFS) to address the limitations of MoS. Extensive empirical analyses confirm our findings and show that against MoS, the proposed MFS achieves two-fold improvements in the perplexity of GPT-2 and BERT.

## introduction recently, researchers have found that transformerbased language models (lms), such as gpt-2, can predict the next/masked word distribution better as their sizes grow @xcite @xcite . compared to greedily outputting the most probable next word, sampling the next word from the predicted distribution allows a lm to generate more diverse and high-quality text sequences @xcite . by autoregressively sampling the next word according to its predicted probability, large lms can be used to assist creative writing @xcite , reduce the cost of building datasets @xcite , generate codes @xcite , solve math problems @xcite , etc. as a result, one natural question arises: do modern language modeling architectures still have restrictions in their ability to represent the appropriate distribution over next words or masked words? in this paper, we discover that, when predicting the next word probabilities given an ambiguous context, gpt-2 is often incapable of assigning the highest probabilities to the appropriate nonsynonym candidates. for example, given the input prompt "after debating whether to bow to the woman or the king first, the jester decided on the [mask]", we would expect the distribution over the @xcite fillers to put high probabilities on woman or king or their synonyms. however, gpt-2 might incorrectly assign the second-highest probability to "queen" as in figure 1 . in the final softmax layer of gpt-2, the log probabilities of the woman and king are computed based on the dot product between a single hidden state embedding and the global word embeddings of the woman and king, respectively. to have the highest but similar dot products for the two options, the transformer encoder in gpt-2 wants to output the hidden state that is close to the average of the woman embedding and the king embedding. however, the words queen, king, woman, and man tend to form a parallelogram in the embedding space @xcite @xcite foot_0 , which means the man and queen also have a similar average. therefore, gpt-2 is forced to also output man or queen when it wants to output woman or king. the problem not only happens to gpt-2 or the words whose embeddings form a parallelogram shape. even though the hidden state embeddings of lms are contextualized, the embedding of each output word embedding space ## multi-facet softmax using multiple embeddings is a natural solution for modeling a multimodal distribution @xcite . for instance, we can use three embeddings to capture the high probability on the woman and king but low probability on the man and queen in figure 1 . inspired by our geometric analysis on the limitation of the single embedding, we improve the state-of-the-art multiple embedding solution, mixture of softmax (mos) @xcite by two enhancements: multiple input hidden states and multiple partitions on the vocabulary. @xcite propose mixture of softmax (mos) to allow a lstm-based @xcite lm to produce more linearly independent log probability distributions of the output words given different contexts. as in figure 2 (c), the mos first uses multiple linear layers l f k to project a hidden state h ct into multiple facet embeddings f ct,@xmath0. 3 the multiple facets f ct,k and softmaxes would lead to multiple probability distributions, and output probability is the weighted average of the distributions: ## language modeling experiments we evaluate different lm architectures by comparing their capability of predicting the next word in wikipedia 2021 and a subset of openweb-text @xcite . in addition to perplexity, we also compare their mean reciprocal ranks (mrr) in appendix c.1. the size of the training, validation, and testing set are 96%, 2%, and 2% of the whole corpus, respectively. after loading the pre-trained gpt-2 models, we train the gpt-2 small for 1 epoch and gpt-2 medium for 0.4 epochs. we also test our methods on bert in appendix b.2. please see appendix g for more details of our experiment setup. ## evaluation on ambiguous templates we synthesize a dataset using templates @xcite to verify whether the softmax layer in the original gpt-2 really has difficulty in learning to output the bimodal distribution in figure 1 and whether the multiple embedding methods could overcome the problem. first, we collect the four words with semantic analogy relations in google analogy dataset @xcite . next, we insert two out of the four words into our manually written templates to form the contexts and the templates we used could be found in appendix g.3. for example, given the context "i went to paris and germany before, and i love one of the places more, which is", the gpt-2 learns to predict either paris or germany. the two words can be either the diagonal words (e.g., king and woman) or the edge word (e.g., king and queen) in the parallelogram. finally, we create a dataset with 122k training contexts, 250k validation contexts, and 122k testing contexts, where the word pairs in the testing set are unseen in the training set to see whether the model could learn to output the bimodal distribution in a general way. 5 5 the setting is realistic because any related words could become the next word in some ambiguous contexts and all we load the models pre-trained on openweb-text and continue fine-tuning the models on the last word of each sentence for 10 epochs. we report the testing performances of the best model selected by the validation loss. since the sets of the word pairs in the training and testing set are disjoint, updating the output word embedding would make gpt-2 solve the task by memorizing/overfitting the training set quickly and lead to much worse testing performances. thus, we freeze the output word embedding during the training. we visualize the predictions of the paris-germany example in the last column of quantitatively, the words are related in a certain way @xcite . we cannot expect the training corpora to contain the ambiguous contexts with so many possible next words. ## answering ambiguous questions protoqa (boratko et al., 2020) is a questionanswering dataset built for evaluating the commonsense reasoning ability of language models. each question in protoqa is ambiguous and leads to a distribution of possible answers. for instance, the answer to "name something that people usually do before they leave for work?" is "shower 0.43, breakfast 0.30, ...". the paper discovers that by reformulating the question-answering task as a context (e.g., "one thing people usually do before they leave for work is ..."), gpt-2 could generate the possible answers by sampling the next words from its word prediction distribution. the dataset gives us a chance to directly compare the quality of the distributions generated by different lms in we can see that the multiple softmaxes, input hidden states, and partitions usually improve the quality of prediction distribution, and the proposed mfs, which combines all modifications, achieves the best performances. a limitation of the aforementioned previous work is that they do not tell us which kinds of sentences would be affected by the bottleneck more and whether the order of the top few next words would be affected, which are the main research questions of our work. contrary to the previous belief that a large hidden state dimension would eliminate the softmax bottleneck, our theorems suggest that some words in a low dimensional subspace could still make the single embedding in the softmax layer become a bottleneck of arbitrarily ranking the output words. furthermore, our geometric analyses provide an intuitive explanation about why breaking the bottleneck using multiple embeddings leads to better performances compared to only adding the non-linearity. demeter et al. (2020) also analyze the structural weakness of the softmax layer from a geometric perspective. they discover that the words with high prior frequencies could stop the lms from assigning the high probabilities to rare words, which can be viewed as a special case of our theory (see appendix e). for instance, our work shows that the softmax layer could still prevent the lms from outputting some top words even if all the possible next words have the same prior frequency. our theory is deeply connected to the mathematical work that counts the number of possible rankings of points in an embedding space @xcite . compared to the studies, our work focuses more on analyzing the multimodal distribution in the word embedding space and its implication to language models. an alternative to model the multimodal distribution is to use multiple embeddings to represent each output word @xcite . compared to mos or our approach that use multiple embeddings to represent each hidden state of the context, their method requires many extra parameters to store different senses of each output word. another type of related model @xcite ## conclusion when the ideal distribution in the output word embedding space has multiple modes, gpt-2 cannot learn to correctly rank the words in all the modes as the top next words. this shows that the single embedding in the softmax layer, which is used nearly universally by current lms, constitutes a performance upper bound of predicting the next/masked word. to address the systematic failure caused by these structural weaknesses, we propose multifacet softmax (mfs). in our experiments, we confirm that the mfs significantly outperforms the standard softmax layer and alleviates the softmax bottleneck in the transformer-based lms such as gpt-2 better than mixture of softmax (mos). this work studies a general limitation of lms and proposes solutions. the proposed theory can help us to understand that some types of hallucinations, mistakes, or biases of lms could come from softmax bottleneck and their incapability of modeling the correct distribution. for example, there are 60% of male characters and 40% of female characters in our training corpus. the language generation model might be forced to assign more than 60% probability to male characters as being much more likely to output king than woman in figure 1 . recently, @xcite ; anonymous (2021) show that mos is one of the few architecture modifications of transformer-based lm that can provide consistent improvements in downstream applications. our work provides a fundamental reason why the multiple embedding representation is better, which could inspire more future studies that propose a better multiple-embedding architecture to improve lms (e.g., multi-lingual bert) or downstream applications. as examples, we list several possible future directions in appendix h. finally, a better lm could lead to both positive and negative societal impacts, but they are not the focus of this paper. generally speaking, this paper deepens our understanding of the weaknesses of modern lms and we believe the knowledge can help us to design a better lm that increases the positive impacts and reduces the negative impacts in the future.