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

Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 8048 - 8073 May 22-27, 2022 c⃝2022 Association for Computational Linguistics Softmax Bottleneck Makes Language Models Unable to Represent Multi-mode Word Distributions Haw-Shiuan Chang and Andrew McCallum CICS, University of Massachusetts, Amherst 140 Governors Dr., Amherst, MA, USA 01003 {hschang,mccallum}@cs.umass.edu 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 vo- cabulary. However, we discover that this sin- gle hidden state cannot produce all probabil- ity distributions regardless of the LM size or training data size because the single hidden state embedding cannot be close to the embed- dings of all the possible next words simulta- neously 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 soft- max bottleneck and mixture of softmax (MoS) but also inspires us to propose multi-facet soft- max (MFS) to address the limitations of MoS. Extensive empirical analyses confirm our find- ings and show that against MoS, the proposed MFS achieves two-fold improvements in the perplexity of GPT-2 and BERT. “The greater the ambiguity, the greater the plea- sure.” — Milan Kundera 1 Introduction Recently, researchers have found that transformer- based language models (LMs), such as GPT-2, can predict the next/masked word distribution better as their sizes grow (Radford et al., 2019; Brown et al., 2020; Kaplan et al., 2020). Compared to greedily outputting the most probable next word, sampling the next word from the predicted distri- bution allows a LM to generate more diverse and high-quality text sequences (Holtzman et al., 2020). By autoregressively sampling the next word ac- cording to its predicted probability, large LMs can be used to assist creative writing (Akoury et al., 2020), reduce the cost of building datasets (West et al., 2021; Liu et al., 2022), generate codes (Li et al., 2022), solve math problems (Cobbe et al., 2021), etc. As a result, one natural question arises: Do modern language modeling architectures still have restrictions in their ability to represent the ap- propriate distribution over next words or masked words? In this paper, we discover that, when predict- ing the next word probabilities given an ambigu- ous context, GPT-2 is often incapable of assigning the highest probabilities to the appropriate non- synonym candidates. For example, given the in- put 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 [MASK] 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 prob- abilities of …

Content truncated to 500 words