Training loop for next-token prediction (reuses Part 1 loop)

The problem, from first principles

[03-dl-training/03-training-loop/02-assemble-training-loop] established the pattern every training loop in this curriculum follows: forward pass, compute loss and its gradient, backward pass, update the weights, repeat, demonstrated there for one linear layer trained with MSE. This question applies the EXACT SAME pattern to [04-full-forward-pass]'s language model, with one deliberate, explicitly-scoped simplification: it trains [01-output-projection]'s output head via REAL backpropagation (reusing [03-dl-training/02-layers/02-linear-backward]'s linear_backward directly, since the output projection IS an ordinary linear layer), while treating the Transformer blocks' output (hidden_states) as FIXED, given input features, exactly the same "linear probe on frozen features" setup widely used in practice when fine-tuning only a model's final layer.

This scoping is a genuine, honest choice, not a shortcut: full end-to-end backpropagation through [01-transformer-block]'s attention and feed-forward sublayers is real, well-defined math, but implementing it requires a backward pass for every piece in that stack, work this curriculum's Transformer-block questions deliberately left as FORWARD-only (mirroring how they were actually built and tested throughout this Part). What this question DOES demonstrate, completely and correctly, is the entire remaining training-loop machinery, loss computation, gradient computation, and a genuine weight update that provably reduces the loss, applied to language modeling specifically.

From theory to code

Implement train_output_head_one_step(hidden_states, token_ids, output_weight, lr): [01-output-projection]'s forward pass, [03-next-token-cross-entropy]'s loss (via the lower-level cross_entropy_forward/cross_entropy_backward), [02-layers/02-linear-backward]'s backward pass, and an SGD update, and train_output_head(hidden_states, token_ids, output_weight, lr, num_steps), looping that single step num_steps times.

Constraints

• hidden_states and token_ids are shifted by one position first, exactly [03-next-token-cross-entropy]'s pattern: the LAST position's hidden state has no next-token target and must be excluded before computing the loss/gradient.
• The gradient with respect to output_weight comes from [02-layers/02-linear-backward]'s linear_backward, called with the (shifted, flattened) hidden states as its x argument, exactly like [02-layers/02-linear-backward]'s own linear-layer backward pass.
• The weight update is output_weight - lr * grad_output_weight, ordinary SGD, exactly [03-dl-training/03-training-loop/02-assemble-training-loop]'s update rule.
• train_output_head_one_step returns the loss computed BEFORE the update (what the weight actually achieved going into this step), not after.

Hints

predicted_hidden = hidden_states[..., :-1, :]
targets = token_ids[..., 1:]
flat_hidden = predicted_hidden.reshape(-1, d_model)
flat_targets = targets.reshape(-1)

logits = output_projection(flat_hidden, output_weight)
loss = cross_entropy_forward(logits, flat_targets)
grad_logits = cross_entropy_backward(logits, flat_targets)