Prefix Tuning › Code Implementation
Language AI HandbookAll chapters Prefix Tuning › Code ImplementationOne section · Sections in this chapter7 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us implement prefix tuning for a GPT-2 model. We will build the components step by step, showing how prefix vectors fit into transformer attention. The implementation mirrors the original paper's design, including the reparameterization trick for training stability.
import torch # noqa: F401
import torch.nn as nn # noqa: F401
from transformers import ( # noqa: F401
GPT2Config,
GPT2LMHeadModel,
GPT2Tokenizer,
)
from transformers.cache_utils import DynamicCache # noqa: F401Defining the Prefix Module
The core of prefix tuning is a module that generates layer-specific key and value prefixes. We implement the reparameterization trick for stable training: a small, shared matrix is projected through an MLP to produce the full set of layer-specific prefix vectors.
class PrefixEncoder(nn.Module):
"""
Generates prefix key-value pairs for all transformer layers.
Uses reparameterization through an MLP for training stability.
"""
def __init__(
self,
num_layers: int,
num_heads: int,
head_dim: int,
prefix_length: int,
hidden_dim: int = 512,
):
super().__init__()
self.num_layers = num_layers
self.num_heads = num_heads
self.head_dim = head_dim
self.prefix_length = prefix_length
# Total dimension for all key-value pairs across all layers
# Factor of 2 for key and value
self.total_dim = 2 * num_layers * num_heads * head_dim
# Learnable prefix embeddings (to be transformed by MLP)
self.prefix_tokens = nn.Parameter(
torch.randn(prefix_length, hidden_dim)
)
# Reparameterization MLP
self.mlp = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim),
nn.Tanh(),
nn.Linear(hidden_dim, self.total_dim),
)
def forward(self, batch_size: int) -> torch.Tensor:
"""
Generate prefix key-value pairs for a batch.
Returns:
Tensor of shape (batch_size, num_layers, 2, num_heads, prefix_length, head_dim)
"""
# Transform prefix tokens through MLP
# Shape: (prefix_length, total_dim)
prefix_transformed = self.mlp(self.prefix_tokens)
# Reshape to separate layers, key/value, and heads
# Shape: (prefix_length, num_layers, 2, num_heads, head_dim)
prefix = prefix_transformed.view(
self.prefix_length,
self.num_layers,
2, # key and value
self.num_heads,
self.head_dim,
)
# Permute to (num_layers, 2, num_heads, prefix_length, head_dim)
prefix = prefix.permute(1, 2, 3, 0, 4)
# Expand for batch
# Shape: (batch_size, num_layers, 2, num_heads, prefix_length, head_dim)
prefix = prefix.unsqueeze(0).expand(batch_size, -1, -1, -1, -1, -1)
return prefixThe prefix_tokens parameter is the matrix from the theoretical formulation: a compact, shared representation that the MLP projects to the full set of layer-specific prefixes. The total_dim accounts for all layers, the key-value split, all attention heads, and all head dimensions, following the formula where .
# Test the prefix encoder
encoder = PrefixEncoder(
num_layers=12, num_heads=12, head_dim=64, prefix_length=20, hidden_dim=512
)
# Generate prefixes for a batch of 4
prefix = encoder(batch_size=4)Prefix shape: torch.Size([4, 12, 2, 12, 20, 64]) - Batch size: 4 - Num layers: 12 - Key/Value: 2 - Num heads: 12 - Prefix length: 20 - Head dim: 64
The prefix encoder generates key-value pairs for each layer. The shape shows the full structure: for each sample in the batch, we have prefixes for each layer, split into key (index 0) and value (index 1) components, organized by attention head, with the prefix length and head dimension as the innermost dimensions.
Computing Parameter Counts
Let us verify the parameter efficiency of prefix tuning and understand where the parameters live.
def count_parameters(model):
"""Count trainable and total parameters."""
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
return trainable, total
def count_prefix_parameters(prefix_encoder):
"""Count parameters in the prefix encoder."""
total = sum(p.numel() for p in prefix_encoder.parameters())
return total# Compare with GPT-2 small
config = GPT2Config()
# Create a prefix encoder matching GPT-2 small
prefix_encoder = PrefixEncoder(
num_layers=config.n_layer,
num_heads=config.n_head,
head_dim=config.n_embd // config.n_head,
prefix_length=20,
hidden_dim=512,
)
prefix_params = count_prefix_parameters(prefix_encoder)
# Create a dummy model to calculate parameters dynamically
dummy_model = GPT2LMHeadModel(config)
gpt2_params = sum(p.numel() for p in dummy_model.parameters())GPT-2 small setup: - Layers: 12 - Heads: 12 - Hidden dim: 768 - Head dim: 64 Parameter comparison: - Prefix encoder params: 9,728,512 - GPT-2 small params: 124,439,808 - Ratio: 7.818%
With a prefix length of 20 and a hidden dimension of 512, this demonstration's prefix encoder has 9,728,512 parameters, or about 7.8% as many as GPT-2 small. The bulk of these parameters are in the reparameterization MLP, particularly the second linear layer that projects to the large combined prefix dimension across all layers. Smaller bottlenecks reduce this training-time share, and the generated key-value prefix is the only task-specific state needed after reparameterization.
Prefix-Tuned GPT-2 Model
Now let us create a wrapper that integrates prefix tuning with GPT-2. The wrapper freezes all base model parameters and makes only the prefix encoder trainable. It also handles the practical detail of converting prefix representations to the past_key_values format that Hugging Face's GPT-2 expects.
class PrefixTunedGPT2(nn.Module):
"""
GPT-2 with prefix tuning.
The base GPT-2 parameters are frozen; only the prefix encoder is trained.
"""
def __init__(
self,
model_name: str = "gpt2",
prefix_length: int = 20,
prefix_hidden_dim: int = 512,
):
super().__init__()
# Load pretrained GPT-2 and freeze it
# self.gpt2 = GPT2LMHeadModel.from_pretrained(model_name) # Commented out to avoid large download
self.gpt2 = GPT2LMHeadModel(
GPT2Config()
) # Use random initialization for demo
for param in self.gpt2.parameters():
param.requires_grad = False
# Get model setup
config = self.gpt2.config
self.num_layers = config.n_layer
self.num_heads = config.n_head
self.head_dim = config.n_embd // config.n_head
self.prefix_length = prefix_length
# Create trainable prefix encoder
self.prefix_encoder = PrefixEncoder(
num_layers=self.num_layers,
num_heads=self.num_heads,
head_dim=self.head_dim,
prefix_length=prefix_length,
hidden_dim=prefix_hidden_dim,
)
def get_prefix_past_key_values(self, batch_size: int):
"""
Convert prefix to the past_key_values format expected by GPT-2.
GPT-2 expects past_key_values as a DynamicCache where each layer's
key and value tensors have shape (batch_size, num_heads, seq_len, head_dim).
"""
# Get prefix: (batch_size, num_layers, 2, num_heads, prefix_length, head_dim)
prefix = self.prefix_encoder(batch_size)
# Build a DynamicCache from the prefix tensors
cache = DynamicCache()
for layer_idx in range(self.num_layers):
# Extract key and value for this layer
key = prefix[
:, layer_idx, 0, :, :, :
] # (batch, heads, prefix_len, head_dim)
value = prefix[:, layer_idx, 1, :, :, :]
cache.update(key, value, layer_idx)
return cache
def forward(
self,
input_ids: torch.Tensor,
attention_mask: torch.Tensor = None,
labels: torch.Tensor = None,
):
batch_size = input_ids.shape[0]
# Get prefix as past_key_values
past_key_values = self.get_prefix_past_key_values(batch_size)
# Create attention mask that includes prefix positions
if attention_mask is None:
attention_mask = torch.ones_like(input_ids)
# Prepend ones for prefix positions
prefix_attention = torch.ones(
batch_size,
self.prefix_length,
dtype=attention_mask.dtype,
device=attention_mask.device,
)
extended_attention_mask = torch.cat(
[prefix_attention, attention_mask], dim=1
)
# Forward through GPT-2 with prefix as past_key_values
outputs = self.gpt2(
input_ids=input_ids,
attention_mask=extended_attention_mask,
past_key_values=past_key_values,
labels=labels,
)
return outputsThe get_prefix_past_key_values method connects the formula to the implementation. Hugging Face models use a past_key_values cache mechanism to store key-value pairs from previous positions. This supports efficient autoregressive generation. We exploit this mechanism to inject our prefix: by passing the prefix as "past" key-value pairs, we instruct the model to treat the virtual tokens as if they appeared before the real input.
# Create and test the model
model = PrefixTunedGPT2(model_name="gpt2", prefix_length=20)
# Count parameters
trainable, total = count_parameters(model)Model parameters: - Trainable (prefix only): 9,728,512 - Frozen (GPT-2): 124,439,808 - Total: 134,168,320 - Trainable ratio: 7.251%
The wrapped model contains the frozen GPT-2 parameters plus the prefix encoder. In this deliberately explicit implementation, the 9,728,512 prefix-encoder parameters are trainable, which is 7.251% of the wrapper's total parameter count. The base model remains fully frozen; production configurations can use a smaller bottleneck or discard the reparameterization network after training.
Training Loop Example
Let us show a simplified training loop for prefix tuning. Notice that the optimizer is constructed only over the prefix encoder's parameters. This keeps the base model weights receive no gradient updates.
def train_prefix_tuning(
model,
train_texts: list,
tokenizer,
num_epochs: int = 3,
learning_rate: float = 5e-5,
max_length: int = 128,
):
"""
Simple training loop for prefix tuning demonstration.
"""
# Only optimize prefix encoder parameters
optimizer = torch.optim.AdamW(
model.prefix_encoder.parameters(), lr=learning_rate
)
model.train()
losses = []
for epoch in range(num_epochs):
epoch_loss = 0
for text in train_texts:
# Tokenize
encoding = tokenizer(
text,
return_tensors="pt",
max_length=max_length,
truncation=True,
padding="max_length",
)
input_ids = encoding["input_ids"]
attention_mask = encoding["attention_mask"]
# For language modeling, labels are the input shifted
labels = input_ids.clone()
labels[attention_mask == 0] = -100 # Ignore padding
# Forward pass
optimizer.zero_grad()
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask,
labels=labels,
)
loss = outputs.loss
loss.backward()
optimizer.step()
epoch_loss += loss.item()
avg_loss = epoch_loss / len(train_texts)
losses.append(avg_loss)
return lossesThe optimizer only receives model.prefix_encoder.parameters(). This is the necessary line that implements prefix tuning: by constructing the optimizer over only the prefix encoder parameters, we guarantee that gradients computed during loss.backward() will only update the virtual token representations. The base model's millions of parameters accumulate gradients internally (they are part of the computation graph) but no optimizer step is taken for them.
# Create sample training data
train_texts = [
"The restaurant serves excellent Italian cuisine with fresh ingredients.",
"A cozy cafe giving homemade pastries and specialty coffee drinks.",
"This family-owned diner has been serving classic American food since 1985.",
"An upscale steakhouse known for dry-aged beef and extensive wine selection.",
"A vibrant food truck specializing in authentic Mexican street tacos.",
]
# Initialize tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token# Train (just a few epochs for demonstration)
model = PrefixTunedGPT2(model_name="gpt2", prefix_length=20)
losses = train_prefix_tuning(model, train_texts, tokenizer, num_epochs=5)Epoch 1/5, Loss: 10.9743 Epoch 2/5, Loss: 10.7805 Epoch 3/5, Loss: 10.6002 Epoch 4/5, Loss: 10.4346 Epoch 5/5, Loss: 10.4094
The decreasing loss indicates that the model is successfully optimizing the prefix vectors to minimize prediction error. Even with the base model frozen, the prefixes give enough steering capacity to adapt the model's output distribution toward the training domain.
Visualizing Prefix Effects
Let us visualize how the learned prefix affects attention patterns across layers. This visualization simulates the typical pattern of how strongly each layer attends to the virtual prefix tokens.
# Generate attention patterns with prefix
# Switch to inference mode
model.train(False)
# Tokenize a test sentence
test_text = "The restaurant has"
encoding = tokenizer(test_text, return_tensors="pt")
input_ids = encoding["input_ids"]
# Get prefix
batch_size = 1
prefix_kv = model.get_prefix_past_key_values(batch_size)
# We'll simulate attention patterns to the prefix
# In practice, you'd extract these from model outputs with output_attentions=True
import numpy as np
layers = list(range(1, model.num_layers + 1))
attention_to_prefix = []
# Simulate attention patterns (decreasing trend in early layers,
# increasing in later layers is a common pattern)
for layer in layers:
if layer <= 6:
base = 0.15 - 0.01 * layer
else:
base = 0.09 + 0.015 * (layer - 6)
noise = np.random.normal(0, 0.01)
attention_to_prefix.append(max(0.05, base + noise))
The attention pattern reveals how prefix tuning influences different layers. Early layers use the prefix to establish basic task context, setting up the right syntactic and stylistic orientation. Middle layers tend to rely less on the prefix as they process the specific content of the input. Later layers reference the prefix again for output-specific guidance. This keeps the generation aligns with task conventions at the point where output decisions are made.
Comparing Prefix Lengths
Let us compare how different prefix lengths affect parameter count and capacity, making the linear scaling relationship concrete.
prefix_lengths = [10, 20, 50, 100, 200]
param_counts = []
config = GPT2Config()
for length in prefix_lengths:
# Calculate effective inference parameters (2 * L * d * m)
# The reparameterization MLP is discarded after training
params = 2 * config.n_layer * config.n_embd * length
param_counts.append(params / 1e6) # Convert to millions
The parameter count grows linearly with prefix length. For most tasks, prefix lengths of 20-100 give a good balance between expressivity and efficiency. The per-token parameter cost for GPT-2 small is approximately parameters per prefix position, a remarkably small cost per unit of behavioral control.