Transformer Language Models Handle Word Frequency in Prediction Head (2023)

Abstract: Prediction head is a crucial component of Transformer language models. Despite its direct impact on prediction, this component has often been overlooked in analyzing Transformers.In this study, we investigate the inner workings of the prediction head, specifically focusing on bias parameters. Our experiments with BERT and GPT-2 models reveal that the biases in their word prediction heads play a significant role in the models’ ability to reflect word frequency in a corpus, aligning with the logit adjustment method commonly used in long-tailed learning. We also quantify the effect of controlling the biases in practical auto-regressive text generation scenarios;under a particular setting, more diverse text can be generated without compromising text quality.

## introduction transformer language models (tlms) @xcite are now fundamental to natural language processing (nlp) techniques, including text generation. owing to this success, extensive research has been conducted to analyze their inner workings @xcite . in this study, we shed light on the operation of the prediction head, the last block of the tlms. despite its direct impact on tlms' output, its characteristics have been overlooked in previous analyses. our experiments with bert and gpt-2 reveal that a particular bias parameter in the prediction head adjusts the model's output toward word frequency in a corpus. particularly, the bias increases the prediction probability for highfrequency words and vice versa (figure 1 ). we further explore this phenomenon from several perspectives. first, we analyze the geometric characteristics of this phenomenon, which show that word frequency is encoded in a specific direction in the output embedding space. second, we analyze the behavioral impact of controlling their frequency biases on text generation. the results demonstrate that the model's text generation can be made more diverse while maintaining the fluency by adequately decaying the bias parameters, suggesting that models can more or less isolate word frequency knowledge from other text generation ability. third, we discuss the potential connection between our findings and the logit adjustment method that is typically used in the machine learning field to address the class imbalance problem. ## background: prediction head tlms have a stack of transformer layers on top of the embedding layer. these components update hidden token representations (figure 2 ). prediction head, which is our target of analysis, is the last, top-most component in tlms. the prediction head in a tlm computes the prediction probabilities for all vocabulary v based on the hidden state in the last transformer layer. formally, the prediction head receives, for each token, the last transformer layer's hidden state @xmath0. the prediction head computes the probability distribution @xmath1, in the case of gpt-2: both prediction heads contain the bias b ln , and the bert head additionally contains the biases b fc and b last . as the first step in analyzing the prediction head, we focus on these three biases because they can easily be mapped to the output space. drawing on the existing findings about the frequency-related workings of several components in the transformer @xcite , we analyze the model behavior with respect to word frequency. ## experiments first, we show that the bias parameters are related to word frequency. next, we analyze their proper-ties from two perspectives: (i) geometric characteristics and (ii) text generation. model: we used bert (cased) @xcite in two different sizes (base and large) and gpt-2 @xcite in four different sizes (small, medium, large, and xl). data: we used 5,000 sequences from the test set of the gpt-2 pre-trainng corpus, openwebtext corpus (gokaslan and cohen, 2019) 1 . each sequence was fed into bert after some tokens were replaced with [mask] 2 , and fed into gpt-2 as it was. further, word frequencies were calculated from the corpus used for training each of bert and gpt-2. 3 ## related work transformer layers (e.g., attention patterns) have been the major focus of tlm analysis @xcite @xcite . the first embedding layer, especially positional encoding, has also been studied @xcite . this study sheds light on the prediction head, the last block of a tlm, and provides new insights into the working mechanisms of tlms. notably, previous studies have reported that words having a similar frequency are clustered in the embedding spaces of various deep nlp models @xcite @xcite ; our observation agrees with theirs. in addition to this, we newly discovered that a particular bias parameter in the tlm prediction head corresponds to "word frequency direction" in the word embedding space. ## conclusions in this study, we explored the workings of bias parameters in the prediction head of tlms. our experiments with bert and gpt-2 showed that the biases adjust the model's prediction with respect to word frequency. we further explored this phenomenon and provided the following insights: (i) word frequency is encoded in a specific direction (the bias direction) in the output embedding space, (ii) properly controlling the bias's effect can encourage more diverse language generation without compromising quality, and (iii) tlms are implicitly trained to be potentially consistent with the logit adjustment method. in future work, we will analyze larger tlms, e.g., open pre-trained transformers @xcite . further, we will analyze the weight parameters in the prediction head in addition to the bias parameters.