Long Short-Term Memory (LSTM) Networks

The Story That Explains LSTMs

An LSTM is a special type of Recurrent Neural Network (RNN) designed to learn long-range dependencies in sequential data. Introduced by Hochreiter & Schmidhuber in 1997, it solved the infamous vanishing gradient problem that made vanilla RNNs useless for long sequences. Today, LSTMs power speech recognition, machine translation, time-series forecasting, and music generation.

Why Vanilla RNNs Fail — The Vanishing Gradient

To understand why LSTMs exist, you must feel the pain of what came before. A Vanilla RNN processes sequences step by step, passing a hidden state from one time step to the next. The hidden state is the only memory it has.

LSTM Cell Anatomy — The Animated Working Diagram

An LSTM cell takes three inputs — the current input xₜ, the previous hidden state h₍ₜ₋₁₎, and the previous cell state C₍ₜ₋₁₎ — and produces two outputs: a new hidden state hₜ and a new cell state Cₜ. Inside, four learned transformations govern everything.

🖱️ Click the diagram to pause/resume the animation. The blue highway at the top is the cell state — information can flow across the entire sequence largely untouched.

The Four Equations — What Actually Happens Inside

Every LSTM step runs exactly four vector-valued equations. Everything else is just wiring those equations together. Learn these four, and you understand the entire model.

Gate Intuition — The Three Security Guards

LSTM Unrolled Through Time

An LSTM doesn't just process one step — it processes a sequence. The same cell (same weights) is applied repeatedly, with the hidden state and cell state passed forward at each step. This is called unrolling.

LSTM Variants — The Family Tree

Key Hyperparameters — What You Actually Tune

Python Implementation — From Scratch to Production (PyTorch)

We'll build an LSTM for stock price forecasting — a classic time-series task. The model predicts the next day's closing price given the last 60 days of price data.

Step 1 — Imports

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt

# Reproducibility
torch.manual_seed(42)
np.random.seed(42)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

Step 2 — Data Preparation

# Load CSV with columns: Date, Close
df = pd.read_csv('stock_prices.csv', parse_dates=['Date'], index_col='Date')
prices = df[['Close']].values  # shape: (N, 1)

# Normalise to [0, 1] — critical for stable LSTM training
scaler = MinMaxScaler(feature_range=(0, 1))
prices_scaled = scaler.fit_transform(prices)

# Train / test split (80% / 20%) — NO shuffling for time series!
split = int(len(prices_scaled) * 0.8)
train_data = prices_scaled[:split]
test_data  = prices_scaled[split:]

print(f"Train samples: {len(train_data)} | Test samples: {len(test_data)}")

Step 3 — Sliding Window Dataset

class StockDataset(Dataset):
    def __init__(self, data, seq_len=60):
        self.seq_len = seq_len
        X, y = [], []
        for i in range(len(data) - seq_len):
            X.append(data[i : i + seq_len])        # shape: (60, 1)
            y.append(data[i + seq_len])             # next value
        self.X = torch.FloatTensor(np.array(X))  # (N, 60, 1)
        self.y = torch.FloatTensor(np.array(y))  # (N, 1)

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]

SEQ_LEN   = 60
BATCH     = 32

train_ds  = StockDataset(train_data, SEQ_LEN)
test_ds   = StockDataset(test_data,  SEQ_LEN)
train_dl  = DataLoader(train_ds, batch_size=BATCH, shuffle=True)
test_dl   = DataLoader(test_ds,  batch_size=BATCH, shuffle=False)

Step 4 — The LSTM Model

class LSTMForecaster(nn.Module):
    def __init__(self, input_size=1, hidden_size=128,
                 num_layers=2, dropout=0.2, output_size=1):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers  = num_layers

        # Core LSTM — 2 stacked layers with dropout between them
        self.lstm = nn.LSTM(
            input_size  = input_size,
            hidden_size = hidden_size,
            num_layers  = num_layers,
            dropout     = dropout,        # between layers only
            batch_first = True           # input: (batch, seq, features)
        )

        # Prediction head
        self.fc = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(hidden_size, 64),
            nn.ReLU(),
            nn.Linear(64, output_size)
        )

    def forward(self, x):
        # x: (batch, seq_len, input_size)
        # Initialise h₀ and C₀ to zeros
        h0 = torch.zeros(self.num_layers, x.size(0),
                          self.hidden_size).to(x.device)
        c0 = torch.zeros(self.num_layers, x.size(0),
                          self.hidden_size).to(x.device)

        # out: (batch, seq_len, hidden_size)
        # hn, cn: (num_layers, batch, hidden_size)
        out, (hn, cn) = self.lstm(x, (h0, c0))

        # Take ONLY the last time step's hidden state
        last_hidden = out[:, -1, :]  # (batch, hidden_size)
        return self.fc(last_hidden)   # (batch, 1)

model = LSTMForecaster(
    input_size  = 1,
    hidden_size = 128,
    num_layers  = 2,
    dropout     = 0.2,
    output_size = 1
).to(device)

total_params = sum(p.numel() for p in model.parameters())
print(f"Model parameters: {total_params:,}")

Step 5 — Training Loop with Gradient Clipping

criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, patience=5, factor=0.5, verbose=True
)

EPOCHS       = 50
CLIP_NORM    = 1.0   # gradient clipping threshold
train_losses = []
val_losses   = []

for epoch in range(1, EPOCHS + 1):

    # ── Training phase ──────────────────────────────────────
    model.train()
    epoch_loss = 0.0
    for X_batch, y_batch in train_dl:
        X_batch = X_batch.to(device)
        y_batch = y_batch.to(device)

        optimizer.zero_grad()
        preds = model(X_batch)                       # forward pass
        loss  = criterion(preds, y_batch)

        loss.backward()                              # BPTT
        nn.utils.clip_grad_norm_(model.parameters(),
                                   CLIP_NORM)        # ← critical
        optimizer.step()
        epoch_loss += loss.item() * len(X_batch)

    train_loss = epoch_loss / len(train_ds)
    train_losses.append(train_loss)

    # ── Validation phase ─────────────────────────────────────
    model.eval()
    val_loss = 0.0
    with torch.no_grad():
        for X_val, y_val in test_dl:
            X_val = X_val.to(device)
            y_val = y_val.to(device)
            val_preds = model(X_val)
            val_loss += criterion(val_preds, y_val).item() * len(X_val)
    val_loss /= len(test_ds)
    val_losses.append(val_loss)

    scheduler.step(val_loss)

    if epoch % 10 == 0:
        print(f"Epoch {epoch:3d}/{EPOCHS} | Train MSE: {train_loss:.6f} | Val MSE: {val_loss:.6f}")

Step 6 — Inference and Inverse Transform

model.eval()
all_preds, all_true = [], []

with torch.no_grad():
    for X_batch, y_batch in test_dl:
        X_batch = X_batch.to(device)
        pred = model(X_batch).cpu().numpy()
        all_preds.append(pred)
        all_true.append(y_batch.numpy())

preds_scaled = np.concatenate(all_preds)       # (N, 1)
true_scaled  = np.concatenate(all_true)

# Inverse MinMax transform → real price values
preds_price  = scaler.inverse_transform(preds_scaled)
true_price   = scaler.inverse_transform(true_scaled)

mae  = np.mean(np.abs(preds_price - true_price))
rmse = np.sqrt(np.mean((preds_price - true_price) ** 2))

print(f"MAE:  ${mae:.2f}")
print(f"RMSE: ${rmse:.2f}")

# Save model
torch.save(model.state_dict(), 'lstm_forecaster.pth')
print("Model saved.")

Keras / TensorFlow Implementation

The same model in Keras — fewer lines, identical concept. Keras is often preferable for rapid prototyping and research iteration.

import numpy as np
from tensorflow import keras
from tensorflow.keras import layers, callbacks

# Build model — functional API for clarity
inputs  = keras.Input(shape=(60, 1))          # (seq_len, features)
x       = layers.LSTM(128, return_sequences=True,
                       dropout=0.2)(inputs)    # layer 1 → passes full sequence
x       = layers.LSTM(64,  return_sequences=False,
                       dropout=0.2)(x)          # layer 2 → returns last step
x       = layers.Dense(32, activation='relu')(x)
outputs = layers.Dense(1)(x)

model = keras.Model(inputs, outputs)
model.compile(
    optimizer = keras.optimizers.Adam(lr=1e-3, clipnorm=1.0),  # ← gradient clip
    loss      = 'mse',
    metrics   = ['mae']
)
model.summary()

# Callbacks
cbs = [
    callbacks.EarlyStopping(patience=10, restore_best_weights=True),
    callbacks.ReduceLROnPlateau(patience=5, factor=0.5),
    callbacks.ModelCheckpoint('best_lstm.keras', save_best_only=True)
]

history = model.fit(
    X_train, y_train,
    validation_data = (X_test, y_test),
    epochs          = 100,
    batch_size      = 32,
    callbacks       = cbs,
    verbose         = 1
)

LSTM vs Transformer — When to Use Which

Real-World Applications — Where LSTMs Live in Production

Golden Rules — Non-Negotiable LSTM Practices

Quick Reference — LSTM Cheat Sheet