Chapter 10 - Modern Recurrent Neural Networks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1f0c9f1f",
   "metadata": {},
   "source": [
    "# Chapter 10: Modern Recurrent Neural Networks\n",
    "\n",
    "### Dhuvi karthikeyan\n",
    "\n",
    "03/01/2023"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1360f84d",
   "metadata": {},
   "source": [
    "* Long-short term memory (LSTM) architecture was first produced in 1997 and has been one of the two most important development in sequence modeling.\n",
    "* Bidirectional recurrent neural networks also introduced in 1997 are the result of the second development which resulted in a huge boost in performance. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70724d83",
   "metadata": {},
   "source": [
    "## 10.1 Long Short-Term Memory (LSTM)\n",
    "\n",
    "While gradient clipping helped control the issue of exploding gradients, the issue of vanishing gradients and the modeling of long-term dependencies in the sequence learning was addressed with a more complicated architectural variation. \n",
    "\n",
    "### 10.1.1 Gated Memory Cells\n",
    "\n",
    "Core to the advancement is the memory cell which replaces the standard Recurrent node. The recurrent node being the unit of computation that takes in a hidden state from time step t-1 and input token at time t and passes the output and hidden state. The RNN has parameter sharing where the parameters are the same at every recurrent cell. \n",
    "\n",
    "Memory Cells consist of:\n",
    "* Internal state C_t that handles information flow by allowing 1 of three things to happend to the input:\n",
    "    1. Input gate: Allow input to enter internal state\n",
    "    2. Forget gate: Reduce internal state to 0\n",
    "    3. Output gate: Allow internal state to affect output\n",
    "\n",
    "\n",
    "**Note:** Gating the hidden state allows us to explicitly allow for mechanisms to decide when the hidden state should be updated, remain the same, or reset.\n",
    "\n",
    "#### Input Gate, Forget Gate, Output Gate:\n",
    "\n",
    "With the input token and the hidden state of the previous timestep, the learnable Forget gate, input gate and output gate are all learnable. This means that theres associated weights and biases with this part of the model. In practice each of these gates is its own MLP with sigmoid activations. \n",
    "\n",
    "* Input node: a tanh activation function that learns how much of the inputs to incldue in the internal state.\n",
    "\n",
    "$$ I_t = \\sigma (X_tW_{xi} + H_{t-1}W_{hi} + b_i)$$\n",
    "$$ F_t = \\sigma (X_tW_{xf} + H_{t-1}W_{hf} + b_f)$$\n",
    "$$ O_t = \\sigma (X_tW_{xo} + H_{t-1}W_{ho} + b_o)$$\n",
    "\n",
    "These gates all have dimensionality of $\\mathbb{R}^{n x h}$ for batch size n and hidden dimension h. The input weight matrices are all $\\in \\mathbb{R}^{d x h}$ whereas the hidden state matrices are all $\\in \\mathbb{R}^{h x h}$. The sigmoid activation is used to squash values to 0,1. \n",
    "\n",
    "#### Input Node\n",
    "\n",
    "The input node is defined as **C** and can be thought of as the context with a tanh activation function mapping values to (-1, 1) . \n",
    "\n",
    "$$ \\hat{C_t} = tanh(X_t W_{xc} + H_{t-1}W_{hc} + b_c)$$\n",
    "\n",
    "At this point the hidden state and input are passed as input to the input node, the forget gate, input gate, and output gate. Each of which computes its output through a different activation function but similar method of adding the input and hidden states in h dimensional tensors. The input node is designed to scale and flip the sign of the coming input. \n",
    "\n",
    "#### Memory Cell Internal State\n",
    "\n",
    "$$ C_t = F_t \\odot C_{t-1} + I_t \\odot \\hat{C_t}$$\n",
    "\n",
    "Input gate $I_t$ dictates how much of the current input token to allow into the internal state and the forget gate controls how much of the previous internal state to use in computation of the current internal state. \n",
    "\n",
    "#### Hidden State\n",
    "\n",
    "The output of the memory cell are the internal state of the memory cell along with the hidden state of the sequence. \n",
    "\n",
    "$$ H_t = O_t \\odot tanh(C_t) $$\n",
    "\n",
    "The tanh nonlinearity in conjunction wiht the sigmoid applied on the output gate ensures the values of H_t are between (-1,1). This flow of information results in the memory cell being able to selectively affect the network with the output gate being close to zero for much of the early timesteps and then flipping the switch to go from zero to one and include more of the internal state in the hidden state.\n",
    "\n",
    "### 10.1.4 Summary\n",
    "\n",
    "LSTMS are the archetypal latent variable autoregressive model with nontrivial state control. Although they do have costly training due to long range dependency. Proposed in 1997 but took off in the 2000s, becoming the de-facto sequence model for much of 2010s until transformers. While the internal state and hidden state are used by hidden layers, only the latter is passed to the output layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "id": "2f6c6cb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch\n",
    "import torch.functional as F\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "class LSTMcell(nn.Module):\n",
    "    \"\"\"\n",
    "    Implement the LSTM cell from \"scratch\".\n",
    "    s/o to Torch community for handling the \n",
    "    heavy lifting.\n",
    "    \"\"\"\n",
    "    def __init__(self, input_dim, hidden_dim, sigma=.01):\n",
    "        super(LSTMcell, self).__init__()\n",
    "        # Save H-params\n",
    "        self.input_dim = input_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.sigma = sigma\n",
    "        # Initialize Net Parameters\n",
    "        self.W_xi, self.W_hi, self.b_i = self.gate_init()\n",
    "        self.W_xf, self.W_hf, self.b_f = self.gate_init()\n",
    "        self.W_xo, self.W_ho, self.b_o = self.gate_init()\n",
    "        self.W_xc, self.W_hc, self.b_c = self.gate_init()\n",
    "    \n",
    "    def gate_init(self):\n",
    "        gate_params =  (self.weight_init(self.input_dim, self.hidden_dim),\n",
    "                self.weight_init(self.hidden_dim, self.hidden_dim),\n",
    "                self.bias_init())\n",
    "        return gate_params\n",
    "    \n",
    "    def weight_init(self, *shape):\n",
    "        return nn.Parameter(torch.randn(shape)*self.sigma)\n",
    "    \n",
    "    def bias_init(self):\n",
    "        return nn.Parameter(torch.zeros(self.hidden_dim))\n",
    "    \n",
    "    def init_hidden(self):\n",
    "        return (torch.zeros(self.hidden_dim), torch.zeros(self.hidden_dim))\n",
    "    \n",
    "    def forward(self, x_t, state):\n",
    "        h_t, c_t = state\n",
    "        \n",
    "        input_gate = torch.sigmoid(torch.matmul(x_t, self.W_xi)\n",
    "                                  + torch.matmul(h_t, self.W_hi)\n",
    "                                  + self.b_i)\n",
    "        forget_gate = torch.sigmoid(torch.matmul(x_t, self.W_xf)\n",
    "                                  + torch.matmul(h_t, self.W_hf)\n",
    "                                  + self.b_f)\n",
    "        output_gate = torch.sigmoid(torch.matmul(x_t, self.W_xo)\n",
    "                                  + torch.matmul(h_t, self.W_ho)\n",
    "                                  + self.b_o)\n",
    "        input_node = torch.sigmoid(torch.matmul(x_t, self.W_xc)\n",
    "                                  + torch.matmul(h_t, self.W_hc)\n",
    "                                  + self.b_c)\n",
    "\n",
    "        c_t = forget_gate * c_t + input_gate * input_node\n",
    "        h_t = output_gate * torch.tanh(c_t)\n",
    "        \n",
    "        return (h_t, c_t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "id": "b2bcb248",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LSTM(nn.Module):\n",
    "    \"\"\"\n",
    "    Implement a Deep LSTM that can take an arbitrary number of LSTM cells.\n",
    "    \"\"\"\n",
    "    def __init__(self, input_dim, hidden_dim, m_layers):\n",
    "        super(LSTM, self).__init__()\n",
    "        self.input_dim = input_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.m_layers = m_layers\n",
    "        self.layers = [LSTMcell(input_dim, hidden_dim)] + [LSTMcell(hidden_dim, hidden_dim)]*(m_layers-1)\n",
    "        \n",
    "        \n",
    "    def forward(self, X):\n",
    "        outs = []\n",
    "        # Initialize the hidden states as a list of tuples\n",
    "        hiddens = [(torch.zeros((X.shape[1], self.hidden_dim)), torch.zeros((X.shape[1], self.hidden_dim)))]*self.m_layers \n",
    "        \n",
    "        # Perform recurrent pass\n",
    "        for i, x_t in enumerate(X):\n",
    "            for j, layer in enumerate(self.layers):\n",
    "                h_t, c_t = layer(x_t, hiddens[j])\n",
    "                # Update the hidden states for that layer @this timestep\n",
    "                hiddens[j] = (h_t, c_t)\n",
    "                # Update the input to the next layer @this timestep\n",
    "                x_t = h_t\n",
    "            # At the end of layer append the hidden state (output)\n",
    "            outs.append(h_t)\n",
    "        return torch.stack(outs, dim=0), (h_t, c_t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "id": "d4cab366",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test the LSTMCell\n",
    "rnn = LSTMcell(10, 20) # (input_size, hidden_size)\n",
    "inputs = torch.randn(5, 3, 10) # (time_steps, batch, input_size)\n",
    "hx = torch.randn(3, 20) # (batch, hidden_size)\n",
    "cx = torch.randn(3, 20)\n",
    "output = []\n",
    "\n",
    "for i in range(inputs.size()[0]):\n",
    "    hx, cx = rnn(inputs[i], (hx,cx))\n",
    "    output.append(hx)\n",
    "output = torch.stack(output, dim=0)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "id": "bf118f0c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[0.1214, 0.1232, 0.1220, 0.1228, 0.1231, 0.1231, 0.1215, 0.1219,\n",
       "           0.1231, 0.1223, 0.1217, 0.1222, 0.1242, 0.1228, 0.1227, 0.1214],\n",
       "          [0.1214, 0.1232, 0.1220, 0.1228, 0.1231, 0.1230, 0.1215, 0.1219,\n",
       "           0.1231, 0.1223, 0.1218, 0.1223, 0.1242, 0.1228, 0.1227, 0.1214],\n",
       "          [0.1214, 0.1232, 0.1220, 0.1228, 0.1231, 0.1230, 0.1215, 0.1219,\n",
       "           0.1231, 0.1223, 0.1217, 0.1223, 0.1242, 0.1228, 0.1227, 0.1214]],\n",
       " \n",
       "         [[0.1760, 0.1813, 0.1779, 0.1802, 0.1813, 0.1809, 0.1785, 0.1779,\n",
       "           0.1792, 0.1782, 0.1774, 0.1785, 0.1815, 0.1793, 0.1798, 0.1772],\n",
       "          [0.1760, 0.1813, 0.1780, 0.1802, 0.1813, 0.1809, 0.1786, 0.1779,\n",
       "           0.1792, 0.1783, 0.1774, 0.1785, 0.1815, 0.1793, 0.1797, 0.1773],\n",
       "          [0.1760, 0.1813, 0.1779, 0.1802, 0.1813, 0.1809, 0.1785, 0.1779,\n",
       "           0.1792, 0.1783, 0.1774, 0.1785, 0.1815, 0.1793, 0.1798, 0.1772]],\n",
       " \n",
       "         [[0.2008, 0.2089, 0.2040, 0.2073, 0.2091, 0.2086, 0.2052, 0.2038,\n",
       "           0.2053, 0.2041, 0.2035, 0.2046, 0.2080, 0.2058, 0.2066, 0.2033],\n",
       "          [0.2008, 0.2089, 0.2040, 0.2073, 0.2091, 0.2085, 0.2052, 0.2038,\n",
       "           0.2053, 0.2042, 0.2035, 0.2046, 0.2080, 0.2058, 0.2066, 0.2034],\n",
       "          [0.2007, 0.2089, 0.2040, 0.2073, 0.2091, 0.2085, 0.2052, 0.2038,\n",
       "           0.2054, 0.2041, 0.2035, 0.2045, 0.2081, 0.2058, 0.2066, 0.2033]],\n",
       " \n",
       "         [[0.2123, 0.2223, 0.2165, 0.2205, 0.2227, 0.2221, 0.2179, 0.2162,\n",
       "           0.2179, 0.2164, 0.2161, 0.2169, 0.2207, 0.2185, 0.2196, 0.2159],\n",
       "          [0.2123, 0.2223, 0.2165, 0.2204, 0.2227, 0.2221, 0.2179, 0.2161,\n",
       "           0.2180, 0.2164, 0.2161, 0.2169, 0.2207, 0.2185, 0.2196, 0.2158],\n",
       "          [0.2123, 0.2223, 0.2165, 0.2205, 0.2227, 0.2221, 0.2179, 0.2161,\n",
       "           0.2180, 0.2164, 0.2161, 0.2169, 0.2207, 0.2185, 0.2196, 0.2158]],\n",
       " \n",
       "         [[0.2178, 0.2289, 0.2226, 0.2269, 0.2295, 0.2289, 0.2241, 0.2221,\n",
       "           0.2241, 0.2224, 0.2223, 0.2228, 0.2268, 0.2247, 0.2260, 0.2220],\n",
       "          [0.2178, 0.2289, 0.2226, 0.2269, 0.2295, 0.2289, 0.2241, 0.2221,\n",
       "           0.2241, 0.2224, 0.2223, 0.2228, 0.2268, 0.2247, 0.2260, 0.2220],\n",
       "          [0.2177, 0.2289, 0.2226, 0.2269, 0.2295, 0.2289, 0.2241, 0.2221,\n",
       "           0.2241, 0.2224, 0.2223, 0.2228, 0.2269, 0.2247, 0.2259, 0.2220]]],\n",
       "        grad_fn=<StackBackward0>),\n",
       " (tensor([[0.2178, 0.2289, 0.2226, 0.2269, 0.2295, 0.2289, 0.2241, 0.2221, 0.2241,\n",
       "           0.2224, 0.2223, 0.2228, 0.2268, 0.2247, 0.2260, 0.2220],\n",
       "          [0.2178, 0.2289, 0.2226, 0.2269, 0.2295, 0.2289, 0.2241, 0.2221, 0.2241,\n",
       "           0.2224, 0.2223, 0.2228, 0.2268, 0.2247, 0.2260, 0.2220],\n",
       "          [0.2177, 0.2289, 0.2226, 0.2269, 0.2295, 0.2289, 0.2241, 0.2221, 0.2241,\n",
       "           0.2224, 0.2223, 0.2228, 0.2269, 0.2247, 0.2259, 0.2220]],\n",
       "         grad_fn=<MulBackward0>),\n",
       "  tensor([[0.4706, 0.4901, 0.4847, 0.4850, 0.4908, 0.4927, 0.4777, 0.4780, 0.4831,\n",
       "           0.4794, 0.4818, 0.4750, 0.4860, 0.4898, 0.4835, 0.4802],\n",
       "          [0.4706, 0.4901, 0.4847, 0.4850, 0.4908, 0.4927, 0.4777, 0.4780, 0.4831,\n",
       "           0.4794, 0.4818, 0.4751, 0.4860, 0.4899, 0.4835, 0.4803],\n",
       "          [0.4706, 0.4902, 0.4847, 0.4850, 0.4908, 0.4927, 0.4777, 0.4780, 0.4831,\n",
       "           0.4794, 0.4818, 0.4750, 0.4862, 0.4898, 0.4834, 0.4802]],\n",
       "         grad_fn=<AddBackward0>)))"
      ]
     },
     "execution_count": 130,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test the arbitrary LSTM\n",
    "inputs = torch.randn(5, 3, 8)       \n",
    "lm = LSTM(8, 16, 2)\n",
    "lm(inputs)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e041e2f6",
   "metadata": {},
   "source": [
    "## 10.2 Gated Recurrence Units (GRU)\n",
    "\n",
    "Proposed in (2014) the GRU was a simplification on the LSTM designed to speed up computation. \n",
    "\n",
    "### 10.2.1 Reset and Update Gates\n",
    "\n",
    "3-1=2 gates compared to the LSTM. Each of the gates here have the same sigmoid function which restricts their values to (0,1) corresponding to how much the previous hidden state matters.\n",
    "\n",
    "$$ R_t = \\sigma (X_tW_{xr} + H_{t-1}W_{hr} + b_r)$$\n",
    "$$ Z_t = \\sigma (X_tW_{xz} + H_{t-1}W_{hz} + b_z)$$\n",
    "\n",
    "### 10.2.2 Candidate Hidden State\n",
    "\n",
    "$$ \\hat{H_t} = tanh(X_tW_{xh} + (R_t \\odot H_{t-1})W_{hh} + b_h) $$ \n",
    "\n",
    "This candidate hidden state is generated using just tihe information from the previous hidden state (scaled based on the reset gate) along with the input. The last part is to combine it wiht the update gate to get the actual hidden state of the GRU cell.\n",
    "\n",
    "### 10.2.3 Hidden State\n",
    "\n",
    "$$ H_t = Z_t \\odot H_{t-1} + (1 - Z_t) \\odot \\hat{H_t} $$\n",
    "\n",
    "We can see here that the elementwise convex combination of the previous hidden state and the current candidate hidden state via the z_t and 1-z_t update. Learning the fraction z_t is what the network must learn.\n",
    "\n",
    "Resets capture the short term dependencies while updates capture the long term dependencies.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "7e2b83d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GRUcell(nn.Module):\n",
    "    \"\"\"\n",
    "    Implement the GRU cell from \"scratch\".\n",
    "    s/o to Torch community for handling the \n",
    "    heavy lifting.\n",
    "    \"\"\"\n",
    "    def __init__(self, input_dim, hidden_dim, sigma=.01):\n",
    "        super(GRUcell, self).__init__()\n",
    "        # Save H-params\n",
    "        self.input_dim = input_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.sigma = sigma\n",
    "        # Initialize Net Parameters\n",
    "        self.W_xr, self.W_hr, self.b_r = self.gate_init()\n",
    "        self.W_xz, self.W_hz, self.b_z = self.gate_init()\n",
    "        self.W_xh, self.W_hh, self.b_h = self.gate_init()\n",
    "        \n",
    "    def gate_init(self):\n",
    "        gate_params =  (self.weight_init(self.input_dim, self.hidden_dim),\n",
    "                self.weight_init(self.hidden_dim, self.hidden_dim),\n",
    "                self.bias_init())\n",
    "        return gate_params\n",
    "    \n",
    "    def weight_init(self, *shape):\n",
    "        return nn.Parameter(torch.randn(shape)*self.sigma)\n",
    "    \n",
    "    def bias_init(self):\n",
    "        return nn.Parameter(torch.zeros(self.hidden_dim))\n",
    "    \n",
    "    def init_hidden(self):\n",
    "        return (torch.zeros(self.hidden_dim), torch.zeros(self.hidden_dim))\n",
    "    \n",
    "    def forward(self, x_t, h_t):\n",
    "        \n",
    "        reset_gate = torch.sigmoid(torch.matmul(x_t, self.W_xr)\n",
    "                                  + torch.matmul(h_t, self.W_hr)\n",
    "                                  + self.b_r)\n",
    "        update_gate = torch.sigmoid(torch.matmul(x_t, self.W_xz)\n",
    "                                  + torch.matmul(h_t, self.W_hz)\n",
    "                                  + self.b_z)\n",
    "        candidate_h = torch.tanh(torch.matmul(x_t, self.W_xh)\n",
    "                                  + torch.matmul((update_gate * h_t), self.W_hh)\n",
    "                                  + self.b_h)\n",
    "        \n",
    "        h_t = update_gate * h_t + (1-update_gate) * candidate_h\n",
    "        \n",
    "        return h_t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "id": "4dd5726c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GRU(nn.Module):\n",
    "    \"\"\"\n",
    "    Implement a Deep GRU that can take an arbitrary number of GRU cells.\n",
    "    \"\"\"\n",
    "    def __init__(self, input_dim, hidden_dim, m_layers):\n",
    "        super(GRU, self).__init__()\n",
    "        self.input_dim = input_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.m_layers = m_layers\n",
    "        self.layers = [GRUcell(input_dim, hidden_dim)] + [GRUcell(hidden_dim, hidden_dim)]*(m_layers-1)\n",
    "        \n",
    "        \n",
    "    def forward(self, X):\n",
    "        outs = []\n",
    "        # Initialize the hidden states as a list of tuples\n",
    "        hiddens = [torch.zeros((X.shape[1], self.hidden_dim))]*self.m_layers \n",
    "        \n",
    "        # Perform recurrent pass\n",
    "        for i, x_t in enumerate(X):\n",
    "            for j, layer in enumerate(self.layers):\n",
    "                h_t = layer(x_t, hiddens[j])\n",
    "                # Update the hidden states for that layer @this timestep\n",
    "                hiddens[j] = h_t\n",
    "                # Update the input to the next layer @this timestep\n",
    "                x_t = h_t\n",
    "            # At the end of layer append the hidden state (output)\n",
    "            outs.append(h_t)\n",
    "        return torch.stack(outs, dim=0), h_t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "id": "47e53482",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[-6.1351e-04,  2.1541e-04, -1.8708e-04,  8.3099e-05, -2.9001e-04,\n",
       "            2.7070e-04,  3.4477e-04,  1.5560e-04,  3.8333e-05,  4.6603e-04,\n",
       "            6.3583e-04,  1.1855e-04, -8.6820e-05, -2.6103e-04, -4.8864e-04,\n",
       "           -1.1058e-05],\n",
       "          [-4.9974e-04,  1.8529e-04, -2.4628e-04,  3.4985e-04, -2.5507e-04,\n",
       "            2.9341e-04,  6.7379e-05, -9.8408e-05,  1.6230e-04,  5.6795e-04,\n",
       "            8.8003e-04, -1.1572e-04, -5.0920e-05, -2.6074e-04, -3.8427e-04,\n",
       "           -9.8128e-05],\n",
       "          [-1.1470e-04,  1.2553e-04, -4.1958e-05,  1.7969e-04, -1.4367e-04,\n",
       "            1.0470e-04, -3.2690e-05,  6.4227e-05,  1.1511e-04,  1.7195e-04,\n",
       "            2.5168e-04, -8.4800e-05, -8.9615e-05, -2.3025e-04,  3.6465e-05,\n",
       "           -1.3526e-04]],\n",
       " \n",
       "         [[-3.9307e-04,  1.8178e-04, -6.8292e-04,  4.4217e-04, -1.9730e-04,\n",
       "           -2.7677e-05,  3.2114e-04,  8.6098e-04,  3.2028e-04,  1.4633e-04,\n",
       "            1.6423e-04,  1.5344e-04, -4.0722e-04, -7.2992e-04, -1.7060e-04,\n",
       "            4.8205e-05],\n",
       "          [ 3.8441e-04, -4.9425e-04, -2.2993e-04,  1.0948e-04,  1.1909e-05,\n",
       "            3.6348e-04, -2.1107e-04, -9.6491e-05,  2.4946e-04,  2.7302e-06,\n",
       "            8.3686e-04, -3.1962e-05, -1.3385e-04, -1.8463e-04, -5.7846e-04,\n",
       "            6.1344e-04],\n",
       "          [ 7.0126e-04, -3.2190e-04, -2.5857e-04,  3.8488e-04,  1.5990e-05,\n",
       "           -1.5779e-04, -1.4874e-04,  3.4264e-04,  2.2098e-04, -3.7849e-04,\n",
       "           -7.1480e-04, -1.7400e-04, -6.6910e-05, -2.0981e-04,  8.5368e-04,\n",
       "            4.2285e-04]],\n",
       " \n",
       "         [[-8.3934e-04,  2.6431e-04,  1.4632e-05, -2.2593e-04, -4.2948e-04,\n",
       "            2.7919e-04, -5.5903e-05,  8.9361e-05,  3.2008e-05,  2.7320e-04,\n",
       "            2.4301e-04,  2.6426e-04, -1.0594e-04,  8.4984e-05, -1.9812e-04,\n",
       "           -2.0145e-04],\n",
       "          [ 1.2920e-03, -7.7082e-04,  2.1631e-04, -2.5095e-04,  3.8753e-04,\n",
       "            2.9257e-04, -4.7690e-04, -1.2168e-04,  9.7643e-05, -5.9894e-04,\n",
       "            2.6088e-04, -8.0958e-05, -3.3962e-05,  6.5514e-06, -1.7511e-04,\n",
       "            7.6452e-04],\n",
       "          [ 6.4895e-04, -3.0780e-04, -3.5833e-04,  4.0782e-04,  1.8818e-05,\n",
       "           -2.5563e-04,  1.1555e-04,  4.9002e-04,  1.6050e-04, -3.8019e-04,\n",
       "           -9.0529e-04, -1.0926e-04, -4.5675e-05, -2.4772e-04,  8.6187e-04,\n",
       "            5.1891e-04]],\n",
       " \n",
       "         [[-7.1610e-04,  1.4424e-04, -1.0527e-04,  2.3563e-04, -4.4320e-04,\n",
       "            1.8959e-04,  9.1296e-05, -3.9872e-04,  1.0743e-04,  6.5109e-04,\n",
       "            6.9484e-04, -4.9912e-05,  3.5687e-05,  2.1548e-05, -3.5474e-04,\n",
       "           -8.4298e-05],\n",
       "          [ 1.8279e-03, -1.3024e-03,  3.2614e-04, -4.1400e-04,  5.4400e-04,\n",
       "            1.9205e-04, -5.9197e-04, -5.2038e-04,  1.6594e-05, -8.8521e-04,\n",
       "           -6.7602e-05, -9.9064e-05,  1.7245e-04,  4.5979e-04, -7.6508e-05,\n",
       "            1.3154e-03],\n",
       "          [ 1.6642e-05, -1.4200e-04,  5.4837e-05, -2.2670e-05, -1.6926e-04,\n",
       "            1.8038e-04,  2.1870e-04, -1.8341e-04, -1.5163e-04,  1.2613e-04,\n",
       "           -1.9901e-04, -2.3118e-05,  2.7924e-04,  2.7391e-04,  2.8518e-04,\n",
       "            4.1497e-04]],\n",
       " \n",
       "         [[-2.7128e-04, -1.8483e-04,  2.8501e-05, -7.0900e-05, -1.7275e-04,\n",
       "            5.6852e-05, -2.4620e-05, -4.8973e-04,  4.7283e-05,  2.5671e-04,\n",
       "            5.3961e-04,  6.4830e-05,  1.7815e-05,  2.0441e-04, -5.8806e-04,\n",
       "            1.4440e-04],\n",
       "          [ 1.6757e-03, -1.1266e-03,  3.4995e-04, -5.4215e-04,  4.6665e-04,\n",
       "            1.5504e-04, -4.6193e-04, -1.7039e-04, -6.3427e-05, -9.9068e-04,\n",
       "           -5.4697e-04,  3.3137e-05,  1.4164e-04,  4.4116e-04,  2.5253e-04,\n",
       "            1.2231e-03],\n",
       "          [-1.8485e-04,  2.1113e-04,  7.6093e-05, -1.2045e-04, -6.0658e-05,\n",
       "            2.7617e-04,  3.8078e-04,  1.1505e-04, -3.2509e-04,  1.6598e-04,\n",
       "           -3.1489e-04,  2.1333e-05,  3.4537e-04,  2.2762e-04,  3.8583e-04,\n",
       "            1.5764e-04]]], grad_fn=<StackBackward0>),\n",
       " tensor([[-2.7128e-04, -1.8483e-04,  2.8501e-05, -7.0900e-05, -1.7275e-04,\n",
       "           5.6852e-05, -2.4620e-05, -4.8973e-04,  4.7283e-05,  2.5671e-04,\n",
       "           5.3961e-04,  6.4830e-05,  1.7815e-05,  2.0441e-04, -5.8806e-04,\n",
       "           1.4440e-04],\n",
       "         [ 1.6757e-03, -1.1266e-03,  3.4995e-04, -5.4215e-04,  4.6665e-04,\n",
       "           1.5504e-04, -4.6193e-04, -1.7039e-04, -6.3427e-05, -9.9068e-04,\n",
       "          -5.4697e-04,  3.3137e-05,  1.4164e-04,  4.4116e-04,  2.5253e-04,\n",
       "           1.2231e-03],\n",
       "         [-1.8485e-04,  2.1113e-04,  7.6093e-05, -1.2045e-04, -6.0658e-05,\n",
       "           2.7617e-04,  3.8078e-04,  1.1505e-04, -3.2509e-04,  1.6598e-04,\n",
       "          -3.1489e-04,  2.1333e-05,  3.4537e-04,  2.2762e-04,  3.8583e-04,\n",
       "           1.5764e-04]], grad_fn=<AddBackward0>))"
      ]
     },
     "execution_count": 138,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test the arbitrary LSTM\n",
    "inputs = torch.randn(5, 3, 8)       \n",
    "lm = GRU(8, 16, 2)\n",
    "lm(inputs)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95b691b2",
   "metadata": {},
   "source": [
    "## 10.3 Deep Recurrent Neural Networks\n",
    "\n",
    "Depth in recurrent neural networks is often thought of how long the unrolled RNN is O(seq_length). However, depth can also be added in terms of how long the chain from input to output is per timestep. In the latter case, the linear RNN becomes like a matrix and the hidden state from the current timestep gets passed to the next timestep as well as up to the next layer as input. This propagates through the depth of the layers at single timestep and the hidden states are aggregated before passing the vector of hidden states to the next timestep.\n",
    "\n",
    "\n",
    "$$ H_t^{(l)} = \\phi_l (H_t^{(l-1)}W_xh^{(l)} + H_t^{(l)}W_hh^{(l)} + b_h^{(l)})$$\n",
    "\n",
    "Where $H_t^{(0)}$ is $X_t$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62e8e0ac",
   "metadata": {},
   "source": [
    "## 10.4 Bidirectional Recurrent Neural Networks\n",
    "\n",
    "To address the problem of unidirectionality and the limitations it poses, in 1997, the bidirectional RNN was born as an intuitive means of alleviating the aforementioned issue. The implementation of the bi-RNN is simply the combination of two uni-directional RNNs where the first RNN reads the sequence in the forwards direction and the second processes the sequence in the reverse direction. This allows for both contexts to be embedded in the final hidden representation of each NN. Instead of reducing the tensor outputs from each timestep the bi-RNN concates them which ensures that the learned vectors do not destructively interfere with each other. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "id": "475ecd2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BiDirectional(nn.Module):\n",
    "    \"\"\"\n",
    "    Take any implementation of a deep RNN and return the concatenated outputs\n",
    "    \"\"\"\n",
    "    def __init__(self, RNN, input_dim, hidden_dim, n_layers):\n",
    "        super(BiDirectional, self).__init__()\n",
    "        self.fc = RNN(input_dim, hidden_dim, n_layers)\n",
    "        self.rc = RNN(input_dim, hidden_dim, n_layers)\n",
    "        self.hidden_dim = hidden_dim\n",
    "        \n",
    "        \n",
    "    def forward(self, X):\n",
    "        forward_outs, H_f = self.fc(X)\n",
    "        reverse_outs, H_r = self.rc(reversed(X))\n",
    "        concat_outs = torch.cat((forward_outs, reversed(reverse_outs)), dim=2)\n",
    "        return concat_outs, (H_f, H_r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "id": "ad63468e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 3, 32])"
      ]
     },
     "execution_count": 159,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#biLSTM = BiDirectional(LSTM, 8, 16, 10)\n",
    "#biLSTM(inputs)[0].shape\n",
    "\n",
    "biGRU = BiDirectional(GRU, 8, 16, 10)\n",
    "biGRU(inputs)[0].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fab9cca",
   "metadata": {},
   "source": [
    "lm(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "id": "59f94160",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 3, 32])"
      ]
     },
     "execution_count": 145,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cat((lm(inputs)[0], lm(inputs)[0]), dim=2).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "707754af",
   "metadata": {},
   "source": [
    "Dynamic algorithms to get the probilistic sequence:\\\n",
    "    * almost always get EOS\\\n",
    "    * are smaller TCR sequences more likely?\\\n",
    "    * what is the probability of TCR generation\\\n",
    "    * top-k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df7d1092",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}