LSTM Gradient Flow › Worked Numerical Example: Tracing a Gradient Through the Cell State

Abstract derivations are easier to internalize when paired with a concrete example. Let's trace a gradient backward through a small LSTM for exactly three timesteps, keeping all numbers visible.

Consider an LSTM with a single hidden dimension ($h = 1$) processing a sequence of three timesteps. This is a toy model, but it captures the essential structure. At each step, the gate values are scalars rather than vectors.

Suppose the learned (or assumed for this example) gate values are:

• Timestep 1: $f_1 = 0.95$, $i_1 = 0.3$, $o_1 = 0.8$
• Timestep 2: $f_2 = 0.90$, $i_2 = 0.5$, $o_2 = 0.7$
• Timestep 3: $f_3 = 0.85$, $i_3 = 0.4$, $o_3 = 0.9$

And the candidate cell state values (the $\tanh$ of the gate pre-activation) are:

• Timestep 1: $\tilde{c}_1 = 0.6$
• Timestep 2: $\tilde{c}_2 = -0.3$
• Timestep 3: $\tilde{c}_3 = 0.7$

We start with $c_0 = 0$.

The cell state evolves as $c_t = f_t \cdot c_{t-1} + i_t \cdot \tilde{c}_t$:

Step 1: $c_1 = 0.95 \cdot 0 + 0.3 \cdot 0.6 = 0.18$

Step 2: $c_2 = 0.90 \cdot 0.18 + 0.5 \cdot (-0.3) = 0.162 - 0.15 = 0.012$

Step 3: $c_3 = 0.85 \cdot 0.012 + 0.4 \cdot 0.7 = 0.0102 + 0.28 = 0.2902$

The cell state carries information across time, here starting at 0, rising to 0.18, dipping to 0.012 (the negative candidate overwriting most of the stored value), then jumping to 0.29.

Suppose the loss at timestep 3 has gradient $\partial L / \partial c_3 = 1.0$ with respect to the cell state.

We ask: what is the gradient with respect to $c_0$? In the cell state path, this is:

$$\frac{\partial L}{\partial c_0} = \frac{\partial L}{\partial c_3} \cdot f_3 \cdot f_2 \cdot f_1 = 1.0 \cdot 0.85 \cdot 0.90 \cdot 0.95 = 0.727$$

Nearly 73% of the original gradient reaches $c_0$ through the cell state path. Now compare what would happen in a vanilla RNN. The multiplicative factor at each step would be $\tanh'(z_t) \cdot w_{hh}$. For typical values, suppose these factors are $0.7$, $0.6$, $0.65$. Then:

$$\frac{\partial L}{\partial h_0} \approx 1.0 \cdot 0.65 \cdot 0.60 \cdot 0.70 = 0.273$$

Only 27% survives after three steps, and the decay rate is steeper per step. Extrapolating to 30 steps: if the RNN decay factor per step averages $0.64$ and the LSTM forget gate product per step averages $0.90$, then after 30 steps the RNN retains $0.64^{30} \approx 10^{-5}$ of the gradient while the LSTM cell retains $0.90^{30} \approx 0.042$. That is roughly a 4,000-fold difference in gradient magnitude.

The gradient $\partial L / \partial c_0$ tells the optimizer how to adjust anything that influenced $c_0$ in order to reduce the loss. Since $c_0 = 0$ by convention (the initial cell state), this gradient acts on the initial state itself. In practice, many implementations make $c_0$ trainable: the optimizer can learn a better starting cell state for the task.

More importantly, the cell state gradient also flows back through the input gate and candidate cell state at step 1 to reach the weights that processed the very first input. This is how the LSTM achieves long-range credit assignment: the cell state acts as a near-lossless wire carrying the loss signal backward to wherever it needs to go.