MLP Language Model: Turn a Context Window into One Prediction

A bigram model cannot distinguish ab from cb when predicting after b. This model reads a fixed window of T tokens. It embeds them, concatenates their features in order, and applies a nonlinear classifier. The training and sampling objective stays next-token prediction.

You will implement MLPLanguageModel, a single-example cross-entropy, a batch training step, and rolling-window generation. The most important change from the preceding module is the target shape: one context produces one logit vector and one target token.

Build the examples explicitly

For IDs [0,1,2,0,2] and context length T=2, the complete examples are:

[0,1] -> 2
[1,2] -> 0
[2,0] -> 2

There are N-T examples for a stream of length N. The batch input xs has shape (B,T); targets ys have shape (B,). Do not pass the sequence of per-position targets from the tokenization exercise directly to this API. You can derive a single MLP target from a sampled next-token window by taking its last target ID.

forward requires exactly T IDs and returns V logits. Batch iteration belongs in train_step, which averages the B single-example losses. Require nonempty matching input and target lists so zip cannot silently discard training examples.

Follow every dimension

With embedding width D and hidden width M:

Stage	Shape
Context IDs	(T,)
Selected token rows	(T,D)
Concatenation in position order	(T*D,)
First affine layer	(M,)
ReLU	(M,)
Second affine layer	(V,)

Flatten with position as the outer loop and feature as the inner loop. The first D entries describe the oldest context token; the last D entries describe the most recent. Keeping the order fixed matters because different columns of the first weight matrix belong to different positions.

Work a complete forward pass

Use V=3, T=2, D=1, M=2, and embeddings E[0]=[1], E[1]=[2], E[2]=[-1]. For context [0,1], concatenation gives [1,2].

Let the first affine layer have rows [1,1] and [1,-1], with zero bias:

preactivation = [1+2, 1-2] = [3,-1]
hidden        = ReLU([3,-1]) = [3,0]

Let the output layer have rows [1,0], [0,1], and [-1,1], with zero bias:

logits = [3,0,-3]

For target 1, stable cross-entropy is log(exp(0)+exp(-3)+exp(-6)) - (-3), approximately 3.050946. The output favors token 0 even though token 1 is correct, so the loss is high.

Reverse the context to [1,0]: its concatenation is [2,1], hidden becomes [3,1], and logits become [3,1,-2]. Order now affects the prediction without a separate positional table. The first layer's position-specific columns provide that distinction.

Parameter sharing and capacity

The token table is shared across positions. Position-specific behavior comes from the classifier after concatenation. Parameter count is:

V*D + (T*D)*M + M + M*V + V

For the numeric example this is 3+4+2+6+3=18. Increasing T widens the first layer, so an already constructed model cannot accept a different context length. A bigram table has V*V parameters; an MLP can share lower-dimensional token features, but it is not automatically smaller for every choice of widths.

ReLU makes the classifier nonlinear. A hidden unit with negative preactivation has zero local derivative for that example; zero gradients there are expected. All-zero gradients throughout the model are not.

Reuse the objective without weakening it

Compute each loss with shifted log-sum-exp, as in the bigram lesson. Keep the loss as a Value until after backward. Sum the B losses, divide by B, clear parameter gradients, backpropagate once, and apply one SGD update per unique parameter.

Duplicating every example in a batch should preserve the mean gradient. Check one weight with a finite difference on a fixed batch, choosing values away from a ReLU kink. Rebuild the forward graph for each perturbation.

Generation needs a full initial window

This exercise has no padding token or shorter-context branch. Require at least T seed IDs. At each step, use tokens[-T:], sample from the one output vector, and append to a copied token list. Longer seeds are preserved in the returned sequence, but only their most recent T tokens affect the prediction. A positive finite temperature controls sampling.

Checkpoint 1. Does a seed of length 1 work when T=3?

Check your reasoning

no; reject it instead of silently padding with an ordinary vocabulary ID.

Checkpoint 2. What information does the model lose after a token falls outside the window?

Check your reasoning

that token no longer participates in the forward pass. It cannot affect the current prediction through a hidden recurrent state because this model has none.

Attention will next keep a vector at every position and let each position form its own weighted mixture of visible context, instead of collapsing the entire window into one classifier input.