BERT Pre-training › The MLM Objective in Practice
Language AI HandbookAll chapters BERT Pre-training › The MLM Objective in PracticeOne section · Sections in this chapter3 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
We covered masked language modeling conceptually in the previous chapter. Here we see exactly how BERT applies it during pre-training.
For each training sequence, BERT randomly selects 15% of tokens for prediction. Of these selected tokens:
- 80% are replaced with the
[MASK]token - 10% are replaced with a random vocabulary token
- 10% are left unchanged
The model must predict the original token at all selected positions, regardless of how they were corrupted.
from transformers import BertTokenizer
# Load BERT tokenizer
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
def apply_bert_masking(input_ids, tokenizer, mask_prob=0.15):
"""
Apply BERT's MLM masking strategy.
Args:
input_ids: Token IDs to mask
tokenizer: Tokenizer with mask_token_id and vocab
mask_prob: Probability of selecting each token
Returns:
masked_ids: Token IDs after masking
labels: Original IDs at masked positions, -100 elsewhere
"""
labels = input_ids.clone()
masked_ids = input_ids.clone()
# Probability matrix for masking
probability_matrix = torch.full(input_ids.shape, mask_prob)
# Don't mask special tokens
special_tokens_mask = torch.tensor(
[
1
if token_id
in [
tokenizer.cls_token_id,
tokenizer.sep_token_id,
tokenizer.pad_token_id,
]
else 0
for token_id in input_ids.tolist()
],
dtype=torch.bool,
)
probability_matrix.masked_fill_(special_tokens_mask, 0.0)
# Sample masked positions
masked_indices = torch.bernoulli(probability_matrix).bool()
# Labels are -100 for non-masked (ignored in loss)
labels[~masked_indices] = -100
# 80% -> [MASK]
indices_replaced = (
torch.bernoulli(torch.full(input_ids.shape, 0.8)).bool()
& masked_indices
)
masked_ids[indices_replaced] = tokenizer.mask_token_id
# 10% -> random token
indices_random = (
torch.bernoulli(torch.full(input_ids.shape, 0.5)).bool()
& masked_indices
& ~indices_replaced
)
random_words = torch.randint(
len(tokenizer), input_ids.shape, dtype=torch.long
)
masked_ids[indices_random] = random_words[indices_random]
# Remaining 10% stay unchanged but still have labels
return masked_ids, labels# Demonstrate masking on a sample sentence
text = "The quick brown fox jumps over the lazy dog."
encoding = tokenizer(text, return_tensors="pt", padding=False)
input_ids = encoding["input_ids"].squeeze()
masked_ids, labels = apply_bert_masking(input_ids, tokenizer)Original text: The quick brown fox jumps over the lazy dog. Token-level view: Position Original Masked Label -------------------------------------------------- 0 [CLS] [CLS] ignore 1 the the ignore 2 quick quick ignore 3 brown brown ignore 4 fox fox ignore 5 jumps jumps ignore 6 over over ignore 7 the the ignore 8 lazy lazy ignore 9 dog [MASK] dog 10 . . ignore 11 [SEP] [SEP] ignore
The output shows which tokens were selected for masking and how each was treated. Positions with "ignore" labels don't contribute to the MLM loss. The model learns from the positions that do have labels, whether they show [MASK], a random token, or the original token.
Whole Word Masking
The original BERT masked individual WordPiece tokens. This created a problem: when a word like "playing" tokenizes to ["play", "##ing"], masking only "##ing" gives the model an easy hint.
BERT later adopted whole word masking (WWM), where all tokens from the same word are masked together. This provides a stronger learning signal because the model cannot peek at sibling tokens.
def apply_whole_word_masking(input_ids, tokenizer, mask_prob=0.15):
"""
Apply whole word masking for BERT.
When any subword of a word is selected, all subwords are masked.
"""
labels = input_ids.clone()
masked_ids = input_ids.clone()
# Get tokens and identify word boundaries
tokens = tokenizer.convert_ids_to_tokens(input_ids.tolist())
# Group tokens into words
word_indices = [] # List of (start_idx, end_idx) for each word
current_word_start = None
for i, token in enumerate(tokens):
if token in [
tokenizer.cls_token,
tokenizer.sep_token,
tokenizer.pad_token,
]:
continue
if not token.startswith("##"):
# Start of new word
if current_word_start is not None:
word_indices.append((current_word_start, i))
current_word_start = i
# Don't forget last word
if current_word_start is not None:
word_indices.append((current_word_start, len(tokens)))
# Select words to mask (15% of words, not tokens)
num_to_mask = max(1, int(len(word_indices) * mask_prob))
words_to_mask = random.sample(
word_indices, min(num_to_mask, len(word_indices))
)
# Apply masking to all tokens in selected words
for start, end in words_to_mask:
for i in range(start, end):
labels[i] = input_ids[i]
rand = random.random()
if rand < 0.8:
masked_ids[i] = tokenizer.mask_token_id
elif rand < 0.9:
masked_ids[i] = random.randint(0, len(tokenizer) - 1)
# else: keep original (10%)
# Set non-masked positions to -100
for i in range(len(labels)):
is_masked = any(start <= i < end for start, end in words_to_mask)
if not is_masked:
labels[i] = -100
return masked_ids, labels# Demonstrate whole word masking
text_wwm = "The unbelievable transformation surprised everyone."
encoding_wwm = tokenizer(text_wwm, return_tensors="pt", padding=False)
input_ids_wwm = encoding_wwm["input_ids"].squeeze()
masked_ids_wwm, labels_wwm = apply_whole_word_masking(input_ids_wwm, tokenizer)Original text: The unbelievable transformation surprised everyone. Tokenization: [CLS] the unbelievable transformation surprised everyone . [SEP] After whole word masking: [CLS] the unbelievable transformation surprised [MASK] . [SEP]
Notice how multi-token words get fully masked. The model cannot see any part of the word "unbelievable" or "transformation" if they are selected for masking. This gives the model representations that depend less on subword fragments than token-level masking does.
