Transformer Series 4: Fine tuning the model

In the previous post, we discussed training the Transformer model, monitoring loss, and ensuring the model improves over time. Now, it’s time to take a step forward: evaluating the model’s performance and fine tuning it to make it more efficient and accurate.

In this post, we’ll focus on:

1. Evaluating the model on unseen data.
2. Using evaluation metrics such as accuracy and loss.
3. Fine-tuning the model for improved performance.
4. Understanding the output and how to interpret it.

---

Evaluating the Model on Unseen Data

Once a model has been trained, it’s important to test its performance on a validation or test dataset that it hasn’t seen during training. This gives us a good estimate of how well the model can generalize to new data.

Here’s a function that evaluates the model on a test set:

def evaluate(model, input_ids, target_ids, attention_mask):
    model.eval()  # Set the model to evaluation mode (no backpropagation)
    total_loss = 0
    with torch.no_grad():  # Disable gradient calculation for efficiency
        outputs = model(input_ids)
        outputs = outputs.view(-1, outputs.size(-1))  # Reshape for loss calculation
        target_ids = target_ids.view(-1)

        # Calculate loss
        loss = loss_fn(outputs, target_ids)
        total_loss += loss.item()

    print(f"Evaluation Loss: {total_loss}")

What’s happening here:

• The model is set to evaluation mode, meaning no updates to the model parameters will be made (no backpropagation).
• We use the cross-entropy loss to compare the model’s predictions to the actual targets.
• The loss value is printed after the evaluation.

Using Evaluation Metrics: Accuracy and Loss

In addition to loss, we can track other evaluation metrics like accuracy. For language models, accuracy is calculated by checking how often the model correctly predicts the next word in the sequence.

Here’s an evaluation function that calculates both accuracy and loss:

def evaluate_with_accuracy(model, input_ids, target_ids, attention_mask):
    model.eval()
    total_loss = 0
    correct_predictions = 0
    total_predictions = 0

    with torch.no_grad():
        outputs = model(input_ids)
        outputs = outputs.view(-1, outputs.size(-1))  # Flatten for loss calculation
        target_ids = target_ids.view(-1)

        # Calculate loss
        loss = loss_fn(outputs, target_ids)
        total_loss += loss.item()

        # Get the predicted token by taking the index of the highest logit
        predictions = torch.argmax(outputs, dim=1)

        # Calculate accuracy
        correct_predictions += (predictions == target_ids).sum().item()
        total_predictions += target_ids.size(0)

    accuracy = correct_predictions / total_predictions
    print(f"Evaluation Loss: {total_loss}")
    print(f"Evaluation Accuracy: {accuracy * 100:.2f}%")

What’s happening here:

• We calculate the total loss and accuracy by comparing the predicted tokens to the actual tokens.
• Accuracy is computed as the number of correct predictions divided by the total number of predictions.

Fine-Tuning the Model

Fine-tuning involves training the model for additional epochs, typically on a smaller, domain-specific dataset. This helps the model become more specialized in a particular task.

We can fine-tune the Transformer model using the same training loop we used in Post 3, but we may want to:

• Lower the learning rate to prevent overfitting.
• Use fewer epochs, since fine-tuning usually requires less training time.

Here’s an updated fine-tuning loop:

def fine_tune(model, input_ids, target_ids, attention_mask, epochs=5):
    optimizer = optim.Adam(model.parameters(), lr=1e-5)  # Lower learning rate for fine-tuning
    model.train()

    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = model(input_ids)
        outputs = outputs.view(-1, outputs.size(-1))
        target_ids = target_ids.view(-1)

        loss = loss_fn(outputs, target_ids)
        loss.backward()
        optimizer.step()

        print(f"Fine-tuning Epoch {epoch + 1}, Loss: {loss.item()}")

What’s happening here:

• We reduce the learning rate to 1e-5 to make smaller updates to the model weights.
• The model is trained for a few epochs (usually 2-5) on the new data.

---

Explanation of the Output

Let’s look at what kind of output you’ll see during evaluation and fine tuning:

Evaluation Output

After running the evaluation function, you will see both the evaluation loss and accuracy printed out:

Example output:

Evaluation Loss: 3.672
Evaluation Accuracy: 85.23%

• Evaluation Loss: This tells us how far off the model’s predictions were from the actual targets.
• Evaluation Accuracy: This measures how often the model predicted the correct next word in the sequence.

Fine-Tuning Output

During fine-tuning, you’ll see the loss for each fine-tuning epoch. Example output:

Fine-tuning Epoch 1, Loss: 2.514
Fine-tuning Epoch 2, Loss: 1.978
Fine-tuning Epoch 3, Loss: 1.612

• Fine-Tuning Loss: This shows how much the model is improving during the fine-tuning process.

Visualizing Loss and Accuracy

You can also visualize the loss and accuracy during evaluation:

import matplotlib.pyplot as plt

def evaluate_and_plot(model, input_ids, target_ids, attention_mask):
    model.eval()
    total_loss = 0
    correct_predictions = 0
    total_predictions = 0

    with torch.no_grad():
        outputs = model(input_ids)
        outputs = outputs.view(-1, outputs.size(-1))
        target_ids = target_ids.view(-1)

        loss = loss_fn(outputs, target_ids)
        total_loss += loss.item()

        predictions = torch.argmax(outputs, dim=1)
        correct_predictions += (predictions == target_ids).sum().item()
        total_predictions += target_ids.size(0)

    accuracy = correct_predictions / total_predictions

    # Plot loss and accuracy
    plt.figure(figsize=(10, 5))

    plt.subplot(1, 2, 1)
    plt.title("Evaluation Loss")
    plt.plot([total_loss], 'r-')
    plt.xlabel("Evaluation")
    plt.ylabel("Loss")

    plt.subplot(1, 2, 2)
    plt.title("Evaluation Accuracy")
    plt.plot([accuracy * 100], 'g-')
    plt.xlabel("Evaluation")
    plt.ylabel("Accuracy (%)")

    plt.show()

# Run evaluation and plot
evaluate_and_plot(model, input_ids, target_ids, attention_mask)

This will create two plots:

• Loss over the evaluation.
• Accuracy over the evaluation.

---

Fine-Tuning Strategies

Here are some key tips to fine-tune your Transformer model effectively:

• Use domain-specific datasets: For example, if you’re fine-tuning for a chatbot, use conversational datasets.
• Monitor overfitting: Fine-tuning for too many epochs can lead to overfitting, so track the loss carefully.
• Experiment with learning rates: A lower learning rate (1e-5 or lower) often works best for fine-tuning.

---

Conclusion

In this post, we covered how to evaluate and fine-tune your Transformer model. To summarize:

1. We used a validation dataset to calculate the loss and accuracy of the model.
2. We implemented a fine tuning process to adapt the model to new tasks or domains.
3. We explained and visualized the model’s output during evaluation and fine-tuning.

Full file

With these techniques, you can now evaluate your model’s performance and fine-tune it to achieve better results. In the next post, we’ll look at deploying your model and making it available for real-time use.