notebooks/custom-gradient.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom Gradient\n",
    "\n",
    "This is a brief demonstration of tensorflow [custom gradients](https://www.tensorflow.org/api_docs/python/tf/custom_gradient)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chain rule\n",
    "\n",
    "Lets say we have a function $f(x) = x^2$. If we now compose this function such that $y = f(f(f(x)))$. Now we want to find the gradient $\\frac{dy}{dx}$.\n",
    "\n",
    "We first decompose this into\n",
    "\n",
    "\\begin{align*}\n",
    "  y &= x_0^2\\\\\n",
    "  x_0 &= x_1^2\\\\\n",
    "  x_1 &= x^2\n",
    "\\end{align*}\n",
    "\n",
    "On taking first order derivative, we get\n",
    "\n",
    "\\begin{align*}\n",
    "  \\frac{dy}{dx_0} &= 2x_0\\\\\n",
    "  \\frac{dx_0}{dx_1} &= 2x_1\\\\\n",
    "  \\frac{dx_1}{dx} &= 2x\n",
    "\\end{align*}\n",
    "\n",
    "Using chain rule\n",
    "\n",
    "\\begin{equation*}\n",
    "  \\frac{dy}{dx} = \\frac{dy}{dx_0}  \\frac{dx_0}{dx_1}  \\frac{dx_1}{dx}\n",
    "\\end{equation*}\n",
    "\n",
    "To generalize\n",
    "\n",
    "\\begin{equation*}\n",
    "  \\frac{dy}{dx} = \\frac{dy}{dx_{0}}  ...  \\frac{dx_i}{dx_{i+1}}  ...  \\frac{dx_n}{dx}\n",
    "\\end{equation*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In tensorflow the `upstream` gradient is passed as an argument to the inner function `grad`.\n",
    "\n",
    "\\begin{equation*}\n",
    "  upstream = \\frac{dx_{i+1}}{dx_{i+2}}  ...  \\frac{dx_n}{dx}\n",
    "\\end{equation*}\n",
    "\n",
    "Now we can multiply this upstream gradient to the gradient of the current function $\\frac{dx_{i}}{dx_{i+1}}$ and pass it downstream.\n",
    "\n",
    "\\begin{equation*}\n",
    "  downstream = \\frac{dx_{i}}{dx_{i+1}}  * upstream\n",
    "\\end{equation*}\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": [
      "x=2.0\ty=4.0\n",
      "x=4.0\ty=16.0\n",
      "x=16.0\ty=256.0\n",
      "x=16.0\tupstream=1.0\tcurrent=32.0\t\tdownstream=32.0\n",
      "x=4.0\tupstream=32.0\tcurrent=8.0\t\tdownstream=256.0\n",
      "x=2.0\tupstream=256.0\tcurrent=4.0\t\tdownstream=1024.0\n",
      "\n",
      "final dy/dx=1024.0\n"
     ]
    }
   ],
   "source": [
    "# @tf.function\n",
    "@tf.custom_gradient\n",
    "def foo(x):\n",
    "    tf.debugging.assert_rank(x, 0)\n",
    "\n",
    "    def grad(dy_dx_upstream):\n",
    "        dy_dx = 2 * x\n",
    "        dy_dx_downstream = dy_dx * dy_dx_upstream\n",
    "        tf.print(f'x={x}\\tupstream={dy_dx_upstream}\\tcurrent={dy_dx}\\t\\tdownstream={dy_dx_downstream}')\n",
    "        return dy_dx_downstream\n",
    "    \n",
    "    y = x ** 2\n",
    "    tf.print(f'x={x}\\ty={y}')\n",
    "    \n",
    "    return y, grad\n",
    "\n",
    "\n",
    "x = tf.constant(2.0, dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    tape.watch(x)\n",
    "    y = foo(foo(foo(x))) # y = x ** 8\n",
    "\n",
    "tf.print(f'\\nfinal dy/dx={tape.gradient(y, x)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gradients with multiple variables\n",
    "\n",
    "If the function takes multiple variables, then the gradient for each variable has to be returned as demonstrated in the example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6\n",
      "3\n",
      "2\n",
      "None\n"
     ]
    }
   ],
   "source": [
    "@tf.custom_gradient\n",
    "def bar(x, y):\n",
    "    tf.debugging.assert_rank(x, 0)\n",
    "    tf.debugging.assert_rank(y, 0)\n",
    "\n",
    "    def grad(upstream):\n",
    "        dz_dx = y\n",
    "        dz_dy = x\n",
    "        return upstream * dz_dx, upstream * dz_dy\n",
    "    \n",
    "    z = x * y\n",
    "    \n",
    "    return z, grad\n",
    "\n",
    "x = tf.constant(2.0, dtype=tf.float32)\n",
    "y = tf.constant(3.0, dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    tape.watch(x)\n",
    "    tape.watch(y)\n",
    "    z = bar(x, y)\n",
    "\n",
    "tf.print(z)\n",
    "tf.print(tape.gradient(z, x))\n",
    "tf.print(tape.gradient(z, y))\n",
    "tf.print(tape.gradient(x, y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Application of custom gradients\n",
    "\n",
    "### Toy example: Differentiable approximation of non-differentiable functions\n",
    "\n",
    "We take the sign function as an example\n",
    "\n",
    "\\begin{equation}\n",
    "sign(x)= \\\\\n",
    "\\begin{cases}\n",
    "  -1, & \\text{if}\\ x<0 \\\\\n",
    "  0, & \\text{if}\\ x=0 \\\\\n",
    "  1, & \\text{if}\\ x>0 \\\\\n",
    "\\end{cases}\\end{equation}\n",
    "  \n",
    "By implementing a custom gradient, we can continue to have the $sign(x)$ function in forward pass but a differentiable approximation in the backward pass. In this case we approximate $sign(x)$ with the sigmoid function $ \\sigma(x)$\n",
    "\n",
    "\\begin{equation}\n",
    "\\frac{dsign_{approx}(x)}{dx} = \\sigma(x) (1 - \\sigma(x)) \\\\\n",
    "sign_{approx}(x) = \\sigma(x) + C \\\\\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "0.0451766551\n",
      "2\n",
      "0.0903533101\n"
     ]
    }
   ],
   "source": [
    "# @tf.function\n",
    "@tf.custom_gradient\n",
    "def differentiable_sign(x):\n",
    "    tf.debugging.assert_rank(x, 0)\n",
    "\n",
    "    def grad(upstream):\n",
    "        dy_dx = tf.math.sigmoid(x) * (1 - tf.math.sigmoid(x))\n",
    "        return upstream * dy_dx\n",
    "    \n",
    "    if x > tf.constant(0.0):\n",
    "        return tf.constant(1.0), grad\n",
    "    else:\n",
    "        return tf.constant(-1.0), grad\n",
    "\n",
    "\n",
    "x = tf.constant(3.0, dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    tape.watch(x)\n",
    "    y = differentiable_sign(x)\n",
    "    loss = tf.nn.l2_loss(y - tf.constant(-1.0))\n",
    "    \n",
    "tf.print(y)\n",
    "tf.print(tape.gradient(y, x))\n",
    "tf.print(loss)\n",
    "tf.print(tape.gradient(loss, x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 2 -0.393223882 -0.89999783 -1\n",
      "10 0 0 0.0995881185 1\n",
      "20 0 0 0.6450876 1\n",
      "30 0 0 0.855591536 1\n",
      "40 0 0 0.938390434 1\n",
      "50 0 0 0.970632374 1\n",
      "60 0 0 0.983009696 1\n",
      "70 0 0 0.987700462 1\n",
      "80 0 0 0.989459515 1\n",
      "90 0 0 0.990113616 1\n"
     ]
    }
   ],
   "source": [
    "x = tf.Variable(-1.0)\n",
    "opt = tf.keras.optimizers.Adam(1e-1)\n",
    "# opt = tf.keras.optimizers.SGD(1)\n",
    "\n",
    "def train_step():\n",
    "    with tf.GradientTape() as tape:\n",
    "        y = differentiable_sign(x)\n",
    "        loss = tf.nn.l2_loss(y - tf.constant(1.0))\n",
    "    grads = tape.gradient(loss, x)\n",
    "    opt.apply_gradients(zip([grads], [x]))\n",
    "    return loss, y, grads\n",
    "\n",
    "for i in range(100):\n",
    "    loss, y, grads = train_step()\n",
    "    if i % 10 == 0:\n",
    "        tf.print(i, loss, grads, x, y)"
   ]
  },
  {
   "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
}