MLP Language Model (Bengio 2003)

The problem

A bigram table cannot look further back than one token, and a bigger table grows out of control. Bengio et al. (2003) replaced the table with a small neural network, an MLP (multilayer perceptron), that reads the last $n$ tokens and outputs one score (logit) for every token of the vocabulary. You will write its forward pass, its loss, and the function that cuts a text into training examples.

The forward pass

For a batch of $B$ contexts, ids of shape $(B, n)$ (each row holds $n$ token ids, oldest first):

1. Embed. Look up each token in the embedding table $C$ of shape $(V, D)$. Result: shape $(B, n, D)$, one vector of $D$ numbers per token.
2. Concatenate. Put the $n$ vectors of each context side by side, oldest first. Result: $e$, shape $(B, nD)$.
3. Hidden layer. $h = \tanh(e W_1 + b_1)$, with $W_1$ of shape $(nD, H)$: $H$ numbers per context.
4. Logits. $z = h W_2 + b_2$, with $W_2$ of shape $(H, V)$: one score per token of the vocabulary.

where $W_1$, $b_1$, $W_2$, $b_2$ are the weights and biases of the two layers.

Example. $V = 3$ tokens, $D = 2$, $n = 2$, $H = 2$, and a batch of $B = 2$ contexts, ids = [[0, 2], [1, 0]].

1. With $C$ = [[1, 0], [0, 1], [1, 1]], context 0 looks up rows 0 and 2: [[1, 0], [1, 1]]. Context 1 looks up rows 1 and 0: [[0, 1], [1, 0]].
2. Concatenated: $e$ = [[1, 0, 1, 1], [0, 1, 1, 0]], shape $(2, 4)$. The first two numbers of each row belong to the oldest token.
3. With $W_1 = 0.5 \cdot$ [[1, 0], [0, 1], [1, 0], [0, 1]] and $b_1 = 0$: $eW_1$ = [[1, 0.5], [0.5, 0.5]], and $h = \tanh$ of it $\approx$ [[0.762, 0.462], [0.462, 0.462]].
4. With $W_2$ = [[1, 0, -1], [0, 1, 1]] and $b_2 = 0$: $z \approx$ [[0.762, 0.462, -0.299], [0.462, 0.462, 0]], shape $(2, 3)$.

The loss

The loss is the mean cross-entropy of the logits against targets, the right next token of each context (section 2): for each row, $\mathrm{LSE}(z_b) - z_{b, y_b}$, averaged over the batch.

Example. With targets = [2, 1]: row 0 has loss about $1.797$ (the model gave token 2 the lowest score), row 1 about $0.967$, and the mean is about $1.382$.

Training examples from a text

To train, a long text is cut into windows: every run of $n$ consecutive tokens is a context, and the token right after it is its target.

Example. build_contexts([5, 6, 7, 8, 9], n=2) returns ids = [[5, 6], [6, 7], [7, 8]] and targets = [7, 8, 9]. A text of $T$ tokens gives $T - n$ examples.

What to implement

• mlp_lm_logits(ids, params): params is a dict with keys "C", "W1", "b1", "W2", "b2". Return the $(B, V)$ logits.
• mlp_lm_loss(ids, targets, params): the mean cross-entropy of the logits against targets (shape $(B,)$), computed stably.
• build_contexts(tokens, n): return (ids, targets), where row $t$ of ids is tokens[t : t + n] and targets[t] is tokens[t + n], for every $t$ that fits.

Requirements

• Concatenation order matters: the first $D$ columns of $e$ are the oldest token.
• No Python loops in mlp_lm_logits and mlp_lm_loss. build_contexts may use a loop, or sliding_window_view from np.lib.stride_tricks.

Hints

• Step 2 is a single reshape of the $(B, n, D)$ lookup to $(B, nD)$: reshape reads the last axis fastest, so each token's $D$ numbers stay together, in order.
• You already wrote every piece: the embedding lookup (section 4), a linear layer (section 1) and cross-entropy from logits (section 2).