Long Short-Term Memory (LSTM) Networks

Simple Recurrent Neural Networks (RNNs) struggle with learning dependencies over long sequences. Their backpropagation process involves repeated multiplications through time, causing gradients to either shrink towards zero (vanish) or grow uncontrollably (explode). The vanishing gradient problem is particularly troublesome because it prevents the network from learning correlations between events that are far apart in the sequence.

Long Short-Term Memory (LSTM) networks were specifically designed to combat these issues. Introduced by Hochreiter and Schmidhuber in 1997, LSTMs incorporate a more complex internal structure than simple RNNs, featuring a dedicated cell state and gates that regulate the flow of information. This architecture allows LSTMs to maintain information over extended time intervals, effectively capturing long-range dependencies.

The Core Idea: Cell State and Gates

Think of the standard RNN hidden state $h_t$ as a working memory that gets overwritten at each time step. LSTMs introduce an additional component, the cell state $C_t$, which acts like a conveyor belt or a memory highway. Information can flow along this highway relatively unchanged, making it easier to preserve context over long durations.

The main innovation of LSTMs lies in their ability to selectively add or remove information from the cell state using structures called gates. Gates are composed of a sigmoid neural network layer and a pointwise multiplication operation. The sigmoid layer outputs numbers between 0 and 1, describing how much of each component should be let through. A value of 0 means "let nothing through," while a value of 1 means "let everything through."

An LSTM cell typically has three main gates:

1. Forget Gate: Decides what information to throw away from the cell state.
2. Input Gate: Decides which new information to store in the cell state.
3. Output Gate: Decides what part of the cell state to output.

Let's examine how these components interact at a single time step $t$, given the input $x_t$, the previous hidden state $h_{t-1}$, and the previous cell state $C_{t-1}$.

Forget Gate ($f_t$)

The first step is to decide what information we're going to discard from the previous cell state $C_{t-1}$. This decision is made by the forget gate layer. It looks at $h_{t-1}$ and $x_t$, and outputs a number between 0 and 1 for each number in the cell state $C_{t-1}$.

$$f_t = \sigma(W_f [h_{t-1}, x_t] + b_f)$$

Here, $\sigma$ is the sigmoid activation function, $W_f$ represents the weights, and $b_f$ the biases for the forget gate. The notation $[h_{t-1}, x_t]$ indicates that the previous hidden state and the current input are concatenated.

Input Gate ($i_t$) and Candidate Cell State ($\tilde{C}_t$)

Next, we need to decide what new information we're going to store in the cell state. This involves two parts. First, an input gate layer (another sigmoid layer) decides which values we'll update.

$$i_t = \sigma(W_i [h_{t-1}, x_t] + b_i)$$

Second, a tanh layer creates a vector of new candidate values, $\tilde{C}_t$, that could be added to the state.

$$\tilde{C}_t = \tanh(W_C [h_{t-1}, x_t] + b_C)$$

Updating the Cell State ($C_t$)

Now, we update the old cell state $C_{t-1}$ into the new cell state $C_t$. We multiply the old state by $f_t$, forgetting the things we decided to forget earlier. Then we add $i_t * \tilde{C}_t$. This is the new candidate information, scaled by how much we decided to update each state value.

$$C_t = f_t * C_{t-1} + i_t * \tilde{C}_t$$

The use of addition here is significant. Unlike the repeated multiplications in simple RNNs, this additive interaction makes it much easier for gradients to flow backward through time without vanishing as quickly.

Output Gate ($o_t$) and Hidden State ($h_t$)

Finally, we need to decide what we're going to output. This output will be based on our cell state, but will be a filtered version. First, we run a sigmoid layer which decides what parts of the cell state we’re going to output.

$$o_t = \sigma(W_o [h_{t-1}, x_t] + b_o)$$

Then, we put the cell state through tanh (to push the values to be between -1 and 1) and multiply it by the output of the sigmoid gate, so that we only output the parts we decided to. This final output is the hidden state $h_t$.

$$h_t = o_t * \tanh(C_t)$$

The hidden state $h_t$ is then passed to the next time step, along with the new cell state $C_t$.

Simplified view of information flow within an LSTM cell at time step $t$. The cell state $C_t$ acts as a conveyor belt, modified by the forget and input gates. The output gate filters the cell state to produce the hidden state $h_t$.

Implementation in PyTorch

Deep learning frameworks like PyTorch provide efficient implementations of LSTM layers. Using torch.nn.LSTM abstracts away the gate calculations.

import torch
import torch.nn as nn

# Example parameters
input_size = 10   # Dimension of input features x_t
hidden_size = 20  # Dimension of hidden state h_t and cell state C_t
num_layers = 1    # Number of stacked LSTM layers
batch_size = 5
seq_len = 7       # Length of the input sequence

# Create an LSTM layer
lstm_layer = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)

# Dummy input data (batch_size, seq_len, input_size)
input_seq = torch.randn(batch_size, seq_len, input_size)

# Initial hidden and cell states (num_layers, batch_size, hidden_size)
# If not provided, they default to zeros.
h0 = torch.randn(num_layers, batch_size, hidden_size)
c0 = torch.randn(num_layers, batch_size, hidden_size)

# Forward pass
# output contains the hidden state h_t for each time step in the sequence
# hn contains the final hidden state for the last time step
# cn contains the final cell state for the last time step
output, (hn, cn) = lstm_layer(input_seq, (h0, c0))

print("Input shape:", input_seq.shape)
print("Output shape (all hidden states):", output.shape)
print("Final hidden state shape (hn):", hn.shape)
print("Final cell state shape (cn):", cn.shape)

# --- Output ---
# Input shape: torch.Size([5, 7, 10])
# Output shape (all hidden states): torch.Size([5, 7, 20])
# Final hidden state shape (hn): torch.Size([1, 5, 20])
# Final cell state shape (cn): torch.Size([1, 5, 20])

Notice that the hidden state and cell state are returned as a tuple (hn, cn). This reflects the two distinct internal states maintained by the LSTM throughout the sequence processing. The output tensor provides the hidden state $h_t$ at every time step, which is often useful in sequence-to-sequence models.

By using gates to control information flow and an additive update mechanism for the cell state, LSTMs effectively mitigate the vanishing gradient problem and allow for the learning of dependencies across much longer time spans compared to simple RNNs. This capability made them a dominant architecture for many NLP tasks before the rise of attention-based models. While Gated Recurrent Units (GRUs), which we explore next, offer a slightly simpler gated mechanism, the fundamental principles introduced by LSTMs remain influential.