Medium End-to-End

LSTM Sequence Forward Pass

An LSTM carries two states through a sequence: a cell state c and a hidden state h. The cell state combines retained memory with a new candidate. The hidden state exposes a gated, tanh-transformed version of the updated cell state and feeds into the next step's gate calculation.

Your task

Implement lstm_forward(x, W_x, W_h, b, h0, c0) for a complete sequence:

x has shape (seq_len, input_dim).
W_x has shape (input_dim, 4 * hidden_dim).
W_h has shape (hidden_dim, 4 * hidden_dim).
b has shape (4 * hidden_dim,).
h0 and c0 each have shape (hidden_dim,).

All arrays are floating-point and all dimensions are positive. There is no batch axis. Initialize h from h0 and c from c0, then process x from first to last using the same weights and bias at every step. Keep inputs unchanged.

Return only the final hidden state, with shape (hidden_dim,). Do not return the cell state, a pair of states, or the full sequence of hidden states.

Gate order

Let H = hidden_dim. At each step, compute one length-4H vector from the current input and previous hidden state:

gates = x[t] @ W_x + h @ W_h + b
f = sigmoid(gates[0:H])       # forget gate
i = sigmoid(gates[H:2*H])     # input gate
g = tanh(gates[2*H:3*H])      # candidate
o = sigmoid(gates[3*H:4*H])   # output gate

The packing order is forget, input, candidate, output. Split into four contiguous length-H slices, not groups of four values per hidden feature. Use this order even if a framework's built-in LSTM uses another convention. Use stable sigmoid operations.

Update the states

After computing all four gate values from the old h, update:

new_c = f * c + i * g
new_h = o * tanh(new_c)

All multiplications here are elementwise. f controls the retained part of the old cell state; i controls the new candidate contribution. o controls how much of tanh(new_c) is exposed as the hidden state. In particular, the hidden update uses the new cell state, not the old one.

Set c to new_c and h to new_h before moving to the next input. There is no extra projection, normalization or dropout.

Example

With all weights and biases zero, f, i and o are 0.5 and g is zero. Starting with c0 = [1], the first step gives c = [0.5] and h = [0.5 * tanh(0.5)], approximately [0.231059]. A second step gives c = [0.25] and final output approximately [0.122459].

Here h0 has no effect because W_h is zero. With nonzero recurrent weights, h0 and each updated h affect the next gates, while c is carried separately.

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…