Module 5 — LSTM: forget, input and output gates

Module 4 ended with the diagnosis: a simple RNN cannot preserve information across hundreds of time steps because its gradient is a product of Jacobians. The long short-term memory cell (Hochreiter and Schmidhuber, 1997) is the architectural answer that has powered speech recognition, machine translation and most sequence work up to the Transformer era. This module reads the cell equation by equation, explains why the gradient survives, and shows what the extra parameters buy on the electricity forecasting task.

Two states, not one

An LSTM cell carries two states from one step to the next:

• The cell state $c_t$: a vector that acts as long-term memory, updated by mostly additive operations.
• The hidden state $h_t$: the actual output of the cell, computed from $c_t$ through the output gate.

The cell state is the crucial addition. Its update law is:

$$c_t = f_t \odot c_{t-1} + i_t \odot \tilde{c}_t$$

where $f_t$ and $i_t$ are gates in $[0, 1]$ and $\odot$ is the element-wise product. Notice what is not there: no multiplication by a recurrent weight matrix. The cell state flows from $t-1$ to $t$ through a mostly additive path, which means the Jacobian $\partial c_t / \partial c_{t-1}$ is roughly $\operatorname{diag}(f_t)$. As long as $f_t$ stays close to 1 on the coordinates that matter, the gradient survives across hundreds of steps.

The three gates

Each gate is a small sigmoidal network that reads $x_t$ and $h_{t-1}$ and outputs a vector in $[0, 1]^H$:

$$f_t = \sigma(W_f x_t + U_f h_{t-1} + b_f) \quad \text{(forget gate)}$$ $$i_t = \sigma(W_i x_t + U_i h_{t-1} + b_i) \quad \text{(input gate)}$$ $$o_t = \sigma(W_o x_t + U_o h_{t-1} + b_o) \quad \text{(output gate)}$$

And the candidate cell content:

$$\tilde{c}_t = \tanh(W_c x_t + U_c h_{t-1} + b_c)$$

Finally, the hidden state:

$$h_t = o_t \odot \tanh(c_t)$$

Reading this in words:

• The forget gate decides, per coordinate, which parts of the previous cell state to keep. A value of 1 means "keep everything on this coordinate", 0 means "erase it".
• The input gate decides which parts of the new candidate $\tilde{c}_t$ to write into the cell.
• The output gate decides which parts of the cell state to expose as the hidden state (and therefore as the layer's output, and as input to the next step).

The gates are learned like everything else, but they are gates in the electrical sense: they let signal pass or block it, without deforming it linearly the way a matrix multiplication would.

Why gradient survives: the additive highway

Take the derivative of $c_t$ with respect to $c_{t-1}$:

$$\frac{\partial c_t}{\partial c_{t-1}} = \operatorname{diag}(f_t) + \text{terms from } i_t, \tilde{c}_t, h_{t-1}$$

The dominant term is $\operatorname{diag}(f_t)$. If $f_t$ is close to 1, the product $\prod_{i} \partial c_i / \partial c_{i-1}$ stays close to 1 too, and the gradient does not vanish. If $f_t$ is close to 0, information is deliberately forgotten and the gradient stops there, which is the correct behaviour — nothing depends on a state that has been discarded.

This is why the forget gate bias is often initialised to 1 or 2: a positive bias makes $\sigma(b_f) \approx 0.73$ or $0.88$ at the start of training, biasing the network toward remembering. Keras does this automatically with unit_forget_bias=True, which is the default.

Parameter count

For an LSTM with input size $D$ and hidden size $H$, each of the four internal operations (three gates plus the candidate) has:

• an input matrix of shape $(D, H)$
• a recurrent matrix of shape $(H, H)$
• a bias of shape $(H,)$

Total per LSTM cell:

$$4 \times (D H + H^2 + H) = 4 H (D + H + 1)$$

For the electricity series with $D = 1$ and $H = 32$: $4 \times 32 \times (1 + 32 + 1) = 4\,352$ parameters, four times more than a SimpleRNN(32). This is the price of the gates, and it is well spent for any sequence longer than a few dozen steps.

from tensorflow.keras import layers, Sequential

model = Sequential([
    layers.Input(shape=(168, 1)),
    layers.LSTM(32),
    layers.Dense(24),
])
model.summary()
# LSTM: 4 * 32 * (1 + 32 + 1) = 4 352 parameters

LSTM on the electricity series

Replacing SimpleRNN with LSTM in the module 2 model is a one-line change, and it usually improves the 24-hour forecast noticeably:

import numpy as np
from tensorflow.keras import layers, Sequential, callbacks

series = np.loadtxt("electricity_hourly.csv").astype("float32")
mu, sigma = series[:6000].mean(), series[:6000].std()
series_s = (series - mu) / sigma

def make_xy(s, lookback=168, horizon=24):
    x, y = [], []
    for t in range(len(s) - lookback - horizon + 1):
        x.append(s[t:t + lookback])
        y.append(s[t + lookback:t + lookback + horizon])
    return np.array(x)[..., None], np.array(y)

x_train, y_train = make_xy(series_s[:6000])
x_val,   y_val   = make_xy(series_s[6000 - 168:])

model = Sequential([
    layers.Input(shape=(168, 1)),
    layers.LSTM(32),
    layers.Dense(24),
])
model.compile(optimizer="adam", loss="mse")
model.fit(
    x_train, y_train,
    validation_data=(x_val, y_val),
    epochs=20, batch_size=64,
    callbacks=[callbacks.EarlyStopping(patience=3, restore_best_weights=True)],
    verbose=2,
)

On a typical office building, this LSTM improves validation MAE by 15 to 30 % over the SimpleRNN of module 2, mostly because it can carry the weekend signal across five weekdays without letting the tanh saturate.

return_state and manual initialisation

Sometimes a downstream architecture needs the final states — the encoder-decoder of module 8 is the paradigmatic example. Setting return_state=True gives back both:

lstm = layers.LSTM(64, return_state=True, return_sequences=True)
sequence_output, final_h, final_c = lstm(x)

Passing an initial state is symmetric: layers.LSTM(64)(x, initial_state=[h0, c0]). This is how a decoder receives an encoder's context.

What the LSTM does not fix

Two limitations survive the LSTM:

• Memory still scales linearly with sequence length (module 3). The gates are cheap in FLOPs but the intermediate activations still have to be stored for BPTT.
• Very long dependencies (thousands of steps) remain hard even with LSTMs, because the forget gate has to stay open on the right coordinates for a very long time. This is one of the reasons Transformers eventually replaced RNNs on long-context tasks.

For sequences up to a few hundred steps — most tabular time series, most sensor logs, all of the electricity forecasting problem — LSTMs remain the strong default.

A gate saturated at 0 or 1 stops learning

If a forget gate is always 0 or always 1, its gradient is dead. Monitoring the histogram of gate activations during training (via TensorBoard, see course 08) is the fastest way to spot a saturated LSTM before it wastes an entire training budget.

Use LSTM(unit_forget_bias=True) and leave it alone

The forget-gate bias initialised to 1 is one of the few defaults you should not touch. It comes from Gers et al. (2000) and it is worth several epochs of warm-up.

In summary

• An LSTM carries two states: a cell state $c_t$ that flows additively (memory) and a hidden state $h_t$ that is the visible output.
• Three gates — forget, input, output — control what is kept, what is written and what is exposed; each is a sigmoid layer over $x_t$ and $h_{t-1}$.
• The additive cell update gives the gradient an almost-identity Jacobian when the forget gate stays near 1, which is why LSTMs handle long dependencies.
• Parameter count is $4 H (D + H + 1)$, four times a SimpleRNN; that is the price of the gates and it is almost always worth paying.

Next module: the GRU, which merges some of these gates into a smaller cell with comparable performance and 25 % fewer parameters.