LSTM Gradient Flow › Stacked and Bidirectional LSTMs

Single-layer LSTMs are the building block, but real applications often use deeper or bidirectional variants. Understanding how gradient flow changes in these configurations is important for practical deployment.

A stacked LSTM consists of multiple LSTM layers where the hidden state output of one layer becomes the input sequence for the next. If we have $L$ layers, the output of layer $l$ at each timestep feeds into layer $l+1$:

$$\mathbf{h}_t^{(l)} = \text{LSTM}_l(\mathbf{h}_t^{(l-1)}, \mathbf{h}_{t-1}^{(l)}, \mathbf{c}_{t-1}^{(l)})$$

Each layer has its own cell state and its own forget gate, so each layer has its own gradient highway. The key question is how gradients flow vertically (through layers) versus horizontally (through time).

Vertically, gradients flow from layer $L$ down to layer 1, passing through the hidden state outputs at each layer. This path does not have the cell state protection: the hidden state $\mathbf{h}_t^{(l)}$ is a compressed version of the cell state, and gradients flowing through the hidden state encounter the full multiplicative barriers of weight matrices and nonlinearities at each layer transition.

Horizontally, within each layer, gradients flow through that layer's cell state with the same forget gate protection as in a single-layer LSTM. So a stacked LSTM has strong horizontal gradient flow within each layer but potentially weaker vertical gradient flow across layers.

This suggests that for stacked LSTMs on very long sequences, the outer layers may struggle to learn from early parts of the sequence because their gradients must pass through multiple vertical layers before reaching layer 1. In practice, residual connections between LSTM layers can alleviate vertical gradient attenuation in stacked configurations: adding a skip connection $\mathbf{h}_t^{(l)} = \text{LSTM}_l(\cdot) + \mathbf{h}_t^{(l-1)}$ creates an additive path from lower to upper layers, replicating the ResNet benefit in the depth dimension.

A bidirectional LSTM processes the sequence in both directions: one LSTM goes left to right, another goes right to left. Their outputs at each position are concatenated (or summed) to produce a representation that captures context from both past and future.

From a gradient flow perspective, the two LSTMs are independent. Each has its own cell state, its own forget gate, and its own gradient highway. Gradients from the loss flow back through both LSTMs simultaneously (in parallel) during backpropagation. There is no interaction between the two gradient flows except through the shared loss function.

Bidirectional LSTMs are standard for sequence labeling tasks like named entity recognition and part-of-speech tagging, where the label at each position depends on both the preceding and following context. They are also useful for document classification when you want the model to have access to the full sequence before making a prediction.

The key practical consideration for bidirectional LSTMs is that they cannot be used for online (streaming) prediction, where you need to produce an output before seeing the entire sequence. This limits their applicability to generation tasks, where you're restricted to causal (left-to-right) models. For tasks like translation and summarization where the full input is available before generating any output, bidirectional encoders paired with unidirectional decoders are a natural choice, and this is exactly the architecture used in the original sequence-to-sequence models.