{'fraction_non_alphanumeric': 0.04959441076197034, 'fraction_numerical': 0.014745646003728397, 'mean_word_length': 3.7734487734487736, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 27, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Why Do Pretrained Language Models Help in Downstream Tasks? An Analysis of Head and Prompt Tuning

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.


Introduction
Natural language processing (NLP) has been revolutionized by large-scale pretrained language models such as BERT [4] and GPT [27], 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,37,36,11,29,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 and prompt tuning, which both freeze all pretrained parameters and allow us to treat the pretrained model as a black box. Head tuning [24] trains task-specific heads on top of the pretrained model outputs. Prompt tuning [33,20,9,22] 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 [20,9,22].
Our work contrasts with prior theoretical work [30], 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 (G i pxq " P rX i | 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 [20,9,22]. 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 longrange 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 attention-based 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. 1

Related works
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 [37,11]) by evaluating the performance of simple classifiers, or probes, on the representations [7,12,36,29,15]. Other papers ablate various aspects of pretraining, such as changing the masking scheme [14,21,42] or permuting the word order [34].
In comparison, theoretical analysis of pretrained language models is limited. Besides [30], which we discussed in Section 1, Zhang and Hashimoto [42] 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 [33,20,9,22,13,6,43,2,25], 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. [33] employ task-specific prompts that are optimized over the discrete token space. Schick and Schütze [31,32] reformulate natural language tasks as cloze-style phrases to enable few-shot learning. Subsequent methods [20,9,22] optimize "soft" prompts, or continuous embedding vectors. Lester et al. [20] employ soft prompts on pretrained large-scale T5 [28] 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 [22] propose prefix tuning, which prepends a trainable prefix embedding sequence to all layers of the transformer.
More broadly, Lee et al. [19] 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,38,40,39,23], but they are not directly relevant for our particular setting.

Formulations and notations
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 " pX 1 , . . . , X T 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 " pG 1 pxq, G 2 pxq, . . .q denote the masked language model which predicts a probability vector for each timestep in the input x. Our theoretical abstraction is that G i perfectly computes the distribution of X i , the i-th token, conditioned on all other tokens: G i pxq " P rX i |X´i " x´is. Here P rX i | X´i " x´is P ∆ |X | is a probability vector. In particular, G i pxq does not depend on x i . 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 " epx 1 q, . . . , epx t q denote the token embeddings, we have Gpepxqq " Gpxq. Soft prompt tuning concatenates a trainable prompt u so that the model output is Gppu, 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 s u,v " PrpU " u | V " vq.
For a sequence v " pv 1 , . . . , v t q, we use the notation v i:j for i ď j to denote pv i , . . . , v j q, and v´i to denote pv 1:i´1 , v i`1:t q. We let 1 denote the indicator function. For set V, we let V˚" V 1 Y V 2 Y¨¨¨denote variablelength sequences of elements of V. Let d denote elementwise product. Let 1 d , 0 d denote the d-dimensional To abstractify soft prompt tuning, we note that every token has a natural embedding, the corresponding row of the emission probability matrix. We view prompt tuning as adding a fake token r z to the vocabulary, assigning it a row u in the emission matrix, and prepending it to the input embedding sequence. More details are provided in Section 3.1.
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 R d , we let a{b denote their element-wise division. We use supppaq to denote the set of indices where vector a is non-zero.

Analysis for Hidden Markov Models
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. [26,18,3]).

Data distribution.
Let H denote the hidden state space of the HMM. We use H " pH 0 , H 1 , . . . , H T q P Ht o denote the sequence of hidden states. For all timesteps i ą 0, the transition probabilities are time-invariant, i.e. P rH i | H i´1 s " A for A P R |H|ˆ|H| . For each timestep i ě 1, tokens X i are emitted following some time-invariant probability: P rX i | H i s " W for W P R |X |ˆ|H| . The joint probability of X, H is Downstream tasks. We assume that H 0 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 rH 0 | X 1:T " xs: 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 1pb J G 1 pxq ě 0q for b P R |X | . The following non-degeneracy condition is crucial for our recovery result in this setting. We also require the following regularity conditions on H 0 and the state transitions. 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 on G applied to a shifted sequence. That is, there exists linear head weights b P R |X | such that for all x P supppP rXsq, where x 1 " p∅, x 1:t q is the concatenation of a special token ∅ with x. 2 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. 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 G 1 px 1 q " W P rH 1 |X 2:T`1 " xs ùñ P rH 1 |X 2:T`1 " xs " W : G 1 px 1 q where W : is the left inverse for W , guaranteed to exist by Assumption 3.1. This lets us recover P rH 1 |X 2:T`1 " xs by applying a linear function to G 1 px 1 q. Additional linear functions will be sufficient to obtain µ J P rH 0 |X 1:T " xs from P rH 1 |X 2:T`1 " xs. We provide the full proof in Section A. Proposition 3.4 is reminiscent of the arguments of [19], 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 rH 0 | X 1: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].

Relaxed non-degeneracy assumptions via prompt tuning
In this section, we study applying soft, or continuous, prompt tuning [20,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 e 1 , . . . , e t to conditional probabilities G 1 pxq, . . . , G t pxq 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 rX i " z | H i s. 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 rX i | X´i " x´is only depends on x 1 , . . . , x t through their embeddings epxq " pepx 1 q, . . . , epx t qq. In other words, there exists a function G i such that G i px 1 , . . . , x t q " G i pepx 1 q, . . . , epx t qq 2 We note that G 1 px 1 q does not depend on x 1 1 and therefore x 1 1 can be any token.
In particular, we let G i compute the standard message passing algorithm [16] that computes the conditional probability of HMMs. This ensures that G i is well defined on all sequences of nonnegative vectors in r0, 1s |H| , beyond sequences of proper embeddings.We assume that pretraining produces this G i , 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 r z. Concretely, consider adding a new token r z to the vocabulary X , and changing the emission probability at position 1 to satisfy P rX 1 " r z | H 1 s " u and for all z ‰ r z, P rX 1 " z | H 1 s9p1´uqdepzq. Then G i pu, epx 1 q, . . . , epx t qq precisely computes the conditional probability P rX i | X´i " pr z, x 1 , . . . , x t q´is under the modified HMM. We refer the readers to Section B for the formal definition of G i 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 and outputs the prediction F pxq " 1pb J G 2 pp epxqq ě 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 :,h u hPH ‹ , are linearly independent. Furthermore, H ‹ covers all meaningful information for the downstream tasks: supppµq Ď H ‹ .
In addition, a last technical requirement on H ‹ is as follows: there exists a set B Ď H such that H ‹ " Y hPB supppA :,h q. 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, where p e 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 [22]. 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 rH 0 | X 1:T " xs is impossible without strong nondegeneracy conditions (Assumption 3.1), we can aim to recover P rH 0 | X 1:T " xs on the subset of essential states H ‹ defined in Assumption 3.5, which suffices for computing µ J P rH 0 | X 1:T " xs, since H ‹ Ě supppµq. To recover P rH 0 | X 1:T " xs on H ‹ , we observe in Lemma B.2 that prepending the prompt u is equivalent to introducing a modified random sequence p X and fake token r z which influences the posterior of H 2 as follows: for invertible diagonal matrix D and positive scalar r x . We choose u such that the vector P rH 2 | p X 1 " r zs d P rH 0 | X 1:T " xs is supported only on H ‹ . Because corresponding columns of W are linearly independent by Assumption 3.5, we can then recover PrpH 0 " h | X 1:T " xq for h P H ‹ by applying a linear function to G 2 pp epxqq. This suffices for computing µ J P rH 0 | X 1:T " xs. More details are in Section B.

Analysis for memory-augmented Hidden Markov Models
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 [41,35] 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 H 0 , H 1 , . . ., meant to model the evolution of syntax, and a persistent "memory" M " pM 1 , . . . , M N q with N total cells, where each M i takes values in a finite set M. The full joint probability is as follows: 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 H i " pJ i , S i q and H " rN sˆS. To sample the token at timestep i given the hidden state H i " pJ i , S i q, we first use J i to index the memory M , obtaining the random variable M Ji . X i is then sampled according to some time-invariant probability depending on M Ji , J i , S i : 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 X i 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 pH 0 , M q, pH 1 , M q, . . . 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." M 1 could store the subject ("cow"), M 2 the location ("pasture"), M 3 the sentiment ("happily"), and S i could determine part-of-speech. For timesteps where "cow" and "rolled" are emitted J i " 1 because we emit information related to the sentence subject. Timesteps for "pasture" and "grass" would have J i " 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: 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).

Tuning attention head for recovering ground-truth downstream labels
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 G i pxq and trainable token embeddings epx i q, the attention head Attnp¨q applies key and value functions K, V to compute the output as follows: 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 KpG i pxqq " Θ pKq G i pxq. 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 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,hq u mPM,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,hq u mPM,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,hq u mPM,hPH ‹ q X spanptW :,pm,h 1 q u mPM,hPHzH ‹ q " t0 |X | u.
Note that the non-degeneracy condition of Theorem 3.3 would require tW :,pm,hq u mPM,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 " pj i , s i q for which the set tW :,pm,hq u mPM of token emission probabilities is fundamentally not very diverse, and therefore not linearly independent. For example, if the syntax s i indicates "article", i.e. words such as "a", "an", and "the", the token emission probabilities would carry little information about M ji because the choice of article does not depend much on semantics, so columns corresponding to s i " "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 J i " 2 (location), S i " "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,h u mPM is linearly independent for this hidden state.
Thus, our aim is to focus on recovering information for the downstream task from positions i where H i P H ‹ . Formally, we define the following set of input sequences containing positions i where the posterior of H i given x´i concentrates on H ‹ : 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 H i 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. . Then there exist an attention head on Gpxq and token embeddings epx i q such that the following holds for any x P R: 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 rH i | 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 G i pxq at such i, the value function in the attention will be able to recover P rM j ‹ | X 1:T " xs. A full proof is provided in Section C.1.

Guarantees for prompt-tuning
Though the generative modeling assumptions in this section already allowed us to relax the 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µ J P rM |X 1:T " xs ě 0q. This simplified setting also doesn't require the explicit disentanglement between J i and S i in H i . 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 rX i " z|M, H i s " W J z,: . In Section C.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 p epxq fi pu, epx 1 q, . . . , epx t qq (4.4) The downstream model applies an attention head to the output of G: F pxq " 1pAttnppG i pp epxqq, p e i pxqq t`1 i"1 q ě 0q, where Attn is defined in (4.2). An additional stationarity assumption on P rH 0 s will simplify the recovery procedure (though it can be removed). 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. Furthermore, the span of these vectors must be disjoint from the span of token emission probabilities from M ‹ˆp HzH ‹ q: spanptW :,pm,hq u mPM ‹ ,hPH ‹ q X spanptW :,pm,h 1 q u mPM ‹ ,hPHzH ‹ q " t0 |X | u.
We note that Assumption 4.5, and Assumption C.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 P M. 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 Gpp epxqq and the token embeddings which can compute the ground-truth F ‹ pxq for any x P R, defined in (4.3): where p e 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 G i pp epxqq, 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 C.3

Simulations
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. Our code is available at the following URL: https://github.com/ sangmichaelxie/pretraining_analysis. 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 X 1:T of length 129, where the first token X 1 " [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 D.
Head vs. prompt tuning. We compare head and prompt tuning as the hidden state size of the datagenerating HMM varies. The downstream label is generated by computing µ J P rH 1 | X´1 " x´1s, where µ is a random ground-truth linear weight. Head tuning learns a linear head on top of the softmax probabilities  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 µ J P 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 H 1 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.

Conclusion
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 generative distributions studied here are meant to provide a first-cut result -we also conjecture similar theorems to hold for other generative models, which we leave as an interesting direction for future work.
Another direction for future work is to analyze finetuning. Existing work analyzes finetuning for linear neural networks and obtains empirically useful insights [17], but analyzing neural networks with nonlinear activations is very challenging. Our analysis of head and prompt tuning treats the model as a black box. Analyzing finetuning requires understanding how to open up the black box, which is a major open question.

A Proofs for Section 3
We provide the formal proof of Theorem 3.3 based on the sketch in Section 3. The following lemma will be useful in our analysis.
Claim A.1. In the setting of Section 3, suppose that Assumption 3.2 holds. Fix any timestep i ě 1. Then there exists a diagonal matrix D such that for all x P supppP rXsq, P rH i | X i`1:T`i " xs " r x DP rH 0 | X 1:T " xs where r x ą 0 is a positive scalar.
Proof. First, we note that by Assumption 3.2, P rH i s has full support. As a consequence, PrpX i`1:t`i " xq ą 0. By Bayes' rule,

by Markovian property of HMMs)
" P rH 0 | X 1:T " xs d P rH i s P rH 0 s¨P Note that the vector P rHis P rH0s has finite and positive entries. The same applies to the ratio r x fi PrpX 1:T "xq PrpX i`1:T`i "xq . Thus, we get the desired statement.

The proof of Theorem 3.3 follows below.
Proof of Theorem 3.3. By definition, G 1 px 1 q " P rX 1 | X 2:T`1 " xs. Therefore, our goal is to rewrite P rH 0 | X 1:T " xs as a linear function of P rX 1 |X 2:T`1 " xs (up to a scaling which won't affect the linear head prediction). Concretely, we will show P rH 0 | X 1:T " xs " r x BP rX 1 | X 2:T`1 " xs (A.1) for a scalar r x ě 0. With this equation, taking b " µ J B will give the desired result.

Proof of Proposition 3.4. We write
" P rU | ZsP rZ | V " vs B Formal abstraction for prompt tuning and proofs for Section 3.1 We first formalize the definition of the model G described in Section 3.1. The model G takes a sequence of embedding vectors v " pv 1 , . . . , v t q as input and implements message passing to compute a sequence of t outputs. We first define left and right messages Ð Ý δ i`1Ñi pvq and Ý Ñ δ i´1Ñi pvq for i P rts, as follows: Next, we define the aggregated message at timestep i by Note that if Assumption 3.2 holds about the Markov chain H 0 , H 1 , . . ., τ i pvq is always well-defined because P rH i s will have full support. Note that for the proper embeddings epx i q " P rX i " x i | H i s, where for x " px 1 , . . . , x t q, we use epxq " pepx 1 q, . . . , epx t qq, we can check via classical results on message passing [16] that Finally, we let the model model G compute There is an edge case where the demoninator is 0, i.e. }τ i pvq} 1 " 0. To make the behavior of G well-defined, in this case we set G i pvq " 0 |X | . We observe that if the input embedding are obtained by epxq, G i pvq indeed computes the desired conditional probability vector for x P supppP rXsq:

B.1 Proof of Theorem 3.6
First we formalize the observation that soft prompt tuning is equivalent to adding a fake token r z to the vocabulary with emission probabilities at timestep 1 given by u, and letting G compute conditional probabilities for this new distribution over sequences.
Lemma B.1. In the setting of Theorem 3.6, fix any prompt vector u P r0, 1s |H| . Define the random variable p X with the same emission probabilities as X for i ą 1: P r p X i | H i s " P rX i | H i s. For timestep 1, we define the emission probabilities of p X 1 as follows: In the above equations, r z is a fake token added to the vocabulary at timestep 1. It follows that for any i, defining τ i as in (B.1) As a consequence, it follows that for i ą 1 and any x such that pr z, ∅, xq´i P supppP r p X´isq, For any x with pr z, ∅, xq´i R supppP r p X´isq, G i pp epxqq " 0.
Next, the following lemma disentangles the influences of the fake token r z and the input sequence on the posterior distribution of the hidden variable.
Lemma B.2. In the setting above, there exists an invertible diagonal matrix D such that for all x such that pr z, xq P supppP r p X´2sq, the following equation holds: Here r x ą 0 is a positive scalar.
We now complete the proof of Theorem 3.6.
Proof of Theorem 3.6. Let B be the set defined in Assumption 3.5 and define u such that u h " 1 if h P B and u h " 0 otherwise. First, we restrict our focus to x such that pr z, xq P supppP r p X´2sq. For these x, we can apply Lemma B.1 and Lemma B.2 in the manner described in the proof sketch. This gives G 2 pp epxqq " r x W Dv for v fi pApu d P rH 1 sqq d P rH 0 | X 1:T " xs. By definition of B, we have supppApu d P rH 1 sqq " H ‹ , so supppDvq Ď H ‹ . Thus, there is a matrix y W : such that The existence of y W : is due to the fact that tW :,h u hPH ‹ is a linearly independent set of vectors, and supppDvq Ď H ‹ whenever x satisfies pr z, xq P supppP r p X´2sq. Next, we note that a matrix B exists such that pBDvq h " PrpH 0 " h | X 1:T " xq for h P H ‹ and pBDvq h " 0 otherwise. This is because D is invertible, and supppApu d P rH 1 sqq " H ‹ , so we can recover P rH 0 | X 1:T " xs on coordinates in H ‹ by applying another coordinate-wise scaling. It follows that we can set b " µ J B y W : . With this choice of b, we compute where the last equality follows because supppµq Ď H ‹ . This completes the case where pr z, xq P supppP r p X´2sq.
Otherwise, for pr z, xq R supppP r p X´2sq, by the behavior of G in Lemma B.1, G 2 pp epxqq " 0, so any linear head must output b J G 2 pp epxqq " 0. Furthermore, by the conditional independence structure in p X, we must also have supppP rH 2 , p X 1 " r zsq X supppP rH 2 , p X 3:T`2 " xsq " H. As supppµq Ď supppP rH 2 , p X 1 " r zsq, this must also mean supppµqXsupppP rH 2 , p X 3:T`2 " xsq " H. However, we also have P rH 2 , p X 3:T`2 " xs " P rH 2 , X 3:T`2 " xs by the definition of p X, and this must have the same support as P rH 0 | X 1:T " xs by applying Claim A.1 and the fact that x P supppP rXsq. It follows that for this choice of x, µ J P rH 0 | X 1:T " xs " 0, so the desired statement still stands.
We fill in the proofs of the lemmas below.
Proof of Lemma B.1. First, we note that (B.2) follows directly from the derivation of τ , and well-known results about message passing [16]. Next, it suffices to consider the case where pr z, ∅, xq´i R supppP r p X´isq, as the other case follows directly from the definition of G in terms of τ . In this case, we observe that τ i pp epxqq " P rH i , p X´i " pr z, ∅, xq´is " 0. It follows that }τ i pp epxqq} 1 " 0. Thus, from our definition of G, we must have G i pp epxqq " 0.
Proof of Lemma B.2. By the conditional independence relations in a HMM, p X 1 K p X 3:T`2 | H 2 . Using Bayes' rule, we obtain Prp p X 1 " r z, p X 3:T`2 " xq (by definition of p X and the Markovian property) " r x P r p X 1 " r z, H 2 s d P rH 0 | X 1:T " xs d 1 P rH 0 s Where we define r x fi . We note that r x is positive and well-defined by the conditions of the lemma and Theorem 3.6. We can set D to be the matrix diagp 1 P rH0s q, which has finite positive entries on the diagonal by Assumption 3.2.

C Proofs for Section 4
First, we introduce a proposition which is generally useful for proving the theorems in Section 4.
Proposition C.1. In the setting of Section 4, it holds that P rX i | X´i " x´is " P rX i | M Ji , J i , S i sP rM Ji , J i , S i , X´i " x´is Equivalently, we have the expansion Proof. An alternative interpretation of this statement is that X i is conditionally independent from everything else given M Ji , J i , S i . However, we will prove this statement algebraically. We compute P rXi | X´i " x´is " Throughout this section, we use M Ji to denote the random variable obtained by indexing M by J i , both of which are themselves random variables. Let p I denote the set of indices i where supppP rJ i | X´i " x´isq " tj ‹ u and supppP rS i | X´i " x´isq Ď S ‹ . We will first construct the key function K and query q such that the set of I of attended-to positions (4.2) is precisely p I. This construction does not require the position embeddings β 1 , . . . , β t , so we set them to 0.
The following lemma demonstrates the existence of K and q such that I " p I.
Lemma C.2. In the setting of Theorem 4.3, define p I fi ti : supppP rJ i | X´i " x´isq " tj ‹ u and supppP rS i | X´i " x´isq Ď S ‹ u. Then there exist query q P R |H| and key K parameterized by Θ pKq P R |H|ˆ|X | , such that when x P supppP rXsq and p I is nonempty, the set I of attended-to positions satisfies I " p I.
The proof of Lemma C.2 requires the following claim.
Claim C.3. In the setting of Theorem 4.3, there is a matrix Θ p1q P R |H|ˆ|X | such that for all x P supppP rXsq and s P S ‹ , pΘ p1q G i pxqq pj ‹ ,sq " P rH i " pj ‹ , sq | X´i " x´is. Furthermore, }Θ p1q G i pxq} 1 " 1. In addition, for s P S ‹ , there exists Θ p2,sq P R |M|ˆ|X | such that for all x P supppP rXsq, Proof. We have, by Proposition C.1, In the last equality, we defined ν phq to be the expression in the parentheses. Note that ν phq P V phq fi spanptW :,pm,hq u mPM q. Furthermore, for h R H ‹ , ν phq P s V fi spanptW :,pm,hq u mPM,hPHzH ‹ q. As the spans pV phq q hPH ‹ and s V are all pairwise disjoint, by Assumption 4.2, for each h P H ‹ , we can recover Likewise, we can obtain Now we have, for h P H ‹ , Likewise, the same reasoning gives 1 J ř hRH ‹ ν phq " ř hRH ‹ PrpH i " h | X´i " x´iq. Thus, we can choose Θ p1q to be the matrix with rows Θ We set all other rows to 0, and we can check that this satisfies the lemma requirements.
We now construct Θ p2,hq . We can express ν phq in a vectorized manner by writing ν phq " W :,pM,hq P rM j , H i " h | X´i " x´is where W :,pM,hq P R |X |ˆ|M| has columns tW :,pm,hq u mPM . Note that for j " j ‹ , s P S ‹ , the non-degeneracy assumptions imply that W :,pM,j ‹ ,sq has left inverse W : :,pM,j ‹ ,sq . Thus, we set Θ p2,sq " W : :,pM,j ‹ ,sq B pj ‹ ,sq to obtain for s P S ‹ , Θ p2,sq G i pxq " W : :,pM,j ‹ ,sq B pj ‹ ,sq P rX i | X´i " x´is " W : :,pM,j ‹ ,sq W :,pM,j ‹ ,sq P rM j ‹ , H i " pj ‹ , sq | X´i " x´is " P rM j ‹ , H i " pj ‹ , sq | X´i " x´is This gives the desired result.
Proof of Lemma C.2. We choose the first |H| entries of q such that q h " 1 if h " pj ‹ , sq for s P S ‹ , and q h " 0 otherwise. The last entry is 0. Next, we choose Θ pKq so that the first |H| rows are Θ p1q , and the last row is all zeros. where Θ p1q is defined in Claim C.3. With this choice of Θ pKq , KpG i pxqq h " PrpH i " h|X´i " x´iq for h P H ‹ . Furthermore, }KpG i pxqq} 1 " 1, by Claim C.3. Now we note that for all i, 1 " }KpG i pxqq} 1 ě q J KpG i pxqq, and for i P p I, q J KpG i pxqq " ř sPS ‹ PrpH i " pj ‹ , sq|X´i " x´iq " 1 by definition of q and p I. This implies that positions i P p I do indeed achieve the maximum attention scores.
Next, we also require a construction of the value function such that it computes the correct prediction for all i P p I.
Lemma C.4. In the setting of Theorem 4.3, let p I be defined as in Lemma C.2. We can choose the parameters of the value function V , Θ pV q P R |M||H|ˆ|X | , b P R |M||H| , such that when x P supppP rXsq and p I is nonempty, for all i P p I, where r x,i ą 0 is a positive scalar.
Proof. We first choose Θ pV q such that the rows satisfy Θ pV q pm,j ‹ ,sq,: " Θ p2,sq m,: when s P S ‹ for Θ p2,sq constructed in Claim C.3, and Θ pV q pm,j,sq,: " 0 |X | otherwise for j ‰ j ‹ or s R S ‹ .
We claim that for i P p I, This is because for s P S ‹ , Θ p2,sq G i pxq " P rM j ‹ , H i " pj ‹ , sq | X´i " x´is by Claim C.3, and for h " pj, sq for j ‰ j ‹ or s R S ‹ , Note that this last equality followed because PrpH i " h | X´i " x´iq " 0 for the choice of h and i P p I. By construction of Θ pV q , these computations imply that (C.2) does indeed hold. The embedding can be chosen such that epx i q " P rX i " x i | M Ji , J i , S i s. Thus, we have for i P p I: The last equality followed from applying the same reasoning as in Proposition C.1.
Now we let B P R |M|ˆ|M||H| be the matrix such that Now we pick the last linear weight in the value function by b " B J µ. It follows that for i P p I, We obtained the last equality by observing that Proof of Theorem 4.3. By applying Lemmas C.2 and C.4, we constructed key, query, and value functions for the attention head such that for all x P supppP rXsq with p I (defined in Lemma C.2) nonempty, the attended-to positions I satisfy I " p I, and V pG i pxq, epx i qq " r x,i µ J P rM j ‹ | X 1:T " xs. As the attention head computes the average of V pG i pxq, epx i qq over attended-to positions, and r x,i is positive for all i P p I, we obtain the desired result.
We note that this proof also works for the case where there is a single memory cell, as that is a special case where J i " j ‹ always, and we only need to consider the evolution of S i . Now we formalize the model G. G will take embedding vectors v " pv 1 , . . . , v t q with v i P R | r H| as follows. We define left and right messages Ð Ý δ i`1Ñi pvq and Ý Ñ δ i´1Ñi pvq for i P rts via: We observe that this definition almost matches Section B, except it replaces H with r H. Next, we define the aggregated message at timestep i by In the edge case where P rM s does not have full support, the coordinate-wise division in the definition above would sometimes divide by 0. However, for all these cases both of the corresponding terms in the numerator must also be 0, so we can simply set the value of τ i in this coordinate to 0. We will see that this preserves the meaning of the message τ i , which for the proper embeddings epx i q " P rX i " x i | r H i s, with epxq " pepx 1 q, . . . , epx t qq, computes We can now define the reverse lifting function φ : R | r H|Ñ|M||H| as follows: We observe that φpτ i pepxqqq " P rM J i ,Ji,Si,X´i"x´is |M| N´1 . We now compute the model output as follows: In the edge case where }φpτ i pvqq} 1 " 0, we again define Gpvq " 0 |X | . We can observe that G i pepxqq " P rX i | X´i " x´is.
The downstream classifier uses the embedding p epxq defined as follows: p epxq " pu, epx 1 qq, . . . , epx t qqq with a tunable prompt embedding u P R | r H| . We also require a slightly modified attention head. The value function V in the attention head is slightly modified to accomodate the new embedding dimension. Letting The dimensions of the parameters b, Θ pV q remain unchanged. Note that when there is just a single memory, this reduces to the case in Section 4.

C.3 Analysis for prompt tuning in the multiple memory setting
We will state and prove our result for the prompt tuning setting with multiple memories. For the multiple memory setting, the downstream classifier uses the following embedding function p e: p epxq " pu, ηpepx 1 qq, . . . , ηpepx t qqq with a tunable prompt embedding u P R | r H| . The attention head is changed so that the value function takes a larger dimensional embedding: where φ is defined in (C.4). The following assumption extends Assumption 4.5 to the multiple memory case.
Assumption C.5 (Multiple memories version of Assumption 4.5). 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,hq u mPM ‹ ,hPH ‹ , is a linearly independent set of vectors.
Furthermore, define the following span of vectors: s V fi spanptW :,pm,j ‹ ,sq u mPM ‹ ,sPSzS ‹ Y tW :,pm,j,sq u mPM,j‰j ‹ ,sPS q Then s V must be disjoint from the span of token emission probabilities from M ‹ˆH‹ : Note that Assumption C.5 reduces to Assumption 4.5 the case where N , the number of memory cells, is 1. In any case, it is a relaxation of Assumption 4.2.
We now state and prove the result for multiple memories.
Theorem C.6. In the setting above, suppose that non-degeneracy Assumption C.5 and holds. In addition, suppose that Assumption 4.4 (stationarity) holds. Then there exists a prompt u and attention head on Gpp epxqq and the token embeddings which can compute the ground-truth F ‹ pxq for any x P R, defined in (4.3): The proof of Lemma C.7 mirrors the proof of Lemma B.1, so we omit it here.
In particular, throughout the proof we will use the following prompt u: We will also use the notation p x fi pr z, x 1 , . . . , x t q. The following lemma considers behaviors in edge cases with this choice of u.
Towards our proofs, the following result is useful.
Proposition C.8. In the setting of Theorem C.6, where P rH 0 s is the stationary distributions satisfying P rH 0 s " AP rH 0 s, it holds that P rM, H i , X i`1:i`t s " P rM, H 0 , X 1:t s for any t ě 1, i ě 1.
Proof. Because P rH 0 s is stationary, we observe that P rM, H i s " P rM, H 0 s for all i. We write We will now restrict our focus to the set of inputs Z fi tx : Prp p X´i " pr z, xq´iq ą 0 @i P rtsu (C.7) We also define the set p I fi ti`1 : supppP rS i |X´i " x´isq Ď S ‹ , supppP rJ i |X´i " x´isq Ď tj ‹ u, i P rtsu (C.8) Here S ‹ is defined in the non-degeneracy assumption. We will first construct key and query parameters such that the set of attended-to positions is precisely p I, following the proof of Theorem 4.3.
Lemma C.9 (Analogue to Lemma C.2). In the setting of Theorem C.6 and above, define u as in (C.6).
There are parameters Θ pKq P R p|H|`1qˆ|X | , q P R |H|`1 , and β 1 , β 2 , . . . P R |H|`1 such that for any x P Z where p I is nonempty, the set of attended-to positions I (defined in (4.1)) satisfies I " p I.
Towards proving Lemma C.9, the following construction will be useful.
Claim C.10 (Analogue of Claim C.3). In the setting of Theorem C.6, define H ‹ as in Assumption C.5. There is a matrix Θ p1q P R |H|ˆ|X | such that for all x P supppP rXsq, and i ą 1 with Prp p X´i " p x´iq ą 0, X´i " p x´is for any h P H ‹ . Furthermore, }Θ p1q G i pp epxqq} 1 " 1.
In addition, for s P S ‹ , there exists Θ p2,sq P R |M|ˆ|X | such that for all i ą 1 and x with Prp p X´i " p x´iq ą 0, Θ p2,sq G i pp epxqq " P rM j ‹ , H i " pj ‹ , sq | p X´i " p x´is Our proof will require the following result which shows that the distribution of M j ‹ has limited support.
Proposition C.11. In the setting of Theorem C.6 and Lemma C.7, let u be defined as in (C.6). Then for all i ą 1, supppP rM j ‹ | p X´i " p x´isq Ď supppµq if Prp p X´i " p x´iq ą 0.
First, if pr z, xq R supppP r p Xsq, by Claim C.15, we have µ J P rM j ‹ | X 1:T " xs " 0. The expression above must also equal 0, as pr z, xq R supppP r p Xsq. Otherwise, we have V pG i pp epxqq, p e i pxqq " µ J P rM j ‹ | p X " p xsPrp p X i " p x i | p X´i " p x´iq Now we apply Claim C.14 to get the desired result in this case. A additional case is when Prp p X´i " p x´iq " 0. In this case, Claim C.15 shows that G i pp epxqq " 0, so it follows that the value function also computes 0 in this case.
Finally, we need to check the case where p x R supppP r p Xsq, and we want to show V pG i pp epxqq, p e i pxqq " 0 for all i ą 1. The case where Prp p X´i " p x´iq " 0 is already handled above. In the case where Prp p X´i " p x´iq ą 0, we can apply Claim C.10 to our construction for Θ pV q to get pΘ pV q G i pp epxqqq m,h " # P rM j ‹ " m, H i " pj ‹ , sq | p X´i " p x´is if h " pj ‹ , sq for s P S ‹ 0 otherwise Thus, taking the element-wise product with φpepx i qq " P r p X i " p x i | M Ji , J i , S i s, we must have, by Proposition C.1, Both of these terms must be 0 since p x R supppP r p Xsq, giving the desired result.
Now we are ready to prove Theorem C.6.
Proof of Theorem C.6. The first case we consider is when x P Z, defined in (C.7). By applying Lemmas C.9 and C.13, we constructed key, query, and value functions for the attention head such that when p I (C.8) is nonempty, the attended-to positions I satisfy I " p I. In addition, by applying Lemma C.13, we also obtain that for x P supppP rXsq, V pG i pp epxqq, p e i pxqq " r x,i µ J P rM j ‹ | X 1:T " xs. As the attention head averages V pG i pp epxqq, p e i pxqq over the attended-to positions, and r x,i is positive for all i P p I, we obtain the desired result.
In the second case, x R Z, so pr z, xq R supppP r p Xsq. By Lemma C.13, for all i ą 1, the value function outputs 0. However, by the construction in Lemma C.9, the attention will only attend to i ą 1. Thus, the output of the attention head is 0. However, Claim C.15 also implies that µ J P rM j ‹ | X 1:T " xs " 0, giving the desired result.

D Experimental details
Generating HMM parameters. For all experiments, we randomly generated the parameters of an HMM with 10 output symbols in its vocabulary. We generate a random transition matrix by taking a random convex combination of random permutation matrices. We mix as many permutation matrices as there are hidden states; i.e. if there are 4 hidden states, then we mix 4 random permutation matrices. The mixing weights are generated by sampling logits IID from a uniform distribution on r0, 1s and then taking a softmax with temperature 0.01. Although this is a small temperature, the transition probabilities can still be around 0.7 for some transitions. The start distribution is also sampled in the same way, but with softmax temperature 10.0. The rows of the emission probability matrix is also sampled the same way with temperature 0.01.
Pretrain model. The pretrained model follows the BERT-base architecture, except with 6 layers and a much smaller vocab size.
Pretrain data and task. The pretraining data consists of 5000 sequences (documents) generated from the HMM, each with length 10240. We pretrain on this data by doing 5% masked LM on chunks of length 512. Pretraining runs for 3 epochs and takes about 5 hours on a single NVIDIA Tesla K80 GPU on 16-bit precision. We use an internal cluster for all experiments. Pretraining uses batch size 8 and learning rate 1e-5 with a linear warmup of 500 steps and linear decay schedule after 500 steps. We generated 20 pretraining (and downstream) datasets for each problem instance and average over the 20 runs in the vanilla HMM comparison, while the memory-based distributions are run for 5 trials of pretraining and finetuning. Downstream. The downstream task samples a sparse ground truth linear weight µ with 6 nonzero elements. Positions for nonzero entries are sampled uniformly at random and values are sampled i.i.d. from a standard normal distribution. Although we do binary classification, we sample µ with 2 rows and take the label to be the argmax of the two scores, instead of having 1 row and taking the sign. We find that this results in less degenerate datasets (datasets where all labels are the same).
We generate 5000 training, 500 validation and 1000 test examples for the downstream tasks. Downstream training uses learning rate 0.01 for both prompt tuning and head tuning, with a linear warmup/decay schedule, for 5 epochs over the downstream data. We take the model returned at the last checkpoint as the result (no early stopping). We found that it was important to train prompt tuning with full precision, since the gradients are relatively small and become zero with discretization.
We used message passing in the HMM to compute the posterior distributions of the latent variables analytically.
Prompt tuning. We prepended a length 20 continuous prompt to each sequence of input word embeddings. We initialize elements of the prompt vectors IID from the uniform distribution on r´0.5, 0.5s. Our implementation for prompt tuning used the code of [20], available at https://github.com/kipgparker/ soft-prompt-tuning.