LSTM Gradient Flow › The Residual Connection Analogy

Understanding the LSTM cell state gradient path is greatly aided by comparing it to residual networks (ResNets), introduced by He et al. (2016). The two architectures solve the same fundamental problem using the same mathematical trick: creating an additive shortcut for gradient flow.

In a standard deep network, each layer transforms its input: $\mathbf{x}_{l+1} = F(\mathbf{x}_l)$ where $F$ includes weight matrices and nonlinearities. During backpropagation, the gradient through $n$ such layers involves a product of $n$ Jacobians, which can vanish.

ResNets add a skip connection: $\mathbf{x}_{l+1} = F(\mathbf{x}_l) + \mathbf{x}_l$. Now the gradient is:

$$\frac{\partial \mathbf{x}_{l+1}}{\partial \mathbf{x}_l} = \frac{\partial F(\mathbf{x}_l)}{\partial \mathbf{x}_l} + \mathbf{I}$$

The identity matrix $\mathbf{I}$ ensures that even if the network learns $F(\mathbf{x}_l) \approx 0$ (effectively turning off the learned transformation), the gradient still flows backward with magnitude 1. The addition creates a guaranteed gradient highway.

The LSTM cell state update:

$$\mathbf{c}_t = \mathbf{f}_t \odot \mathbf{c}_{t-1} + \mathbf{i}_t \odot \tilde{\mathbf{c}}_t$$

is structurally analogous to the ResNet update:

$$\mathbf{x}_{l+1} = F(\mathbf{x}_l) + \mathbf{x}_l$$

In both cases, the output is the sum of the input (scaled by a gate in the LSTM case, or passed through unchanged in ResNet) plus a learned transformation. In both cases, the gradient of the output with respect to the input contains a term that doesn't pass through the learned transformation.

The key difference is that the LSTM's skip connection is learnable (the forget gate scales the cell state), while ResNet's skip connection passes the input unchanged. The LSTM therefore has more flexibility: it can open or close the gradient highway depending on what it has learned, while ResNet always keeps the highway fully open.

This difference has tradeoffs. ResNet's unconditional skip connection guarantees gradient flow without any training: even an untrained ResNet has a clear gradient path. The LSTM's gated skip connection needs to be learned into the "open" position before it starts helping. This is one reason why forget gate bias initialization matters so much for LSTMs: it puts the gate in an approximately "open" position before training starts, mimicking the ResNet's unconditional highway.

This analogy clarifies why the cell state is the "memory" component and the hidden state is the "communication" component. The cell state is the gradient highway, designed for long-range information storage. The hidden state is exposed to the outside world (fed to output layers, used by attention mechanisms, passed to the next layer), but it is more "narrow" in gradient flow because it passes through additional nonlinearities.

Later architectures like the GRU (which we'll explore in the next chapter) and Transformer architectures incorporate the residual connection insight even more directly. Transformers use explicit skip connections between every sublayer. This ensures that gradient flow is guaranteed at every depth level. In a sense, the LSTM's cell state foreshadowed this design pattern: create an additive path that bypasses the heavy computations so that gradient can flow freely. The progression from LSTM to ResNet to Transformer represents an evolution in how clearly and explicitly this design principle is applied.