Bigram Language Model: Make the First Training Step Explainable

The first language model predicts the next token using only the current token. Its simplicity makes it a useful reference point: if its loss or gradients are wrong, a deeper network will only hide the cause.

You will implement a table of logits, stable cross-entropy, one SGD step, and autoregressive sampling. By the end you should be able to compute one update by hand and explain what extra context this model can never use.

The model is a transition table

For vocabulary size V, BigramLM owns a (V,V) table. Row t contains unnormalized scores for the token following t. Input [0,1,0] returns three logit rows of shape (3,V), with the first and third rows referencing the same parameters.

Rows are not probabilities until softmax is applied. Adding the same constant to every logit in a row changes no probabilities:

p[k] = exp(z[k] - m) / sum_j exp(z[j] - m)
m = max_j z[j]

Subtracting the maximum prevents exponentiating a large positive number. Keep the subtraction, exponential, sum, and division as Value operations. Reading .data to choose the shift is safe here because a uniform shift cancels from the softmax expression.

Derive a loss that remains meaningful

For target y, cross-entropy is -log(p[y]). Algebraically:

loss(z,y) = log(sum_j exp(z[j]-m)) - (z[y]-m)

Implement this expression directly. Do not first compute a tiny target probability and then clamp it before taking log. With logits [1000,-1000] and target 1, the correct loss is approximately 2000. The target probability may underflow to zero, but the shifted log-sum expression remains finite and its gradient is approximately [1,-1].

For T positions, average the T losses. Inputs must have matching nonzero lengths: a silent zip truncation would train on a different objective. This exercise receives raw logits, matching the input convention described in CrossEntropyLoss.

Work one SGD update

Take two tokens, a=0 and b=1. Start row a at [0,0] and train on one transition a -> b:

probabilities = [0.5, 0.5]
loss          = -log(0.5) = 0.693147...
grad logits   = [0.5, -0.5]

With learning rate 0.1, SGD subtracts lr*grad:

new row a = [-0.05, 0.05]
P(b | a)  = exp(0.05)/(exp(-0.05)+exp(0.05)) ≈ 0.524979
new loss  ≈ 0.644397

Row b is unchanged because it was never used as an input. The target ID selects a loss term; it does not directly select a parameter row for updating. When the same input token occurs repeatedly, its row receives all contributions, each scaled by 1/T for the mean loss.

The derivative p[k] - 1[k=y] is also an independent oracle for testing the autograd implementation. Check it on moderate and extreme logits; merely checking that probabilities sum to one misses a broken loss.

Keep the training step ordered

A step must build logits and mean loss on a fresh graph, zero all parameter gradients, call loss.backward(), and update each unique parameter once. Return the numeric loss from before the update. Read .data at the logging and optimizer boundaries, not inside the differentiable objective.

Use a fixed tiny transition sequence before random batches. Uniform logits give mean loss log(V) regardless of the target IDs. A persistently unchanged loss suggests disconnected or undiscovered parameters. A plausible loss with enormous gradients suggests sum-vs-mean confusion or duplicate updates.

Sample a sequence

Generation begins with a nonempty seed. Copy the seed, take the last token's row, divide logits by a finite positive temperature, softmax, sample an ID, append it, and repeat. Return the seed plus exactly max_new_tokens samples. The seed argument must remain unchanged.

Temperature 1 keeps the learned distribution. A lower positive temperature sharpens it; a higher one flattens it. Temperature 0 is not supported as an alias for greedy decoding. Set the random seed when debugging stochastic behavior, but do not reset it inside the generation loop.

For sampling, you may operate on numeric logits: subtract their maximum before dividing by temperature, then exponentiate. These unnormalized weights are sufficient for random.choices and avoid an overflowing reciprocal at tiny positive temperatures. No gradient is needed through the sampling decision.

Training uses the observed preceding token; generation feeds the model its own previous samples. The loss is an average over provided targets, while generation uses only the last prediction each step.

Identify the model's ceiling

Suppose the corpus has both "abx" and "cby". After b, a bigram model returns the same distribution regardless of whether a or c preceded it. More training cannot recover information absent from its input. A fixed-window MLP will address this limitation next.

Checkpoint. If the same training sequence is duplicated in a single mean-loss batch, should the gradient double?

Check your reasoning

no. Twice as many identical terms are divided by twice as many positions. With a summed loss it would double.

Use this as a reduction test, then verify that the one-step numeric example changes only the intended table row.