How to Fine-Tune Language Models: First Principles to Scalable Performance

Author(s): Ehssan

Originally published on Towards AI.

In this article, we’ll explore the process of fine-tuning language models for text classification. We’ll do so in three levels: first, by manually adding a classification head in PyTorch* and training the model so you can see the full process; second, by using the Hugging Face* Transformers library to streamline the process; and third, by leveraging PyTorch Lightning* and accelerators to optimize training performance. By the end of this guide, you’ll have a well-rounded understanding of the fine-tuning workflow.

Introduction

The idea behind using fine-tuning in Natural Language Processing (NLP) was borrowed from Computer Vision (CV). The CV models were first trained on large datasets such as ImageNet to teach them the basic features of images such as edges or colors. These pretrained models were then fine-tuned on a downstream task such as classifying birds with a relatively small number of labeled examples.

Fine-tuned models typically achieved a higher accuracy than supervised models trained from scratch on the same amount of labeled data.

Despite the popularity and success of transfer learning in CV, for many years it wasn’t clear what the analogous pretraining process was for NLP. Consequently, NLP applications required large amounts of labeled data to achieve high performance.

How is Fine-tuning Different from Pretraining?

With pretraining, language models gain a general understanding of languages. During this process, they learn language patterns but typically are not capable of following instructions or answering questions. In the case of GPT models, this self-supervised learning includes predicting the next word (unidirectional) based on their training data, which is often webpages. In the case of BERT (Bidirectional Encoder Representations from Transformers), learning involves predicting randomly masked words (bidirectional) and sentence-order prediction. But how can we adapt language models for our own data or our own tasks?

Fine-tuning continues training a pretrained model to increase its performance on specific tasks. For instance, through “instruction fine-tuning” you can teach a model to behave more like a chatbot. This is the process for specializing a general purpose model like OpenAI* GPT-4 into an application like ChatGPT* or GitHub* Copilot. By fine-tuning your own language model, you can increase the reliability, performance, and privacy of your model while reducing the associated inference costs compared to subscription-based services, especially if you have a large volume of data or frequent requests.

Fine-Tuning a Language Model for Text Classification

Preprocessing and Preparing DataLoader

Feel free to skip this section if you’re comfortable with preprocessing data. Throughout we assume that we have our labeled data saved in train, validation, and test csv files each with a text and a label column. For training, the labels should be numeric, so if that’s not the case, you
can use a label_to_id dictionary such as {"negative": 0, "positive": 1} and do a mapping to get the desired format.

For concreteness, we will use BERT as the base model and set the number of classification labels to 4. After running the code below, you are encouraged to swap BERT for DistilBERT which reduces the size of the BERT model by 40%, speeding inference by 60%, while retaining 97% of BERT’s language understanding capabilities.

A Quick Look at BERT

BERT was introduced by Google in 2018 and has since revolutionized the field of NLP. Unlike traditional models that process text in a unidirectional manner, BERT is designed to understand the context of a word in a sentence by looking at both its left and right surroundings. This bidirectional approach allows BERT to capture the nuances of language more effectively.

Key Features of BERT

1. Pretraining: BERT is pretrained on a massive corpus of text,
   including the entire Wikipedia and BookCorpus. The pretraining
   involves two tasks: Masked Language Modeling (MLM) and Next
   Sentence Prediction (NSP).
2. Architecture: BERT_BASE has 12 layers (transformer blocks), 768
   hidden units, and 12 attention heads, totaling 110 million
   parameters.

You can run this tutorial on Intel® Tiber™ AI Cloud, using an Intel® Xeon® CPU instance. This platform provides ample computing resources for smooth execution of our code.

import os

import torch
from torch.utils.data import DataLoader, Dataset
from transformers import AutoTokenizer
import pandas as pd

# Parameters
model_ckpt = "bert-base-uncased"
num_labels = 4
batch_size = 16
num_workers = 6

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)


# Custom Dataset class
class TextDataset(Dataset):
 def __init__(self, dataframe, tokenizer, max_length=512):
 self.data = dataframe
 self.tokenizer = tokenizer
 self.max_length = max_length

 def __len__(self):
 return len(self.data)

 def __getitem__(self, idx):
 row = self.data.iloc[idx]
 text = row["text"] # Replace "text" with your actual column name for text
 label = row["label"] # Replace "label" with your actual column name for labels

 # Tokenize the input text
 encoding = self.tokenizer(
 text,
 max_length=self.max_length,
 padding="max_length",
 truncation=True,
 return_tensors="pt",
 )

 return {
 "input_ids": encoding["input_ids"].squeeze(0), # Remove batch dimension with squeeze
 "attention_mask": encoding["attention_mask"].squeeze(0),
 "label": torch.tensor(label, dtype=torch.long),
 }

os.environ["TOKENIZERS_PARALLELISM"] = "false"

# Load csv files
train_df = pd.read_csv("train.csv")
val_df = pd.read_csv("val.csv")
test_df = pd.read_csv("test.csv")

# Create Dataset objects
train_dataset = TextDataset(train_df, tokenizer)
val_dataset = TextDataset(val_df, tokenizer)
test_dataset = TextDataset(test_df, tokenizer)

# Create DataLoaders
train_loader = DataLoader(
 train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers
)
val_loader = DataLoader(val_dataset, batch_size=batch_size, num_workers=num_workers)
test_loader = DataLoader(test_dataset, batch_size=batch_size, num_workers=num_workers)

The Classification Token [CLS]

The [CLS] token is typically added at the beginning of the input sequence in transformer-based models such as BERT and its variants. During fine-tuning, the model learns to assign meaningful information to the [CLS] token, which aggregates the input sequence’s context. The last hidden state corresponding to the [CLS] token is then used as a representation of the entire input, which can be passed through a classifier layer for downstream tasks like sentiment analysis, topic categorization, or any task requiring a decision based on the entire sequence. This mechanism allows the model to focus on both the global understanding of the text and task-specific features for accurate predictions.

Unlike traditional models that may rely on static embeddings (like word2vec), transformers generate contextualized embeddings, so that the meaning of a token depends on the tokens around it. The [CLS] token, as it passes through the layers, becomes increasingly “aware” of the entire sequence’s meaning, which makes it a good summary representation for downstream tasks. For some tasks, especially those requiring finer-grained understanding, other strategies might be employed. For instance, for document classification, where every word contributes equally, some models use mean pooling over all token embeddings.

Level 1: PyTorch

In this section, we manually add a classification head to the base model and do the fine-tuning. We achieve this using the AutoModelclass which converts the tokens (or rather token encodings) to embeddings and then then feeds them through the encoder stack to return the hidden states. While AutoModel is helpful for understanding the idea behind what we’re doing, to fine-tune for text classification it’s better practice to work withAutoModelForSequenceClassification instead, as we discuss below.

import torch
from torch import nn
from transformers import AutoModel


# Load the base model with AutoModel and add a classifier
class CustomModel(nn.Module):
 def __init__(self, model_ckpt, num_labels):
 super(CustomModel, self).__init__()
 self.model = AutoModel.from_pretrained(model_ckpt) # Base transformer model
 self.classifier = nn.Linear(
 self.model.config.hidden_size, num_labels
 ) # Classification head. The 1st parameter equals 768 for BERT as discussed above

 def forward(self, input_ids, attention_mask):
 # Forward pass through the transformer model
 outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)

 # Use the [CLS] token (0-th token in the sequence) for classification
 cls_output = outputs.last_hidden_state[
 :, 0, :
 ] # Shape: (batch_size, hidden_size)

 # Pass through the classifier head
 logits = self.classifier(cls_output)
 return logits


# Initialize the model
model = CustomModel(model_ckpt, num_labels)

# Loss function and optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)


# Training function
def train(model, optimizer, train_loader, loss_fn):
 model.train()
 total_loss = 0

 for batch in train_loader:
 optimizer.zero_grad()

 # Unpack the batch data
 input_ids = batch["input_ids"]
 attention_mask = batch["attention_mask"]
 label = batch["label"]

 # Forward pass
 output = model(input_ids, attention_mask)

 # Compute loss
 loss = loss_fn(output, label)
 loss.backward()

 # Update the model parameters
 optimizer.step()

 total_loss += loss.item()

 print(f"Train loss: {total_loss / len(train_loader):.2f}")


import torch


def evaluate(model, test_loader, loss_fn):
 model.eval() # Set model to evaluation mode
 total_loss = 0
 total_acc = 0
 total_samples = 0

 with torch.no_grad(): # No gradient computation needed during evaluation
 for batch in test_loader:
 input_ids = batch["input_ids"]
 attention_mask = batch["attention_mask"]
 labels = batch["label"]

 # Forward pass
 output = model(input_ids, attention_mask)

 # Compute loss
 loss = loss_fn(output, labels)
 total_loss += loss.item()

 # Compute accuracy
 predictions = torch.argmax(output, dim=1)
 total_acc += torch.sum(predictions == labels).item()
 total_samples += labels.size(0)

 # Calculate average loss and accuracy
 avg_loss = total_loss / len(test_loader)
 avg_acc = total_acc / total_samples * 100

 print(f"Test loss: {avg_loss:.2f}, Test acc: {avg_acc:.2f}%")

Finally, we can train, evaluate, and save the model.

num_epochs = 3

for epoch in range(num_epochs):
 train(model, optimizer, train_loader, loss_fn)

evaluate(model, test_loader, loss_fn)

torch.save(model.state_dict(), "./fine-tuned-model.pt")

Level 2: Hugging Face Transformers

Now, we use the convenience of AutoModelForSequenceClassification class that will add the classification head to the base model automatically. Compare this against what we did with the AutoModel class in the previous section!

Also note that the Trainer class from Hugging Face’s Transformers
library can directly handle Dataset objects without needing a
DataLoader, as it automatically handles batching and shuffling for
you.

from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments

model = AutoModelForSequenceClassification.from_pretrained(
 model_ckpt, num_labels=num_labels
)


training_args = TrainingArguments(
 output_dir="./results",
 num_train_epochs=3,
 warmup_steps=500,
 weight_decay=0.01,
 logging_dir="./logs",
 logging_steps=10, # Log every 10 steps
 evaluation_strategy="steps",
 save_steps=500, # Save model checkpoint every 500 steps
 load_best_model_at_end=True, # Load the best model at the end of training
 metric_for_best_model="accuracy",
)

# Train the model
trainer = Trainer(
 model=model,
 args=training_args,
 train_dataset=train_dataset,
 eval_dataset=test_dataset,
)

trainer.train()
trainer.evaluate(test_dataset)

Level 3: PyTorch Lightning

Lightning is, in the words of its documentation, “the deep learning framework with ‘batteries included’ for professional AI researchers and machine learning engineers who need maximal flexibility while super-charging performance at scale.”

As we shall see, with a bit of additional organizational code, the Lightning Trainer automates the following:

1. Epoch and batch iteration
2. optimizer.step(), loss.backward(), optimizer.zero_grad() calls
3. Calling of model.eval(), enabling and disabling grads during
   evaluation
4. Checkpoint Saving and Loading
5. Logging
6. Accelerator, Multi-GPU, and TPU Support (No .to(device) calls required.)
7. Mixed-precision training

You can accelerate training with Intel Gaudi processors, which allow you to conduct more deep learning training at a lower expense. You can try an Intel Gaudi instance for free on Intel® Tiber™ AI Cloud.

import torchmetrics
import lightning as L
from lightning.pytorch.callbacks import ModelCheckpoint
from lightning.pytorch.loggers import TensorBoardLogger
from transformers import AutoModelForSequenceClassification


# A LightningModule is a torch.nn.Module with added functionality.
# It wraps around a regular PyTorch model.
class LightningModel(L.LightningModule):
 def __init__(self, model, learning_rate=5e-5):
 super().__init__()

 self.learning_rate = learning_rate
 self.model = model

 self.val_acc = torchmetrics.Accuracy(task="multiclass", num_classes=num_labels)
 self.test_acc = torchmetrics.Accuracy(task="multiclass", num_classes=num_labels)

 def forward(self, input_ids, attention_mask, labels):
 return self.model(input_ids, attention_mask=attention_mask, labels=labels)

 def _shared_step(self, batch, batch_idx):
 outputs = self(
 batch["input_ids"],
 attention_mask=batch["attention_mask"],
 labels=batch["label"],
 )
 return outputs

 def training_step(self, batch, batch_idx):
 outputs = self._shared_step(batch, batch_idx)
 self.log("train_loss", outputs["loss"])
 return outputs["loss"]

 def validation_step(self, batch, batch_idx):
 outputs = self._shared_step(batch, batch_idx)
 self.log("val_loss", outputs["loss"], prog_bar=True)
 logits = outputs["logits"]
 self.val_acc(logits, batch["label"])
 self.log("val_acc", self.val_acc, prog_bar=True)

 def test_step(self, batch, batch_idx):
 outputs = self._shared_step(batch, batch_idx)
 logits = outputs["logits"]
 self.test_acc(logits, batch["label"])
 self.log("accuracy", self.test_acc, prog_bar=True)

 def configure_optimizers(self):
 optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
 return optimizer


model = AutoModelForSequenceClassification.from_pretrained(
 model_ckpt, num_labels=num_labels
)
lightning_model = LightningModel(model)

callbacks = [
 ModelCheckpoint(save_top_k=1, mode="max", monitor="val_acc") # Save top 1 model
]
logger = TensorBoardLogger(save_dir="./logs", name="fine-tuned-model")

trainer = L.Trainer(
 max_epochs=3,
 callbacks=callbacks,
 accelerator="hpu",
 precision="bf16-mixed", # By default, HPU training uses 32-bit precision. To enable mixed precision, set the precision flag.
 devices="auto",
 logger=logger,
 log_every_n_steps=10,
)

trainer.fit(lightning_model, train_dataloaders=train_loader, val_dataloaders=val_loader)

trainer.test(lightning_model, train_loader, ckpt_path="best")
trainer.test(lightning_model, val_loader, ckpt_path="best")
trainer.test(lightning_model, test_loader, ckpt_path="best")

While the Transformers Trainer class supports distributed training, it doesn’t offer the same level of integration and flexibility as Lightning when it comes to advanced features like custom callbacks, logging, and seamless scaling across multiple GPUs or nodes.

Practical Advice

Now that you’re familiar with the fine-tuning process, you might
wonder how you can apply it to your specific task. Here’s some
practical advice:

1. Collect real data for your task, or generate synthetic data. See, for instance, Synthetic Data Generation with Language Models: A Practical Guide.
2. Fine-tune a relatively small model.
3. Evaluate your language model on your test set, and on a benchmark if available for your task.
4. Increase the training dataset size, base model size, and, if necessary, task complexity.

Keep in mind that the standard or conventional fine-tuning of language
models as described in this writing can be expensive. Rather than
updating all the weights and biases, we could update only the last
layer as follows:

# Freeze all layers
for param in model.parameters():
 param.requires_grad = False

# Unfreeze last layer
for param in model.pre_classifier.parameters():
 param.requires_grad = True

for param in model.classifier.parameters():
 param.requires_grad = True

In future articles, we shall discuss more efficient fine-tuning techniques, so stay tuned!

For more AI development how-to content, visit Intel® AI Development
Resources.

Acknowledgments

The author thanks Jack Erickson for providing detailed feedback on an earlier draft of this work.

Suggested Reading

• Synthetic Data Generation with Language Models: A Practical Guide
• Natural Language Processing with Transformers
• Intel Gaudi AI Accelerators Overview
• Lightning support for AI Accelerators

*Other names and brands may be claimed as the property of others.

Join thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming a sponsor.

Published via Towards AI