Prompt tuning and prefix tuning, explained

Prompt tuning and prefix tuning adapt a frozen language model by learning a handful of continuous vectors that are fed to it as extra context, instead of changing any of its weights. Prompt tuning learns $p$ "virtual token" embeddings that are prepended to the input sequence; prefix tuning learns key and value vectors that are prepended inside every attention layer. Both train somewhere between a few thousand and a few million parameters, store a task as a file measured in kilobytes or megabytes, and let one served copy of the model handle many tasks in the same batch. In the post-training stack they sit in the parameter-efficient fine-tuning (PEFT) family next to LoRA and IA3. Like those methods they are parameterisations, not objectives: the loss is usually the SFT cross-entropy, although any differentiable objective works. Today they are niche for LLM post-training, and the reasons why are as instructive as the methods themselves.

The problem and where it came from#

GPT-3 showed that a frozen model could be steered with a written prompt, but discrete prompts are brittle: small wording changes swing accuracy, and searching over tokens is a combinatorial problem. Fine-tuning, on the other hand, needed a full copy of the model per task. Two papers from 2021 proposed the middle ground of optimising the prompt itself in continuous embedding space.

Prefix tuning came first. Li and Liang (2021) froze GPT-2 and BART and trained a prefix of activations that every later token could attend to "as if it were virtual tokens". With 0.1% of the parameters it matched fine-tuning on table-to-text generation, beat it in their low-data settings, and extrapolated better to unseen topics. On XSUM summarisation it fell short (ROUGE-L 35.05 against 37.25 for full fine-tuning). They also found that optimising the prefix directly was unstable and sensitive to learning rate and initialisation, so they reparameterised it through an MLP during training.

Prompt tuning was the simplification. Lester, Al-Rfou and Constant (2021) learned only the input embeddings of the virtual tokens, with no per-layer parameters, on T5. Their headline was a scaling result: on SuperGLUE, prompt tuning trailed full model tuning badly at small sizes but "closes the gap" as models pass roughly 10 billion parameters, at under 0.01% task-specific parameters for models over a billion. Two side results hold up well: a prompt-tuned model transferred better under domain shift (a 12.5 point F1 advantage on TextbookQA after training on SQuAD), and five prompts trained on one frozen T5-XXL formed an ensemble that beat the best single prompt (91.3 against 91.0 SuperGLUE dev), at the cost of one model in memory.

The gap at small scale was then closed by going back to depth. P-Tuning v2 (Liu et al., 2021) applied learned prompts at every layer, essentially prefix tuning adapted for understanding tasks, and matched fine-tuning from 330M to 10B parameters with 0.1% to 3% tuned parameters, including hard sequence-labelling tasks where shallow prompt tuning failed.

The math, with every symbol defined#

Prompt tuning

Let a frozen decoder with parameters $\theta$ have token embedding matrix $E\in {\mathbb{R}}^{|V|\times d}$, where $d$ is the model width. For an input of $n$ tokens with embeddings $X_e\in {\mathbb{R}}^{n\times d}$ and a target $y$, prompt tuning introduces a trainable matrix $P_e\in {\mathbb{R}}^{p\times d}$ of $p$ soft-prompt vectors and feeds the concatenation $[P_e;X_e]\in {\mathbb{R}}^{(p+n)\times d}$ to the first layer. The objective is ordinary maximum likelihood, optimised over $P_e$ only:

The trainable parameter count is $pd$: for $p=20$ and $d=4096$ that is 81,920 numbers. The gradient reaches $P_e$ only through the network. Write $h_j^{(0)}$ for the input to layer 1 at prompt position $j$; then $\partial \mathcal{L}/\partial P_e[j]=\partial \mathcal{L}/\partial h_j^{(0)}$, the full backpropagated signal at that position. Every layer must be traversed backwards to compute it, which is why prompt tuning saves optimiser memory but not activation memory or backward-pass compute.

Prefix tuning

Prefix tuning acts on attention directly. In layer $\ell$, a head with query $q\in {\mathbb{R}}^{d_h}$ normally attends over the content keys $K\in {\mathbb{R}}^{n\times d_h}$ and values $V\in {\mathbb{R}}^{n\times d_h}$. Prefix tuning prepends trainable rows $P_K^{(\ell )},P_V^{(\ell )}\in {\mathbb{R}}^{p\times d_h}$ (per head, or per key-value group in grouped-query attention):

He, Zhou, Ma, Berg-Kirkpatrick and Neubig (2021) split this softmax into the part that lands on the prefix and the part that lands on the content, which turns prefix tuning into a gated additive update:

This identity is the most useful way to reason about what these methods can and cannot do. The prefix contributes a vector drawn from the span of the learned values, mixed in with a query-dependent weight $\lambda$. The relative weights the head puts on the real tokens are unchanged, because a common factor $1-\lambda$ rescales them all. Petrov, Torr and Bibi (2023) turned this into a theorem: soft prompts and prefixes cannot change the relative attention pattern over the content and can only bias each attention layer's output, so they can elicit skills the pretrained model already has but may be unable to learn tasks that need new attention patterns. Prompt tuning is weaker still: its vectors enter only at layer 0, and the deeper keys and values they produce are fixed functions of those inputs.

Side by side with its siblings

The counts assume 32 layers, $d=4096$, key-value width $d_{kv}=1024$ and MLP width 14,336. The ordering is the usual one for expressiveness as well as cost, with one caveat worth keeping: parameter count is not the right measure of capacity here. A prompt or prefix can be large and still unable to represent an update that a rank-1 LoRA represents easily, because of the structural limit above.

Worked example: a prefix in one attention head#

Take one head with $d_h=2$, and skip the $1/\sqrt{d_h}$ scale to keep the arithmetic clean. The query is $q=(1,0)$. There are two content tokens with keys $k_1=(1,0)$, $k_2=(0,1)$ and values $v_1=(1,0)$, $v_2=(0,1)$. Without a prefix, the scores are $q\cdot k_1=1$ and $q\cdot k_2=0$, the softmax weights are $e/(e+1)=0.731$ and $0.269$, and the head outputs $(0.731,0.269)$.

Now add one prefix slot with key $P_K=(2,0)$ and value $P_V=(0,-1)$. The scores become $(2,1,0)$ and the softmax gives weights $(0.665,0.245,0.090)$, so the output is $0.665\cdot (0,-1)+0.245\cdot (1,0)+0.090\cdot (0,1)=(0.245,-0.575)$. Using the gated form: $\lambda =e^2/(e^2+e+1)=7.389/11.107=0.665$, and $(1-\lambda )(0.731,0.269)+\lambda (0,-1)=(0.245,-0.575)$, the same number. Notice the ratio of weights on the two content tokens: $0.245/0.090=2.718=e$ with the prefix, $0.731/0.269=e$ without. The prefix moved the output a long way, but it did not change which content token the head prefers. If a new task needs this head to look at $k_2$ instead of $k_1$, no prefix can make it do so; a weight update to $W_Q$ or $W_K$ can.

A toy run: eliciting a skill versus learning a new one#

To see the elicit-versus-teach distinction directly, we pretrained a small decoder (2 layers, width 64, 4 heads, PyTorch, seed 0) on six-digit sequences preceded by a task token for copy, reverse or sort, with a random-length filler prefix so that the content positions did not depend on the context length. After 1,200 steps it scored 99.9%, 99.95% and 96.8% exact match on the three tasks. Without the task token it reversed only 1.0% of inputs correctly. We then froze it and adapted it, with no task token, to two tasks: reverse, which it already knows, and rotate-by-one (output the digits shifted left by one place), which it never saw and solved 0.05% of the time. Prompt tuning used 10 soft tokens initialised from vocabulary embeddings (640 parameters), prefix tuning 10 key and value slots per layer (2,560), and LoRA rank 4 on every linear map (9,216), each for 1,500 Adam steps at batch 128.

All three methods recovered the existing skill completely: 100% on reverse. The new task separated them. Prompt tuning reached 26.8% at our default settings, and a sweep over learning rate (1e-2 to 1e-1) and length (up to 4,000 steps) topped out at 30.9%, with training loss stuck between 0.35 and 0.75. Prefix tuning and LoRA both reached 100%. The prefix result is worth a moment, because it seems to contradict the expressiveness limit above. It does not: the theorem constrains one attention layer, and in a 2-layer model the prefix's bias in layer 1 changes the inputs from which layer 2 computes its queries and keys, so the second layer's attention pattern can change after all. Prompt tuning's vectors enter only at the input, where they compete with the content for attention in layer 1 and must route all their influence through it. The lesson for practice is the ordering, not the specific numbers: soft prompts are an excellent switch for behaviour the model has, prefixes buy real extra capacity through depth, and weight updates are the safe choice for anything new.

The algorithm and training loop#

load base model; freeze every parameter
prompt tuning:  P_e = embeddings of p real tokens (task description or sampled vocab)   # p x d
prefix tuning:  P'   = small matrix (p x d'), MLP: d' -> 2 * L * d_kv                    # reparameterisation
optimizer = Adam or Adafactor over P_e (or P' and MLP), lr high (1e-3 to 0.3)

for step in range(total_steps):
    x, y = next(batch)
    if prompt tuning: inputs = concat(P_e, embed(x))            # shift attention mask and positions by p
    if prefix tuning: past_kv = reshape(MLP(P'))                # p entries per layer, per KV head
    loss = cross_entropy(model(inputs, past_kv)[response positions], y)
    loss.backward()                                             # gradients traverse every frozen layer
    optimizer.step(); optimizer.zero_grad()

after training (prefix): compute P_K, P_V = MLP(P') once, drop the MLP, store p x 2 x L x d_kv numbers
serve: prepend the stored vectors per request; different requests in a batch can use different prompts

Three implementation details decide whether a run works. First, initialisation: Lester et al. found that random initialisation lagged behind initialising from sampled vocabulary or from the embeddings of the class labels, with the difference vanishing at T5-XXL; Li and Liang found that initialising the prefix with activations of real words was significantly better than random. Initialising from the embeddings of a short instruction that describes the task is the sensible default for LLMs. Second, positions and masks: soft tokens occupy positions, so the attention mask, position ids and loss mask must all be shifted by $p$. Prefix entries live in the KV cache and do not consume input positions in most implementations. Third, the reparameterisation for prefix tuning, which smooths an otherwise ill-conditioned optimisation.

Data requirements are those of SFT. Prompt tuning is at its best with a few hundred to a few tens of thousands of examples of one narrow task (classification, extraction, a fixed output format) on a large model. It needs more steps than LoRA to converge: Lester et al. trained for 30,000 steps at a constant learning rate of 0.3 with batch size 32 and Adafactor, and that is a fair indication of the optimisation budget.

Hyperparameters that matter#

Tune the learning rate first, on a log grid spanning at least two orders of magnitude, because the right value depends heavily on the embedding scale of the model. Then tune $p$. If doubling $p$ does not help, more parameters will not help either; switch to prefix tuning or LoRA.

Compute and memory cost#

With $P$ frozen base parameters in bf16, training memory is roughly

where $N_{\text{soft}}$ is the number of soft parameters (tiny) and $M_{\text{act}}(n+p)$ is the activation memory at the lengthened sequence. The optimiser state is negligible, so the bill is the frozen model plus activations, the same profile as LoRA. Backward compute is also similar to LoRA: activation gradients through every layer, no weight gradients, about $4PD$ FLOPs for $D$ tokens against $6PD$ for full fine-tuning. Prompt tuning lengthens every sequence by $p$ tokens, which costs roughly $p/n$ extra compute per sequence in training and serving (1% for $p=20$ and 2,000-token inputs). Prefix tuning lengthens the KV cache by $p$ entries per layer, but adds no forward computation for those positions.

Storage is where these methods win outright. A 20-token prompt for an 8B model is 81,920 numbers, 160 KB in bf16. A prefix of 20 slots is about 2.6 MB. A rank-16 LoRA on all linear layers is about 84 MB. If you need to store tens of thousands of per-user or per-tenant adaptations, that difference decides the architecture.

Failure modes and how to detect them#

The loss falls slowly and plateaus high

The most common outcome on mid-sized models. T-Few (Liu et al., 2022) was unable to reach reasonable accuracy with prompt tuning on T0-3B and saw validation accuracy fluctuate wildly. Diagnose with a learning rate sweep and a LoRA baseline at matched training steps. If LoRA reaches a much lower loss quickly and prompt tuning does not at any learning rate, the task needs a capability change, not a cue.

Instability and seed sensitivity

Different seeds give very different validation scores, and loss curves show spikes. Li and Liang's reparameterisation exists for this reason. Use the MLP reparameterisation for prefixes, initialise from real token embeddings, clip gradients, and report results over three seeds.

Format failures from the base model's prior

Lester et al. found that T5 checkpoints trained only on span corruption often never learned to output a legal class label under prompt tuning, copying input spans or predicting an empty string instead; a short LM-adaptation stage fixed it. The LLM analogue is a base model whose prior strongly favours some output style. Check sampled outputs early, not just the loss.

Silent mask and position bugs

Loss computed over the soft-token positions, or attention masks not extended by $p$, produce runs that train but evaluate badly. Assert the shapes of the extended mask and the label mask in the first batch.

When to use it, and what to use instead#

Prompt tuning is the right tool when the task is something the model can already do and needs to be reliably cued into, when storage per task must be tiny, or when many tasks must share one batch on one frozen model. Prefix tuning is the better variant when you want more capacity at the same deployment shape. For post-training an LLM to follow instructions, reason, prefer one answer over another, or learn new knowledge, low-rank weight updates are more expressive per parameter and easier to optimise, and they dominate open-source practice.

Minimal implementation#

Prompt tuning on a Hugging Face causal LM, written out so the masks and shifts are visible. The PEFT library implements the same thing with PromptTuningConfig and prefix tuning with PrefixTuningConfig.

import torch, torch.nn as nn

class SoftPrompt(nn.Module):
    def __init__(self, model, tokenizer, init_text="Classify the sentiment of the review:", p=20):
        super().__init__()
        emb = model.get_input_embeddings()
        ids = tokenizer(init_text, add_special_tokens=False).input_ids
        ids = (ids * (p // len(ids) + 1))[:p]                       # repeat the text to fill p slots
        self.prompt = nn.Parameter(emb.weight[ids].detach().clone().float())
        self.model, self.p = model, p
        for prm in model.parameters():
            prm.requires_grad_(False)

    def forward(self, input_ids, attention_mask, labels=None):
        b = input_ids.shape[0]
        x = self.model.get_input_embeddings()(input_ids)
        soft = self.prompt.to(x.dtype).unsqueeze(0).expand(b, -1, -1)
        x = torch.cat([soft, x], dim=1)
        mask = torch.cat([attention_mask.new_ones(b, self.p), attention_mask], dim=1)
        if labels is not None:                                       # never compute loss on soft positions
            labels = torch.cat([labels.new_full((b, self.p), -100), labels], dim=1)
        return self.model(inputs_embeds=x, attention_mask=mask, labels=labels)

# sp = SoftPrompt(model, tok); opt = torch.optim.Adam([sp.prompt], lr=3e-3)
# with PEFT:
# from peft import PromptTuningConfig, PromptTuningInit, get_peft_model
# cfg = PromptTuningConfig(task_type="CAUSAL_LM", num_virtual_tokens=20,
#                          prompt_tuning_init=PromptTuningInit.TEXT,
#                          prompt_tuning_init_text="Classify the sentiment of the review:",
#                          tokenizer_name_or_path=model_id)
# model = get_peft_model(model, cfg)
# prefix tuning: PrefixTuningConfig(task_type="CAUSAL_LM", num_virtual_tokens=20, prefix_projection=True)

Keep the soft prompt in fp32 even when the model runs in bf16; the updates are small relative to the embedding norm and bf16 rounding can swallow them. In POFE, prompt tuning is a deliberate choice for narrow, high-cardinality adaptation, and the recipe runs a LoRA baseline alongside it so the capability question gets answered on the first run.

Further reading#

Li and Liang (2021) and Lester, Al-Rfou and Constant (2021) are the two founding papers and both remain readable; the second has the scaling curves. Liu et al. (2021, P-Tuning v2) explains why depth closes the small-model gap, and Liu et al. (2021, "GPT Understands, Too") is the earlier P-Tuning paper with a learned prompt encoder. He et al. (2021) gives the gated-attention identity that unifies prefix tuning with adapters and LoRA. Petrov, Torr and Bibi (2023) proves the expressiveness limits. Liu et al. (2022, T-Few) is a useful counterweight: a careful few-shot comparison where prompt tuning lost to IA3 and LoRA.

Frequently asked questions#

References#

1. Xiang Lisa Li, Percy Liang (2021). Prefix-Tuning: Optimizing Continuous Prompts for Generation. arxiv.org/abs/2101.00190
2. Brian Lester, Rami Al-Rfou, Noah Constant (2021). The Power of Scale for Parameter-Efficient Prompt Tuning. arxiv.org/abs/2104.08691
3. Xiao Liu, Kaixuan Ji, Yicheng Fu, Weng Lam Tam, Zhengxiao Du, Zhilin Yang, Jie Tang (2021). P-Tuning v2: Prompt Tuning Can Be Comparable to Fine-tuning Universally Across Scales and Tasks. arxiv.org/abs/2110.07602
4. Xiao Liu, Yanan Zheng, Zhengxiao Du, Ming Ding, Yujie Qian, Zhilin Yang, Jie Tang (2021). GPT Understands, Too. arxiv.org/abs/2103.10385
5. Junxian He, Chunting Zhou, Xuezhe Ma, Taylor Berg-Kirkpatrick, Graham Neubig (2021). Towards a Unified View of Parameter-Efficient Transfer Learning. arxiv.org/abs/2110.04366
6. Aleksandar Petrov, Philip H. S. Torr, Adel Bibi (2023). When Do Prompting and Prefix-Tuning Work? A Theory of Capabilities and Limitations. arxiv.org/abs/2310.19698
7. Haokun Liu, Derek Tam, Mohammed Muqeeth, Jay Mohta, Tenghao Huang, Mohit Bansal, Colin Raffel (2022). Few-Shot Parameter-Efficient Fine-Tuning is Better and Cheaper than In-Context Learning. arxiv.org/abs/2205.05638

Related reading