SetFit: Efficient Few-Shot Learning Without Prompts

Introduction

Related work

In-context learning

Pattern Exploiting Training (PET)

PET was described in the paper Exploiting Cloze Questions for Few Shot Text Classification and Natural Language Inference. ADAPET is a improved version of PET which allowed it to train more easily and requires less task specific data.

Parameter Efficient Fine-tuning (PEFT)

SetFit: Sentence Transformer Fine-Tuning

SetFit's fine-tuning and training block diagram

ST fine-tuning

To handle the limited amount of labelled data, authors adopt a contrastive training approach often used for image similarity.
Given a small set of $K$ labeled examples $D = {\{(x_i, y_i)\}}$, where $x_i$ and $y_i$ are sentences and their class labels, respectively, for each class label $c \in C$, the authors generate a set of $R$ positive triplets; $T^c_p = {\{(x_i, x_j, 1)\}}$ (pairs of sentences randomly chosen from the same class $c$).
Similarly, a set of $R$ negative triplets are generated as $T^c_n = {\{(x_i, x_j, 0)\}}$ (pairs of sentences randomly chosen from different classes).
Finally, the contrastive fine-tuning data set $T$ is produced by concatenating the positive and negative tripletes across all class labels; $T=\left\{\left(T_p^0, T_n^0\right),\left(T_p^1, T_n^1\right), \ldots,\left(T_p^{|C|}, T_n^{|C|}\right)\right\}$, where $|C|$ is the number of class labels, $|T| = 2R|C|$ is the number of pairs.
We can see that this technique of creating data enlarges the data in few-shot settings. For example: assuming that a small number $(K)$ of labeled examples are given for a binary classification task, the potential size of the ST fine-tuning set $T$ is derived from the number of unique sentence pairs that can be generated, namely $K(K-1)/2$, which is significantly larger than $K$.
In the process of contrastive fine-tuning, the functionality of ST is changed from generating sentence embedding to topic embedding.

Classification head training

In the second stage (after the ST is fine-tuned contrastively), the original labeled training data ${\{x_i}\}$ is encoded into a vector (sentence embedding) per training sample; $Emb^{x_i} = ST(x_i)$ where $ST()$ is the function representing the fine-tuned ST.
The embeddings, along with their class labels, constitute the training set for the classification head $T^{CH} = {\{(Emb^{x_i}, y_i \}}$.
The authors used a logistic regression model as the classification head (however, we can use anything we want).

Inference

How to train a SetFit model for few-shot classification

from datasets import load_dataset
from sentence_transformers.losses import CosineSimilarityLoss

# Load a dataset from the Hugging Face Hub
dataset = load_dataset("sst2")

# Simulate the few-shot regime by sampling 8 examples per class
train_dataset = sample_dataset(dataset["train"], label_column="label", num_samples=8)
eval_dataset = dataset["validation"]

# Load a SetFit model from Hub
model = SetFitModel.from_pretrained("sentence-transformers/paraphrase-mpnet-base-v2")

It takes in the model, train dataset and evaluation dataset.
loss_class: these are the loss classes from the sentence_transformer package used for contrastive learning. In contrastive learning, there are different losses that have been proposed for how you measure the distance between positive and negative samples. The authors found out that CosineSimilarityLoss works quite well.
We need to also provide the metric in the metric field to evaluate our models on. We need to just pass the name of the metric in the form of string and it'll be loaded from the Evaluate library of HuggingFace. You can have a look at different metrics it supports here.
num_iterations: this specifies how many times the sampling process should be carried out to create the augmented data for contrastive learning (This is the $R$ hyperparameter as discussed above in the explanation of the SetFit method).
num_epochs: the number of epochs to use for contrastive learning.
column_mapping: the SetFit expects your texts and labels into columns named as text and label respectively. If the dataset doesn't have the column named the same way, we can provide the mapping to align the column names.

# Create trainer
trainer = SetFitTrainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss_class=CosineSimilarityLoss,
    metric="accuracy",
    batch_size=16,
    num_iterations=20, # The number of text pairs to generate for contrastive learning
    num_epochs=1, # The number of epochs to use for constrastive learning
    column_mapping={"sentence": "text", "label": "label"} # Map dataset columns to text/label expected by trainer
)

# Train and evaluate
trainer.train()
metrics = trainer.evaluate()

# Push model to the Hub
trainer.push_to_hub("my-awesome-setfit-model")

# Download from Hub and run inference
model = SetFitModel.from_pretrained("<your_hf_username>/my-awesome-setfit-model")
# Run inference
preds = model(["i loved the spiderman movie!", "pineapple on pizza is the worst 🤮"]) 

Using a end-to-end network with differentiable head.

num_classes = 2
# Load a SetFit model from Hub (with differentiable head)
model = SetFitModel.from_pretrained(
    "sentence-transformers/paraphrase-mpnet-base-v2",
    use_differentiable_head=True,
    head_params={"out_features": num_classes},
)

# Train and evaluate
trainer.freeze() # Freeze the head
trainer.train() # Train only the body

# Unfreeze the head and freeze the body -> head-only training
trainer.unfreeze(keep_body_frozen=True)
# or
# Unfreeze the head and unfreeze the body -> end-to-end training
trainer.unfreeze(keep_body_frozen=False)

trainer.train(
    num_epochs=25, # The number of epochs to train the head or the whole model (body and head)
    batch_size=16,
    body_learning_rate=1e-5, # The body's learning rate
    learning_rate=1e-2, # The head's learning rate
    l2_weight=0.0, # Weight decay on **both** the body and head. If `None`, will use 0.01.
)

Scikit-learn head vs Differentiable head

SentEval-CR
CR
sst2
sst5
bbc-news
onestop_english
amazon_counterfactual_en
enron_spam
subj

Colab notebook for comparing Scikit-learn head and differentiable head ->

Hyperparameter search using W&B Sweeps

Parallel coordinates charts summarize the relationship between large numbers of hyperparameters and model metrics at a glance.
Parameter Importance chart tells which of your hyperparameters were the best predictors of, and highly correlated to desirable values of your metrics. To learn about how to read these charts, you can refer the documentation here.

Key Insights from our Sweep

batch_size have a strong negative correlation with the accuracy as seen in the parallel coordinates and parameter importance chart. Higher batch size leads to a lower accuracy.
learning_rate too have a strong negative correlation with the accuracy. Higher learning rates degrade the performance.

Compressing a SetFit model with knowledge distillation

from datasets import load_dataset
from sentence_transformers.losses import CosineSimilarityLoss
from setfit import SetFitModel, SetFitTrainer, DistillationSetFitTrainer, sample_dataset

# Load a dataset from the Hugging Face Hub
dataset = load_dataset("ag_news")
# Create a sample few-shot dataset to train the teacher model
train_dataset_teacher = sample_dataset(dataset["train"], label_column="label", num_samples=16)
# Create a dataset of unlabeled examples to train the student
train_dataset_student = dataset["train"].shuffle(seed=0).select(range(500))
# Dataset for evaluation
eval_dataset = dataset["test"]

# Load teacher model
teacher_model = SetFitModel.from_pretrained(
    "sentence-transformers/paraphrase-mpnet-base-v2"
)
# Create trainer for teacher model
teacher_trainer = SetFitTrainer(
    model=teacher_model,
    train_dataset=train_dataset_teacher,
    eval_dataset=eval_dataset,
    loss_class=CosineSimilarityLoss,
)
# Train teacher model
teacher_trainer.train()

The student model is trained using sentence pairs and the level of similarity between each pair as input. The similarity is generated by using the underlying sentence transformer (ST) of the teacher model (which is used in the first stage of the training, before the classification head) to produce semantic embeddings for each pair and to calculate cosine-similarity between them.
The underlying ST of the student model is trained to mimic the ST of the teacher output by minimizing the error between the teacher model produced cosine-similarity and its output.
The classification head of the student model is then trained using the embeddings produced by the student's ST and the logits produced by the SetFit teacher classification head.
The baseline student is trained to mimic the teacher output by minimizing the error between logits produced by the SetFit teacher classification head and its output.

# Load small student model
student_model = SetFitModel.from_pretrained("paraphrase-MiniLM-L3-v2")

# Create trainer for knowledge distillation
student_trainer = DistillationSetFitTrainer(
    teacher_model=teacher_model,
    train_dataset=train_dataset_student,
    student_model=student_model,
    eval_dataset=eval_dataset,
    loss_class=CosineSimilarityLoss,
    metric="accuracy",
    batch_size=16,
    num_iterations=20,
    num_epochs=1,
)

# Train student with knowledge distillation
student_trainer.train()

SetFit Efficiency

Conclusion

References

SetFit Paper
SetFit Video on HuggingFace YouTube Channel
SetFit Official GitHub Repository