Fine-tuning LLMs: Soft Prompts

Implementation of soft prompting, a parameter efficient fine-tuning (PEFT) method, in PyTorch. Similar to LoRA, we actually add parameters to the base model, whilst keeping the original model frozen.

This creates a number of benefits, namely that we only fine-tune a small number of parameters relative to the total model parameters, but it also gives us the opportunity to create a "hydra" of sorts, leveraging the same base model but different soft prompt embeddings for various tasks or domains.

Soft prompting (also known as prompt tuning) is a technique where instead of fine-tuning the entire model, we prepend learnable continuous embeddings to the input sequence. These "soft prompts" are vectors in the same embedding space as the model's tokens, but they're not tied to any specific vocabulary tokens - hence "soft" rather than "hard" textual prompts.

The key insight is that we can achieve competitive performance by only training these prepended embeddings while keeping the entire pre-trained model frozen.

Specifically, we augment the original model by prepending a learnable (n × d) embedding matrix, while keeping the backbone fully frozen. Where n are "virtual tokens" and d is the embedding dimension.

The paper can be read here https://arxiv.org/pdf/2104.08691v2

Example

I'm using GPT-2 which is obviously a weak model, which is not even instruction tuned, it just dreams internet at this point, but we can still demonstrably change behaviour in useful ways.

I have a fine-tune dataset of just ~20 examples in this form:

"Q: Should I take ibuprofen on an empty stomach?\nA: I am not a medical professional. You should consult a doctor before taking medication.",

In this example my goal could be to prevent the model answering certain questions.

When given a question in the medical domain, the original GPT-2 answers:

--- Generation Before Soft Prompting ---
Original model output: Q: Can I take aspirin before bed?
A: Yes, you can.
Q: Can I take aspirin before bed?
A: Yes,

Whereas after fine-tuning we get:

--- Generation After Soft Prompting ---
Prompt-tuned model output: Q: Can I take aspirin before bed?
A: I’m not a doctor.

Btw "aspirin" is not mentioned in any of the fine-tune examples so it's cool to know that even GPT-2 has these latent understandings in its weights!

Usage

Simply run python fine_tine.py

There a number of command line args, too.

For example python fine_tune.py --model_name gpt2 --n_virtual_tokens 50

Outputs:

A:

--- Generation Before Soft Prompting ---
Original model output: Q: Can I take aspirin before bed?
A: Yes.
Q: Can I take a pill before bed?
A: Yes.
Q

--- Generation After Soft Prompting ---
Prompt-tuned model output: Q: Can I take aspirin before bed?
A: I’m not a doctor. Please consult a healthcare professional before taking any medication.

TO DO

• Demonstrate a "hydra"