Whole Word Masking › Using WWM with Hugging Face
Language AI HandbookAll chapters Whole Word Masking › Using WWM with Hugging FaceOne section · Sections in this chapter9 / 13 ↓
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
The Hugging Face transformers library provides built-in support for whole word masking through the DataCollatorForLanguageModeling class with whole_word_mask=True. This collator handles all the boundary detection and masking logic automatically, making it straightforward to incorporate WWM into any training pipeline.
from transformers import DataCollatorForLanguageModeling
# Create WWM data collator
wwm_collator = DataCollatorForLanguageModeling(
tokenizer=tokenizer, mlm=True, mlm_probability=0.15, whole_word_mask=True
)Input IDs shape: torch.Size([2, 12]) Labels shape: torch.Size([2, 12]) First sentence after WWM: Input: ['[CLS]', 'the', 'quick', 'brown', 'fox', 'jumps', 'over', 'the', 'lazy', '[MASK]', '[MASK]', '[SEP]'] Labels: ['_', '_', '_', '_', '_', '_', '_', '_', '_', 'dog', '.', '_']
The output shows tensors with matching shapes for inputs and labels. The [MASK] tokens appear in the input sequence, while the labels tensor stores the original token IDs at those positions and -100 everywhere else. The Hugging Face collator uses the offset_mapping to determine which tokens in the tokenized output correspond to the same span in the original text. This provides a reliable source of word boundary information that works correctly for diverse inputs including text with unusual punctuation, numbers, or mixed-script content.
One practical note: the collator applies masking dynamically at batch creation time rather than pre-computing masked versions of the dataset. This means each training epoch sees a different random masking of the same sentences, which is a form of data augmentation that helps the model generalize. It also means the collator introduces a small computational overhead per batch, though this is negligible compared to the forward and backward passes through the model.
Training with WWM
To train a model with WWM, simply use the collator in your training pipeline. The rest of the setup is identical to standard MLM pre-training:
# Training setup (not executed; shown for reference)
# from transformers import Trainer, TrainingArguments
#
# training_args = TrainingArguments(
# output_dir="./bert-wwm",
# per_device_train_batch_size=16,
# num_train_epochs=3,
# learning_rate=5e-5,
# )
#
# trainer = Trainer(
# model=model,
# args=training_args,
# train_dataset=train_dataset,
# data_collator=wwm_collator, # WWM applied here
# )
#
# trainer.train()
print("Training configuration shown above (not executed in notebook).")The training loop applies WWM dynamically to each batch. This ensures different masking patterns on each epoch. Because masking is stochastic and applied at collation time, the same sentence will present different masked tokens across different training steps, effectively multiplying the number of distinct training examples. A corpus of 1 million sentences, trained for 3 epochs, exposes the model to roughly 3 million distinct masking patterns rather than 3 repetitions of the same 1 million patterns.
One configuration consideration: when using WWM for continued pre-training (fine-tuning a pre-trained BERT on a domain-specific corpus), you may want to use a lower learning rate and fewer epochs than for pre-training from scratch. The model has already developed strong general representations; continued pre-training with WWM should refine them for the domain vocabulary rather than overwrite them.