We can't find the internet
Attempting to reconnect
Something went wrong!
Attempting to reconnect
← Parameter-Efficient FT step 3 of 4
Prefix Tuning
Prefix tuning adapts a model by learning extra vectors while keeping the base model's parameters fixed. This exercise implements one attention forward pass with supplied prefix keys and values. It does not train the prefix or run a full language model.
Implement prefix_tuning_attention(Q, K, V, prefix_K, prefix_V).
Inputs and output
- Q, K and V have shape (S, D): S original positions, each of width D.
- prefix_K and prefix_V have shape (P, D): P matching prefix key/value pairs.
- S and D are positive; P may be zero.
- All inputs share a floating-point dtype and device.
Return shape (S, D), preserving dtype and device and leaving inputs unchanged. The prefix adds keys and values, not queries, so it does not add output rows.
Attention with a prefix
Prepend prefix_K to K along the position axis, and prepend prefix_V to V in the same order. Both combined tensors have shape (P + S, D).
full_K = concatenate([prefix_K, K], axis=0)
full_V = concatenate([prefix_V, V], axis=0)
scores = Q @ full_K.T / sqrt(D) # (S, P + S)
weights = softmax(scores, axis=-1) # normalize each query row
output = weights @ full_V # (S, D)
Each query attends to all prefix and original key/value pairs. There is no causal mask, padding mask, dropout or output projection in this exercise. Normalize across all P+S keys together, not separately for the prefix and original sequence. Use a stable softmax that subtracts each row's maximum before exponentiating.
For one query with equal scores for a prefix value [20] and an original value [10], the weights are [0.5, 0.5] and the output is [15]. A zero query gives equal scores to every key and returns their values' average. With P=0, this reduces to ordinary unmasked attention over K and V.
The key/value tensors are already projected into attention space; do not treat prefix vectors as token IDs or prepend anything to Q.
See Li and Liang, Prefix-Tuning: Optimizing Continuous Prompts for Generation.
Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.
Train a smaller set of parameters
Low-rank adaptation, or LoRA, learns a change to a weight matrix using two
smaller matrices. For row-vector inputs and weight shape (d_in, d_out), one
convention is
$$W_{effective}=W+\frac{\alpha}{r}AB,$$
where A has shape (d_in, r) and B has shape (r, d_out). The base weight
W can remain fixed while A and B learn.
For d_in = d_out = 100 and r = 4, the factors contain 400 + 400 = 800
parameters, compared with 10,000 in the full matrix. The saving comes with a
constraint: the update has rank at most r. Follow the task's matrix orientation
and scale, since implementations use different conventions.
Start with an initialization that allows learning
If both factors are zero, both receive zero gradients through their product. A common initialization makes one factor random and the other zero. The initial weight change is still zero, but the zero factor can receive a nonzero gradient and start learning.
Distinguish other adaptation methods
An adapter applies a small network to intermediate activations, often reducing then restoring the feature width in a residual branch. Prefix tuning learns context such as key/value vectors that tokens can attend to. These approaches modify different parts of the model and are not substitutes for one another.
To freeze base weights, ensure the optimizer cannot move them through momentum or weight decay as well as through current gradients. Check that base values stay unchanged after an update, while the loss still has a gradient path to the trainable factors or prefix.
Your code, saved here
Sign in to access your accepted solutions and saved draft history. You can keep reading the lesson without an account.
Sign in to your accountStuck?
PyTorch reference solution
Sign in to track your attempts toward unlocking the reference solution.