Vanilla RNN: Forward & BPTT - Sequence Models

Vanilla RNN: Forward & BPTT

What you'll build

Two functions: rnn_forward, which runs a vanilla (Elman) RNN over a batch of sequences, and rnn_backward, which computes every gradient by hand with backpropagation through time (BPTT).

Background

An RNN reads a sequence one step at a time and carries a hidden state h that summarizes everything seen so far, reusing the same weights at every step. Writing BPTT yourself shows exactly why plain RNNs suffer from vanishing and exploding gradients, a classic interview topic, and it is the same backward pass that autograd runs for you.

Forward

h_t = tanh(x_t @ W_xh + h_{t-1} @ W_hh + b)        for t = 1..T, starting from h_0

Backward (BPTT)

Given dH = dL/dH_all of shape (N, T, H) (the loss may depend on every hidden state), walk the sequence backwards. The gradient reaching h_t has two parts: the direct one, dH[:, t], and the one flowing back from h_{t+1}:

dh_next = 0
for t = T..1:
    dh   = dH[:, t] + dh_next                # direct + from the future
    da   = dh * (1 - h_t ** 2)               # through tanh (tanh' = 1 - tanh^2)
    dW_xh += x_t.T @ da                      # shared weights: ACCUMULATE over t
    dW_hh += h_{t-1}.T @ da
    db    += da.sum(axis=0)
    dX[:, t] = da @ W_xh.T
    dh_next  = da @ W_hh.T                   # flows into h_{t-1}
dh0 = dh_next

The repeated multiplication by W_hh.T and (1 - h**2) is why gradients vanish or explode over long sequences.

Inputs (batch first)

Tensor	Shape	Meaning
X	(N, T, D)	N sequences of T steps with D features
h0	(N, H)	initial hidden state
W_xh	(D, H)	input-to-hidden weights, used as x @ W_xh
W_hh	(H, H)	hidden-to-hidden weights, used as h @ W_hh
b	(H,)	bias
H_all	(N, T, H)	output of rnn_forward (input to rnn_backward)
dH	(N, T, H)	upstream gradient dL/dH_all

Layout note: this problem uses (in, out) weights and one bias. nn.RNN stores the transpose, weight_ih_l0 of shape (H, D), and has two biases (b_ih + b_hh).

Expected return format

• rnn_forward returns H_all, shape (N, T, H): H_all[:, t] is the state after reading step t. h_0 is not included.
• rnn_backward returns a dictionary with these keys, each the same shape as the matching input:
  • 'dX': (N, T, D)
  • 'dh0': (N, H)
  • 'dW_xh': (D, H)
  • 'dW_hh': (H, H)
  • 'db': (H,)

Edge cases and conventions

• T = 1 works: one step, and dh0 = da @ W_hh.T.
• At t = 0 the previous state is h0, not zeros.
• Recompute nothing: the tanh derivative uses the stored outputs H_all.

Constraints

NumPy only. Loops over time are expected (the recurrence is sequential); vectorise over the batch.

Follow-up

• Why does gradient clipping fix exploding but not vanishing gradients?
• What is truncated BPTT and what does it trade away?

Input: X = np.ones((1, 3, 1)), h0 = [[0.0]], W_xh = [[0.5]], W_hh = [[1.0]], b = [0.0]
Output: [[[0.462117], [0.74522], [0.846938]]]

Input: X: (2, 3, 4), h0: (2, 5), W_xh: (4, 5), W_hh: (5, 5); rnn_backward(...)
Output: {'dX': (2, 3, 4), 'dh0': (2, 5), 'dW_xh': (4, 5), 'dW_hh': (5, 5), 'db': (5,)}

import numpy as np

def rnn_forward(X, h0, W_xh, W_hh, b):
    """
    Run a vanilla RNN over a batch of sequences:
    h_t = tanh(x_t @ W_xh + h_{t-1} @ W_hh + b).

    Args:
        X: Inputs (N, T, D), batch first
        h0: Initial hidden state (N, H)
        W_xh: Input-to-hidden weights (D, H)
        W_hh: Hidden-to-hidden weights (H, H)
        b: Bias (H,)

    Returns:
        H_all: Hidden states for every step (N, T, H), h0 not included
    """
    # Your code here
    pass


def rnn_backward(X, h0, W_xh, W_hh, H_all, dH):
    """
    Backpropagation through time.

    Args:
        X, h0, W_xh, W_hh: Same as rnn_forward
        H_all: Output of rnn_forward (N, T, H)
        dH: Gradient of the loss w.r.t. H_all (N, T, H)

    Returns:
        dict with 'dX' (N, T, D), 'dh0' (N, H), 'dW_xh' (D, H),
        'dW_hh' (H, H) and 'db' (H,)
    """
    # Your code here
    pass