Train MLM Pretraining Step

Implement train_mlm_step: corrupt selected input tokens, predict their original IDs with a bidirectional transformer, and take one SGD step. Compute every parameter gradient by hand.

The training signal comes only from selected positions. Other positions still provide context through attention, so their embeddings can receive gradients. The output head is a separate matrix in this exercise; it is not tied to w_emb.

Inputs and return value

• tokens: (N, T), integer-valued IDs delivered as floats. Cast to integers.
• w_emb: (V, D), token embeddings, with V = vocab_size.
• pos_embed: (T, D), position embeddings shared across examples.
• blocks_weights: (L, 6, D, D), L encoder blocks. Each contains [w_q, w_k, w_v, w_o, w_mlp1, w_mlp2].
• w_head: (D, V), the vocabulary projection.
• num_heads: a positive integer dividing D.
• lr: a finite, nonnegative learning rate.
• mask_token_id: an integer in [0, V), as are all input IDs.
• seed: the integer seed for the random recipe below.

N, T, V and D are positive; L may be zero. Parameters are finite and share a floating-point dtype and device. tokens is on that device. Shapes agree. Leave inputs unchanged and preserve the parameter dtype and device.

Return a one-dimensional tensor containing the updated parameters in this order: w_emb, pos_embed, blocks_weights, w_head. Flatten each in row-major order and concatenate. The length is VD + TD + L6DD + DV. Return weights, not gradients or a loss value. With lr=0, return the original parameters in this order.

1. Choose prediction targets and corrupt the input

Use a private CPU generator and make these three draws in order. Draw whole (N, T) arrays, including entries for positions that will not be selected. Keep the uniform draws in float32 regardless of the parameter dtype.

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, V, (N, T), generator=gen, device="cpu")

Move these arrays to the input device before combining them with tokens. Do not change the global random-generator state. Select positions where select_probs is strictly less than 0.15. Start from a copy of the integer input IDs. For each selected position:

• corrupt_probs < 0.8: replace its ID with mask_token_id.
• 0.8 <= corrupt_probs < 0.9: replace it with the corresponding random_tokens ID.
• corrupt_probs >= 0.9: keep its original ID, but still include it in the loss.

A random replacement may equal the original ID or the mask ID. Unselected positions keep their IDs. If nothing was selected across the entire batch, select position (0, 0) and replace it with mask_token_id. This exercise uses that fallback so the loss always has at least one target. Do not resample.

The target for every selected position is its original ID, before corruption. No IDs have special padding behavior, including zero and mask_token_id.

2. Run the encoder and output head

Start with x = w_emb[corrupted_ids] + pos_embed, broadcasting positions over the batch. Apply the L blocks in order, with this pre-normalization architecture:

h = LN(x)
q, k, v = h @ w_q, h @ w_k, h @ w_v
y = x + multi_head_attention(q, k, v) @ w_o
x = y + GELU(LN(y) @ w_mlp1) @ w_mlp2

LN uses each token's feature mean and population variance: (token - mean) / sqrt(variance + 1e-5). There is no learned scale or bias. Split projected features into consecutive heads of width D/num_heads. Divide q @ k.T by sqrt(D/num_heads), apply stable softmax over key positions and concatenate heads in their original order. All positions in one example may attend to each other; different examples never attend to each other.

Use the elementwise approximation: GELU(u) = 0.5u(1 + tanh(sqrt(2/pi)(u + 0.044715u**3))). There are no biases, dropout, causal masks or final normalization. With L=0, use token-plus-position embeddings directly. Compute logits = x @ w_head.

3. Differentiate the selected-position loss

Let M be the number of selected positions after the fallback. Compute mean cross-entropy against their original IDs, averaging over M across the whole batch. Use stable log-softmax if you calculate the loss. The returned value does not include it.

At each selected row, the logit gradient is (softmax(logits) - one_hot(target)) / M. Other logit rows have zero gradient. There is no further division by N or T.

Backpropagate through the head and then through blocks in reverse order. Each residual addition sends the incoming gradient through both branches. For a shared matrix, flatten batch and position axes before computing its weight gradient; for example, dw_head = x_flat.T @ dlogits_flat.

Sum embedding gradients at every occurrence of each corrupted input ID. Repeated IDs accumulate into the same row. Sum position-embedding gradients across examples. The corruption choices and integer IDs are held fixed during this calculation. An unselected token can receive an embedding gradient through its contribution to a selected token's attention context.

4. Update all parameters once

Compute every gradient using the original weights, then apply new_parameter = parameter - lr * gradient. Do not update one weight early and use it to calculate another gradient. There is no weight decay, clipping, momentum or optimizer state.

Build normalization, GELU, attention, softmax and their derivatives from basic tensor operations. Do not use autograd, backward(), packaged transformer layers or optimizer helpers. Cache the forward intermediates needed for the manual backward pass.