Prefix Tuning › The Prefix Tuning Formulation

A necessary detail distinguishes prefix tuning from simpler approaches: prefix tuning introduces separate prefix vectors at each transformer layer, not just at the input. This distinguishes it from methods that only modify the input embedding layer. For a model with $L$ layers, prefix tuning learns a complete set of layer-specific prefix matrices:

$$\{P_K^{(1)}, P_V^{(1)}, P_K^{(2)}, P_V^{(2)}, \ldots, P_K^{(L)}, P_V^{(L)}\}$$

where:

• $P_K^{(l)}, P_V^{(l)}$: learnable prefix matrices for the $l$-th transformer layer, each of shape $m \times d_k$ and $m \times d_v$ respectively
• $L$: the total number of transformer layers in the frozen pretrained model

This layer-specific design is needed for achieving strong performance, and understanding why requires appreciating how transformers process information hierarchically. Each layer of a transformer captures different levels of abstraction. Early layers (roughly the first quarter) often detect basic linguistic features like word boundaries, morphological patterns, and simple grammatical structures. Middle layers build up more complex representations involving phrases and clauses, including syntactic dependencies. Later layers capture high-level semantic relationships, discourse coherence, and task-needed features such as sentiment or logical structure. By learning prefixes at each layer, prefix tuning can guide representations at every level of the hierarchy. This gives layer-appropriate steering signals that accumulate into coherent task-specific behavior.

Think of the layer-specific design as having a different co-pilot at each altitude of flight. At low altitude, the co-pilot gives instructions about staying in the lane and avoiding obstacles (syntactic details). At cruising altitude, they give navigational guidance about heading and speed (semantic structure). At high altitude, they advise on flight plan and destination (task-level objectives). Each set of instructions is specialized for the altitude at which it operates, and their combined effect is a smooth, coherent journey to the right destination.

If we only added prefixes at the input layer, the steering signal would need to propagate through all subsequent layers without refreshing. Research on information propagation in transformers has shown that signals injected at the input layer tend to attenuate as they pass through successive attention and feed-forward operations. By injecting fresh task guidance at every layer, prefix tuning ensures that the model receives consistent direction throughout its entire processing pipeline. The prefix at layer $l$ does not need to survive unchanged through $L - l$ layers of transformation; instead, it gives appropriate guidance at exactly the point in processing where that guidance is needed.

Directly optimizing the prefix matrices often leads to unstable training. The loss surface can be non-convex, and the prefixes can drift to regions that cause attention weights to saturate or vanish. This instability occurs because the prefix vectors interact with the attention mechanism in nonlinear ways. Recall that attention weights are computed via softmax over dot products of queries and keys. Small changes to key values can produce large changes in attention probabilities when the softmax operates near saturation. A key vector that becomes very large in magnitude will dominate all attention weights, causing the model to ignore all real input tokens and attend only to that virtual token. Conversely, key vectors with small magnitudes will be ignored entirely.

To address this challenge, Li and Liang proposed a reparameterization using a small feed-forward network. The reason is that instead of directly optimizing the high-dimensional prefix vectors in their raw form, we can learn a more compact representation and project it into the required space. This bottleneck structure constrains the optimization to a lower-dimensional subspace, regularizing the solution and smoothing the loss surface. Instead of optimizing layer-specific prefixes directly, we learn a single smaller matrix $P_{\theta}$ and project it through an MLP to generate prefixes for all layers:

$$\{P_K^{(l)}, P_V^{(l)}\}_{l=1}^L = \text{Split}(\text{MLP}(P_{\theta}))$$

where:

• $P_{\theta}$: the learnable matrix of shape $m \times d'$ (where $d' \ll d$), shared across all layers as a compressed task representation
• $\text{MLP}$: a feed-forward network with a hidden layer, projecting the low-dimensional $P_{\theta}$ to the combined dimensionality required for all layer prefixes
• $\text{Split}$: an operation that divides the large MLP output into $2L$ individual matrices, one $P_K^{(l)}$ and one $P_V^{(l)}$ per layer
• $P_K^{(l)}$: the resulting prefix key matrix for layer $l$, of shape $m \times d_k$
• $P_V^{(l)}$: the resulting prefix value matrix for layer $l$, of shape $m \times d_v$

The MLP projects the shared $P_{\theta}$ from the compressed representation space to the high-dimensional space required for all layers. During training, we optimize $P_{\theta}$ and the MLP weights rather than the final prefix vectors directly. The total parameter count for the trainable component is determined by the MLP architecture and the compressed dimension $d'$, not by the full $2Ld \cdot m$ product that we would need if optimizing directly.

A necessary practical benefit: after training, we can discard the MLP entirely. We simply compute the final prefix values by running $P_{\theta}$ through the trained MLP once, storing the resulting $P_K^{(l)}$ and $P_V^{(l)}$ matrices. At inference time, we use these precomputed matrices directly. The MLP is only needed during training; it is a smooth parameterization that makes optimization more stable, after which it is no longer required.

The MLP typically has one hidden layer with a nonlinear activation:

$$\text{MLP}(x) = \tanh(xW_1 + b_1)W_2 + b_2$$

where:

• $x$: a row vector from the prefix matrix $P_{\theta}$, with dimension $d'$
• $W_1$: weight matrix of the first linear layer, of shape $d' \times d_{\text{hidden}}$, performing an initial transformation
• $b_1$: bias vector of the first linear layer, of dimension $d_{\text{hidden}}$
• $\tanh$: the hyperbolic tangent nonlinear activation, bounding outputs to $[-1, 1]$
• $W_2$: weight matrix of the second linear layer, of shape $d_{\text{hidden}} \times d_{\text{out}}$, projecting to the full prefix dimension
• $b_2$: bias vector of the second linear layer, of dimension $d_{\text{out}}$

The hyperbolic tangent activation is particularly well-suited here because it bounds outputs to the range $[-1, 1]$, preventing the prefix values from growing arbitrarily large during training. This bounded activation helps maintain numerical stability throughout the optimization process. In contrast, using a ReLU activation could allow prefix values to grow unboundedly during gradient descent, potentially causing the attention saturation problems we are trying to avoid.

This reparameterization gives several important benefits:

• Training stability: The MLP constrains the prefix space, preventing extreme values that could cause attention saturation and vanishing gradients
• Better generalization: The bottleneck dimension $d'$ acts as regularization, encouraging the model to learn compact, generalizable representations rather than overfitting to training data
• Smoother optimization: Gradients flow through the MLP with better conditioning, giving a more favorable loss surface with fewer sharp local minima and saddle points
• Reduced effective dimensionality: The number of degrees of freedom in the optimization is $m \times d'$ plus MLP parameters, rather than the full $2L \times m \times d$, making the optimization problem more tractable

The total number of trainable parameters in prefix tuning depends on prefix length $m$, model dimension $d$, number of layers $L$, and the reparameterization architecture. Understanding this count helps us compare prefix tuning's efficiency against other methods and plan resource requirements.

Without reparameterization, the inference-time parameter count is the sum of all elements in the layer-specific prefix matrices:

$$\text{Params}_{\text{inference}} = 2 \times m \times d \times L$$

where:

• $m$: prefix length (number of virtual tokens)
• $d$: model hidden dimension
• $L$: number of transformer layers
• $2$: factor for separate key and value prefixes at each layer

The factor of 2 accounts for the fact that we learn independent prefix matrices for both keys and values at each layer. Each of these matrices has $m$ rows (one per prefix position) and $d$ columns (matching the model's hidden dimension, assuming $d_k = d_v = d$ for simplicity).

For a model like GPT-2 Medium with $L = 24$ layers and $d = 1024$, a prefix length of $m = 10$ yields:

$$\begin{aligned} \text{Params} &= 2 \times m \times d \times L \\ &= 2 \times 10 \times 1024 \times 24 \\ &= 491{,}520 \text{ parameters} \end{aligned}$$

This is roughly 0.1% of GPT-2 Medium's 345 million parameters. Learning fewer than half a million parameters can substantially alter the behavior of a model with hundreds of millions of frozen parameters. This efficiency lets practical fine-tuning even in resource-constrained environments and allows a single base model to serve many different tasks simultaneously, each with its own lightweight prefix stored separately.

For training with reparameterization, we add the MLP parameters. If the compressed dimension is $d' = 512$ and the hidden dimension is $d_{\text{hidden}} = 512$, and the output dimension is $d_{\text{out}} = 2 \times L \times d$, then:

$$\text{Params}_{\text{MLP}} = m \times d' + d' \times d_{\text{hidden}} + d_{\text{hidden}} \times d_{\text{out}} + \text{biases}$$

For GPT-2 Medium with $m = 10$, $d' = 512$, $d_{\text{hidden}} = 512$, $L = 24$, $d = 1024$, this adds roughly $512 \times 512 + 512 \times (2 \times 24 \times 1024) \approx 25$ million training-only parameters. These are discarded at inference, so the inference overhead remains the small $491,520$ from the prefix matrices themselves.