SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Connecting Pre-trained Language Models and Downstream Tasks via Properties of Representations
Chenwei Wu
Duke University cwwu@cs.duke.edu
Holden Lee
Johns Hopkins University hlee283@jhu.edu
Rong Ge
rongge@cs.duke.edu
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.
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.
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},\dots,x_{L})$  where the entries come from a dictionary  $\{1,\ldots ,n\}$ , we assume the observation probability of  $x_{i}$  satisfies a log-linear model
$$
p ^ {*} (x _ {i} = j | x _ {- i}) \propto \exp (\langle v _ {- i} ^ {*} (x _ {- i}), v _ {j} ^ {*} \rangle),
$$
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})\propto \exp (\langle v_{-i}(x_{-i}),v_{j}\rangle)$ .
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^{*} \coloneqq (\langle v_{-i}^{*}(x_{-i}), v_{j}^{*}\rangle)_{j=1}^{n}$  (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 := (\langle v_{-i}(x_{-i}), v_{j} \rangle)_{j=1}^{n}$ . 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  $\mathbb{E}_x[D_{\mathrm{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)\approx 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.,  $\sum_{j: z_j^* \text{ 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.
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.
Notations. We use  $[n]$  to denote the set  $\{1,2,\ldots ,n\}$ . For an input sequence  $x = (x_{1},\dots,x_{L})$  we use  $x_{-i}$  to denote the input sequence without the  $i$ -th entry where  $i\in [L]$ , i.e.,  $x_{-i} := (x_{1},\dots,x_{i - 1},x_{i + 1},\dots,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}^{*} \in \mathbb{R}^{d}$  ( $j \in [n]$ ). We refer to each  $v_{j}^{*}$  as an atom and the matrix  $[v_{1}^{*}, \dots, v_{n}^{*}] \in \mathbb{R}^{d \times n}$  as the dictionary. At each position  $i$ , let  $x_{i}$  be the value of the word at that position; then  $x_{i} \in [n]$ . Assume that the probability of  $x_{i}$  given  $x_{-i}$  follows a log-linear model, i.e.,
$$
p ^ {*} \left(x _ {i} = j \mid x _ {- i}\right) \propto \exp \left(\left\langle v _ {- i} ^ {*} \left(x _ {- i}\right), v _ {j} ^ {*} \right\rangle\right), \tag {1}
$$
where  $v_{-i}^{*}(\cdot)$  is a function that encodes the remaining sequence  $x_{-i}$  into a vector in  $\mathbb{R}^d$ .
We also use  $z_{j}^{*}(x,i) \coloneqq \langle v_{-i}^{*}(x_{-i}),v_{j}^{*}\rangle$  to denote the  $j$ -th logit and  $Z^{*}(x,i) \coloneqq \sum_{j = 1}^{n}\exp (z_{j}^{*}(x,i))$  to denote the partition function, i.e., the normalization factor of equation (1). In other words,
$$
\forall j \in [ n ], \quad p ^ {*} (x _ {i} = j | x _ {- i}) = \frac {\exp \left(z _ {j} ^ {*} (x , i)\right)}{Z ^ {*} (x , i)} = \frac {\exp \left(\left\langle v _ {- i} ^ {*} (x _ {- i}) , v _ {j} ^ {*} \right\rangle\right)}{Z ^ {*} (x , i)}. \tag {2}
$$
Student model. We use a black-box neural network model whose penultimate layer outputs a  $d'$ -dimensional vector  $v_{-i}(x_{-i}) \in \mathbb{R}^{d'}$  and the last layer is a fully-connected layer with weight matrix  $[v_1, v_2, \dots, v_n] \in \mathbb{R}^{d' \times n}$  followed by the softmax function. In other words, the model output is
$$
p \left(x _ {i} = j \mid x _ {- i}\right) \propto \exp \left(\left\langle v _ {- i} \left(x _ {- i}\right), v _ {j} \right\rangle\right). \tag {3}
$$
Similar to the ground-truth model, we use  $z_{j}(x,i) \coloneqq \langle v_{-i}(x_{-i}),v_{j}\rangle$  to denote the  $j$ -th logit and  $Z(x,i) \coloneqq \sum_{j = 1}^{n}\exp (z_{j}(x,i))$  to denote the partition function of equation (3), so
$$
\forall j \in [ n ], \quad p \left(x _ {i} = j \mid x _ {- i}\right) = \frac {\exp \left(\left\langle v _ {- i} \left(x _ {- i}\right) , v _ {j}\right)\right)}{Z (x , i)}. \tag {4}
$$
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:
$$
\ell (v _ {- i}) = \mathbb {E} _ {x} [ - p ^ {*} (x _ {i} | x _ {- i}) \log p (x _ {i} | x _ {- i}) ] = \mathbb {E} _ {x} \left[ D _ {K L} \left(p ^ {*} (x _ {i} | x _ {- i}) | | p (x _ {i} | x _ {- i})\right) \right] + \mathbb {E} _ {x} \left[ H \left(p ^ {*} (x _ {i} | x _ {- i})\right) \right].
$$
Note that  $\mathbb{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  $\mathbb{E}_x[D_KL(p^* (x_i|x_{-i})||p(x_i|x_{-i}))]\leq \epsilon_{\mathrm{KL}}$  for some  $\epsilon_{\mathrm{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\in \mathrm{POS}$ ,  $f^{*}(x,i) > 0$  and for all  $x\in \mathrm{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) \coloneqq \langle v_{-i}^{*}(x_{-i}), u^{*} \rangle$  for some  $u^{*} \in \mathbb{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}^{*}\}_{j = 1}^{n}$ ). Assume  $u^{*}$  is a sparse combination of at most  $k$  vectors in  $\{v_{j}^{*}\}_{j = 1}^{n}$  and without loss of generality assume these  $k$  vectors are  $\{v_{1},\ldots ,v_{k}\}$ , i.e., there exist coefficients  $\{c_j^*\}_{j = 1}^k\in \mathbb{R}^k$  such that  $u^{*} = \sum_{j = 1}^{k}c_{j}^{*}v_{j}^{*}$ .
Note that the fixed representations  $\{v_{1},\ldots ,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  $\mu \in (0,1]$  to capture this difference. A smaller  $\mu$  indicates a larger difference between the two distributions, and  $\mu = 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  $\mu$ .
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  $\mu \in (0,1]$  such that
$$
\forall i, \forall x \in \mathrm {P O S} \cup \mathrm {N E G}, p _ {p r e} (x _ {- i}) \geq \mu \cdot p _ {D S} (x _ {- i}).
$$
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  $\gamma \in \mathbb{R}^{+}$  such that at any position  $i$ , if  $x \in \mathrm{POS}$ , then  $f^{*}(x,i) \geq \gamma$ , and if  $x \in \mathrm{NEG}$ , then  $f^{*}(x,i) \leq -\gamma$ , 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\}_{j=1}^k$ . 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 ground-truth 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) = \langle v_{-i}^{*}(x_{-i}),u^{*}\rangle = \sum_{t = 1}^{k}c_{t}^{*}\langle v_{-i}^{*}(x_{-i}),v_{t}^{*}\rangle$ . For  $t^{-} \in [k]$ , define  $p^{-} := p^{*}(x_{i} = t^{-}|x_{-i})$  and assume  $p^{-} \leq \frac{1}{2}$ . Then for all  $s \in \mathbb{R}^+$ , there exist functions  $v_{-i}$  and  $\{v_t\}_{t=1}^k$  such that  $D_{\mathrm{KL}}(p^{*}(x_{i}|x_{-i})||p(x_{i}|x_{-i})) \leq 2sp^{-}$  and  $f(x,i) := \sum_{t=1}^{k}c_{t}^{*}\langle v_{-i}(x_{-i}),v_{t}\rangle \leq f^{*}(x,i) - s \cdot 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  $2p^{-} \cdot \frac{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  $2p^{-} \cdot \frac{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  $(\sigma(x) := \max\{x, 0\})$  is the ReLU function:
$$
f ^ {*} (x, i) = \sum_ {j = 1} ^ {k} a _ {j} ^ {*} \sigma \left(z _ {j} ^ {*} (x, i) - b _ {j} ^ {*}\right) = \sum_ {j = 1} ^ {k} a _ {j} ^ {*} \sigma \left(\left\langle v _ {- i} ^ {*} (x _ {- i}), v _ {j} ^ {*} \right\rangle - b _ {j} ^ {*}\right). \tag {5}
$$
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) = \sum_{j=1}^{n} a_{j}\sigma(z_{j}^{*}(x,i) - b_{j})$ , there exist functions  $\{\hat{z}_{j}(x,i)\}_{j=1}^{n}$  such that for all  $x$  and  $i$ , we have  $\hat{p}(x_{i}|x_{-i}) = p^{*}(x_{i}|x_{-i})$  and  $\hat{f}(x,i) := \sum_{j=1}^{n} a_{j}^{*}\sigma(\hat{z}_{j}(x,i) - b_{j}^{*})$  is always equal to 0. In other words, the pre-training loss of the model  $\{\hat{z}_{j}(x,i)\}_{j=1}^{n}$  is the same as  $\{z_{j}(x,i)\}_{j=1}^{n}$ , but its logits are useless for the downstream task.