LSTM Cell - Sequence Models

LSTM Cell

What you'll build

One time step of an LSTM, matching PyTorch's nn.LSTMCell exactly (same weight layout, same gate order).

Background

The LSTM (Hochreiter & Schmidhuber, 1997) fixes the vanilla RNN's vanishing gradients with a separate cell state c that is updated additively and controlled by gates. When the forget gate is close to 1, dc_t / dc_{t-1} ≈ 1, so gradients flow back through c almost unchanged over hundreds of steps.

Formula

gates = x @ W_ih.T + b_ih + h_prev @ W_hh.T + b_hh      # (N, 4H)
i, f, g, o = split(gates, 4, axis=-1)                   # each (N, H), in THIS order

i = sigmoid(i)        # input gate:  how much new content to write
f = sigmoid(f)        # forget gate: how much of the old cell to keep
g = tanh(g)           # candidate cell content, in (-1, 1)
o = sigmoid(o)        # output gate: how much of the cell to expose

c = f * c_prev + i * g
h = o * tanh(c)

Weight layout and gate order: i, f, g, o

The four gates are stacked along the first axis of each weight, in the order i, f, g, o, exactly like nn.LSTMCell.weight_ih / weight_hh:

Rows of W_ih, W_hh (and entries of the biases)	Gate
0 : H	input gate i
H : 2H	forget gate f
2H : 3H	candidate g
3H : 4H	output gate o

Weights are (out, in) (PyTorch convention), hence x @ W_ih.T. Keras also orders the gates i, f, c, o but stores kernels as (D, 4H).

Inputs

Tensor	Shape
x	(N, D)
h_prev, c_prev	(N, H)
W_ih	(4H, D)
W_hh	(4H, H)
b_ih, b_hh	(4H,)

Expected return format

A tuple (h, c): the new hidden state and the new cell state, each a float array of shape (N, H), in that order.

Edge cases and conventions

• Both biases are added (PyTorch keeps two; they are redundant but part of the layout).
• A sigmoid helper is provided. h always lies in (-1, 1); c is unbounded.

Constraints

NumPy only; vectorised over the batch (one matmul per weight, not one per gate).

Follow-up

• How many parameters does nn.LSTMCell(D, H) have? (4(HD + H^2 + 2H).)
• Why is the forget-gate bias often initialised to 1? What do peephole connections add?

Input: x = [[0.3, -0.1]], h_prev = [[0, 0]], c_prev = [[1.0, -2.0]], all weights and biases zero
Output: (array([[0.231059, -0.380797]]), array([[0.5, -1.0]]))

Input: D = H = 1, x = [[1.0]], h_prev = [[0.0]], c_prev = [[1.0]], W_ih = [[0], [2], [1], [0]], W_hh = 0, biases 0
Output: (array([[0.425752]]), array([[1.261594]]))

import numpy as np

def sigmoid(z):
    return 1 / (1 + np.exp(-z))


def lstm_cell(x, h_prev, c_prev, W_ih, W_hh, b_ih, b_hh):
    """
    One LSTM step (gate order i, f, g, o, as in nn.LSTMCell).

    Args:
        x: Input (N, D)
        h_prev: Previous hidden state (N, H)
        c_prev: Previous cell state (N, H)
        W_ih: Input weights (4H, D), row blocks [i; f; g; o]
        W_hh: Hidden weights (4H, H), row blocks [i; f; g; o]
        b_ih, b_hh: Biases (4H,)

    Returns:
        (h, c): new hidden and cell states, each (N, H)
    """
    # Your code here
    pass