Whole Word Masking › Key Parameters

When implementing whole word masking, several parameters control the masking behavior. Understanding what each parameter does and how to set it helps you tailor WWM to your specific pre-training setup.

• mlm_probability (default: 0.15): The target fraction of tokens to mask. WWM aims to achieve this ratio at the token level while selecting at the word level. Values between 0.10 and 0.20 are reasonable; lower values provide less training signal per sequence but make each prediction task easier, while higher values provide more signal but may make individual predictions too difficult. For continued pre-training on a new domain with a smaller corpus, a slightly higher masking probability (0.20) can extract more signal from limited data.

• mask_token: The special token used for masking (e.g., [MASK] for BERT). This must match the tokenizer's mask token exactly for the model to recognize masked positions during both pre-training and any fine-tuning that uses masking. Using a different token accidentally produces a training/inference mismatch.

• 80-10-10 distribution: Controls how selected tokens are modified. The standard split (80% [MASK], 10% random, 10% unchanged) helps bridge the gap between pre-training and fine-tuning, where [MASK] never appears in input. Research has shown that small changes to this distribution have little effect on downstream performance, so the standard 80-10-10 split is a safe default.

• Word boundary detection: Different tokenizers require different detection logic, and getting this right is critical for correctness:

  • WordPiece (BERT, ALBERT): Continuation tokens start with ##
  • SentencePiece (T5, LLaMA, Mistral): Word-initial tokens start with ▁
  • GPT-2 BPE (GPT-2, RoBERTa): Word-initial tokens start with Ġ
  • Custom tokenizers may use other conventions; always verify with example outputs before deploying WWM

• random_state / seed: Setting a seed ensures reproducible masking patterns during evaluation or debugging. During training, avoid fixed seeds to maximize data diversity across epochs: different masking patterns on each pass through the data provide additional implicit augmentation.

• Word group filtering: Special tokens should be filtered from word grouping. The implementation should explicitly exclude [CLS], [SEP], [PAD], and any other non-content special tokens from both grouping and selection. Failing to do this can produce groups that straddle content tokens and special tokens, leading to incorrect masking patterns.