MLM Forward Pass

Compute vocabulary scores for selected positions using a small bidirectional transformer. Inputs may already contain corrupted tokens. This function does not choose or replace tokens and does not compute a loss or training update.

The model is simplified: pre-normalized blocks, a feed-forward width equal to the model width, and a linear vocabulary head. It is not a complete BERT implementation.

Inputs

Implement mlm_forward(input_ids, mask_indicator, w_emb, pos_embed, blocks_weights, w_head, num_heads):

Input	Shape	Meaning
`input_ids`	`(N, T)`	Integer-valued IDs in `[0, V)`; cast to integers for lookup.
`mask_indicator`	`(N, T)`	Positions selected for output, using `> 0.5`.
`w_emb`	`(V, D)`	Token embeddings.
`pos_embed`	`(T, D)`	Supplied position embeddings, shared across the batch.
`blocks_weights`	`(L, 6, D, D)`	Weights for `L` blocks in execution order.
`w_head`	`(D, V)`	Vocabulary projection.

N, T, V, D and num_heads are positive; D is divisible by num_heads. L may be zero. Preserve every input.

For each block, the six matrices are packed in this order: [w_q, w_k, w_v, w_o, w_mlp1, w_mlp2]. Use them as supplied, with row-vector matrix multiplication. There are no biases, learned normalization parameters, dropout, segment embeddings or final layer normalization.

Forward pass

Start with x = w_emb[input_ids] + pos_embed, broadcasting positions across examples. For each block:

Normalize each token across its D features. Compute the population variance, dividing by D, then use (x - mean) / sqrt(var + 1e-5).
Project this normalized tensor with w_q, w_k and w_v. Split the feature axis into num_heads consecutive chunks of width d_head = D / num_heads.
Within each example and head, compute scores = Q @ K.T / sqrt(d_head). Apply softmax over key positions: subtract each score row's maximum, exponentiate, and divide by the sum of exponentials in that row.
Multiply attention probabilities by V, concatenate the heads in their original order, project with w_o, then add the original x as a residual.
Normalize this updated x again using the same rule. Compute gelu(norm_x @ w_mlp1) @ w_mlp2 and add it to x.

Use the tanh approximation to GELU:

gelu(u) = 0.5 * u * (1 + tanh(sqrt(2/pi) * (u + 0.044715 * u**3)))

Implement attention, softmax, layer normalization and GELU from these formulas, without high-level attention, normalization or activation helpers.

Selected positions are not an attention mask

Every position can attend to every other position in the same example. There is no causal or padding mask, and examples do not attend across the batch. mask_indicator affects only which rows are returned, not which positions supply context. A selected token need not have a special mask ID.

The diagram below compares attention patterns; use the bidirectional pattern for this exercise.

Loading visualization…

Output

After all blocks, compute raw logits x @ w_head. Do not apply a vocabulary softmax. Return only rows where mask_indicator > 0.5, ordered by batch index first, then position within that example.

The shape is (M, V), where M is the count of selected positions. For mask_indicator = [[0, 1, 0], [1, 0, 1]], return rows for positions (0, 1), (1, 0) and (1, 2) in that order. With no selected positions, return shape (0, V). With zero blocks, apply just embedding plus position, the vocabulary projection and selection.

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…