chapter_recurrent-modern/gru.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 0
   },
   "source": [
    "# Gated Recurrent Units (GRU)\n",
    ":label:`sec_gru`\n",
    "\n",
    "In :numref:`sec_bptt`,\n",
    "we discussed how gradients are calculated\n",
    "in RNNs.\n",
    "In particular we found that long products of matrices can lead\n",
    "to vanishing or exploding gradients.\n",
    "Let us briefly think about what such\n",
    "gradient anomalies mean in practice:\n",
    "\n",
    "* We might encounter a situation where an early observation is highly\n",
    "  significant for predicting all future observations. Consider the somewhat\n",
    "  contrived case where the first observation contains a checksum and the goal is\n",
    "  to discern whether the checksum is correct at the end of the sequence. In this\n",
    "  case, the influence of the first token is vital. We would like to have some\n",
    "  mechanisms for storing vital early information in a *memory cell*. Without such\n",
    "  a mechanism, we will have to assign a very large gradient to this observation,\n",
    "  since it affects all the subsequent observations.\n",
    "* We might encounter situations where some tokens carry no pertinent\n",
    "  observation. For instance, when parsing a web page there might be auxiliary\n",
    "  HTML code that is irrelevant for the purpose of assessing the sentiment\n",
    "  conveyed on the page. We would like to have some mechanism for *skipping* such\n",
    "  tokens in the latent state representation.\n",
    "* We might encounter situations where there is a logical break between parts of\n",
    "  a sequence. For instance, there might be a transition between chapters in a\n",
    "  book, or a transition between a bear and a bull market for securities. In\n",
    "  this case it would be nice to have a means of *resetting* our internal state\n",
    "  representation.\n",
    "\n",
    "A number of methods have been proposed to address this. One of the earliest is long short-term memory :cite:`Hochreiter.Schmidhuber.1997` which we\n",
    "will discuss in :numref:`sec_lstm`. The gated recurrent unit (GRU)\n",
    ":cite:`Cho.Van-Merrienboer.Bahdanau.ea.2014` is a slightly more streamlined\n",
    "variant that often offers comparable performance and is significantly faster to\n",
    "compute  :cite:`Chung.Gulcehre.Cho.ea.2014`.\n",
    "Due to its simplicity, let us start with the GRU.\n",
    "\n",
    "## Gated Hidden State\n",
    "\n",
    "The key distinction between vanilla RNNs and GRUs\n",
    "is that the latter support gating of the hidden state.\n",
    "This means that we have dedicated mechanisms for\n",
    "when a hidden state should be *updated* and\n",
    "also when it should be *reset*.\n",
    "These mechanisms are learned and they address the concerns listed above.\n",
    "For instance, if the first token is of great importance\n",
    "we will learn not to update the hidden state after the first observation.\n",
    "Likewise, we will learn to skip irrelevant temporary observations.\n",
    "Last, we will learn to reset the latent state whenever needed.\n",
    "We discuss this in detail below.\n",
    "\n",
    "\n",
    "### Reset Gate and Update Gate\n",
    "\n",
    "The first thing we need to introduce are\n",
    "the *reset gate* and the *update gate*.\n",
    "We engineer them to be vectors with entries in $(0, 1)$\n",
    "such that we can perform convex combinations.\n",
    "For instance,\n",
    "a reset gate would allow us to control how much of the previous state we might still want to remember.\n",
    "Likewise, an update gate would allow us to control how much of the new state is just a copy of the old state.\n",
    "\n",
    "We begin by engineering these gates.\n",
    ":numref:`fig_gru_1` illustrates the inputs for both\n",
    "the reset and update gates in a GRU, given the input\n",
    "of the current time step\n",
    "and the hidden state of the previous time step.\n",
    "The outputs of two gates\n",
    "are given by two fully-connected layers\n",
    "with a sigmoid activation function.\n",
    "\n",
    "![Computing the reset gate and the update gate in a GRU model.](https://raw.githubusercontent.com/d2l-ai/d2l-en/master/img/gru-1.svg)\n",
    ":label:`fig_gru_1`\n",
    "\n",
    "Mathematically,\n",
    "for a given time step $t$,\n",
    "suppose that the input is\n",
    "a minibatch\n",
    "$\\mathbf{X}_t \\in \\mathbb{R}^{n \\times d}$ (number of examples: $n$, number of inputs: $d$) and the hidden state of the previous time step is $\\mathbf{H}_{t-1} \\in \\mathbb{R}^{n \\times h}$ (number of hidden units: $h$). Then, the reset gate $\\mathbf{R}_t \\in \\mathbb{R}^{n \\times h}$ and update gate $\\mathbf{Z}_t \\in \\mathbb{R}^{n \\times h}$ are computed as follows:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\mathbf{R}_t = \\sigma(\\mathbf{X}_t \\mathbf{W}_{xr} + \\mathbf{H}_{t-1} \\mathbf{W}_{hr} + \\mathbf{b}_r),\\\\\n",
    "\\mathbf{Z}_t = \\sigma(\\mathbf{X}_t \\mathbf{W}_{xz} + \\mathbf{H}_{t-1} \\mathbf{W}_{hz} + \\mathbf{b}_z),\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "where $\\mathbf{W}_{xr}, \\mathbf{W}_{xz} \\in \\mathbb{R}^{d \\times h}$ and\n",
    "$\\mathbf{W}_{hr}, \\mathbf{W}_{hz} \\in \\mathbb{R}^{h \\times h}$ are weight\n",
    "parameters and $\\mathbf{b}_r, \\mathbf{b}_z \\in \\mathbb{R}^{1 \\times h}$ are\n",
    "biases. \n",
    "Note that broadcasting (see :numref:`subsec_broadcasting`) is triggered during the summation.\n",
    "We use sigmoid functions (as introduced in :numref:`sec_mlp`) to transform input values to the interval $(0, 1)$.\n",
    "\n",
    "### Candidate Hidden State\n",
    "\n",
    "Next, let us\n",
    "integrate the reset gate $\\mathbf{R}_t$ with\n",
    "the regular latent state updating mechanism\n",
    "in :eqref:`rnn_h_with_state`.\n",
    "It leads to the following\n",
    "*candidate hidden state*\n",
    "$\\tilde{\\mathbf{H}}_t \\in \\mathbb{R}^{n \\times h}$ at time step $t$:\n",
    "\n",
    "$$\\tilde{\\mathbf{H}}_t = \\tanh(\\mathbf{X}_t \\mathbf{W}_{xh} + \\left(\\mathbf{R}_t \\odot \\mathbf{H}_{t-1}\\right) \\mathbf{W}_{hh} + \\mathbf{b}_h),$$\n",
    ":eqlabel:`gru_tilde_H`\n",
    "\n",
    "where $\\mathbf{W}_{xh} \\in \\mathbb{R}^{d \\times h}$ and $\\mathbf{W}_{hh} \\in \\mathbb{R}^{h \\times h}$\n",
    "are weight parameters,\n",
    "$\\mathbf{b}_h \\in \\mathbb{R}^{1 \\times h}$\n",
    "is the bias,\n",
    "and the symbol $\\odot$ is the Hadamard (elementwise) product operator.\n",
    "Here we use a nonlinearity in the form of tanh to ensure that the values in the candidate hidden state remain in the interval $(-1, 1)$.\n",
    "\n",
    "The result is a *candidate* since we still need to incorporate the action of the update gate.\n",
    "Comparing with :eqref:`rnn_h_with_state`,\n",
    "now the influence of the previous states\n",
    "can be reduced with the\n",
    "elementwise multiplication of\n",
    "$\\mathbf{R}_t$ and $\\mathbf{H}_{t-1}$\n",
    "in :eqref:`gru_tilde_H`.\n",
    "Whenever the entries in the reset gate $\\mathbf{R}_t$ are close to 1, we recover a vanilla RNN such as in :eqref:`rnn_h_with_state`.\n",
    "For all entries of the reset gate $\\mathbf{R}_t$ that are close to 0, the candidate hidden state is the result of an MLP with $\\mathbf{X}_t$ as the input. Any pre-existing hidden state is thus *reset* to defaults.\n",
    "\n",
    ":numref:`fig_gru_2` illustrates the computational flow after applying the reset gate.\n",
    "\n",
    "![Computing the candidate hidden state in a GRU model.](https://raw.githubusercontent.com/d2l-ai/d2l-en/master/img/gru-2.svg)\n",
    ":label:`fig_gru_2`\n",
    "\n",
    "\n",
    "### Hidden State\n",
    "\n",
    "Finally, we need to incorporate the effect of the update gate $\\mathbf{Z}_t$. This determines the extent to which the new hidden state $\\mathbf{H}_t \\in \\mathbb{R}^{n \\times h}$ is just the old state $\\mathbf{H}_{t-1}$ and by how much the new candidate state $\\tilde{\\mathbf{H}}_t$ is used.\n",
    "The update gate $\\mathbf{Z}_t$ can be used for this purpose, simply by taking elementwise convex combinations between both $\\mathbf{H}_{t-1}$ and $\\tilde{\\mathbf{H}}_t$.\n",
    "This leads to the final update equation for the GRU:\n",
    "\n",
    "$$\\mathbf{H}_t = \\mathbf{Z}_t \\odot \\mathbf{H}_{t-1}  + (1 - \\mathbf{Z}_t) \\odot \\tilde{\\mathbf{H}}_t.$$\n",
    "\n",
    "\n",
    "Whenever the update gate $\\mathbf{Z}_t$ is close to 1, we simply retain the old state. In this case the information from $\\mathbf{X}_t$ is essentially ignored, effectively skipping time step $t$ in the dependency chain. In contrast, whenever $\\mathbf{Z}_t$ is close to 0, the new latent state $\\mathbf{H}_t$ approaches the candidate latent state $\\tilde{\\mathbf{H}}_t$. These designs can help us cope with the vanishing gradient problem in RNNs and better capture dependencies for sequences with large time step distances.\n",
    "For instance,\n",
    "if the update gate has been close to 1\n",
    "for all the time steps of an entire subsequence,\n",
    "the old hidden state at the time step of its beginning\n",
    "will be easily retained and passed\n",
    "to its end,\n",
    "regardless of the length of the subsequence.\n",
    "\n",
    "\n",
    "\n",
    ":numref:`fig_gru_3` illustrates the computational flow after the update gate is in action.\n",
    "\n",
    "![Computing the hidden state in a GRU model.](https://raw.githubusercontent.com/d2l-ai/d2l-en/master/img/gru-3.svg)\n",
    ":label:`fig_gru_3`\n",
    "\n",
    "\n",
    "In summary, GRUs have the following two distinguishing features:\n",
    "\n",
    "* Reset gates help capture short-term dependencies in sequences.\n",
    "* Update gates help capture long-term dependencies in sequences.\n",
    "\n",
    "## Implementation from Scratch\n",
    "\n",
    "To gain a better understanding of the GRU model, let us implement it from scratch. We begin by reading the time machine dataset that we used in :numref:`sec_rnn_scratch`. The code for reading the dataset is given below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load ../utils/djl-imports\n",
    "%load ../utils/plot-utils\n",
    "%load ../utils/Functions.java\n",
    "%load ../utils/PlotUtils.java\n",
    "\n",
    "%load ../utils/StopWatch.java\n",
    "%load ../utils/Accumulator.java\n",
    "%load ../utils/Animator.java\n",
    "%load ../utils/Training.java\n",
    "%load ../utils/timemachine/Vocab.java\n",
    "%load ../utils/timemachine/RNNModel.java\n",
    "%load ../utils/timemachine/RNNModelScratch.java\n",
    "%load ../utils/timemachine/TimeMachine.java\n",
    "%load ../utils/timemachine/TimeMachineDataset.java"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NDManager manager = NDManager.newBaseManager();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "int batchSize = 32;\n",
    "int numSteps = 35;\n",
    "\n",
    "TimeMachineDataset dataset =\n",
    "        new TimeMachineDataset.Builder()\n",
    "                .setManager(manager)\n",
    "                .setMaxTokens(10000)\n",
    "                .setSampling(batchSize, false)\n",
    "                .setSteps(numSteps)\n",
    "                .build();\n",
    "dataset.prepare();\n",
    "Vocab vocab = dataset.getVocab();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 3
   },
   "source": [
    "### Initializing Model Parameters\n",
    "\n",
    "The next step is to initialize the model parameters.\n",
    "We draw the weights from a Gaussian distribution\n",
    "with standard deviation to be 0.01 and set the bias to 0. The hyperparameter `num_hiddens` defines the number of hidden units.\n",
    "We instantiate all weights and biases relating to the update gate, the reset gate, the candidate hidden state,\n",
    "and the output layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "public static NDArray normal(Shape shape, Device device) {\n",
    "    return manager.randomNormal(0, 0.01f, shape, DataType.FLOAT32, device);\n",
    "}\n",
    "\n",
    "public static NDList three(int numInputs, int numHiddens, Device device) {\n",
    "    return new NDList(\n",
    "            normal(new Shape(numInputs, numHiddens), device),\n",
    "            normal(new Shape(numHiddens, numHiddens), device),\n",
    "            manager.zeros(new Shape(numHiddens), DataType.FLOAT32, device));\n",
    "}\n",
    "\n",
    "public static NDList getParams(int vocabSize, int numHiddens, Device device) {\n",
    "    int numInputs = vocabSize;\n",
    "    int numOutputs = vocabSize;\n",
    "\n",
    "    // Update gate parameters\n",
    "    NDList temp = three(numInputs, numHiddens, device);\n",
    "    NDArray W_xz = temp.get(0);\n",
    "    NDArray W_hz = temp.get(1);\n",
    "    NDArray b_z = temp.get(2);\n",
    "\n",
    "    // Reset gate parameters\n",
    "    temp = three(numInputs, numHiddens, device);\n",
    "    NDArray W_xr = temp.get(0);\n",
    "    NDArray W_hr = temp.get(1);\n",
    "    NDArray b_r = temp.get(2);\n",
    "\n",
    "    // Candidate hidden state parameters\n",
    "    temp = three(numInputs, numHiddens, device);\n",
    "    NDArray W_xh = temp.get(0);\n",
    "    NDArray W_hh = temp.get(1);\n",
    "    NDArray b_h = temp.get(2);\n",
    "\n",
    "    // Output layer parameters\n",
    "    NDArray W_hq = normal(new Shape(numHiddens, numOutputs), device);\n",
    "    NDArray b_q = manager.zeros(new Shape(numOutputs), DataType.FLOAT32, device);\n",
    "\n",
    "    // Attach gradients\n",
    "    NDList params = new NDList(W_xz, W_hz, b_z, W_xr, W_hr, b_r, W_xh, W_hh, b_h, W_hq, b_q);\n",
    "    for (NDArray param : params) {\n",
    "        param.setRequiresGradient(true);\n",
    "    }\n",
    "    return params;\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 6
   },
   "source": [
    "### Defining the Model\n",
    "\n",
    "Now we will define the hidden state initialization function `init_gru_state`. Just like the `init_rnn_state` function defined in :numref:`sec_rnn_scratch`, this function returns a tensor with a shape (batch size, number of hidden units) whose values are all zeros.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "public static NDList initGruState(int batchSize, int numHiddens, Device device) {\n",
    "    return new NDList(manager.zeros(new Shape(batchSize, numHiddens), DataType.FLOAT32, device));\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 9
   },
   "source": [
    "Now we are ready to define the GRU model.\n",
    "Its structure is the same as that of the basic RNN cell, except that the update equations are more complex.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "public static Pair<NDArray, NDList> gru(NDArray inputs, NDList state, NDList params) {\n",
    "    NDArray W_xz = params.get(0);\n",
    "    NDArray W_hz = params.get(1);\n",
    "    NDArray b_z = params.get(2);\n",
    "\n",
    "    NDArray W_xr = params.get(3);\n",
    "    NDArray W_hr = params.get(4);\n",
    "    NDArray b_r = params.get(5);\n",
    "\n",
    "    NDArray W_xh = params.get(6);\n",
    "    NDArray W_hh = params.get(7);\n",
    "    NDArray b_h = params.get(8);\n",
    "\n",
    "    NDArray W_hq = params.get(9);\n",
    "    NDArray b_q  = params.get(10);\n",
    "\n",
    "    NDArray H = state.get(0);\n",
    "    NDList outputs = new NDList();\n",
    "    NDArray X, Y, Z, R, H_tilda;\n",
    "    for (int i = 0; i < inputs.size(0); i++) {\n",
    "        X = inputs.get(i);\n",
    "        Z = Activation.sigmoid(X.dot(W_xz).add(H.dot(W_hz).add(b_z)));\n",
    "        R = Activation.sigmoid(X.dot(W_xr).add(H.dot(W_hr).add(b_r)));\n",
    "        H_tilda = Activation.tanh(X.dot(W_xh).add(R.mul(H).dot(W_hh).add(b_h)));\n",
    "        H = Z.mul(H).add(Z.mul(-1).add(1).mul(H_tilda));\n",
    "        Y = H.dot(W_hq).add(b_q);\n",
    "        outputs.add(Y);\n",
    "    }\n",
    "    return new Pair(outputs.size() > 1 ? NDArrays.concat(outputs) : outputs.get(0), new NDList(H));\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 12
   },
   "source": [
    "### Training and Prediction\n",
    "\n",
    "Training and prediction work in exactly the same manner as in :numref:`sec_rnn_scratch`.\n",
    "After training,\n",
    "we print out the perplexity on the training set\n",
    "and the predicted sequence following\n",
    "the provided prefixes \"time traveller\" and \"traveller\", respectively.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "int vocabSize = vocab.length();\n",
    "int numHiddens = 256;\n",
    "Device device = manager.getDevice();\n",
    "int numEpochs = Integer.getInteger(\"MAX_EPOCH\", 500);\n",
    "\n",
    "int lr = 1;\n",
    "\n",
    "Functions.TriFunction<Integer, Integer, Device, NDList> getParamsFn = (a, b, c) -> getParams(a, b, c);\n",
    "Functions.TriFunction<Integer, Integer, Device, NDList> initGruStateFn =\n",
    "        (a, b, c) -> initGruState(a, b, c);\n",
    "Functions.TriFunction<NDArray, NDList, NDList, Pair<NDArray, NDList>> gruFn = (a, b, c) -> gru(a, b, c);\n",
    "\n",
    "RNNModelScratch model =\n",
    "        new RNNModelScratch(vocabSize, numHiddens, device,\n",
    "                getParamsFn, initGruStateFn, gruFn);\n",
    "TimeMachine.trainCh8(model, dataset, vocab, lr, numEpochs, device, false, manager);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 14
   },
   "source": [
    "## Concise Implementation\n",
    "\n",
    "In high-level APIs,\n",
    "we can directly\n",
    "instantiate a GPU model.\n",
    "This encapsulates all the configuration detail that we made explicit above.\n",
    "The code is significantly faster as it uses compiled operators rather than Python for many details that we spelled out before.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "GRU gruLayer = GRU.builder().setNumLayers(1)\n",
    "        .setStateSize(numHiddens).optReturnState(true).optBatchFirst(false).build();\n",
    "RNNModel modelConcise = new RNNModel(gruLayer,vocab.length());\n",
    "TimeMachine.trainCh8(modelConcise, dataset, vocab, lr, numEpochs, device, false, manager);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 17
   },
   "source": [
    "## Summary\n",
    "\n",
    "* Gated RNNs can better capture dependencies for sequences with large time step distances.\n",
    "* Reset gates help capture short-term dependencies in sequences.\n",
    "* Update gates help capture long-term dependencies in sequences.\n",
    "* GRUs contain basic RNNs as their extreme case whenever the reset gate is switched on. They can also skip subsequences by turning on the update gate.\n",
    "\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Assume that we only want to use the input at time step $t'$ to predict the output at time step $t > t'$. What are the best values for the reset and update gates for each time step?\n",
    "1. Adjust the hyperparameters and analyze the their influence on running time, perplexity, and the output sequence.\n",
    "1. Compare runtime, perplexity, and the output strings for `rnn.RNN` and `rnn.GRU` implementations with each other.\n",
    "1. What happens if you implement only parts of a GRU, e.g., with only a reset gate or only an update gate?\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Java",
   "language": "java",
   "name": "java"
  },
  "language_info": {
   "codemirror_mode": "java",
   "file_extension": ".jshell",
   "mimetype": "text/x-java-source",
   "name": "Java",
   "pygments_lexer": "java",
   "version": "14.0.2+12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}