The 2019 paper “RoBERTa: A Robustly Optimized BERT Pretraining Approach” by Liu et al. (Facebook AI) is essentially a “masterclass in hyperparameter tuning.” The authors found that BERT was significantly undertrained and that several of its design choices were suboptimal.
RoBERTa uses the exact same architecture as BERT but changes the pre-training strategy. To implement RoBERTa in PyTorch, we take our BERT implementation and apply four critical modifications.
1. Dynamic Masking
In the original BERT, masking was done once during data preprocessing (Static Masking). RoBERTa performs masking dynamically every time a sequence is fed into the model during training.
Python
import torch
import torch.nn as nn
class RoBERTaMasker:
def __init__(self, vocab_size, mask_token_id, pad_token_id):
self.vocab_size = vocab_size
self.mask_id = mask_token_id
self.pad_id = pad_token_id
def mask_tokens(self, inputs):
"""
Dynamic Masking: Create a new mask for the batch on the fly.
"""
labels = inputs.clone()
# Sample 15% of tokens for masking
probability_matrix = torch.full(labels.shape, 0.15)
# Don't mask padding
probability_matrix.masked_fill_(labels == self.pad_id, value=0.0)
masked_indices = torch.bernoulli(probability_matrix).bool()
labels[~masked_indices] = -100 # Only calculate loss on masked tokens
# 80% of the time, replace masked input tokens with [MASK]
indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
inputs[indices_replaced] = self.mask_id
# 10% of the time, replace with random token
indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced
random_words = torch.randint(self.vocab_size, labels.shape, dtype=torch.long)
inputs[indices_random] = random_words[indices_random]
# The rest (10%) keep the original token
return inputs, labels
2. Full-Sentences (Removing NSP)
RoBERTa removes the Next Sentence Prediction (NSP) task entirely. Instead, it uses “Full-Sentences”: input sequences are packed with full sentences sampled contiguously from one or more documents until they reach the maximum length (512 tokens).
Python
class RoBERTa(nn.Module):
def __init__(self, vocab_size, hidden_size, n_layers, n_heads, d_ff):
super().__init__()
# Note: No Segment Embeddings needed since NSP is removed
self.embedding = nn.Embedding(vocab_size, hidden_size)
self.pos_embedding = nn.Embedding(512, hidden_size)
self.layers = nn.ModuleList([
EncoderLayer(hidden_size, n_heads, d_ff, 0.1) for _ in range(n_layers)
])
# RoBERTa only has the Masked LM head
self.lm_head = nn.Linear(hidden_size, vocab_size)
self.norm = nn.LayerNorm(hidden_size)
def forward(self, x, mask):
# x: [batch_size, seq_len]
seq_len = x.size(1)
pos = torch.arange(seq_len).unsqueeze(0).to(x.device)
x = self.embedding(x) + self.pos_embedding(pos)
for layer in self.layers:
x = layer(x, mask)
return self.lm_head(x) # [batch_size, seq_len, vocab_size]
3. Byte-Level BPE
BERT used a character-level WordPiece vocabulary (30k). RoBERTa uses a larger Byte-level Byte-Pair Encoding (BPE) vocabulary (50k). This allows the model to learn a representation for every sub-word without ever hitting an “Unknown” ([UNK]) token, even with complex or misspelled text.
4. Comparison: BERT vs. RoBERTa
The architectural changes are subtle, but the training changes are massive.
| Feature | BERT | RoBERTa |
| Masking | Static (Pre-calculated) | Dynamic (Every epoch) |
| Tasks | MLM + NSP | MLM only |
| Batch Size | 256 | 8,000 (Massively Parallel) |
| Training Steps | 1M steps | 500k steps (but much larger batches) |
| Vocabulary | 30k WordPiece | 50k Byte-BPE |
Key Takeaway
RoBERTa proved that the original BERT architecture was not the bottleneck; rather, it was the duration and scale of training. By training on 160GB of text (instead of BERT’s 16GB) and using much larger batches with dynamic masking, RoBERTa significantly outperformed BERT on almost every benchmark, including GLUE and SQuAD.
Leave a Reply