We start by importing a few helpful libraries.
from datasets import load_datasetfrom transformers import AutoTokenizer, AutoModelForSequenceClassification
from transformers import DistilBertForSequenceClassification, DistilBertConfig
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
Then, we load our dataset from the Hugging Face Hub. This includes training (2100 rows), testing (450 rows), and validation (450 rows) sets.
data = load_dataset("shawhin/phishing-site-classification")
Next, we load in our teacher model. To help speed up training, I loaded the model onto a T4 GPU that was freely available on Google Colab.
# use Nvidia GPU
device = torch.device('cuda')# Load teacher model and tokenizer
model_path = "shawhin/bert-phishing-classifier_teacher"
tokenizer = AutoTokenizer.from_pretrained(model_path)
teacher_model = AutoModelForSequenceClassification.from_pretrained(model_path)
.to(device)
The teacher model is a fine-tuned version of Goolge’s bert-base-uncased that performs binary classification on phishing website URLs. The code to train the teacher model is available on GitHub.
For the student model, we initialize a new model from scratch based on distilbert-base-uncased. We modify the architecture by removing two layers and four attention heads from the remaining layers.
# Load student model
my_config = DistilBertConfig(n_heads=8, n_layers=4) # drop 4 heads per layer and 2 layersstudent_model = DistilBertForSequenceClassification
.from_pretrained("distilbert-base-uncased",
config=my_config,)
.to(device)
Before we can train our student model, we will need to tokenize the dataset. This is important because the models expect input text to be represented in a particular way.
Here, I pad examples based on each batch’s longest example. This allows the batches to be represented as a PyTorch tensor.
# define text preprocessing
def preprocess_function(examples):
return tokenizer(examples["text"], padding='max_length', truncation=True)# tokenize all datasetse
tokenized_data = data.map(preprocess_function, batched=True)
tokenized_data.set_format(type='torch',
columns=['input_ids', 'attention_mask', 'labels'])
Another important step before training is defining an evaluation strategy for our models during training. Below, I define a function that computes the accuracy, precision, recall, and F1 score given a model and dataset.
# Function to evaluate model performance
def evaluate_model(model, dataloader, device):
model.eval() # Set model to evaluation mode
all_preds = []
all_labels = []# Disable gradient calculations
with torch.no_grad():
for batch in dataloader:
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['labels'].to(device)
# Forward pass to get logits
outputs = model(input_ids, attention_mask=attention_mask)
logits = outputs.logits
# Get predictions
preds = torch.argmax(logits, dim=1).cpu().numpy()
all_preds.extend(preds)
all_labels.extend(labels.cpu().numpy())
# Calculate evaluation metrics
accuracy = accuracy_score(all_labels, all_preds)
precision, recall, f1, _ = precision_recall_fscore_support(all_labels,
all_preds,
average='binary')
return accuracy, precision, recall, f1
Now, we are ready to begin the training process. To allow our student model to learn from both the ground truth labels in the training set (i.e., hard targets) and the teacher model’s logits (i.e., soft targets), we must construct a special loss function that considers both targets.
This is done by combining the KL divergence of the student and teacher’s output probability distribution with the cross entropy loss of the student’s logits with the ground truth.
# Function to compute distillation and hard-label loss
def distillation_loss(student_logits, teacher_logits,
true_labels, temperature, alpha):
# Compute soft targets from teacher logits
soft_targets = nn.functional.softmax(teacher_logits / temperature, dim=1)
student_soft = nn.functional.log_softmax(student_logits / temperature, dim=1)# KL Divergence loss for distillation
distill_loss = nn.functional.kl_div(student_soft,
soft_targets,
reduction='batchmean') * (temperature ** 2)
# Cross-entropy loss for hard labels
hard_loss = nn.CrossEntropyLoss()(student_logits, true_labels)
# Combine losses
loss = alpha * distill_loss + (1.0 - alpha) * hard_loss
return loss
Next, we define our hyperparameters, optimizer, and train/test datasets.
# hyperparameters
batch_size = 32
lr = 1e-4 #5e-5
num_epochs = 5
temperature = 2.0
alpha = 0.5# define optimizer
optimizer = optim.Adam(student_model.parameters(), lr=lr)
# create training data loader
dataloader = DataLoader(tokenized_data['train'], batch_size=batch_size)
# create testing data loader
test_dataloader = DataLoader(tokenized_data['test'], batch_size=batch_size)
Finally, we train our student model using PyTorch.
# put student model in train mode
student_model.train()# train model
for epoch in range(num_epochs):
for batch in dataloader:
# Prepare inputs
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['labels'].to(device)
# Disable gradient calculation for teacher model
with torch.no_grad():
teacher_outputs = teacher_model(input_ids,
attention_mask=attention_mask)
teacher_logits = teacher_outputs.logits
# Forward pass through the student model
student_outputs = student_model(input_ids,
attention_mask=attention_mask)
student_logits = student_outputs.logits
# Compute the distillation loss
loss = distillation_loss(student_logits, teacher_logits, labels,
temperature, alpha)
# Backpropagation
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Epoch {epoch + 1} completed with loss: {loss.item()}")
# Evaluate the teacher model
teacher_accuracy, teacher_precision, teacher_recall, teacher_f1 =
evaluate_model(teacher_model, test_dataloader, device)
print(f"Teacher (test) - Accuracy: {teacher_accuracy:.4f},
Precision: {teacher_precision:.4f},
Recall: {teacher_recall:.4f},
F1 Score: {teacher_f1:.4f}")
# Evaluate the student model
student_accuracy, student_precision, student_recall, student_f1 =
evaluate_model(student_model, test_dataloader, device)
print(f"Student (test) - Accuracy: {student_accuracy:.4f},
Precision: {student_precision:.4f},
Recall: {student_recall:.4f},
F1 Score: {student_f1:.4f}")
print("\n")
# put student model back into train mode
student_model.train()
The training results are shown in the screenshot below. Remarkably, by the end of training, the student model outperformed the teacher across all evaluation metrics!
As a final step, we can evaluate the models on the independent validation set, i.e., data not used in training model parameters or tuning hyperparameters.
# create testing data loader
validation_dataloader = DataLoader(tokenized_data['validation'], batch_size=8)# Evaluate the teacher model
teacher_accuracy, teacher_precision, teacher_recall, teacher_f1 =
evaluate_model(teacher_model, validation_dataloader, device)
print(f"Teacher (validation) - Accuracy: {teacher_accuracy:.4f},
Precision: {teacher_precision:.4f},
Recall: {teacher_recall:.4f},
F1 Score: {teacher_f1:.4f}")
# Evaluate the student model
student_accuracy, student_precision, student_recall, student_f1 =
evaluate_model(student_model, validation_dataloader, device)
print(f"Student (validation) - Accuracy: {student_accuracy:.4f},
Precision: {student_precision:.4f},
Recall: {student_recall:.4f},
F1 Score: {student_f1:.4f}")
Here, again, we see the student outperform the teacher.
So far, we’ve reduced our model from 109M parameters (438 MB) to 52.8M parameters (211 MB). However, we can go one step further and quantize the student model.
First, we push the model of the Hugging Face Hub.
student_model.push_to_hub("shawhin/bert-phishing-classifier_student")
Then, we can load it back in using 4-bit quantization. For that, we can use the BitsAndBytes integration in the transformers library.
We set up the config to store model parameters using the 4-bit NormalFloat data type described in the QLoRA paper and the bfloat16 for computation [10].
from transformers import BitsAndBytesConfig# load model in model as 4-bit
nf4_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype = torch.bfloat16,
bnb_4bit_use_double_quant=True
)
model_nf4 = AutoModelForSequenceClassification.from_pretrained(model_id,
device_map=device,
quantization_config=nf4_config)
We can then evaluate our quantized model on the validation set.
# Evaluate the student model
quantized_accuracy, quantized_precision, quantized_recall, quantized_f1 =
evaluate_model(model_nf4, validation_dataloader, device)print("Post-quantization Performance")
print(f"Accuracy: {quantized_accuracy:.4f},
Precision: {quantized_precision:.4f},
Recall: {quantized_recall:.4f},
F1 Score: {quantized_f1:.4f}")
Once again, we see a small performance improvement after compression. An intuitive explanation for this is Occam’s Razor principle, which states that simpler models are better.
In this case, the model may be overparameterized for this binary classification task. Thus, simplifying the model results in better performance.
While modern large language models (LLMs) demonstrate impressive performance on various tasks, their scale presents challenges in deploying them in real-world settings.
Recent innovations in model compression techniques help mitigate these challenges by reducing the computational cost of LLM solutions. Here, we discussed three broad categories of compression techniques (Quantization, Pruning, and Knowledge Distillation) and walked through an example implementation in Python.
More on LLMs 👇
Large Language Models (LLMs)
My website: https://www.shawhintalebi.com/