Skip to content

← Tokenization & Embeddings step 12 of 15

Medium End-to-End

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:

  1. Normalize x across its D features at each token. Use population variance (divide by D) and (x - mean) / sqrt(var + 1e-5).
  2. 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.
  3. For each example and head, compute Q @ K.T / sqrt(d_head). Apply softmax over key positions, subtracting the row maximum before exponentiation.
  4. Multiply those attention probabilities by V. Concatenate heads in their original order, multiply by w_o, and add the result to x.
  5. 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.