notebooks/differentiable-indirection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Differentiable Indirection\n",
    "\n",
    "In computer programming, [indirection](https://en.wikipedia.org/wiki/Indirection) (also called dereferencing) is the ability to reference something using a name, reference, or container instead of the value itself.\n",
    "\n",
    "Continious and differentiable implementation of indexing can be divided into two types. In the first type, the data decides how it wants to be addressed. An example of this would be [Transformers](http://papers.nips.cc/paper/7181-attention-is-all-you-need.pdf), where the query and key vectors are generated from the data itself. It can also choose to take no extra information from the data itself, e.g. [Neural Turing Machines](https://arxiv.org/pdf/1410.5401.pdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Linked List\n",
    "\n",
    "In this example we implement a circular [linked list](https://en.wikipedia.org/wiki/Linked_list) with differentiable indirection.\n",
    "\n",
    "![circular linked list](./images/525px-Circularly-linked-list.svg.png)\n",
    "\n",
    "Linked lists are structures which have data and pointer to the next memory location for the next data and pointer is stored."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can think of the RAM as an indexable array of data. Then each pointer address is essentially an integer index for this array.\n",
    "\n",
    "The structs are implicitly defined using two arrays. The data array and the index to the next pointer array.\n",
    "\n",
    "$$ data = [12, 37, 99] $$\n",
    "$$ pointers_{next} = [2, 0, 1] $$\n",
    "\n",
    "The $pointers_{next}$ can be used to traverse the linked list to get the required order of elements.\n",
    "\n",
    "$$ pointers_{next}[0] = 2 $$\n",
    "$$ pointers_{next}[pointers_{next}[0]] = 1 $$\n",
    "$$ pointers_{next}[pointers_{next}[pointers_{next}[0]]] = 0 $$\n",
    "\n",
    "In an non-differentiable setting, we can traverse the linked list like so\n",
    "\n",
    "```python\n",
    "ptr = 0\n",
    "for _ in range(3):\n",
    "    print(data[ptr])\n",
    "    ptr = pointers_next[ptr]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Superpositioned Indexes\n",
    "\n",
    "As discussed in [differentiable-indexed-arrays.ipynb](differentiable-indexed-arrays.ipynb), there are many ways to implement a continuous indexing of arrays. In this example we would use the superposition lookup. \n",
    "\n",
    "Our pointers are one hot vectors corresponding to the index.\n",
    "\n",
    "$$ element_i = data \\cdot p $$\n",
    "$$ element_i = \n",
    "\\begin{bmatrix}\n",
    "12 & 37 & 99\n",
    "\\end{bmatrix} \n",
    "\\begin{bmatrix}\n",
    "1\\\\\n",
    "0\\\\\n",
    "0\n",
    "\\end{bmatrix} \n",
    " = 12 \n",
    "$$\n",
    "\n",
    "This allows for a continuous lookup of elements\n",
    "\n",
    "$$ element_i = \n",
    "\\begin{bmatrix}\n",
    "12 & 37 & 99\n",
    "\\end{bmatrix} \n",
    "\\begin{bmatrix}\n",
    "0.5\\\\\n",
    "0\\\\\n",
    "0.5\n",
    "\\end{bmatrix} \n",
    " = 55.5 \n",
    "$$\n",
    "\n",
    "Similarly, the index array which looked like\n",
    "\n",
    "$$\n",
    "pointers_{next} = \\begin{bmatrix}\n",
    "2 & 0 & 1\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "After converting to one hot vectors, looks like\n",
    "\n",
    "$$\n",
    "pointers_{next} = P = \\begin{bmatrix}\n",
    "0 & 0 & 1\\\\\n",
    "1 & 0 & 0\\\\\n",
    "0 & 1 & 0\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "Thus we can see\n",
    "\n",
    "$$\n",
    "p_{i+1} = p_i P\n",
    "$$\n",
    "\n",
    "$$\n",
    "p_1 = p_0 P =\n",
    "\\begin{bmatrix}\n",
    "1 & 0 & 0\n",
    "\\end{bmatrix} \n",
    "\\begin{bmatrix}\n",
    "0 & 0 & 1\\\\\n",
    "1 & 0 & 0\\\\\n",
    "0 & 1 & 0\n",
    "\\end{bmatrix}\n",
    "=\n",
    "\\begin{bmatrix}\n",
    "0 & 0 & 1\n",
    "\\end{bmatrix} \n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([1. 2. 3.], shape=(3,), dtype=float32)\n",
      "tf.Tensor([1. 1. 1.], shape=(3,), dtype=float32)\n",
      "tf.Tensor(\n",
      "[[3. 4. 5.]\n",
      " [0. 0. 0.]\n",
      " [1. 3. 2.]], shape=(3, 3), dtype=float32)\n",
      "tf.Tensor(0.0, shape=(), dtype=float32)\n",
      "tf.Tensor(\n",
      "[[0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]], shape=(3, 3), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def iterate_over(data, nexts):\n",
    "    data_len = tf.shape(data)[0]\n",
    "    P = nexts\n",
    "    p = tf.one_hot([0], data_len)\n",
    "    \n",
    "    x = tf.expand_dims(data, -1)\n",
    "    y_ = tf.zeros((data_len))\n",
    "    eye = tf.eye(data_len)\n",
    "    \n",
    "    for i in tf.range(data_len):\n",
    "        # The @ token denotes matrix multiplication\n",
    "        x_scalar = tf.squeeze(p @ x)\n",
    "        y_ += eye[i] * x_scalar\n",
    "        \n",
    "        p = p @ P\n",
    "\n",
    "    return y_\n",
    "\n",
    "data  = tf.Variable([1, 3, 2], dtype=tf.float32)\n",
    "target = tf.Variable([1, 2, 3], dtype=tf.float32)\n",
    "data_len = tf.shape(data)[0]\n",
    "nexts = tf.Variable(tf.one_hot([2, 0, 1], data_len), dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    result = iterate_over(data, nexts)\n",
    "    loss = tf.nn.l2_loss(result - target)\n",
    "    \n",
    "print(result)\n",
    "print(tape.gradient(result, data))\n",
    "print(tape.gradient(result, nexts))\n",
    "\n",
    "print(loss)\n",
    "print(tape.gradient(loss, nexts))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Toy problem: Find a pointer array which when traversed, leads to a particular permutation\n",
    "\n",
    "The goal of this toy problem is not to find a permutation matrix $P$ such that $y = Px$ but instead find a matrix such that $y_i = pP^ix$ where $i$ is the $i^{th}$ element in the linked list and $p=onehot(0)$\n",
    "\n",
    "For circular linked lists, there exists some cycle length $n$ where $P^n = I$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(5, 5), dtype=float32, numpy=\n",
       "array([[1., 0., 0., 0., 0.],\n",
       "       [0., 1., 0., 0., 0.],\n",
       "       [0., 0., 1., 0., 0.],\n",
       "       [0., 0., 0., 1., 0.],\n",
       "       [0., 0., 0., 0., 1.]], dtype=float32)>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Q = tf.one_hot([2, 4, 1, 0, 3], 5)\n",
    "Q @ Q @ Q @ Q @ Q"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loss function\n",
    "\n",
    "We want the predicted $\\bar{y}$ to match the real $y$ after one complete traversal of the linked list. So, we add an L2 loss $ | \\bar{y} - y | $. However, this causes the network to learn a $P$ matrix which is a linear combination of the $x$ instead of learning a permutation matrix. Thus we need to add more losses to make sure that the $P$ matrix is a permutation matrix.\n",
    "\n",
    "For permutation matrix loss, we have the following rules\n",
    "* All columns must add up to 1\n",
    "* All rows must add up to 1\n",
    "* All elements must be close to either 0 or 1 (bistable loss)\n",
    "* Cycle loss $P^n = I$ (if cyclic)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "More information about `bistable_loss` can be found [here](notebooks/boolean-satisfiability.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from library.loss import bistable_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(0.0, shape=(), dtype=float32)\n",
      "tf.Tensor(0.0, shape=(), dtype=float32)\n",
      "tf.Tensor(4.0, shape=(), dtype=float32)\n",
      "tf.Tensor(13.0, shape=(), dtype=float32)\n",
      "tf.Tensor(2.9646, shape=(), dtype=float32)\n",
      "tf.Tensor(0.75, shape=(), dtype=float32)\n",
      "tf.Tensor(2.0, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def permute_matrix_loss(P, cycle_length=1, cycle_weight=0):\n",
    "    loss = 0\n",
    "    \n",
    "    P_square = tf.math.square(P)\n",
    "    axis_1_sum = tf.reduce_sum(P_square, axis=1)\n",
    "    axis_0_sum = tf.reduce_sum(P_square, axis=0)\n",
    "    \n",
    "    # Penalize axes not adding up to one\n",
    "    loss += tf.nn.l2_loss(axis_1_sum - 1)\n",
    "    loss += tf.nn.l2_loss(axis_0_sum - 1)\n",
    "    \n",
    "    # Penalize numbers outside [0, 1]\n",
    "    loss += tf.math.reduce_sum(bistable_loss(P))\n",
    "    \n",
    "    # Cycle loss\n",
    "    Q = P\n",
    "    for _ in tf.range(cycle_length - 1):\n",
    "        Q = P @ Q\n",
    "    cycle_loss = tf.nn.l2_loss(Q - tf.eye(tf.shape(Q)[0]))\n",
    "    loss += cycle_loss * cycle_weight\n",
    "    \n",
    "    return loss\n",
    "\n",
    "test1 = tf.constant([\n",
    "    [1,0,0],\n",
    "    [0,1,0],\n",
    "    [0,0,1]\n",
    "],dtype=tf.float32)\n",
    "\n",
    "test2 = tf.constant([\n",
    "    [0,1,0],\n",
    "    [1,0,0],\n",
    "    [0,0,1]\n",
    "],dtype=tf.float32)\n",
    "\n",
    "test3 = tf.constant([\n",
    "    [-1, 0, 0],\n",
    "    [0, 1, 0],\n",
    "    [0, 0, 1],\n",
    "],dtype=tf.float32)\n",
    "\n",
    "test4 = tf.constant([\n",
    "    [2, 0, 0],\n",
    "    [0, 1, 0],\n",
    "    [0, 0, 1],\n",
    "],dtype=tf.float32)\n",
    "\n",
    "test5 = tf.constant([\n",
    "    [0.1, 0, 0],\n",
    "    [0, 0.1, 0],\n",
    "    [0, 0, 0.1],\n",
    "],dtype=tf.float32)\n",
    "\n",
    "test6 = tf.constant([\n",
    "    [0.5, 0.5, 0],\n",
    "    [0.5, 0.5, 0],\n",
    "    [0, 0, 1],\n",
    "],dtype=tf.float32)\n",
    "\n",
    "test7 = tf.constant([\n",
    "    [0, 1, 0],\n",
    "    [1, 0, 0],\n",
    "    [0, 0, 1],\n",
    "],dtype=tf.float32)\n",
    "\n",
    "print(permute_matrix_loss(test1))\n",
    "print(permute_matrix_loss(test2))\n",
    "print(permute_matrix_loss(test3))\n",
    "print(permute_matrix_loss(test4))\n",
    "print(permute_matrix_loss(test5))\n",
    "print(permute_matrix_loss(test6))\n",
    "print(permute_matrix_loss(test7, 1, 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "For convergence, a [softmax](https://en.wikipedia.org/wiki/Softmax_function) operation on $P$ is critical before traversing and computing loss. `TODO: Why?`\n",
    "\n",
    "Initializing $P$ with any invalid permutation matrix leads to a faster convergence than valid permutation matrix. `TODO: WHY?`\n",
    "\n",
    "In order to make sure that our matrix $P$ is learning a permutation matrix and not linear combination of the input, we also print a defuzzified result `y_defuzz`. It is generated by taking the `argmax` of $P$ and iterating over it again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|   loss  |   y_pred  | y_defuzz |   P_pred  |   P_actual   |\n",
      "8.7374239 [1 3 3 3 3] [1 3 3 3 3] [1 1 1 1 1] [2, 4, 1, 0, 3]\n",
      "5.66760302 [1 2 3 4 4] [1 2 3 5 5] [2 3 1 3 3] [2, 4, 1, 0, 3]\n",
      "3.87311959 [1 2 3 4 4] [1 2 3 4 5] [2 4 1 3 3] [2, 4, 1, 0, 3]\n",
      "1.50285113 [1 2 3 4 5] [1 2 3 4 5] [2 4 1 0 3] [2, 4, 1, 0, 3]\n",
      "0.205469772 [1 2 3 4 5] [1 2 3 4 5] [2 4 1 0 3] [2, 4, 1, 0, 3]\n",
      "0.0740788057 [1 2 3 4 5] [1 2 3 4 5] [2 4 1 0 3] [2, 4, 1, 0, 3]\n",
      "0.0350763313 [1 2 3 4 5] [1 2 3 4 5] [2 4 1 0 3] [2, 4, 1, 0, 3]\n",
      "0.0185657572 [1 2 3 4 5] [1 2 3 4 5] [2 4 1 0 3] [2, 4, 1, 0, 3]\n",
      "0.0103952968 [1 2 3 4 5] [1 2 3 4 5] [2 4 1 0 3] [2, 4, 1, 0, 3]\n",
      "0.00600781338 [1 2 3 4 5] [1 2 3 4 5] [2 4 1 0 3] [2, 4, 1, 0, 3]\n",
      "[[-2.18890929 -2.56413817 4.51624966 -2.87268877 -3.21818304]\n",
      " [-2.38402414 -2.36196756 -2.74159765 -1.29020119 4.7228055]\n",
      " [-2.63246131 5.06587076 -2.92895341 -2.58874464 -2.98058629]\n",
      " [3.2075038 -3.05285668 -3.33832383 -1.79198897 -2.805336]\n",
      " [-2.95142794 -2.45104289 -3.21477652 4.53860283 -2.06163]]\n"
     ]
    }
   ],
   "source": [
    "opt = tf.keras.optimizers.Adam()\n",
    "\n",
    "@tf.function\n",
    "def train_step(data, nexts, target_data):\n",
    "    data_length = tf.shape(data)[0]\n",
    "    \n",
    "    with tf.GradientTape() as tape:\n",
    "        nextss = tf.nn.softmax(nexts, axis=1)\n",
    "        actual_data = iterate_over(data, nextss)\n",
    "        loss = tf.nn.l2_loss(actual_data - target_data)\n",
    "        loss += permute_matrix_loss(nextss, data_length, 1)\n",
    "    \n",
    "    grads = tape.gradient(loss, nexts)\n",
    "    opt.apply_gradients(zip([grads], [nexts]))\n",
    "    \n",
    "    return loss, actual_data\n",
    "\n",
    "data  = tf.constant([1, 3, 2, 5, 4], dtype=tf.float32)\n",
    "target_data  = tf.constant([1, 2, 3, 4, 5], dtype=tf.float32)\n",
    "data_len = tf.shape(data)[0]\n",
    "# nexts = tf.Variable(tf.one_hot([2, 4, 1, 0, 3], data_len), dtype=tf.float32)\n",
    "nexts = tf.Variable(tf.one_hot([1, 1, 1, 1, 1], data_len), dtype=tf.float32)\n",
    "# nexts = tf.Variable(tf.one_hot([0,0,0,0,0], data_len), dtype=tf.float32)\n",
    "# nexts = tf.Variable(tf.one_hot([4, 4, 4, 4, 4], data_len), dtype=tf.float32)\n",
    "# nexts = tf.Variable(tf.one_hot([3, 3, 3, 3, 3], data_len), dtype=tf.float32)\n",
    "# nexts = tf.Variable(tf.one_hot([1, 2, 3, 4, 5], data_len), dtype=tf.float32)\n",
    "# nexts = tf.Variable(tf.random.uniform((data_len, data_len), 0, 1))\n",
    "\n",
    "tf.print('|   loss  |   y_pred  | y_defuzz |   P_pred  |   P_actual   |')\n",
    "for i in range(10000):\n",
    "    loss, actual_data = train_step(data, nexts, target_data)\n",
    "    if i % 1000 == 0:\n",
    "        argmax_next = tf.argmax(nexts, 1)\n",
    "        defuzzified = tf.one_hot(argmax_next, data_len)\n",
    "        defuzzified_data = iterate_over(data, defuzzified)\n",
    "        tf.print(loss, tf.round(actual_data), defuzzified_data, argmax_next, [2, 4, 1, 0, 3])\n",
    "        \n",
    "tf.print(nexts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Verifying cyclic permutation\n",
    "\n",
    "We can see that $P^n = I$ for both normal and defuzzified cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2 4 1 0 3]\r\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(5, 5), dtype=float32, numpy=\n",
       "array([[1., 0., 0., 0., 0.],\n",
       "       [0., 1., 0., 0., 0.],\n",
       "       [0., 0., 1., 0., 0.],\n",
       "       [0., 0., 0., 1., 0.],\n",
       "       [0., 0., 0., 0., 1.]], dtype=float32)>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "P = tf.nn.softmax(nexts, axis=1)\n",
    "tf.print(tf.argmax(P,1))\n",
    "tf.round(P @ P @ P @ P @ P)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(5, 5), dtype=float32, numpy=\n",
       "array([[1., 0., 0., 0., 0.],\n",
       "       [0., 1., 0., 0., 0.],\n",
       "       [0., 0., 1., 0., 0.],\n",
       "       [0., 0., 0., 1., 0.],\n",
       "       [0., 0., 0., 0., 1.]], dtype=float32)>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "argmax_next = tf.argmax(P, 1)\n",
    "DQ = tf.one_hot(argmax_next, data_len)\n",
    "DQ @ DQ @ DQ @ DQ @ DQ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}