LSTM Cell Forward Pass

Run an LSTM over a sequence.

At each step the input and previous hidden state produce four gate pre-activations at once:

z = x_t @ W_x + h_(t-1) @ W_h + b        shape (4 * d_hidden,)

Split z into four equal blocks in the order [i, f, g, o], matching PyTorch:

i = sigmoid(z_i)     input gate, how much of the candidate to write
f = sigmoid(z_f)     forget gate, how much of the cell state to keep
g = tanh(z_g)        candidate values
o = sigmoid(z_o)     output gate, how much of the cell to expose

Then update:

c_t = f * c_(t-1) + i * g
h_t = o * tanh(c_t)

Return (hidden_states, c_final) where hidden_states has shape (seq_len, d_hidden) and c_final has shape (d_hidden,).

The cell state is the whole point. c_t is updated by addition rather than by a matrix multiply and a squashing function, so when the forget gate is near one the gradient flows back through many steps almost unchanged. A vanilla RNN multiplies by W_h at every step instead, which is why it cannot remember.

Use a numerically stable sigmoid.

Function signature

def lstm_forward(x, h0, c0, W_x, W_h, b):

Example

Input

x  = np.zeros((2, 1))                  # 2 steps, d_in = 1
h0 = np.zeros(1)
c0 = np.zeros(1)
W_x = np.zeros((1, 4))
W_h = np.zeros((1, 4))
b  = np.array([0.0, 0.0, 1.0, 0.0])    # only the g block is non-zero
lstm_forward(x, h0, c0, W_x, W_h, b)

Output

(array([[0.18169974],
       [0.2581184 ]]), array([0.57119562]))

Reasoning

Every gate pre-activation is 0 except g, so i = f = o = 0.5 and g = tanh(1) = 0.7616. Step 1: c = 0.5 * 0.7616 = 0.3808 and h = 0.5 * tanh(0.3808) = 0.1817. Step 2: c = 0.5 * 0.3808 + 0.3808 = 0.5712 and h = 0.2582.

Constraints and notes

• Split the pre-activations in the order [i, f, g, o].
• hidden_states has shape (seq_len, d_hidden).
• Use a stable sigmoid.