LSTM Gradient Flow › Gradient Flow Through Addition vs. Multiplication

To deeply understand why the cell state path is so much better for gradient flow, we need to compare how addition and multiplication behave during backpropagation. This comparison is the core insight behind LSTMs, ResNets, Highway Networks, and many other modern architectures.

Consider a simple computation: $z = x + y$. The partial derivatives are $\partial z / \partial x = 1$ and $\partial z / \partial y = 1$. During backpropagation, if we receive an upstream gradient $\partial L / \partial z$, then:

$$\frac{\partial L}{\partial x} = \frac{\partial L}{\partial z} \cdot \frac{\partial z}{\partial x} = \frac{\partial L}{\partial z} \cdot 1 = \frac{\partial L}{\partial z}$$

Addition is a gradient router: it copies the incoming gradient unchanged to both inputs. This means that in a long chain of additions, the gradient at the start of the chain equals the gradient at the end. No decay, no amplification.

This property has a direct implication. If you can route a gradient signal through an all-addition path from a loss function back to an early parameter, that parameter receives the same magnitude of gradient signal regardless of how many "steps" it took to get there. The distance in time (or depth) becomes irrelevant. The gradient doesn't know how far it has traveled through an additive path.

Now consider $z = f \cdot c$ where $f$ is a learnable gate value. The partial derivative with respect to $c$ is $\partial z / \partial c = f$. During backpropagation:

$$\frac{\partial L}{\partial c} = \frac{\partial L}{\partial z} \cdot f$$

A single element-wise multiplication scales the gradient by $f$. Chaining $n$ such multiplications:

$$\frac{\partial L}{\partial c_0} = \frac{\partial L}{\partial c_n} \cdot \prod_{t=1}^{n} f_t$$

If each $f_t$ is 0.9, then after 50 timesteps: $0.9^{50} \approx 0.005$. After 100 steps: $0.9^{100} \approx 0.00003$. The gradient has decayed to essentially nothing. But if each $f_t$ is close to 1, say 0.99, then $0.99^{100} \approx 0.37$: still 37% of the original signal survives.

This is the key contrast with vanilla RNNs. In a vanilla RNN, the multiplicative factor at each step is $\text{diag}(\tanh'(\mathbf{z}_t)) \cdot \mathbf{W}_{hh}$, whose spectral norm is typically well below 1. In an LSTM, the multiplicative factor along the cell state path is just $\mathbf{f}_t$, a vector of sigmoid outputs. The LSTM can be trained to keep these values near 1 for dimensions where long-term memory is needed.

The main difference lies in who controls the multiplier, rather than in its magnitude. In the vanilla RNN, the gradient multiplier is determined by the weight matrix and the input at each step. The network has no way to "decide" to preserve gradients in certain dimensions; the gradient multiplier is the byproduct of the same weights that compute the hidden state transformation. In the LSTM, the forget gate is a learned, dedicated mechanism for controlling gradient flow. The network can separately learn "how to transform information" (through the candidate cell state) and "how to protect gradient flow" (through the forget gate). This separation of concerns is architecturally elegant and practically powerful.

Let's derive the full gradient through the cell state more carefully. The partial derivative of $\mathbf{c}_t$ with respect to $\mathbf{c}_{t-1}$ is:

$$\frac{\partial \mathbf{c}_t}{\partial \mathbf{c}_{t-1}} = \text{diag}(\mathbf{f}_t) + \text{diag}(\mathbf{i}_t) \frac{\partial \tilde{\mathbf{c}}_t}{\partial \mathbf{c}_{t-1}} + \text{diag}(\mathbf{c}_{t-1}) \frac{\partial \mathbf{f}_t}{\partial \mathbf{c}_{t-1}} + \text{diag}(\tilde{\mathbf{c}}_t) \frac{\partial \mathbf{i}_t}{\partial \mathbf{c}_{t-1}}$$

This looks complex, but the dominant term in practice is $\text{diag}(\mathbf{f}_t)$. The other terms involve derivatives of the gates with respect to the cell state, which in standard LSTMs (without peephole connections) are zero: the gates depend on $\mathbf{h}_{t-1}$ and $\mathbf{x}_t$, not directly on $\mathbf{c}_{t-1}$.

For a standard LSTM:

$$\frac{\partial \mathbf{c}_t}{\partial \mathbf{c}_{t-1}} = \text{diag}(\mathbf{f}_t) + \text{diag}(\mathbf{i}_t) \frac{\partial \tilde{\mathbf{c}}_t}{\partial \mathbf{c}_{t-1}}$$

Since $\tilde{\mathbf{c}}_t = \tanh(\mathbf{W}_c \mathbf{x}_t + \mathbf{U}_c \mathbf{h}_{t-1} + \mathbf{b}_c)$ and $\mathbf{h}_{t-1} = \mathbf{o}_{t-1} \odot \tanh(\mathbf{c}_{t-1})$, there is an indirect path through $\mathbf{h}_{t-1}$. But this path still passes through a tanh nonlinearity and an additional gate multiplication, making it substantially smaller than the direct $\mathbf{f}_t$ path.

The primary gradient flow through the cell state is thus:

$$\frac{\partial L}{\partial \mathbf{c}_k} \approx \frac{\partial L}{\partial \mathbf{c}_T} \odot \prod_{t=k+1}^{T} \mathbf{f}_t$$

This approximation captures the essence of the constant error carousel.