15 RNN / LSTM — First Principles¶

Goal. Build vanilla RNN and LSTM cells from scratch in NumPy, demonstrate vanishing gradients, and verify against PyTorch.

Prerequisites. 13 Neural Networks — forward pass, backpropagation, gradient descent.

Theory. See theory.md for derivations of BPTT, vanishing gradient analysis, and LSTM gate equations.

In [1]:

import random

import matplotlib.pyplot as plt
import numpy as np

%load_ext autoreload
%autoreload 2

SEED = 42
random.seed(SEED)
rng = np.random.default_rng(SEED)

import random import matplotlib.pyplot as plt import numpy as np %load_ext autoreload %autoreload 2 SEED = 42 random.seed(SEED) rng = np.random.default_rng(SEED)

1. Problem Setup — WHY¶

Feed-forward networks treat each input independently. But many tasks have sequential structure: the output at time $t$ depends on inputs at times $1, 2, \ldots, t$. We need a model that:

1. Handles variable-length sequences.
2. Maintains a running summary (hidden state) of past inputs.
3. Shares parameters across time steps.

Demo: sequence pattern prediction¶

We'll work with a simple task: given a binary sequence, predict whether the first element was 0 or 1 after reading the entire sequence. This requires the model to remember information across all time steps.

In [2]:

def generate_remember_first(n_samples, seq_len, rng):
    """Generate sequences where the target is the first element.

    X: (n_samples, seq_len, 1) — binary sequences
    y: (n_samples,) — value of x[0] for each sequence
    """
    X = rng.integers(0, 2, size=(n_samples, seq_len, 1)).astype(np.float64)
    y = X[:, 0, 0]  # target = first element
    return X, y


# Visualize a few sequences
X_demo, y_demo = generate_remember_first(6, 20, rng)

fig, axes = plt.subplots(2, 3, figsize=(12, 5), sharex=True, sharey=True)
for i, ax in enumerate(axes.flat):
    ax.step(range(20), X_demo[i, :, 0], where='mid', color='steelblue')
    ax.set_title(f'target = {int(y_demo[i])} (first element)', fontsize=10)
    ax.set_ylim(-0.2, 1.2)
    if i >= 3:
        ax.set_xlabel('time step')
    if i % 3 == 0:
        ax.set_ylabel('value')
fig.suptitle('"Remember the First" Task — model must retain first element', fontsize=12)
plt.tight_layout()
plt.show()

def generate_remember_first(n_samples, seq_len, rng): """Generate sequences where the target is the first element. X: (n_samples, seq_len, 1) — binary sequences y: (n_samples,) — value of x[0] for each sequence """ X = rng.integers(0, 2, size=(n_samples, seq_len, 1)).astype(np.float64) y = X[:, 0, 0] # target = first element return X, y # Visualize a few sequences X_demo, y_demo = generate_remember_first(6, 20, rng) fig, axes = plt.subplots(2, 3, figsize=(12, 5), sharex=True, sharey=True) for i, ax in enumerate(axes.flat): ax.step(range(20), X_demo[i, :, 0], where='mid', color='steelblue') ax.set_title(f'target = {int(y_demo[i])} (first element)', fontsize=10) ax.set_ylim(-0.2, 1.2) if i >= 3: ax.set_xlabel('time step') if i % 3 == 0: ax.set_ylabel('value') fig.suptitle('"Remember the First" Task — model must retain first element', fontsize=12) plt.tight_layout() plt.show()

2. Mathematical Core — WHAT¶

Vanilla RNN¶

$$h_t = \tanh(W_{xh}\, x_t + W_{hh}\, h_{t-1} + b_h)$$ $$\hat{y} = W_{hy}\, h_T + b_y \quad \text{(many-to-one: use final hidden state)}$$

LSTM¶

$$f_t = \sigma(W_f\, [h_{t-1}, x_t] + b_f) \quad \text{(forget gate)}$$ $$i_t = \sigma(W_i\, [h_{t-1}, x_t] + b_i) \quad \text{(input gate)}$$ $$\tilde{c}_t = \tanh(W_c\, [h_{t-1}, x_t] + b_c)$$ $$c_t = f_t \odot c_{t-1} + i_t \odot \tilde{c}_t$$ $$o_t = \sigma(W_o\, [h_{t-1}, x_t] + b_o) \quad \text{(output gate)}$$ $$h_t = o_t \odot \tanh(c_t)$$

3. Solution Method — HOW¶

Training loop:

1. Forward pass: unroll the RNN/LSTM for $T$ steps, cache intermediates.
2. Loss: compute loss from the final hidden state (many-to-one) or all outputs (many-to-many).
3. BPTT: backpropagate through the unrolled graph, accumulating gradients for shared parameters at each step.
4. Update: gradient descent (with optional gradient clipping).

Gradient clipping (norm-based):

$$g \leftarrow g \cdot \min\!\left(1,\; \frac{\theta}{\|g\|}\right)$$

4. Implementation — BUILD¶

4.1 Vanilla RNN Cell¶

In [3]:

def sigmoid(z):
    """Numerically stable sigmoid."""
    return np.where(z >= 0,
                    1 / (1 + np.exp(-z)),
                    np.exp(z) / (1 + np.exp(z)))


class RNNCell:
    """Vanilla RNN cell with BPTT.

    h_t = tanh(W_xh @ x_t + W_hh @ h_{t-1} + b_h)
    """

    def __init__(self, input_dim, hidden_dim, rng):
        scale = np.sqrt(2.0 / (input_dim + hidden_dim))
        self.W_xh = rng.normal(0, scale, (hidden_dim, input_dim))
        self.W_hh = rng.normal(0, scale, (hidden_dim, hidden_dim))
        self.b_h = np.zeros(hidden_dim)
        self.hidden_dim = hidden_dim

    def forward(self, X):
        """Forward pass through the full sequence.

        Args:
            X: (seq_len, input_dim) — one sequence

        Returns:
            h_all: (seq_len, hidden_dim) — hidden states at each step
        """
        T = X.shape[0]
        h = np.zeros(self.hidden_dim)
        self.cache = {'X': X, 'h': [h.copy()], 'z': []}  # h[0] = h_0

        h_all = []
        for t in range(T):
            z = self.W_xh @ X[t] + self.W_hh @ h + self.b_h
            h = np.tanh(z)
            self.cache['z'].append(z)
            self.cache['h'].append(h.copy())
            h_all.append(h)

        return np.array(h_all)

    def backward(self, dh_all):
        """BPTT: backpropagate through the unrolled graph.

        Args:
            dh_all: (seq_len, hidden_dim) — gradient of loss w.r.t. each h_t

        Returns:
            None (gradients stored as self.dW_xh, self.dW_hh, self.db_h)
        """
        X = self.cache['X']
        T = X.shape[0]

        self.dW_xh = np.zeros_like(self.W_xh)
        self.dW_hh = np.zeros_like(self.W_hh)
        self.db_h = np.zeros_like(self.b_h)

        dh_next = np.zeros(self.hidden_dim)  # gradient from future steps

        for t in reversed(range(T)):
            dh = dh_all[t] + dh_next  # combine direct + future gradient

            # tanh derivative: d(tanh(z))/dz = 1 - tanh^2(z)
            h_t = self.cache['h'][t + 1]  # h[t+1] = tanh(z[t])
            dz = dh * (1 - h_t ** 2)

            # Accumulate parameter gradients
            self.dW_xh += np.outer(dz, X[t])
            self.dW_hh += np.outer(dz, self.cache['h'][t])
            self.db_h += dz

            # Gradient flowing to h_{t-1}
            dh_next = self.W_hh.T @ dz


# Quick test: forward pass shapes
rnn_cell = RNNCell(input_dim=1, hidden_dim=8, rng=rng)
X_test = rng.normal(size=(10, 1))
h_out = rnn_cell.forward(X_test)
assert h_out.shape == (10, 8), f'Expected (10, 8), got {h_out.shape}'

# Test backward pass
dh_test = np.zeros_like(h_out)
dh_test[-1] = rng.normal(size=8)  # gradient only at last step (many-to-one)
rnn_cell.backward(dh_test)
assert rnn_cell.dW_xh.shape == (8, 1)
assert rnn_cell.dW_hh.shape == (8, 8)
print(f'RNNCell forward output shape: {h_out.shape}')
print(f'RNNCell backward: dW_xh norm = {np.linalg.norm(rnn_cell.dW_xh):.6f}')
print('RNNCell forward + backward: ✓')

def sigmoid(z): """Numerically stable sigmoid.""" return np.where(z >= 0, 1 / (1 + np.exp(-z)), np.exp(z) / (1 + np.exp(z))) class RNNCell: """Vanilla RNN cell with BPTT. h_t = tanh(W_xh @ x_t + W_hh @ h_{t-1} + b_h) """ def __init__(self, input_dim, hidden_dim, rng): scale = np.sqrt(2.0 / (input_dim + hidden_dim)) self.W_xh = rng.normal(0, scale, (hidden_dim, input_dim)) self.W_hh = rng.normal(0, scale, (hidden_dim, hidden_dim)) self.b_h = np.zeros(hidden_dim) self.hidden_dim = hidden_dim def forward(self, X): """Forward pass through the full sequence. Args: X: (seq_len, input_dim) — one sequence Returns: h_all: (seq_len, hidden_dim) — hidden states at each step """ T = X.shape[0] h = np.zeros(self.hidden_dim) self.cache = {'X': X, 'h': [h.copy()], 'z': []} # h[0] = h_0 h_all = [] for t in range(T): z = self.W_xh @ X[t] + self.W_hh @ h + self.b_h h = np.tanh(z) self.cache['z'].append(z) self.cache['h'].append(h.copy()) h_all.append(h) return np.array(h_all) def backward(self, dh_all): """BPTT: backpropagate through the unrolled graph. Args: dh_all: (seq_len, hidden_dim) — gradient of loss w.r.t. each h_t Returns: None (gradients stored as self.dW_xh, self.dW_hh, self.db_h) """ X = self.cache['X'] T = X.shape[0] self.dW_xh = np.zeros_like(self.W_xh) self.dW_hh = np.zeros_like(self.W_hh) self.db_h = np.zeros_like(self.b_h) dh_next = np.zeros(self.hidden_dim) # gradient from future steps for t in reversed(range(T)): dh = dh_all[t] + dh_next # combine direct + future gradient # tanh derivative: d(tanh(z))/dz = 1 - tanh^2(z) h_t = self.cache['h'][t + 1] # h[t+1] = tanh(z[t]) dz = dh * (1 - h_t ** 2) # Accumulate parameter gradients self.dW_xh += np.outer(dz, X[t]) self.dW_hh += np.outer(dz, self.cache['h'][t]) self.db_h += dz # Gradient flowing to h_{t-1} dh_next = self.W_hh.T @ dz # Quick test: forward pass shapes rnn_cell = RNNCell(input_dim=1, hidden_dim=8, rng=rng) X_test = rng.normal(size=(10, 1)) h_out = rnn_cell.forward(X_test) assert h_out.shape == (10, 8), f'Expected (10, 8), got {h_out.shape}' # Test backward pass dh_test = np.zeros_like(h_out) dh_test[-1] = rng.normal(size=8) # gradient only at last step (many-to-one) rnn_cell.backward(dh_test) assert rnn_cell.dW_xh.shape == (8, 1) assert rnn_cell.dW_hh.shape == (8, 8) print(f'RNNCell forward output shape: {h_out.shape}') print(f'RNNCell backward: dW_xh norm = {np.linalg.norm(rnn_cell.dW_xh):.6f}') print('RNNCell forward + backward: ✓')

RNNCell forward output shape: (10, 8)
RNNCell backward: dW_xh norm = 2.672323
RNNCell forward + backward: ✓

4.2 Vanishing Gradient Demonstration¶

We send a gradient signal from the last time step and measure how its norm decays as it propagates backward through the sequence. For a vanilla RNN, the gradient should decay exponentially.

In [4]:

def measure_gradient_flow(cell_class, seq_lengths, input_dim, hidden_dim, rng):
    """Measure gradient norm at first step vs sequence length."""
    grad_norms = []
    for T in seq_lengths:
        cell = cell_class(input_dim, hidden_dim, rng)
        X = rng.normal(size=(T, input_dim)) * 0.5
        h_out = cell.forward(X)

        # Send unit gradient from last step only
        dh = np.zeros_like(h_out)
        dh[-1] = np.ones(hidden_dim)
        cell.backward(dh)

        # Measure: how much gradient reached W_xh from early steps?
        # We re-run backward tracking per-step gradient norms
        dh_next = np.zeros(hidden_dim)
        norms = []
        for t in reversed(range(T)):
            dh_t = dh[t] + dh_next
            h_t = cell.cache['h'][t + 1]
            dz = dh_t * (1 - h_t ** 2)
            dh_next = cell.W_hh.T @ dz
            norms.append(np.linalg.norm(dh_next))

        # Norm at the first time step (after propagating through T-1 steps)
        grad_norms.append(norms[-1] if norms else 0.0)

    return grad_norms


seq_lengths = [5, 10, 20, 30, 50, 75, 100]
rng_vanish = np.random.default_rng(123)
vanilla_norms = measure_gradient_flow(RNNCell, seq_lengths, 1, 16, rng_vanish)

fig, ax = plt.subplots(figsize=(8, 5))
ax.semilogy(seq_lengths, vanilla_norms, 'o-', color='crimson', label='Vanilla RNN')
ax.set_xlabel('Sequence length T')
ax.set_ylabel('Gradient norm at first time step (log scale)')
ax.set_title('Vanishing Gradients: gradient decays exponentially with sequence length')
ax.legend()
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

print(f'Gradient norm ratio (T=100)/(T=5): {vanilla_norms[-1] / (vanilla_norms[0] + 1e-30):.2e}')
print('This exponential decay is the vanishing gradient problem.')

def measure_gradient_flow(cell_class, seq_lengths, input_dim, hidden_dim, rng): """Measure gradient norm at first step vs sequence length.""" grad_norms = [] for T in seq_lengths: cell = cell_class(input_dim, hidden_dim, rng) X = rng.normal(size=(T, input_dim)) * 0.5 h_out = cell.forward(X) # Send unit gradient from last step only dh = np.zeros_like(h_out) dh[-1] = np.ones(hidden_dim) cell.backward(dh) # Measure: how much gradient reached W_xh from early steps? # We re-run backward tracking per-step gradient norms dh_next = np.zeros(hidden_dim) norms = [] for t in reversed(range(T)): dh_t = dh[t] + dh_next h_t = cell.cache['h'][t + 1] dz = dh_t * (1 - h_t ** 2) dh_next = cell.W_hh.T @ dz norms.append(np.linalg.norm(dh_next)) # Norm at the first time step (after propagating through T-1 steps) grad_norms.append(norms[-1] if norms else 0.0) return grad_norms seq_lengths = [5, 10, 20, 30, 50, 75, 100] rng_vanish = np.random.default_rng(123) vanilla_norms = measure_gradient_flow(RNNCell, seq_lengths, 1, 16, rng_vanish) fig, ax = plt.subplots(figsize=(8, 5)) ax.semilogy(seq_lengths, vanilla_norms, 'o-', color='crimson', label='Vanilla RNN') ax.set_xlabel('Sequence length T') ax.set_ylabel('Gradient norm at first time step (log scale)') ax.set_title('Vanishing Gradients: gradient decays exponentially with sequence length') ax.legend() ax.grid(True, alpha=0.3) plt.tight_layout() plt.show() print(f'Gradient norm ratio (T=100)/(T=5): {vanilla_norms[-1] / (vanilla_norms[0] + 1e-30):.2e}') print('This exponential decay is the vanishing gradient problem.')

Gradient norm ratio (T=100)/(T=5): 7.56e-03
This exponential decay is the vanishing gradient problem.

4.3 LSTM Cell¶

In [5]:

class LSTMCell:
    """LSTM cell with forward pass and BPTT.

    Gates: forget (f), input (i), output (o), candidate (g).
    Cell update: c_t = f_t * c_{t-1} + i_t * g_t
    Hidden: h_t = o_t * tanh(c_t)
    """

    def __init__(self, input_dim, hidden_dim, rng):
        self.hidden_dim = hidden_dim
        concat_dim = hidden_dim + input_dim
        scale = np.sqrt(2.0 / concat_dim)

        # All four gate weight matrices stacked: f, i, g, o
        self.W = rng.normal(0, scale, (4 * hidden_dim, concat_dim))
        self.b = np.zeros(4 * hidden_dim)
        # Forget gate bias = 1.0 to encourage remembering initially
        self.b[:hidden_dim] = 1.0

    def forward(self, X):
        """Forward pass through the sequence.

        Args:
            X: (seq_len, input_dim)

        Returns:
            h_all: (seq_len, hidden_dim)
        """
        T, d = X.shape
        m = self.hidden_dim
        h = np.zeros(m)
        c = np.zeros(m)

        # Cache for backward pass
        self.cache = {
            'X': X, 'h': [h.copy()], 'c': [c.copy()],
            'f': [], 'i': [], 'g': [], 'o': [], 'concat': []
        }

        h_all = []
        for t in range(T):
            concat = np.concatenate([h, X[t]])
            gates = self.W @ concat + self.b

            f = sigmoid(gates[:m])        # forget gate
            i = sigmoid(gates[m:2*m])     # input gate
            g = np.tanh(gates[2*m:3*m])   # candidate
            o = sigmoid(gates[3*m:])      # output gate

            c = f * c + i * g             # cell state update
            h = o * np.tanh(c)            # hidden state

            self.cache['concat'].append(concat)
            self.cache['f'].append(f)
            self.cache['i'].append(i)
            self.cache['g'].append(g)
            self.cache['o'].append(o)
            self.cache['c'].append(c.copy())
            self.cache['h'].append(h.copy())
            h_all.append(h)

        return np.array(h_all)

    def backward(self, dh_all):
        """BPTT for LSTM.

        Args:
            dh_all: (seq_len, hidden_dim)
        """
        T = dh_all.shape[0]
        m = self.hidden_dim

        self.dW = np.zeros_like(self.W)
        self.db = np.zeros_like(self.b)

        dh_next = np.zeros(m)
        dc_next = np.zeros(m)

        for t in reversed(range(T)):
            dh = dh_all[t] + dh_next

            f = self.cache['f'][t]
            i = self.cache['i'][t]
            g = self.cache['g'][t]
            o = self.cache['o'][t]
            c = self.cache['c'][t + 1]   # c after update at step t
            c_prev = self.cache['c'][t]   # c before update at step t
            concat = self.cache['concat'][t]

            # dh -> dc through o * tanh(c)
            tanh_c = np.tanh(c)
            dc = dh * o * (1 - tanh_c ** 2) + dc_next

            # Gate gradients
            df = dc * c_prev
            di = dc * g
            dg = dc * i
            do = dh * tanh_c

            # Through gate nonlinearities
            df_raw = df * f * (1 - f)     # sigmoid derivative
            di_raw = di * i * (1 - i)
            dg_raw = dg * (1 - g ** 2)    # tanh derivative
            do_raw = do * o * (1 - o)

            dgates = np.concatenate([df_raw, di_raw, dg_raw, do_raw])

            # Accumulate parameter gradients
            self.dW += np.outer(dgates, concat)
            self.db += dgates

            # Gradient to previous hidden state and cell state
            d_concat = self.W.T @ dgates
            dh_next = d_concat[:m]
            dc_next = dc * f  # gradient highway through forget gate!


# Quick test
lstm_cell = LSTMCell(input_dim=1, hidden_dim=8, rng=rng)
h_lstm = lstm_cell.forward(X_test)
assert h_lstm.shape == (10, 8), f'Expected (10, 8), got {h_lstm.shape}'

dh_lstm = np.zeros_like(h_lstm)
dh_lstm[-1] = rng.normal(size=8)
lstm_cell.backward(dh_lstm)
assert lstm_cell.dW.shape == (32, 9)  # 4*8 x (8+1)
print(f'LSTMCell forward output shape: {h_lstm.shape}')
print(f'LSTMCell backward: dW norm = {np.linalg.norm(lstm_cell.dW):.6f}')
print('LSTMCell forward + backward: ✓')

class LSTMCell: """LSTM cell with forward pass and BPTT. Gates: forget (f), input (i), output (o), candidate (g). Cell update: c_t = f_t * c_{t-1} + i_t * g_t Hidden: h_t = o_t * tanh(c_t) """ def __init__(self, input_dim, hidden_dim, rng): self.hidden_dim = hidden_dim concat_dim = hidden_dim + input_dim scale = np.sqrt(2.0 / concat_dim) # All four gate weight matrices stacked: f, i, g, o self.W = rng.normal(0, scale, (4 * hidden_dim, concat_dim)) self.b = np.zeros(4 * hidden_dim) # Forget gate bias = 1.0 to encourage remembering initially self.b[:hidden_dim] = 1.0 def forward(self, X): """Forward pass through the sequence. Args: X: (seq_len, input_dim) Returns: h_all: (seq_len, hidden_dim) """ T, d = X.shape m = self.hidden_dim h = np.zeros(m) c = np.zeros(m) # Cache for backward pass self.cache = { 'X': X, 'h': [h.copy()], 'c': [c.copy()], 'f': [], 'i': [], 'g': [], 'o': [], 'concat': [] } h_all = [] for t in range(T): concat = np.concatenate([h, X[t]]) gates = self.W @ concat + self.b f = sigmoid(gates[:m]) # forget gate i = sigmoid(gates[m:2*m]) # input gate g = np.tanh(gates[2*m:3*m]) # candidate o = sigmoid(gates[3*m:]) # output gate c = f * c + i * g # cell state update h = o * np.tanh(c) # hidden state self.cache['concat'].append(concat) self.cache['f'].append(f) self.cache['i'].append(i) self.cache['g'].append(g) self.cache['o'].append(o) self.cache['c'].append(c.copy()) self.cache['h'].append(h.copy()) h_all.append(h) return np.array(h_all) def backward(self, dh_all): """BPTT for LSTM. Args: dh_all: (seq_len, hidden_dim) """ T = dh_all.shape[0] m = self.hidden_dim self.dW = np.zeros_like(self.W) self.db = np.zeros_like(self.b) dh_next = np.zeros(m) dc_next = np.zeros(m) for t in reversed(range(T)): dh = dh_all[t] + dh_next f = self.cache['f'][t] i = self.cache['i'][t] g = self.cache['g'][t] o = self.cache['o'][t] c = self.cache['c'][t + 1] # c after update at step t c_prev = self.cache['c'][t] # c before update at step t concat = self.cache['concat'][t] # dh -> dc through o * tanh(c) tanh_c = np.tanh(c) dc = dh * o * (1 - tanh_c ** 2) + dc_next # Gate gradients df = dc * c_prev di = dc * g dg = dc * i do = dh * tanh_c # Through gate nonlinearities df_raw = df * f * (1 - f) # sigmoid derivative di_raw = di * i * (1 - i) dg_raw = dg * (1 - g ** 2) # tanh derivative do_raw = do * o * (1 - o) dgates = np.concatenate([df_raw, di_raw, dg_raw, do_raw]) # Accumulate parameter gradients self.dW += np.outer(dgates, concat) self.db += dgates # Gradient to previous hidden state and cell state d_concat = self.W.T @ dgates dh_next = d_concat[:m] dc_next = dc * f # gradient highway through forget gate! # Quick test lstm_cell = LSTMCell(input_dim=1, hidden_dim=8, rng=rng) h_lstm = lstm_cell.forward(X_test) assert h_lstm.shape == (10, 8), f'Expected (10, 8), got {h_lstm.shape}' dh_lstm = np.zeros_like(h_lstm) dh_lstm[-1] = rng.normal(size=8) lstm_cell.backward(dh_lstm) assert lstm_cell.dW.shape == (32, 9) # 4*8 x (8+1) print(f'LSTMCell forward output shape: {h_lstm.shape}') print(f'LSTMCell backward: dW norm = {np.linalg.norm(lstm_cell.dW):.6f}') print('LSTMCell forward + backward: ✓')

LSTMCell forward output shape: (10, 8)
LSTMCell backward: dW norm = 1.028193
LSTMCell forward + backward: ✓

5. Library Comparison¶

We verify our from-scratch implementations by comparing forward pass outputs with PyTorch's nn.RNN and nn.LSTM, using identical weights.

In [6]:

try:
    import torch
    import torch.nn as nn
    HAS_TORCH = True
    torch.manual_seed(SEED)
except ImportError:
    HAS_TORCH = False
    print('PyTorch not available — skipping library comparison.')

if HAS_TORCH:
    # ---- Vanilla RNN comparison ----
    input_dim, hidden_dim = 3, 5
    seq_len = 7

    # Create our RNN with known weights
    rng_cmp = np.random.default_rng(77)
    our_rnn = RNNCell(input_dim, hidden_dim, rng_cmp)

    # Create PyTorch RNN with same weights
    pt_rnn = nn.RNN(input_dim, hidden_dim, batch_first=False, bias=True)
    with torch.no_grad():
        pt_rnn.weight_ih_l0.copy_(torch.from_numpy(our_rnn.W_xh))
        pt_rnn.weight_hh_l0.copy_(torch.from_numpy(our_rnn.W_hh))
        pt_rnn.bias_ih_l0.zero_()
        pt_rnn.bias_hh_l0.copy_(torch.from_numpy(our_rnn.b_h))

    X_cmp = rng_cmp.normal(size=(seq_len, input_dim))
    h_ours = our_rnn.forward(X_cmp)

    X_torch = torch.from_numpy(X_cmp).unsqueeze(1)  # (T, 1, d)
    h0 = torch.zeros(1, 1, hidden_dim)
    with torch.no_grad():
        h_pt, _ = pt_rnn(X_torch, h0)
    h_pt_np = h_pt.squeeze(1).numpy()

    rnn_match = np.allclose(h_ours, h_pt_np, atol=1e-6)
    print(f'RNN output match (atol=1e-6): {rnn_match}')
    print(f'  Max absolute difference: {np.max(np.abs(h_ours - h_pt_np)):.2e}')
    assert rnn_match, 'RNN outputs do not match PyTorch!'

    # ---- LSTM comparison ----
    rng_cmp2 = np.random.default_rng(88)
    our_lstm = LSTMCell(input_dim, hidden_dim, rng_cmp2)

    pt_lstm = nn.LSTM(input_dim, hidden_dim, batch_first=False, bias=True)
    with torch.no_grad():
        # PyTorch LSTM stacks gates as [i, f, g, o] — we use [f, i, g, o]
        # Reorder our weights to match PyTorch convention
        m = hidden_dim
        # Our order: f, i, g, o -> PyTorch order: i, f, g, o
        reorder = np.concatenate([
            our_lstm.W[m:2*m],    # i
            our_lstm.W[:m],       # f
            our_lstm.W[2*m:3*m],  # g
            our_lstm.W[3*m:],     # o
        ])
        reorder_b = np.concatenate([
            our_lstm.b[m:2*m],
            our_lstm.b[:m],
            our_lstm.b[2*m:3*m],
            our_lstm.b[3*m:],
        ])
        concat_dim = hidden_dim + input_dim
        pt_lstm.weight_ih_l0.copy_(torch.from_numpy(reorder[:, hidden_dim:]))
        pt_lstm.weight_hh_l0.copy_(torch.from_numpy(reorder[:, :hidden_dim]))
        pt_lstm.bias_ih_l0.zero_()
        pt_lstm.bias_hh_l0.copy_(torch.from_numpy(reorder_b))

    h_lstm_ours = our_lstm.forward(X_cmp)

    h0_lstm = torch.zeros(1, 1, hidden_dim)
    c0_lstm = torch.zeros(1, 1, hidden_dim)
    with torch.no_grad():
        h_pt_lstm, _ = pt_lstm(X_torch, (h0_lstm, c0_lstm))
    h_pt_lstm_np = h_pt_lstm.squeeze(1).numpy()

    lstm_match = np.allclose(h_lstm_ours, h_pt_lstm_np, atol=1e-6)
    print(f'\nLSTM output match (atol=1e-6): {lstm_match}')
    print(f'  Max absolute difference: {np.max(np.abs(h_lstm_ours - h_pt_lstm_np)):.2e}')
    assert lstm_match, 'LSTM outputs do not match PyTorch!'
    print('\nBoth RNN and LSTM match PyTorch. ✓')

try: import torch import torch.nn as nn HAS_TORCH = True torch.manual_seed(SEED) except ImportError: HAS_TORCH = False print('PyTorch not available — skipping library comparison.') if HAS_TORCH: # ---- Vanilla RNN comparison ---- input_dim, hidden_dim = 3, 5 seq_len = 7 # Create our RNN with known weights rng_cmp = np.random.default_rng(77) our_rnn = RNNCell(input_dim, hidden_dim, rng_cmp) # Create PyTorch RNN with same weights pt_rnn = nn.RNN(input_dim, hidden_dim, batch_first=False, bias=True) with torch.no_grad(): pt_rnn.weight_ih_l0.copy_(torch.from_numpy(our_rnn.W_xh)) pt_rnn.weight_hh_l0.copy_(torch.from_numpy(our_rnn.W_hh)) pt_rnn.bias_ih_l0.zero_() pt_rnn.bias_hh_l0.copy_(torch.from_numpy(our_rnn.b_h)) X_cmp = rng_cmp.normal(size=(seq_len, input_dim)) h_ours = our_rnn.forward(X_cmp) X_torch = torch.from_numpy(X_cmp).unsqueeze(1) # (T, 1, d) h0 = torch.zeros(1, 1, hidden_dim) with torch.no_grad(): h_pt, _ = pt_rnn(X_torch, h0) h_pt_np = h_pt.squeeze(1).numpy() rnn_match = np.allclose(h_ours, h_pt_np, atol=1e-6) print(f'RNN output match (atol=1e-6): {rnn_match}') print(f' Max absolute difference: {np.max(np.abs(h_ours - h_pt_np)):.2e}') assert rnn_match, 'RNN outputs do not match PyTorch!' # ---- LSTM comparison ---- rng_cmp2 = np.random.default_rng(88) our_lstm = LSTMCell(input_dim, hidden_dim, rng_cmp2) pt_lstm = nn.LSTM(input_dim, hidden_dim, batch_first=False, bias=True) with torch.no_grad(): # PyTorch LSTM stacks gates as [i, f, g, o] — we use [f, i, g, o] # Reorder our weights to match PyTorch convention m = hidden_dim # Our order: f, i, g, o -> PyTorch order: i, f, g, o reorder = np.concatenate([ our_lstm.W[m:2*m], # i our_lstm.W[:m], # f our_lstm.W[2*m:3*m], # g our_lstm.W[3*m:], # o ]) reorder_b = np.concatenate([ our_lstm.b[m:2*m], our_lstm.b[:m], our_lstm.b[2*m:3*m], our_lstm.b[3*m:], ]) concat_dim = hidden_dim + input_dim pt_lstm.weight_ih_l0.copy_(torch.from_numpy(reorder[:, hidden_dim:])) pt_lstm.weight_hh_l0.copy_(torch.from_numpy(reorder[:, :hidden_dim])) pt_lstm.bias_ih_l0.zero_() pt_lstm.bias_hh_l0.copy_(torch.from_numpy(reorder_b)) h_lstm_ours = our_lstm.forward(X_cmp) h0_lstm = torch.zeros(1, 1, hidden_dim) c0_lstm = torch.zeros(1, 1, hidden_dim) with torch.no_grad(): h_pt_lstm, _ = pt_lstm(X_torch, (h0_lstm, c0_lstm)) h_pt_lstm_np = h_pt_lstm.squeeze(1).numpy() lstm_match = np.allclose(h_lstm_ours, h_pt_lstm_np, atol=1e-6) print(f'\nLSTM output match (atol=1e-6): {lstm_match}') print(f' Max absolute difference: {np.max(np.abs(h_lstm_ours - h_pt_lstm_np)):.2e}') assert lstm_match, 'LSTM outputs do not match PyTorch!' print('\nBoth RNN and LSTM match PyTorch. ✓')

PyTorch not available — skipping library comparison.

6. Experiments and Failures — VERIFY¶

6.1 RNN vs LSTM on the "Remember First" task¶

This task tests long-range dependency: the model must remember the first element of a sequence after processing all subsequent elements.

In [7]:

class SequenceClassifier:
    """Many-to-one classifier: RNN/LSTM + linear output layer."""

    def __init__(self, cell, hidden_dim, output_dim, rng):
        self.cell = cell
        scale = np.sqrt(2.0 / hidden_dim)
        self.W_out = rng.normal(0, scale, (output_dim, hidden_dim))
        self.b_out = np.zeros(output_dim)

    def forward(self, X):
        """Forward pass: run cell, take last hidden state, linear + sigmoid."""
        h_all = self.cell.forward(X)
        self.h_last = h_all[-1]
        logit = self.W_out @ self.h_last + self.b_out
        self.y_hat = sigmoid(logit)
        return self.y_hat

    def backward(self, y_true):
        """Backward pass with BCE loss."""
        # BCE gradient: d(L)/d(logit) = y_hat - y
        d_logit = self.y_hat - y_true

        # Output layer gradients
        self.dW_out = np.outer(d_logit, self.h_last)
        self.db_out = d_logit

        # Gradient to last hidden state
        dh_last = self.W_out.T @ d_logit

        # BPTT through the cell
        T = self.cell.cache['X'].shape[0]
        dh_all = np.zeros((T, self.cell.hidden_dim))
        dh_all[-1] = dh_last
        self.cell.backward(dh_all)

    def bce_loss(self, y_hat, y_true):
        """Binary cross-entropy loss."""
        eps = 1e-12
        return -np.mean(
            y_true * np.log(y_hat + eps) + (1 - y_true) * np.log(1 - y_hat + eps)
        )

    def clip_gradients(self, max_norm):
        """Clip all gradients by global norm."""
        all_grads = [self.dW_out, self.db_out]
        if isinstance(self.cell, RNNCell):
            all_grads += [self.cell.dW_xh, self.cell.dW_hh, self.cell.db_h]
        elif isinstance(self.cell, LSTMCell):
            all_grads += [self.cell.dW, self.cell.db]

        total_norm = np.sqrt(sum(np.sum(g ** 2) for g in all_grads))
        if total_norm > max_norm:
            scale = max_norm / total_norm
            for g in all_grads:
                g *= scale

    def update(self, lr):
        """SGD update."""
        self.W_out -= lr * self.dW_out
        self.b_out -= lr * self.db_out
        if isinstance(self.cell, RNNCell):
            self.cell.W_xh -= lr * self.cell.dW_xh
            self.cell.W_hh -= lr * self.cell.dW_hh
            self.cell.b_h -= lr * self.cell.db_h
        elif isinstance(self.cell, LSTMCell):
            self.cell.W -= lr * self.cell.dW
            self.cell.b -= lr * self.cell.db

class SequenceClassifier: """Many-to-one classifier: RNN/LSTM + linear output layer.""" def __init__(self, cell, hidden_dim, output_dim, rng): self.cell = cell scale = np.sqrt(2.0 / hidden_dim) self.W_out = rng.normal(0, scale, (output_dim, hidden_dim)) self.b_out = np.zeros(output_dim) def forward(self, X): """Forward pass: run cell, take last hidden state, linear + sigmoid.""" h_all = self.cell.forward(X) self.h_last = h_all[-1] logit = self.W_out @ self.h_last + self.b_out self.y_hat = sigmoid(logit) return self.y_hat def backward(self, y_true): """Backward pass with BCE loss.""" # BCE gradient: d(L)/d(logit) = y_hat - y d_logit = self.y_hat - y_true # Output layer gradients self.dW_out = np.outer(d_logit, self.h_last) self.db_out = d_logit # Gradient to last hidden state dh_last = self.W_out.T @ d_logit # BPTT through the cell T = self.cell.cache['X'].shape[0] dh_all = np.zeros((T, self.cell.hidden_dim)) dh_all[-1] = dh_last self.cell.backward(dh_all) def bce_loss(self, y_hat, y_true): """Binary cross-entropy loss.""" eps = 1e-12 return -np.mean( y_true * np.log(y_hat + eps) + (1 - y_true) * np.log(1 - y_hat + eps) ) def clip_gradients(self, max_norm): """Clip all gradients by global norm.""" all_grads = [self.dW_out, self.db_out] if isinstance(self.cell, RNNCell): all_grads += [self.cell.dW_xh, self.cell.dW_hh, self.cell.db_h] elif isinstance(self.cell, LSTMCell): all_grads += [self.cell.dW, self.cell.db] total_norm = np.sqrt(sum(np.sum(g ** 2) for g in all_grads)) if total_norm > max_norm: scale = max_norm / total_norm for g in all_grads: g *= scale def update(self, lr): """SGD update.""" self.W_out -= lr * self.dW_out self.b_out -= lr * self.db_out if isinstance(self.cell, RNNCell): self.cell.W_xh -= lr * self.cell.dW_xh self.cell.W_hh -= lr * self.cell.dW_hh self.cell.b_h -= lr * self.cell.db_h elif isinstance(self.cell, LSTMCell): self.cell.W -= lr * self.cell.dW self.cell.b -= lr * self.cell.db

In [8]:

def train_model(model, X_train, y_train, n_epochs, lr, clip_norm=5.0):
    """Train a SequenceClassifier on individual sequences."""
    losses = []
    accuracies = []
    n = X_train.shape[0]

    for epoch in range(n_epochs):
        epoch_loss = 0.0
        correct = 0
        for idx in range(n):
            y_hat = model.forward(X_train[idx])  # (seq_len, input_dim)
            loss = model.bce_loss(y_hat, y_train[idx:idx+1])
            epoch_loss += loss

            model.backward(y_train[idx:idx+1])
            model.clip_gradients(clip_norm)
            model.update(lr)

            pred = (y_hat[0] > 0.5).astype(float)
            correct += (pred == y_train[idx])

        losses.append(epoch_loss / n)
        accuracies.append(correct / n)

    return losses, accuracies


# Generate data with a moderately long sequence
SEQ_LEN = 25
N_TRAIN = 100
N_EPOCHS = 40
HIDDEN = 16
LR = 0.01

rng_exp = np.random.default_rng(42)
X_train, y_train = generate_remember_first(N_TRAIN, SEQ_LEN, rng_exp)
X_val, y_val = generate_remember_first(50, SEQ_LEN, rng_exp)

# Train vanilla RNN
rng_rnn = np.random.default_rng(42)
rnn_model = SequenceClassifier(
    RNNCell(1, HIDDEN, rng_rnn), HIDDEN, 1, rng_rnn
)
rnn_losses, rnn_accs = train_model(rnn_model, X_train, y_train, N_EPOCHS, LR)

# Train LSTM
rng_lstm = np.random.default_rng(42)
lstm_model = SequenceClassifier(
    LSTMCell(1, HIDDEN, rng_lstm), HIDDEN, 1, rng_lstm
)
lstm_losses, lstm_accs = train_model(lstm_model, X_train, y_train, N_EPOCHS, LR)

print(f'RNN  — final loss: {rnn_losses[-1]:.4f}, final accuracy: {rnn_accs[-1]:.2%}')
print(f'LSTM — final loss: {lstm_losses[-1]:.4f}, final accuracy: {lstm_accs[-1]:.2%}')

def train_model(model, X_train, y_train, n_epochs, lr, clip_norm=5.0): """Train a SequenceClassifier on individual sequences.""" losses = [] accuracies = [] n = X_train.shape[0] for epoch in range(n_epochs): epoch_loss = 0.0 correct = 0 for idx in range(n): y_hat = model.forward(X_train[idx]) # (seq_len, input_dim) loss = model.bce_loss(y_hat, y_train[idx:idx+1]) epoch_loss += loss model.backward(y_train[idx:idx+1]) model.clip_gradients(clip_norm) model.update(lr) pred = (y_hat[0] > 0.5).astype(float) correct += (pred == y_train[idx]) losses.append(epoch_loss / n) accuracies.append(correct / n) return losses, accuracies # Generate data with a moderately long sequence SEQ_LEN = 25 N_TRAIN = 100 N_EPOCHS = 40 HIDDEN = 16 LR = 0.01 rng_exp = np.random.default_rng(42) X_train, y_train = generate_remember_first(N_TRAIN, SEQ_LEN, rng_exp) X_val, y_val = generate_remember_first(50, SEQ_LEN, rng_exp) # Train vanilla RNN rng_rnn = np.random.default_rng(42) rnn_model = SequenceClassifier( RNNCell(1, HIDDEN, rng_rnn), HIDDEN, 1, rng_rnn ) rnn_losses, rnn_accs = train_model(rnn_model, X_train, y_train, N_EPOCHS, LR) # Train LSTM rng_lstm = np.random.default_rng(42) lstm_model = SequenceClassifier( LSTMCell(1, HIDDEN, rng_lstm), HIDDEN, 1, rng_lstm ) lstm_losses, lstm_accs = train_model(lstm_model, X_train, y_train, N_EPOCHS, LR) print(f'RNN — final loss: {rnn_losses[-1]:.4f}, final accuracy: {rnn_accs[-1]:.2%}') print(f'LSTM — final loss: {lstm_losses[-1]:.4f}, final accuracy: {lstm_accs[-1]:.2%}')

RNN  — final loss: 0.0020, final accuracy: 100.00%
LSTM — final loss: 0.0162, final accuracy: 100.00%

In [9]:

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5))

ax1.plot(rnn_losses, 'o-', color='crimson', markersize=3, label='Vanilla RNN')
ax1.plot(lstm_losses, 's-', color='steelblue', markersize=3, label='LSTM')
ax1.set_xlabel('Epoch')
ax1.set_ylabel('BCE Loss')
ax1.set_title(f'"Remember First" (T={SEQ_LEN}) — Training Loss')
ax1.legend()
ax1.grid(True, alpha=0.3)

ax2.plot(rnn_accs, 'o-', color='crimson', markersize=3, label='Vanilla RNN')
ax2.plot(lstm_accs, 's-', color='steelblue', markersize=3, label='LSTM')
ax2.axhline(0.5, color='gray', ls='--', alpha=0.5, label='Random baseline')
ax2.set_xlabel('Epoch')
ax2.set_ylabel('Accuracy')
ax2.set_title(f'"Remember First" (T={SEQ_LEN}) — Training Accuracy')
ax2.legend()
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

print('\nLSTM\'s gated cell state allows it to remember across long sequences.')
print('Vanilla RNN struggles due to vanishing gradients.')

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5)) ax1.plot(rnn_losses, 'o-', color='crimson', markersize=3, label='Vanilla RNN') ax1.plot(lstm_losses, 's-', color='steelblue', markersize=3, label='LSTM') ax1.set_xlabel('Epoch') ax1.set_ylabel('BCE Loss') ax1.set_title(f'"Remember First" (T={SEQ_LEN}) — Training Loss') ax1.legend() ax1.grid(True, alpha=0.3) ax2.plot(rnn_accs, 'o-', color='crimson', markersize=3, label='Vanilla RNN') ax2.plot(lstm_accs, 's-', color='steelblue', markersize=3, label='LSTM') ax2.axhline(0.5, color='gray', ls='--', alpha=0.5, label='Random baseline') ax2.set_xlabel('Epoch') ax2.set_ylabel('Accuracy') ax2.set_title(f'"Remember First" (T={SEQ_LEN}) — Training Accuracy') ax2.legend() ax2.grid(True, alpha=0.3) plt.tight_layout() plt.show() print('\nLSTM\'s gated cell state allows it to remember across long sequences.') print('Vanilla RNN struggles due to vanishing gradients.')

LSTM's gated cell state allows it to remember across long sequences.
Vanilla RNN struggles due to vanishing gradients.

6.2 Failure: Exploding gradients without clipping¶

Without gradient clipping, the vanilla RNN can experience exploding gradients when the spectral norm of $W_{hh}$ is large.

In [10]:

# Demonstrate exploding gradients with large weight initialization
rng_exp2 = np.random.default_rng(55)
rnn_big = RNNCell(1, 16, rng_exp2)
rnn_big.W_hh *= 3.0  # scale up to trigger explosion

X_long = rng_exp2.normal(size=(50, 1)) * 0.1
h_long = rnn_big.forward(X_long)

dh_long = np.zeros_like(h_long)
dh_long[-1] = np.ones(16)

# Track gradient norms at each time step during BPTT
grad_norms_per_step = []
dh_next = np.zeros(16)
for t in reversed(range(50)):
    dh = dh_long[t] + dh_next
    h_t = rnn_big.cache['h'][t + 1]
    dz = dh * (1 - h_t ** 2)
    dh_next = rnn_big.W_hh.T @ dz
    grad_norms_per_step.append(np.linalg.norm(dh_next))

grad_norms_per_step.reverse()  # now indexed by time step

fig, ax = plt.subplots(figsize=(8, 5))
ax.semilogy(range(50), grad_norms_per_step, 'o-', color='darkorange', markersize=3)
ax.set_xlabel('Time step (backward from T=50)')
ax.set_ylabel('Gradient norm (log scale)')
ax.set_title('Exploding Gradients: large W_hh causes exponential gradient growth')
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

spectral_norm = np.linalg.norm(rnn_big.W_hh, ord=2)
print(f'Spectral norm of W_hh: {spectral_norm:.4f}')
print(f'Gradient norm at step 0: {grad_norms_per_step[0]:.2e}')
print(f'Gradient norm at step 49: {grad_norms_per_step[49]:.2e}')
print(f'Ratio: {grad_norms_per_step[0] / (grad_norms_per_step[49] + 1e-30):.2e}')

# Demonstrate exploding gradients with large weight initialization rng_exp2 = np.random.default_rng(55) rnn_big = RNNCell(1, 16, rng_exp2) rnn_big.W_hh *= 3.0 # scale up to trigger explosion X_long = rng_exp2.normal(size=(50, 1)) * 0.1 h_long = rnn_big.forward(X_long) dh_long = np.zeros_like(h_long) dh_long[-1] = np.ones(16) # Track gradient norms at each time step during BPTT grad_norms_per_step = [] dh_next = np.zeros(16) for t in reversed(range(50)): dh = dh_long[t] + dh_next h_t = rnn_big.cache['h'][t + 1] dz = dh * (1 - h_t ** 2) dh_next = rnn_big.W_hh.T @ dz grad_norms_per_step.append(np.linalg.norm(dh_next)) grad_norms_per_step.reverse() # now indexed by time step fig, ax = plt.subplots(figsize=(8, 5)) ax.semilogy(range(50), grad_norms_per_step, 'o-', color='darkorange', markersize=3) ax.set_xlabel('Time step (backward from T=50)') ax.set_ylabel('Gradient norm (log scale)') ax.set_title('Exploding Gradients: large W_hh causes exponential gradient growth') ax.grid(True, alpha=0.3) plt.tight_layout() plt.show() spectral_norm = np.linalg.norm(rnn_big.W_hh, ord=2) print(f'Spectral norm of W_hh: {spectral_norm:.4f}') print(f'Gradient norm at step 0: {grad_norms_per_step[0]:.2e}') print(f'Gradient norm at step 49: {grad_norms_per_step[49]:.2e}') print(f'Ratio: {grad_norms_per_step[0] / (grad_norms_per_step[49] + 1e-30):.2e}')

Spectral norm of W_hh: 9.0217
Gradient norm at step 0: 5.43e+06
Gradient norm at step 49: 4.52e+00
Ratio: 1.20e+06

6.3 Failure: Sequential bottleneck¶

RNNs process sequences step-by-step — each $h_t$ depends on $h_{t-1}$. This means:

1. No parallelism across time — unlike CNNs or Transformers.
2. Fixed-size bottleneck — all history must fit in $h_t \in \mathbb{R}^m$.

This is why Transformers (which attend to all positions in parallel) have largely replaced RNNs for sequence modeling tasks.

7. Connections¶

From	To	Relationship
MLP (topic 13)	RNN	RNN = MLP with weight sharing across time
Gradient Descent (topic 02)	BPTT	BPTT = backprop applied to unrolled graph
CNN (topic 14)	RNN	Spatial vs temporal parameter sharing
RNN/LSTM	Transformer (topic 16)	Attention replaces recurrence for parallelism
Vanishing gradients	LSTM gates	Additive cell update creates gradient highway

Takeaway¶

Recurrent networks handle variable-length sequences by sharing parameters across time and maintaining a hidden state. The vanilla RNN suffers from vanishing/exploding gradients due to repeated Jacobian products in BPTT. LSTM solves this with gated additive cell updates that act as gradient highways. Despite this, RNNs are inherently sequential and have been largely superseded by Transformers, which use attention for direct access to all positions with full parallelism.

Next: 16 Transformer — attention is all you need.