MLM Evaluation: Masked Accuracy

Measure how often a small masked language model predicts the original token at selected positions. Run the supplied model once, choose its highest-scoring vocabulary item at each selected position, and report the fraction correct. There is no token corruption, loss calculation or parameter update to perform.

The model uses the same simplified pre-normalized architecture as MLM Forward Pass. Its complete computation is specified below; no helper from another exercise is available automatically. This is not a full BERT implementation.

Inputs

Implement mlm_eval_accuracy(input_ids, original_ids, mask_indicator, w_emb, pos_embed, blocks_weights, w_head, num_heads).

Input	Shape	Meaning
`input_ids`	`(N, T)`	Supplied, possibly corrupted token IDs.
`original_ids`	`(N, T)`	Target IDs before corruption.
`mask_indicator`	`(N, T)`	Score positions whose value is strictly greater than 0.5.
`w_emb`	`(V, D)`	Token embedding table.
`pos_embed`	`(T, D)`	Position embeddings, shared across the batch.
`blocks_weights`	`(L, 6, D, D)`	Weights of L transformer blocks.
`w_head`	`(D, V)`	Vocabulary projection.

IDs are integer-valued floats in [0, V); cast them to integers for lookup and comparison. N, T, V, D and num_heads are positive, D is divisible by num_heads, and L may be zero. Leave the inputs unchanged.

Forward pass

Start with x = w_emb[input_ids] + pos_embed. In each block, unpack [w_q, w_k, w_v, w_o, w_mlp1, w_mlp2] in that order and perform these steps:

Normalize x across its D features at each token. Use population variance (divide by D) and (x - mean) / sqrt(var + 1e-5).
Project the normalized values with w_q, w_k and w_v. Split each projection into consecutive feature chunks of width d_head = D / num_heads.
For each example and head, compute Q @ K.T / sqrt(d_head). Apply softmax over key positions, subtracting the row maximum before exponentiation.
Multiply those attention probabilities by V. Concatenate heads in their original order, multiply by w_o, and add the result to x.
Normalize this updated x by the same rule. Add gelu(norm_x @ w_mlp1) @ w_mlp2 to x.

Use gelu(u) = 0.5 * u * (1 + tanh(sqrt(2/pi) * (u + 0.044715 * u**3))). After all blocks, compute logits = x @ w_head. With L=0, this is just the embedding-plus-position tensor projected to vocabulary scores.

There are no biases, learned normalization parameters, dropout or final layer normalization. Attention is bidirectional within each example, with no causal or padding mask. mask_indicator selects positions to score, not positions to hide from attention. Selected inputs need not contain a special mask ID.

Accuracy

At each position where mask_indicator > 0.5, choose the vocabulary index with the largest logit. Break exact ties by choosing the lowest index. Compare that index with original_ids, not the corrupted input ID.

Return a Python float: total correct predictions divided by total selected positions across the whole batch. Do not average per-example accuracies or include unselected positions. Do not apply vocabulary softmax or sample tokens.

For four selected positions with predictions [2, 0, 1, 3] and targets [2, 1, 1, 0], return 0.5. If no positions are selected, return 0.0 as this exercise's convention; it does not mean the model was evaluated and got every prediction wrong. Return the float itself, not a dictionary wrapper.

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…