How prompt tuning works

Prompt tuning is a technique for adapting large language models (LLMs) without needing to fine-tune all the model's parameters. Instead, you add a learnable input (called "prompt parameters") to the model’s input to guide the model toward better performance on a specific task. Prompt tuning is more lightweight and efficient compared to full fine-tuning, as it doesn't modify the core model itself but optimizes only the prompt, which is often much smaller in size.

Why should you use prompt tuning?

Developers and data scientists should use prompt tuning for these two key reasons:

• Efficiency: Instead of training billions of parameters (like in full fine-tuning), we only need to optimize the prompt, which typically involves a few thousand parameters. This makes the approach much lighter and faster.
• Generalizability: The core model remains unchanged, preserving the general knowledge of the large language model while adapting it for specific tasks with a minimal overhead.

Mathematical explanation of prompt tuning

Consider a pre-trained LLM f (x; θ) where:

• x is the input text.
• θ represents the fixed parameters (weights and biases) of the large language model.
• f (x; θ) outputs a probability distribution over the next word or generates a continuation of the input text.

To understand the benefits of prompt tuning, let's compare it with fine tuning. See the following schematic that explains the difference between prompt tuning and fine tuning, from a mathematical point of view.

How prompt tuning really works

Let's dive deeper into how prompt tuning works in the context of a specific example. Suppose we are working with a LLM that is already trained on general language tasks. We want to adapt it for a specific task like sentiment analysis.

• Task: Our task in this example is: "Classify the sentiment of a given text as positive or negative."
• Input: Our original input text (x) is: "The movie was fantastic."
• The learnable prompt (p): This prompt is a learnable set of tokens or embeddings that we will tune to steer the LLM toward the sentiment classification task.
• The true label (y) is "Positive" or "1".
• Objective: To use the large language model to classify the sentiment of the input text, but instead of retraining or fine-tuning the entire model, we will learn an optimal prompt p to adjust the model’s behavior.

Let's now explore this example with a mathematical deep dive.

Step 1. Model input

Instead of directly feeding the input sentence x ("The movie was fantastic") into the model, we prepend a learnable prompt p to it. This can be thought of as adding a few extra "virtual tokens" at the start of the sentence. These prompt tokens p are vectors (embeddings) that will be learned over time to improve the task-specific performance.

So, the actual input to the model becomes a concatenation of p and x:

• p₁, p₂..., p_n are the prompt tokens, which are learned during the training process.
• "The movie was fantastic" is the input text x.
• x′ represents the full input to the language model (a combination of both p and x).

So, x′ = [p, x] = [p₁, p₂.., p_n, "The movie was fantastic"]. The following figure shows our example (for illustration purposes only) with 3-dimensional token embeddings. Each word (or token) is represented by a set of numbers called "word embeddings" (in this example, it is 3-dimensional). For example, the token p1 is represented as [0.5 0.4 0.2], the word ‘movie’ is represented by [0.4 0.3 0.1], and so on.

The true label y is explained in Step 2.

Step 2. Model operation

The LLM then processes this concatenated input x′ to produce output, ŷ:

The result of which is [0.3,0.7].

Step 3. Loss function

The model's output ŷ is compared to the true label y (either "positive" or "negative"). We compute the loss function L based on the difference between ŷ and y.

With y=1 and ŷ_positive = 0.7

The the loss function result is L=-(1)log(0.7)=0.357.

Step 4. Gradient descent and updating the prompt

The optimization process involves using gradient descent to adjust the learnable prompt p. We update the prompt embeddings p based on the gradients of the loss function with respect to p:

Where:

• η is the learning rate (a small step size that controls how much we adjust the prompt).
• ∇pL is the gradient of the loss with respect to the prompt p.

Step 5. Convergence

After enough iterations, the prompt p will converge to a set of embeddings that guide the model to classify the sentiment of the input text more accurately. Once optimized, this prompt will:

• Steer the model to classify positive sentiment correctly when it sees sentences like "The food was delicious."
• Steer the model to classify negative sentiment for sentences like "The food was terrible."

How the optimized prompt influences the output

The key idea here is that the optimized prompt p has learned task-specific information during training. This prompt p provides an additional signal to the model, guiding it toward sentiment-specific patterns in the input text. Let's break this down mathematically.

• During training, the prompt p was optimized to minimize the loss for sentiment classification. The gradient updates adjusted p in such a way that it influences the LLM to recognize patterns that correlate with positive or negative sentiment.
• The optimized prompt p essentially encodes prior knowledge about the task (sentiment classification) When the new input "The food was delicious" is concatenated with p, the model already has a "head start" in recognizing that the input might be positive due to the optimized prompt.
• Specifically, the embeddings in p might "nudge" the model toward focusing on words like "fantastic" (from previous inputs) or "delicious" (in this new input) as indicators of positive sentiment.

Implementing prompt tuning with Python

Below is a Python code example that shows the optimization of the prompt p in a simplified manner. We'll simulate a basic example where we minimize a loss function using gradient descent to update p while keeping the model parameters θ fixed.

For simplicity, we'll assume:

• Some random vector embeddings of the Input data (embedding for the sentence "The food is delicious"): [[0.2, 0.8], [0.5, 0.4], [0.9, 0.1], [0.6, 0.7]]
• Tokens have 2D embeddings for simplicity as shown above (you can use higher dimensions in practice)
• We will build a simple neural network model instead of a LLM just to show the concept of how prompt tuning works.

Let's write a few lines of code with a step-by-step approach to better understand the concept.

First, we must import the necessary libraries:

#import the necessary libraries
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, models
import matplotlib.pyplot as plt

Step 1. Define the simple model using Keras

This code block defines a simple neural network model using TensorFlow's Keras API, specifically by subclassing tf.keras.Model. This subclassed model, SimpleModel, has two fully connected (dense) layers, and it’s structured to perform binary classification, such as predicting if an input belongs to one of two classes.

• Takes an input vector of size input_size.
• Passes it through a hidden layer with hidden_size neurons and ReLU activation.
• Outputs two raw scores representing the likelihood of two classes, suitable for binary classification)

class SimpleModel(tf.keras.Model):
    def __init__(self, input_size, hidden_size):
        super(SimpleModel, self).__init__()
        self.fc1 = layers.Dense(hidden_size, activation='relu', input_shape=(input_size,))
        self.fc2 = layers.Dense(2)  # Binary classification: positive or negative

    def call(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x

Step 2. Model input and output for training

The sentence_embedding and prompt_embedding are used as input representations. The model will adjust prompt_embedding during training to better match the target label, through a loss function that measures how well the combination of sentence and prompt embeddings predicts the target sentiment

# Input data (embedding for the sentence "The food is delicious")
# Assuming 2D embeddings for simplicity (you can use higher dimensions in practice)
sentence_embedding = tf.constant([[0.2, 0.8], [0.5, 0.4], [0.9, 0.1], [0.6, 0.7]], dtype=tf.float32)

# Learnable prompt embeddings (p), we will optimize this
prompt_embedding = tf.Variable([[0.1, 0.5], [-0.4, 0.9]], dtype=tf.float32, trainable=True)

# Target label (1 for positive sentiment, 0 for negative)
target = tf.constant([1], dtype=tf.int32)  # Positive sentiment

Step 3. Model training

This code trains the model to optimize the embeddings. The model takes a combined input from "prompt embedding" and "sentence embedding" and aims to adjust the prompt embedding to achieve a certain target output over multiple training epochs.

input_size = 2  # Embedding size
hidden_size = 4
model = SimpleModel(input_size, hidden_size)
# Define the loss function (SparseCategoricalCrossentropy is equivalent to CrossEntropyLoss in PyTorch)
loss_function = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
# Define the Adam optimizer
optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
# Store the loss values for plotting
loss_values = []
# Training loop to optimize the prompt
for epoch in range(5000):
    with tf.GradientTape() as tape:
        # Concatenate prompt with sentence embedding to form the model input
        model_input = tf.concat([prompt_embedding, sentence_embedding], axis=0)
       # Forward pass: Simplified by averaging the input embeddings
        output = model(tf.reduce_mean(model_input, axis=0, keepdims=True))
        # Compute the loss
        loss = loss_function(target, output)
    # Backward pass and optimization
    gradients = tape.gradient(loss, [prompt_embedding])  # Compute gradients only for the prompt
    optimizer.apply_gradients(zip(gradients, [prompt_embedding]))  # Update the prompt embeddings

    # Store the loss value
    loss_values.append(loss.numpy())
    if epoch % 10 == 0:
        print(f'Epoch {epoch}, Loss: {loss.numpy()}, Prompt: {prompt_embedding.numpy()}')
# Plot the loss values
plt.plot(loss_values)
plt.title('Loss Function Over Epochs')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.grid(True)
plt.show()
print("Optimized Prompt:", prompt_embedding.numpy())

Output of our Python code

Here is the output of the Python code, the loss function and the optimized prompt.

Conclusion

We saw mathematically, prompt tuning adapts LLMs by introducing a set of learnable parameters p that modify the input instead of updating the model's core weights. This technique allows for task-specific adaptation with much lower computational cost than full fine-tuning.

TL;DR

• Prompt tuning is done for a specific task like fine tuning.
• Unlike fine tuning, the model is frozen (that is, the weights and biases are not changed).
• Learnable input or soft prompts are added to the input text.
• The input prompt is modified so that the model produces a desired output.