Whole Word Masking for BERT Pretraining

When BERT tokenizes "undeniably" into ["un", "##deni", "##ably"] and then masks only "##ably", something goes wrong. The model sees "un" and "##deni" in the clear, giving it strong hints about the masked portion. The prediction task becomes almost trivial: what word starts with "undeni-" and ends with a common suffix? This partial visibility undermines the entire learning objective and turns a sophisticated language modeling task into something closer to a morphology lookup table.

Whole Word Masking (WWM) fixes this by treating subword tokens as parts of atomic units. If any subword of a word is selected for masking, all subwords of that word are masked together. The model must now predict the entire word from surrounding context alone, without peeking at sibling subwords. This change improves downstream task performance, particularly on reading comprehension benchmarks and tasks that require semantic interpretation rather than pattern-matching at the subword level.

The idea behind WWM reflects a deeper principle: the unit of prediction should match the unit of meaning. Words carry meaning; subword fragments typically do not. When you mask "##ation" while leaving "transform" visible, you are asking the model to guess a suffix, not a meaning. When you mask "transformation" entirely, you are asking the model to reason about what concept fits the surrounding sentence. These are fundamentally different cognitive tasks, and only the second one trains the kind of contextual understanding that makes BERT useful downstream.

In this chapter, we examine why subword masking creates problems, how whole word masking works algorithmically at a precise level, and how to implement it correctly for different tokenizer families. We also look at the 80-10-10 masking rule and how it extends naturally to word boundaries, the challenge of maintaining a consistent masking ratio when word lengths vary, and the practical question of what happens in morphologically rich languages where word boundaries carry even more semantic weight. We then compare WWM against random subword masking with real BERT predictions to see the empirical differences in prediction difficulty.

Modern language models use subword tokenization to handle vocabulary efficiently. Algorithms like WordPiece and BPE, along with SentencePiece, break words into smaller units based on frequency or likelihood, balancing vocabulary size against the ability to represent any input without unknown-token fallbacks. Common words like "the", "is", and "run" remain intact as single tokens. Moderate-frequency words decompose into recognizable pieces: "transformation" becomes ["transform", "##ation"]. Rare or highly specific words may split into many fragments: "photosynthetically" might become four or five subword pieces.

This design is essential for handling open vocabularies and morphologically diverse text. A fixed-word vocabulary of 30,000 entries would fail on technical jargon, proper nouns, or even common words in inflected forms. Subword tokenization sidesteps this problem elegantly. But subword tokenization interacts badly with random masking because it creates a fundamental asymmetry between single-token and multi-token words.

Think of it as two different exam questions for the same concept. The first question shows you a word with the last three letters blanked out: "transfor___". The second question hides the entire word: "The _________ was complete." Both questions test knowledge about the same word, but they test very different things. The first tests your knowledge of English morphology, the second tests your understanding of the sentence's meaning. When you train BERT with random subword masking, a significant fraction of the "questions" the model faces are of the first type: morphology lookups rather than semantic reasoning.

This decomposition creates an asymmetry in the masking process. When we randomly select 15% of tokens for masking, we are selecting subword tokens, not words. A multi-token word might have some subwords masked while others remain visible. The visible subwords leak information about the masked ones. The degree of leakage depends on how much information a prefix or suffix carries about the rest of the word, and for many common English morphological patterns, this information is substantial.

Consider how BERT processes the sentence "The transformation was remarkable." Using WordPiece tokenization:

If random masking selects only "##ation", the input becomes:

The model sees "transform" immediately adjacent to the mask. How many English words start with "transform-"? Only a handful: transformation, transformed, transforming, transformer. The prediction task collapses from choosing among 30,000 vocabulary items to distinguishing between 3-4 suffixes. With a bit of context bias (this is probably a noun form given the sentence structure), the model can essentially narrow it down to one answer without consulting the rest of the sentence at all.

This is problematic for several reasons:

• Weak learning signal: The model learns to pattern-match subword combinations rather than understand context deeply. It does not need to reason about meaning, just morphology. Gradient updates during backpropagation reflect this: they encode "##ation follows transform-" rather than "this word fits here because of its semantic role."

• Uneven difficulty: Single-token words like "was" face the full prediction challenge when masked, while multi-token words get easy hints. The model develops uneven representations, with single-token words anchored more firmly in their semantic contexts and multi-token words understood more through their morphological structure.

• Distributional mismatch: During fine-tuning, the model sees complete words. Pre-training on partial words creates a subtle distribution shift. The model has learned to rely on subword co-occurrence statistics that simply do not apply during inference, where every word appears complete.

• Gradient dilution across multi-token words: When only one subword of a word is masked, only that subword's prediction contributes to the loss. The model updates its parameters based on a small, easy prediction, rather than the harder full-word prediction that would produce richer gradient signal.

The problem compounds across training. Early in training, the model might struggle with suffix prediction. But it quickly learns the most common suffix patterns for each root, at which point those examples become nearly loss-free and contribute almost nothing to further learning. The model has overfit to morphological patterns and underlearned semantic patterns, precisely the opposite of what we want for a general-purpose language model.

To understand exactly how much information leaks when sibling subwords are visible, we need a way to measure uncertainty. How hard is the prediction task? Information theory gives us a precise tool: entropy.

The intuition behind entropy

Imagine you are playing a guessing game. If someone picks a number between 1 and 1,000, you have high uncertainty: you need many guesses or much information to identify the answer. If they pick between 1 and 2, you have low uncertainty: one yes-or-no question resolves it. Entropy quantifies this precisely: it measures how many "bits" of information you need to identify the answer. More possible answers with similar probabilities means higher entropy; fewer possibilities or more concentrated probability means lower entropy.

For language model predictions, entropy captures how "spread out" the model's probability distribution is over the vocabulary. When the model is very confident about one token, entropy is low. When it is nearly equally uncertain between many tokens, entropy is high. A good training signal comes from high-entropy situations: the model must work hard to make the right prediction, and the resulting gradient update carries rich information about what the model still needs to learn.

The entropy formula

Given a masked position, the model outputs a probability distribution over all possible tokens. We want to quantify how uncertain that distribution is. The conditional entropy of this prediction is:

where:

• $H(x_{\text{mask}} \mid x_{\text{context}})$: the conditional entropy of the masked token given visible context, measured in bits (if using $\log_2$) or nats (if using natural log)
• $x_{\text{mask}}$: the masked token the model must predict
• $x_{\text{context}}$: the visible tokens surrounding the masked position
• $V$: the vocabulary of all possible tokens
• $P(v \mid x_{\text{context}})$: the probability the model assigns to token $v$ being the correct prediction given the context

The formula works by taking each possible token, weighting its log-probability by the probability itself, and summing. The negative sign ensures we get a positive number (since log probabilities are negative). When probability concentrates on one token, that token's contribution $P \log P$ is large and negative, but all others contribute near zero, so the total sum (before negation) is a large negative number, and after negation the entropy is small. When probability spreads uniformly across many tokens, many moderate contributions sum to a large negative number, and entropy is high.

Why does this formula make sense? Notice that $-\log P(v)$ measures how "surprising" token $v$ would be: common tokens have probabilities near 1, making their log near 0 (unsurprising), while rare tokens have small probabilities and large negative logs (very surprising). Entropy is the expected surprise under the model's own distribution. High entropy means the model expects to be frequently surprised, so it is uncertain and must rely on context to narrow down the answer.

Entropy bounds tell us about task difficulty

Two extreme cases illuminate the formula's behavior:

1. Maximum entropy: When all $|V|$ tokens are equally likely, each has probability $1/|V|$, and entropy reaches $\log |V|$. For a vocabulary of 30,000 tokens, that is about 14.9 bits: a difficult prediction requiring deep contextual reasoning.

2. Minimum entropy: When one token has probability 1.0 and all others have probability 0, entropy equals 0. The prediction is trivial; the model is certain before even reasoning about context.

The visualization shows how dramatically entropy drops when the effective vocabulary shrinks. Partial masking constrains predictions to just a few valid suffixes, collapsing entropy from nearly 15 bits to around 2.6 bits. This roughly 6x reduction in uncertainty means the model can achieve low loss without deep contextual understanding. The shaded region at the lower left represents the regime of weak training signal where partial masking operates: probability mass concentrated among a tiny number of morphologically constrained options.

Why partial masking destroys the training signal

Now we can quantify the information leakage problem precisely. When "transform" is visible and we mask only "##ation", what happens to the entropy?

The visible prefix constrains the possibilities dramatically. The only valid continuations are tokens that can follow "transform-" in English words: {##ation, ##ed, ##ing, ##er, ##s, ##able}. The effective vocabulary shrinks from 30,000 to perhaps 6 options. Even with a uniform distribution over these 6, entropy drops from 14.9 bits to about 2.6 bits: a five-fold reduction in prediction difficulty before the model even looks at the surrounding sentence.

In practice, it is even worse. The model learns that ##ation is by far the most common suffix after "transform-", so $P(\texttt{\#\#ation})$ might be 0.7 or higher. The entropy collapses further to perhaps 1.0-1.5 bits. The model has learned a morphological lookup table and uses it to bypass contextual prediction.

This explains why partial masking produces weak gradients. The model achieves low loss not by understanding context, but by memorizing subword co-occurrence patterns. It learns morphology instead of semantics. After enough training, the attention mechanism learns to route heavily toward the visible sibling subwords of any masked token, essentially building a lookup mechanism that short-circuits contextual reasoning. Whole Word Masking closes this shortcut entirely by removing the sibling subwords from the visible context.

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:

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.

Let's walk through a complete WWM application on a concrete sentence to solidify the algorithm. We will use the sentence "The transformation was undeniably remarkable." and trace every step by hand.

Step 1: Tokenize

First, WordPiece tokenizes the sentence:

We have 9 tokens. The 15% target means we want to mask approximately $0.15 \times 9 \approx 1.35$ tokens. Since we operate at the word level, we will aim for 1-2 tokens of masking.

Step 2: Identify word groups

Scanning for ## prefixes, we group tokens into words:

Six words total, with lengths [1, 2, 1, 2, 2, 1]. Total tokens: 9.

Step 3: Shuffle and greedily select words

Suppose the random shuffle produces the order: [Word 4, Word 1, Word 5, Word 2, Word 3, Word 6].

We step through this order, accumulating token count until we reach approximately 15% (about 1-2 tokens):

• Consider Word 4 ("undeni", "##ably"): total masked would be 2 tokens (22%). This is close to our target and within acceptable variance, so we select it and stop.

The selected word is Word 4. This gives us a masking ratio of 2/9 = 22%, slightly higher than 15%, which is acceptable variance for a short sentence.

Step 4: Apply the 80-10-10 rule

Roll the dice for Word 4: suppose the outcome is 0.73 (less than 0.80), so we apply masking. Both "undeni" and "##ably" become [MASK]:

The label array stores the original token IDs at masked positions and -100 everywhere else. During training, the loss is computed only at positions with non-(-100) labels.

Step 5: Observe the learning signal

The model receives the masked sequence and must predict the tokens at positions 4 and 5. The visible context is: "The transformation was [MASK] [MASK] remarkable." To predict what fills the blank, the model must reason about what adverb-adjective combination could fit between "was" and "remarkable". This requires contextual interpretation: something that means "completely, without doubt, certainly" makes sense here. The model cannot rely on morphological patterns because neither the stem nor the suffix is visible. This is precisely the learning signal we want.

Let's implement whole word masking step by step. We will start with the core logic for identifying word boundaries, then build the complete masking function.

The first task is grouping subword tokens into words. WordPiece tokenizers mark continuation tokens with a ## prefix. The grouping logic traverses the token list and creates a new group whenever it encounters a token without the ## prefix, which signals the start of a new word.

The function correctly identifies word boundaries by detecting the ## prefix. "The" remains a single-token word, while "transformation" is grouped as ["transform", "##ation"] and "undeniably" as ["undeni", "##ably"]. Each tuple represents the start and end indices in the token list as a half-open interval (start inclusive, end exclusive), making it easy to slice and mask all subwords of a word together using standard Python indexing.

Notice that the function handles both single-token and multi-token words uniformly. A single-token word like "was" produces a group spanning a single index. A three-token word would produce a group spanning three indices. This uniformity simplifies the masking code that follows, since the same loop structure works regardless of word length.

Now we implement the complete WWM function. We will use the greedy word selection approach (shuffle words, accumulate until we hit the token budget) and apply the 80-10-10 rule at the word level.

The function masks entire words together. When "jumps" is selected, both the root and any continuation subwords receive the same treatment. The labels array stores the original token IDs for computing the loss, with -100 marking positions to ignore during loss computation. This -100 convention is standard in PyTorch's cross-entropy loss, which skips positions with this sentinel value automatically.

The masking loop applies the 80-10-10 rule independently to each subword token within a selected word. This means there is a small chance that different subwords of the same word get different treatments (one masked, another randomly replaced). The more principled approach applies one dice roll per word, but the difference in practice is negligible and the independent-per-token approach is what many library implementations use.

Let's implement random subword masking for comparison. This is the original BERT approach: select individual tokens uniformly at random, with no regard for word boundaries.

In random subword masking, you might see "transform" with "##ational" masked, revealing the word structure through the visible prefix. In WWM, the entire word is hidden, forcing the model to reason about what concept fits the context rather than simply completing a visible stem.

The output also reveals another practical difference: random subword masking creates unpredictable patterns where some multi-token words have both subwords masked (by chance), some have none masked, and some have a mix. WWM produces cleaner masking patterns aligned with linguistic units. This alignment makes the masked language modeling task more consistent across different word types and lengths.

Let's create a visualization comparing how both masking strategies affect a set of sentences with multi-token words. We use a higher masking probability (30%) to make the patterns more visible in the diagram.

The visualization highlights how random subword masking creates fragmented masking patterns that leak information. In "transformation", random masking might leave "transform" visible while masking "##ation", handing the model a near-trivial suffix prediction task. In "internationalization", which tokenizes into many fragments, random masking might mask only a middle fragment while leaving both prefix and suffix visible, making prediction even easier. WWM eliminates this entirely by treating words as atomic units: all fragments of a selected word are masked simultaneously, or none are.

Different tokenizers use different conventions for marking subword boundaries, and implementing WWM correctly requires adapting to each convention. The underlying principle is the same across all tokenizer families: identify which tokens belong to the same word and treat them atomically during masking. But the detection mechanism differs based on how the tokenizer encodes word boundary information.

Understanding these conventions matters for implementing WWM from scratch and for debugging unexpected tokenizer outputs. If your tokenizer produces tokens in an unusual format (custom special characters, byte-level encoding, etc.), you need to verify that your boundary detection logic correctly identifies word groups before applying masking.

WordPiece uses ## to prefix continuation tokens. A word starts with a token lacking ##, and subsequent ##-prefixed tokens continue it. This convention makes boundary detection straightforward: a token is the start of a new word if and only if it does not begin with ##.

The WordPiece convention is clean and explicit. Every continuation token carries its ## marker, so boundary detection requires no lookahead and no ambiguity. The edge cases are special tokens like [CLS], [SEP], and [PAD], which do not follow the ## convention but also should not be masked. Production implementations filter these out before applying WWM grouping.

One subtlety: tokenizers applied to pre-tokenized text (text that has already been split into words by whitespace or another tokenizer) can produce different results than tokenizers applied to raw text. If you pre-tokenize first and then apply WordPiece to each word separately, the resulting tokens may not have the same ## prefix structure as if you applied WordPiece to the full raw string. Always verify that your tokenization pipeline produces the expected token format before implementing WWM on top of it.

SentencePiece uses the special character ▁ (U+2581, a lower one-eighth block) to mark word boundaries. Unlike WordPiece, this character appears at the start of words that follow a space, not as a continuation marker. Tokens without the ▁ prefix are continuations of the previous word:

The SentencePiece convention is logically inverted from WordPiece: instead of marking continuations, it marks new starts. A token without ▁ is a continuation of the previous word. This design reflects SentencePiece's unified approach to treating text as a sequence of Unicode characters, where spaces are just another character that happens to signal word boundaries in most languages.

For SentencePiece-tokenized text, the first token in a sequence (or the first token after a special prefix token) typically lacks the ▁ marker regardless of whether it starts a word in the traditional sense. This can create an edge case for the first word in a sequence, which you may need to handle explicitly in your grouping logic.

GPT-2's BPE tokenizer uses a different convention: it adds Ġ (representing a space character in the tokenizer's byte-level encoding) at the start of tokens that begin a new word. Tokens without the leading Ġ are continuations:

GPT-2's byte-level BPE encodes text at the byte level, using a bijective mapping from bytes to Unicode characters to avoid the need for unknown tokens. The Ġ character (Unicode U+0120) represents the byte 0x20, which is a space. So Ġ at the start of a token indicates that a space preceded this token in the original text, which means this token starts a new word. This is logically equivalent to the SentencePiece ▁ approach but uses a different Unicode codepoint.

We can create a universal function that detects the tokenizer type and applies the appropriate grouping logic. This is useful when building tools that need to work across multiple model families without requiring users to specify the tokenizer convention manually.

Each tokenizer produces different subword splits for "unhappiness", but the universal function correctly groups them into a single word for all three. BERT uses ## prefixes to mark the continuation pieces. T5 produces a different subword segmentation using ▁. GPT-2's byte-level BPE produces yet another segmentation, but the Ġ detection correctly identifies the word boundary. In each case, the result is a single group spanning all the subword tokens of "unhappiness", which is exactly what we need for whole word masking.

The fallback case (treating each token as a separate word) is safe but conservative: it effectively disables whole word masking for unknown tokenizer formats. A better fallback for production systems might raise an exception to alert the user that their tokenizer format is not recognized, preventing silent degradation to per-token masking.

Let's compare the prediction difficulty under both masking strategies using a real BERT model. We will measure how much information leakage helps the model make correct predictions. This provides concrete evidence for the entropy analysis from earlier sections.

The contrast is striking. With partial masking, the model assigns overwhelming probability to ##ation because it simply completes the visible prefix "transform-". The top prediction dominates, and the remaining vocabulary items receive negligible probability. With whole word masking, the model must consider what noun could fit the context "The [MASK] was remarkable," leading to an uncertain distribution across many plausible words. This uncertainty is the source of rich gradient signal: during training, when the model makes the wrong prediction under WWM, the resulting gradient update carries far more information about the model's limitations than a gradient update from an easy suffix prediction.

We can quantify the difficulty difference precisely using prediction entropy. This connects our theoretical analysis to concrete numbers with a real pre-trained BERT model.

The entropy measurements confirm our intuition quantitatively. Partial masking consistently produces lower entropy because the visible root token constrains the prediction to a handful of morphologically valid suffixes. Whole word masking forces higher entropy in every case, indicating that the model is more uncertain and must use more context to produce a prediction. This uncertainty is not a weakness of WWM; it is the feature. Higher-entropy predictions generate richer gradient updates, which drive the model to build more powerful contextual representations.

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.

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.

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:

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.

Whole Word Masking is a well-motivated improvement over random subword masking, but it comes with trade-offs and boundary conditions that practitioners should understand before deciding whether and how to apply it.

Language-dependent effectiveness: WWM assumes that word boundaries carry semantic significance. This assumption holds well for English and other European languages that use spaces to delimit words, but it becomes complicated for languages without clear word boundaries. Chinese and Japanese are the most prominent examples: Chinese text typically has no spaces between characters, so identifying word boundaries requires running a separate word segmentation system before tokenization. This adds a preprocessing dependency that can introduce segmentation errors, particularly for domain-specific vocabulary, technical terms, or proper nouns that the segmenter was not trained on. For Chinese BERT, the Google team used an external word segmenter to identify word boundaries, then applied WWM on the resulting word-segmented text. Errors in segmentation propagate directly into errors in word grouping, and the effectiveness of WWM depends on the quality of the underlying segmenter.

Inconsistent masking ratio: Because words have variable lengths in subword tokens, the actual percentage of tokens masked varies by sentence even with the greedy selection approach. A sentence composed primarily of long, morphologically complex words (common in technical writing, or in German which frequently uses compound nouns) might see 20% or more of its tokens masked when a few words are selected. A sentence of short common words might see only 10% masking. Over a full training run, these variations average out, but they can cause instability early in training when the model is still adapting to the masking distribution.

Computational overhead: Identifying word boundaries adds preprocessing cost. For every sequence in every training batch, the data collator must scan the token list, detect boundary markers, build the word groups, shuffle them, and apply the greedy selection algorithm. For production training at scale with billions of tokens, this overhead is negligible compared to the GPU compute for model forward and backward passes. But for research experiments with limited hardware, or for fast iteration during development, the additional collation cost can slow the overall training loop if the collator is not efficiently implemented.

Morphology complications in agglutinative languages: Some morphological systems do not interact cleanly with the "mask the whole word" approach. Finnish and Turkish, along with Hungarian, are agglutinative languages where a single word might consist of many morphemes stacked together, each carrying distinct semantic content. The Finnish word "talossanikin" means "also in my house" and could tokenize into many subword pieces. Masking the entire word removes substantial semantic content in a single masking event, potentially making the prediction task too difficult or providing too little residual context for the model to learn from. For these languages, span masking or character-level approaches might be more appropriate than either standard MLM or WWM.

Does not fix all information leakage: WWM eliminates within-word information leakage, but other forms of context leakage remain. If a word appears multiple times in the same sentence, masking one occurrence while leaving the other visible still provides strong evidence about the masked token. Similarly, syntactic cues (a masked position in a role that only certain word types can fill) reduce the effective prediction difficulty even without visible subword siblings. WWM is a targeted fix for one specific source of leakage; it does not create a uniformly difficult prediction task across all sentence types.

Interaction with special tokens: Special tokens like [CLS], [SEP], [PAD], and entity markers (used in some BERT variants) must be excluded from WWM grouping. Implementations that naively group all tokens can accidentally include these special tokens in word groups, producing incorrect groupings for the first or last real word in a sequence. Production implementations need explicit handling to skip special tokens before building word groups.

Despite these limitations, WWM has proven to be a durable and valuable technique with clear empirical benefits and lasting influence on how the field approaches masked language model pre-training.

Improved downstream performance: Google released official BERT-wwm models alongside the Chinese BERT paper, and subsequent benchmarking showed consistent improvements on SQuAD and other reading comprehension benchmarks. For English, the improvements are modest but reliable: typically 0.5-1.0 F1 points on SQuAD, which sounds small but represents a meaningful reduction in error rate when considered relative to state-of-the-art performance. For Chinese, the improvements are much larger because character-level masking in standard BERT creates severe information leakage at the finest possible granularity. WWM, combined with proper word segmentation for Chinese, produces substantially better models on Chinese NLP tasks.

Better morphological understanding: WWM forces models to learn word-level rather than subword-level patterns. When the model must predict entire words from context, it cannot rely on within-word co-occurrence statistics. The resulting representations tend to capture more of the word's semantic role and less of its morphological structure. Words with shared roots form more semantically coherent clusters in the embedding space because the model has learned what those roots mean in context rather than just what tokens tend to follow them. This benefit compounds over training: as the model develops richer word-level representations, its attention mechanism learns to perform more sophisticated semantic matching between words at different positions.

Standard practice for non-English models: WWM has become the default for training BERT models in morphologically rich languages. German and Arabic BERT variants, together with Turkish, Korean, plus Chinese variants, almost universally use WWM because random subword masking loses too much training signal to information leakage in these languages. The community has effectively concluded that WWM should be the default masking strategy for any new BERT pre-training effort, regardless of language, with random subword masking only as a fallback when word boundary information is unavailable.

Foundation for span masking: WWM paved the way for more sophisticated masking strategies like span corruption in T5 and SpanBERT's span masking approach. The insight that masking units should respect linguistic boundaries generalizes beyond single words to phrases and sentences. T5's sentinel token approach, where entire spans of consecutive tokens are masked and replaced with a single sentinel, is essentially WWM generalized to arbitrary spans. SpanBERT demonstrated that masking contiguous spans of 2-10 tokens (often corresponding to phrases or named entities) produces better representations for span-based tasks like named entity recognition and relation extraction. These techniques built directly on the intuition first formalized in WWM: the unit of masking should match the unit of prediction, and both should be linguistically meaningful.

Continued relevance in the era of larger models: Even as language models have scaled dramatically in size (from BERT's 110M parameters to models with hundreds of billions of parameters), the masked language modeling objective and its variants remain important for encoder pre-training. Models like RoBERTa and DeBERTa, along with ALBERT, all use masking-based objectives, and WWM or span masking variants are standard components of their pre-training procedures. The core insight of WWM, that masking granularity should match semantic granularity, remains valid regardless of model size.

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.

Whole Word Masking addresses a fundamental flaw in applying masked language modeling to subword-tokenized text. When subword tokens are masked independently, visible sibling tokens leak information about masked positions, reducing the prediction task to morphological pattern matching rather than contextual reasoning. This information leakage is not a small perturbation: it can reduce prediction entropy by a factor of five or more, dramatically weakening the gradient signal that drives learning.

WWM's core contributions are straightforward in principle and impactful in practice:

• Preserves word boundaries: All subwords of a word are masked together, eliminating within-word information leakage. The model must predict complete words from external context alone.

• Strengthens learning signal: Because the model cannot rely on visible sibling subwords, it must reason about semantic roles, syntactic structure, and pragmatic context to predict masked words. This produces richer gradient updates and stronger contextual representations.

• Adapts to tokenizer conventions: Different tokenizers mark word boundaries differently (## for WordPiece, ▁ for SentencePiece, Ġ for GPT-2 BPE), and WWM implementations should detect and handle each convention automatically.

• Integrates cleanly with existing frameworks: Libraries like Hugging Face provide ready-to-use WWM data collators that handle all implementation details, making it straightforward to add WWM to any existing MLM training pipeline.

• Has proven empirical value: Consistent improvements on reading comprehension and other benchmarks, particularly for morphologically rich languages, have made WWM the standard approach for non-English BERT pre-training.

The technique is particularly important for languages with rich morphology (German, Turkish, Finnish, Arabic) and for any domain where word-level understanding matters more than subword pattern matching (scientific text, legal documents, technical documentation). While the improvements over standard MLM are modest for English, they compound with other techniques and become essential for many non-English languages.

The next chapter explores span corruption, which extends the WWM insight further by masking contiguous spans of multiple words together, creating even more challenging prediction tasks that encourage models to learn longer-range dependencies and develop representations suited to sequence-to-sequence tasks.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about Whole Word Masking.