RoBERTa vs ALBERT vs DistilBERT

The Origin Story — Why BERT Needed Successors

BERT (Bidirectional Encoder Representations from Transformers) was introduced by Google in 2018 and fundamentally changed NLP. But its training recipe had real flaws. RoBERTa, ALBERT, and DistilBERT each attacked one or more of those flaws — and in doing so, taught us deep lessons about what actually matters in pre-training language models.

BERT in 5 Minutes — The Foundation You Must Know

Before understanding the successors, you need a clear mental model of what BERT is actually doing.

BERT's Two Pre-training Tasks

RoBERTa — A Robustly Optimised BERT Pretraining Approach

RoBERTa was published by Facebook AI in 2019. The paper's central claim was audacious: BERT was significantly undertrained, and fixing the training procedure — without changing the architecture — produced state-of-the-art results on many benchmarks.

RoBERTa's Five Key Changes to BERT Training

RoBERTa Architecture Diagram

ALBERT — A Lite BERT for Self-Supervised Learning

ALBERT (A Lite BERT) was published by Google Research and the Toyota Technological Institute in 2019. Its contribution is not about training data or masking — it's about parameter efficiency. ALBERT introduced two architectural innovations and one new training objective.

ALBERT's Three Core Innovations

ALBERT Model Variants vs BERT

ALBERT Architecture Diagram

DistilBERT — Knowledge Distillation from BERT

DistilBERT was published by HuggingFace in 2019. It uses knowledge distillation — a technique where a small "student" model is trained to mimic the behaviour of a large "teacher" model, not just its final predictions, but the soft probability distributions it produces at every step.

Knowledge Distillation — The Core Mechanism

DistilBERT vs BERT — Size, Speed, Performance

Side-by-Side Comparison — RoBERTa vs ALBERT vs DistilBERT

Fine-Tuning RoBERTa for Text Classification

All three models share the same fine-tuning paradigm via HuggingFace Transformers. We'll fine-tune RoBERTa on sentiment classification (SST-2) as our first example.

# ── RoBERTa Fine-tuning for Sentiment Classification (SST-2) ──
from transformers import (
    RobertaTokenizer,
    RobertaForSequenceClassification,
    Trainer,
    TrainingArguments
)
from datasets import load_dataset
import numpy as np
from sklearn.metrics import accuracy_score, f1_score

# ── 1. Load data ───────────────────────────────────────────────
dataset = load_dataset("sst2")
tokeniser = RobertaTokenizer.from_pretrained("roberta-base")

# ── 2. Tokenise ────────────────────────────────────────────────
def tokenise(batch):
    return tokeniser(
        batch["sentence"],
        truncation=True,
        padding="max_length",
        max_length=128
    )

tokenised = dataset.map(tokenise, batched=True)
tokenised = tokenised.rename_column("label", "labels")
tokenised.set_format("torch",
    columns=["input_ids", "attention_mask", "labels"])

# ── 3. Load pre-trained model ──────────────────────────────────
model = RobertaForSequenceClassification.from_pretrained(
    "roberta-base",
    num_labels=2     # binary: positive / negative
)

# ── 4. Metrics ─────────────────────────────────────────────────
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=-1)
    return {
        "accuracy": accuracy_score(labels, preds),
        "f1":       f1_score(labels, preds, average="binary")
    }

# ── 5. Training arguments ──────────────────────────────────────
args = TrainingArguments(
    output_dir="./roberta-sst2",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    learning_rate=2e-5,
    weight_decay=0.01,
    warmup_ratio=0.1,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    fp16=True            # mixed precision: faster on GPU
)

# ── 6. Train ───────────────────────────────────────────────────
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=tokenised["train"],
    eval_dataset=tokenised["validation"],
    compute_metrics=compute_metrics
)

trainer.train()
results = trainer.evaluate()
print(f"Validation Accuracy: {results['eval_accuracy']:.4f}")
print(f"Validation F1:       {results['eval_f1']:.4f}")

Fine-Tuning ALBERT for Named Entity Recognition

ALBERT shines in memory-constrained settings. Here we fine-tune it on NER (token classification), labelling each token as PER, ORG, LOC, or O (not an entity).

# ── ALBERT Fine-tuning for NER (CoNLL-2003) ───────────────────
from transformers import (
    AlbertTokenizerFast,
    AlbertForTokenClassification,
    DataCollatorForTokenClassification,
    Trainer, TrainingArguments
)
from datasets import load_dataset
import numpy as np

dataset   = load_dataset("conll2003")
tokeniser = AlbertTokenizerFast.from_pretrained("albert-base-v2")

label_list = dataset["train"].features["ner_tags"].feature.names
# ['O','B-PER','I-PER','B-ORG','I-ORG','B-LOC','I-LOC','B-MISC','I-MISC']
num_labels = len(label_list)

# ── Align labels with subword tokens ──────────────────────────
def tokenise_and_align(examples):
    tok = tokeniser(
        examples["tokens"],
        truncation=True,
        is_split_into_words=True
    )
    aligned_labels = []
    for i, label in enumerate(examples["ner_tags"]):
        word_ids   = tok.word_ids(batch_index=i)
        prev_word  = None
        label_ids  = []
        for wid in word_ids:
            if wid is None:
                label_ids.append(-100)   # ignore [CLS], [SEP]
            elif wid != prev_word:
                label_ids.append(label[wid])
            else:
                label_ids.append(-100)   # ignore sub-tokens
            prev_word = wid
        aligned_labels.append(label_ids)
    tok["labels"] = aligned_labels
    return tok

tokenised  = dataset.map(tokenise_and_align, batched=True)
collator   = DataCollatorForTokenClassification(tokeniser)

model = AlbertForTokenClassification.from_pretrained(
    "albert-base-v2",
    num_labels=num_labels
)

args = TrainingArguments(
    output_dir="./albert-ner",
    num_train_epochs=5,
    per_device_train_batch_size=32,   # large batch — ALBERT is small!
    learning_rate=3e-5,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    fp16=True
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=tokenised["train"],
    eval_dataset=tokenised["validation"],
    data_collator=collator,
    tokenizer=tokeniser
)
trainer.train()

Fine-Tuning DistilBERT for Question Answering

DistilBERT's strength is production speed. Here we fine-tune it on extractive QA (SQuAD-style) — given a passage and a question, find the answer span within the passage.

# ── DistilBERT Fine-tuning for Extractive QA (SQuAD v1) ───────
from transformers import (
    DistilBertTokenizerFast,
    DistilBertForQuestionAnswering
)
from datasets import load_dataset
import torch

dataset   = load_dataset("squad", split="train[:5000]")
tokeniser = DistilBertTokenizerFast.from_pretrained("distilbert-base-uncased")

# ── Tokenise + find start/end positions ───────────────────────
def preprocess(examples):
    inputs = tokeniser(
        examples["question"],
        examples["context"],
        truncation="only_second",
        max_length=384,
        stride=128,
        return_overflowing_tokens=True,
        return_offsets_mapping=True,
        padding="max_length"
    )

    offset_mapping = inputs.pop("offset_mapping")
    sample_map     = inputs.pop("overflow_to_sample_mapping")
    answers        = examples["answers"]
    start_positions, end_positions = [], []

    for i, offset in enumerate(offset_mapping):
        sample_idx = sample_map[i]
        answer     = answers[sample_idx]
        start_char = answer["answer_start"][0]
        end_char   = start_char + len(answer["text"][0])

        seq_ids = inputs.sequence_ids(i)
        # Find context token range
        ctx_start = next(j for j, s in enumerate(seq_ids) if s == 1)
        ctx_end   = next(j for j in range(len(seq_ids)-1, -1, -1) if seq_ids[j] == 1)

        if offset[ctx_start][0] > start_char or offset[ctx_end][1] < end_char:
            start_positions.append(0); end_positions.append(0)
        else:
            start_positions.append(next(
                j for j in range(ctx_start, ctx_end+1)
                if offset[j][0] <= start_char < offset[j][1] or offset[j][0] == start_char
            ))
            end_positions.append(next(
                j for j in range(ctx_end, ctx_start-1, -1)
                if offset[j][1] >= end_char
            ))

    inputs["start_positions"] = start_positions
    inputs["end_positions"]   = end_positions
    return inputs

tokenised = dataset.map(preprocess, batched=True, remove_columns=dataset.column_names)
model     = DistilBertForQuestionAnswering.from_pretrained("distilbert-base-uncased")

# ── Training ───────────────────────────────────────────────────
from transformers import Trainer, TrainingArguments
args = TrainingArguments(
    output_dir="./distilbert-squad",
    num_train_epochs=2,
    per_device_train_batch_size=16,
    learning_rate=5e-5,
    fp16=True
)
trainer = Trainer(model=model, args=args, train_dataset=tokenised)
trainer.train()

# ── Inference ──────────────────────────────────────────────────
from transformers import pipeline

qa_pipe = pipeline(
    "question-answering",
    model="./distilbert-squad",
    tokenizer=tokeniser
)
result = qa_pipe({
    "question": "Where did the Battle of Hastings take place?",
    "context":  "The Battle of Hastings was fought on 14 October 1066 near Hastings, East Sussex, England between the Norman-French army of William the Conqueror and the English army of King Harold II."
})
print(f"Answer:     {result['answer']}")
print(f"Confidence: {result['score']:.4f}")

Advanced — Building Your Own Distillation Pipeline

You are not limited to the pre-distilled DistilBERT. You can distil from any fine-tuned teacher — for example, a RoBERTa teacher fine-tuned on your domain data, distilled into a smaller student model for production deployment.

# ── Custom Knowledge Distillation: RoBERTa → Small Student ────
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import (
    RobertaForSequenceClassification,
    DistilBertForSequenceClassification,
    RobertaTokenizer
)

# Assume teacher is already fine-tuned on your task
teacher = RobertaForSequenceClassification.from_pretrained("./roberta-finetuned")
student = DistilBertForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=2
)
teacher.eval()   # teacher is frozen during distillation

class DistillationLoss(nn.Module):
    def __init__(self, temperature=4.0, alpha=0.5):
        super().__init__()
        self.T     = temperature   # soften probability distributions
        self.alpha = alpha         # weight: distil vs hard label loss

    def forward(self, student_logits, teacher_logits, true_labels):
        # Soft distillation loss (KL divergence)
        soft_loss = F.kl_div(
            F.log_softmax(student_logits / self.T, dim=-1),
            F.softmax(teacher_logits / self.T, dim=-1),
            reduction="batchmean"
        ) * (self.T ** 2)

        # Hard label (cross-entropy) loss
        hard_loss = F.cross_entropy(student_logits, true_labels)

        # Combined loss: alpha controls the trade-off
        return self.alpha * soft_loss + (1 - self.alpha) * hard_loss

loss_fn   = DistillationLoss(temperature=4.0, alpha=0.7)
optimiser = torch.optim.AdamW(student.parameters(), lr=3e-5)

# ── Training loop ──────────────────────────────────────────────
for epoch in range(3):
    for batch in train_loader:
        input_ids      = batch["input_ids"]
        attention_mask = batch["attention_mask"]
        labels         = batch["labels"]

        # Teacher forward pass (no gradient)
        with torch.no_grad():
            teacher_out = teacher(
                input_ids=input_ids,
                attention_mask=attention_mask
            )

        # Student forward pass
        student_out = student(
            input_ids=input_ids,
            attention_mask=attention_mask
        )

        loss = loss_fn(
            student_out.logits,
            teacher_out.logits,
            labels
        )

        optimiser.zero_grad()
        loss.backward()
        optimiser.step()

    print(f"Epoch {epoch+1}: loss = {loss.item():.4f}")

When Things Go Wrong — Common Mistakes and Fixes

Production Inference — Loading and Using All Three Models

# ── Unified Inference Examples — All Three Models ─────────────
from transformers import pipeline

# ── RoBERTa: Sentiment Classification ──────────────────────────
roberta_pipe = pipeline(
    "text-classification",
    model="cardiffnlp/twitter-roberta-base-sentiment-latest"
)
print(roberta_pipe("This model is absolutely brilliant!"))
# [{'label': 'positive', 'score': 0.9876}]

# ── ALBERT: Zero-Shot Classification ───────────────────────────
albert_pipe = pipeline(
    "zero-shot-classification",
    model="cross-encoder/nli-albert-small-v2"
)
result = albert_pipe(
    "The Federal Reserve raised interest rates by 25 basis points.",
    candidate_labels=["finance", "sports", "politics", "technology"]
)
print(f"Label: {result['labels'][0]}, Score: {result['scores'][0]:.4f}")
# Label: finance, Score: 0.9543

# ── DistilBERT: Fast Feature Extraction ────────────────────────
distil_pipe = pipeline(
    "feature-extraction",
    model="distilbert-base-uncased",
    return_tensors=True
)
# Get sentence embedding (mean pool the [CLS] vector)
import torch
outputs = distil_pipe("Knowledge distillation compresses large models.")
embedding = torch.tensor(outputs[0]).mean(dim=1)   # [1, 768]
print(f"Sentence embedding shape: {embedding.shape}")
# Sentence embedding shape: torch.Size([1, 768])

# ── Benchmark Inference Speed ──────────────────────────────────
import time
texts = ["This is a test sentence for benchmarking."] * 100

for name, pipe in [(
    "RoBERTa", roberta_pipe), ("DistilBERT", distil_pipe
)]:
    start = time.time()
    _ = [pipe(t) for t in texts]
    elapsed = time.time() - start
    print(f"{name}: {elapsed:.2f}s for 100 sentences ({elapsed*10:.0f}ms/sentence)")

Golden Rules — Non-Negotiable Practices

Final Decision Framework