Prompt Tuning › Code Implementation
Language AI HandbookAll chapters Prompt Tuning › Code ImplementationOne section · Sections in this chapter8 / 13 ↓
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's implement prompt tuning for a text classification task to see how it works in practice. We will build a prompt-tunable wrapper around a frozen transformer model, step by step, explaining each component as we go.
import numpy as np
import torch
import torch.nn as nn # noqa: F401
import torch.nn.functional as F # noqa: F401
from transformers import AutoModel, AutoTokenizer # noqa: F401
# Set random seeds for reproducibility
torch.manual_seed(42)
np.random.seed(42)Prompt Tuning Module
The core of prompt tuning is a learnable embedding matrix that gets prepended to the input. We create a dedicated module that handles the soft prompt generation and concatenation. This encapsulation makes the approach easy to attach to any transformer model.
The key design decisions are: how to initialize the prompt embeddings (random or from vocabulary), and how to expand a single shared prompt matrix into a batch of prompts, since each example in the batch needs its own copy of the prompts (though all copies share the same learned values).
class PromptTuningEmbedding(nn.Module):
"""
Learnable soft prompts for prompt tuning.
This module creates and manages trainable soft prompt embeddings
that are prepended to the input sequence. These embeddings serve
as virtual tokens that the frozen model processes normally.
"""
def __init__(
self,
num_prompt_tokens: int,
embedding_dim: int,
init_from_vocab: bool = False,
vocab_embeddings: torch.Tensor = None,
):
super().__init__()
self.num_prompt_tokens = num_prompt_tokens
self.embedding_dim = embedding_dim
# Initialize the soft prompt embeddings
if init_from_vocab and vocab_embeddings is not None:
# Sample random tokens from vocabulary for initialization
vocab_size = vocab_embeddings.size(0)
random_indices = torch.randint(0, vocab_size, (num_prompt_tokens,))
init_embeddings = vocab_embeddings[random_indices].detach().clone()
self.prompt_embeddings = nn.Parameter(init_embeddings)
else:
# Random initialization with small scale matching embedding norms
self.prompt_embeddings = nn.Parameter(
torch.randn(num_prompt_tokens, embedding_dim) * 0.02
)
def forward(self, batch_size: int):
"""
Generate soft prompts for a batch.
Returns:
Tensor of shape (batch_size, num_prompt_tokens, embedding_dim)
"""
# Expand prompts for the batch: all examples share the same prompts
return self.prompt_embeddings.unsqueeze(0).expand(batch_size, -1, -1)Full Prompt Tuning Classifier
Now we build the complete classifier that combines the soft prompts with a frozen pretrained model. The necessary implementation detail is the attention mask: we must extend it to cover the soft prompt positions, telling the model to always attend to these virtual tokens.
class PromptTuningClassifier(nn.Module):
"""
Text classifier using prompt tuning.
Prepends learnable soft prompts to input embeddings while keeping
the pretrained transformer frozen. Only the soft prompt parameters
and the classification head receive gradient updates during training.
"""
def __init__(
self,
pretrained_model: nn.Module,
tokenizer,
num_prompt_tokens: int = 20,
num_classes: int = 2,
init_from_vocab: bool = True,
):
super().__init__()
# Get model setup
self.embedding_dim = pretrained_model.config.hidden_size
self.num_prompt_tokens = num_prompt_tokens
# Freeze the pretrained model completely
self.pretrained_model = pretrained_model
for param in self.pretrained_model.parameters():
param.requires_grad = False
# Get the embedding layer for vocabulary initialization
vocab_embeddings = None
if init_from_vocab:
vocab_embeddings = (
self.pretrained_model.embeddings.word_embeddings.weight.data
)
# Create learnable soft prompts
self.soft_prompts = PromptTuningEmbedding(
num_prompt_tokens=num_prompt_tokens,
embedding_dim=self.embedding_dim,
init_from_vocab=init_from_vocab,
vocab_embeddings=vocab_embeddings,
)
# Classification head: maps final hidden state to class logits
self.classifier = nn.Linear(self.embedding_dim, num_classes)
# Store tokenizer for reference
self.tokenizer = tokenizer
def forward(self, input_ids, attention_mask):
batch_size = input_ids.size(0)
# Step 1: Get input embeddings from the frozen model's embedding layer
input_embeddings = self.pretrained_model.embeddings.word_embeddings(
input_ids
)
# Step 2: Get soft prompts expanded for this batch
prompt_embeddings = self.soft_prompts(batch_size)
# Step 3: Concatenate: [soft_prompts; input_embeddings]
# Shape: (batch_size, num_prompt_tokens + seq_len, embedding_dim)
combined_embeddings = torch.cat(
[prompt_embeddings, input_embeddings], dim=1
)
# Step 4: Extend attention mask for soft prompt positions
# Soft prompts are always attended to (mask = 1 everywhere for them)
prompt_attention = torch.ones(
batch_size,
self.num_prompt_tokens,
device=attention_mask.device,
dtype=attention_mask.dtype,
)
extended_attention_mask = torch.cat(
[prompt_attention, attention_mask], dim=1
)
# Step 5: Forward through frozen transformer using inputs_embeds
outputs = self.pretrained_model(
inputs_embeds=combined_embeddings,
attention_mask=extended_attention_mask,
)
# Step 6: Use the [CLS] token representation for classification
pooled_output = outputs.last_hidden_state[:, 0, :]
# Step 7: Project to class logits
logits = self.classifier(pooled_output)
return logits
def count_trainable_parameters(self):
"""Count parameters that receive gradient updates."""
return sum(p.numel() for p in self.parameters() if p.requires_grad)
def count_total_parameters(self):
"""Count all parameters including frozen ones."""
return sum(p.numel() for p in self.parameters())Creating the Model
Let's instantiate the model and examine the parameter efficiency to see the large difference between trainable and total parameters.
# Load a pretrained model and tokenizer
model_name = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
base_model = AutoModel.from_pretrained(model_name)
# Create prompt tuning classifier with 20 soft prompt tokens
prompt_tuning_model = PromptTuningClassifier(
pretrained_model=base_model,
tokenizer=tokenizer,
num_prompt_tokens=20, # 20 virtual tokens prepended to input
num_classes=2, # Binary sentiment classification
init_from_vocab=True, # Initialize from vocabulary embeddings
)
trainable_params = prompt_tuning_model.count_trainable_parameters()
total_params = prompt_tuning_model.count_total_parameters()
param_percentage = 100 * trainable_params / total_params
# Verify the model is created successfully
print(f"Model created with {trainable_params:,} trainable parameters")Trainable parameters: 16,898 Total parameters: 66,379,778 Trainable percentage: 0.0255%
With just 20 soft prompt tokens on DistilBERT (embedding dimension ), the model trains approximately 0.02% of the total parameters. The soft prompt embeddings contribute parameters, and the classification head adds parameters. The roughly 66 million DistilBERT parameters remain completely frozen throughout training, never receiving any gradient signal.
Training Loop
Let's create a training example showing how prompt tuning works in practice. We use synthetic sentiment data to keep the example self-contained and reproducible.
# Create synthetic training data for demonstration
texts = [
"This movie was absolutely wonderful and entertaining!",
"Terrible film, complete waste of time and money.",
"I loved every moment of this beautiful story.",
"Boring and predictable, would not recommend.",
"A masterpiece of modern cinema, truly inspiring.",
"Awful acting and terrible script throughout.",
"Heartwarming and delightful from start to finish.",
"One of the worst movies I have ever seen.",
] * 10 # Repeat for more training data
labels = [1, 0, 1, 0, 1, 0, 1, 0] * 10 # 1 = positive, 0 = negative
# Tokenize the texts
encodings = tokenizer(
texts, padding=True, truncation=True, max_length=64, return_tensors="pt"
)
input_ids = encodings["input_ids"]
attention_mask = encodings["attention_mask"]
labels_tensor = torch.tensor(labels)# Training setup: only soft prompt and classifier parameters get updated
optimizer = torch.optim.AdamW(
prompt_tuning_model.parameters(), # Only trainable params get gradients
lr=3e-4,
weight_decay=0.01,
)
num_epochs = 50
batch_size = 16
losses = []
prompt_tuning_model.train()
for epoch in range(num_epochs):
# Simple full-batch training for demonstration
for i in range(0, len(texts), batch_size):
batch_input_ids = input_ids[i : i + batch_size]
batch_attention_mask = attention_mask[i : i + batch_size]
batch_labels = labels_tensor[i : i + batch_size]
optimizer.zero_grad()
logits = prompt_tuning_model(batch_input_ids, batch_attention_mask)
loss = F.cross_entropy(logits, batch_labels)
loss.backward()
optimizer.step()
losses.append(loss.item())
The training loss decreases smoothly. This shows that the model can learn the classification task by updating only the soft prompt embeddings. Despite training just 0.02% of parameters, the model successfully adapts to the sentiment classification task. This smooth convergence is a sign that vocabulary initialization provided a good starting point: the optimization surface near well-initialized embeddings is more amenable to gradient descent than the surface near randomly initialized ones.
Inference and Evaluation
# Evaluate on held-out examples to verify the model learned correctly
prompt_tuning_model.eval()
test_texts = [
"An incredible journey that touched my heart.",
"Disappointing and forgettable experience.",
"Best movie I've watched this year!",
"Poorly executed with no redeeming qualities.",
]
test_encodings = tokenizer(
test_texts,
padding=True,
truncation=True,
max_length=64,
return_tensors="pt",
)
with torch.no_grad():
test_logits = prompt_tuning_model(
test_encodings["input_ids"], test_encodings["attention_mask"]
)
predictions = torch.argmax(test_logits, dim=-1)Predictions: 'An incredible journey that touched my heart....' -> Positive 'Disappointing and forgettable experience....' -> Negative 'Best movie I've watched this year!...' -> Negative 'Poorly executed with no redeeming qualities....' -> Negative
The model correctly classifies the test examples, identifying positive and negative sentiments accurately. The learned soft prompts have effectively steered the frozen DistilBERT model to perform sentiment classification using only the 16,898 parameters that were allowed to change.
Comparing Initialization Strategies
Let's compare random versus vocabulary initialization to see the effect on early training dynamics. This experiment runs two otherwise identical models with different initialization strategies and tracks their loss curves.
def train_and_record(model, num_epochs=30):
"""Train model and return loss history."""
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
history = []
model.train()
for epoch in range(num_epochs):
optimizer.zero_grad()
logits = model(input_ids[:32], attention_mask[:32])
loss = F.cross_entropy(logits, labels_tensor[:32])
loss.backward()
optimizer.step()
history.append(loss.item())
return history
# Random initialization: prompts start as small Gaussian vectors
base_model_rand = AutoModel.from_pretrained(model_name)
model_random = PromptTuningClassifier(
pretrained_model=base_model_rand,
tokenizer=tokenizer,
num_prompt_tokens=20,
num_classes=2,
init_from_vocab=False, # Random initialization
)
# Vocabulary initialization: prompts start as existing token embeddings
base_model_vocab = AutoModel.from_pretrained(model_name)
model_vocab = PromptTuningClassifier(
pretrained_model=base_model_vocab,
tokenizer=tokenizer,
num_prompt_tokens=20,
num_classes=2,
init_from_vocab=True, # Vocabulary initialization
)
# Train both with identical seeds for fair comparison
torch.manual_seed(42)
history_random = train_and_record(model_random)
torch.manual_seed(42)
history_vocab = train_and_record(model_vocab)
The vocabulary-initialized model shows faster initial convergence, reaching low loss values several epochs before the randomly initialized model. Both approaches eventually converge to similar performance on this simple synthetic dataset, but the difference becomes more pronounced with smaller models, limited training data, or stricter compute budgets where fewer training steps are possible.
Examining Learned Prompts
One fascinating property of prompt tuning is that we can analyze the learned soft prompts by finding their nearest vocabulary neighbors. This offers a window into what semantic concepts the optimization has encoded, even though the prompts themselves are not discrete tokens.
def find_nearest_tokens(prompt_embedding, vocab_embeddings, tokenizer, top_k=5):
"""Find vocabulary tokens nearest to a soft prompt embedding."""
# Compute cosine similarity between prompt embedding and all vocab tokens
prompt_norm = prompt_embedding / prompt_embedding.norm()
vocab_norm = vocab_embeddings / vocab_embeddings.norm(dim=1, keepdim=True)
similarities = torch.matmul(vocab_norm, prompt_norm)
# Get top-k most similar vocabulary tokens
top_values, top_indices = torch.topk(similarities, top_k)
tokens = [tokenizer.decode([idx.item()]) for idx in top_indices]
scores = top_values.tolist()
return list(zip(tokens, scores))
# Get the learned soft prompts after training
learned_prompts = prompt_tuning_model.soft_prompts.prompt_embeddings.data
vocab_embeds = base_model.embeddings.word_embeddings.weight.data
nearest_tokens_list = []
for i in range(min(5, learned_prompts.size(0))): # First 5 prompt tokens
tokens_scores = find_nearest_tokens(
learned_prompts[i], vocab_embeds, tokenizer, top_k=3
)
nearest_tokens_list.append((i, tokens_scores))Nearest vocabulary tokens to learned soft prompts:
Prompt token 0: [('glance', 0.9519124031066895), ('glances', 0.7681981325149536), ('glancing', 0.7599889039993286)]
Prompt token 1: [('transition', 0.9572118520736694), ('transitions', 0.7650735378265381), ('transitioned', 0.7035654783248901)]
Prompt token 2: [('##dit', 0.9801098704338074), ('344', 0.6925739049911499), ('288', 0.6913912892341614)]
Prompt token 3: [('utilized', 0.960684061050415), ('utilize', 0.8766919374465942), ('utilizing', 0.8532299399375916)]
Prompt token 4: [('良', 0.9699046015739441), ('保', 0.8295562267303467), ('昭', 0.8277311325073242)]The nearest vocabulary tokens to each soft prompt reveal the semantic neighborhood that training has pushed the prompt embeddings into. For a sentiment classification task, you may observe tokens related to evaluation, opinion, or sentiment polarity near the learned prompts. However, this analysis has important limitations: soft prompts often drift to regions of embedding space that are not tightly clustered around any single vocabulary token, meaning the nearest-neighbor interpretation can be misleading. The prompt's actual influence comes from its continuous vector, not from any discrete token it happens to be nearest to.
