LSTM Cell Step

Carry two state vectors through one step

An LSTM cell updates a hidden state h and a cell state c. Both have shape (H,), but they play different roles: the cell state accumulates information, while the hidden state is a gated, squashed view of it. The next step receives both returned vectors.

Implement one step for input x of shape (D,). This exercise uses explicit Linen parameters, with no batch axis, sequence loop, peephole connections or hidden-state projection.

Compute three gates and a candidate

Concatenate x before h_prev and multiply from the right:

xh = concatenate([x, h_prev])       # (D + H,)
f  = sigmoid(xh @ W_f + b_f)        # forget gate (H,)
i  = sigmoid(xh @ W_i + b_i)        # input gate  (H,)
o  = sigmoid(xh @ W_o + b_o)        # output gate (H,)
g  = tanh(xh @ W_g + b_g)           # candidate   (H,)

c_new = f * c_prev + i * g
h_new = o * tanh(c_new)

The sigmoid gates choose how much to retain, write and expose. The candidate uses tanh, so it can add positive or negative values. Every multiplication between state or gate vectors is elementwise.

Use the supplied c_prev; it need not be zero. Compute c_new before h_new, and apply tanh to the new cell state. Unlike the hidden state, the cell state is not restricted to the interval from -1 to 1.

What the memory path does and does not guarantee

The term f * c_prev gives the cell state a direct path to the next step. Holding x and h_prev fixed, the derivative of c_new with respect to c_prev is a diagonal matrix with f on the diagonal. A forget gate near one can preserve information along this path.

That is not a guarantee against vanishing gradients. For example, a fixed forget gate of 0.5 multiplies the direct path by 0.5 at each step, giving 0.5 ** N after N steps. In a complete recurrence, dependencies through the hidden state and gates also contribute to gradients. The cell update can retain, decay or accumulate values; it does not necessarily leave memory nearly unchanged.

Keep the parameter shapes explicit

Each of W_f, W_i, W_o and W_g has shape (D + H, H). Each corresponding bias has shape (H,). These are eight parameter arrays, containing 4 * ((D + H) * H + H) scalar parameters. For D=3 and H=4, the count is 128.

Concatenation is a storage choice. Splitting a weight into input and hidden rows gives the same affine map:

xh @ W + b = x @ W[:D] + h_prev @ W[D:] + b

Likewise, packing the four matrices side by side can compute the four projections together if you copy the same weights and split in the same order. Freshly initializing a packed matrix does not reproduce this exercise's four separately named random matrices.

A step with known parameters

Set all matrices to zero. Zero gate biases make f, i and o equal to 0.5; choose a candidate bias that makes g=0.5:

import jax
import jax.numpy as jnp

x = jnp.array([1.])
h_prev = jnp.array([.2, -.1])
c_prev = jnp.array([2., -2.])
xh = jnp.concatenate([x, h_prev])
W = jnp.zeros((3, 2))
gate_bias = jnp.zeros(2)
candidate_bias = jnp.full(2, jnp.arctanh(.5))
f = jax.nn.sigmoid(xh @ W + gate_bias)
i = jax.nn.sigmoid(xh @ W + gate_bias)
o = jax.nn.sigmoid(xh @ W + gate_bias)
g = jnp.tanh(xh @ W + candidate_bias)
c_new = f * c_prev + i * g
h_new = o * jnp.tanh(c_new)
result = jnp.concatenate([h_new, c_new])

assert jnp.allclose(c_new, jnp.array([1.25, -.75]))
assert jnp.allclose(h_new, .5 * jnp.tanh(jnp.array([1.25, -.75])))
assert result.shape == (4,)
assert jnp.array_equal(result[2:], c_new)

The equal matrices in this example make the arithmetic simple. The exercise still creates four independently named weight/bias pairs.