Comparing the Performance of LLMs: A Deep Dive into Roberta, Llama 2, and Mistral for Disaster Tweets Analysis with Lora
mehdiiraqui
<!-- TOC -->

- [Comparing the Performance of LLMs: A Deep Dive into Roberta, Llama 2, and Mistral for Disaster Tweets Analysis with LoRA](#comparing-the-performance-of-llms-a-deep-dive-into-roberta-llama-2-and-mistral-for-disaster-tweets-analysis-with-lora)
    - [Introduction](#introduction)
    - [Hardware Used](#hardware-used)
    - [Goals](#goals)
    - [Dependencies](#dependencies)
    - [Pre-trained Models](#pre-trained-models)
        - [RoBERTa](#roberta)
        - [Llama 2](#llama-2)
        - [Mistral 7B](#mistral-7b)
    - [LoRA](#lora)
    - [Setup](#setup)
    - [Data preparation](#data-preparation)
        - [Data loading](#data-loading)
        - [Data Processing](#data-processing)
    - [Models](#models)
        - [RoBERTa](#roberta)
            - [Load RoBERTA Checkpoints for the Classification Task](#load-roberta-checkpoints-for-the-classification-task)
            - [LoRA setup for RoBERTa classifier](#lora-setup-for-roberta-classifier)
        - [Mistral](#mistral)
            - [Load checkpoints for the classfication model](#load-checkpoints-for-the-classfication-model)
            - [LoRA setup for Mistral 7B classifier](#lora-setup-for-mistral-7b-classifier)
        - [Llama 2](#llama-2)
            - [Load checkpoints for the classification mode](#load-checkpoints-for-the-classfication-mode)
            - [LoRA setup for Llama 2 classifier](#lora-setup-for-llama-2-classifier)
    - [Setup the trainer](#setup-the-trainer)
        - [Evaluation Metrics](#evaluation-metrics)
        - [Custom Trainer for Weighted Loss](#custom-trainer-for-weighted-loss)
        - [Trainer Setup](#trainer-setup)
            - [RoBERTa](#roberta)
            - [Mistral-7B](#mistral-7b)
            - [Llama 2](#llama-2)
    - [Hyperparameter Tuning](#hyperparameter-tuning)
    - [Results](#results)
    - [Conclusion](#conclusion)
    - [Resources](#resources)

<!-- /TOC -->



## Introduction 

In the fast-moving world of Natural Language Processing (NLP), we often find ourselves comparing different language models to see which one works best for specific tasks. This blog post is all about comparing three models: RoBERTa, Mistral-7b, and Llama-2-7b. We used them to tackle a common problem - classifying tweets about disasters. It is important to note that Mistral and Llama 2 are large models with 7 billion parameters. In contrast, RoBERTa-large (355M parameters) is a relatively smaller model used as a baseline for the comparison study.

In this blog, we used PEFT (Parameter-Efficient Fine-Tuning) technique: LoRA (Low-Rank Adaptation of Large Language Models) for fine-tuning the pre-trained model on the sequence classification task. LoRa is designed to significantly reduce the number of trainable parameters while maintaining strong downstream task performance. 

The main objective of this blog post is to implement LoRA fine-tuning for sequence classification tasks using three pre-trained models from Hugging Face: [meta-llama/Llama-2-7b-hf](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf), [mistralai/Mistral-7B-v0.1](https://huggingface.co/mistralai/Mistral-7B-v0.1), and [roberta-large](https://huggingface.co/roberta-large)

## Hardware Used 

- Number of nodes: 1 
- Number of GPUs per node: 1
- GPU type: A6000 
- GPU memory: 48GB


## Goals

- Implement fine-tuning of pre-trained LLMs using LoRA PEFT methods.
- Learn how to use the HuggingFace APIs ([transformers](https://huggingface.co/docs/transformers/index), [peft](https://huggingface.co/docs/peft/index), and [datasets](https://huggingface.co/docs/datasets/index)).
- Setup the hyperparameter tuning and experiment logging using [Weights & Biases](https://wandb.ai).



## Dependencies

```bash
datasets
evaluate
peft
scikit-learn
torch
transformers
wandb 
```
Note: For reproducing the reported results, please check the pinned versions in the [wandb reports](#resources).


## Pre-trained Models

### [RoBERTa](https://arxiv.org/abs/1907.11692)

RoBERTa (Robustly Optimized BERT Approach) is an advanced variant of the BERT model proposed by Meta AI research team. BERT is a transformer-based language model using self-attention mechanisms for contextual word representations and trained with a masked language model objective. Note that BERT is an encoder only model used for natural language understanding tasks (such as sequence classification and token classification).

RoBERTa is a popular model to fine-tune and appropriate as a baseline for our experiments. For more information, you can check the Hugging Face model [card](https://huggingface.co/docs/transformers/model_doc/roberta).


### [Llama 2](https://arxiv.org/abs/2307.09288)

Llama 2 models, which stands for Large Language Model Meta AI, belong to the family of large language models (LLMs) introduced by Meta AI. The Llama 2 models vary in size, with parameter counts ranging from 7 billion to 65 billion.

Llama 2 is an auto-regressive language model, based on the transformer decoder architecture. To generate text, Llama 2 processes a sequence of words as input and iteratively predicts the next token using a sliding window.
Llama 2 architecture is slightly different from models like GPT-3. For instance, Llama 2 employs the SwiGLU activation function rather than ReLU and opts for rotary positional embeddings in place of absolute learnable positional embeddings. 
 
The recently released Llama 2 introduced architectural refinements to better leverage very long sequences by extending the context length to up to 4096 tokens, and using grouped-query attention (GQA) decoding. 

### [Mistral 7B](https://arxiv.org/abs/2310.06825)

Mistral 7B v0.1, with 7.3 billion parameters, is the first LLM introduced by Mistral AI.
The main novel techniques used in Mistral 7B's architecture are: 
- Sliding Window Attention: Replace the full attention (square compute cost) with a sliding window based attention where each token can attend to at most 4,096 tokens from the previous layer (linear compute cost). This mechanism enables Mistral 7B to handle longer sequences, where higher layers can access historical information beyond the window size of 4,096 tokens. 
- Grouped-query Attention: used in Llama 2 as well, the technique optimizes the inference process (reduce processing time) by caching the key and value vectors for previously decoded tokens in the sequence.  

## [LoRA](https://arxiv.org/abs/2106.09685)

PEFT, Parameter Efficient Fine-Tuning, is a collection of techniques (p-tuning, prefix-tuning, IA3, Adapters, and LoRa) designed to fine-tune large models using a much smaller set of training parameters while preserving the performance levels typically achieved through full fine-tuning. 

LoRA, Low-Rank Adaptation, is a PEFT method that shares similarities with Adapter layers. Its primary objective is to reduce the model's trainable parameters. LoRA's operation involves 
learning a low rank update matrix while keeping the pre-trained weights frozen.

![image](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/Lora-for-sequence-classification-with-Roberta-Llama-Mistral/lora.png)

## Setup

RoBERTa has a limitatiom of maximum sequence length of 512, so we set the `MAX_LEN=512` for all models to ensure a fair comparison.

```python
MAX_LEN = 512 
roberta_checkpoint = "roberta-large"
mistral_checkpoint = "mistralai/Mistral-7B-v0.1"
llama_checkpoint = "meta-llama/Llama-2-7b-hf"
```

## Data preparation
### Data loading

We will load the dataset from Hugging Face:
```python
from datasets import load_dataset
dataset = load_dataset("mehdiiraqui/twitter_disaster")
```
 Now, let's split the dataset into training and validation datasets. Then add the test set:

```python
from datasets import Dataset
# Split the dataset into training and validation datasets
data = dataset['train'].train_test_split(train_size=0.8, seed=42)
# Rename the default "test" split to "validation"
data['val'] = data.pop("test")
# Convert the test dataframe to HuggingFace dataset and add it into the first dataset
data['test'] = dataset['test']
```

Here's an overview of the dataset:

```bash
DatasetDict({
    train: Dataset({
        features: ['id', 'keyword', 'location', 'text', 'target'],
        num_rows: 6090
    })
    val: Dataset({
        features: ['id', 'keyword', 'location', 'text', 'target'],
        num_rows: 1523
    })
    test: Dataset({
        features: ['id', 'keyword', 'location', 'text', 'target'],
        num_rows: 3263
    })
})
```

Let's check the data distribution:

```python
import pandas as pd

data['train'].to_pandas().info()
data['test'].to_pandas().info()
```

- Train dataset

```<class 'pandas.core.frame.DataFrame'>
RangeIndex: 7613 entries, 0 to 7612
Data columns (total 5 columns):
 #   Column    Non-Null Count  Dtype