Explain Each LSTM Gate With a Concrete Example and Diagnose a Forgetful Model
An LSTM-based code-completion model must track "am I currently inside a string literal, and which quote character opened it?" across lines of code. After training it handles 10-token strings but forgets the quote type in 100-token strings. You inspect it and find that the forget gate averages 0.55 on the memory unit that carries the quote type.
- Write the LSTM equations and explain the role of the forget, input and output gates and of the candidate, using the string-literal unit as a running example. What should each gate do on the opening quote, on ordinary characters inside the string, and on the closing quote?
- Why does the average forget value of 0.55 explain the failure? Compute the surviving fraction of the memory after 10 and after 100 steps at 0.55 and at 0.99.
- Why is the cell-state path described as a "gradient highway", and what does the 0.55 imply about the gradient for this unit?
- Give two concrete changes that you would make before collecting more data. Then say how the GRU relates and whether you would switch to it.
1. Equations and roles.
For the string-literal unit (one coordinate of c):
| Moment | Forget f | Input i and candidate \tilde c | Output o |
|---|---|---|---|
Opening quote " |
low: erase the previous value ("not in a string") | high i, \tilde c \approx +1 for a double quote (write "in a double-quoted string") | may be low or high, the next layer needs it soon |
| Ordinary characters inside the string | about 1: keep the memory | about 0: nothing new to write | high whenever the next-token prediction needs it (for example when deciding whether " closes the string or ' is literal text) |
| Closing quote | low: erase the memory | high i, \tilde c \approx 0 or -1 ("not in a string") | high so downstream layers see the change |
The forget gate decides how much of the old memory survives, the input gate how much of the candidate is written, the candidate what is written, and the output gate how much of the memory is exposed in h_t (the cell can hold information the layers above never see until it matters).
2. Why 0.55 explains it. Within a string the forget gate should be near 1. At 0.55 the memory is multiplied by 0.55 each step: 0.55^{10} = 0.0025, 0.55^{100} \approx 10^{-26}, so the quote type is gone within ten or so characters even though the gate "works" on short strings. At 0.99, 0.99^{10} = 0.90 and 0.99^{100} = 0.37: still a usable signal after 100 steps. The model has learned that remembering matters only for the short cases it saw most; nothing in the loss pushed the gate higher for long ones, and the gradient to do so is tiny (next part).
3. Gradient highway. The direct path of the gradient along the cell is \partial c_t/\partial c_{t-1} = \operatorname{diag}(f_t): no weight matrix and no activation derivative, only the gate. If f \approx 1 the error signal of a late step passes back almost unchanged (the same role as the identity in a residual connection). At f = 0.55 the gradient shrinks by 0.55 per step, 0.55^{100} \approx 10^{-26}, so there is no signal telling the model it should have remembered: a self-reinforcing failure.
4. Changes before more data.
- Forget-gate bias initialization to about 1 (or 2-3 for very long dependencies) so the cell remembers by default; sigmoid(3) = 0.95. Re-train from this initialization. This is cheap and often the biggest single fix.
- Training sequence length and truncation: make sure the truncated-BPTT window is longer than the dependency (at least 100 steps) and include long strings in training (curriculum), otherwise no gradient ever asks for long memory. Also clip gradients and check that the data contain enough long strings of both quote types.
- If long-range structure matters a lot, add attention (or move to a Transformer): an LSTM's reliable range is on the order of tens to a few hundred steps.
The GRU has one update gate z_t and a reset gate, merging the cell and hidden state: h_t = (1 - z_t) \odot n_t + z_t \odot h_{t-1}, the same additive highway with 75% of the LSTM's parameters. It would not by itself fix a forget-gate initialization problem (a GRU's update gate also starts near 0.5), so switching is an experiment, not the remedy; I would first fix initialization and the training window.
Share this question