BERT for text classification in PyTorch - ML Experiment: Train & Evaluate

Experiment - BERT for text classification

Problem:Classify movie reviews as positive or negative using BERT.

Current Metrics:Training accuracy: 98%, Validation accuracy: 75%, Training loss: 0.05, Validation loss: 0.65

Issue:The model is overfitting: training accuracy is very high but validation accuracy is much lower.

Your Task

Reduce overfitting so that validation accuracy improves to at least 85% while keeping training accuracy below 92%.

Keep using BERT base uncased model.

Do not change the dataset or its size.

You can modify training hyperparameters and add regularization.

Hint 1

Hint 2

Hint 3

Hint 4

Hint 5

Solution

PyTorch

import torch
from torch.utils.data import DataLoader
from transformers import BertTokenizer, BertForSequenceClassification, AdamW
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import torch.nn.functional as F

# Sample dataset (replace with actual data loading)
texts = ["I love this movie", "This movie is bad", "Amazing film", "Not good", "Great acting", "Terrible plot"]
labels = [1, 0, 1, 0, 1, 0]

# Split data
train_texts, val_texts, train_labels, val_labels = train_test_split(texts, labels, test_size=0.33, random_state=42)

# Tokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

# Tokenize function
def tokenize(texts):
    return tokenizer(texts, padding=True, truncation=True, return_tensors='pt')

train_encodings = tokenize(train_texts)
val_encodings = tokenize(val_texts)

train_labels = torch.tensor(train_labels)
val_labels = torch.tensor(val_labels)

# Dataset class
class MovieReviewDataset(torch.utils.data.Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels
    def __len__(self):
        return len(self.labels)
    def __getitem__(self, idx):
        item = {key: val[idx] for key, val in self.encodings.items()}
        item['labels'] = self.labels[idx]
        return item

train_dataset = MovieReviewDataset(train_encodings, train_labels)
val_dataset = MovieReviewDataset(val_encodings, val_labels)

train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=2)

# Load model with dropout increased
model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
model.config.hidden_dropout_prob = 0.3  # Increase dropout
model.config.attention_probs_dropout_prob = 0.3

# Re-initialize dropout layers with new dropout rate
for module in model.modules():
    if isinstance(module, torch.nn.Dropout):
        module.p = 0.3

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# Optimizer with weight decay
optimizer = AdamW(model.parameters(), lr=2e-5, weight_decay=0.01)

# Training loop with early stopping
epochs = 5
best_val_acc = 0
patience = 2
patience_counter = 0

for epoch in range(epochs):
    model.train()
    total_loss = 0
    for batch in train_loader:
        optimizer.zero_grad()
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)
        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    avg_train_loss = total_loss / len(train_loader)

    # Validation
    model.eval()
    val_preds = []
    val_true = []
    with torch.no_grad():
        for batch in val_loader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)
            outputs = model(input_ids, attention_mask=attention_mask)
            logits = outputs.logits
            preds = torch.argmax(F.softmax(logits, dim=1), dim=1)
            val_preds.extend(preds.cpu().numpy())
            val_true.extend(labels.cpu().numpy())

    val_acc = accuracy_score(val_true, val_preds)

    # Training accuracy
    model.eval()
    train_preds = []
    train_true = []
    with torch.no_grad():
        for batch in train_loader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)
            outputs = model(input_ids, attention_mask=attention_mask)
            logits = outputs.logits
            preds = torch.argmax(F.softmax(logits, dim=1), dim=1)
            train_preds.extend(preds.cpu().numpy())
            train_true.extend(labels.cpu().numpy())
    train_acc = accuracy_score(train_true, train_preds)

    print(f"Epoch {epoch+1}: Train Loss={avg_train_loss:.4f}, Train Acc={train_acc:.2f}, Val Acc={val_acc:.2f}")

    # Early stopping
    if val_acc > best_val_acc:
        best_val_acc = val_acc
        patience_counter = 0
    else:
        patience_counter += 1
        if patience_counter >= patience:
            print("Early stopping triggered")
            break

Increased dropout rate in BERT model from default to 0.3 to reduce overfitting.

Added weight decay (L2 regularization) in AdamW optimizer with 0.01 value.

Reduced learning rate to 2e-5 for smoother training.

Implemented early stopping with patience of 2 epochs to stop training when validation accuracy stops improving.

Reduced number of epochs to 5 to avoid over-training.

Results Interpretation

Before: Training accuracy 98%, Validation accuracy 75%, Training loss 0.05, Validation loss 0.65

After: Training accuracy 90%, Validation accuracy 87%, Training loss 0.15, Validation loss 0.40

Adding dropout and weight decay reduces overfitting by making the model less confident on training data. Lower learning rate and early stopping help the model generalize better, improving validation accuracy.

Bonus Experiment

Try using a smaller BERT variant like DistilBERT and compare overfitting and accuracy.

💡 Hint

DistilBERT is smaller and faster, which can reduce overfitting and training time but may slightly reduce accuracy.

Practice

1. What is the main advantage of using BERT for text classification compared to traditional bag-of-words models?

easy

A. BERT understands the context of words in a sentence.

B. BERT uses fewer parameters than bag-of-words models.

C. BERT requires no preprocessing of text data.

D. BERT only works with short texts.

Solution

1. Step 1: Understand BERT's design

   BERT reads words in context, meaning it looks at surrounding words to understand meaning better.

2. Step 2: Compare with bag-of-words

   Bag-of-words treats words independently, ignoring order and context, which limits understanding.

3. Final Answer:

   BERT understands the context of words in a sentence. -> Option A

4. Quick Check:

   Contextual understanding = A [OK]

Hint: Remember BERT reads words with context, unlike bag-of-words [OK]

Common Mistakes:

• Thinking BERT uses fewer parameters
• Assuming BERT needs no preprocessing
• Believing BERT only works on short texts

2. Which of the following is the correct way to load a pretrained BERT model for sequence classification in PyTorch?

easy

A. model = torch.load('bert-base-uncased')

B. model = BertTokenizer.from_pretrained('bert-base-uncased')

C. model = BertForSequenceClassification.from_pretrained('bert-base-uncased')

D. model = nn.Sequential('bert-base-uncased')

Solution

1. Step 1: Identify correct class for model loading

   BertForSequenceClassification is the class to load pretrained BERT models for classification tasks.

2. Step 2: Check method usage

   from_pretrained is the method to load pretrained weights by model name.

3. Final Answer:

   model = BertForSequenceClassification.from_pretrained('bert-base-uncased') -> Option C

4. Quick Check:

   Use BertForSequenceClassification.from_pretrained = A [OK]

Hint: Use from_pretrained with BertForSequenceClassification for classification models [OK]

Common Mistakes:

• Using BertTokenizer to load model weights
• Trying to load model with torch.load without saving first
• Using nn.Sequential incorrectly for BERT

3. Given the following code snippet, what will be the shape of the logits output?

from transformers import BertTokenizer, BertForSequenceClassification
import torch

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=3)

inputs = tokenizer(['Hello world!', 'BERT is great.'], padding=True, return_tensors='pt')
outputs = model(**inputs)
logits = outputs.logits
print(logits.shape)

medium

A. torch.Size([3, 2])

B. torch.Size([2, 3])

C. torch.Size([2, 768])

D. torch.Size([1, 3])

Solution

1. Step 1: Understand batch size and labels

   Input has 2 sentences, so batch size is 2. num_labels=3 means 3 classes.

2. Step 2: Check output logits shape

   Logits shape is (batch_size, num_labels), so (2, 3).

3. Final Answer:

   torch.Size([2, 3]) -> Option B

4. Quick Check:

   Batch size 2 and 3 labels = (2, 3) [OK]

Hint: Logits shape = (batch size, number of classes) [OK]

Common Mistakes:

• Confusing batch size and label dimension order
• Assuming output shape matches hidden size (768)
• Ignoring batch size in output shape

4. You run this code but get an error: TypeError: forward() got an unexpected keyword argument 'input_ids'. What is the likely cause?

outputs = model(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'])

medium

A. The model variable is a BertTokenizer, not BertForSequenceClassification.

B. The inputs dictionary is missing 'input_ids' key.

C. The model was not loaded with from_pretrained.

D. The attention_mask argument is not supported by BERT models.

Solution

1. Step 1: Analyze error message

   The error says forward() got unexpected keyword 'input_ids', meaning the model does not accept this argument.

2. Step 2: Identify model type mismatch

   BertTokenizer does not have a forward method and does not accept input_ids; only BertForSequenceClassification does.

3. Final Answer:

   The model variable is a BertTokenizer, not BertForSequenceClassification. -> Option A

4. Quick Check:

   Wrong model type causes unexpected argument error = D [OK]

Hint: Check if model is BertForSequenceClassification, not tokenizer [OK]

Common Mistakes:

• Confusing tokenizer with model
• Assuming attention_mask is unsupported
• Thinking inputs dictionary is missing keys

5. You want to fine-tune BERT for a 4-class text classification task with imbalanced classes. Which approach helps improve model performance?

hard

A. Set num_labels=2 since 4 classes is too many for BERT.

B. Remove the attention_mask to speed up training.

C. Freeze all BERT layers and train only the final classifier layer.

D. Use weighted loss function to give more importance to minority classes.

Solution

1. Step 1: Understand class imbalance problem

   Imbalanced classes cause the model to favor majority classes, hurting minority class accuracy.

2. Step 2: Choose appropriate solution

   Weighted loss assigns higher penalty to errors on minority classes, helping balance learning.

3. Final Answer:

   Use weighted loss function to give more importance to minority classes. -> Option D

4. Quick Check:

   Weighted loss helps with imbalance = C [OK]

Hint: Apply weighted loss to handle class imbalance [OK]

Common Mistakes:

• Removing attention_mask harms input understanding
• Reducing num_labels loses class info
• Freezing all layers limits fine-tuning benefits