{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  },
  "colab": {
   "provenance": []
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cell-0000",
   "metadata": {},
   "source": [
    "# LSTM & GRU\n",
    "\n",
    "Gated memory cells that beat the vanishing gradient — how LSTM's forget/input/output gates work, GRU's streamlined variant, and both in PyTorch.\n",
    "\n",
    "*Part of the free [Deep Learning with PyTorch](https://ramadnsyh.dev/courses/deep-learning) course by [Muhammad Ramadiansyah](https://ramadnsyh.dev). This notebook is generated from the interactive lesson — [read it online](https://ramadnsyh.dev/courses/deep-learning/lstm-gru).*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0001",
   "metadata": {},
   "source": [
    "The last lesson ended on a problem: a vanilla RNN's gradient shrinks\n",
    "geometrically as it flows back through time, so the network can't learn\n",
    "dependencies more than a couple of dozen steps long. This lesson covers the\n",
    "fix that carried sequence modeling for two decades — the **LSTM** and its\n",
    "lighter sibling the **GRU** — and puts both to work in PyTorch."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0002",
   "metadata": {},
   "source": [
    "## Why \"I have a pen\" isn't enough\n",
    "\n",
    "Consider predicting the last word of: \"I grew up in France, moved away for\n",
    "work, lived in three other countries, and after all these years I still speak\n",
    "fluent ___\". The answer, *French*, depends on a word from 25 steps earlier.\n",
    "A vanilla RNN squashes its entire memory through a tanh at every step —\n",
    "relevant or not — so by the time \"French\" is needed, \"France\" has been\n",
    "overwritten dozens of times. What we need is a memory that is **kept by\n",
    "default** and only changed **on purpose**."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0003",
   "metadata": {},
   "source": [
    "## The conveyor belt\n",
    "\n",
    "The LSTM (Long Short-Term Memory, 1997) adds a second track of state: the\n",
    "**cell state** `c_t`, alongside the usual hidden state `h_t`. Think of the\n",
    "cell state as a conveyor belt running the length of the sequence. At each\n",
    "step, information rides along largely untouched; the network can *remove*\n",
    "something from the belt or *place* something new on it, but only through\n",
    "small, learned, elementwise adjustments:\n",
    "\n",
    "**c_t = f_t · c_(t-1) + i_t · g_t**\n",
    "\n",
    "Read it as: keep a fraction `f_t` of the old memory, and add `i_t` worth of\n",
    "new content `g_t`. The update is **additive**, not a full rewrite — and if\n",
    "that reminds you of ResNet's skip connection, it should. Both create a path\n",
    "along which gradients flow without being squashed at every step, and both were\n",
    "invented to cure the same disease: signals dying in deep compositions. Watch\n",
    "the difference numerically:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0004",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "h = 1.0     # vanilla RNN memory: squashed through tanh every step\n",
    "c = 1.0     # LSTM cell state: forget gate ~0.98, nothing new written\n",
    "\n",
    "print(\" step |  vanilla h  |  gated c\")\n",
    "print(\"------+-------------+----------\")\n",
    "for t in range(1, 51):\n",
    "    h = np.tanh(0.6 * h)      # typical effective recurrent factor < 1\n",
    "    c = 0.98 * c              # forget gate stays close to 1\n",
    "    if t in (1, 5, 10, 25, 50):\n",
    "        print(f\"  {t:3d} |   {h:9.5f} |  {c:7.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0005",
   "metadata": {},
   "source": [
    "The vanilla state collapses toward zero within a dozen steps; the gated cell\n",
    "still holds most of its signal after fifty. Gradients flowing backward enjoy\n",
    "the same protection."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0006",
   "metadata": {},
   "source": [
    "## Three gates: erase, write, reveal\n",
    "\n",
    "Who decides what to keep and what to add? **Gates** — tiny learned networks,\n",
    "each a sigmoid over the current input `x_t` and previous hidden state\n",
    "`h_(t-1)`, producing values between 0 and 1 that act as soft switches:\n",
    "\n",
    "- **Forget gate `f_t` — what to erase.** Multiplies the old cell state\n",
    "  elementwise. A value near 1 means \"keep this memory slot\", near 0 means\n",
    "  \"wipe it\". Seeing a new subject in a sentence might trigger forgetting the\n",
    "  old subject's gender.\n",
    "- **Input gate `i_t` — what to write.** Controls how much of the freshly\n",
    "  proposed content `g_t` (a tanh layer) is added onto the belt.\n",
    "- **Output gate `o_t` — what to reveal.** The cell state is private. The\n",
    "  hidden state that other layers actually see is a filtered view:\n",
    "  `h_t = o_t · tanh(c_t)`. The network can carry a memory for hundreds of\n",
    "  steps without exposing it until it's needed.\n",
    "\n",
    "Each gate has its own weights, all learned by backprop like everything else.\n",
    "That's the whole trick: memory management is not hard-coded — the network\n",
    "*learns* what's worth remembering, for how long, and when to use it."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0007",
   "metadata": {},
   "source": [
    "## GRU: the streamlined variant\n",
    "\n",
    "The **Gated Recurrent Unit** (2014) asks: do we really need three gates and\n",
    "two states? It keeps just two gates and folds the cell state back into `h_t`:\n",
    "\n",
    "- **Update gate `z_t`** merges forget and input into one decision — whatever\n",
    "  fraction of memory you erase is exactly replaced by new content:\n",
    "  `h_t = (1 − z_t) · h_(t-1) + z_t · h̃_t`.\n",
    "- **Reset gate `r_t`** controls how much of the previous state is consulted\n",
    "  when proposing new content.\n",
    "\n",
    "The result has roughly 25% fewer parameters than an LSTM of the same hidden\n",
    "size, trains a bit faster, and performs comparably on most tasks.\n",
    "\n",
    "**Which one?** Honest answer: it rarely matters much. Reasonable defaults —\n",
    "start with GRU for smaller datasets or when speed matters (fewer parameters,\n",
    "less overfitting); reach for LSTM on larger datasets and longer sequences,\n",
    "where its separate cell state sometimes gives it the edge. If the choice is\n",
    "decisive for your problem, you'll only find out by benchmarking both."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0008",
   "metadata": {},
   "source": [
    "## LSTM and GRU in PyTorch\n",
    "\n",
    "Both are drop-in replacements for `nn.RNN` — same constructor, same\n",
    "`batch_first`, same output shapes. The one difference: `nn.LSTM`'s second\n",
    "return value is a *tuple* `(h_n, c_n)` because of the extra cell state. Here\n",
    "is the sine next-step task from last lesson, with all three cells racing\n",
    "under identical conditions (Colab):"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0009",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn, optim\n",
    "\n",
    "torch.manual_seed(0)\n",
    "\n",
    "# data: sine wave -> (window, next value) pairs\n",
    "t = torch.linspace(0, 60, 600)\n",
    "series = torch.sin(t)\n",
    "seq_len = 20\n",
    "X = torch.stack([series[i:i + seq_len] for i in range(len(series) - seq_len)])\n",
    "X, y = X.unsqueeze(-1), series[seq_len:].unsqueeze(-1)\n",
    "\n",
    "CELLS = {\"rnn\": nn.RNN, \"lstm\": nn.LSTM, \"gru\": nn.GRU}\n",
    "\n",
    "class SeqModel(nn.Module):\n",
    "    def __init__(self, cell, hidden_size=32):\n",
    "        super().__init__()\n",
    "        self.rnn = CELLS[cell](1, hidden_size, batch_first=True)\n",
    "        self.fc = nn.Linear(hidden_size, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out, _ = self.rnn(x)          # for LSTM, \"_\" is the tuple (h_n, c_n)\n",
    "        return self.fc(out[:, -1, :])\n",
    "\n",
    "for cell in CELLS:\n",
    "    torch.manual_seed(0)\n",
    "    model = SeqModel(cell)\n",
    "    optimizer = optim.Adam(model.parameters(), lr=1e-3)\n",
    "    criterion = nn.MSELoss()\n",
    "    for epoch in range(200):\n",
    "        loss = criterion(model(X), y)\n",
    "        loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "    print(f\"{cell:>4} after 200 epochs: MSE = {loss.item():.6f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0010",
   "metadata": {},
   "source": [
    "On this short, clean sequence all three learn — but the gated cells typically\n",
    "converge faster and to a lower loss, and the gap widens dramatically as\n",
    "`seq_len` grows (try 100). The exercise below makes you measure exactly that."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0011",
   "metadata": {},
   "source": [
    "## Gradient clipping\n",
    "\n",
    "Gates solve vanishing gradients; **exploding** gradients get a blunter tool.\n",
    "Before each optimizer step, rescale the gradient vector if its norm exceeds a\n",
    "threshold:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0012",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "loss.backward()\n",
    "torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n",
    "optimizer.step()\n",
    "optimizer.zero_grad()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0013",
   "metadata": {},
   "source": [
    "It's one line between `backward()` and `step()`, and it turns \"loss went NaN\n",
    "at epoch 37\" into a non-event. Clipping is near-universal practice when\n",
    "training recurrent networks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0014",
   "metadata": {},
   "source": [
    "## Do RNNs still matter?\n",
    "\n",
    "Honesty time: since ~2018, **transformers ate NLP**. Attention sees every\n",
    "position at once, trains in parallel instead of step by step, and scales to\n",
    "billions of parameters — no recurrent model competes on language benchmarks\n",
    "today. But recurrence is far from dead. RNNs process a stream with **constant\n",
    "memory per step** and no need to store a growing context window, which keeps\n",
    "them relevant for streaming and low-latency inference, wake-word detection\n",
    "and other tiny on-device models, and plenty of time-series work where a\n",
    "50K-parameter GRU beats an over-parameterized transformer on 3,000 data\n",
    "points. The ideas you just learned — gating, additive state, learned\n",
    "forgetting — also live on inside modern state-space models. Learn the\n",
    "concepts; they keep resurfacing."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0015",
   "metadata": {},
   "source": [
    "### 🏋️ Exercise — Stress-test the memory\n",
    "\n",
    "3}:  {line}\")\n",
    "`}\n",
    ">\n",
    "Take the three-cell comparison from this lesson and turn it into a function of\n",
    "`seq_len`. Run it at `seq_len=20` and `seq_len=100` (same seed, epochs, and\n",
    "hidden size throughout) and report the final MSE for RNN, LSTM, and GRU at\n",
    "each length. Which cell degrades most as the sequence gets longer — and does\n",
    "that match the vanishing-gradient story?"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0016",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your solution here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0017",
   "metadata": {},
   "source": [
    "<details><summary>✅ Show solution</summary>\n",
    "\n",
    "```python\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "\n",
    "CELLS = {\"rnn\": nn.RNN, \"lstm\": nn.LSTM, \"gru\": nn.GRU}\n",
    "\n",
    "class SeqModel(nn.Module):\n",
    "    def __init__(self, cell, hidden_size=32):\n",
    "        super().__init__()\n",
    "        self.rnn = CELLS[cell](1, hidden_size, batch_first=True)\n",
    "        self.fc = nn.Linear(hidden_size, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out, _ = self.rnn(x)\n",
    "        return self.fc(out[:, -1, :])\n",
    "\n",
    "def benchmark(seq_len, epochs=200):\n",
    "    t = torch.linspace(0, 60, 600)\n",
    "    series = torch.sin(t)\n",
    "    X = torch.stack([series[i:i + seq_len]\n",
    "                     for i in range(len(series) - seq_len)]).unsqueeze(-1)\n",
    "    y = series[seq_len:].unsqueeze(-1)\n",
    "\n",
    "    results = {}\n",
    "    for cell in CELLS:\n",
    "        torch.manual_seed(0)\n",
    "        model = SeqModel(cell)\n",
    "        opt = optim.Adam(model.parameters(), lr=1e-3)\n",
    "        crit = nn.MSELoss()\n",
    "        for _ in range(epochs):\n",
    "            loss = crit(model(X), y)\n",
    "            loss.backward()\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "            opt.step()\n",
    "            opt.zero_grad()\n",
    "        results[cell] = loss.item()\n",
    "    return results\n",
    "\n",
    "for seq_len in (20, 100):\n",
    "    res = benchmark(seq_len)\n",
    "    line = \"  \".join(f\"{c}={v:.6f}\" for c, v in res.items())\n",
    "    print(f\"seq_len={seq_len:>3}:  {line}\")\n",
    "```\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0018",
   "metadata": {},
   "source": [
    "Next: putting sequence models to work — bidirectional and stacked RNNs,\n",
    "windowing and scaling real series, and a full forecasting pipeline."
   ]
  }
 ]
}