MLM Masking Strategy

Prepare inputs and a loss-selection indicator for masked language modeling. A selected token is a prediction target even if its input value stays unchanged. This indicator is not an attention mask: it tells the loss which positions to use. Keep the original tokens separately as the targets.

Your task

Implement apply_bert_masking(tokens, vocab_size, mask_token_id, seed, mask_prob=0.15). tokens has shape (N, T) with positive dimensions and integer-valued token IDs in [0, vocab_size). Cast IDs to integers before processing them. vocab_size is positive, mask_token_id is a valid ID, seed is an integer and 0 <= mask_prob <= 1.

Every position is eligible in this simplified exercise. There is no special-token or padding exclusion. Preserve the input tokens.

Select each position when its selection draw is strictly less than mask_prob. For selected positions, use a separate corruption draw:

Corruption draw	Input token after corruption
Less than `0.8`	`mask_token_id`
At least `0.8`, less than `0.9`	The supplied random vocabulary draw
At least `0.9`	Original token

Unselected positions keep their original tokens. These are probabilities, not exact quotas: a small batch need not have exactly 15% selected positions or an exact 80/10/10 split. A random replacement may equal the original token or the mask ID; do not resample it.

Reproducible draws

Draw full (N, T) arrays, including positions that will not be selected. Do not draw replacements only for selected positions or reseed between draws. The two frameworks use different random streams and have separate expected outputs.

For PyTorch, use one local CPU generator and draw in this order:

gen = torch.Generator(device="cpu").manual_seed(int(seed))
select_probs = torch.rand(N, T, generator=gen, dtype=torch.float32, device="cpu")
corrupt_probs = torch.rand(N, T, generator=gen, dtype=torch.float32, device="cpu")
random_tokens = torch.randint(0, int(vocab_size), (N, T), generator=gen, device="cpu")

Move the completed draw arrays to the tokens' device if needed. Generating on a different device can change the random stream.

In JAX

For JAX, use this split order with the runtime's jax_threefry_partitionable=True:

key = jax.random.PRNGKey(int(seed))
_, k1, k2, k3 = jax.random.split(key, 4)
select_probs = jax.random.uniform(k1, (N, T), dtype=jnp.float32)
corrupt_probs = jax.random.uniform(k2, (N, T), dtype=jnp.float32)
random_tokens = jax.random.randint(k3, (N, T), 0, int(vocab_size), dtype=jnp.int32)

Return format

Return a float32 tensor of shape (N, T, 2). At each position, channel 0 is the corrupted token ID and channel 1 is 1.0 if selected, otherwise 0.0. Test IDs fit exactly in float32. Do not infer selection by checking whether a token changed: a selected, unchanged token still has indicator 1.0.

For example, an original token 7 that is selected and kept becomes [7.0, 1.0]; the same token when unselected becomes [7.0, 0.0]. With mask_prob=0, all indicators are zero. With mask_prob=1, all are one, but not every token must become the mask ID.

The selection/replacement idea comes from BERT. This exercise specifies its own exact random draws and simplified eligibility rules.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Match each prediction to the token it should predict

A causal language model predicts the next token from the tokens before it. For sequence [A, B, C], logits produced after reading A are trained against B, and logits after reading B are trained against C. This is why input and target sequences are usually shifted by one position. A causal attention mask prevents access to later tokens that would reveal the answer.

A masked language model instead predicts selected tokens from a corrupted sequence and can use context on both sides. The supervised positions are not necessarily just those showing a mask token. A corruption rule may leave some selected tokens unchanged or replace them randomly while still training on them. Compute the loss over the specified supervised positions only.

Teacher forcing supplies the correct previous tokens

In encoder-decoder training, teacher forcing feeds the decoder the correct previous target tokens. During generation, it usually receives its own previous choices instead. A mistake can therefore change later inputs, which helps explain why training loss alone does not describe generation quality.

Convert mean token loss to perplexity

For M valid target tokens, mean negative log likelihood is

$$L=-\frac1M\sum_t\log p(y_t).$$

With natural logarithms, perplexity is exp(L). A model that gives the correct token probability 1/4 at every step has perplexity 4. A uniform distribution over a vocabulary of size V has perplexity V.

When batches contain different numbers of valid tokens, add their token losses and valid-token counts before dividing. Averaging batch means would give small batches too much weight. Exclude padding and other unsupervised positions as specified, and define what happens when no tokens are supervised.

Perplexity also depends on tokenization. Scores based on different token units are not directly comparable just because both are called perplexity.

Detecting runtime…