Prefix Tuning

Prefix Tuning offers a distinct approach within parameter-efficient fine-tuning (PEFT). Unlike methods such as LoRA or Adapter Modules, which involve modifying or adding weights directly within the model architecture, Prefix Tuning keeps the original Large Language Model (LLM) parameters completely frozen. Instead of learning embeddings solely added to the input sequence, this method introduces learnable parameters that directly influence the hidden states within the transformer layers.

The core idea is to generate a small sequence of continuous vectors, called the prefix, which are prepended to the ($\mathbf{K}$) and Value ($\mathbf{V}$) matrices in the multi-head self-attention mechanisms of the transformer blocks. This prefix essentially provides learnable, task-specific context that the model can attend to at each layer, guiding its activations towards the desired output for the fine-tuning task.

Mechanism: Injecting Context into Attention

Recall the standard self-attention calculation where Query ($\mathbf{Q}$), Key ($\mathbf{K}$), and Value ($\mathbf{V}$) matrices are derived from the layer's input hidden states ($\mathbf{H}$) using learned weight matrices ($\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V$):

$$\mathbf{Q} = \mathbf{H} \mathbf{W}_Q$$ $$\mathbf{K} = \mathbf{H} \mathbf{W}_K$$ $$\mathbf{V} = \mathbf{H} \mathbf{W}_V$$

The attention scores and output are then computed based on these matrices.

In Prefix Tuning, we introduce a learnable prefix tensor $\mathbf{P}$, which holds the parameters for the prefix vectors. Let the prefix have a length $L_p$. This tensor generates layer-specific prefix vectors for keys, $\mathbf{P}_K$, and values, $\mathbf{P}_V$, both of shape ($L_p \times d_{model}$), where $d_{model}$ is the hidden dimension size of the model. These prefix vectors are then concatenated with the original $\mathbf{K}$ and $\mathbf{V}$ matrices along the sequence length dimension before the attention calculation:

$$\mathbf{K}' = \text{concat}([\mathbf{P}_K, \mathbf{K}])$$ $$\mathbf{V}' = \text{concat}([\mathbf{P}_V, \mathbf{V}])$$

The attention mechanism then operates using these modified key and value matrices:

$$\text{Attention}(\mathbf{Q}, \mathbf{K}', \mathbf{V}') = \text{softmax}\left(\frac{\mathbf{Q} (\mathbf{K}')^T}{\sqrt{d_k}}\right) \mathbf{V}'$$

where $d_k$ is the dimension of the keys.

This modification happens within each transformer layer (or a subset of layers), allowing the learned prefix to influence the model's internal representations throughout its depth. Crucially, the original weights of the LLM ($\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V$, feed-forward layers, etc.) remain unchanged during training. Only the parameters defining the prefix $\mathbf{P}$ are updated.

Flow of Prefix Tuning within a Transformer layer's attention mechanism. Learnable prefix vectors ($P_K, P_V$) derived from trained parameters are concatenated with the original Key ($K_{orig}$) and Value ($V_{orig}$) matrices before the attention computation. The base model weights ($W_Q, W_K, W_V$, etc.) remain frozen.

Parameterization and Efficiency

Directly optimizing the prefix vectors $\mathbf{P}_K$ and $\mathbf{P}_V$ for each layer can still involve a considerable number of parameters ($2 \times \text{num\_layers} \times L_p \times d_{model}$). To enhance parameter efficiency further, Prefix Tuning often employs a reparameterization technique.

Instead of directly learning the full prefix tensor $\mathbf{P}$ of shape ($L_p \times d_{model}$), a smaller matrix $\mathbf{P}'$ of shape ($L_p \times k$) is learned, where $k$ is a small intermediate dimension ($k \ll d_{model}$). This smaller matrix is then projected up to the full dimension using a simple feed-forward network (often just a linear layer, or sometimes a two-layer MLP with a non-linearity), whose weights are also learned:

$$\mathbf{P} = \text{MLP}(\mathbf{P}')$$

or simply

$$\mathbf{P} = \mathbf{P}' \mathbf{W}_{up}$$

Only the parameters of $\mathbf{P}'$ and the small projection network (e.g., $\mathbf{W}_{up}$) are trained. This significantly reduces the number of trainable parameters compared to learning $\mathbf{P}$ directly, making the method highly efficient, akin to the low-rank factorization principle used in LoRA. The total number of trainable parameters becomes independent of the number of layers if the projection network is shared, or scales linearly but with a very small constant factor otherwise.

Comparison with Prompt Tuning

Prefix Tuning shares the goal of adapting a frozen LLM by prepending learnable vectors, but differs from Prompt Tuning in important ways:

1. Location of Injection: Prompt Tuning adds learnable embeddings only to the input sequence embeddings. Prefix Tuning injects learnable vectors (prefixes) into the hidden states (specifically, Key and Value matrices) within each transformer layer's attention mechanism.
2. Expressiveness: By influencing attention computations at every layer, Prefix Tuning can potentially exert stronger control over the model's internal processing and generation compared to Prompt Tuning, whose influence might attenuate in deeper layers. This often translates to better performance, especially on complex generation tasks.
3. Number of Parameters: Prefix Tuning generally requires slightly more parameters than Prompt Tuning (due to layer-wise prefixes, even with reparameterization), but still represents a tiny fraction (often <0.1%) of the total parameters of the base LLM. Both are much more efficient than full fine-tuning.

Training and Implementation

The training process involves:

1. Initializing the prefix parameters (e.g., the small matrix $\mathbf{P}'$ and the projection network weights).
2. Freezing all parameters of the pre-trained LLM.
3. Performing standard supervised fine-tuning (e.g., minimizing cross-entropy loss on a task-specific dataset), where gradients are calculated only with respect to the prefix parameters.

Libraries like Hugging Face's PEFT provide convenient implementations. Configuring Prefix Tuning typically involves specifying the peft_type="PREFIX_TUNING", the num_virtual_tokens (which corresponds to the prefix length $L_p$), and potentially parameters related to the reparameterization network.

# Example using Hugging Face PEFT
from peft import get_peft_model, PrefixTuningConfig, TaskType
from transformers import AutoModelForCausalLM

# Load the base pre-trained model
model_name = "meta-llama/Llama-2-7b-hf"
base_model = AutoModelForCausalLM.from_pretrained(model_name)

# Configure Prefix Tuning
peft_config = PrefixTuningConfig(
    task_type=TaskType.CAUSAL_LM, # Specify task type
    num_virtual_tokens=20,       # Define prefix length (Lp)
    # Additional options for reparameterization might be available
    # e.g., prefix_projection=True
    inference_mode=False         # Set to True for inference later
)

# Wrap the base model with PEFT configuration
prefix_tuned_model = get_peft_model(base_model, peft_config)

# Print trainable parameters - notice how few they are!
prefix_tuned_model.print_trainable_parameters()
# trainable params: 9,830,400 || all params: 6,748,211,200 || trainable%: 0.145679...

# Proceed with standard training loop using 'prefix_tuned_model'
# ... (define optimizer, dataloader, training steps) ...
# Only the prefix parameters will receive gradient updates.

Advantages

Advantages:

• High Parameter Efficiency: Trains a very small number of parameters compared to the base model size.
• Resource Savings: Significantly reduces memory requirements for training and storage (only the small prefix needs saving per task).
• Preserves Base Model: The original LLM remains untouched, facilitating reuse and serving multiple tasks with a single base model instance.
• Strong Performance: Can achieve results comparable to full fine-tuning on many tasks, particularly generative ones, often outperforming Prompt Tuning due to its layer-wise influence.

Key Points:

• Hyperparameter Tuning: The prefix length (num_virtual_tokens) is a significant hyperparameter that requires tuning for optimal performance. The structure of the reparameterization network can also impact results.
• Interpretability: The learned prefix vectors lack direct human-readable meaning, unlike discrete prompts.
• Implementation Complexity: Modifying the attention mechanism might be slightly more involved than simply prepending embeddings to the input, although libraries abstract this away.

Prefix Tuning presents an option within PEFT, offering a balance between parameter efficiency and expressive power by allowing learned context to guide the internal workings of a frozen LLM at each layer. It is particularly well-suited for adapting large models to specific generative tasks when computational resources are constrained.