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

Introduction

Natural language processing (NLP) has been revolutionized by large-scale pretrained language models such as BERT and GPT , 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 , 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 trains task-specific heads on top of the pretrained model outputs. Prompt tuning 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 .

Our work contrasts with prior theoretical work , 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}(x)=P[X_{i}\,|\,x_{-i}]$) 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 . 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 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.Code is available at https://github.com/sangmichaelxie/pretraining_analysis.

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 ) by evaluating the performance of simple classifiers, or probes, on the representations . Other papers ablate various aspects of pretraining, such as changing the masking scheme or permuting the word order .

In comparison, theoretical analysis of pretrained language models is limited. Besides , which we discussed in Section 1, Zhang and Hashimoto 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 , 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. employ task-specific prompts that are optimized over the discrete token space. Schick and Schütze reformulate natural language tasks as cloze-style phrases to enable few-shot learning. Subsequent methods optimize “soft” prompts, or continuous embedding vectors. Lester et al. employ soft prompts on pretrained large-scale T5 models and show that as the model size increases, prompt tuning performance can eventually match finetuning. Hambardzumyan et al. applies a variant of soft prompt tuning to MLM models. Li and Liang propose prefix tuning, which prepends a trainable prefix embedding sequence to all layers of the transformer.

More broadly, Lee et al. 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 , but they are not directly relevant for our particular setting.

Formulations and notations

We analyze models pretrained on masked language modeling (MLM) objectives. Let $\mathcal{X}$ denote a finite vocabulary of input tokens, $\mathcal{X}^{*}$ the set of variable-length sequences of tokens, and $X=(X_{1},\ldots,X_{T})\in\mathcal{X}^{*}$ a random sequence of $T$ tokens. Let $\Delta^{|\mathcal{X}|}$ denote the space of probability distributions over tokens.

Let $G(x)=(G_{1}(x),G_{2}(x),\ldots)$ 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}(x)=P[X_{i}|X_{-i}=x_{-i}]$. Here $P[X_{i}\,|\,X_{-i}=x_{-i}]\in\Delta^{|\mathcal{X}|}$ is a probability vector. In particular, $G_{i}(x)$ does not depend on $x_{i}$. The downstream task involves labeled examples $(x,F^{\star}(x))\in\mathcal{X}^{*}\times{\mathcal{Y}}$, where $F^{\star}:\mathcal{X}^{*}\to{\mathcal{Y}}$ provides ground-truth downstream labels and ${\mathcal{Y}}$ is a discrete set of labels for classification.

Head and prompt tuning.

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. ).

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^{\star}$ is assumed to be a linear classifier on the posterior $P[H_{0}\,|\,X_{1:T}=x]$:

The token emission probability matrix $W$ has linearly independent columns.

We also require the following regularity conditions on $H_{0}$ and the state transitions.

The Markov chain $H_{0},H_{1},\ldots$ is ergodic, and $P[H_{0}]$ has full support.

We show that if $W$ has linearly independent columns, a linear head fits downstream labels.

where $x^{\prime}=(\varnothing,x_{1:t})$ is the concatenation of a special token $\varnothing$ with $x$.We note that $G_{1}(x^{\prime})$ does not depend on $x_{1}^{\prime}$ and therefore $x_{1}^{\prime}$ can be any token.

The key for the proof is to leverage the following general statement about random variables $U,V,Z$ such that $U\perp V\,|\,Z$, which decomposes the expression for $P[U\,|\,V]$.

Let $U,V,Z$ be random variables such that $U\perp V\,|\,Z$. Then for any $v$, $P[U\,|\,V=v]=P[U\,|\,Z]\cdot P[Z\,|\,V=v]$. Thus, if $P[U\,|\,Z]$ has a left inverse $(P[U\,|\,Z])^{\dagger}$, then $P[Z\,|\,V=v]=(P[U\,|\,Z])^{\dagger}P[U\,|\,V=v]$.

By the conditional independence structure of the HMM, Proposition 3.4 immediately implies

where $W^{\dagger}$ is the left inverse for $W$, guaranteed to exist by Assumption 3.1. This lets us recover $P[H_{1}|X_{2:T+1}=x]$ by applying a linear function to $G_{1}(x^{\prime})$. Additional linear functions will be sufficient to obtain $\mu^{\top}P[H_{0}|X_{1:T}=x]$ from $P[H_{1}|X_{2:T+1}=x]$. We provide the full proof in Section A.

Proposition 3.4 is reminiscent of the arguments of , 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 $|\mathcal{H}|\leq|\mathcal{X}|$. Without this assumption, it is unclear how to recover $P[H_{0}\,|\,X_{1:T}=x]$ from $G(x)$ alone. However, in realistic settings we would expect $|\mathcal{H}|>|\mathcal{X}|$, as increasing the size of the hidden state space improves language modeling capabilities of HMMs .

In this section, we study applying soft, or continuous, prompt tuning to the setting above. We show that by using soft prompt tuning, we can recover $F^{\star}$ 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 $\overline{G}$ that maps a sequence of embeddings $e_{1},\dots,e_{t}$ to conditional probabilities $G_{1}(x),\dots,G_{t}(x)$ as follows. We observe that each token $z$ in the vocabulary $\mathcal{X}$ naturally corresponds to a $|\mathcal{H}|$-dimensional vector: the $z$-th row of the emission probability matrix $W$, or equivalently, $P[X_{i}=z\,|\,H_{i}]$. We denote this embedding by $e(z)$ and call the family of embeddings $\{e(z):z\in\mathcal{X}\}$ proper embeddings. A fundamental property of HMMs is that the conditional probability $P[X_{i}\,|\,X_{-i}=x_{-i}]$ only depends on $x_{1},\dots,x_{t}$ through their embeddings $e(x)=(e(x_{1}),\dots,e(x_{t}))$. In other words, there exists a function $\overline{G}_{i}$ such that

In particular, we let $\overline{G}_{i}$ compute the standard message passing algorithm that computes the conditional probability of HMMs. This ensures that $\overline{G}_{i}$ is well defined on all sequences of nonnegative vectors in $^{|\mathcal{H}|}$, beyond sequences of proper embeddings.We assume that pretraining produces this $\overline{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\in^{|\mathcal{H}|}$ to $\overline{G}$ in the first argument and proper embeddings at positions $i>1$. We can interpret $u$ as the embedding of a fake token $\widetilde{z}$. Concretely, consider adding a new token $\widetilde{z}$ to the vocabulary $\mathcal{X}$, and changing the emission probability at position 1 to satisfy $P[X_{1}=\widetilde{z}\,|\,H_{1}]=u$ and for all $z\neq\widetilde{z}$, $P[X_{1}=z\,|\,H_{1}]\propto(1-u)\odot e(z)$. Then $\overline{G}_{i}(u,e(x_{1}),\ldots,e(x_{t}))$ precisely computes the conditional probability $P[X_{i}\,|\,X_{-i}=(\widetilde{z},x_{1},\dots,x_{t})_{-i}]$ under the modified HMM. We refer the readers to Section B for the formal definition of $\overline{G}_{i}$ and formal proofs of the interpretation above.

There exists a set of essential hidden states $\mathcal{H}^{\star}\subseteq\mathcal{H}$, so that the columns of $W$ corresponding to $\mathcal{H}^{\star}$, $\{W_{:,h}\}_{h\in\mathcal{H}^{\star}}$ , are linearly independent. Furthermore, $\mathcal{H}^{\star}$ covers all meaningful information for the downstream tasks: $\textup{supp}(\mu)\subseteq\mathcal{H}^{\star}$.

In addition, a last technical requirement on $\mathcal{H}^{\star}$ is as follows: there exists a set ${\mathcal{B}}\subseteq\mathcal{H}$ such that $\mathcal{H}^{\star}=\cup_{h\in{\mathcal{B}}}\textup{supp}(A_{:,h})$. In other words, $\mathcal{H}^{\star}$ must be the set of all states reachable by starting from some state in ${\mathcal{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 $\mathcal{H}^{\star}$ of essential states. In the setting where $|\mathcal{H}|>|\mathcal{X}|$, the condition for Theorem 3.3 can never hold. On the other hand, Assumption 3.5 could still hold, for example, if $|\textup{supp}(\mu)|<|\mathcal{X}|$ and the set of columns of $W$ corresponding to hidden states in $\textup{supp}(\mu)$ 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^{\star}$ using soft prompt tuning with a linear head.

where $\widehat{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 . 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[H_{0}\,|\,X_{1:T}=x]$ is impossible without strong non-degeneracy conditions (Assumption 3.1), we can aim to recover $P[H_{0}\,|\,X_{1:T}=x]$ on the subset of essential states $\mathcal{H}^{\star}$ defined in Assumption 3.5, which suffices for computing $\mu^{\top}P[H_{0}\,|\,X_{1:T}=x]$, since $\mathcal{H}^{\star}\supseteq\textup{supp}(\mu)$. To recover $P[H_{0}\,|\,X_{1:T}=x]$ on $\mathcal{H}^{\star}$, we observe in Lemma B.2 that prepending the prompt $u$ is equivalent to introducing a modified random sequence $\widehat{X}$ and fake token $\widetilde{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[H_{2}\,|\,\widehat{X}_{1}=\widetilde{z}]\odot P[H_{0}\,|\,X_{1:T}=x]$ is supported only on $\mathcal{H}^{\star}$. Because corresponding columns of $W$ are linearly independent by Assumption 3.5, we can then recover $\textup{Pr}(H_{0}=h\,|\,X_{1:T}=x)$ for $h\in\mathcal{H}^{\star}$ by applying a linear function to $\overline{G}_{2}(\widehat{e}(x))$. This suffices for computing $\mu^{\top}P[H_{0}\,|\,X_{1:T}=x]$. 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 and is closely related to Hidden Topic Markov Models . There are two sets of latent variables in the memory-augmented HMM: a Markov chain on hidden states $H_{0},H_{1},\ldots$, meant to model the evolution of syntax, and a persistent “memory” $M=(M_{1},\ldots,M_{N})$ with $N$ total cells, where each $M_{i}$ takes values in a finite set ${\mathcal{M}}$. The full joint probability is as follows:

The hidden state is modified to explicitly consist of a disentangled cell index $J\in[N]$ and syntax state $S\in{\mathcal{S}}$, such that $H_{i}=(J_{i},S_{i})$ and $\mathcal{H}=[N]\times{\mathcal{S}}$. To sample the token at timestep $i$ given the hidden state $H_{i}=(J_{i},S_{i})$, we first use $J_{i}$ to index the memory $M$, obtaining the random variable $M_{J_{i}}$. $X_{i}$ is then sampled according to some time-invariant probability depending on $M_{J_{i}},J_{i},S_{i}$:

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$.

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 $\{W_{:,(m,h)}\}_{m\in{\mathcal{M}},h\in\mathcal{H}^{\star}}$ to be linearly independent in a subset $\mathcal{H}^{\star}$ of “recoverable” hidden states, whereas Assumption 3.1 required all columns to be linearly independent.

There exists a set of recoverable hidden states $\mathcal{H}^{\star}=\{j^{\star}\}\times{\mathcal{S}}^{\star}$, such that the collection of token emission probabilities from ${\mathcal{M}}\times\mathcal{H}^{\star}$, $\{W_{:,(m,h)}\}_{m\in{\mathcal{M}},h\in\mathcal{H}^{\star}}$, is a linearly independent set of vectors.

Furthermore, the span of these vectors must be disjoint from the span of token emission probabilities from ${\mathcal{M}}\times(\mathcal{H}\setminus\mathcal{H}^{\star})$: $\textup{span}(\{W_{:,(m,h)}\}_{m\in{\mathcal{M}},h\in\mathcal{H}^{\star}})\cap\textup{span}(\{W_{:,(m,h^{\prime})}\}_{m\in{\mathcal{M}},h\in\mathcal{H}\setminus\mathcal{H}^{\star}})=\{\mathbf{0}_{|\mathcal{X}|}\}$.

Note that the non-degeneracy condition of Theorem 3.3 would require $\{W_{:,(m,h)}\}_{m\in{\mathcal{M}},h\in\mathcal{H}}$ to be linearly independent, whereas Assumption 4.2 only requires linear independence for $h\in\mathcal{H}^{\star}$. The second condition states that $\mathcal{H}^{\star}$ and $\mathcal{H}\setminus\mathcal{H}^{\star}$ 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=(j_{i},s_{i})$ for which the set $\{W_{:,(m,h)}\}_{m\in{\mathcal{M}}}$ 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_{j_{i}}$ because the choice of article does not depend much on semantics, so columns corresponding to $s_{i}=\textup{``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 $\mathcal{H}\setminus\mathcal{H}^{\star}$, a set of hidden states which we can ignore, and only including hidden states which carry a lot of information about $M$ in $\mathcal{H}^{\star}$. In Example 4.1, when $J_{i}=2$ (location), $S_{i}=\textup{``noun''}$, the position $i$ should convey a lot about the location (in this case, “pasture”), so it is more reasonable to assume that $\{W_{:,m,h}\}_{m\in{\mathcal{M}}}$ 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}\in\mathcal{H}^{\star}$. Formally, we define the following set of input sequences containing positions $i$ where the posterior of $H_{i}$ given $x_{-i}$ concentrates on $\mathcal{H}^{\star}$:

The following theorem shows that under Assumption 4.2, we can recover $F^{\star}$ using the attention head described above, if $x\in{\mathcal{R}}$ is nonempty. Note that ${\mathcal{R}}$ is nonempty if the posterior of $H_{i}$ concentrates on $\mathcal{H}^{\star}$ 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.

Assume that non-degeneracy (Assumption 4.2) and regularity (Assumption 3.2) hold. Define ${\mathcal{R}}$ as in (4.3). Then there exist an attention head on $G(x)$ and token embeddings $e(x_{i})$ such that the following holds for any $x\in{\mathcal{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 $\textup{supp}(P[H_{i}\,|\,X_{-i}=x_{-i}])\subseteq\mathcal{H}^{\star}$. The intuition of Assumption 4.2 is that such positions are more informative for recovering the latent posteriors; indeed, from the outputs $G_{i}(x)$ at such $i$, the value function in the attention will be able to recover $P[M_{j^{\star}}\,|\,X_{1:T}=x]$. A full proof is provided in Section C.1.

2 Guarantees for prompt-tuning

Assumption 3.2 holds on the Markov chain $H_{0},H_{1},\ldots$. Furthermore, $P[H_{0}]$ is the stationary distribution: $P[H_{0}]=AP[H_{0}]$, where $A$ is the transition matrix.

As before, we assume sparsity of $\mu$ and some non-degeneracy of $W$, though the assumption is more relaxed and easier to state compared to the vanilla HMM setting.

Let ${\mathcal{M}}^{\star}\triangleq\textup{supp}(\mu)$ denote the set of non-zero coordinates in $\mu$. There exists a set of recoverable hidden states $\mathcal{H}^{\star}$, such that the collection of token emission probabilities from ${\mathcal{M}}^{\star}\times\mathcal{H}^{\star}$, $\{W_{:,(m,h)}\}_{m\in{\mathcal{M}}^{\star},h\in\mathcal{H}^{\star}}$, is linearly independent.

Furthermore, the span of these vectors must be disjoint from the span of token emission probabilities from ${\mathcal{M}}^{\star}\times(\mathcal{H}\setminus\mathcal{H}^{\star})$: $\textup{span}(\{W_{:,(m,h)}\}_{m\in{\mathcal{M}}^{\star},h\in\mathcal{H}^{\star}})\cap\textup{span}(\{W_{:,(m,h^{\prime})}\}_{m\in{\mathcal{M}}^{\star},h\in\mathcal{H}\setminus\mathcal{H}^{\star}})=\{\mathbf{0}_{|\mathcal{X}|}\}$.

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 $\textup{supp}(\mu)$, whereas Assumption 4.2 considers all $m\in{\mathcal{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^{\star}$ with soft prompt tuning and an attention head.

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 $\overline{G}(\widehat{e}(x))$ and the token embeddings which can compute the ground-truth $F^{\star}(x)$ for any $x\in{\mathcal{R}}$, defined in (4.3):

where $\widehat{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 $\textup{supp}(\mu)$. As a result, all irrelevant information to the task is removed from $\overline{G}_{i}(\widehat{e}(x))$, 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) 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}=\texttt{[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 data-generating HMM varies. The downstream label is generated by computing $\mu^{\top}P[H_{1}\,|\,X_{-1}=x_{-1}]$, where $\mu$ 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 and preprends it 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 $|\mathcal{H}|=4$ and varying memory sizes $|{\mathcal{M}}|\in\{2,3,5,7\}$. The downstream label is generated by computing $\mu^{\top}P[M\,|\,X_{-1}=x_{-1}]$, where $\mu$ denotes the ground-truth weights. Viewing the memory HMM as a HMM where the component on ${\mathcal{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 , 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.

Acknowledgements

We thank Percy Liang, Tianyi Zhang, and Nelson Liu for helpful discussions. 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.

References

Appendix 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.

In the setting of Section 3, suppose that Assumption 3.2 holds. Fix any timestep $i\geq 1$. Then there exists a diagonal matrix $D$ such that for all $x\in\textup{supp}(P[X])$,

First, we note that by Assumption 3.2, $P[H_{i}]$ has full support. As a consequence, $\textup{Pr}(X_{i+1:t+i}=x)>0$. By Bayes’ rule,

Note that the vector $\frac{P[H_{i}]}{P[H_{0}]}$ has finite and positive entries. The same applies to the ratio $r_{x}\triangleq\frac{\textup{Pr}(X_{1:T}=x)}{\textup{Pr}(X_{i+1:T+i}=x)}$. Thus, we get the desired statement. ∎

By definition, $G_{1}(x^{\prime})=P[X_{1}\,|\,X_{2:T+1}=x]$. Therefore, our goal is to rewrite $P[H_{0}\,|\,X_{1:T}=x]$ as a linear function of $P[X_{1}|X_{2:T+1}=x]$ (up to a scaling which won’t affect the linear head prediction). Concretely, we will show

for a scalar $r_{x}\geq 0$. With this equation, taking $b=\mu^{\top}B$ will give the desired result.

First, observe that $P[X_{1}\,|\,X_{2:T+1}=x]=WP[H_{1}\,|\,X_{2:T+1}=x]$ by Proposition 3.4. Next, we apply Claim A.1 to obtain an invertible matrix $D$ such that for all $x\in\textup{supp}(P[X])$, $P[H_{1}|X_{2:T+1}=x]=r_{x}DP[H_{0}|X_{1:T}=x]$, where $r_{x}>0$ is a scalar.

If $W$ has full row rank, it has a left inverse $W^{\dagger}$ with $W^{\dagger}W=I_{|\mathcal{H}|\times|\mathcal{H}|}$. Choosing $b=\mu D^{-1}W^{\dagger}$, we obtain

Next, we complete the proof of Proposition 3.4.

Appendix B Formal abstraction for prompt tuning and proofs for Section 3.1

We first formalize the definition of the model $\overline{G}$ described in Section 3.1. The model $\overline{G}$ takes a sequence of embedding vectors $v=(v_{1},\ldots,v_{t})$ as input and implements message passing to compute a sequence of $t$ outputs. We first define left and right messages $\overleftarrow{\delta}_{i+1\to i}(v)$ and $\overrightarrow{\delta}_{i-1\to i}(v)$ for $i\in[t]$, 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},\ldots$, $\tau_{i}(v)$ is always well-defined because $P[H_{i}]$ will have full support. Note that for the proper embeddings $e(x_{i})=P[X_{i}=x_{i}\,|\,H_{i}]$, where for $x=(x_{1},\ldots,x_{t})$, we use $e(x)=(e(x_{1}),\ldots,e(x_{t}))$, we can check via classical results on message passing that

Finally, we let the model model $\overline{G}$ compute

There is an edge case where the demoninator is 0, i.e. $\|\tau_{i}(v)\|_{1}=0$. To make the behavior of $\overline{G}$ well-defined, in this case we set $\overline{G}_{i}(v)=\mathbf{0}_{|\mathcal{X}|}$. We observe that if the input embedding are obtained by $e(x)$, $\overline{G}_{i}(v)$ indeed computes the desired conditional probability vector for $x\in\textup{supp}(P[X])$:

First we formalize the observation that soft prompt tuning is equivalent to adding a fake token $\widetilde{z}$ to the vocabulary with emission probabilities at timestep 1 given by $u$, and letting $\overline{G}$ compute conditional probabilities for this new distribution over sequences.

In the setting of Theorem 3.6, fix any prompt vector $u\in^{|\mathcal{H}|}$. Define the random variable $\widehat{X}$ with the same emission probabilities as $X$ for $i>1$: $P[\widehat{X}_{i}\,|\,H_{i}]=P[X_{i}\,|\,H_{i}]$. For timestep 1, we define the emission probabilities of $\widehat{X}_{1}$ as follows:

In the above equations, $\widetilde{z}$ is a fake token added to the vocabulary at timestep 1. It follows that for any $i$, defining $\tau_{i}$ as in (B.1)

As a consequence, it follows that for $i>1$ and any $x$ such that $(\widetilde{z},\varnothing,x)_{-i}\in\textup{supp}(P[\widehat{X}_{-i}])$,

For any $x$ with $(\widetilde{z},\varnothing,x)_{-i}\notin\textup{supp}(P[\widehat{X}_{-i}])$, $\overline{G}_{i}(\widehat{e}(x))=\mathbf{0}$.

Next, the following lemma disentangles the influences of the fake token $\widetilde{z}$ and the input sequence on the posterior distribution of the hidden variable.

In the setting above, there exists an invertible diagonal matrix $D$ such that for all $x$ such that $(\widetilde{z},x)\in\textup{supp}(P[\widehat{X}_{-2}])$, the following equation holds:

We now complete the proof of Theorem 3.6.

Let ${\mathcal{B}}$ be the set defined in Assumption 3.5 and define $u$ such that $u_{h}=1$ if $h\in{\mathcal{B}}$ and $u_{h}=0$ otherwise. First, we restrict our focus to $x$ such that $(\widetilde{z},x)\in\textup{supp}(P[\widehat{X}_{-2}])$. For these $x$, we can apply Lemma B.1 and Lemma B.2 in the manner described in the proof sketch. This gives $\overline{G}_{2}(\widehat{e}(x))=r_{x}WDv$ for $v\triangleq(A(u\odot P[H_{1}]))\odot P[H_{0}\,|\,X_{1:T}=x]$. By definition of ${\mathcal{B}}$, we have $\textup{supp}(A(u\odot P[H_{1}]))=\mathcal{H}^{\star}$, so $\textup{supp}(Dv)\subseteq\mathcal{H}^{\star}$. Thus, there is a matrix $\widehat{W^{\dagger}}$ such that

The existence of $\widehat{W^{\dagger}}$ is due to the fact that $\{W_{:,h}\}_{h\in\mathcal{H}^{\star}}$ is a linearly independent set of vectors, and $\textup{supp}(Dv)\subseteq\mathcal{H}^{\star}$ whenever $x$ satisfies $(\widetilde{z},x)\in\textup{supp}(P[\widehat{X}_{-2}])$. Next, we note that a matrix $B$ exists such that $(BDv)_{h}=\textup{Pr}(H_{0}=h\,|\,X_{1:T}=x)$ for $h\in\mathcal{H}^{\star}$ and $(BDv)_{h}=0$ otherwise. This is because $D$ is invertible, and $\textup{supp}(A(u\odot P[H_{1}]))=\mathcal{H}^{\star}$, so we can recover $P[H_{0}\,|\,X_{1:T}=x]$ on coordinates in $\mathcal{H}^{\star}$ by applying another coordinate-wise scaling. It follows that we can set $b=\mu^{\top}B\widehat{W^{\dagger}}$. With this choice of $b$, we compute

where the last equality follows because $\textup{supp}(\mu)\subseteq\mathcal{H}^{\star}$. This completes the case where $(\widetilde{z},x)\in\textup{supp}(P[\widehat{X}_{-2}])$.

Otherwise, for $(\widetilde{z},x)\notin\textup{supp}(P[\widehat{X}_{-2}])$, by the behavior of $\overline{G}$ in Lemma B.1, $\overline{G}_{2}(\widehat{e}(x))=\mathbf{0}$, so any linear head must output $b^{\top}\overline{G}_{2}(\widehat{e}(x))=\mathbf{0}$. Furthermore, by the conditional independence structure in $\widehat{X}$, we must also have $\textup{supp}(P[H_{2},\widehat{X}_{1}=\widetilde{z}])\cap\textup{supp}(P[H_{2},\widehat{X}_{3:T+2}=x])=\emptyset$. As $\textup{supp}(\mu)\subseteq\textup{supp}(P[H_{2},\widehat{X}_{1}=\widetilde{z}])$, this must also mean $\textup{supp}(\mu)\cap\textup{supp}(P[H_{2},\widehat{X}_{3:T+2}=x])=\emptyset$. However, we also have $P[H_{2},\widehat{X}_{3:T+2}=x]=P[H_{2},X_{3:T+2}=x]$ by the definition of $\widehat{X}$, and this must have the same support as $P[H_{0}\,|\,X_{1:T}=x]$ by applying Claim A.1 and the fact that $x\in\textup{supp}(P[X])$. It follows that for this choice of $x$, $\mu^{\top}P[H_{0}\,|\,X_{1:T}=x]=0$, so the desired statement still stands. ∎

We fill in the proofs of the lemmas below.

First, we note that (B.2) follows directly from the derivation of $\tau$, and well-known results about message passing . Next, it suffices to consider the case where $(\widetilde{z},\varnothing,x)_{-i}\notin\textup{supp}(P[\widehat{X}_{-i}])$, as the other case follows directly from the definition of $\overline{G}$ in terms of $\tau$. In this case, we observe that $\tau_{i}(\widehat{e}(x))=P[H_{i},\widehat{X}_{-i}=(\widetilde{z},\varnothing,x)_{-i}]=\mathbf{0}$. It follows that $\|\tau_{i}(\widehat{e}(x))\|_{1}=0$. Thus, from our definition of $\overline{G}$, we must have $\overline{G}_{i}(\widehat{e}(x))=\mathbf{0}$. ∎

By the conditional independence relations in a HMM, $\widehat{X}_{1}\perp\widehat{X}_{3:T+2}\,|\,H_{2}$. Using Bayes’ rule, we obtain

Where we define $r_{x}\triangleq\frac{\textup{Pr}(X_{1:T}=x)}{\textup{Pr}(\widehat{X}_{1}=\widetilde{z},\widehat{X}_{3:T+2}=x)}$. 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 $\textup{diag}(\frac{\mathbf{1}}{P[H_{0}]})$, which has finite positive entries on the diagonal by Assumption 3.2. ∎

Appendix C Proofs for Section 4

First, we introduce a proposition which is generally useful for proving the theorems in Section 4.

In the setting of Section 4, it holds that

An alternative interpretation of this statement is that $X_{i}$ is conditionally independent from everything else given $M_{J_{i}},J_{i},S_{i}$. However, we will prove this statement algebraically. We compute

Throughout this section, we use $M_{J_{i}}$ to denote the random variable obtained by indexing $M$ by $J_{i}$, both of which are themselves random variables. Let $\widehat{{\mathcal{I}}}$ denote the set of indices $i$ where $\textup{supp}(P[J_{i}\,|\,X_{-i}=x_{-i}])=\{j^{\star}\}$ and $\textup{supp}(P[S_{i}\,|\,X_{-i}=x_{-i}])\subseteq{\mathcal{S}}^{\star}$. We will first construct the key function $K$ and query $q$ such that the set of ${\mathcal{I}}$ of attended-to positions (4.2) is precisely $\widehat{{\mathcal{I}}}$. This construction does not require the position embeddings $\beta_{1},\ldots,\beta_{t}$, so we set them to $\mathbf{0}$.

The following lemma demonstrates the existence of $K$ and $q$ such that ${\mathcal{I}}=\widehat{{\mathcal{I}}}$.

The proof of Lemma C.2 requires the following claim.

In the last equality, we defined $\nu^{(h)}$ to be the expression in the parentheses. Note that $\nu^{(h)}\in{\mathcal{V}}^{(h)}\triangleq\textup{span}(\{W_{:,(m,h)}\}_{m\in{\mathcal{M}}})$. Furthermore, for $h\notin\mathcal{H}^{\star}$, $\nu^{(h)}\in\widebar{{\mathcal{V}}}\triangleq\textup{span}(\{W_{:,(m,h)}\}_{m\in{\mathcal{M}},h\in\mathcal{H}\setminus\mathcal{H}^{\star}})$. As the spans $({\mathcal{V}}^{(h)})_{h\in\mathcal{H}^{\star}}$ and $\widebar{{\mathcal{V}}}$ are all pairwise disjoint, by Assumption 4.2, for each $h\in\mathcal{H}^{\star}$, we can recover

Now we have, for $h\in\mathcal{H}^{\star}$,

Likewise, the same reasoning gives $1^{\top}\sum_{h\notin\mathcal{H}^{\star}}\nu^{(h)}=\sum_{h\notin\mathcal{H}^{\star}}\textup{Pr}(H_{i}=h\,|\,X_{-i}=x_{-i})$. Thus, we can choose $\Theta^{(1)}$ to be the matrix with rows $\Theta^{(1)}_{h,:}=\mathbf{1}^{\top}B^{(h)}$ when $h\in\mathcal{H}^{\star}$, and for some arbitrary $\widebar{h}\notin\mathcal{H}^{\star}$, $\Theta^{(1)}_{\widebar{h},:}=\mathbf{1}^{\top}\widebar{B}$. We set all other rows to $\mathbf{0}$, and we can check that this satisfies the lemma requirements.

We now construct $\Theta^{(2,h)}$. We can express $\nu^{(h)}$ in a vectorized manner by writing

We choose the first $|\mathcal{H}|$ entries of $q$ such that $q_{h}=1$ if $h=(j^{\star},s)$ for $s\in{\mathcal{S}}^{\star}$, and $q_{h}=0$ otherwise. The last entry is 0. Next, we choose $\Theta^{(K)}$ so that the first $|\mathcal{H}|$ rows are $\Theta^{(1)}$, and the last row is all zeros. where $\Theta^{(1)}$ is defined in Claim C.3. With this choice of $\Theta^{(K)}$, $K(G_{i}(x))_{h}=\textup{Pr}(H_{i}=h|X_{-i}=x_{-i})$ for $h\in\mathcal{H}^{\star}$. Furthermore, $\|K(G_{i}(x))\|_{1}=1$, by Claim C.3.

Now we note that for all $i$, $1=\|K(G_{i}(x))\|_{1}\geq q^{\top}K(G_{i}(x))$, and for $i\in\widehat{{\mathcal{I}}}$, $q^{\top}K(G_{i}(x))=\sum_{s\in{\mathcal{S}}^{\star}}\textup{Pr}(H_{i}=(j^{\star},s)|X_{-i}=x_{-i})=1$ by definition of $q$ and $\widehat{{\mathcal{I}}}$. This implies that positions $i\in\widehat{{\mathcal{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\in\widehat{{\mathcal{I}}}$.

We first choose $\Theta^{(V)}$ such that the rows satisfy $\Theta^{(V)}_{(m,j^{\star},s),:}=\Theta^{(2,s)}_{m,:}$ when $s\in{\mathcal{S}}^{\star}$ for $\Theta^{(2,s)}$ constructed in Claim C.3, and $\Theta^{(V)}_{(m,j,s),:}=\mathbf{0}_{|\mathcal{X}|}$ otherwise for $j\neq j^{\star}$ or $s\notin{\mathcal{S}}^{\star}$.

We claim that for $i\in\widehat{{\mathcal{I}}}$,

This is because for $s\in{\mathcal{S}}^{\star}$, $\Theta^{(2,s)}G_{i}(x)=P[M_{j^{\star}},H_{i}=(j^{\star},s)\,|\,X_{-i}=x_{-i}]$ by Claim C.3, and for $h=(j,s)$ for $j\neq j^{\star}$ or $s\notin{\mathcal{S}}^{\star}$,

Note that this last equality followed because $\textup{Pr}(H_{i}=h\,|\,X_{-i}=x_{-i})=0$ for the choice of $h$ and $i\in\widehat{{\mathcal{I}}}$. By construction of $\Theta^{(V)}$, these computations imply that (C.2) does indeed hold. The embedding can be chosen such that $e(x_{i})=P[X_{i}=x_{i}\,|\,M_{J_{i}},J_{i},S_{i}]$. Thus, we have for $i\in\widehat{I}$:

The last equality followed from applying the same reasoning as in Proposition C.1.

Now we pick the last linear weight in the value function by $b=B^{\top}\mu$. It follows that for $i\in\widehat{{\mathcal{I}}}$,

We obtained the last equality by observing that $\sum_{s}P[X_{i}=x_{i},M_{j^{\star}},J_{i}=j^{\star},S_{i}=s\,|\,X_{-i}=x_{-i}]=P[M_{j^{\star}},X_{i}=x_{i}\,|\,X_{-i}=x_{-i}]$ for $i\in\widehat{{\mathcal{I}}}$, as the distribution of $H_{i}$ must concentrate where $J_{i}=j^{\star}$. Finally, we observe that $\mu^{\top}P[M_{j^{\star}},X_{i}=x_{i}\,|\,X_{-i}=x_{-i}]=\mu^{\top}P[M_{j^{\star}}\,|\,X_{1:T}=x]\textup{Pr}(X_{i}=x_{i}\,|\,X_{-i}=x_{-i})$, so setting $r_{x,i}=\textup{Pr}(X_{i}=x_{i}\,|\,X_{-i}=x_{-i})$ completes the proof. ∎

Now we can complete the 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\in\textup{supp}(P[X])$ with $\widehat{{\mathcal{I}}}$ (defined in Lemma C.2) nonempty, the attended-to positions ${\mathcal{I}}$ satisfy ${\mathcal{I}}=\widehat{{\mathcal{I}}}$, and $V(G_{i}(x),e(x_{i}))=r_{x,i}\mu^{\top}P[M_{j^{\star}}\,|\,X_{1:T}=x]$. As the attention head computes the average of $V(G_{i}(x),e(x_{i}))$ over attended-to positions, and $r_{x,i}$ is positive for all $i\in\widehat{{\mathcal{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^{\star}$ always, and we only need to consider the evolution of $S_{i}$.

C.2 Formal abstraction for prompt tuning in Section 4.2

We will work directly in the case with multiple memories, as the single memory case is captured in this setting. We follow the construction in Section B. our message passing formulation requires the augmented Markov chain $\widetilde{H}_{0}\triangleq(M_{1},\ldots,M_{N},H_{0}),\widetilde{H}_{1}\triangleq(M_{1},\ldots,M_{N},H_{1}),...$, which uses the following transition probabilities:

We observe that $\eta(P[X_{i}=x_{i}\,|\,M_{J_{i}},(J_{i},S_{i})])=P[X_{i}=x_{i}\,|\,\widetilde{H}_{i}]$.

We observe that this definition almost matches Section B, except it replaces $H$ with $\widetilde{H}$. Next, we define the aggregated message at timestep $i$ by

In the edge case where $P[M]$ 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 $\tau_{i}$ in this coordinate to 0. We will see that this preserves the meaning of the message $\tau_{i}$, which for the proper embeddings $e(x_{i})=P[X_{i}=x_{i}\,|\,\widetilde{H}_{i}]$, with $e(x)=(e(x_{1}),\ldots,e(x_{t}))$, computes

We observe that $\phi(\tau_{i}(e(x)))=\frac{P[M_{J_{i}},J_{i},S_{i},X_{-i}=x_{-i}]}{|{\mathcal{M}}|^{N-1}}$. We now compute the model output as follows:

In the edge case where $\|\phi(\tau_{i}(v))\|_{1}=0$, we again define $\overline{G}(v)=\mathbf{0}_{|\mathcal{X}|}$. We can observe that $\overline{G}_{i}(e(x))=P[X_{i}\,|\,X_{-i}=x_{-i}]$.

The downstream classifier uses the embedding $\widehat{e}(x)$ defined as follows:

The dimensions of the parameters $b,\Theta^{(V)}$ 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 $\widehat{e}$:

where $\phi$ is defined in (C.4). The following assumption extends Assumption 4.5 to the multiple memory case.

Let ${\mathcal{M}}^{\star}\triangleq\textup{supp}(\mu)$ denote the set of non-zero coordinates in $\mu$. There exists a set of recoverable hidden states $\mathcal{H}^{\star}$, such that the collection of token emission probabilities from ${\mathcal{M}}^{\star}\times\mathcal{H}^{\star}$, $\{W_{:,(m,h)}\}_{m\in{\mathcal{M}}^{\star},h\in\mathcal{H}^{\star}}$, is a linearly independent set of vectors.

Furthermore, define the following span of vectors:

Then $\widebar{{\mathcal{V}}}$ must be disjoint from the span of token emission probabilities from ${\mathcal{M}}^{\star}\times\mathcal{H}^{\star}$:

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.

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 $\overline{G}(\widehat{e}(x))$ and the token embeddings which can compute the ground-truth $F^{\star}(x)$ for any $x\in{\mathcal{R}}$, defined in (4.3):

Here $\widehat{e}$ is the embedding in (4.4) and Attn is defined in (4.2).

We begin by rigorously stating the observation that soft prompt tuning is equivalent to adding a fake token $\widetilde{z}$ to the vocabulary and modifying the token emission probabilities at timestep 1, analogous to Lemma B.1.

In the setting of Theorem C.6, define $\widetilde{H}$ as in Section C.2. Fix any prompt vector $u\in^{|\widetilde{\mathcal{H}}|}$. Define the random variable $\widehat{X}$ with the same emission probabilities as $X$ for $i>1$: $P[\widehat{X}_{i}\,|\,\widetilde{H}_{i}]=P[X_{i}\,|\,\widetilde{H}_{i}]$. For timestep 1, we define the emission probabilities of $\widehat{X}_{1}$ as follows:

In the above equations, $\widetilde{z}$ is a fake token added to the vocabulary at timestep 1. It follows that for any $i$, defining $\tau_{i}$ as in (C.3)

As a consequence, it follows that for $i>1$ and any $x$ such that $(\widetilde{z},x)_{-i}\in\textup{supp}(P[\widehat{X}_{-i}])$,

For any $i$ and $x$ with $(\widetilde{z},x)_{-i}\notin\textup{supp}(P[\widehat{X}_{-i}])$, $\overline{G}_{i}(\widehat{e}(x))=\mathbf{0}$.

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 $\widehat{x}\triangleq(\widetilde{z},x_{1},\ldots,x_{t})$. The following lemma considers behaviors in edge cases with this choice of $u$.