LSTM and GRU: gates that remember

01 · Concept

Concept

Lesson 3.2 left a specific diagnosis, not a vague complaint. The plain RNN’s problem is that its state passes through the same multiplication and the same saturating nonlinearity at every step, so credit across $k$ steps is a product of $k$ Jacobians whose factors the model does not control. Two repairs follow from that diagnosis directly. Give the state a path that is added to rather than transformed. And let the model decide, per step and per feature, how much of the old state to keep.

Both repairs are in the LSTM, and both are visible in one equation. Alongside the hidden output $h_t$, the cell maintains a separate cell state $c_t$, updated as

followed by $h_t=o_t\odot\tanh(c_t)$. Here $f_t$, $i_t$, and $o_t$ are the forget, input, and output gates: sigmoid outputs between zero and one, computed from $x_t$ and $h_{t-1}$ through learned affine maps, and multiplied elementwise. The candidate $\tilde c_t$ is the new content on offer.

Read the cell update as an instruction rather than a formula. The forget gate says, coordinate by coordinate, what fraction of the existing ledger survives. The input gate says how much of the new candidate gets written. The two are added, not composed. A coordinate whose forget gate sits at one and whose input gate sits at zero carries its value forward completely untouched — no matrix multiply, no nonlinearity, nothing.

That changes the gradient path in exactly the way lesson 3.2 asked for. Holding the gate values fixed, the direct cell-state path from $c_k$ to $c_T$ contributes $\prod_{t=k+1}^{T} f_t$: the product of the forget gates. The weight matrices and saturating derivative are absent from this direct route, and what remains is a quantity the network computes for itself. This product is not the whole total derivative, however: the gates and candidate also depend on the previous hidden state, creating additional gradient paths.

Now work out what that product actually buys, because this is where honest accounting matters. Suppose a coordinate’s forget gate settles at $0.99$. Over 100 steps the retained fraction is $0.99^{100}\approx0.366$ — about a third of the signal survives, which is entirely workable. Drop the gate to $0.9$ and the same 100 steps give $0.9^{100}\approx2.7\times10^{-5}$: gone. The gate is not a switch between remember and forget; it sets an exponential half-life, and small differences in the gate value produce enormous differences in reach.

This makes initialization consequential in a way that surprises people. Leave the forget-gate bias at zero and the sigmoid starts near $0.5$, so at the beginning of training the cell retains $0.5^{100}\approx8\times10^{-31}$ over a hundred steps — a half-life of one step. The model must first learn to remember before it can learn what to remember, and the gradient that would teach it has already vanished. The standard fix is to initialize the forget-gate bias positive: at a bias of 3 the gate starts near $0.953$, and $0.953^{100}\approx0.008$, which is small but no longer astronomically dead, so the long-range gradient exists to be improved. The wrong turn here is diagnosing the resulting failure as a learning-rate or capacity problem and tuning for days; the real cause is a bias initialization that made the memory path start in the off position.

A GRU simplifies the design by merging cell and hidden state. An update gate interpolates between the previous state and a candidate, and a reset gate controls how much of the previous state contributes while the candidate is formed. Fewer gates, fewer parameters, often comparable quality. Exact equations vary by convention — notably whether the reset is applied before or after a matrix multiplication — and those variants are not numerically equivalent, so checkpoint compatibility depends on the precise definition rather than the name.

Gate visualizations invite a second error worth naming. A coordinate showing a high forget value on a particular token does not prove that unit “stores” a human-nameable fact. State features are distributed, gates interact across thousands of coordinates, and a downstream projection can amplify or discard whatever was preserved. Causal intervention — clamp the coordinate, measure the behavioral change — is evidence; a colorful heatmap is a hypothesis.

To compare recurrent cells honestly, match total parameter count by adjusting hidden width rather than holding hidden size fixed — equal hidden size flatters the smaller GRU on parameter cost, while equal parameters changes the width and therefore the representation. Report latency and memory at the sequence lengths you actually care about, and test streaming state resets explicitly: state carried across unrelated sequences produces impressive-looking numbers that are leakage rather than memory.