Skip to content

← Tokenization & Embeddings step 8 of 15

Medium End-to-End

MLM Masking Strategy

Implement BERT’s masked language model (MLM) masking strategy from Devlin et al. 2018.

Before pre-training, BERT corrupts ~15% of input tokens using an 80/10/10 split:

  • 80% of selected positions → replaced with [MASK] token id.
  • 10% of selected positions → replaced with a random vocabulary token (forces the model to be robust to noise at every position).
  • 10% of selected positions → kept as the original token (forces representations to remain useful even at unmasked positions).

Only selected positions contribute to the pre-training loss. The output includes both the corrupted token ids and a binary indicator mask so the training loop can compute loss only on the selected positions.

The 15% rate

The 15% target rate balances enough signal for the model to learn vs not corrupting too much of the input context. Too low and the model rarely gets training signal; too high and the corrupted context becomes uninformative.

Implementation

Use a single seeded torch.Generator for all three random draws in order:

gen = torch.Generator().manual_seed(seed)
select_probs  = torch.rand(N, T, generator=gen)       # which positions
corrupt_probs = torch.rand(N, T, generator=gen)       # 80/10/10 split
random_tokens = torch.randint(0, vocab_size, (N, T), generator=gen)

In JAX

JAX’s PRNG does not match PyTorch’s, so the recipe above is PyTorch’s alone and the expected outputs for the JAX cases come from this exact derivation. Follow it literally or the values will not match:

key = jax.random.PRNGKey(int(seed))
key, k1, k2, k3 = jax.random.split(key, 4)
select_probs  = jax.random.uniform(k1, shape=(N, T))
corrupt_probs = jax.random.uniform(k2, shape=(N, T))
random_tokens = jax.random.randint(k3, shape=(N, T), minval=0, maxval=int(vocab_size))
  • Position selected if select_prob < mask_prob.
  • For selected positions:
    • corrupt_prob < 0.8 → replace with mask_token_id
    • 0.8 ≤ corrupt_prob < 0.9 → replace with random_token
    • corrupt_prob ≥ 0.9 → keep original (but still mark as selected in the indicator)

Output shape

Return a tensor of shape (N, T, 2):

  • Channel 0: corrupted token ids (cast to float).
  • Channel 1: 1.0 if the position was selected for loss, 0.0 otherwise.

Reference

Devlin et al., “BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding”, NAACL 2019.