PyTorch LSTM

LSTM stands for Long Short-Term Memory. It is an improved RNN designed to remember information over long sequences. A standard RNN forgets information from early steps because gradients vanish during backpropagation. The LSTM introduces a cell state — a separate memory lane — controlled by gates that decide what to keep, what to forget, and what to output.

The Core Problem LSTM Solves

Standard RNN: "The cat sat on the mat"
  → After 100 words, the network has nearly forgotten the subject "cat"
  → Gradient has decayed to almost zero through 100 time steps

LSTM:
  → Cell state acts as long-term memory
  → Gates control what enters, exits, and persists in memory
  → Gradient can flow through the cell state without decaying

LSTM Has Two States

Standard RNN:  one state   → h_t (hidden state)
LSTM:          two states  → h_t (hidden state) + c_t (cell state)

h_t = short-term memory  (what to output right now)
c_t = long-term memory   (what to remember over time)

The Three Gates

Gates are learned sigmoid functions (output between 0 and 1). A value near 0 means "block this." A value near 1 means "let this through."

Forget Gate  (f_t):
  Decides what to erase from cell state c_{t-1}
  f_t ≈ 0 → forget everything  |  f_t ≈ 1 → keep everything

Input Gate   (i_t):
  Decides what new information to write into cell state
  Works together with a candidate cell value ĉ_t

Output Gate  (o_t):
  Decides what part of the cell state to expose as hidden state h_t

LSTM Step Diagram

Inputs:  x_t (current input), h_{t-1} (prev hidden), c_{t-1} (prev cell)

  ┌──────────────────────────────────────────────────────┐
  │ Forget gate:   f_t = σ(W_f·[h_{t-1}, x_t] + b_f)     │
  │ Input gate:    i_t = σ(W_i·[h_{t-1}, x_t] + b_i)     │
  │ Candidate:     ĉ_t = tanh(W_c·[h_{t-1}, x_t] + b_c)  │
  │ Cell update:   c_t = f_t ⊙ c_{t-1} + i_t ⊙ ĉ_t     │    
  │ Output gate:   o_t = σ(W_o·[h_{t-1}, x_t] + b_o)     │
  │ Hidden state:  h_t = o_t ⊙ tanh(c_t)                │
  └──────────────────────────────────────────────────────┘

⊙ = element-wise multiplication
σ = sigmoid

PyTorch nn.LSTM

import torch
import torch.nn as nn

lstm = nn.LSTM(
    input_size=10,
    hidden_size=32,
    num_layers=1,
    batch_first=True
)

x  = torch.rand(8, 20, 10)        # batch=8, seq_len=20, features=10
h0 = torch.zeros(1, 8, 32)        # initial hidden state
c0 = torch.zeros(1, 8, 32)        # initial cell state

output, (h_n, c_n) = lstm(x, (h0, c0))

print(output.shape)   # torch.Size([8, 20, 32])  ← hidden at every step
print(h_n.shape)      # torch.Size([1, 8, 32])   ← final hidden state
print(c_n.shape)      # torch.Size([1, 8, 32])   ← final cell state

Stacked LSTM

lstm = nn.LSTM(
    input_size=10,
    hidden_size=64,
    num_layers=2,     # 2 stacked LSTM layers
    batch_first=True,
    dropout=0.3       # dropout between stacked layers (not on last layer)
)

x  = torch.rand(8, 20, 10)
h0 = torch.zeros(2, 8, 64)    # 2 layers
c0 = torch.zeros(2, 8, 64)

output, (h_n, c_n) = lstm(x, (h0, c0))
print(h_n.shape)   # torch.Size([2, 8, 64])  ← one per layer

LSTM for Text Classification

class LSTMClassifier(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_size, num_classes):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.lstm      = nn.LSTM(embed_dim, hidden_size, batch_first=True)
        self.fc        = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        emb = self.embedding(x)               # [B, seq_len, embed_dim]
        _, (h_n, _) = self.lstm(emb)          # h_n: [1, B, hidden_size]
        return self.fc(h_n.squeeze(0))        # [B, num_classes]

model = LSTMClassifier(vocab_size=5000, embed_dim=64,
                       hidden_size=128, num_classes=2)

# Fake tokenized text: batch of 16 sentences, each 50 tokens long
tokens = torch.randint(0, 5000, (16, 50))
output = model(tokens)
print(output.shape)   # torch.Size([16, 2])  ← binary classification

When to Use Last Hidden State vs All Outputs

Use h_n (final hidden state) when:
  → Classifying the entire sequence (sentiment, topic)
  → One label for the whole sequence

Use output (all hidden states) when:
  → Predicting a label for each step (named entity recognition)
  → Building a sequence-to-sequence model
  → Applying attention over the sequence

Bidirectional LSTM

bilstm = nn.LSTM(input_size=10, hidden_size=32,
                 batch_first=True, bidirectional=True)

x = torch.rand(8, 20, 10)
output, (h_n, c_n) = bilstm(x)

print(output.shape)   # torch.Size([8, 20, 64])  ← 32 forward + 32 backward
print(h_n.shape)      # torch.Size([2, 8, 32])   ← 2 directions

Summary

LSTM extends the RNN with a cell state and three gates — forget, input, and output — that control memory flow. The forget gate decides what to erase. The input gate decides what to add. The output gate decides what to expose as the hidden state. This architecture allows gradients to flow without vanishing over long sequences. Use nn.LSTM with batch_first=True so your input has shape [batch, seq_len, features]. The output tuple contains the full output sequence and the final hidden and cell states. Use the final hidden state for sequence classification tasks. Use all outputs when you need predictions at every time step.