Why Do Pretrained Language Models Help in Downstream Tasks? An Analysis of Head and Prompt Tuning
Year: 2021

Abstract:
Pretrained language models have achieved state-of-the-art performance when adapted to a downstream NLP task. However, theoretical analysis of these models is scarce and challenging since the pretraining and downstream tasks can be very different. We propose an analysis framework that links the pretraining and downstream tasks with an underlying latent variable generative model of text — the downstream classifier must recover a function of the posterior distribution over the latent variables. We analyze head tuning (learning a classifier on top of the frozen pretrained model) and prompt tuning in this setting. The generative model in our analysis is either a Hidden Markov Model (HMM) or an HMM augmented with a latent memory component, motivated by long-term dependencies in natural language. We show that 1) under certain non-degeneracy conditions on the HMM, simple classification heads can solve the downstream task, 2) prompt tuning obtains downstream guarantees with weaker non-degeneracy conditions, and 3) our recovery guarantees for the memory-augmented HMM are stronger than for the vanilla HMM because task-relevant information is easier to recover from the long-term memory. Experiments on synthetically generated data from HMMs back our theoretical findings.

heading: 1 Introduction
text: Natural language processing (NLP) has been revolutionized by large-scale pretrained language models such as BERT [4] and GPT [25], which are adapted to a variety of downstream NLP tasks. Although a large body of empirical work seeks to understand the effectiveness of pretrained models [7, 5, 12, 35, 34, 11, 27, 15], theoretical understanding is scarce. Theoretically analyzing the relationship between the pretraining and downstream tasks is challenging because pretraining and downstream settings can greatly differ.
The key starting point for our analysis is to link the pretraining and downstream settings through an underlying generative model of the data. We model the data distribution as a latent variable model and the downstream task as a function of the latent variables. Assuming that pretraining on a large corpus allows us to learn the generative model, the conditional token probabilities predicted by the pretrained model carry information about the hidden variables. In downstream adaptation, we aim to recover this information to solve the downstream task.
Though full finetuning is the de facto empirical standard, analyzing it is challenging because it requires characterizing the weights of the pretrained model. In this paper, we focus on head tuning
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
and prompt tuning, which both freeze all pretrained parameters and allow us to treat the pretrained model as a black box. Head tuning [22] trains task-specific heads on top of the pretrained model outputs. Prompt tuning [31, 19, 9, 21] optimizes a task-specific “prompt” that is concatenated to the model input. Studying prompt tuning is particularly interesting since it can match the performance of full finetuning with less computation time [19, 9, 21].
Our work contrasts with prior theoretical work [28], which assumes that downstream labels are recoverable via a linear head applied to the conditional token probabilities, and analyze how errors in pretraining or model misspecification propagate downstream. We consider specific generative distributions for which we can prove these assumptions, showing that head and prompt tuning can recover the downstream labels.
Our analysis considers two data-generating distributions with increasing realism. First, we consider data generated from a Hidden Markov Model (HMM), where the downstream task is to learn a linear classifier on the posterior distribution over the hidden states (Section 3). We prove that, under strong non-degeneracy conditions on token emission probabilities, a linear head applied to a pretrained model G which outputs exact conditional token probabilities (Gipxq “ P rXi |x´is) can recover the downstream label (Theorem 3.3). Furthermore, we can prove better recovery guarantees with relaxed non-degeneracy assumptions (Assumption 3.1) by using continuous prompt tuning (Theorem 3.6), reflecting the strong empirical performance of prompt tuning [19, 9, 21]. Intuitively, prompt tuning conditions the latent variables so that nonessential information for the downstream task can be ignored during the tuning phase, making task-essential information easier to recover.
Second, we also strengthen our analysis by leveraging additional structure in the data. Motivated by long-range dependences in natural language, we analyze HMM variants with additional latent “memory” variables that can store long-term information more easily than vanilla HMMs (Section 4). Here, the downstream task is to learn a linear classifier on the posterior distribution of the memory variables. We show that, under weaker non-degeneracy conditions than the first setting, an attentionbased classification head can recover ground-truth downstream labels from pretrained model outputs (Theorem 4.3). Intuitively, our recovery guarantees improve because the classification head can focus on the persistent, task-essential information in the memory while ignoring other transient and nonessential aspects of the latent variables. As with the vanilla HMM, we analyze prompt tuning for relaxing the non-degeneracy conditions even further (Theorem 4.6).
In summary, we relate the pretraining and downstream tasks by assuming that the downstream task is to learn a classifier on the posterior distributions of the latent variables defined by an underlying generative model of text. Our theoretical contributions are: 1) in this setting we analyze an HMM generative model show that simple classification heads can recover the true downstream labels under certain non-degeneracy assumptions, 2) we prove that soft prompt tuning can relax the non-degeneracy assumptions needed for downstream recovery making it easier to extract task-specific information, and 3) our recovery guarantees are stronger for memory-augmented HMMs in comparison to the vanilla HMM when tuning an attention-based classfication head.
We empirically evaluate our theoretical results with language models pretrained on synthetically generated data from HMMs. We find that prompt tuning obtains good downstream performance when our non-degeneracy conditions are relaxed, whereas head tuning performs poorly. Furthermore, we show that head tuning obtains better downstream performance when data is generated from a memory-augmented HMM, compared to a vanilla HMM, as is predicted by our theory.

heading: 1.1 Related works
text: The black box nature of BERT and related models has inspired a variety of empirical works which seek to understand them. Probing papers study whether a pretrained model computes various types of structured information (e.g., syntactic [35, 11]) by evaluating the performance of simple classifiers, or probes, on the representations [7, 12, 34, 27, 15]. Other papers ablate various aspects of pretraining, such as changing the masking scheme [14, 20, 40] or permuting the word order [32].
In comparison, theoretical analysis of pretrained language models is limited. Besides [28], which we discussed in Section 1, Zhang and Hashimoto [40] analyze using a linear classifier to approximately recover the latent variable in a Gaussian graphical model with sparse dependencies between observed variables. However, their analysis and setting are focused towards understanding syntactic dependencies between tokens, whereas we directly model and analyze downstream performance.
Prompt-based tuning [31, 19, 9, 21, 13, 6, 41, 2, 23], which has improved empirical downstream performance for lightweight adaptation methods beyond head tuning to approach full finetuning, is an important focus of our theoretical analysis. Shin et al. [31] employ task-specific prompts that are optimized over the discrete token space. Schick and Schütze [29, 30] reformulate natural language tasks as cloze-style phrases to enable few-shot learning. Subsequent methods [19, 9, 21] optimize “soft” prompts, or continuous embedding vectors. Lester et al. [19] employ soft prompts on pretrained large-scale T5 [26] models and show that as the model size increases, prompt tuning performance can eventually match finetuning. Hambardzumyan et al. [9] applies a variant of soft prompt tuning to MLM models. Li and Liang [21] propose prefix tuning, which prepends a trainable prefix embedding sequence to all layers of the transformer.
More broadly, Lee et al. [18] analyze reconstruction-based self-supervised learning methods in a general setting and show that under certain conditional independence assumptions, predicting one observed variable from another allows recovery of the latent with a linear head. Other theoretical works analyzing self-supervised or constrastive learning include [1, 10, 36, 38, 37], but they do not directly relate to our particular setting.

heading: 2 Formulations and notations
text: We analyze models pretrained on masked language modeling (MLM) objectives. Let X denote a finite vocabulary of input tokens, X ˚ the set of variable-length sequences of tokens, and X “ pX1, . . . , XT q P X ˚ a random sequence of T tokens. Let ∆|X | denote the space of probability distributions over tokens.
Pretraining and downstream task. Let Gpxq “ pG1pxq, G2pxq, . . .q denote the masked language model which predicts a probability vector for each timestep in the input x. Our theoretical abstraction is that Gi perfectly computes the distribution of Xi, the i-th token, conditioned on all other tokens: Gipxq “ P rXi|X´i “ x´is. Here P rXi |X´i “ x´is P ∆|X | is a probability vector. In particular, Gipxq does not depend on xi. The downstream task involves labeled examples px, F ‹pxqq P X ˚ˆY , where F ‹ : X ˚ Ñ Y provides ground-truth downstream labels and Y is a discrete set of labels for classification.
Head and prompt tuning. Head tuning trains a classification head f on top of fixed model outputs, resulting in the classifier F pxq “ 1pfpGpxqq ě 0q. We expect f to be a simple function such as a linear or one layer attention model. We also analyze variants where f also takes the tokens x or embeddings of x as input, which provides additional information. Soft prompt tuning requires viewing the pretrained model G as a function of the token embeddings; we refer to this model by G. Letting epxq “ epx1q, . . . , epxtq denote the token embeddings, we have Gpepxqq “ Gpxq. Soft prompt tuning concatenates a trainable prompt u so that the model output isGppu, epxqq. We consider simultaneously training the prompt parameter u and a classification head to fit the downstream task.
Notations. Let ∆d denote the space of d-dimensional probability vectors. We work with discrete random variables V taking values in a finite set V . We use P rV s P ∆|V| to denote the distribution of V and P rU |V “ vs P R|U | the conditional distribution of U given V “ v. PrpV “ vq P r0, 1s will denote the probability that V takes values v. We also let P rU “ u |V s P R|V| denote the vector with entries PrpU “ u |V “ vq. P rU |V s P R|U |ˆ|V| will describe the matrix with entries P rU |V su,v “ PrpU “ u |V “ vq. For a sequence v “ pv1, . . . , vtq, we use the notation vi:j for i ď j to denote pvi, . . . , vjq, and v´i to denote pv1:i´1, vi`1:tq. We let 1 denote the indicator function. For set V , we let V˚ “ V1YV2Y¨ ¨ ¨ denote variable-length sequences of elements of V . Let d denote elementwise product. Let 1d,0d denote the d-dimensional all-1’s and all-0’s vector. We omit the subscript if the dimension is clear from context. For two vectors a, b P Rd, we let a{b denote their element-wise division. We use supppaq to denote the set of indices where vector a is non-zero.

heading: 3 Analysis for Hidden Markov Models
text: Defining a relation between pretraining and downstream tasks is the foremost challenge for analysis. We propose to link the two via latent variable generative assumptions on the input distribution. We
model the downstream task as a function of the posterior distribution of the latent variables. Towards a first result, this section studies the case where inputs are generated by HMMs (see Figure 1 (left)), which have been well-studied in the context of language and speech processing (see e.g. [24, 17, 3]).
Data distribution. Let H denote the hidden state space of the HMM. We use H “ pH0, H1, . . . ,HT q P H˚ to denote the sequence of hidden states. For all timesteps i ą 0, the transition probabilities are time-invariant, i.e. P rHi |Hi´1s “ A for A P R|H|ˆ|H|. For each timestep i ě 1, tokens Xi are emitted following some time-invariant probability: P rXi |His “W for W P R|X |ˆ|H|. The joint probability of X,H is
PrpX,H “ x, h |T “ tq “ PrpH0 “ h0q t ź
i“1 PrpHi “ hi |Hi´1 “ hi´1qPrpXi “ xi |Hi “ hiq.
Downstream tasks. We assume that H0 has the meaningful information for the downstream task, which is a binary classification task where the ground-truth labeling F ‹ is assumed to be a linear classifier on the posterior P rH0 |X1:T “ xs:
F ‹pxq “ 1pµJP rH0 |X1:T “ xs ě 0q (3.1)
for µ P R|H|. Our results are easily extended to the multiclass setting. We consider tuning a linear head for the downstream classifier, which formally computes 1pbJG1pxq ě 0q for b P R|X |. The following non-degeneracy condition is crucial for our recovery result in this setting. Assumption 3.1 (Non-degeneracy, vanilla HMM). The token emission probability matrix W has linearly independent columns.
We also require the following regularity conditions on H0 and the state transitions. Assumption 3.2 (Regularity). The Markov chain H0, H1, . . . is ergodic, and P rH0s has full support.
We show that if W has linearly independent columns, a linear head fits downstream labels. Theorem 3.3. Assume that non-degeneracy (Assumption 3.1) and regularity (Assumption 3.2) hold. Then any downstream task F ‹pxq of the form (3.1) can be computed by a linear head onG applied to a shifted sequence. That is, there exists linear head weights b P R|X | such that for all x P supppP rXsq,
F ‹pxq “ 1pbJG1px1q ě 0q
where x1 “ p∅, x1:tq is the concatenation of a special token ∅ with x.1
The key for the proof is to leverage the following general statement about random variables U, V, Z such that U K V |Z, which decomposes the expression for P rU |V s.
1We note that G1px1q does not depend on x11 and therefore x11 can be any token.
Proposition 3.4. Let U, V, Z be random variables such that U K V |Z. Then for any v, P rU |V “ vs “ P rU |Zs ¨ P rZ |V “ vs. Thus, if P rU |Zs has a left inverse pP rU |Zsq:, then P rZ |V “ vs “ pP rU |Zsq:P rU |V “ vs.
By the conditional independence structure of the HMM, Proposition 3.4 immediately implies
G1px1q “WP rH1|X2:T`1 “ xs ùñ P rH1|X2:T`1 “ xs “W :G1px1q where W : is the left inverse for W , guaranteed to exist by Assumption 3.1. This lets us recover P rH1|X2:T`1 “ xs by applying a linear function to G1px1q. Additional linear functions will be sufficient to obtain µJP rH0|X1:T “ xs from P rH1|X2:T`1 “ xs. We provide the full proof in Section B.
Proposition 3.4 is reminiscent of the arguments of [18], which leverages the independence structure in the same way. Subsequent sections will require more complicated analyses and recovery procedures.
A drawback of Theorem 3.3 is that it relies heavily on assuming W has full column rank, which implies the necessary condition that |H| ď |X |. Without this assumption, it is unclear how to recover P rH0 |X1:T “ xs from Gpxq alone. However, in realistic settings we would expect |H| ą |X |, as increasing the size of the hidden state space improves language modeling capabilities of HMMs [3].

heading: 3.1 Relaxed non-degeneracy assumptions via prompt tuning
text: In this section, we study applying soft, or continuous, prompt tuning [19, 9] to the setting above. We show that by using soft prompt tuning, we can recover F ‹ using a linear head on G for HMMs where the non-degeneracy assumptions on W are relaxed. Our analysis provides insight into the empirical successes of prompt-tuning: intuitively, prompt tuning enables better recovery of the downstream task by conditioning the output of G to only contain task-specific information.
Soft prompt tuning trains task-specific embedding vectors, but analyzing how the model processes embedding vectors is challenging because it requires opening up the black box of the pretrained model. Thus, we require additional abstractions about how the pretrained model processes the embedding vectors. We will extend the mask language model G to a model G that maps a sequence of embeddings e1, . . . , et to conditional probabilities G1pxq, . . . , Gtpxq as follows. We observe that each token z in the vocabulary X naturally corresponds to a |H|-dimensional vector: the z-th row of the emission probability matrix W , or equivalently, P rXi “ z |His. We denote this embedding by epzq and call the family of embeddings tepzq : z P X u proper embeddings. A fundamental property of HMMs is that the conditional probability P rXi |X´i “ x´is only depends on x1, . . . , xt through their embeddings epxq “ pepx1q, . . . , epxtqq. In other words, there exists a function Gi such that
Gipx1, . . . , xtq “ Gipepx1q, . . . , epxtqq In particular, we let Gi compute the standard message passing algorithm [16] that computes the conditional probability of HMMs. This ensures thatGi is well defined on all sequences of nonnegative vectors in r0, 1s|H|, beyond sequences of proper embeddings.We assume that pretraining produces this Gi, which we treat as a blackbox for prompt tuning.
In particular, for prompt tuning we can consider the case where we pass an arbitrary nonnegative vector u P r0, 1s|H| to G in the first argument and proper embeddings at positions i ą 1. We can interpret u as the embedding of a fake token rz. Concretely, consider adding a new token rz to the vocabulary X , and changing the emission probability at position 1 to satisfy P rX1 “ rz |H1s “ u and for all z ‰ rz, P rX1 “ z |H1s9p1 ´ uq d epzq. Then Gipu, epx1q, . . . , epxtqq precisely computes the conditional probability P rXi |X´i “ prz, x1, . . . , xtq´is under the modified HMM. We refer the readers to Section C for the formal definition of Gi and formal proofs of the interpretation above.
We consider a downstream training algorithm which trains the prompt tuning parameter u described above and a linear classification head. Letting u denote the trainable prompt parameter and b P R|X | the trainable linear head weights, the model uses the embedding sequence
pepxq fi pu, ep∅q, epx1q, . . . , epxtqq (3.2) and outputs the prediction F pxq “ 1pbJG2ppepxqq ě 0q. We can provide recovery guarantees for this model if the ground-truth classifier weights µ (defined in (3.1)) and columns of the HMM transition matrix A satisfy the following relaxation of the requirement in Theorem 3.3 that W is nondegenerate.
Assumption 3.5 (Relaxed non-degeneracy condition). There exists a set of essential hidden states H‹ Ď H, so that the columns of W corresponding to H‹, tW:,huhPH‹ , are linearly independent. Furthermore, H‹ covers all meaningful information for the downstream tasks: supppµq Ď H‹.

heading: In addition, a last technical requirement on H‹ is as follows: there exists a set B Ď H such that
text: H‹ “ YhPBsupppA:,hq. In other words, H‹ must be the set of all states reachable by starting from some state in B and transitioning one step in the hidden Markov chain.
Compared to Assumption 3.1, which required that all columns of W are linearly independent, Assumption 3.5 only requires linear independence on a subset H‹ of essential states. In the setting where |H| ą |X |, the condition for Theorem 3.3 can never hold. On the other hand, Assumption 3.5 could still hold, for example, if |supppµq| ă |X | and the set of columns of W corresponding to hidden states in supppµq is linearly independent. The last technical requirement in Assumption 3.5 is also required, which could be satisfied if columns of A are sparse. The following theorem shows that when Assumption 3.5 holds, we can recover F ‹ using soft prompt tuning with a linear head.
Theorem 3.6. In the above setting, assume that Assumptions 3.2 and 3.5 hold. Then F ‹ can be computed using soft prompt tuning with a linear head on G. Concretely, there is a continuous prompt parameter u P R|H| and weight vector b P R|X |, such that for all x P supppP rXsq,
F ‹pxq “ 1pbJG2ppepxqq ě 0q
where pe prepends u to the input embedding sequence, as defined in (3.2).
Theorem 3.6 provides a stronger recovery result than Theorem 3.3, which only used a linear head. This is also reflected in our synthetic experiments (Section 5), and prior work which shows that variants of prompt tuning can perform much better than only training the last few layers of the model [21]. Our theory suggests that prompt tuning could help by conditioning the hidden variables to remove nonessential information for the task from the output of G. This makes task-essential information easier to recover.
The key proof intuition is that although recovering P rH0 |X1:T “ xs is impossible without strong non-degeneracy conditions (Assumption 3.1), we can aim to recover P rH0 |X1:T “ xs on the subset of essential states H‹ defined in Assumption 3.5, which suffices for computing µJP rH0 |X1:T “ xs, since H‹ Ě supppµq. To recover P rH0 |X1:T “ xs on H‹, we observe in Lemma C.2 that prepending the prompt u is equivalent to introducing a modified random sequence pX and fake token rz which influences the posterior of H2 as follows:
G2ppepxqq “ rxWDpP rH2 | pX1 “ rzs d P rH0 |X1:T “ xsq (3.3)
for invertible diagonal matrix D and positive scalar rx. We choose u so P rH2 | pX1 “ rzs d P rH0 |X1:T “ xs is supported only on H‹. As corresponding columns of W are linearly independent (Assumption 3.5), we recover PrpH0 “ h |X1:T “ xq for h P H‹ via a linear function of G2ppepxqq. This suffices for computing µJP rH0 |X1:T “ xs. For more details, see Section C.

heading: 4 Analysis for memory-augmented Hidden Markov Models
text: We study a memory-augmented HMM which explicitly disentangles the evolution of hidden states from a persistent “memory” variable. Inspired by natural sentences, this model is intended to better capture the distinction between syntax, which constantly evolves, and semantics, which changes less. This additional structure in the generative model allows us to strengthen our results by relaxing the non-degeneracy conditions on W , the token emission probabilities. Thus, both head and prompt tuning are more powerful in this setting compared to Section 3 and can recover the downstream label with weaker non-degeneracy assumptions on W . In Section 4.2, we show that soft prompt tuning also provides an advantage over head tuning alone.
Data distribution. The memory-augmented HMM, depicted in Figure 2, can be viewed as a generative variant of memory networks [39, 33] and is closely related to Hidden Topic Markov Models [8]. There are two sets of latent variables in the memory-augmented HMM: a Markov chain on hidden states H0, H1, . . ., meant to model the evolution of syntax, and a persistent “memory” M “ pM1, . . . ,MN q with N total cells, where each Mi takes values in a finite set M. The full joint
probability is as follows: PrpX,H,M “ x, h,m|T “ tq “
PrpM “ mqPrpH0 “ h0q t ź
i“1 PrpHi “ hi|Hi´1 “ hi´1qPrpXi “ xi|M “ m,Hi “ hiq
The hidden state is modified to explicitly consist of a disentangled cell index J P rN s and syntax state S P S, such that Hi “ pJi, Siq and H “ rN s ˆ S. To sample the token at timestep i given the hidden state Hi “ pJi, Siq, we first use Ji to index the memory M , obtaining the random variable MJi . Xi is then sampled according to some time-invariant probability depending on MJi , Ji, Si: P rXi |M “ m,Hi “ pj, sqs “ P rXi |MJi “ mj , Hi “ pj, sqs “W:,pmj ,j,sq Here W P R|X |ˆ|M||H| stores the emission probabilities for each choice of memory cell value and hidden state. Note that in particular, the conditional probabilities for Xi only depend on a single memory cell for each timestep. We also note that memory-augmented HMMs can be viewed as vanilla HMMs with structured transitions because pH0,Mq, pH1,Mq, . . . can be viewed as a Markov chain where the memory component does not change. Example 4.1 (Generating natural sentence with memory-augmented HMM). We consider how this model may generate the sentence “The cow in the pasture rolled on the grass’ happily.” M1 could store the subject (“cow”), M2 the location (“pasture”), M3 the sentiment (“happily”), and Si could determine part-of-speech. For timesteps where “cow” and “rolled” are emitted Ji “ 1 because we emit information related to the sentence subject. Timesteps for “pasture” and “grass” have Ji “ 2.
Downstream tasks. We consider downstream tasks where ground-truth labels are obtained via a linear classifier on the posterior distribution of a particular memory cell j‹ P rN s: F ‹pxq “ 1pµJP rMj‹ |X1:T “ xs ě 0q, where µ P R|M|. Intuitively, this formulation models downstream tasks which depend on a particular aspect of the semantics but not on syntax (e.g. in the setting of Example 4.1, if j‹ “ 3, the task is sentiment analysis).

heading: 4.1 Tuning attention head for recovering ground-truth downstream labels
text: To recover the downstream labeling, we require an attention-based classification head, which is a function of both the input embeddings and outputs of G. Formally, let q P R|H|`1 denote a query parameter and β1, . . . , βt P R|H|`1 denote trainable position embeddings. Given pretrained model outputs Gipxq and trainable token embeddings epxiq, the attention head Attnp¨q applies key and value functions K,V to compute the output as follows:
I fi arg max i
tqJpKpGipxqq ` βiqu (4.1)
AttnppGipxq, epxiqqti“1q fi 1 |I| ÿ
iPI V pGipxq, epxiqq (4.2)
where arg max refers to the set of indices achieving the maximum in (4.1). We note that standard attention heads in practice rely on the softmax function, but the expression based on arg max above captures the limiting behavior as }q}2 Ñ8. We consider linear key functions given by KpGipxqq “ ΘpKqGipxq. The value function V : R|X | ˆ R|M||H| Ñ R uses parameters ΘpV q P R|M||H|ˆ|X | and b P R|M||H| and computes V pGipxq, epxiqq “ bJppΘpV qGipxqq d epxiqq. Because our generative model disentangles H and M , we can relax the non-degeneracy assumption on the token emission probabilities W , compared to Theorem 3.3. The relaxed assumption only requires the columns tW:,pm,hqumPM,hPH‹ to be linearly independent in a subset H‹ of “recoverable” hidden states, whereas Assumption 3.1 required all columns to be linearly independent. Assumption 4.2 (Existence of “recoverable” hidden states). There exists a set of recoverable hidden states H‹ “ tj‹u ˆ S‹, such that the collection of token emission probabilities from M ˆ H‹, tW:,pm,hqumPM,hPH‹ , is a linearly independent set of vectors. Furthermore, the span of these vectors must be disjoint from the span of token emission probabilities from Mˆ pHzH‹q: spanptW:,pm,hqumPM,hPH‹q X spanptW:,pm,h1qumPM,hPHzH‹q “ t0|X |u.
Note that the non-degeneracy condition of Theorem 3.3 would require tW:,pm,hqumPM,hPH to be linearly independent, whereas Assumption 4.2 only requires linear independence for h P H‹. The second condition states that H‹ and HzH‹ are distinguishable by the token emission probabilities. We explain Assumption 4.2 in the setting of Example 4.1. For natural language, there might be choices of h “ pji, siq for which the set tW:,pm,hqumPM of token emission probabilities is fundamentally not very diverse, and therefore not linearly independent. For example, if the syntax si indicates “article”, i.e. words such as “a”, “an”, and “the”, the token emission probabilities would carry little information about Mji because the choice of article does not depend much on semantics, so columns corresponding to si “ “article” would not be linearly independent, violating Assumption 3.1. However, Assumption 4.2 allows us to avoid this issue by placing such h in HzH‹, a set of hidden states which we can ignore, and only including hidden states which carry a lot of information about M in H‹. In Example 4.1, when Ji “ 2 (location), Si “ “noun”, the position i should convey a lot about the location (in this case, “pasture”), so it is more reasonable to assume that tW:,m,humPM is linearly independent for this hidden state.
Thus, our aim is to focus on recovering information for the downstream task from positions i where Hi P H‹. Formally, we define the following set of input sequences containing positions i where the posterior of Hi given x´i concentrates on H‹:
R fi tpx1, . . . , xtq P supppP rXsq : Di with supppP rHi |X´i “ x´isq Ď H‹u (4.3)
The following theorem shows that under Assumption 4.2, we can recover F ‹ using the attention head described above, if x P R is nonempty. Note that R is nonempty if the posterior of Hi concentrates on H‹ for some i. For natural language, it is realistic to assume this can occur because syntactic aspects of a sentence are typically low-entropy when the full sentence is observed. Theorem 4.3. Assume that non-degeneracy (Assumption 4.2) and regularity (Assumption 3.2) hold. Define R as in (4.3). Then there exist an attention head on Gpxq and token embeddings epxiq such that the following holds for any x P R:
F ‹pxq “ 1pAttnppGipxq, epxiqqti“1q ě 0q where the function Attn is in the form described in (4.2).
The idea is to use the attention mechanism to attend to positions i where supppP rHi |X´i “ x´isq Ď H‹. The intuition of Assumption 4.2 is that such positions are more informative for recovering the latent posteriors; indeed, from the outputs Gipxq at such i, the value function in the attention will be able to recover P rMj‹ |X1:T “ xs. A full proof is provided in Section D.1.

heading: 4.2 Guarantees for prompt-tuning
text: Though the generative modeling assumptions in this section already allowed relaxed non-degeneracy assumptions, applying soft prompt tuning allows us to relax them even further. For simplicity, we consider the setting where there is a single memory cell, so M P M, and the downstream task is a linear classifier on the posterior of the memory: F ‹pxq “ 1pµJP rM |X1:T “ xs ě 0q. This
simplified setting doesn’t require the explicit disentanglement between Ji and Si in Hi. We analyze continuous prompt-tuning in a setting where the pretrained model G follows the same abstraction as in Section 3.1. We modify the model to take |M||H|-dimensional vectors, so the proper embedding for token z is given by epzq “ P rXi “ z|M,His “ WJz,:. In Section D.3, we describe the formal construction and interpretation of G in the more general setting with more memories.
Letting u P R|M||H| denote the trainable prompt parameter, we define the input embeddings pepxq fi pu, epx1q, . . . , epxtqq (4.4)
The downstream model applies an attention head to the output of G: F pxq “ 1pAttnppGippepxqq, peipxqqt`1i“1q ě 0q, where Attn is defined in (4.2). An additional stationarity assumption on P rH0s will simplify the recovery procedure (though it can be removed). Assumption 4.4 (Stationarity). Assumption 3.2 holds on the Markov chainH0, H1, . . .. Furthermore, P rH0s is the stationary distribution: P rH0s “ AP rH0s, where A is the transition matrix.
As before, we assume sparsity of µ and some non-degeneracy of W , though the assumption is more relaxed and easier to state compared to the vanilla HMM setting. Assumption 4.5 (Relaxed version of Assumption 4.2). Let M‹ fi supppµq denote the set of non-zero coordinates in µ. There exists a set of recoverable hidden states H‹, such that the collection of token emission probabilities from M‹ ˆH‹, tW:,pm,hqumPM‹,hPH‹ , is linearly independent. Furthermore, the span of these vectors must be disjoint from the span of token emission probabilities from M‹ ˆ pHzH‹q: spanptW:,pm,hqumPM‹,hPH‹q X spanptW:,pm,h1qumPM‹,hPHzH‹q “ t0|X |u.
We note that Assumption 4.5, and Assumption D.5 for multiple memories, are relaxations of Assumption 4.2, as they only consider memory values in supppµq, whereas Assumption 4.2 considers all m PM. An additional advantage of the memory-augmented HMM is that Assumption 4.2 is simpler than Assumption 3.1 and does not require any conditions on the transition matrix A. We now state our result for recovering F ‹ with soft prompt tuning and an attention head. Theorem 4.6. In the setting above, suppose that non-degeneracy Assumption 4.5 and stationarity Assumption 4.4 hold. Then there exists a prompt u and attention head on Gppepxqq and the token embeddings which can compute the ground-truth F ‹pxq for any x P R, defined in (4.3):
F ‹pxq “ 1pAttnppGippepxqq, peipxqqt`1i“1q ě 0q where pe is the embedding in (4.4) and Attn is defined in (4.2).
The intuition for this proof is similar to Theorem 3.6: the soft prompt conditions the memory M to concentrate on supppµq. As a result, all irrelevant information to the task is removed from Gippepxqq, making it easier to recover the task-specific information about the posterior of M . A more general theorem statement for the multiple memories setting, and the full proof, is provided in Section D.3

heading: 5 Simulations
text: We empirically evaluate our theoretical results by pretraining a BERT-like masked language model (MLM) [4] on synthetic data generated by an HMM. Our goal is to verify key implications of our theory in a more realistic setting where some assumptions, such as that G outputs exact conditional probabilities, may not hold. First, we compare head and prompt tuning and show that prompt tuning improves downstream performance, especially when the recovery problem is degenerate. Second, we compare the effect of changing the data distribution from vanilla HMMs to memory-augmented HMMs on head tuning with an attention layer. We find that the downstream performance improves when the data has a long-term memory component. These observations support our theory.
Pretraining data and downstream task. We generate pretraining data from an HMM with randomly generated transition matrix, emission probabilities, and start distributions. In all experiments, the HMMs have 10 vocabulary symbols, while the hidden state size varies. The downstream task uses input sequences X1:T of length 129, where the first token X1 “ [MASK]. We consider binary classifcation where labels are generated using linear functions of the analytically-computed posteriors in the HMMs. In all experiments, the ground truth linear weight is sparse with 6 nonzero entries at uniformly random locations with Gaussian values. More details are in Appendix E.
Head vs. prompt tuning. We compare head and prompt tuning as the hidden state size of the HMM varies. The downstream label is computed via µJP rH1 |X´1 “ x´1s, where µ is a random ground-truth linear weight. Head tuning learns a linear head on top of the softmax probabilities predicted by the pretrained model for filling in the first [MASK] token. Prompt tuning uses the same setup but also optimizes a length 20 continuous embedding prepended to the input sequence.
Figure 3 (left) shows that prompt tuning improves downstream performance substantially across all hidden state sizes ({4,8,10,15,25,30}). Prompt tuning improves especially when the hidden state size increases beyond the vocabulary size, which makes the recovery problem degenerate. Thus, as suggested by Theorem 3.6, prompt tuning helps relax the non-degeneracy conditions.
Memory-augmented HMMs. We investigate the effect of augmenting the data-generating HMM with a long-term memory. We consider the single memory case with |H| “ 4 and varying memory sizes |M| P t2, 3, 5, 7u. The downstream label is generated by computing µJP rM |X´1 “ x´1s, where µ denotes the ground-truth weights. Viewing the memory HMM as a HMM where the component on M never changes, we can compare against the vanilla HMMs from the previous setting. For the memory-augmented HMM, we use head tuning with a single-cell attention layer on the entire sequence of softmax probability outputs. For the vanilla HMM in the comparison, we use a linear head on the output at the first position, as an attention head would perform worse since the downstream task depends only on H1 and not any other timesteps.
Figure 3 (right) verifies that head tuning recovers the downstream task better when there is more structure in the data, as predicted by Theorem 4.3. Head tuning achieves near 100% downstream accuracy on all hidden state sizes.

heading: 6 Conclusion
text: We analyze how pretraining on generic language modeling tasks can improve performance on diverse downstream tasks. In our analysis framework, the downstream task requires predicting properties of the posterior distribution over latent variables in an underlying generative model. When the generative model is a standard HMM, downstream recovery is possible with a simple classification head under strong non-degeneracy assumptions. We also show that we can relax the non-degeneracy conditions by changing the generative model to a memory-augmented HMM or using prompt tuning. The distributions studied here are meant to provide a first-cut result – we also expect similar theorems to hold for other generative models, which we leave as an interesting direction for future work.

heading: Acknowledgements
text: We thank Percy Liang, Tianyi Zhang, and Nelson Liu for helpful discussions.
Funding statement
CW was supported by a NSF Graduate Research Fellowship. SMX was supported by a NDSEG Fellowship. TM acknowledges support of Google Faculty Award, NSF IIS 2045685, and JD.com. Additional revenue: CW received an honorarium for a talk at G-Research.

Review summaries:
[UPDATE: Thanks to the authors for engaging with the comments in the review.  I hope to see them fully addressed in the revision and look forward to seeing the next version of the paper!]  The paper investigates how pretrained representations might be expected to help prediction on a downstream task.  In some sense, they should "obviously" help: the downstream predictor can learn to exploit any mutual information between its target output and the pretrained representations.  That's the standard motivation for multitask learning.  However, the paper exhibits specific settings in which the target output can be *exactly* predicted from the pretrained representations, using only a *simple* probe such as a linear classifier.    The paper is primarily theoretical, although a few experiments are included by way of illustration.  It is not clear what the experiments are intended to verify.  To make the analysis possible, the pretrained representations in the paper are not underlying representations such as word embeddings.  Rather, they are identifiable given the pretraining data, which makes it clear what "successful pretraining" means.  For example, the representation of word x_i in a sentence is the conditional distribution of x_i given all other words in the sentence, as predicted by a pretrained cloze language model.  The analysis assumes (though the experiments do not check) that the predicted conditional distribution actually matches the true conditional distribution.  This should be the case in the limit of infinite training data, provided that the cloze language model has adequate capacity and can be trained to a global optimum of log-loss.  It is shown that if the vocabulary is large enough, or under other assumptions, then this representation of x_1 is enough to recover the posterior distributions over h_1 and h_0, which are assumed to be sufficient for the downstream task.  Caveat: In the analysis, the observed sentences are assumed to be generated by a latent-variable model.  One would then expect the downstream task target output to be predictable from the actual latent variables.  However, the paper assumes more strongly that the target is a function of the *posterior distribution* of the latent variables given the observed sentences.  In other words, the goal of the downstream task is not to recover what the speaker *actually* meant, but only to recover what a hearer who knows the true generating distribution *thinks* it meant (given only the sentence and no other context).  Two structures are considered for the latent-variable model and its relationship to the downstream task.  For the more complex version, the predictor is also more complex, requiring argmax attention over all positions i.  In both cases, the paper considers allowing the predictor to prepend an arbitrary word embedding x_1 to the input.  This soft prompt can serve to shift the posterior predictions of x_i so as to be more informative for the downstream task.   

This work puts forward a theoretical analysis of pre-trained language models and head and prompt tuning on downstream tasks. In particular it studies the assumptions on the distribution of language and downstream tasks under which the downstream label can be recovered from a pretrained LM by using head or prompt tuning, by considering a generative model of language that can be represented as a HMM. The key components of the analysis include: 1. The data distribution is generated from a HMM (or a memory augmented HMM). 2. The initial state of the HMM, $H_0$, contains all the information necessary to recover the downstream label for the sequence.  Under this setup, for head tuning, the downstream labels are recoverable under the condition that the token emission probability matrix $W$ has linearly independent columns (implying that the dimensionality of the HMM hidden states is smaller than the vocabulary size of the data generating distribution). On the other hand, prompt tuning requires only a relaxed version of this condition, only the columns corresponding to the support of the downstream task need to be independent, eliminating the condition on the hidden state dimensionality.  Under memory augmented HMMs, for downstream tasks that can be represented as a linear function of the memory (with a fixed number of cells, $N$), the non-degeneracy condition can be relaxed even for head tuning. For prompt tuning, the independence assumption can be further relaxed to hold only for the set of memory cells that the downstream task depends on.  The authors further perform simulation experiments to validate the findings on different settings of the data generating distribution. Simulation results agree with the findings: 1. Prompt tuning outperforms head tuning in the setting where hidden state size is significantly larger than the token vocabulary size. 2. When the data is generated using memory-augmented HMMs, and the downstream task information can be recovered entirely from the memory, even head tuning gets close to perfect performance for all hidden state sizes.

The paper studies the question why head and prompt tuning work for pre-trained LM. Analyzing complicated neural networks is a daunting task. The paper offers an analysis framework based on the Hidden Markov Model (HMM) and memory augmented HMM. Under this framework, the paper makes a connection with masked LM (BERT) that at time step i, the output of BERT(`[mask]_i`) represents the marginal P(X_i) of the HMM. The framework further allows for analysis of a binary classification task using head tuning or prompt tuning via the posterior distribution of a HMM. To this end, the paper runs a simulation study with data generated from HMM and memory augmented HMM. 

This paper proposes an analysis framework to link the pre-training and downstream tasks. In this framework, the downstream task needs to predict the properties of the posterior distribution over latent variables in an underlying generative model. And When the generative model is a standard HMM, the downstream recovery is possible with a simple classification head under strong non-degeneracy assumptions. It further shows that changing the generative model to a memory-augmented HMM or using prompt tuning can relax the non-degeneracy conditions.