Hard End-to-End

One RNN Language-Model Training Step

Train a small RNN to predict the next token in a sequence. Implement one forward pass, manual backpropagation through time (BPTT), and one simultaneous SGD update. Derive the gradients yourself; do not use autograd or a built-in RNN module in your solution.

Inputs and return value

Implement train_rnn_lm_step(tokens, w_emb, w_h, w_out, h0, lr).

tokens has shape (T,), with T at least 2. Entries are integer-valued token IDs in [0, V), supplied as numbers; cast them to integers for indexing.
w_emb has shape (V, d_emb).
w_h has shape (d_emb + d_h, d_h).
w_out has shape (d_h, V).
h0 has shape (d_h,).
lr is a nonnegative learning rate.

All dimensions are positive. The weights and initial state are floating-point arrays. This is one sequence, with no batch axis, biases or dropout. Work on local parameter copies and leave the input arrays unchanged.

Compute updates for w_emb, w_h and w_out, treating h0 as fixed. Return only the updated w_out flattened in row-major order, with shape (d_h * V,). This output format checks the output-projection update; it does not independently verify the embedding or recurrent-weight updates. Implement the full backward pass described below as part of the exercise.

Align inputs and targets

There are n_steps = T - 1 predictions. Input tokens are tokens[:-1] and targets are tokens[1:]. For [2, 0, 2, 1], train on the pairs 2 -> 0, 0 -> 2, and 2 -> 1. The last token is a target, not another input step.

Starting from h0, compute and save each step's input embedding, previous hidden state, updated hidden state and probabilities:

embedding = w_emb[tokens[t]]
joined = concatenate([embedding, previous_h])
h = tanh(joined @ w_h)
logits = h @ w_out
probs = softmax(logits)

Use stable softmax over the V vocabulary entries. The objective is mean cross-entropy over all n_steps, not a sum. You do not return the loss; if you calculate it for debugging, use stable log-softmax rather than taking the log of a probability that may have underflowed.

Backpropagate from last step to first

Initialize all parameter gradients and dh_next to zero. At each step t, use saved forward values and the original, pre-update weights:

dlogits = (probs[t] - one_hot(tokens[t+1])) / n_steps
dw_out += outer(h[t], dlogits)
dh = w_out @ dlogits + dh_next
dpre = dh * (1 - h[t]**2)
dw_h += outer(joined[t], dpre)
djoined = w_h @ dpre
dw_emb[tokens[t]] += djoined[:d_emb]
dh_next = djoined[d_emb:]

Here h[t] means the updated hidden state after input t. For t=0, joined[t] contains h0; later steps contain the preceding updated state. dh_next is the contribution arriving through the recurrence from later predictions. Add it to the current output contribution before differentiating tanh.

Accumulate gradients rather than overwriting them. Repeated input tokens contribute multiple times to the same embedding row. The division by n_steps already accounts for the mean loss; do not divide the accumulated gradients again. The recurrence uses matrix products as well as tanh derivatives, so its gradients can shrink or grow rather than always vanishing.

Apply one update

Only after the full backward pass, compute new_parameter = parameter - lr * gradient for all three weight arrays. Updating parameters inside the time loop would change the gradients being computed. Return the flattened new w_out.

With lr=0, the result must equal the original flattened w_out. Allocate helper arrays with the appropriate input dtype and device in PyTorch.

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

Update a state as you read a sequence

A recurrent neural network carries a hidden state from one time step to the next. A basic cell computes

$$h_t=\tanh(x_tW_x+h_{t-1}W_h+b).$$

x_t is the current input and h_(t-1) is the previous state. The weights and bias are shared across time steps. The hidden state changes with the input; it is not a separate learned parameter at each step.

Even when the input and previous state are zero, a nonzero bias can produce a nonzero new state. For sequence classification, a classifier can read the final state. For token labeling or language modeling, it may need an output at every step instead.

Gates control how information is carried forward

An LSTM maintains a cell state as well as a hidden state. Input, forget, and output gates control how information enters, remains in, and is read from that cell state. A GRU uses update and reset gates with a candidate hidden state.

Follow the given equations for gate order and placement. In particular, applying a reset gate before a matrix multiplication is generally different from applying it afterward.

A bidirectional RNN reads the sequence in both directions. Align the two states with their original token positions before concatenating them. The reverse recurrence can use later tokens, so its representation cannot be used unchanged for a causal next-token predictor.

Add gradient contributions across time

Backpropagation through time treats the recurrence as repeated uses of the same cell. Each use of a weight contributes to that weight's total gradient. A hidden state can receive gradient from its current output and from later states.

Repeated multiplication by derivatives can shrink or amplify gradients. Gradient clipping bounds the gradient passed to the optimizer; the resulting parameter step also depends on the optimizer and its state. Gates provide other paths for information and gradients, but do not replace correct handling of padding, initial states, and sequence lengths.

Detecting runtime…