NALU.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Neural Arithmatic Logic Units\n",
    "Google DeepMind's research paper: https://arxiv.org/abs/1808.00508"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![title](images/naluandnac.png)\n",
    "**The Neural Accumulator (NAC)** is a linear transformation of its inputs.\n",
    "And the tranformation matrix is elementwise product of **tanh(W)** and **sigmoid(M)**\n",
    "\n",
    "**The Neural Arithmetic Logic Unit (NALU)** uses two NACs with tied weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The Neural Arithmetic Logic Unit\n",
    "def NALU(in_dim, out_dim):\n",
    "\n",
    "    shape = (int(in_dim.shape[-1]), out_dim)\n",
    "    epsilon = 1e-7 \n",
    "    \n",
    "    # NAC\n",
    "    W_hat = tf.Variable(tf.truncated_normal(shape, stddev=0.02))\n",
    "    M_hat = tf.Variable(tf.truncated_normal(shape, stddev=0.02))\n",
    "    G = tf.Variable(tf.truncated_normal(shape, stddev=0.02))\n",
    "        \n",
    "    W = tf.tanh(W_hat) * tf.sigmoid(M_hat)\n",
    "    # Forward propogation\n",
    "    a = tf.matmul(in_dim, W)\n",
    "    \n",
    "    # NALU  \n",
    "    m = tf.exp(tf.matmul(tf.log(tf.abs(in_dim) + epsilon), W))\n",
    "    g = tf.sigmoid(tf.matmul(in_dim, G))\n",
    "    y = g * a + (1 - g) * m\n",
    "    \n",
    "    return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_dataset(size=10000):\n",
    "    # input data\n",
    "    X = np.random.randint(9, size=(size,2))\n",
    "    # output data (labels)   \n",
    "    Y = np.prod(X, axis=1, keepdims=True)\n",
    "\n",
    "        \n",
    "    return X, Y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train NALU on generated data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters\n",
    "EPOCHS = 200\n",
    "LEARNING_RATE = 1e-3\n",
    "BATCH_SIZE = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create dataset\n",
    "X_data, Y_data = generate_dataset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define placeholders and network\n",
    "X = tf.placeholder(tf.float32, shape=[BATCH_SIZE, 2])\n",
    "\n",
    "Y_true = tf.placeholder(tf.float32, shape=[BATCH_SIZE, 1])\n",
    "\n",
    "Y_pred = NALU(X, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = tf.nn.l2_loss(Y_pred - Y_true) \n",
    "    \n",
    "optimizer = tf.train.AdamOptimizer(LEARNING_RATE).minimize(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create session\n",
    "sess = tf.Session()\n",
    "# create writer to store tensorboard graph   \n",
    "writer = tf.summary.FileWriter('/tmp', sess.graph)\n",
    "    \n",
    "init = tf.global_variables_initializer()\n",
    "    \n",
    "sess.run(init)\n",
    "\n",
    "# Run training loop\n",
    "for i in range(EPOCHS):\n",
    "    j = 0\n",
    "    g = 0\n",
    "        \n",
    "    while j < len(X_data):\n",
    "        xs, ys = X_data[j:j + BATCH_SIZE], Y_data[j:j + BATCH_SIZE]\n",
    "\n",
    "        _, ys_pred, l = sess.run([optimizer, Y_pred, loss], \n",
    "                    feed_dict={X: xs, Y_true: ys})\n",
    "            \n",
    "        # calculate number of correct predictions from batch\n",
    "        g += np.sum(np.isclose(ys, ys_pred, atol=1e-4, rtol=1e-4)) \n",
    "\n",
    "        j += BATCH_SIZE\n",
    "\n",
    "    acc = g / len(Y_data)\n",
    "        \n",
    "    print(f'epoch {i}, loss: {l}, accuracy: {acc}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Uncomment to run TensorBoard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !tensorboard --logdir /tmp"
   ]
  }
 ],
 "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}