Transformer Language Models Handle Word Frequency in Prediction Head
Year: 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.  https://github.com/gorokoba560/ transformer-lm-word-freq-bias

heading: 1 Introduction
text: Transformer language models (TLMs) (Devlin et al., 2019; Radford et al., 2019) 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 (Rogers et al., 2020; Geva et al., 2021).
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
4523
BERT GPT-2
Attention Mechanism
Layer Normalization
Layer Normalization
Embedding Layer
Fully Connected Layer
Layer Normalization
Projection onto
Adding
! ×
Transformer
Layer
Prediction
Head
Attention Mechanism Layer Normalization
Layer Normalization
Embedding Layer
Layer Normalization
Projection onto
(w/ ) (w/ )
(w/ )
Feed-forward Network
Feed-forward Network
Figure 2: Architecture overview of BERT and GPT-2.
in the case of GPT-2:
p = softmax ( LN(x)W emb ) (1)
LN(x) := x−m(x)
s(x) ⊙ γ + bLN ∈ Rd, (2)
where W emb ∈ Rd×|V| denotes the word embedding matrix, m(x) and s(x) denote the elementwise mean and standard deviation, respectively, and ⊙ denotes the element-wise product. γ and bLN ∈ Rd are learnable parameters.
For BERT, there is an additional fully connected layer (FC). The prediction head computes the probability distribution p for the hidden state x that corresponds to the [MASK] token as follows:
p = softmax ( LN(x′)W emb + blast ) (3)
x′ = GELU ( xW FC + bFC ) ∈ Rd, (4)
where W FC ∈ Rd×d denotes the learnable weight matrix, and bFC ∈ Rd and blast ∈ R|V| denote the learnable bias parameters. GELU (Hendrycks and Gimpel, 2016) is the activation function.
Both prediction heads contain the bias bLN, and the BERT head additionally contains the biases bFC and blast. 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 (Voita et al., 2019; Kobayashi et al., 2020), we analyze the model behavior with respect to word frequency.

heading: 3 Experiments
text: 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) (Devlin et al., 2019) in two different sizes (base and large) and GPT-2 (Radford et al., 2019) 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

heading: 3.1 Impact of biases on prediction distribution
text: We compared the models’ word prediction with and without each bias. Specifically, we once obtained word prediction distributions p̂ ∈ R|V| from a model for each time step across the test data. The average of these distributions are referred to as model’s word prediction distribution henceforth.
Bias adjusts the model’s prediction distribution closer to the corpus frequency distribution: Figure 1 shows changes in the model’s word prediction distribution before and after the bias bLN is removed.4 The removal of bLN increases the probability of the model predicting low-frequency words (right side of the figures) and vice versa, which results in a word prediction distribution that approaches a flat (UNIFORM in the figure). In other words, the bias bLN adjusts the models’ word prediction distribution to be closer to the corpus word frequency distribution (UNIGRAM in the figure). This finding can be generalized across all model sizes (Appendix A).
To quantify the above effect, we calculated the Kullback–Leibler (KL) divergence between the model’s word prediction distribution and the corpus word frequency distribution (UNIGRAM). Note that a higher value indicates that the model’s prediction distribution has more discrepancy with that in
1webtext.test.jsonl published in https://github. com/openai/gpt-2-output-dataset was used.
2Following Devlin et al. (2019), 15% of tokens were replaced with [MASK] 80% of the time.
3BERT was trained on Wikipedia and BooksCorpus (Zhu et al., 2015), and GPT-2 was trained on OpenWebText Corpus. We reproduced them using Datasets (Lhoest et al., 2021).
4We created bins to divide the corpus word frequencies into constant intervals and plotted each bin’s geometric mean and standard deviation for the word prediction probabilities.
the pretraining corpus. Table 1 shows that removing bLN always results in a higher value, which indicates that bLN indeed adjusts the prediction distribution to be closer to the corpus frequency distribution. The biases bFC and blast in BERT also exert a similar effect, but it is weaker than that of bLN; we focus on bLN in the following. We also observe that larger models have less change of frequency biases due to bLN.

heading: 3.2 Geometric observations
text: We observed the geometric properties of the bias bLN and the output embedding space of the TLMs.
Word frequency is encoded in the bias vector’s direction in the output embedding space: The observation that the bias vector shifts predictions according to word frequency suggests that word frequency is encoded in the output embedding space W emb, and the bias vector bLN is a good projection to extract this frequency information. In fact, the inner product of bLN and each word embedding wi in the embedding layer correlates well with the word frequency5 (Figure 3).
Furthermore, we observed that removing the bias direction (≈ frequency direction) from the embedding matrix W emb improved the isotropy (uniformity in direction, e.g., Ethayarajh, 2019) in the output embedding space. Formally, we removed the bias direction using wi ← wi−⟨wi, bLN∥bLN∥⟩ bLN ∥bLN∥ ; then, the average value 1 n2 ∑ i ∑ j cos(wi,wj) decreased from 0.15 to 0.09 in BERT base. This observation shows that the anisotropy in the output space is, more or less, caused by the frequency direction.
We further observed that hidden states htoken before bLN was added were almost orthogonal
5Spearman’s ρ was 0.78 on GPT-2 small.
Corpus frequency of word
In n er
p ro
d u ct
sm al
l. hb
L N , w
ii
wi
Figure 3: Relationship between the corpus word frequency and inner product of bLN and each output word embedding wi in GPT-2 small.
to bLN (≈ word frequency direction); in particular, Etoken |cos (htoken, bLN)| = 0.08 ≪ 1.0 in BERT-base. This corroborates that the frequency bias injected in the prediction head indeed does not exist in the hidden states before the prediction head.
Word frequency encoded on the bias vector is shifted via fine-tuning:
We also inspected whether the model’s word prediction distribution is shifted to that in the target domain during fine-tuning to enhance the generality of our observation. Specifically, we fine-tuned GPT-2 small on a dataset consisting of abstracts from papers in the machine learning field6, whose word frequency distribution is different from the pretraining data. After fine-tuning, the inner product of bLN and each word embedding wi correlated more with the additional fine-tuning corpus after fine-tuning (the Spearman’s ρ changed from 0.38 to 0.62) and slightly less with the pre-training corpus (the Spearman’s ρ changed from 0.78 to 0.73). This suggests that frequency information captured by the bias bLN is updated during fine-tuning.

heading: 3.3 Impact of bias on text generation
text: We next demonstrate that controlling the bias bLN can lead to more diverse text generation without significant harm to the quality of the text. We hope that quantifying the effect of such control using metrics for the evaluation of text generation (e.g., n-gram diversity) will enhance the connection between the language generation field and the field of probing/interpreting LMs’ internals.
6CShorten/ML-ArXiv-Papers published in https:// huggingface.co/datasets/CShorten/ML-ArXiv-Papers on Datasets (Lhoest et al., 2021) was used.
Procedure: We adjusted bLN during text generation by GPT-2, and we then evaluated the generated text. Specifically, we introduced an adjustment coefficient λ ∈ [0, 1] and replaced bLN with λbLN. We report the evaluation scores by varying λ. The results generated with the top-p sampling strategy (Fan et al., 2018) are reported in this section. The results for other decoding settings are in Appendix A; we found similar results for the top-p and top-k sampling but found degradation with the vanilla sampling setting. The details of the settings are described in Appendix B.
Evaluation methods: Text generated by each model was evaluated from two perspectives: diversity and quality. For the diversity evaluation, Distinct-n Dn (Li et al., 2016) and N-gram diversity D (Meister et al., 2022) were used. These measures of n-gram overlap in generated texts were calculated as follows:
Dn(texts) := # Unique n-grams in texts
# n-grams in texts (5)
D(texts) := 1
4
4∑
n=1
Dn(texts). (6)
For the quality evaluation, MAUVE (Pillutla et al., 2021) and Perplexity (PPL) were used. MAUVE evaluates how similar a given text generation model is to humans by comparing human-written texts and model-generated texts according to the difference in their distributions in a sentence embedding space. PPL evaluates how well models can predict words in human-written texts.
Results: Table 2 shows the results. Weakening the bias bLN (λ < 1) increased the diversity of the generated text but decreased the PPL score, exhibiting a general trade-off between them. Nevertheless, for the larger models, GPT-2 large (λ = 0.5) and xl (λ = 0.7), there was a sweet spot, where the diversity and the MAUVE score improved with little decrease in PPL. This observation can be interpreted as follows. The larger models were able to predict the context-dependent probability of low-frequency words as precisely as that of high-frequency words, so promoting lowfrequency words with those models improved the diversity while maintaining the quality of the text. The smaller models were equally accurate in predicting the probability of high-frequency words but tended to be inaccurate for low-frequency words, so promoting low-frequency words degraded the
quality of the text. This interpretation is also consistent with the class imbalance problem, which will be discussed in Section 4.1. From the application perspective, this observation also suggests that the lexical diversity in text generation can be improved simply by modifying particular parameters in the prediction head.
We also show several samples of the generated texts (Appendix C). We generally observed that overly decreasing λ incurs (i) more proper nouns, (ii) more repetitions of the same words or similar phrases, and (iii) the generation of ungrammatical text, especially for the small models. This may also be related to the suppression of the punctuation and end-of-sequence token, which are highly frequent.

heading: 4 Discussion

heading: 4.1 Connection with logit adjustment methods
text: We revealed that adding the bias bLN (which was performed immediately before the logit was computed) encourages TLMs to generate highfrequency words, and de-biasing promotes diversity. This can also be seen as analogous to logit adjustment, which is a common technique for addressing the class imbalance problem, where the label (the word in text generation) frequency distribution is long-tailed (Provost, 2000; Zhou and Liu, 2006; Collell et al., 2016; Menon et al., 2021). In particular, Menon et al. (2021) proposed to minimize the balanced error (i.e., an average of per-class errors) by directly adding the label frequency distribution to logits during training but not during inference.
One can find an analogy between the modification of bLN and their method: (i) adding the frequencyshifting bias bLN corresponds to the operation of adding the class-frequency-based margins to the logits; (ii) promoting low-frequency words by removing bLN during inference corresponds to the way logit adjustment encourages low-class prediction. In other words, interestingly, TLMs seem to implicitly learn something similar to balanced error minimization without being explicitly designed to do so (e.g., loss function).

heading: 4.2 Connection with a technique to initialize bias parameter with class frequency
text: In training neural classification models, using class frequency to initialize the last bias to be added to the logit is a well-known and efficient technique (Karpathy, 2019). Therefore, our observation that the bias vector at the prediction head (i.e., the last block) encodes word frequency might seem somewhat obvious. However, our experimental results showed peculiar trends that might be stemmed from the inductive bias of TLMs. First, although the initialization technique implies the relationship between the last bias blast and the corpus word frequency, we found that the bias bLN ∈ Rd, which is further away from the output and less expressive than blast ∈ R|V|, plays the role in encoding the frequency in BERT (Table 1). For GPT-2, not even blast exists. Second, even bLN plays a weak role in encoding the frequency in larger models (Table 1). These findings suggest that neural models dynamically determines the role of each internal module according to various factors such as parameter size and architecture. When and under what conditions the short vector bLN strongly encodes the frequency is an interesting question and left to future research.

heading: 5 Related work
text: Transformer layers (e.g., attention patterns) have been the major focus of TLM analysis (Clark et al., 2019; Mareček and Rosa, 2019; Kobayashi et al., 2021; Dai et al., 2022). The first embedding layer, especially positional encoding, has also been studied (Wang et al., 2021; Kiyono et al., 2021). 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 (Mu and Viswanath, 2018; Gong et al., 2018; Provilkov et al., 2020; Liang et al., 2021); 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.

heading: 6 Conclusions
text: 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 (Zhang et al., 2022). Further, we will analyze the weight parameters in the prediction head in addition to the bias parameters.
Limitations
There are mainly two limitations in this study. First, we still do not consider components other than the bias parameters in the prediction head. For example, the weight parameters of the prediction head, i.e., γ and W FC, can also affect a model’s prediction. Second, our findings do not cover the Transformer language models other than BERT (base and large) and GPT-2 (small, medium, large, and xl). Consistent findings were obtained for the two main architectures (i.e., encoder-based masked, and decoder-based causal language models) and for various model sizes, although future research is needed to show whether the findings can be generalized to RoBERTa (Liu et al., 2019), Open Pre-trained Transformer Language Models (OPT, Zhang et al., 2022), and other variants. Considering Transformer encoder-decoder models, such as neural machine translation models and T5 (Raffel et al., 2020), would also be an interesting future direction.
Ethics Statement
This paper sheds light on the workings of the prediction head of the fundamental models in NLP. In recent years, unintended biases (e.g., gender bias) in neural network models have been problematic. This paper may help in this direction by encouraging researchers to analyze the prediction head as well as Transformer layers.

heading: Acknowledgements
text: We would like to thank the members of the Tohoku NLP Group for their insightful comments, particularly Shiki Sato for his valuable suggestions regarding experimental settings. This work was supported by JSPS KAKENHI Grant Number JP22J21492, JP22H05106; JST CREST Grant Number JPMJCR20D2, Japan; and JST ACT-X Grant Number JPMJAX200S, Japan.

heading: A Experimental results in other settings
text: In Section 3.1, we presented the changes in the word prediction distribution before and after removing the bias bLN of BERT base and GPT-2 small in Figure 1. The results of the other models are shown in Figures 4 to 7.
In Section 3.2, we showed that the inner product of bLN and each output word embedding wi correlated well with the word frequency for GPT-2 small (Figure 3). The results for the other models are shown in Figures 8 to 12. The Spearman’s correlation coefficient is listed in Table 3.
In Section 3.3, we showed the effect of controlling the bias bLN on GPT-2’s text generation with a top-p sampling strategy. We also conducted experiments with other sampling strategies: top-k sampling (Holtzman et al., 2020) and vanilla sampling. The results of these two sampling strategies are listed in Tables 4 and 5. We found that the results of top-k sampling were similar to those of top-p sampling; for the larger models, GPT-2 large (λ = 0.3) and xl (λ = 0.5), there also was a sweet spot, where diversity and MAUVE improved with little decrease in PPL. In contrast, with vanilla sampling, both MAUVE and PPL decreased consistently and quickly.
Figure 4: Changes in word prediction probabilities due to biasbLN removal on BERT large.
Figure 5: Changes in word prediction probabilities due to biasbLN removal on GPT-2 medium.
Figure 6: Changes in word prediction probabilities due to biasbLN removal on GPT-2 large.
Figure 8: Relationship between the corpus word frequency and the inner product ofbLN and each output word embeddingw i in BERT base.
Figure 9: Relationship between the corpus word frequency and the inner product ofbLN and each output word embeddingw i in BERT large.
Figure 10: Relationship between the corpus word frequency and the inner product ofbLN and each output word embeddingw i in GPT-2 medium.

heading: B Detailed experimental settings
text: To observe the TLM word prediction distribution (the main experiments in Section 3.1 and the measurement of PPL in Section 3.3), we let BERT predict words corresponding to [MASK] tokens, and we let GPT-2 predict the second and subsequent words in each sequence. If the length of an input sequence was greater than the maximum input length k of the model, only the first k words were used.
To evaluate the TLM text generation (Section 3.3), the first 10 words of each sequence were fed into to GPT-2, and subsequent words were generated until the length of the sequence reached 1,024 words or the end-of-sequence token was generated. For GPT-2 small and medium, we varied λ in increments of 0.1 to control the bias bLN. For GPT-2 large and xl, we first checked the results for 100 samples and obtained the values with some kind of trends; we then varied λ in {0, 0.3, 0.5, 0.7, 1.0} for the entire dataset, including the values.
We experimented with three decoding strategies: vanilla sampling, top-k sampling, and top-p sampling. In the top-k sampling, k was set to 50. In the top-p sampling, p was set to 0.9. Furthermore, before we evaluated the model-generated texts with the N-gram based diversity metrics, we applied the word tokenizer provided by NLTK (Bird and Loper, 2004).

heading: C Examples of generated text
text: Table 6 shows examples of text generated by GPT-2 small and large while controlling the bias bLN with λ.
are highlighted with
:::: wavy ::::::::: underlines. Note that the first 10 words are given to the model as context.
ACL 2023 Responsible NLP Checklist

heading: A For every submission:

heading: 3 A1. Did you describe the limitations of your work?
text: Limitations section after Conclusions (Section 5).

heading: 3 A2. Did you discuss any potential risks of your work?
text: We discussed the generalizability of our findings in Limitations section after Conclusions (Section 5).

heading: 3 A3. Do the abstract and introduction summarize the paper’s main claims?
text: Abstract and Introduction (Section 1) summarize our main claims.

heading: 3 A4. Have you used AI writing assistants when working on this paper?
text: We used DeepL and Langsmith. DeepL is a machine translation tool, and Langsmith is a rephrasing tool trained on scientific text. We used them only to refine the English of our submission. Neither service is to copy the work of others nor make novel ideas or claims.
B 3 Did you use or create scientific artifacts? Section 3. We used published pre-trained models and corpus.

heading: 3 B1. Did you cite the creators of artifacts you used?
text: Section 3.
7 B2. Did you discuss the license or terms for use and / or distribution of any artifacts? We have already checked their licenses and the artifacts are enough popular not to need to discuss their license in the paper.
7 B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? It is obvious that all the artifacts we used were created in the context of the research and we also used them for the research purpose.
7 B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? Since we used the dataset for evaluating models’ workings, discussing about its specific content is not so essential for our paper. In addition, the dataset we used is the training corpus of the models. It is natural to use as is.
7 B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and linguistic phenomena, demographic groups represented, etc.? All the artifacts we used are enough popular not to need to provide documentation.
3 B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. Section 3.
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.
C 3 Did you run computational experiments? Section 3.
7 C1. Did you report the number of parameters in the models used, the total computational budget (e.g., GPU hours), and computing infrastructure used? Since the purpose of our study is to analyze models’ inner workings, the details of computational environment are not necessary.

heading: 3 C2. Did you discuss the experimental setup, including hyperparameter search and best-found
text: hyperparameter values? Section 3 and Appendix B.

heading: 3 C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary
text: statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? Section 3.

heading: 3 C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did
text: you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? Section 3.
D 7 Did you use human annotators (e.g., crowdworkers) or research with human participants? Left blank.
D1. Did you report the full text of instructions given to participants, including e.g., screenshots, disclaimers of any risks to participants or annotators, etc.? Not applicable. Left blank.
D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students) and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? Not applicable. Left blank.
D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to crowdworkers explain how the data would be used? Not applicable. Left blank.
D4. Was the data collection protocol approved (or determined exempt) by an ethics review board? Not applicable. Left blank.
D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data? Not applicable. Left blank.

heading: Future Work
text: In future work, we will analyze larger TLMs, e. g. , Open Pre-trained Transformers (Zhang et al. , 2022). Further, we will analyze the weight parameters in the prediction head in addition to the bias parameters. Limitations
There are mainly two limitations in this study. First, we still do not consider components other than the bias parameters in the prediction head. For example, the weight parameters of the prediction head, i. e. , γ and W FC, can also affect a model’s prediction. Second, our findings do not cover the Transformer language models other than BERT (base and large) and GPT-2 (small, medium, large, and xl). Consistent findings were obtained for the two main architectures (i. e. , encoder-based masked, and decoder-based causal language models) and for various model sizes, although future research is needed to show whether the findings can be generalized to RoBERTa (Liu et al. , 2019), Open Pre-trained Transformer Language Models (OPT, Zhang et al. , 2022), and other variants. Considering Transformer encoder-decoder models, such as neural machine translation models and T5 (Raffel et al. , 2020), would also be an interesting future direction