Tag: machine-learning

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:

1. 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.
2. 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.