Whole Word Masking › The Whole Word Masking Procedure

Whole Word Masking preserves word boundaries during the masking process. The algorithm requires knowing which subword tokens belong to the same original word, a task that different tokenizers signal differently. WordPiece uses ## prefixes for continuation tokens. SentencePiece uses special Unicode characters at word boundaries. BPE-based tokenizers like GPT-2's use space characters embedded into token representations. In each case, the tokenizer encodes enough information to reconstruct word boundaries, and WWM exploits this encoding to group subwords before selecting masking targets.

Think of the procedure as adding a layer of organization on top of ordinary MLM. Instead of randomly pointing at individual tokens and flipping a coin to decide whether to mask each one, WWM first organizes the tokens into word-groups, then flips coins at the word level. Once a word is selected, all its constituent tokens are masked simultaneously, maintaining the integrity of the word boundary as an atomic unit of prediction.

The procedure works cleanly with the rest of BERT's pre-training setup. The model architecture is unchanged. The input format, with [CLS] at the start and [SEP] at the end, is unchanged. The loss function, which computes cross-entropy only over masked positions, is unchanged. The only modification is in the data collation step: how you decide which positions to mask before feeding a sequence to the model.

The WWM procedure works in three steps:

1. Group subwords into words: Traverse the token sequence and collect consecutive subword tokens that belong to the same word. A new word starts when a token lacks the continuation marker, or equivalently, when a token carries the word-start marker. Special tokens like [CLS] and [SEP] are typically excluded from masking consideration entirely.

2. Select words for masking: Choose words (not tokens) to mask based on the masking probability. The selection targets approximately 15% of the total tokens while operating at the word level. This requires a greedy or probabilistic approach that accounts for the variable number of tokens each word contributes.

3. Mask all subwords together: For each selected word, replace all its constituent subword tokens according to the 80-10-10 rule (discussed below). Every subword of the selected word receives the same treatment. This ensures no partial information leaks through visible siblings.

The key insight is that we are changing the unit of selection from subword tokens to whole words while maintaining approximately the same masking ratio in terms of total tokens masked. The model's loss computation does not change: it still predicts masked tokens individually, one position at a time. But the masking pattern is now constrained to respect word boundaries, which dramatically changes the difficulty and nature of the prediction task.

Standard MLM masks exactly 15% of tokens per sequence. With WWM, we shift from selecting individual tokens to selecting whole words, but we still want approximately 15% of tokens to end up masked. This creates an interesting problem: words have different lengths in subword tokens, so selecting a fixed fraction of words does not yield a fixed fraction of tokens.

The variable-length complication

Consider a concrete example. You have a sentence with 20 tokens forming 10 words:

• 6 words are single tokens (like "the", "is", "a")
• 3 words are two tokens each (like "trans##form", "learn##ing")
• 1 word is four tokens (like "un##believ##ab##ly")

If you randomly select 15% of words (roughly 1-2 words), you might mask anywhere from 1 token (if you pick "the") to 4 tokens (if you pick the long word). The actual token masking ratio becomes unpredictable: sometimes 5%, sometimes 20%, depending entirely on which words were chosen.

This variance in masking rate matters because the training objective's difficulty scales directly with how many tokens must be predicted. A sequence with only 5% masking has few prediction targets and contributes little to learning. A sequence with 20% masking asks the model to reconstruct many tokens simultaneously, which may be too hard and produce noisy gradients. The 15% target exists precisely because it was found empirically to balance learning efficiency against task difficulty.

A length-weighted selection approach

One principled solution weights word selection by length. The probability of masking each word becomes proportional to how many tokens it contributes, normalized to hit the 15% target overall:

$$P(\text{mask word } w) = \frac{0.15 \times |w|}{\sum_{i} |w_i|}$$

where:

• $P(\text{mask word } w)$: the probability of selecting word $w$ for masking
• $|w|$: the number of subword tokens in word $w$
• $\sum_{i} |w_i|$: the total number of subword tokens in the sentence (summed across all words)
• $0.15$: the target masking ratio

The formula works by giving each word a selection probability proportional to its token count. A 4-token word is four times more likely to be selected than a 1-token word. This compensates for the fact that selecting the long word contributes four times as many masked tokens. In expectation, exactly 15% of all tokens get masked.

Why does this formula make sense? Notice that if all words were single tokens, every word would have probability $0.15 / \sum_i |w_i|$, and exactly 15% of words (and therefore tokens) would be selected in expectation. The length weighting generalizes this logic: each token has the same expected probability of being masked, regardless of which word it belongs to. Long words appear more often in the selection set, but each appearance masks more tokens, so the per-token masking probability stays constant.

Why this approach has drawbacks

The length-weighted approach achieves precise 15% token masking in expectation, but it introduces a systematic bias: longer words get masked more often. Since longer words tend to be rarer and more complex (like "internationalization" or "photosynthetically"), the model sees these words masked disproportionately often during training. Common short words like "the", "is", "in" are rarely masked. The model might develop weaker contextual representations for common words and over-rely on positional or structural cues to predict them when they do appear masked.

There is also a practical implementation complexity: computing individual masking probabilities for each word requires knowing the full sequence length upfront, which complicates streaming or batched data loading.

The practical solution

Production implementations typically use a simpler greedy approach that avoids both the variance problem and the complexity of length-weighted probabilities:

1. Shuffle the list of words randomly
2. Add words to the masking set one by one, in shuffled order
3. Stop when the total number of selected tokens reaches or exceeds approximately 15% of the sequence length

This method is unbiased across word lengths (each word has an equal chance of being near the front of the shuffled list) and easy to implement in any data loading framework. The masking ratio varies slightly per sequence (sometimes 12%, sometimes 18% depending on where the cutoff lands), but these small variations average out across a training batch of 32 or more sequences. Empirically, this variance does not destabilize model training.

The greedy approach also has a nice practical property: it terminates cleanly. Length-weighted selection can theoretically select all words (if all are independently drawn with probability proportional to their length), creating a sequence that is 100% masked. The greedy approach hits the target and stops, maintaining a controllable masking budget per sequence.

BERT's original masking strategy applies a probabilistic treatment after the binary mask/do-not-mask selection. Of the 15% of tokens selected for masking:

• 80%: Replace with [MASK]
• 10%: Replace with a random token from the vocabulary
• 10%: Keep the original token unchanged

This three-way split serves an important purpose. During fine-tuning and inference, [MASK] tokens never appear in the input. If pre-training used only [MASK] replacements, the model would develop a representation for the special [MASK] token that does not correspond to any real-world token distribution it encounters during fine-tuning. The 10% random replacement forces the model to maintain context-sensitive representations even for non-mask tokens (since any token might secretly be the wrong one). The 10% keep-original forces the model to produce useful representations for correct tokens (since it does not know whether any given token has been replaced).

With WWM, this rule extends naturally to word boundaries. The key design decision is whether all subwords of a selected word receive the same treatment or independent treatments. The standard approach applies the same treatment to all subwords of a word: if the word is selected for random replacement, all its subwords are replaced with random tokens. If selected for keep-original, all subwords remain unchanged. If selected for masking (the 80% case), all subwords become [MASK].

This consistency within words is important. Consider what happens if we apply the 80-10-10 rule independently to each subword of "transformation". We might mask "##ation" but keep "transform" original, reproducing exactly the information leakage problem we set out to solve. By applying the same decision to the entire word, we maintain the integrity of word-level masking throughout the distribution.

A subtle implementation note: the 80-10-10 decision is applied once per selected word, not once per selected token. This means a word with three subword tokens does not make three independent rolls of the dice; it makes one roll and all three subwords follow that outcome. The resulting token-level statistics differ slightly from the original BERT specification (where each masked position independently makes the 80-10-10 decision), but the difference is negligible in practice.