Following Sepp Hochreiter’s 1991 thesis, he collaborated with Jürgen Schmidhuber to publish “Long Short-Term Memory” (1997). This paper introduced the Constant Error Carousel (CEC) and Gating Units (specifically the Input and Output gates) to solve the vanishing gradient problem.
Note: The “Forget Gate” was actually added later (Gers et al., 2000), but I will include it here as it is standard in modern PyTorch implementations.
1. The LSTM Cell Architecture
The core idea is the Cell State ($c_t$), which acts as a “conveyor belt” allowing information to flow through time with only minor linear interactions.
Python
import torch
import torch.nn as nn
import math
class LSTMCell(nn.Module):
def __init__(self, input_size, hidden_size):
super(LSTMCell, self).__init__()
self.input_size = input_size
self.hidden_size = hidden_size
# We combine all gates into one weight matrix for efficiency
# Gates: Input (i), Forget (f), Cell (g), Output (o)
self.x2h = nn.Linear(input_size, 4 * hidden_size)
self.h2h = nn.Linear(hidden_size, 4 * hidden_size)
self.reset_parameters()
def reset_parameters(self):
std = 1.0 / math.sqrt(self.hidden_size)
for w in self.parameters():
w.data.uniform_(-std, std)
def forward(self, x, states):
# states is a tuple (h_prev, c_prev)
h_prev, c_prev = states
# Batch computation of gate linear transformations
gates = self.x2h(x) + self.h2h(h_prev)
# Split the tensor into the 4 gate components
i_gate, f_gate, g_gate, o_gate = gates.chunk(4, 1)
# Apply activations
i_gate = torch.sigmoid(i_gate) # Input gate
f_gate = torch.sigmoid(f_gate) # Forget gate
o_gate = torch.sigmoid(o_gate) # Output gate
g_gate = torch.tanh(g_gate) # New candidate cell state
# Update Cell State: f_t * c_{t-1} + i_t * g_t
# This addition is the "Constant Error Carousel" logic
c_next = (f_gate * c_prev) + (i_gate * g_gate)
# Update Hidden State
h_next = o_gate * torch.tanh(c_next)
return h_next, c_next
2. The LSTM Model (Handling Sequences)
To process a full sentence or time series, we wrap the cell in a loop. PyTorch’s built-in nn.LSTM does this in C++ for speed, but here is the manual implementation for clarity:
Python
class LSTMModel(nn.Module):
def __init__(self, input_size, hidden_size, output_size):
super(LSTMModel, self).__init__()
self.hidden_size = hidden_size
self.cell = LSTMCell(input_size, hidden_size)
self.fc = nn.Linear(hidden_size, output_size)
def forward(self, x):
# x shape: [batch_size, sequence_length, input_size]
batch_size, seq_len, _ = x.size()
# Initialize h and c states with zeros
h = torch.zeros(batch_size, self.hidden_size).to(x.device)
c = torch.zeros(batch_size, self.hidden_size).to(x.device)
# Loop through the sequence
for t in range(seq_len):
h, c = self.cell(x[:, t, :], (h, c))
# Use the final hidden state for prediction
out = self.fc(h)
return out
Why this solved the problem
As highlighted in the 1997 paper, the LSTM avoids vanishing gradients because:
- Additive Updates: The cell state update $c_t = f_t c_{t-1} + i_t g_t$ uses addition rather than multiplication. When the forget gate $f_t$ is near 1, the gradient can flow backward through the cell state indefinitely without shrinking.
- Gating: The gates act as “valves” that learn when to let information in, when to clear the memory, and when to pass the memory to the next layer.
| Component | Function | Math Symbol |
| Forget Gate | Decides what to discard from the cell state | $f_t$ |
| Input Gate | Decides which new values to update | $i_t$ |
| Cell State | The long-term memory (CEC) | $c_t$ |
| Output Gate | Decides what parts of the cell state to output | $o_t$ |