RNNs & LSTMs

RNNs process sequences by maintaining a hidden state that gets updated at each timestep. The same weights are applied at every step — weight sharing across time.

Problem: vanilla RNNs suffer from vanishing gradients over long sequences. Gradients shrink exponentially as they backpropagate through time (BPTT), making it impossible to learn long-range dependencies.

LSTM (Long Short-Term Memory) fixes this with a cell state — a highway for gradients that runs through the sequence with only multiplicative interactions (gates). Three gates:
• Forget gate → what to erase from cell state
• Input gate  → what new info to write
• Output gate → what to expose as hidden state

In practice today: LSTMs are still used for time series and streaming tasks. For NLP, transformers replaced them entirely.

1import torch
2import torch.nn as nn
3
4# ── Vanilla RNN ───────────────────────────────────────────────────
5rnn = nn.RNN(
6    input_size=64,     # features per timestep
7    hidden_size=128,   # hidden state size
8    num_layers=2,      # stacked RNN layers
9    batch_first=True,  # (batch, seq, feature)
10    dropout=0.2,
11)
12
13x  = torch.randn(32, 20, 64)              # (batch=32, seq=20, features=64)
14h0 = torch.zeros(2, 32, 128)             # (num_layers, batch, hidden)
15out, hn = rnn(x, h0)
16print(f"RNN output:  {out.shape}")        # (32, 20, 128) — all timesteps
17print(f"RNN hidden:  {hn.shape}")         # (2, 32, 128)  — last hidden
18
19# ── LSTM ──────────────────────────────────────────────────────────
20lstm = nn.LSTM(
21    input_size=64,
22    hidden_size=128,
23    num_layers=2,
24    batch_first=True,
25    dropout=0.2,
26    bidirectional=False,   # set True for bi-LSTM
27)
28
29h0 = torch.zeros(2, 32, 128)
30c0 = torch.zeros(2, 32, 128)             # cell state (LSTM-specific)
31out, (hn, cn) = lstm(x, (h0, c0))
32print(f"LSTM output: {out.shape}")        # (32, 20, 128)
33
34# ── Sequence classifier using LSTM ───────────────────────────────
35class LSTMClassifier(nn.Module):
36    def __init__(self, vocab_size, embed_dim, hidden, num_classes):
37        super().__init__()
38        self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
39        self.lstm  = nn.LSTM(embed_dim, hidden, batch_first=True,
40                             num_layers=2, dropout=0.3, bidirectional=True)
41        self.head  = nn.Linear(hidden * 2, num_classes)  # *2 for bidirectional
42
43    def forward(self, x):
44        e         = self.embed(x)                # (B, T, E)
45        out, _    = self.lstm(e)                 # (B, T, H*2)
46        pooled    = out.mean(dim=1)              # mean pool over timesteps
47        return self.head(pooled)
48
49model = LSTMClassifier(vocab_size=10000, embed_dim=128,
50                       hidden=256, num_classes=5)
51tokens = torch.randint(0, 10000, (16, 50))      # (batch=16, seq_len=50)
52out    = model(tokens)
53print(f"Classifier output: {out.shape}")         # (16, 5)