LSTM and GRU cells, and beam search, from scratch

Why gates

A plain recurrent network updates its state as h′ = tanh(Wx + Uh + b). Back-propagating through many steps multiplies many Jacobians together, and the gradient either shrinks to nothing or blows up. Gated cells add a path along which information is copied rather than transformed, so the gradient can travel a long way unchanged.

The LSTM step

An LSTM keeps two vectors: the cell state c, which carries memory, and the hidden state h, which is the output. One matrix multiplication on the concatenation [x, h] produces four blocks of pre-activations:

z = W @ np.concatenate([x, h]) + b          # shape (4H,)
i, f, g, o = sigmoid(z[:H]), sigmoid(z[H:2*H]), np.tanh(z[2*H:3*H]), sigmoid(z[3*H:])
c_new = f * c + i * g                       # forget part of the old memory, write some new
h_new = o * np.tanh(c_new)                  # expose part of it

The line f * c is the copy path: when the forget gate is near 1, the cell state, and its gradient, pass through almost untouched. Libraries differ in the order of the four blocks. PyTorch uses input, forget, candidate, output; others put the forget gate first. Getting the order wrong gives code that runs and is wrong.

The GRU step, as PyTorch computes it

A GRU merges the two states into one and uses two gates. PyTorch keeps separate weights for the input and the hidden state, with gates in the order reset, update, new, and applies the reset gate after the hidden-state matrix multiplication:

r = sigmoid(Wir x + bir + Whr h + bhr)
z = sigmoid(Wiz x + biz + Whz h + bhz)
n = np.tanh(Win x + bin + r * (Whn h + bhn))     # r scales the result, not h
h_new = (1 - z) * n + z * h

The original paper applies r to h before the multiplication. The two are not equal, and an interviewer who asks you to match a library expects the library's version.

A sigmoid that never overflows

1 / (1 + exp(−x)) overflows for large negative x. Use exp(x) / (1 + exp(x)) when x is negative and the usual form when it is not; each branch only ever exponentiates a non-positive number.

Decoding: greedy and beam search

Generating a sequence means choosing a token at each step. Greedy decoding takes the most likely token every time and can miss a better sequence that starts with a less likely token. Beam search keeps the B best partial sequences, extends each by every token, and keeps the B best of the extensions. Width 1 is greedy decoding; a width as large as the number of sequences is exhaustive search. Scores are sums of log-probabilities, so longer sequences always score lower, which is why production decoders divide by a length penalty.

What interviewers listen for

• The gate order stated before any code, with the library it matches.
• A stable sigmoid, and the reason it is needed.
• Log-probabilities added, never probabilities multiplied.
• A deterministic rule for ties among beams.