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