MLM Forward with Tied Output Head

Implement mlm_forward_tied: run a bidirectional transformer encoder and return vocabulary logits at selected positions. Use the same embedding table for both input lookup and the output projection.

This is weight tying: a hidden vector of width D is scored against every row of w_emb by multiplying it by w_emb.T. There is no separate vocabulary projection matrix. Compared with adding an independent (D, V) head, this shares V*D parameters between the two uses.

Inputs and output

input_ids: (N, T), integer-valued IDs in [0, V), delivered as floats.
mask_indicator: (N, T), finite values selecting rows where the value is greater than 0.5.
w_emb: (V, D), the shared token embedding table.
pos_embed: (T, D), position embeddings shared across the batch.
blocks_weights: (L, 6, D, D), containing L pre-normalization encoder blocks.
num_heads: a positive integer dividing D.

N, T, V and D are positive; L may be zero. Shapes are consistent. Model parameters are finite and share a floating-point dtype and device. IDs and selection indicators are on the same device. Leave all inputs unchanged.

Return raw logits of shape (M, V), where M is the count of indicator values strictly greater than 0.5. Preserve the parameter dtype and device. Visit examples in batch order and positions from left to right within each example. If no positions are selected, return shape (0, V).

Selection is separate from attention

Cast input_ids to integers and use them exactly as supplied. Any replacement with a mask token happens before this function. mask_indicator only selects which output rows to return; it does not change embeddings or block attention.

Every position can attend to all positions in its own sequence. Do not add a causal mask or treat token ID zero as padding. Keep different batch examples separate.

Embeddings and encoder blocks

Start with x = w_emb[input_ids] + pos_embed, broadcasting positions over the batch to obtain shape (N, T, D).

Apply blocks from 0 through L-1. For block i, the six matrices are [w_q, w_k, w_v, w_o, w_mlp1, w_mlp2], each of shape (D, D). The feed-forward width is D in this exercise. Each block performs:

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 normalizes each token over its D features using population variance: (token - mean) / sqrt(variance + 1e-5), with no learned scale or bias. For attention, split projected features into consecutive heads of width D/num_heads. Divide scores by sqrt(D/num_heads), apply stable softmax over key positions, and concatenate head outputs in their original order.

Use the tanh GELU approximation elementwise:

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

Build normalization, attention and GELU from basic tensor operations in either framework, without packaged attention or layer-normalization helpers. A stable softmax helper is allowed. There are no linear-layer biases, dropout or an additional final normalization.

Tied projection and row selection

After the last block, compute logits_all = x @ w_emb.T, then return the rows where mask_indicator > 0.5. Do not apply a vocabulary softmax or choose a predicted token. With L=0, project the token-plus-position embeddings directly.

For mask_indicator=[[0,1,1],[1,0,0]], return rows for (example, position) (0,1), (0,2), then (1,0). An indicator equal to 0.5 is not selected. The same w_emb supplies both the input vectors and every vocabulary output score.

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

Look up a learned vector for each token

An embedding table E has shape (vocabulary_size, embedding_dim). Integer token ID i selects row E_i. IDs are labels for rows, not measurements: token 100 does not have twice the meaning of token 50.

A lookup is equivalent to multiplying a one-hot token vector by E, but avoids constructing the mostly zero vector. Looking up IDs [2, 2, 5] returns rows [E_2, E_2, E_5]. Both uses of row 2 contribute gradients to that row during training, subject to any scaling or special padding behavior in the task.

Averaging token embeddings gives one fixed-width vector for a sequence. It loses order information: the average for tokens [2, 5] is the same as for [5, 2].

Reuse the embedding weights to predict tokens

Weight tying uses the same table for input embeddings and output scores. For a hidden vector h with width embedding_dim, the product hE^T gives one score, or logit, for every vocabulary entry.

An input lookup selects one row. The output projection compares the hidden vector with every row. Sharing the table saves a separate output weight matrix, but means both uses must contribute gradients to the same parameter. Two copies that are updated independently are no longer tied.

Check that token IDs are in range and have the required integer dtype. When pooling padded sequences, exclude padding from both the sum and the count. Even a zero padding vector lowers an ordinary average if it is included in the denominator. Follow the specified behavior for a sequence with no valid tokens.

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…