{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# The gradient highway\n", "\n", "What the cell state replaces the Jacobian product with, and what one number does to it.\n", "\n", "From chapter 13, [Gated Recurrence: LSTM and GRU](https://nlp.jcrlabz.com/book/gated/), of the course notes.\n", "\n", "Source: `book/code/worked_examples/gated.py`" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## What you need\n", "\n", "`pip install numpy`. The next cell does it. Colab usually has numpy already, but nothing here assumes that." ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "%pip install -q numpy" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## The script\n", "\n", "Everything the example defines, in one cell. Run it and the definitions\n", "are live in the notebook." ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "\"\"\"LSTM and GRU, worked as Chapter 13 works them.\n", "\n", "Four things, all reproducing the chapter's tables:\n", "\n", " the gradient highway why f_t multiplying beats tanh' times U\n", " the forget gate bias the initialisation that decides whether it works\n", " one LSTM step all four gates, with numbers\n", " the parameter count what a GRU saves by dropping a gate\n", "\n", " python3 gated.py # all four\n", " python3 gated.py --highway # the decay comparison\n", " python3 gated.py --bias # what b_f does to the decay\n", " python3 gated.py --step # one LSTM cell, every gate shown\n", " python3 gated.py --params # LSTM against GRU against plain RNN\n", "\n", "Install: pip install numpy\n", "\"\"\"\n", "\n", "import argparse\n", "import math\n", "\n", "try:\n", " import numpy as np\n", "except ImportError:\n", " raise SystemExit(\"this one needs numpy: pip install numpy\")\n", "\n", "\n", "def sigmoid(z):\n", " return 1 / (1 + np.exp(-z))\n", "\n", "\n", "# ------------------------------------------------------------- the highway\n", "\n", "def show_highway():\n", " print(\"\\nSTEP 1 two routes for a gradient\\n\")\n", " print(\" Plain RNN. Every step multiplies by diag(tanh') U, so what\")\n", " print(\" survives is a product of matrices you do not control.\\n\")\n", " print(\" dE/dh_0 = dE/dh_tau * product of diag(tanh') U\\n\")\n", " print(\" LSTM. The cell state has its own path, and along that path the\")\n", " print(\" only thing multiplying the gradient is the forget gate:\\n\")\n", " print(\" dE/dC_{t-1} = f_t * dE/dC_t\\n\")\n", " print(\" That is the whole idea. f_t is a number the network learns, so\")\n", " print(\" it can choose to keep a memory alive. tanh' times U was never\")\n", " print(\" under anyone's control.\\n\")\n", "\n", " print(f\" {'distance':>9}{'RNN, factor 0.5':>18}{'LSTM, f = 0.9':>16}\"\n", " f\"{'LSTM, f = 0.99':>17}\")\n", " print(\" \" + \"-\" * 62)\n", " for tau in (1, 10, 47, 100, 500):\n", " print(f\" {tau:>9}{0.5 ** tau:>18.3e}{0.9 ** tau:>16.3e}\"\n", " f\"{0.99 ** tau:>17.3e}\")\n", "\n", " print(\"\\n At f = 0.99 a gradient still has 1 per cent of its strength\")\n", " print(\" after 500 steps. At 0.5 it is gone before step 50.\")\n", " print(\"\\n Note what the LSTM did not do. It did not remove the decay.\")\n", " print(\" It put the decay rate under the network's control, and that is\")\n", " print(\" enough, because the network can now learn to set it near 1.\")\n", "\n", "\n", "# ------------------------------------------------------- the forget gate bias\n", "\n", "def show_bias():\n", " print(\"\\n\\nSTEP 2 the initialisation that decides whether any of it works\\n\")\n", " print(\" f_t = sigma(W_f [h_{t-1}, x_t] + b_f). At the start of training\")\n", " print(\" the weights are small and random, so f_t is roughly sigma(b_f).\\n\")\n", " print(f\" {'b_f':>6}{'sigma(b_f)':>13}{'after 10':>12}{'after 47':>12}\"\n", " f\"{'after 100':>13}\")\n", " print(\" \" + \"-\" * 58)\n", " for b in (0.0, 1.0, 2.0, 3.0, 5.0):\n", " f = float(sigmoid(b))\n", " print(f\" {b:>6.1f}{f:>13.4f}{f ** 10:>12.3e}{f ** 47:>12.3e}\"\n", " f\"{f ** 100:>13.3e}\")\n", "\n", " f0, f2 = float(sigmoid(0.0)), float(sigmoid(2.0))\n", " print(f\"\\n The default is b_f = 0, which gives f = {f0:.2f}. That decays\")\n", " print(f\" at exactly the rate the plain RNN did, so an LSTM initialised\")\n", " print(f\" this way looks like it cannot learn long dependencies at all.\")\n", " print(f\"\\n Set b_f = 2 and f starts at {f2:.4f}. After 47 steps the\")\n", " print(f\" gradient retains {f2 ** 47:.3e} instead of {f0 ** 47:.3e}, which\")\n", " print(f\" is {f2 ** 47 / f0 ** 47:.3e} times more signal.\")\n", " print(f\"\\n The architecture was never the problem in that case. One\")\n", " print(f\" scalar was. This is the most commonly skipped line in an LSTM\")\n", " print(f\" implementation, and skipping it looks exactly like the model\")\n", " print(f\" being incapable.\")\n", "\n", "\n", "# ------------------------------------------------------------- one LSTM step\n", "\n", "def show_step():\n", " \"\"\"One cell, small enough to print. Two units, input of size two.\"\"\"\n", " print(\"\\n\\nSTEP 3 one LSTM cell, every gate shown\\n\")\n", " h_prev = np.array([0.10, -0.20])\n", " C_prev = np.array([0.50, 0.30])\n", " x = np.array([0.60, 0.40])\n", " q = np.concatenate([h_prev, x]) # [h_{t-1}, x_t]\n", "\n", " Wf = np.array([[0.3, -0.2, 0.5, 0.1], [0.1, 0.4, -0.3, 0.2]])\n", " Wi = np.array([[-0.2, 0.3, 0.4, -0.1], [0.5, -0.1, 0.2, 0.3]])\n", " Wc = np.array([[0.4, 0.1, -0.2, 0.5], [-0.3, 0.2, 0.6, -0.1]])\n", " Wo = np.array([[0.2, 0.5, 0.1, -0.3], [0.4, -0.2, 0.3, 0.2]])\n", " bf = np.array([1.0, 1.0]) # the recommended init\n", " bi = bo = np.zeros(2)\n", "\n", " print(f\" h_(t-1) = {h_prev}, C_(t-1) = {C_prev}, x_t = {x}\")\n", " print(f\" b_f = {bf}, the initialisation Step 2 argues for\\n\")\n", "\n", " f = sigmoid(Wf @ q + bf)\n", " i = sigmoid(Wi @ q + bi)\n", " Ct = np.tanh(Wc @ q)\n", " o = sigmoid(Wo @ q + bo)\n", " C = f * C_prev + i * Ct\n", " h = o * np.tanh(C)\n", "\n", " def row(name, formula, v):\n", " print(f\" {name:<12}{formula:<34}{'(' + ', '.join(f'{a:+.4f}' for a in v) + ')'}\")\n", "\n", " print(f\" {'gate':<12}{'formula':<34}{'value'}\")\n", " print(\" \" + \"-\" * 68)\n", " row(\"forget\", \"f_t = sigma(W_f q + b_f)\", f)\n", " row(\"input\", \"i_t = sigma(W_i q + b_i)\", i)\n", " row(\"candidate\", \"C~_t = tanh(W_C q)\", Ct)\n", " row(\"output\", \"o_t = sigma(W_o q + b_o)\", o)\n", " print(\" \" + \"-\" * 68)\n", " row(\"cell\", \"C_t = f*C_(t-1) + i*C~_t\", C)\n", " row(\"hidden\", \"h_t = o * tanh(C_t)\", h)\n", "\n", " print(f\"\\n Read the cell update one term at a time.\")\n", " print(f\"\\n kept from the past: f * C_(t-1) =\"\n", " f\" ({f[0] * C_prev[0]:+.4f}, {f[1] * C_prev[1]:+.4f})\")\n", " print(f\" added this step: i * C~_t =\"\n", " f\" ({i[0] * Ct[0]:+.4f}, {i[1] * Ct[1]:+.4f})\")\n", " print(f\" new cell state: C_t =\"\n", " f\" ({C[0]:+.4f}, {C[1]:+.4f})\")\n", "\n", " print(f\"\\n The forget gate is at {f[0]:.4f} and {f[1]:.4f}, so most of\")\n", " print(f\" the old memory survives. That is the b_f = 1 initialisation\")\n", " print(f\" doing its job on the very first step.\")\n", " print(f\"\\n Note that C_t is reached by addition, not by matrix\")\n", " print(f\" multiplication. That is the structural difference from the\")\n", " print(f\" plain RNN, and it is the reason the gradient survives.\")\n", "\n", " print(f\"\\n\\n THE SAME STEP AS A GRU\\n\")\n", " Wz = np.array([[0.3, -0.2, 0.5, 0.1], [0.1, 0.4, -0.3, 0.2]])\n", " Wr = np.array([[-0.2, 0.3, 0.4, -0.1], [0.5, -0.1, 0.2, 0.3]])\n", " Wh = np.array([[0.4, 0.1, -0.2, 0.5], [-0.3, 0.2, 0.6, -0.1]])\n", " z = sigmoid(Wz @ q)\n", " r = sigmoid(Wr @ q)\n", " ht = np.tanh(Wh @ np.concatenate([r * h_prev, x]))\n", " hg = (1 - z) * h_prev + z * ht\n", "\n", " print(f\" {'gate':<12}{'formula':<34}{'value'}\")\n", " print(\" \" + \"-\" * 68)\n", " row(\"update\", \"z_t = sigma(W_z q)\", z)\n", " row(\"reset\", \"r_t = sigma(W_r q)\", r)\n", " row(\"candidate\", \"h~_t = tanh(W [r*h_(t-1), x])\", ht)\n", " print(\" \" + \"-\" * 68)\n", " row(\"hidden\", \"h_t = (1-z)*h_(t-1) + z*h~_t\", hg)\n", "\n", " print(f\"\\n No separate cell state. The GRU keeps one vector and splits\")\n", " print(f\" it with a single gate: (1-z) of the old, z of the new.\")\n", " print(f\" One gate does the job the LSTM gave to two.\")\n", "\n", "\n", "# ------------------------------------------------------------ the parameters\n", "\n", "def show_params(d=100, H=500):\n", " print(f\"\\n\\nSTEP 4 what a gate costs (input d = {d}, hidden H = {H})\\n\")\n", " print(\" Every gate reads the concatenation [h_(t-1), x_t], of length\")\n", " print(\" H + d, and produces H numbers. So one gate costs H(H + d) + H.\\n\")\n", " per = H * (H + d) + H\n", " rows = [(\"plain RNN\", 1, \"one state update\"),\n", " (\"GRU\", 3, \"update, reset, candidate\"),\n", " (\"LSTM\", 4, \"forget, input, candidate, output\")]\n", " print(f\" {'model':<12}{'gates':>7}{'parameters':>14}{'vs RNN':>9} what they are\")\n", " print(\" \" + \"-\" * 76)\n", " for name, n, what in rows:\n", " print(f\" {name:<12}{n:>7}{n * per:>14,}{n:>8}x {what}\")\n", "\n", " print(f\"\\n A GRU is {3 / 4:.0%} the size of an LSTM, which on this\")\n", " print(f\" configuration is {per:,} fewer parameters.\")\n", " print(f\"\\n That is the trade. The LSTM separates what to forget from\")\n", " print(f\" what to add, and keeps a cell state distinct from the output.\")\n", " print(f\" The GRU ties forgetting to adding with a single z, and exposes\")\n", " print(f\" its whole state.\")\n", " print(f\"\\n On most tasks the two score within noise of each other, so\")\n", " print(f\" the GRU's smaller size and faster step often decide it.\")\n", "\n", "\n", "def main():\n", " ap = argparse.ArgumentParser(description=__doc__)\n", " for f in (\"highway\", \"bias\", \"step\", \"params\"):\n", " ap.add_argument(f\"--{f}\", action=\"store_true\")\n", " a = ap.parse_args()\n", " picked = a.highway or a.bias or a.step or a.params\n", " if a.highway or not picked:\n", " show_highway()\n", " if a.bias or not picked:\n", " show_bias()\n", " if a.step or not picked:\n", " show_step()\n", " if a.params or not picked:\n", " show_params()\n", " print()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Running it\n", "\n", "The script is a command line tool, so each run below sets the arguments\n", "and calls `main()`. The heading shows the equivalent shell command." ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "import sys\n", "\n", "def run(*args):\n", " sys.argv = [\"gated.py\", *args]\n", " main()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### `python3 gated.py`\n", "\n", "The highway, the bias, one LSTM step and one GRU step." ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "run()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### `python3 gated.py --highway`\n", "\n", "The forget gate alone carries the gradient. At f = 0.99 one per cent of it survives 500 steps." ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "run(\"--highway\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### `python3 gated.py --bias`\n", "\n", "b_f = 0 gives f = 0.5, which decays exactly as the plain RNN did. Set b_f = 2 and 47 steps back is 2.6e-03 instead of 7.1e-15." ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "run(\"--bias\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "---\n", "\n", "Read the chapter this comes from: [Gated Recurrence: LSTM and GRU](https://nlp.jcrlabz.com/book/gated/)." ] } ], "metadata": { "colab": { "name": "gated.ipynb", "provenance": [], "toc_visible": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }