neural_turing_machine.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import argparse\n",
    "import sys\n",
    "\n",
    "from tensorflow.examples.tutorials.mnist import input_data\n",
    "import tensorflow as tf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "batch_size = 100\n",
    "n_itr = 50000\n",
    "T = 10\n",
    "eps = 1e-7\n",
    "lr = 1e-3\n",
    "N = 128\n",
    "M = 20\n",
    "controller_size = 100\n",
    "shift_stride = 1\n",
    "num_shift = 2*shift_stride + 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "memory = tf.Variable(tf.fill([N, M], eps))\n",
    "# for feed forward controller\n",
    "#controller_W = tf.Variable(tf.zeros([N, controller_size]))\n",
    "#controller_b = tf.Variable(tf.zeros([controller_size]))\n",
    "\n",
    "k_W = tf.Variable(tf.zeros([controller_size, M]))\n",
    "k_b = tf.Variable(tf.zeros([M]))\n",
    "beta_W = tf.Variable(tf.zeros([controller_size, 1]))\n",
    "beta_b = tf.Variable(tf.zeros([1]))\n",
    "g_W = tf.Variable(tf.zeros([controller_size, 1]))\n",
    "g_b = tf.Variable(tf.zeros([1]))\n",
    "s_W = tf.Variable(tf.zeros([controller_size, num_shift]))\n",
    "s_b = tf.Variable(tf.zeros([num_shift]))\n",
    "gamma_W = tf.Variable(tf.zeros([controller_size, 1]))\n",
    "gamma_b = tf.Variable(tf.zeros([1]))\n",
    "\n",
    "a_W = tf.Variable(tf.zeros([controller_size, M]))\n",
    "a_b = tf.Variable(tf.zeros([M]))\n",
    "e_W = tf.Variable(tf.zeros([controller_size, M]))\n",
    "e_b = tf.Variable(tf.zeros([M]))\n",
    "\n",
    "controller = tf.nn.rnn_cell.LSTMCell(controller_size)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def read(w, mem): #w_t => Nx1\n",
    "    r = tf.matmul(w, mem)\n",
    "    return r\n",
    "\n",
    "def write(w, h, mem_prev): #w_t => Nx1 , e_t => 1xM, a_t = 1xM\n",
    "    e = tf.nn.relu(tf.matmul(h, e_W) + e_b)\n",
    "    a = tf.nn.relu(tf.matmul(h, a_W) + a_b)\n",
    "    \n",
    "    erase_factor = 1 - tf.multiply(w, e)\n",
    "    memory_cap = tf.multiply(mem_prev, erase_factor)\n",
    "    add_factor = tf.multiply(w, a)\n",
    "    memory_next = memory_cap + add_factor\n",
    "    assign_op_memory = mem_prev.assign(memory_next)\n",
    "    return assign_op_memory \n",
    "\n",
    "def content_addressing(k, beta, mem_prev): # k => 1xM\n",
    "    norm_k = tf.nn.l2_normalize(k, dim=1)\n",
    "    norm_memory = tf.nn.l2_normalize(mem_prev, dim=1)\n",
    "    cosine_similarity = tf.matmul(norm_memory, norm_k)\n",
    "    \n",
    "    w_c = tf.nn.softmax(tf.multiply(beta, cosine_similarity))\n",
    "    return w_c\n",
    "\n",
    "def interpolation(w_prev, w_c, g):\n",
    "    w_g = tf.multiply(w_prev, g) + tf.multiply(w_c, 1 - g)\n",
    "    return w_g\n",
    "\n",
    "def convolution_shift(g_w, s):\n",
    "    for i in xrange(N):\n",
    "        for j in xrange(num_shift):\n",
    "            idx = shift_stride + i - j + 1\n",
    "            if idx < 0: idx = idx + N\n",
    "            if idx >= N: idx = idx - N\n",
    "            w = tf.gather(g_w, int(idx)) * tf.gather(s, j)\n",
    "            output = tf.scatter_add(output, [i], tf.reshape(w, [1, -1]))\n",
    "\n",
    "def sharpening(w_cap, gamma):\n",
    "    w = tf.pow(w_cap, gamma)\n",
    "    w = w / tf.maximum(tf.reduce_sum(w, 2, keep_dims=True), eps)\n",
    "    return w\n",
    "\n",
    "def addressing(h, w_prev, mem_prev):\n",
    "    '''\n",
    "    The key vector, kt, and key strength, βt, are used to perform \n",
    "    content-based addressing of the memory matrix, Mt. \n",
    "    The resulting content-based weighting is interpolated with \n",
    "    the weighting from the previous time step based on the value \n",
    "    of the interpolation gate, gt. The shift weighting, st, \n",
    "    determines whether and by how much the weighting is rotated. \n",
    "    Finally, depending on γt, the weighting is sharpened \n",
    "    and used for memory access.\n",
    "    \n",
    "    k => key vector  # 1xM\n",
    "    beta => key strength # 1\n",
    "    g => interpolaton gate # 1, (0, 1)\n",
    "    s => location shift weight # num_shiftx1\n",
    "    gamma => sharpening wt # 1, >= 1\n",
    "    '''\n",
    "    k = tf.nn.relu(tf.matmul(h, k_W) + k_b)\n",
    "    beta = tf.nn.relu(tf.matmul(h, beta_W) + beta_b)\n",
    "    g = tf.sigmoid(tf.matmul(h, g_W) + g_b)\n",
    "    s = tf.nn.softmax(tf.matmul(h, s_W) + s_b)\n",
    "    gamma = 1 + tf.nn.relu(tf.matmul(h, gamma_W) + gamma_b)\n",
    "    \n",
    "    w_c = content_addressing(k, beta, mem_prev)\n",
    "    w_g = interpolation(w_prev, w_c, g)\n",
    "    w_cap = convolution_shift(w_g, s)\n",
    "    w = sharpening(w_cap, gamma)\n",
    "    return w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "lstm_state = controller.zero_state(batch_size, tf.float32) \n",
    "\n",
    "h0 = tf.zeros([batch_size, controller_size])\n",
    "h = h0\n",
    "w = w0\n",
    "for t in range(T):\n",
    "    r = read(w, memory)\n",
    "    # input to hidden\n",
    "    controller_in = tf.concat(1,[x, r, h])\n",
    "    with tf.variable_scope(\"controller\", reuse=(t != 0)):\n",
    "        h, controller_state = rnn_enc(controller_in, lstm_state)\n",
    "    w = addressing(h, w, memory)\n",
    "    \n",
    "    write(w, e, a, memory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [conda root]",
   "language": "python",
   "name": "conda-root-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}