Prefix Tuning › Worked Example: Numerical Walk-Through

To build concrete intuition for prefix tuning, let us trace through a simplified numerical example. We will work with a tiny model to keep the numbers manageable: a single transformer layer with $d = 4$ dimensions, 1 attention head, and a prefix of length $m = 2$.

Setup. Our input is a sequence of 3 tokens with representations:

$$X = \begin{bmatrix} 1 & 0 & 0 & 0 \\ 0 & 1 & 0 & 0 \\ 0 & 0 & 1 & 0 \end{bmatrix} \in \mathbb{R}^{3 \times 4}$$

For simplicity, assume the projection matrices are the identity: $W_Q = W_K = W_V = I_4$. This means $Q = K = V = X$.

Without prefix. Standard attention computes:

$$QK^T = X \cdot X^T = \begin{bmatrix} 1 & 0 & 0 \\ 0 & 1 & 0 \\ 0 & 0 & 1 \end{bmatrix}$$

After scaling by $\sqrt{d_k} = \sqrt{4} = 2$ and applying softmax, each token attends primarily to itself (the diagonal dominates), giving a nearly diagonal attention matrix. The output is approximately:

$$\text{Attention}(Q, K, V) \approx X$$

The model retrieves mostly its own representation with slight spreading, capturing no task-specific behavior.

Adding prefix vectors. Suppose we have learned prefix matrices:

$$P_K = \begin{bmatrix} 2 & 1 & 0 & 0 \\ 0 & 2 & 1 & 0 \end{bmatrix} \in \mathbb{R}^{2 \times 4}$$ $$P_V = \begin{bmatrix} 0 & 0 & 1 & 2 \\ 0 & 0 & 2 & 1 \end{bmatrix} \in \mathbb{R}^{2 \times 4}$$

These prefix vectors stand for learned task guidance: $P_K$ encodes "key patterns to match against" while $P_V$ encodes "information to retrieve when matched."

Extended attention computation. The extended key matrix is:

$$K' = [P_K; K] = \begin{bmatrix} 2 & 1 & 0 & 0 \\ 0 & 2 & 1 & 0 \\ 1 & 0 & 0 & 0 \\ 0 & 1 & 0 & 0 \\ 0 & 0 & 1 & 0 \end{bmatrix} \in \mathbb{R}^{5 \times 4}$$

Now compute $Q \cdot (K')^T$ for the first query $q_1 = [1, 0, 0, 0]$:

$$q_1 \cdot (K')^T = [2 \cdot 1 + 1 \cdot 0 + 0 + 0, \; 0, \; 1, \; 0, \; 0] = [2, 0, 1, 0, 0]$$

After scaling by $1/\sqrt{4} = 0.5$: $[1.0, 0.0, 0.5, 0.0, 0.0]$.

Applying softmax: the first prefix position (score 1.0) receives the highest attention, followed by the third position (original first token, score 0.5). The learned prefix key $[2, 1, 0, 0]$ matches strongly with query $[1, 0, 0, 0]$ because their dot product is 2, larger than the self-attention score of 1. This shows how the prefix can "outbid" real tokens for attention: the virtual token's key vector is specifically designed to attract attention from queries that look like the first token.

Retrieving prefix values. Since the first prefix position receives the most attention weight, the output for query 1 will be a weighted combination that includes significant contribution from $P_V[0] = [0, 0, 1, 2]$. This prefix value vector injects information that was not present in the original input: the components in positions 3 and 4 that the original tokens $X$ had no content in. This is how the prefix "writes" task-specific information into the model's representations, even though the base weights remain frozen.

The key takeaway. The prefix vectors act like a learned attention bias. By engineering the key vectors to attract attention from specific query patterns and the value vectors to carry task-needed information, prefix tuning can redirect information flow within the frozen model. The model's projection matrices remain untouched, but the data flowing through attention is enriched with task-specific signals at every layer.