[
  {
    "path": "README.md",
    "content": "# Implementation of Differentiable Neural Computers (DNC) in Chainer\n\nDifferentiable Neural Computers (DNC) is a neural network architecture proposed by DeepMind in [their third paper on Nature](http://www.nature.com/articles/nature20101.epdf?author_access_token=ImTXBI8aWbYxYQ51Plys8NRgN0jAjWel9jnR3ZoTv0MggmpDmwljGswxVdeocYSurJ3hxupzWuRNeGvvXnoO8o4jTJcnAyhGuZzXJ1GEaD-Z7E6X_a9R-xqJ9TfJWBqz).\nI have implemented DNC in [Chainer](http://chainer.org/), a flexible framework of neural networks developped by [Preferred Networks](https://www.preferred-networks.jp/en/).\n\n# What is DNC ?\nDNC is a newly proposed neural network. In their paper, DNC learns well in several complex tasks, including finding shortest path in a graph and solving a block puzzle game. It is expected to have the capacity to solve complex, structured tasks that are inaccessible to previous neural networks.\n\nDNC consists of a RNN (recurrent neural network) and a \"memory matrix\", with some heads for reading and writing to it. The RNN can control the heads at will; it can manipulate the heads in a predetermined fashion to read out the content of the memory and write some data to the memory.\n\nIn each timestep, a vanilla RNN receives some external input and yields some output (and refreshes its internal state). In contrast, a RNN in a DNC recieves \"data read by the read head at the previous timestep\" together with external input, and yields \"memory manipulation command\" in addition to output data. In accordance with this command, the heads are moved, the memory content at the write head is edited, and the memory contents at the read heads are fetched. Fetched data is input to RNN at the next timestep (together with external input data).\n\nA RNN in DNC learns so that it achieves appropriate input-output relationship in the situation that \"the memory\" --- a convenient tool to compute --- is given to use freely. How the RNN utilizes the tool depends on its learning.\n\nAlthough \"read-write memory\" seems to be very special, it can be regarded as a form of internal state of a RNN(*); in other words, DNC is an RNN that has non-trivial internal-state dynamics like LSTM, but the dynamics are very complicated. This memory enables the RNN to perform complicated information processings. Moreover, equipped with \"read-write memory\", which is fairly convenient for every kind of information processing, I expect that DNC has high versatility --- to perform various types of tasks reasonably well.\n\nNote that their [NTM (Neural Turing Machine)](https://arxiv.org/pdf/1410.5401v2.pdf) proposed in 2014 has similar structure to DNC. The difference between DNC and NTM is that DNC has more reasonable memory heads' movement. (For datails see the Methods in the DNC paper.)\n\n(*): They call the memory as \"external\". They say that is because \"The behaviour of the network is independent of the memory size as long as the memory is not filled to capacity\".\n\n\n# About my code\n\nIn my code, a very small-scale DNC learns a very easy \"repeat after me\" task. It seems to learn correctly without errors, but it does not necessarily mean that this program correctly performs DNC. If you have any comments about my code, please feel free to contact @yos1up (twitter).\n\nThe Supplementary Material of their paper is very useful to implement DNC. It contains ALL variables used in the model and ALL equations to construct the computational graph of the model in two pages. Most of the names of the variables shown in my code coincide with that in their paper.\n"
  },
  {
    "path": "main.py",
    "content": "import numpy as np\r\nimport math\r\nimport chainer\r\nfrom chainer import functions as F\r\nfrom chainer import links as L\r\nfrom chainer import \\\r\n     cuda, gradient_check, optimizers, serializers, utils, \\\r\n     Chain, ChainList, Function, Link, Variable\r\n\r\n\r\ndef onehot(x,n):\r\n    ret = np.zeros(n).astype(np.float32)\r\n    ret[x] = 1.0\r\n    return ret\r\n\r\ndef overlap(u, v): # u, v: (1 * -) Variable  -> (1 * 1) Variable\r\n    denominator = F.sqrt(F.batch_l2_norm_squared(u) * F.batch_l2_norm_squared(v))\r\n    if (np.array_equal(denominator.data, np.array([0]))):\r\n        return F.matmul(u, F.transpose(v))\r\n    return F.matmul(u, F.transpose(v)) / F.reshape(denominator,(1,1))\r\n\r\n\r\ndef C(M, k, beta):\r\n    # (N * W), (1 * W), (1 * 1) -> (N * 1)\r\n    # (not (N * W), ({R,1} * W), (1 * {R,1}) -> (N * {R,1}))\r\n    W = M.data.shape[1]    \r\n    ret_list = [0] * M.data.shape[0]\r\n    for i in range(M.data.shape[0]):\r\n        ret_list[i] = overlap(F.reshape(M[i,:], (1, W)), k) * beta # pick i-th row\r\n    return F.transpose(F.softmax(F.transpose(F.concat(ret_list, 0)))) # concat vertically and calc softmax in each column\r\n\r\n\r\n\r\ndef u2a(u): # u, a: (N * 1) Variable\r\n    N = len(u.data)\r\n    phi = np.argsort(u.data.reshape(N)) # u.data[phi]: ascending\r\n    a_list = [0] * N    \r\n    cumprod = Variable(np.array([[1.0]]).astype(np.float32)) \r\n    for i in range(N):\r\n        a_list[phi[i]] = cumprod * (1.0 - F.reshape(u[phi[i],0], (1,1)))\r\n        cumprod *= F.reshape(u[phi[i],0], (1,1))\r\n    return F.concat(a_list, 0) # concat vertically\r\n\r\n\r\n\r\nclass DeepLSTM(Chain): # too simple?\r\n    def __init__(self, d_in, d_out):\r\n        super(DeepLSTM, self).__init__(\r\n            l1 = L.LSTM(d_in, d_out),\r\n            l2 = L.Linear(d_out, d_out),)\r\n    def __call__(self, x):\r\n        self.x = x\r\n        self.y = self.l2(self.l1(self.x))\r\n        return self.y\r\n    def reset_state(self):\r\n        self.l1.reset_state()\r\n\r\n\r\n    \r\nclass DNC(Chain):\r\n    def __init__(self, X, Y, N, W, R):\r\n        self.X = X # input dimension\r\n        self.Y = Y # output dimension\r\n        self.N = N # number of memory slot\r\n        self.W = W # dimension of one memory slot\r\n        self.R = R # number of read heads\r\n        self.controller = DeepLSTM(W*R+X, Y+W*R+3*W+5*R+3)\r\n        \r\n        super(DNC, self).__init__(\r\n            l_dl = self.controller,\r\n            l_Wr = L.Linear(self.R * self.W, self.Y) # nobias=True ? \r\n            )# <question : should all learnable weights be here??>\r\n        self.reset_state()\r\n    def __call__(self, x):\r\n        # <question : is batchsize>1 possible for RNN ? if No, I will implement calculations without batch dimension.>\r\n        self.chi = F.concat((x, self.r))\r\n        (self.nu, self.xi) = \\\r\n                  F.split_axis(self.l_dl(self.chi), [self.Y], 1)\r\n        (self.kr, self.betar, self.kw, self.betaw,\r\n         self.e, self.v, self.f, self.ga, self.gw, self.pi\r\n         ) = F.split_axis(self.xi, np.cumsum(\r\n             [self.W*self.R, self.R, self.W, 1, self.W, self.W, self.R, 1, 1]), 1)\r\n\r\n        self.kr = F.reshape(self.kr, (self.R, self.W)) # R * W\r\n        self.betar = 1 + F.softplus(self.betar) # 1 * R\r\n        # self.kw: 1 * W\r\n        self.betaw = 1 + F.softplus(self.betaw) # 1 * 1\r\n        self.e = F.sigmoid(self.e) # 1 * W\r\n        # self.v : 1 * W\r\n        self.f = F.sigmoid(self.f) # 1 * R\r\n        self.ga = F.sigmoid(self.ga) # 1 * 1\r\n        self.gw = F.sigmoid(self.gw) # 1 * 1\r\n        self.pi = F.softmax(F.reshape(self.pi, (self.R, 3))) # R * 3 (softmax for 3)\r\n\r\n        # self.wr : N * R\r\n        self.psi_mat = 1 - F.matmul(Variable(np.ones((self.N, 1)).astype(np.float32)), self.f) * self.wr # N * R\r\n        self.psi = Variable(np.ones((self.N, 1)).astype(np.float32)) # N * 1\r\n        for i in range(self.R):\r\n            self.psi = self.psi * F.reshape(self.psi_mat[:,i],(self.N,1)) # N * 1\r\n\r\n        # self.ww, self.u : N * 1\r\n        self.u = (self.u + self.ww - (self.u * self.ww)) * self.psi\r\n        \r\n        self.a = u2a(self.u) # N * 1\r\n        self.cw = C(self.M, self.kw, self.betaw) # N * 1\r\n        self.ww = F.matmul(F.matmul(self.a, self.ga) + F.matmul(self.cw, 1.0 - self.ga), self.gw) # N * 1\r\n        self.M = self.M * (np.ones((self.N, self.W)).astype(np.float32) - F.matmul(self.ww, self.e)) + F.matmul(self.ww, self.v) # N * W\r\n\r\n        self.p = (1.0 - F.matmul(Variable(np.ones((self.N,1)).astype(np.float32)), F.reshape(F.sum(self.ww),(1,1)))) \\\r\n                  * self.p + self.ww # N * 1\r\n        self.wwrep = F.matmul(self.ww, Variable(np.ones((1, self.N)).astype(np.float32))) # N * N\r\n        self.L = (1.0 - self.wwrep - F.transpose(self.wwrep)) * self.L + F.matmul(self.ww, F.transpose(self.p)) # N * N\r\n        self.L = self.L * (np.ones((self.N, self.N)) - np.eye(self.N)) # force L[i,i] == 0   \r\n\r\n        self.fo = F.matmul(self.L, self.wr) # N * R\r\n        self.ba = F.matmul(F.transpose(self.L), self.wr) # N * R\r\n\r\n        self.cr_list = [0] * self.R\r\n        for i in range(self.R):\r\n            self.cr_list[i] = C(self.M, F.reshape(self.kr[i,:],(1, self.W)),\r\n                                F.reshape(self.betar[0,i],(1, 1))) # N * 1\r\n        self.cr = F.concat(self.cr_list) # N * R\r\n\r\n        self.bacrfo = F.concat((F.reshape(F.transpose(self.ba),(self.R,self.N,1)),\r\n                               F.reshape(F.transpose(self.cr),(self.R,self.N,1)),\r\n                               F.reshape(F.transpose(self.fo) ,(self.R,self.N,1)),),2) # R * N * 3\r\n        self.pi = F.reshape(self.pi, (self.R,3,1)) # R * 3 * 1\r\n        self.wr = F.transpose(F.reshape(F.batch_matmul(self.bacrfo, self.pi), (self.R, self.N))) # N * R\r\n            \r\n        self.r = F.reshape(F.matmul(F.transpose(self.M), self.wr),(1, self.R * self.W)) # W * R (-> 1 * RW)\r\n        \r\n        self.y = self.l_Wr(self.r) + self.nu # 1 * Y\r\n        return self.y\r\n    def reset_state(self):\r\n        self.l_dl.reset_state()\r\n        self.u = Variable(np.zeros((self.N, 1)).astype(np.float32))\r\n        self.p = Variable(np.zeros((self.N, 1)).astype(np.float32))\r\n        self.L = Variable(np.zeros((self.N, self.N)).astype(np.float32))                           \r\n        self.M = Variable(np.zeros((self.N, self.W)).astype(np.float32))\r\n        self.r = Variable(np.zeros((1, self.R*self.W)).astype(np.float32))\r\n        self.wr = Variable(np.zeros((self.N, self.R)).astype(np.float32))\r\n        self.ww = Variable(np.zeros((self.N, 1)).astype(np.float32))\r\n        # any variable else ?\r\n\r\nX = 5\r\nY = 5\r\nN = 10\r\nW = 10\r\nR = 2\r\nmdl = DNC(X, Y, N, W, R)\r\nopt = optimizers.Adam()\r\nopt.setup(mdl)\r\ndatanum = 100000\r\nloss = 0.0\r\nacc = 0.0\r\nfor datacnt in range(datanum):\r\n    lossfrac = np.zeros((1,2))\r\n    # x_seq = np.random.rand(X,seqlen).astype(np.float32)\r\n    # t_seq = np.random.rand(Y,seqlen).astype(np.float32)\r\n    # t_seq = np.copy(x_seq)\r\n\r\n    contentlen = np.random.randint(3,6)\r\n    content = np.random.randint(0,X-1,contentlen)\r\n    seqlen = contentlen + contentlen\r\n    x_seq_list = [float('nan')] * seqlen\r\n    t_seq_list = [float('nan')] * seqlen    \r\n    for i in range(seqlen):\r\n        if (i < contentlen):\r\n            x_seq_list[i] = onehot(content[i],X)\r\n        elif (i == contentlen):\r\n            x_seq_list[i] = onehot(X-1,X)\r\n        else:\r\n            x_seq_list[i] = np.zeros(X).astype(np.float32)\r\n            \r\n        if (i >= contentlen):\r\n            t_seq_list[i] = onehot(content[i-contentlen],X)    \r\n    \r\n    mdl.reset_state()\r\n    for cnt in range(seqlen):\r\n        x = Variable(x_seq_list[cnt].reshape(1,X))\r\n        if (isinstance(t_seq_list[cnt], np.ndarray)):\r\n            t = Variable(t_seq_list[cnt].reshape(1,Y))\r\n        else:\r\n            t = []\r\n            \r\n        y = mdl(x)\r\n        if (isinstance(t,chainer.Variable)):\r\n            loss += (y - t)**2\r\n            print y.data, t.data, np.argmax(y.data)==np.argmax(t.data)\r\n            if (np.argmax(y.data)==np.argmax(t.data)): acc += 1\r\n        if (cnt+1==seqlen):\r\n            mdl.cleargrads()\r\n            loss.grad = np.ones(loss.data.shape, dtype=np.float32)\r\n            loss.backward()\r\n            opt.update()\r\n            loss.unchain_backward()\r\n            print '(', datacnt, ')', loss.data.sum()/loss.data.size/contentlen, acc/contentlen\r\n            lossfrac += [loss.data.sum()/loss.data.size/seqlen, 1.]\r\n            loss = 0.0\r\n            acc = 0.0\r\n"
  }
]