Title: Connecting Pre-trained Language Model and Downstream Task via Properties of Representation
Venue: NeurIPS 2023
Accepted: yes
Keywords: language model representation

downstream performance

deep learning theory
Abstract: Recently, researchers have found that representations learned by large-scale pre-trained language models are useful in various downstream tasks. However, there is little theoretical understanding of how pre-training performance is related to downstream task performance. In this paper, we analyze how this performance transfer depends on the properties of the downstream task and the structure of the representations. We consider a log-linear model where a word can be predicted from its context through a network having softmax as its last layer. We show that even if the downstream task is highly structured and depends on a simple function of the hidden representation, there are still cases when a low pre-training loss cannot guarantee good performance on the downstream task. On the other hand, we propose and empirically validate the existence of an ``anchor vector'' in the representation space, and show that this assumption, together with properties of the downstream task,  guarantees performance transfer.

Connecting Pre-trained Language Models and Downstream Tasks via Properties of Representations

Abstract

Recently, researchers have found that representations learned by large-scale pretrained language models are useful in various downstream tasks. However, there is little theoretical understanding of how pre-training performance is related to downstream task performance. In this paper, we analyze how this performance transfer depends on the properties of the downstream task and the structure of the representations. We consider a log-linear model where a word can be predicted from its context through a network having softmax as its last layer. We show that even if the downstream task is highly structured and depends on a simple function of the hidden representation, there are still cases when a low pre-training loss cannot guarantee good performance on the downstream task. On the other hand, we propose and empirically validate the existence of an 'anchor vector' in the representation space, and show that this assumption, together with properties of the downstream task, guarantees performance transfer.

1 Introduction

Large-scale pre-trained language models have achieved strong performance in a wide range of downstream tasks, including natural language inference [Devlin et al., 2018] and reading comprehension [Brown et al., 2020]. For many of these tasks, training a linear classifier on top of the hidden-layer representations generated by the pre-trained models can already provide near state-of-the-art results [Belinkov et al., 2017]. Despite many empirical investigations about the zero-shot applications of these pre-trained models, there is little theoretical understanding of their empirical success. In this paper, we aim to theoretically investigate this core question:

When can the representations from pre-trained models transfer to downstream tasks that are very different from the pre-training task?

This is a fundamental question in understanding why good performance in pre-training leads to good performance on downstream tasks. Unlike the notion of generalization in traditional learning theory where the models are evaluated in the same task and the test data are sampled from the same distribution as the training data, here the downstream tasks are usually very different from pre-training. For instance, people can pre-train large language models using cross-entropy loss on a language modeling task with webpage data, and evaluate the models using classification accuracy on text classification in news articles. The differences between pre-training and downstream tasks make it challenging to explain the success of these language models.

To overcome this challenge, we need a way to model the relationship between the pre-training and downstream tasks. Previous research has taken several approaches in this direction: Wei et al. [2021] assumes a latent-variable generative model for the data and a downstream task depending on the latent variables; Saunshi et al. [2021] formulates the downstream classification task as a language modeling task which is similar to the pre-training task. These works either rely on strong explicit assumptions about the structure of the data (i.e., assuming the data is generated from a simple generative model) or treat the entire pre-trained model as a black box.

1.1 Our contributions

In this paper, we consider a very general model for the data and open the black box of the pre-trained model at the last layer. Specifically, for an input sequence x = ( x 1 , . . . , x L ) where the entries comes from a dictionary { 1 , . . . , n } , we assume the observation probability of x i satisfies a log-linear model

where x -i is the sequence x without x i , v ∗ j is a vector only depending on word j , and v ∗ -i can be an arbitrary function. This aligns with commonly used networks whose last layer is usually a softmax layer. Moreover, since our model does not put any constraint on the function v ∗ -i , it can be arbitrarily complicated, such as a huge transformer model such as BERT [Devlin et al., 2018] or GPT-3 [Brown et al., 2020]. We also allow the distribution of the input to be different in pre-training and downstream tasks. This makes our setting more general than previous latent models Wei et al. [2021], Arora et al. [2016].

We assume the pre-training task is to predict a word from its context. During pre-training, for every input sequence x , we want our model to predict the 'label' x i from x -i . The model we use in training (which we call the 'student model') has the same log-linear structure: p ( x i = j | x -i ) ∝ exp( ⟨ v -i ( x -i ) , v j ⟩ ) .

For the downstream task, for simplicity we focus on binary sequence classification, e.g., sentiment classification. To define downstream tasks, let the 'logits' of the ground-truth model be defined as z ∗ := ( ⟨ v ∗ -i ( x -i ) , v ∗ j ⟩ ) n j =1 (these are just the outputs before the softmax computation), and assume that the downstream task is specified by a function of the logits, f ∗ ( z ∗ ) .

In reality, we do not have access to the ground-truth model v ∗ -i and v ∗ j . Instead, we only have access to the student model v -i and v j that achieves low pre-training loss. We can define the student logits as z := ( ⟨ v -i ( x -i ) , v j ⟩ ) n j =1 . In some sense, z is the representation learned by the pre-training step. A natural idea for solving the downstream task would be to learn a function f ( z ) . More details about this model will be provided in Section 2.

Our goal is to understand the properties of the learned representation. In order to simplify the problem, we assume that the student model can be optimized to achieve a small KL divergence with the true word probabilities during training. Under this setting, the question we ask above becomes:

If the downstream task depends on a simple function of the logits f ∗ ( z ∗ ) and we have access to a student model p such that E x [ D KL ( p ∗ ( x i | x -i ) || p ( x i | x -i ))] is small, under what conditions is there a function f such that f ( z ) ≈ f ∗ ( z ∗ ) ?

A priori, one might imagine if the function f ∗ is very simple (e.g., a linear function), then it should be easy to find a function f . However, we give two counter-examples that show there are additional properties that f ∗ needs to satisfy: (i) f ∗ should not distinguish between words with small probabilities and words with super-small probabilities, and (ii) the hidden representations must have some structure that deals with the shift-invariance of the softmax function (that is, the result of softmax does not change if all the logits are shifted by the same constant). We present the counterexamples in Section 3.

To further investigate the structure of the hidden representations and see how we can deal with the shift-invariance property of softmax, in Section 4, we propose and empirically verified the 'anchor vector hypothesis': there exists an 'anchor vector' in the representation space that can be used to estimate the bulk partition function , which we define to be the sum of the exponential of all logits except the largest few, i.e., ∑ j : z ∗ j not large e z ∗ j . We show how the anchor vector can be used to address the shift-invariance of softmax.

Based on the observation that anchor vectors exist, in Section 5, we give sufficient conditions that shows when a sparse one-hidden-layer ReLU network f ∗ can be learned by our student model f . Specifically, assuming that f ∗ is a one-hidden-layer ReLU network depending on a small set of words and the downstream task is binary classification depending on f ∗ ( z ∗ ) , the existence of the anchor vector enables us to upper bound the loss on the downstream task attained by the student model f ( z ) , in terms of its KL divergence in pre-training. In other words, a small pre-training loss is guaranteed to transfer to a small downstream classification error.

1.2 Related works

Theoretical understanding why pre-training helps downstream tasks : Most of the relevant existing works rely on latent variable models and show that pre-training could recover some form of the latent variables. Arora et al. [2016] proposed the RAND-WALK latent model and explains the empirical success of word embedding methods such as word2vec [Mikolov et al., 2013] and GloVe [Pennington et al., 2014]. Arora et al. [2017] extended the previous model to justify sentence embeddings, and Arora et al. [2018] explained sentence embedding via compressed sensing. Other models are also used in this line of work, e.g., hidden Markov models [Wei et al., 2021] and graphical models [Zhang and Hashimoto, 2021]. Lee et al. [2021] and Tosh et al. [2021] assume conditional independence or multi-view structure in the pre-training data and prove that training an additional linear layer on top of learned representations can perform well in downstream tasks.

The problem setting in our paper is similar to that of Saunshi et al. [2021], which also analyzes the performance transfer of pre-trained language models to binary classification downstream tasks. They treat the pre-trained model as a black box and assume that the downstream task can be formulated as a sentence completion task, while we open the black box at the last layer and connect pre-training with downstream tasks by the 'anchor vector' and function f ∗ . Moreover, they focus on the prediction probabilities of the pre-trained model while we instead focus on the representations.

Applications and analysis of hidden representations from large-scale language models : The hidden representations produced by large language models such as BERT [Devlin et al., 2018] or ELMo [Peters et al., 2018] have been very useful in various NLP tasks. A standard method is to train a linear classifier on these representations, though there are other methods such as using the normalized mean of concatenated word embeddings [Tanaka et al., 2020]. To understand why these word embeddings are useful, people have empirically showed that BERT word embeddings contain information about sentence-level context [Miaschi and Dell'Orletta, 2020], word sense [Wiedemann et al., 2019], and syntactic phenomena [Tenney et al., 2019] including parse trees [Hewitt and Manning, 2019, Kim et al., 2020]. Other empirical explanations include the flatness of local minima achieved by pre-training [Hao et al., 2019], connection to deep metric learning [Tschannen et al., 2019], and the attention patterns across different heads [Kovaleva et al., 2019].

Language modeling can also be considered as a way of using the hidden representations because the next word probability is usually the softmax of the product of the representation and the dictionary matrix. Therefore, any zero-shot application of pre-trained auto-regressive language models, e.g., GPT-3 [Brown et al., 2020] and T5 [Raffel et al., 2020], is a specific method of using the hidden representations.

Some previous works have found that the word embeddings learned by language models can lie in a narrow cone: Gao et al. [2019] empirically found this phenomena for the learned word embeddings in LSTM and vanilla transformers. Ethayarajh [2019], Cai et al. [2020] found similar phenomena for contextualized embeddings in BERT-like models.

2 Problem setup

Notations. We use [ n ] to denote the set { 1 , 2 , . . . , n } . For an input sequence x = ( x 1 , . . . , x L ) , we use x -i to denote the input sequence without the i -th entry where i ∈ [ L ] , i.e., x -i := ( x 1 , . . . , x i -1 , x i +1 , . . . , x L ) . We let D KL ( p || q ) be the KL-divergence between distributions p and q and define H ( p ) to be the entropy of distribution p .

Ground-truth model. We consider the following model: There is a set of words [ n ] , each with a fixed corresponding vector v ∗ j ∈ R d ( j ∈ [ n ]) . We refer to each v ∗ j as an atom and the matrix [ v ∗ 1 , · · · , v ∗ n ] ∈ R d × n as the dictionary. At each position i , let x i be the value of the word at that position; then x i ∈ [ n ] . Assume that the probability of x i given x -i follows a log-linear model, i.e.,

where v ∗ -i ( · ) is a function that encodes the remaining sequence x -i into a vector in R d .

We also use z ∗ j ( x, i ) := ⟨ v ∗ -i ( x -i ) , v ∗ j ⟩ to denote the j -th logit and Z ∗ ( x, i ) := ∑ n j =1 exp( z ∗ j ( x, i )) to denote the partition function, i.e., the normalization factor of equation (1). In other words,

Student model. We use a black-box neural network model whose penultimate layer outputs a d ′ -dimensional vector v -i ( x -i ) ∈ R d ′ and the last layer is a fully-connected layer with weight matrix [ v 1 , v 2 , · · · , v n ] ∈ R d ′ × n followed by the softmax function. In other words, the model output is

Similar to the ground-truth model, we use z j ( x, i ) := ⟨ v -i ( x -i ) , v j ⟩ to denote the j -th logit and Z ( x, i ) := ∑ n j =1 exp( z j ( x, i )) to denote the partition function of equation (3), so

Pre-training. For self-supervised pre-training, we are given (data, 'label') pairs ( x -i , x i ) , and want our model to predict the 'label' x i given x -i . The pre-training loss we use here is cross-entropy loss:

Note that E x [ H ( p ∗ ( x i | x -i ))] is a constant, and we assume that our student model achieves a small loss value so that the KL-divergence term E x [ D KL ( p ∗ ( x i | x -i ) || p ( x i | x -i ))] ≤ ϵ KL for some ϵ KL .

Downstream task. The downstream task we are considering is binary sequence classification, e.g., sentiment classification. For instance, given a sentence, we (probably after adding some prompts) use our pre-trained model to predict the missing word x i from the given input x -i . We assume that there is a perfect classifier f ∗ ( x, i ) only depending on v ∗ -i that can distinguish between positive and negative samples. In other words, for all x ∈ POS , f ∗ ( x, i ) > 0 and for all x ∈ NEG , f ∗ ( x, i ) < 0 . Here POS and NEG are the set of positive and negative input sequences, respectively. We choose to focus on binary classification for theoretical analysis for simplicity, but most of the ideas can be extended to a multi-class classification setting.

A simple downstream task is one whose classifier is linear in v ∗ -i , that is, f ∗ ( x, i ) := ⟨ v ∗ -i ( x -i ) , u ∗ ⟩ for some u ∗ ∈ R d . One might also expect more structures in the vectors u ∗ and v ∗ -i . For u ∗ , the downstream task usually depends on only a small set of the words which are related to this task. For example, for sentiment classification, the sentiment of a sentence depends mostly on words similar to 'positive' or 'negative'. This can be formalized as the following assumption:

Assumption 1 ( u ∗ is a k -sparse combination of { v ∗ j } n j =1 ) . Assume u ∗ is a sparse combination of at most k vectors in { v ∗ j } n j =1 and without loss of generality assume these k vectors are { v 1 , . . . , v k } , i.e., there exist coefficients { c ∗ j } k j =1 ∈ R k such that u ∗ = ∑ k j =1 c ∗ j v ∗ j .

Note that the fixed representations { v 1 , . . . , v k } correspond to the weights of the last layer in large language models instead of the token representations v ∗ -i ( x -i ) . In other words, the token representations can still be different depending on their context.

There are usually differences between input distributions of pre-training and downstream tasks. The difference may be due to different data sources. For example, the samples in the downstream task

may only contain movie reviews while the pre-training dataset can include all kinds of texts on the Internet. It can also result from prompting, which has become the dominant way of using large language models for downstream tasks [Brown et al., 2020, Radford et al., 2019]. For instance, for movie review classification, appending 'This movie was' to the original input could improve the classification accuracy. Similar to the assumption made in Saunshi et al. [2021], we use µ ∈ (0 , 1] to capture this difference. A smaller µ indicates a larger difference between the two distributions, and µ = 1 if and only if these two distributions are the same. In most tasks, the two distributions are not too different and we would expect a reasonable value of µ .

Assumption 2 (Difference between pre-training and downstream distribution) . Let p pre and p DS be the probability density functions of the pre-training and downstream task, respectively. We assume that there exists µ ∈ (0 , 1] such that

If the ground truth classifier f ∗ can be too close to 0, it would not be robust to small perturbations. We use the standard margin assumption to avoid such cases:

Assumption 3 (Margin for downstream task) . There exists a margin γ ∈ R + such that at any position i , if x ∈ POS , then f ∗ ( x, i ) ≥ γ , and if x ∈ NEG , then f ∗ ( x, i ) ≤ -γ , where POS and NEG are the sets of positive and negative samples, respectively.

Ideally, we want to show that such simple downstream tasks can also be solved well with the representations learned by our student model, i.e., v -i and { v j } k j =1 . However, as we will see in Section 3, our current model with this margin assumption still doesn't guarantee good downstream task performances. More structures in the model and the downstream task is necessary to make sure that the pre-trained representations are useful for the downstream task.

3 Cases when learned representations are insufficient for downstream tasks

The problem setting in Section 2 seems reasonable at first sight, but in the following subsections, we will show that this model is not enough to guarantee good pre-training performance to generalize to downstream tasks. In other words, there are ways for the student model to approximate the groundtruth probabilities very well in terms of KL divergence but perform very badly at the downstream task. Therefore, we need to put further constraints on the ground-truth model and the downstream task.

3.1 Downstream tasks sensitive to words with super-small probability

Intuitively, KL divergence is a weighted log probability difference between two distributions where the weight is the ground-truth probability. Therefore, for the entries with small ground-truth probabilities, a large log probability difference will not result in a large KL divergence. However, the log probability difference is proportional to the difference in the value of f ∗ ( x, i ) . This makes it possible for the student model to flip the sign of f ∗ ( x, i ) without incurring a large KL divergence, as presented in Theorem 1 whose proof is given in Appendix A.

Theorem 1. Suppose the downstream task performance depends only on a function f ∗ ( x, i ) = ⟨ v ∗ -i ( x -i ) , u ∗ ⟩ = ∑ k t =1 c ∗ t ⟨ v ∗ -i ( x -i ) , v ∗ t ⟩ . For t -∈ [ k ] , define p -:= p ∗ ( x i = t -| x -i ) , and assume p -≤ 1 2 . Then for all s ∈ R + , there exist functions v -i and { v t } k t =1 such that D KL ( p ∗ ( x i | x -i ) || p ( x i | x -i )) ≤ 2 sp -and f ( x, i ) := ∑ k t =1 c ∗ t ⟨ v -i ( x -i ) , v t ⟩ ≤ f ∗ ( x, i ) -s · c ∗ t -.

Theorem 1 shows that if there is some word of interest t -that has a small probability p -, then it is possible to have a model with small KL divergence in pre-training but bad downstream performance. This is because changing the KL divergence by only 2 p -· f ∗ ( x,i ) c ∗ t -is enough to change the label of the downstream prediction. In other words, as long as the KL divergence is higher than the threshold 2 p -· f ∗ ( x,i ) c ∗ t -, we cannot distinguish between the case where the student model makes an already small probability even smaller (which can hurt the downstream task performance) and the case where random approximation errors are spread across the entries. In this case, a small KL divergence does not necessarily imply good downstream performance.

Note that this sensitivity of the downstream task to very small logits is not natural. For the downstream tasks in practice, after conditioning on the context, whether a word has a probability of 10 -5 or

10 -10 should not influence the label of the sequence. Thus, we need to impose additional structure on our model. We make the downstream task ignore super-small entries by setting a threshold for the logits and ignoring the logits smaller than that threshold. In this case, making the logits smaller when they are already small will have no influence on the downstream task performance. Concretely, the enhanced model will be ( σ ( x ) := max { x, 0 } is the ReLU function):

3.2 Representations are not shift-invariant

The softmax function is invariant under shift, i.e., the output stays the same if we add the same value to every coordinate of the input. In the current model, we have no control over the shift of student model logits on unseen data. Consequently, even if we get a student model that performs well on the training data for the downstream task, we cannot guarantee the performance of this model on new data. This can be formalized in the following theorem.

Theorem 2. Assume z ∗ ( x, i ) is bounded. For any function f ∗ ( x, i ) = ∑ n j =1 a j σ ( z ∗ j ( x, i ) -b j ) , there exist functions { ˆ z j ( x, i ) } n j =1 such that for all x and i , we have ˆ p ( x i | x -i ) = p ∗ ( x i | x -i ) and ˆ f ( x, i ) := ∑ n j =1 a ∗ j σ (ˆ z j ( x, i ) -b ∗ j ) is always equal to 0 . In other words, the pre-training loss of the model { ˆ z j ( x, i ) } n j =1 is the same as { z j ( x, i ) } n j =1 , but its logits are useless for the downstream task.

Proof. We choose τ ∈ R such that ∀ x, i, τ < min j ∈ [ n ] b ∗ j -max j ∈ [ n ] z ∗ j ( x, i ) , and ∀ x, i, ∀ j ∈ [ n ] , we set ˆ z j ( x, i ) := z ∗ j ( x, i ) + τ , then

Theorem 2 indicates that without any structure in the representations, the student model is able to shift the logits for any sample and keep the pre-training loss unchanged. In the worst case, it can shift the logits for unseen data drastically, resulting in a bad downstream performance. Therefore, a theoretical guarantee for downstream performance requires structure in the representations learned by the pre-trained model.

4 'Anchor vector' hypothesis and empirical verifications

In Section 3.2, we showed that the shift-invariance of the softmax function can potentially make the student logits useless for the downstream task. Therefore, to understand why the downstream tasks benefit from representations from the pre-trained model, we need to understand the structure of these representations, and this structure must be able to handle the shift-invariance problem.

4.1 'Anchor vector' hypothesis

There are different ways to prevent the shift-invariance of softmax from influencing the performance of the downstream tasks. One way of doing this is to keep the partition function stable. Recall that in (4), the probability of a word is the exponential of the corresponding logit divided by the partition function. If the partition function is constant for different samples, the logits can be uniquely determined by the probabilities, which solves the shift-invariance problem. Arora et al. [2016] showed that when both the word embeddings and the latent representation are uniformly distributed on a sphere, the partition function is close to a constant with high probability. They also empirically verified this uniformity of word embeddings trained using GloVe [Pennington et al., 2014] and word2vec [Mikolov et al., 2013] on English Wikipedia. However, as we will see in later experiments in Sections 4.2 and F, this is not true for recent large-scale pre-trained language models.

Instead of uniformity of word embeddings, in large pre-trained models such as GPT-2 [Radford et al., 2019], we observe that if we remove several most frequent words from the computation of the log partition function, the remaining part can be well approximated by the inner product between the

hidden representation v -i ( x ) and a fixed vector. This motivates us to have the following 'anchor vector' hypothesis:

Definition 1. For a sample x and position i , we could select a set of bulk words B ( x, i ) ⊂ [ n ] , and we define the bulk partition functions as Z ∗ bulk ( x, i ) := ∑ j ∈ B ( x,i ) exp( ⟨ v ∗ -i ( x -i ) , v ∗ j ⟩ ) and Z bulk ( x, i ) := ∑ j ∈ B ( x,i ) exp( ⟨ v -i ( x -i ) , v j ⟩ ) .

The selection of bulk words B ( x, i ) can usually be selected manually or by a simple algorithm. For instance, we can construct B ( x, i ) by taking out the words corresponding to the largest entries in p ( x i | x -i ) . We can also manually select all the words that are irrelevant to the downstream task.

Hypothesis 1 ('Anchor vector' hypothesis) . There exists v 0 ∈ R d such that

If our hypothesis holds, we can use v 0 as an anchor for the logits because it can be used to estimate the partition function, and this does not change the form of the downstream task by much. In the next subsection, we will show that Hypothesis 1 holds for most popular language models and provide some discussions. More details about these models are provided in Appendix E.

4.2 Empirical verification of 'anchor vector' hypothesis

Figure 1a plots the mean squared approximation error of the log bulk partition function. We use different versions of pre-trained GPT-2 [Radford et al., 2019] and OPT [Zhang et al., 2022], and use the first 1 / 4 of WikiText-2 [Merity et al., 2016] as the input text. The hidden representations we use in this experiment are the last hidden states of these models, i.e., the output of the penultimate layer. The dimension of the hidden representations ranges from 768 to 2048, and the number of tokens is about 70k. We choose the bulk words to be all the words except those having topk probabilities and compute the optimal anchor vector using the closed-form least-squares solution. In our experiments, we use the mean squared error (MSE) to measure the approximation quality. Formally, the MSE is defined as

The values of the bulk partition functions are around 10 (with comparable standard deviation), and we can see from Figure 1a that the MSE is usually several orders of magnitude smaller. Therefore, the inner product between the hidden representation and the optimal anchor vector can usually approximate the log bulk partition function well, and the approximation improves as k increases, i.e., when we ignore more top words. This validates our 'anchor vector' hypothesis.

100

120

8000

original

：approximation error

100

6000

7000

bullk

80

5000

10-1

60

4000

O

3000

20

2000

0

1000

mean square

10-2

GPT-2 Medium

0

200

400

600

800

10001200

14001600

6

8

10

14

16

18

Z

GPT-2 Large

GPT-2 XL

160

140

10000

original

bulk

OPT-125M

120

100

8000

103

OPT-1.3B

OPT-350M

08

60

6000

4000

40

101

102

103

104

2000

k

0

0

005

1000

1500

2000

0

10

12

14

16

18

OZ

(a) Mean squared error of log bulk partition function with different k in different models. Bulk partition function includes all logits except topk ones. (b) Singular values of dictionary matrices (up: GPT2 XL, down: OPT-1.3B) (c) Histogram of log original/bulk partition functions (GPT-2 XL/OPT-1.3B)

The existence of anchor vector is quite surprising. Although there are some settings where the partition function can become easy to predict, e.g., when the word embeddings are uniformly distributed on a sphere, we have empirically shown that the large-scale language models considered in this paper

do not fall under these settings. Figure 1b shows that the word embeddings are far from uniformly distributed on a sphere, and Figure 1c indicates that the values of log partition functions and log bulk partition functions vary greatly for different samples. More experiments and discussions are provided in Section F.

5 Anchor vector guarantees performance transfer from pre-training to downstream tasks

Based on the counterexamples in Section 3 and the observations in Section 4, we now give sufficient conditions on the downstream tasks so that the downstream classification accuracy provably benefits from pre-trained representations.

5.1 Model and assumptions

For huge language models, the anchor vector hypothesis states that there exists a vector v ∗ 0 such that its inner product with the hidden state v ∗ -i ( x -i ) well approximates the logarithm of the bulk partition function. Therefore, we can use v ∗ 0 as an anchor to handle the shift invariance problem of the softmax function, i.e., we want to subtract ⟨ v ∗ -i ( x -i ) , v ∗ 0 ⟩ from ⟨ v ∗ -i ( x -i ) , v ∗ j ⟩ . As a result, we modify our model for the downstream task to be

The student model is also modified accordingly:

Note that in the above model we have already adapted Assumption 1 which was originally stated for the linear model. Therefore, we also update the assumption and restate it below:

Assumption 4 (At most k words of interest) . Assume there are at most k vectors in { v ∗ j } n j =1 whose logits are relevant to the downstream task and WLOG assume these k vectors are { v ∗ 1 , . . . , v ∗ k } . In other words, we assume there exist coefficients { a ∗ j } k j =1 ∈ R k such that f ∗ ( x, i ) := ∑ k j =1 a ∗ j σ ( ⟨ v ∗ -i ( x -i ) , v ∗ j -v ∗ 0 ⟩ -b ∗ j ) .

We are still assuming the teacher-student setting and the log-linear word production model for the word probabilities, as restated below:

Under the modified model defined above, we will make some additional assumptions. As noticed in the experiments, the log bulk partition function (as defined in Definition 1) can be linearly approximated by the hidden state. Normally, the bulk only contains words that are not related to the downstream task, and every single word usually has a small probability, but the total probability of the bulk words is not negligible, as reflected in the following two assumptions. Note that the set of bulk words can be strictly contained in the complement of the set of words of interest. This is useful especially when we are not certain about which words are important for the downstream task.

Assumption 5 (Bulk contains no words of interest) . For all x and i , B ( x, i ) ∩ { 1 , . . . , k } = ∅ . Assumption 6 (Lower bound of bulk probability) . For all x and i ,

The 'anchor' vector v ∗ 0 ( i ) is used to handle the instability of the partition function, and we formalize our anchor vector hypothesis as the following assumption:

Assumption 7 (Linear approximation of log bulk partition function) . There exist v 0 , v ∗ 0 ∈ R d , ε b ∈ R + , s.t.,

Furthermore, we assume ε b ≤ γ 4 k max j ∈ [ k ] | a ∗ j | .

For notational simplicity, we will use all the notations (including f ∗ , f , v ∗ 0 , and v 0 ) without i when the selection of i is clear from the context.

5.2 Main theorem and interpretations

In our model, the ground-truth function f ∗ contains k terms with coefficients { a ∗ j } k j =1 and we define the margin γ as the margin for f ∗ . If we scale { a ∗ j } k j =1 or increase k by adding duplicated terms to f ∗ , we can scale γ arbitrarily without changing the pre-training performance or the downstream prediction of the student model. To construct a quantity that better indicates the difficulty of the downstream task, we introduce the following definition of the normalized margin that is invariant to the scaling of k and { a ∗ j } k j =1 :

Definition 2. The normalized margin is defined as Γ := γ k max j ∈ [ k ] | a ∗ j | .

Then we are ready to state our main result. The proof and further discussions for this theorem will be provided in Section B.

Theorem 3. Let ϵ KL := E x ∼D pre [ D KL ( p ∗ ( x ) || p ( x ))] be the pre-training loss, and ϵ CLS := Pr x ∼D DS [ f ( x ) · f ∗ ( x ) < 0] be the downstream classification error rate, where D pre and D DS are the input distributions of pre-training and downstream data. Under Assumptions 2-7, further assuming min j ∈ [ k ] b ∗ j ≥ ϵ b -log(4 k ) and 8 ϵ b < Γ < 6 , there exists a set of parameters ( a j ) n j =1 , ( b j ) n j =1 such that

This theorem shows that when an 'anchor vector' exists, we can upper bound the downstream classification error by the KL divergence of the student model during pre-training. This upper bound becomes smaller when the distribution of downstream input is close to that of pre-training, the bulk probability is non-negligible, and the normalized margin of the ground-truth classifier is large. Here we discuss these ways to decrease the upper bound and their corresponding intuitions.

Large µ . µ is larger when the data distributions of pre-training and the downstream task are closer, which helps with the performance transfer.

∗

Large p b . The bulk probability Z bulk Z ∗ is usually at least a constant in practice. When the bulk probability becomes larger, the anchor vector plays a more important role in the partition function and the downstream task.

Large Γ . A larger normalized margin makes it harder for the student model to make mistakes in downstream prediction.

Remark 1. For ease of presentation, we are making additional assumptions in Theorem 3, e.g., min j ∈ [ k ] b ∗ j ≥ ϵ b -log(4 k ) and 8 ϵ b < Γ < 3 . More general cases are covered in Lemma 1. We also discuss potential ways to improve the bound in Section B.1.

6 Conclusions, limitations and future work

In this paper, we analyzed when and how the representations generated by pre-trained large-scale language models can be used in downstream classification tasks. We found two necessary conditions for guaranteeing the usefulness of these representations: the insensitivity of the downstream task to super-small probability words, and the underlying structure in the representations to handle the shiftinvariance of softmax. We also provide a sufficient condition, i.e., the existence of an anchor vector, that can guarantee the representations from pre-trained language model to help with downstream tasks. We verify this existence empirically in various large language models and believe that this is an important reason why recent large-scale language models can adapt to different downstream tasks.

We made some assumptions to provide theoretical understanding on how representations can be helpful for downstream applications, and these assumptions may not always hold in practice. For instance, some downstream tasks require multiple rounds of reasoning and do not rely on a small set of words, and the 'anchor vector' hypothesis can be weaker for some models such as OPT-350M and GPT-2 Medium. Relaxing these assumptions is a meaningful future direction.

While our work showed the existence of the anchor vector, it remains unclear why this vector exists in most large language models. It might be related to the initialization, optimization, and structure of the neural networks, especially transformers, and it could also be related to the underlying structure of the training data. Digging deeper into this may reveal more fundamental properties of these models. Our analysis of downstream tasks focuses on classification tasks. Since these large networks perform well in various types of downstream tasks, another future direction would be to analyze other kinds of downstream tasks.

We only consider using the last layer representation without fine-tuning the whole model, which is usually weaker in performance. Furthermore, we model the network before the softmax function as a black box, ignoring its inner structure. Further opening up this black box and consider fine-tuning the whole model are important for deeper understanding of the structures in the learned representations from pre-training.

References

Arora, Y. Li, Y. Liang, T. Ma, and A. Risteski. A latent variable model approach to pmi-based word embeddings. Transactions of the Association for Computational Linguistics , 4:385-399, 2016.

Arora, Y. Liang, and T. Ma. A simple but tough-to-beat baseline for sentence embeddings. In International conference on learning representations , 2017.

Arora, M. Khodak, N. Saunshi, and K. Vodrahalli. A compressed sensing view of unsupervised text embeddings, bag-of-n-grams, and lstms. In International Conference on Learning Representations , 2018.

Belinkov, N. Durrani, F. Dalvi, H. Sajjad, and J. Glass. What do neural machine translation models learn about morphology? arXiv preprint arXiv:1704.03471 , 2017.

Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Cai, J. Huang, Y. Bian, and K. Church. Isotropy in the contextual embedding space: Clusters and manifolds. In International Conference on Learning Representations , 2020.

Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

Ethayarajh. How contextual are contextualized word representations? comparing the geometry of bert, elmo, and gpt-2 embeddings. arXiv preprint arXiv:1909.00512 , 2019.

Gao, D. He, X. Tan, T. Qin, L. Wang, and T.-Y. Liu. Representation degeneration problem in training natural language generation models. arXiv preprint arXiv:1907.12009 , 2019.

Hao, L. Dong, F. Wei, and K. Xu. Visualizing and understanding the effectiveness of bert. arXiv preprint arXiv:1908.05620 , 2019.

Hewitt and C. D. Manning. A structural probe for finding syntax in word representations. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pages 4129-4138, 2019.

Kim, J. Choi, D. Edmiston, and S.-g. Lee. Are pre-trained language models aware of phrases? simple but strong baselines for grammar induction. arXiv preprint arXiv:2002.00737 , 2020.

Kovaleva, A. Romanov, A. Rogers, and A. Rumshisky. Revealing the dark secrets of bert. arXiv preprint arXiv:1908.08593 , 2019.

D. Lee, Q. Lei, N. Saunshi, and J. Zhuo. Predicting what you already know helps: Provable self-supervised learning. Advances in Neural Information Processing Systems , 34:309-323, 2021.

Merity, C. Xiong, J. Bradbury, and R. Socher. Pointer sentinel mixture models, 2016.

Miaschi and F. Dell'Orletta. Contextual and non-contextual word embeddings: an in-depth linguistic investigation. In Proceedings of the 5th Workshop on Representation Learning for NLP , pages 110-119, Online, July 2020. Association for Computational Linguistics. doi: 10.18653/v1/ 2020.repl4nlp-1.15. URL https://aclanthology.org/2020.repl4nlp-1.15 .

Mikolov, K. Chen, G. Corrado, and J. Dean. Efficient estimation of word representations in vector space. arXiv preprint arXiv:1301.3781 , 2013.

Pennington, R. Socher, and C. D. Manning. Glove: Global vectors for word representation. In Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP) , pages 1532-1543, 2014.

E. Peters, M. Neumann, M. Iyyer, M. Gardner, C. Clark, K. Lee, and L. Zettlemoyer. Deep contextualized word representations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers) , pages 2227-2237, New Orleans, Louisiana, June 2018. Association for Computational Linguistics. doi: 10.18653/v1/N18-1202. URL https://aclanthology.org/N18-1202 .

Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research , 21(1):5485-5551, 2020.

Saunshi, S. Malladi, and S. Arora. A mathematical exploration of why language models help solve downstream tasks. In International Conference on Learning Representations , 2021.

Tanaka, H. Shinnou, R. Cao, J. Bai, and W. Ma. Document classification by word embeddings of bert. In Computational Linguistics: 16th International Conference of the Pacific Association for Computational Linguistics, PACLING 2019, Hanoi, Vietnam, October 11-13, 2019, Revised Selected Papers 16 , pages 145-154. Springer, 2020.

Tenney, P. Xia, B. Chen, A. Wang, A. Poliak, R. T. McCoy, N. Kim, B. Van Durme, S. R. Bowman, D. Das, et al. What do you learn from context? probing for sentence structure in contextualized word representations. arXiv preprint arXiv:1905.06316 , 2019.

Tosh, A. Krishnamurthy, and D. Hsu. Contrastive learning, multi-view redundancy, and linear models. In Algorithmic Learning Theory , pages 1179-1206. PMLR, 2021.

Tschannen, J. Djolonga, P. K. Rubenstein, S. Gelly, and M. Lucic. On mutual information maximization for representation learning. arXiv preprint arXiv:1907.13625 , 2019.

Wei, S. M. Xie, and T. Ma. Why do pretrained language models help in downstream tasks? an analysis of head and prompt tuning. Advances in Neural Information Processing Systems , 34: 16158-16170, 2021.

Wiedemann, S. Remus, A. Chawla, and C. Biemann. Does bert make any sense? interpretable word sense disambiguation with contextualized embeddings. arXiv preprint arXiv:1909.10430 , 2019.

Zhang, S. Roller, N. Goyal, M. Artetxe, M. Chen, S. Chen, C. Dewan, M. Diab, X. Li, X. V. Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068 , 2022.

Zhang and T. Hashimoto. On the inductive bias of masked language modeling: From statistical to syntactic dependencies. arXiv preprint arXiv:2104.05694 , 2021.

A Proof of Theorem 1

Theorem 1. Suppose the downstream task performance depends only on a function f ∗ ( x, i ) = ⟨ v ∗ -i ( x -i ) , u ∗ ⟩ = ∑ k t =1 c ∗ t ⟨ v ∗ -i ( x -i ) , v ∗ t ⟩ . For t -∈ [ k ] , define p -:= p ∗ ( x i = t -| x -i ) , and assume p -≤ 1 2 . Then for all s ∈ R + , there exist functions v -i and { v t } k t =1 such that D KL ( p ∗ ( x i | x -i ) || p ( x i | x -i )) ≤ 2 sp -and f ( x, i ) := ∑ k t =1 c ∗ t ⟨ v -i ( x -i ) , v t ⟩ ≤ f ∗ ( x, i ) -s · c ∗ t -.

Proof. We choose functions v -i and { v k t } r t =1 such that ∀ x, i , ∀ t ∈ [ n ] \ { t -} , ⟨ v -i ( x -i ) , v t ⟩ = ⟨ v ∗ -i ( x -i ) , v ∗ t ⟩ . Besides, ∀ x, i , ⟨ v -i ( x -i ) , v t -⟩ = ⟨ v ∗ -i ( x -i ) , v ∗ t -⟩ -s .

From the definition of p -and t -we know that

Following our construction of the student model, its partition function satisfies

̸

̸

̸

Besides,

Therefore,

Thus,

and

B Proof and discussions for main theorem

Here we provide a proof sketch of Theorem 3.

Theorem 3. Let ϵ KL := E x ∼D pre [ D KL ( p ∗ ( x ) || p ( x ))] be the pre-training loss, and ϵ CLS := Pr x ∼D DS [ f ( x ) · f ∗ ( x ) < 0] be the downstream classification error rate, where D pre and D DS

are the input distributions of pre-training and downstream data. Under Assumptions 2-7, further assuming min j ∈ [ k ] b ∗ j ≥ ϵ b -log(4 k ) and 8 ϵ b < Γ < 6 , there exists a set of parameters ( a j ) n j =1 , ( b j ) n j =1 such that

Proof. We have the following lemma providing a lower bound of the TV distance between p ∗ ( x i | x -i ) and p ( x i | x -i ) , which will be proved in Section C.

Lemma 1. There exists a choice of ( a j ) k j =1 , ( b j ) k j =1 such that if Assumptions 2-7 hold and there exists x ∈ POS ∪ NEG and i such that f ( x, i ) · f ∗ ( x, i ) < 0 , then we must have

From Lemma 1and Pinsker's inequality we can lower bound the KL divergence at any incorrectly classified sample: For all x ∈ POS ∪ NEG ,

Since the pre-training loss is the expected KL divergence for pre-training data, it is lower bounded by the KL divergence on the incorrectly classified samples, which is related to the downstream classification error rate:

Here D pre and D DS are the data distributions of pre-training and downstream task, and the second inequality comes from Assumption 2.

Now we have established the relationship between ϵ KL and ϵ CLS , and the only step left is to bound the minimum over ε p . For notation simplicity, we define Λ := Γ 2 -2 ϵ b .

If ε p > Λ 3 , we know that

Thus,

Plugging this into (28) gives us

Since Γ > 8 ϵ b , we know that Γ 4 < Λ < 3 , so (1 -e -Λ 3 ) 2 > ( Λ 6 ) 2 > Γ 2 576 .

Thus, min { (1 -e -Λ 3 ) 2 , 4 k 2 · min j ∈ [ k ] e 2 b ∗ j -2 ε b · sinh 2 ( Λ 3 ) } > Γ 2 576 . Plugging this into (31) finishes the proof of Theorem 3.

B.1 Further discussions of Theorem 3

Potentially improving the upper bound by better thresholding. We are taking the minimum over two terms in the proof of Theorem 3. These two terms come from two parts of the difference between the ground-truth probability and our learned probability: the bulk part and the top part. Both parts are functions of ϵ p , which measures the approximation error for the bulk probability. A large ϵ p will incur a large error in the bulk part while a smaller ϵ p will make the top part larger. We are setting Λ 3 as the threshold for the current bound in Theorem 3. The bound could be improved by choosing other thresholds if we know more about the values of { b ∗ j } k j =1 and the normalized margin Γ .

More general cases for thresholds { b ∗ j } k j =1 and normalized margin Γ . The thresholds { b ∗ j } k j =1 can be considered as the model's sensitivity to the logits. A smaller threshold indicates a higher sensitivity. In the proof of Theorem 3, we are assuming that the thresholds are not very small. In general cases where the thresholds can be small, an important observation is that: For the same set of samples, decreasing the thresholds by some number s will increase the normalized margin by s as well. Since both the thresholds and the normalized margin are in the exponent in our bound ( sinh function is close to exponential when Γ is large), these two effects can cancel out to a large degree and don't influence the bound by much.

Prompt engineering can help with downstream performance. Prompt engineering can help with the downstream task performance in multiple ways. One direct way is to make the data distribution closer to that of the pre-training stage to increase µ . Furthermore, it can also indirectly help improve the bound by decreasing the number of activated neurons, increasing the margin, etc.

C Proof of Lemma 1

Here we provide a proof of Lemma 1, with the proofs of the necessary technical lemmas postponed to Appendix D.

Lemma 1. There exists a choice of ( a j ) k j =1 , ( b j ) k j =1 such that if Assumptions 2-7 hold and there exists x ∈ POS ∪ NEG and i such that f ( x, i ) · f ∗ ( x, i ) < 0 , then we must have

Proof. For all j ∈ [ n ] , we set a j = a ∗ j and b j = b ∗ j . Since i is fixed in the proof, we will omit i when the selection of i is clear from context.

WLOG, assume x ∈ POS because otherwise we can flip the sign of all a ∗ j 's. In this case, from Assumption 3 we know that f ∗ ( x, i ) ≥ γ . Therefore, f ( x, i ) · f ∗ ( x, i ) < 0 implies f ( x, i ) < 0 , so f ∗ ( x, i ) -f ( x, i ) > γ .

There are two possible ways to make f ( x, i ) smaller than f ∗ ( x, i ) : Making the neurons with positive coefficients smaller or making those with negative coefficients larger. The total difference must be at least γ , so there must exist one sign whose change in value is at least γ 2 . Formally, we decompose the

difference f ∗ ( x, i ) -f ( x, i ) into two terms:

Note that at least one of the two terms on the right-hand side must be at least γ 2 because their sum is at least γ . We first consider the first case, i.e., ∑ j : a ∗ j > 0 a ∗ j ( σ ( ⟨ v ∗ -i ( x ) , v ∗ j -v ∗ 0 ⟩ -b ∗ j ) -σ ( ⟨ v -i ( x ) , v j -v 0 ⟩ -b ∗ j ) ) ≥ γ 2 . The analysis of the second term is almost the same as the first one.

For notational simplicity, we define

(n.b. σ only appears in the second term) which implicitly depends on x , and and define the set

If ∆ j > 0 , we must have ⟨ v ∗ -i ( x ) , v ∗ j -v ∗ 0 ⟩ -b ∗ j > 0 because ReLU is non-negative. This means that the neurons corresponding to the words in S ( x ) must be activated.

Intuitively, S ( x ) is the set of words whose corresponding neuron in our student model has a smaller value than the ground-truth model, and ∆ j for a word j ∈ S ( x ) is the difference between these two neurons. In other words, if we want to make a mistake in classifying a sample x , ∆ j is the obstacle for the neuron to overcome, and the sum of all these obstacles must be at least the margin γ . Formally, we have the following lemma for S ( x ) and lower bound for ∆ j :

̸

Considering the TV distance, we get

where the first inequality comes from Assumption 5, i.e., S ( x ) ⊆ [ k ] ⊆ [ n ] \ B ( x, i ) .

We define ε p := ∣ ∣ ∣ log Z ∗ bulk ( x,i ) Z ∗ ( x,i ) -log Z bulk ( x,i ) Z ( x,i ) ∣ ∣ ∣ , which is the approximation error of the log bulk probability and an important quantity to trade off the two terms in the TV distance. There are two terms in (39). The first term corresponds to the bulk probability and the second term corresponds to the activated neurons. When ε p is large, the bulk probability has a large approximation error, resulting in a large TV distance. We lower bound this term in the following lemma:

Lemma 4.

The right-hand side of the above lemma is an increasing function of ε p , and this lower bound becomes large when the approximation error of the bulk probability is large.

In the other regime when ε p is small, the bulk partition function must be approximated accurately, i.e., Z ∗ bulk ( x,i ) Z ∗ ( x,i ) ≈ Z bulk ( x,i ) Z ( x,i ) . Thus,

From Assumption 7 we know that the bulk partition functions can be accurately approximated using the vectors v ∗ 0 and v 0 . Thus, exp( ⟨ v ∗ -i ( x ) ,v ∗ j ⟩ ) Z ∗ bulk ( x,i ) ≈ exp( ⟨ v ∗ -i ( x ) , v ∗ j -v ∗ 0 ⟩ ) . Therefore, this difference is approximately | exp( ⟨ v ∗ -i ( x ) , v ∗ j -v ∗ 0 ⟩ ) -exp( ⟨ v -i ( x ) , v j -v 0 ⟩ ) | ≈ exp( ⟨ v -i ( x ) , v j -v 0 ⟩ )( e ∆ j -1) . As ε p becomes smaller, the aforementioned approximations become more accurate, making the TV distance suffer from a term that is roughly proportional to ( e ∆ j -1) . This is formalized in the following lemma:

Lemma 5. For all j ∈ S ( x ) ,

Plugging Lemma 4 and 5 into (39) gives us

where the last inequality comes from Jensen's inequality.

Let g ( u, v ) := u · ( e v u -1) for u, v > 0 , then ∂g ∂u = e v u ( 1 -v u -e -v u ) ≤ 0 , so from | S ( x ) | ≤ k we know that

Using the fact that ∑ i σ ( y i ) = ∑ i max { y i , 0 } ≥ max { ∑ i y i , 0 } = σ ( ∑ i y i ) ,

where the last inequality follows from Lemma 3. Therefore,

Since we do not have any assumption for ε p , we take the minimum over all possible values of ε p and get the following bound:

Above we have derived a lower bound for the TV distance between p ∗ and p when the first term of (33) is at least γ 2 . When the second term of (33) is at least γ 2 , the proof is symmetric to that of the first term and we only need to exchange the role of parameters from the teacher model and student model and the related definitions. Therefore, with the second term being at least γ 2 , we have

where S ′ ( x ) := { j : a ∗ j < 0 , ∆ ′ j > 0 } and ∆ ′ j := ( ⟨ v -i ( x ) , v j -v 0 ⟩-b ∗ j ) -σ ( ⟨ v ∗ -i ( x ) , v ∗ j -v ∗ 0 ⟩-b ∗ j ) Since S ( x ) ∪ S ′ ( x ) ⊆ [ k ] , merging (53) and (55) finishes the proof of Lemma 1.

D Detailed proofs of auxilliary lemmas

̸

Proof. Assume by way of contradiction that S ( x ) = ∅ . Then by the definition of S ( x ) , we know that for all j such that a ∗ j > 0 ,

Therefore,

̸

which is a contradiction. Thus, S ( x ) = ∅ .

Proof. Note that for all j ∈ S ( x ) , ⟨ v ∗ -i ( x ) , v ∗ j -v ∗ 0 ⟩ -b ∗ j > σ ( ⟨ v -i ( x ) , v j -v 0 ⟩ -b ∗ j ) ≥ 0 , so σ ( ⟨ v ∗ -i ( x ) , v ∗ j -v ∗ 0 ⟩ -b ∗ j ) = ⟨ v ∗ -i ( x ) , v ∗ j -v ∗ 0 ⟩ -b ∗ j . Therefore,

which implies

Lemma 4.

Proof. If ε p = 0 , both sides of the inequality equal 0, and the inequality holds. When ε p > 0 ,

Lemma 5. For all j ∈ S ( x ) ,

Proof. We know that the word probabilities come from a log-linear model, so

If the probability p ( x i = j | x -i ) is reasonably large, i.e., its corresponding neuron is activated, then we can use that with ∆ j to bound this difference. In other words, when ⟨ v -i ( x ) , v j -v 0 ⟩ ≥ b ∗ j , we know that

Note that the last inequality (72) comes from Assumption 7 and the following property of ∆ j :

Then we bound the difference in probabilities: for all j ∈ S ( x ) ,

Note that

Moreover, when ∆ j -2 ε b -ε p > 0 , we have by (72) that

When ∆ j -2 ε b -ε p ≤ 0 , we have

Thus,

Therefore, when ⟨ v -i ( x ) , v j -v 0 ⟩ ≥ b ∗ j , we must have

In the second case where ⟨ v -i ( x ) , v j -v 0 ⟩ < b ∗ j , we have ⟨ v -i ( x ) , v j ⟩ < ⟨ v -i ( x ) , v 0 ⟩ + b ∗ j . The proof for the second case is very similar to that of the first case, with the main difference being we replace ⟨ v -i ( x ) , v j ⟩ in the first case by ⟨ v -i ( x ) , v 0 ⟩ + b ∗ j in the second case.

In the second case, by the definition of ∆ j we know that ∆ j = ( ⟨ v ∗ -i ( x ) , v ∗ j -v ∗ 0 ⟩ -b ∗ j ) . Similar to the first case, we can get

Therefore, for the difference in probabilities, if ∆ j -2 ε b -ε p > 0 , we have ( ⟨ v ∗ -i ( x ) , v ∗ j ⟩ -log Z ∗ ) > ( ⟨ v -i ( x ) , v 0 ⟩ + b ∗ j -log Z ) , so

If ∆ j -2 ε b -ε p ≤ 0 , the absolute value of probability difference can be trivially bounded below by 0. Thus,

Since

we finally get the same bound as the first case:

Merging the two cases, we know that for all j ∈ S ( x ) ,

This finishes the proof of Lemma 5.

E Experiment details

Language models. We use various versions of GPT-2 Radford et al. [2019] and OPT Zhang et al. [2022] with number of parameters ranging from 125M to 1.5B. We use the base, medium, large, xl version of GPT-2 which have #parameters 124M / 355M / 774M / 1.5B and hidden dimension 768 / 1024 / 1280 / 1600. For OPT, we use three versions, with #parameters 125M / 350M / 1.3B and hidden dimension 768 / 1024 / 2048. For all these models, the word probability is the softmax of the product of the penultimate layer representation and the dictionary. This is consistent with our theoretical model introduced in Section 2. The parameter settings and performances of these models are shown in Table 1.

Dataset. We use WikiText-2 Merity et al. [2016] as the text corpus. WikiText-2 has about 280k tokens, and we only use the first 1/4 of it for computational efficiency. The perplexities of the language models on this corpus are shown in Table 1.

F Existence of anchor vector is not trivial

In this section, we will empirically verify that the existence of the anchor vector in large-scale language models does not arise from the trivial reasons that we could think of.

Linear approximation is not accurate for original partition function. Table 2 shows the mean squared approximation error of the log bulk partition function for different models and different k . For the models whose log bulk partition functions can be well approximated linearly, their log original partition function cannot. This confirms the necessity of removing the words of interest in the partition function.

Word embeddings are not uniformly distributed on a sphere. Figure 2 shows the singular values of dictionary matrices from different models. For the four selected models, their word embeddings are close to being low-rank. Therefore, the word embeddings in these models are far from uniformly distributed on a sphere.

(a) GPT-2 Large

(b) GPT-2 XL

(c) OPT-125M

(d) OPT-1.3B

140

120

160

140

120

100

120

100

80

09

80

60

150

100

100

08

60

40

40

50

40

20

20

Z

0

0

0

0

0

00Z

400

600

800

1000

1200

0

200

400

600

800

1000

1200

1400

1600

0

100

200

300

400

500

600

700

800

Q

500

1000

1500

2000

Dictionary atoms are not uniformly distributed on sphere. Figure 3 shows the histogram of the ℓ 2 -norm of the dictionary atoms from different language models, and Figure 4 shows the distribution of the Cosine similarity between two random atoms . The norms of these atoms are somewhat bounded, but their cosine similarity is strongly biased towards the positive part, and the dictionary matrices are close to low rank, indicating that these vectors are far from uniformly distributed on a sphere. Instead, they may concentrate around a cone-shaped region.

(a) GPT-2

(b) GPT-2 Medium

(c) GPT-2 Large

(d) GPT-2 XL

(e) OPT-125M

(f) OPT-350M

(g) OPT-1.3B

17500

17500

25000

15000

15000

20000

0000

12500

10000

12500

10000

15000

15000

00SL

7500

10000

10000

2500

5000

2500

5000

5000

5000

0

2.5

3.0

3.5

4.0

4.5

5.0

5.5

6.0

6.5

0

2.0

2.5

3.0

3.5

4.0

4.5

5.0

5.5

1.0

15

2.0

2.5

3.0

3.5

4.0

4.5

5.0

0

1.0

1.5

2.0

2.5

3.0

3.5

35000

14000

25000

30000

12000

20000

25000

10000

20000

8000

15000

15000

6000

10000

10000

4000

5000

5000

2000

0

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

5.0

0

0.8

1.0

1.2

1.4

1.6

1.8

0

1.0

1.5

2.0

2.5

3.0

3.5

(Bulk) partition functions have large variations. Figure 5 shows the histogram of log partition function and log bulk partition function for the language models, where the bulk words are defined as all the words except the top 100, in terms of logit values. We can see from the figures that both the partition function and the bulk partition function vary a lot depending on the samples. Therefore, in recent large-scale language models, the partition functions are not stable across samples.

(a) GPT-2

(b) GPT-2 Medium

(c) GPT-2 Large

(d) GPT-2 XL

(e) OPT 125M

(f) OPT 350M

(g) OPT 1.3B

5000

5000

4000

4000

4000

3500

4000

000E

3000

000

2500

3000

2000

2000

2000

1500

2000

1000

1000

1000

1000

005

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0

0.2

0.1

0.0

0.1

0.2

0.3

0.4

0

0.20.1

0.0

0.1

0.20.3

0.4

0.5

4000

3500

4000

3500

000E

3000

2500

3000

2500

2000

2000

000

1500

1500

1000

1000

1000

500

500

0.2

0.0

0.2

0.4

0.6

0.8

0

0.2

0.0

0.2

0.4

0.6

0

0.4

0.2

0.0

0.2

0.4

(a) GPT-2 Large

(b) GPT-2 XL

(c) OPT-125M

(d) OPT-1.3B

8000

7000

original

bulk

0008

7000

original

bulk

12000

10000

original

bulk

10000

original

bulk

6000

6000

8000

5000

5000

8000

4000

4000

6000

6000

3000

000E

4000

4000

2000

2000

1000

1000

2000

2000

0

8

10

12

14

16

18

20

0

10

12

14

16

18

20

女

8

10

12

14

16

18

20

6

8

10

12

14

16

18

0Z

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper mainly addresses two conditions that enable the representation of pre-trained LLMs to be transferred effectively to the downstream tasks which are usually different from pre-trained objectives. \n\n1. The insensitivity of the downstream task of \"super-small\" probability words must be guaranteed for good downstream performance. Thus, a function like ReLU is needed as an additional structure of the student model to ignore “small” logits which are considered to have no influence on the downstream task. \n\n2. Using a shift-invariant softmax function can render the logit values meaningless in unseen data. Here, the need for structure in the representation space of the pre-trained model arises and leads to stabilization of the partition function with the “Anchor vector” hypothesis. \n\nThe anchor vector hypothesis has backed up with empirical verification. The authors demonstrate that as more words with top probabilities are excluded to compute the optimal anchor vector, the error of approximated bulk partition function decreases. This finding supports that the anchor vector actually exists and can be utilized to address the shift-invariance of the softmax function. Furthermore, with the anchor vector, the upper bound of the downstream task’s error rate can be made.\n\nSTRENGTHS:\n- The paper is well written with a clear flow. Starting from the problem setting, the properties of the learned representation that affect the performance of the downstream task are well stated and addressed with the anchor vector hypothesis.\n- The empirical verifications provide substantial evidence in support of the anchor vector hypothesis, solidifying its credibility and reinforcing its validity.\n\nWEAKNESSES:\n(1) The empirical verification of the “anchor vector” hypothesis lacks clear explanations and detailed information regarding the experimental setup, leading to confusion and a sense of insufficiency. One particular point that requires clarification is the choice of using auto-regressive models, specifically GPT-2 and OPT, for the verification.\n\nAs mentioned earlier in Line 138, the pre-training task of the language model is described as predicting a “label” given a sequence without “label”, which aligns with the mask infilling task associated with masked language models. However, in section 4.2, the verification process utilizes the auto-regressive models, GPT-2 and OPT, which have a pre-training task of predicting the next token of the given sequence. \n\nBased on the provided context, it appears that a single token which can be an intermediate token of the sequence is removed from the sequence (-i) rather than considering the last token of the sequence. This discrepancy raises questions about the alignment between the pretraining task and the verification task depicted in the provided context in section 4.2. \nAddressing these issues and providing a clear explanation of the rationale behind using auto-regressive models instead of masked language models for the verification process would enhance the overall understanding of the presented claim. \n\n(2) In Figure 1, the labels for the x and y axes should be provided.\n\nQUESTIONS:\n(1) Based on Figure 1 (a), it is evident that OPT-350M and GPT-2 Medium show minimal decrease in mean squared approximation error, while other three models exhibit a significant decrease in mean squared approximation error as k increases. What could be the underlying reasons for the different tendency? It cannot be attributed solely to the difference in model scale (number of parameters), since OPT-125M displays a considerable decrease in error. Without further analysis, the findings from two of five models (OPT-350M and GPT-2 Medium) present weak evidence to support the “anchor vector” hypothesis. This raises some concerns regarding the generalizability of the proposed hypothesis to pretrained LLMs as a whole. \n\n(2) Do you consider including a broader range of activation functions like swish [1] and gelu [2]? Does the proposed hypothesis still hold for these activation functions, other than ReLU, that do not completely ignore (treat as 0) small logit values?\n\n[1] https://arxiv.org/abs/1710.05941 \n[2] https://arxiv.org/abs/1606.08415 \n\n[1] Ramachandran, P., Zoph, B., & Le, Q. V. (2017). Swish: a self-gated activation function. arXiv preprint arXiv:1710.05941.\n[2] Hendrycks, D., & Gimpel, K. (2016). Gaussian Error Linear Units (GELUs). arXiv preprint arXiv:1606.08415.\n\nLIMITATIONS:\nLimitations are listed in section 6: Conclusions and future work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the relationship between language model pretraining and downstream classification tasks. Under certain assumptions, the author theoretically demonstrated that pre-trained models can guarantee performance on downstream tasks with the existence of proposed “anchor vectors”.\n\nSTRENGTHS:\n1. The research question of “why pre-trained models can help with downstream tasks” is quite important.\n2. The paper is generally quite clearly written, the question definition is clear and the mathematical settings are reasonable.\n3. For me, the method to connect minimizing the KL divergence during pre-training and the performance of downstream tasks is interesting and valuable.\n\nWEAKNESSES:\n1. I think there is a difference between “pretrained embedding” (e.g., word2vec) and “large language models” (e.g., BERT/GPT), also the large language model is not equal to “the last layer embeddings”. In this paper, the author takes the model fixed and takes the extracted representations for analysis, however, fine-tuning the whole model is actually a more common way for models such as BERT/GPT on downstream classification tasks.\n2. Some of the claims and assumptions may be incorrect. For example, “We can also manually select all the words that are irrelevant to the downstream task.” in Line 255 and “Assume there are at most $k$ vectors are relevant to the downstream task” in assumption 4. These only hold true for certain tasks like SST-2, which is not generalizable.\n3. There exists more works need to be discussed and compared. e.g.,\n[1] Visualizing and Understanding the Effectiveness of BERT (EMNLP2019)\n[2] Revealing the Dark Secrets of BERT (EMNLP2019)\n[3] On Mutual Information Maximization for Representation Learning (ICLR 2020)\n4. I see some findings from Appendix F (Existence of anchor vector is not trivial) in the supplementary material, however, most of the conclusions are already found from previous research. e.g.,\n[4] How Contextual are Contextualized Word Representations? Comparing the Geometry of BERT, ELMo, and GPT-2 Embeddings (EMNLP 2019)\n[5] Representation Degeneration Problem in Training Natural Language Generation Models (ICLR 2019)\n[6] Isotropy in the Contextual Embedding Space: Clusters and Manifolds (ICLR 2021)\n\nQUESTIONS:\nSee weaknesses.\n\nLIMITATIONS:\nThe downstream tasks discussed in this paper are somewhat limited, e.g. the sentiment classification. Some of the assumptions are also based on the importance of words that are related to the task itself, which makes the overall framework not general enough.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a sequence of assumptions along with their respective conclusions, which advance the objective of comprehending the intricate connection between the performance of pre-training and downstream tasks in language models.\nBy framing the prevailing language models as log-linear models, this paper initially presents a mathematical demonstration showcasing instances where pre-trained models encounter limitations in effectively transferring their knowledge to address downstream tasks.\nAdditionally, the paper enumerates several imperative prerequisites for successful knowledge transfer, such as the requirement for shift invariance. \nFurthermore, the authors introduce the \"anchor vector hypothesis,\" which serves as a crucial framework for elucidating the remarkable adaptability of language models to various tasks.\n\nSTRENGTHS:\n- Attempted to rigorously prove the effectiveness of pre-training language models with a mathematical framework.\n- The claims posited in the paper appear to be reasonable, albeit necessitating a potential round of further verification, provided that readers accept the underlying assumptions upon which the proofs are based.\n- The storyline of the paper is intuitive and persuasive.\n\nWEAKNESSES:\n- Given the non-negligible series of assumptions outlined in the paper, a legitimate concern arises regarding the practical implications associated with the claims put forth. \nFor instance, the authors commence the proof procedures by assuming that \"a downstream task depends on only a small set of words related to this task,\" which may not hold true for tasks that involve complex rounds of reasoning.\n- In addition to Section 4.2, it would be advantageous for the paper to include additional empirical experiments that substantiate the claims presented, thereby strengthening the overall argument.\n\nQUESTIONS:\nIf we examine the standard (masked) language modeling during the pre-training phase, as far as my understanding goes, $p^*(x_i|x_{-i})$ should correspond to an instance of the Dirac delta distribution. If this understanding is correct, I'm wondering the subsequent discussions in the paper, such as the one in Section 3.1 where scenarios are considered where $p^*$ is relatively small but possibly not zero, are still valid.\n\nLIMITATIONS:\nWhile the mathematical framework proposed by the authors is compelling, there is a valid concern regarding the explanatory power of the paper's contents in relation to the inner workings of current language models. This concern arises due to the series of (unrealistic) assumptions upon which the framework is built.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper explores how to connect pretraining performance with downstream task performance (i.e., binary classification). The theoretical analysis is based on token representations. The authors find the ``anchor vector'' in the representation space and bridge pretraining and downstream tasks performance based on it.\n\n---\n\nRebuttal response:\n\nThanks for the clarification! It would be better if these parts can be described clearly in the paper.\n\nAs for the responses to the two questions, I don't think they are convincing. As I mentioned, previous work used ``surprisal'' to measure the information amount brought by that word. “low probability words” could be vital and carry essential information in that sentence. But this is a theoretical analysis paper. So, I'm fine with this assumption. Therefore, I would only slightly raise my score.\n\nSTRENGTHS:\nThis paper discusses a vital problem in language model understanding from a theoretical view. The idea of the proofs is interesting.\n\nWEAKNESSES:\n1. There are many assumptions and some of them may not hold in practice. A1 could fail for contextualized language models (e.g., BERT, GPT-2) since one token could have very different representations with different contexts. A4 may fail since the prediction head of downstream tasks could be nonlinear. Besides, there are some claims that could be wrong (see Q1, Q2).\n\n\n2. The ``anchor vector'' hypothesis could fail. Even if the authors show that the MSE for the approximation is less than 1. However, it is not sufficient to prove the hypothesis empirically. Language models could have quite similar representations for all tokens, and they may be far from the zero point. More baselines should be included to support the hypothesis. For example, calculating the average/minimum MSE distance between two random words.\n\n3. The scope of this paper is not a little limited. Even if in the introduction and empirical section (4.2) popular language models (e.g., GPT-2) are included, the theoretical discussion only considers binary classification as the downstream task. Besides, not sure if I understand it correctly, the analysis of language models assumes that they are uncontextualized (e.g., GloVe and ELMo).\n\nQUESTIONS:\nQ1 (Line 204-206): Why does this claim hold? First of all, the vocabulary size of LM could be very large (e.g., for BLOOM, its vocabulary size is 250,000). A word with probability 1e-5 is not super-small. Second, words with small probability could be vital to the downstream tasks. Some people use ``surprisal'' to measure the information gained from that word. It means words with small probability could convey information.\n\nQ2 (Line 246-248): Does it mean that the frequent words are not that useful? Is the claim contradictory to the aforementioned one?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-18"}