PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Colin Wei, Sang Michael Xie, Tengyu Ma

arXiv:2106.09226v2cs.LGstat.ML

TL;DR

Theoretical understanding of pretrained language models is limited, in part because pretraining and downstream tasks can differ substantially. This paper links them with latent-variable generative models and analyzes head and prompt tuning; its results show recovery under specified HMM conditions, with weaker conditions for prompt tuning and memory-augmented models.

  • Problem

    Theoretical understanding of pretrained language models is scarce, and analyzing how pretraining relates to downstream tasks is challenging because the settings can differ substantially.

    Source: §1 · p. 1

  • Method

    The paper models downstream labels as functions of latent-variable posteriors and analyzes head and prompt tuning with HMM and memory-augmented HMM data models.

    Source: §1 · p. 2

  • Results

    Simple heads recover labels under strong HMM non-degeneracy conditions; prompt tuning relaxes those conditions, and memory-augmented HMMs provide stronger head-recovery guarantees.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The studied settings show downstream recovery is possible with head tuning under specified assumptions, while prompt tuning or memory structure permits weaker conditions.

    Source: §6 · p. 11

    The analysis treats pretrained models as black boxes and does not analyze full fine-tuning, whose nonlinear networks are challenging to characterize.

    Source: §6 · p. 12

Abstract

from arXiv · show

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.

1 Introduction

The paper addresses the limited theoretical understanding of how pretrained language models support downstream tasks by linking pretraining and downstream prediction through latent-variable generative models. It analyzes head and prompt tuning, reports recovery guarantees under different assumptions, and situates the analysis alongside prior empirical and theoretical work.

1 Introduction

  • The framework links pretraining and downstream prediction by modeling labels as functions of latent-variable posteriors and recovering relevant information from pretrained-model outputs.It assumes pretraining learns the generative model, so conditional token probabilities carry information about hidden variables.Source: §1 · pp. 1–2

1.1 Related works

  • Prior theoretical work assumes linear recovery from conditional token probabilities [30], whereas this paper proves such recovery for specified generative distributions.Other cited analysis studies latent recovery in a sparse Gaussian graphical model rather than downstream performance [42].Source: §1 · p. 2, §1.1 · p. 3

1 Introduction

  • Prompt tuning relaxes recovery conditions, while memory-augmented HMMs yield stronger head-tuning guarantees than vanilla HMMs.For the vanilla HMM, simple heads recover labels under strong non-degeneracy conditions; in the memory setting, an attention-based head works under weaker conditions.Source: §1 · p. 2
  • The paper evaluates its theoretical findings using language models pretrained on synthetically generated data.The introduction reports that prompt tuning performs well when non-degeneracy conditions are relaxed, while head tuning performs better with memory-augmented than vanilla HMM data.Source: §1 · p. 2

1.1 Related works

  • Prompt-based methods have improved lightweight adaptation beyond head tuning toward full fine-tuning empirically [33] [20] [9] [22] [6] [43] [2] [25].Prior work also studies discrete task prompts, continuous soft prompts, and trainable prefixes [33] [20] [9] [22].Source: §1.1 · p. 3

2 Formulations and notations

The paper formulates masked-language-model pretraining as predicting each token from the others, then defines downstream prediction and two frozen-model adaptation approaches. Soft prompts are represented as trainable embeddings prepended to the input.

2 Formulations and notations

  • The pretrained masked language model outputs a probability distribution for each token conditioned on all other tokens, and the downstream task assigns labels to token sequences.The formulation assumes exact conditional token probabilities and defines the downstream labels through a ground-truth labeling function.Source: §2 · p. 3
  • Head tuning trains a classifier on fixed model outputs, whereas soft prompt tuning jointly trains a classifier and a prompt prepended to the input embeddings.The head can also receive tokens or their embeddings; prompt tuning views the pretrained model as a function of token embeddings.Source: §2 · p. 3
  • The figure abstracts soft prompting as adding a fake vocabulary token with a trainable emission-matrix row and prepending it to the input sequence.Figure 1 also depicts the HMM and the analysis setup for head and prompt tuning.Source: Figure 1 · p. 4

3 Analysis for Hidden Markov Models

The analysis models downstream labels as functions of latent-state posteriors and proves when frozen-model outputs support recovery. Soft prompts extend recovery to HMMs with weaker emission non-degeneracy conditions by focusing outputs on task-essential information.

3 Analysis for Hidden Markov Models

  • In the vanilla HMM setting, downstream labels are linear classifiers on the posterior of H0, and the head guarantee requires linearly independent emission-matrix columns.The HMM has time-invariant hidden-state transitions and token emissions; the theorem also assumes an ergodic chain with full-support initial state distribution.Source: §3 · pp. 4–5
  • Under these conditions, a linear head on exact conditional-token probabilities recovers the downstream label for every supported input, using a shifted sequence.The proof recovers hidden-state posteriors from model outputs via a left inverse of the emission matrix.Source: §3 · p. 5

3.1 Relaxed non-degeneracy assumptions via prompt tuning

  • Soft prompt tuning relaxes the requirement to linearly independent emissions for essential hidden states containing the label-relevant information.The essential-state set must also satisfy the stated one-step reachability condition.Source: §3.1 · p. 6
  • With those conditions, a continuous prompt prepended to the input and a linear head recover the label; conditioning can remove nonessential information and ease recovery.The prompt lets the analysis recover the posterior on essential states, which suffices for the task's classifier.Source: §3.1 · p. 6

4 Analysis for memory-augmented Hidden Markov Models

The memory-augmented HMM separates evolving hidden states from persistent memory, enabling downstream recovery under weaker emission conditions than the vanilla HMM. Attention heads recover labels from informative positions, and soft prompts relax the conditions further.

4 Analysis for memory-augmented Hidden Markov Models

  • The memory-augmented HMM combines an evolving hidden-state chain with persistent memory, and emissions depend on the indexed memory cell and hidden state.Downstream labels can depend on the posterior of a particular memory cell, such as sentiment rather than syntax.Source: §4 · pp. 7–8

4.1 Tuning attention head for recovering ground-truth downstream labels

  • Its emission condition requires linear independence only for recoverable hidden states, unlike the vanilla HMM's requirement across all hidden states.The condition also requires the emission-vector spans for recoverable and non-recoverable states to be disjoint.Source: §4.1 · p. 8
  • An attention head can recover the label when some input position has a hidden-state posterior concentrated on the recoverable set.The head attends to informative positions and uses model outputs there to recover the target memory posterior.Source: §4.1 · p. 9

4.2 Guarantees for prompt-tuning

  • In the single-memory-cell setting, soft prompt tuning relaxes the emission conditions further by considering only memory values in the label-weight support.The guarantee uses a prompt and attention head, under the stated stationarity and non-degeneracy assumptions.Source: §4.2 · pp. 9–10
  • The prompt conditions memory toward label-relevant values, removing irrelevant information from model outputs and easing recovery of the memory posterior.The theorem's recovery guarantee applies to inputs in the defined set R.Source: §4.2 · p. 10

5 Simulations

Synthetic HMM experiments support the theory: prompt tuning improves downstream accuracy, particularly in degenerate settings, and memory structure improves head-tuning recovery.

5 Simulations

  • The experiments pretrain a BERT-like masked language model on synthetic HMM data to test theoretical implications when exact conditional-probability outputs may not hold.The authors report that both observed patterns support their theory.Source: §5 · p. 10
  • HMM data use 10 vocabulary symbols and length-129 sequences with a masked first token; binary labels are generated from analytically computed posteriors.The ground-truth linear weight has six nonzero entries.Source: §5 · p. 10
  • The prompt variant prepends an optimized length-20 continuous embedding, while the memory experiment varies memory size with |H| fixed at 4.Head tuning uses softmax probabilities at the masked position; the memory task labels depend on the posterior over memory variables.Source: §5 · pp. 10–11
  • Prompt tuning substantially improves downstream performance across tested hidden-state sizes, especially when |H| exceeds |X| and recovery is degenerate.The comparison uses a linear head on synthetically generated HMM data.Source: §5 · p. 11, Figure 3 · p. 11
  • Head tuning recovers the downstream task better with more structure in memory-augmented HMM data than in the vanilla-HMM comparison.The memory-augmented setup uses attention over sequence outputs; the vanilla comparison uses a first-position linear head because its task depends only on H1.Source: §5 · p. 11, Figure 3 · p. 11

6 Conclusion

The paper studies how generic language-model pretraining can support diverse downstream tasks through latent-variable posteriors. Its HMM analysis gives recovery guarantees under non-degeneracy conditions and identifies prompt tuning and memory augmentation as ways to relax them; extending the analysis to finetuning remains open.

6 Conclusion

  • The framework models downstream prediction as recovering properties of latent-variable posteriors, with generic language-model pretraining studied as support for diverse downstream tasks.For standard HMMs, a simple classification head recovers labels under strong non-degeneracy assumptions.Source: §6 · p. 11
  • Prompt tuning or a memory-augmented HMM can relax the non-degeneracy conditions required for downstream recovery.The studied generative distributions are described as a first-cut result; similar theorems for other models are conjectured as future work.Source: §6 · p. 11
  • The analysis covers head and prompt tuning as black-box methods; extending it to finetuning, especially with nonlinear activations, remains a major open question.The paper notes existing analysis of linear neural networks but says opening the pretrained-model black box for finetuning is challenging.Source: §6 · p. 12

A Proofs for Section 3

The proof expresses hidden-state posteriors from future observations as scaled diagonal transforms of the full-sequence posterior, then uses this relation to establish linear-head recovery.

A Proofs for Section 3

  • Under Assumption 3.2, the posterior of Hi given future observations is a positive scalar times a diagonal transform of the posterior of H0 given the full sequence.The claim holds for each timestep and each sequence in the observation support.Source: A Proofs for Section 3 · p. 15
  • The proof rewrites the full-sequence posterior over H0 as a scaled linear function of the pretrained model's conditional-token probabilities.The proof uses the model output for the first token conditioned on the remaining sequence.Source: A Proofs for Section 3 · p. 15
  • When the emission matrix W has full row rank, its left inverse lets the proof construct a head whose sign prediction matches the downstream label.The argument combines the posterior relation with the factorization through W and an invertible diagonal matrix.Source: A Proofs for Section 3 · p. 15

B Formal abstraction for prompt tuning and proofs for Section 3.1

The formal abstraction defines a message-passing model that takes token embeddings and computes conditional-token probabilities through left and right messages and their aggregation.

B Formal abstraction for prompt tuning and proofs for Section 3.1

  • The model G takes a sequence of embedding vectors as input and uses message passing to compute an output at each position, including conditional-token probabilities.The output is defined as the probability of Xi conditioned on all other tokens.Source: B Formal abstraction for prompt tuning… · p. 16
  • At each position, the abstraction propagates left and right messages using adjacent hidden-state transition conditionals and the embedding at that position.The boundary messages are initialized from hidden-state probabilities.Source: B Formal abstraction for prompt tuning… · p. 16
  • The model aggregates messages at each timestep and normalizes them; if the aggregate has zero norm, it defines the output as the zero vector.For proper token embeddings under Assumption 3.2, the construction computes the desired conditional probability vector on the observation support.Source: B Formal abstraction for prompt tuning… · p. 16

B.1 Proof of Theorem 3.6

The proof models soft prompt tuning by adding a fake first-step token, then constructs a linear head that recovers the downstream target from the hidden-state posterior. It handles inputs outside the relevant support separately, where both the head output and target are zero.

B.1 Proof of Theorem 3.6

  • The proof treats soft prompt tuning as adding a fake token rz with prompt-controlled first-step emissions while retaining the original emissions at later steps.The resulting model computes conditional probabilities for the modified sequence distribution.Source: §B.1 · pp. 16–17
  • A posterior identity separates the fake token's effect from the input sequence through an invertible diagonal scaling and positive scalar.This relation connects the hidden-state posterior under the prompted distribution to the original posterior given the sequence.Source: §B.1 · pp. 17–18
  • On inputs in the prompted distribution's support, the proof constructs a linear head whose output equals a positive scalar times the target posterior expectation.The construction uses the non-degeneracy set B to recover the needed posterior coordinates.Source: §B.1 · p. 17
  • Outside that support, message passing gives a zero head representation, and conditional independence implies the target posterior expectation is also zero.Thus the desired equality holds for these inputs as well.Source: §B.1 · pp. 17–18

C Proofs for Section 4

The section establishes a conditional-independence property for the memory-augmented model. Given the indexed memory, its index, and state, the current observation is conditionally independent of the remaining observations.

C Proofs for Section 4

  • The current observation Xi is conditionally independent of the other observations given the indexed memory MJi, its index Ji, and state Si.The proposition also expresses the conditional distribution of Xi as a mixture over these variables.Source: C Proofs for Section 4 · p. 18

C.1 Proof of Theorem 4.3

The proof constructs an attention head that selects positions with a uniquely identified memory index and a restricted state, then recovers a memory-based posterior prediction from them. It also formalizes prompt tuning by augmenting the hidden state with memory variables and applying message passing on that augmented chain.

C.1 Proof of Theorem 4.3

  • The key-query construction attends precisely to positions whose memory index is uniquely determined and whose state lies in S*.The corresponding value function recovers the posterior for the selected memory index at those positions.Source: §C.1 · pp. 19–20
  • The proof's attention construction relies on non-degeneracy assumptions that make the relevant emission-vector spans pairwise disjoint.This separation supports recovery of the hidden-state posterior coordinates.Source: §C.1 · pp. 19–20
  • Theorem 4.3's proof constructs key, query, and value functions so attended positions yield a positive-weighted version of the target memory posterior.For a single memory cell, the construction reduces to tracking the memory state alone.Source: §C.1 · pp. 20–21

C.2 Formal abstraction for prompt tuning in Section 4.2

  • The prompt-tuning abstraction augments the hidden state with memory variables and defines message passing on the resulting Markov chain.A lifting function maps distributions over memory and hidden state to the augmented state representation; the model uses left and right messages to compute its output.Source: §C.2 · pp. 21–22

C.3 Analysis for prompt tuning in the multiple memory setting

This section completes the prompt-tuning analysis for the multiple-memory setting, showing how a prompt and attention head compute the ground-truth downstream function through task-relevant hidden-state and memory information. It also handles zero-probability edge cases and the target memory-posterior quantity.

C.3 Analysis for prompt tuning in the multiple memory setting

  • Assumption C.5 requires linearly independent emissions for supported memory cells and recoverable states, with their span disjoint from a specified span of other emissions.It reduces to Assumption 4.5 with one memory cell and relaxes Assumption 4.2.Source: §C.3 · p. 23
  • Under Assumption C.5 and stationarity, some prompt and attention head compute the ground-truth downstream function for any input.This is Theorem C.6’s guarantee for the multiple-memory setting.Source: §C.3 · p. 23
  • Soft prompt tuning is represented as adding a fake token at timestep 1 and modifying its emission probabilities using the prompt vector.The prompt used in the proof assigns 1 to memory indices in supp(µ) and 0 otherwise.Source: §C.3 · pp. 23–24
  • When pI is nonempty, the set I of attended-to positions (defined in (4.1)) equals pI, giving the exact attention-set result under that condition.For i > 1 with Pr(p X_-i = p x_-i) > 0, the support of Pr(H_i | p X_-i = p x_-i) is contained in the support of Pr(H_{i-1} | X_{-(i-1)} = x_{-(i-1)}).Source: §C.3 · pp. 24–25
  • The value function yields a positive scalar multiple of µ weighted by the target memory-cell posterior, and returns zero for prompted inputs outside its support.For selected positions, the output is proportional to µ^T Pr(M_j* | X=x), including when the prompted context has zero probability.Source: §C.3 · p. 26
  • For prompted sequences in the support, the µ-weighted target-memory posterior equals that of the original sequence.This equality links the constructed value function to the original-input posterior used in the downstream task.Source: §C.3 · p. 26
  • For inputs with zero model probability, the target memory posterior is zero; if any leave-one-out input has zero probability, the corresponding G_i output is also zero.The claim applies to x in the support of the true input distribution when the model assigns zero probability to x.Source: §C.3 · p. 27
  • The constructed attention head selects the designated positions and averages their value outputs to recover the target memory-posterior quantity for inputs in the analyzed cases.For inputs outside the specified set, the head outputs zero, matching the zero target posterior.Source: §C.3 · p. 28

D Experimental details

The experiments use synthetically generated HMM data to pretrain a compact BERT-style model and evaluate downstream head and prompt tuning. The section details the HMM and task construction, training procedures, and posterior computation.

D Experimental details

  • HMMs use a 10-symbol vocabulary, with transition and emission rows sampled from random permutation mixtures and start distributions sampled separately.The transition and emission softmax temperatures are 0.01; the start-distribution temperature is 10.0.Source: D Experimental details · p. 28
  • Pretraining uses a six-layer BERT-base variant on 5,000 HMM-generated sequences of length 10,240, with 5% masked language modeling on 512-token chunks.Training runs for three epochs; vanilla-HMM results average 20 dataset runs, while memory-based distributions use five pretraining and fine-tuning trials.Source: D Experimental details · p. 29
  • Downstream labels come from a sparse two-row linear weight with six nonzero entries, and each task has 5,000 training, 500 validation, and 1,000 test examples.Labels are the argmax of two scores, a choice reported to reduce datasets in which all labels are the same.Source: D Experimental details · p. 29
  • Head tuning and prompt tuning use learning rate 0.01 for five epochs, and prompt tuning requires full precision because discretization zeros its small gradients.The model at the last checkpoint is used without early stopping.Source: D Experimental details · p. 29
  • Message passing in the HMM computes posterior distributions over latent variables analytically for the experiments.Source: D Experimental details · p. 29
  • Prompt tuning prepends 20 continuous vectors initialized independently from a uniform distribution on [-0.5, 0.5].The implementation uses code cited as [20].Source: D Experimental details · p. 29
Loading 2106.09226v2…