What the cell state replaces the Jacobian product with, and what one number does to it.
The notebook carries the source, the install step and every run below, so nothing has to be on your machine. Open it, then choose Runtime › Run all.
Colab opens it read only. Click Copy to Drive to keep
your changes. You can also download
gated.ipynb, or run it
locally: python3 gated.py.
python3 gated.py
The highway, the bias, one LSTM step and one GRU step.
python3 gated.py --highway
The forget gate alone carries the gradient. At f = 0.99 one per cent of it survives 500 steps.
python3 gated.py --bias
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.
Download gated.py
· served verbatim at https://nlp.jcrlabz.com/code/gated.py
"""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()