[PYTHON] Overview of DNC (Differentiable Neural Computers) + Implementation by Chainer

Hello. My name is @ yos1up. My favorite food is green mushrooms.

2016/10/12 to paper DeepMind have posted in Nature, [Hybrid computing using a neural network with dynamic external memory](http://www.nature.com/articles/nature20101.epdf?author_access_token=ImTXBI8aWbYxYQ51Plys8NRgN0jAjWel9jnR3ZoTv0MggmpDmwljGswxVdeocYSurJ3hxupzWuRNeGvvXnoO8o4jTJcnAyhGuZzXJ1GEaD-Z7E6X_a9R-xqJ9TfJWBqz I hurriedly implemented the neural network model ** DNC (Differentiable Neural Computers) ** proposed in) with Chainer.

About DNC

DNC is a new neural network proposed in the above paper, and its high information processing ability is expected. In the paper, tasks that were previously thought to be impossible to learn with neural networks, such as the shortest path task on a graph and a small puzzle task, can be learned with DNC, and its high information processing ability can be seen.

DNC has a structure in which "memory" is attached to RNN (recurrent neural network). This "memory" has a head that reads information and a head that writes information, and you can move the head (according to a limited method) and read and write information about the position of the head.

A normal RNN takes some external input and produces some output (and updates its internal state) every hour. On the other hand, the RNN in the DNC receives "data read from" memory "at the immediately preceding time" in addition to the normal external input at each time. Then, in addition to the normal output, "operation instruction of" memory "" is also output. According to this operation instruction, the position of the memory head changes, the memory at the position of the write head is rewritten, and the memory at the position of the read head is read. The read data will be input to the RNN (along with the normal external input) at the next time.

The RNN in the DNC learns the coupling weight (by the gradient method) so that an appropriate input / output relationship can be realized in the situation given this prop called "memory". How you use (or don't use) this prop depends on your learning.

Although it has a special form of "memory with head", it must be the "internal state" of the RNN, and in that sense, the DNC is an extension of the GRU and LSTM, and the "internal state has become significantly complicated." It can be said that it is "RNN" [^ 2]. Thanks to this "memory", it is possible to realize ** complicated information processing ** that cannot be handled by conventional RNNs with DNC.

[^ 2]: However, they refer to this memory as "external memory" in their dissertation. From the following paper: The behavior of the network is independent of the memory size as long as the memory is not filled to capacity, which is why we view the memory as'external'. The behavior of the network is independent of memory size, so we consider this memory to be "external" memory.)

In addition, the "memory with head" that seems to be reasonably convenient for performing all kinds of information processing gives DNC ** versatility ** that can handle all types of tasks in its own way. I personally expect that it will be there. The wide variety of task genres that DNC solves in the paper also makes us expect high versatility.

The neural network they proposed in December 2014 called NTM (Neural Turing Machine) has a structure similar to that of this DNC. However, DNC is powered up in that it makes it more rational to move the memory head than NTM. (There is a summary of the differences between NTM and DNC in Methods in the paper.)

Implementation

The code implemented this time (Python 2.7) is shown below. (GitHub)

main.py


import numpy as np
import math
import chainer
from chainer import functions as F
from chainer import links as L
from chainer import \
     cuda, gradient_check, optimizers, serializers, utils, \
     Chain, ChainList, Function, Link, Variable


def onehot(x,n):
    ret = np.zeros(n).astype(np.float32)
    ret[x] = 1.0
    return ret

def overlap(u, v): # u, v: (1 * -) Variable  -> (1 * 1) Variable
    denominator = F.sqrt(F.batch_l2_norm_squared(u) * F.batch_l2_norm_squared(v))
    if (np.array_equal(denominator.data, np.array([0]))):
        return F.matmul(u, F.transpose(v))
    return F.matmul(u, F.transpose(v)) / F.reshape(denominator,(1,1))


def C(M, k, beta):
    # (N * W), (1 * W), (1 * 1) -> (N * 1)
    # (not (N * W), ({R,1} * W), (1 * {R,1}) -> (N * {R,1}))
    W = M.data.shape[1]    
    ret_list = [0] * M.data.shape[0]
    for i in range(M.data.shape[0]):
        ret_list[i] = overlap(F.reshape(M[i,:], (1, W)), k) * beta # pick i-th row
    return F.transpose(F.softmax(F.transpose(F.concat(ret_list, 0)))) # concat vertically and calc softmax in each column



def u2a(u): # u, a: (N * 1) Variable
    N = len(u.data)
    phi = np.argsort(u.data.reshape(N)) # u.data[phi]: ascending
    a_list = [0] * N    
    cumprod = Variable(np.array([[1.0]]).astype(np.float32)) 
    for i in range(N):
        a_list[phi[i]] = cumprod * (1.0 - F.reshape(u[phi[i],0], (1,1)))
        cumprod *= F.reshape(u[phi[i],0], (1,1))
    return F.concat(a_list, 0) # concat vertically



class DeepLSTM(Chain): # too simple?
    def __init__(self, d_in, d_out):
        super(DeepLSTM, self).__init__(
            l1 = L.LSTM(d_in, d_out),
            l2 = L.Linear(d_out, d_out),)
    def __call__(self, x):
        self.x = x
        self.y = self.l2(self.l1(self.x))
        return self.y
    def reset_state(self):
        self.l1.reset_state()


    
class DNC(Chain):
    def __init__(self, X, Y, N, W, R):
        self.X = X # input dimension
        self.Y = Y # output dimension
        self.N = N # number of memory slot
        self.W = W # dimension of one memory slot
        self.R = R # number of read heads
        self.controller = DeepLSTM(W*R+X, Y+W*R+3*W+5*R+3)
        
        super(DNC, self).__init__(
            l_dl = self.controller,
            l_Wr = L.Linear(self.R * self.W, self.Y) # nobias=True ? 
            )# <question : should all learnable weights be here??>
        self.reset_state()
    def __call__(self, x):
        # <question : is batchsize>1 possible for RNN ? if No, I will implement calculations without batch dimension.>
        self.chi = F.concat((x, self.r))
        (self.nu, self.xi) = \
                  F.split_axis(self.l_dl(self.chi), [self.Y], 1)
        (self.kr, self.betar, self.kw, self.betaw,
         self.e, self.v, self.f, self.ga, self.gw, self.pi
         ) = F.split_axis(self.xi, np.cumsum(
             [self.W*self.R, self.R, self.W, 1, self.W, self.W, self.R, 1, 1]), 1)

        self.kr = F.reshape(self.kr, (self.R, self.W)) # R * W
        self.betar = 1 + F.softplus(self.betar) # 1 * R
        # self.kw: 1 * W
        self.betaw = 1 + F.softplus(self.betaw) # 1 * 1
        self.e = F.sigmoid(self.e) # 1 * W
        # self.v : 1 * W
        self.f = F.sigmoid(self.f) # 1 * R
        self.ga = F.sigmoid(self.ga) # 1 * 1
        self.gw = F.sigmoid(self.gw) # 1 * 1
        self.pi = F.softmax(F.reshape(self.pi, (self.R, 3))) # R * 3 (softmax for 3)

        # self.wr : N * R
        self.psi_mat = 1 - F.matmul(Variable(np.ones((self.N, 1)).astype(np.float32)), self.f) * self.wr # N * R
        self.psi = Variable(np.ones((self.N, 1)).astype(np.float32)) # N * 1
        for i in range(self.R):
            self.psi = self.psi * F.reshape(self.psi_mat[:,i],(self.N,1)) # N * 1

        # self.ww, self.u : N * 1
        self.u = (self.u + self.ww - (self.u * self.ww)) * self.psi
        
        self.a = u2a(self.u) # N * 1
        self.cw = C(self.M, self.kw, self.betaw) # N * 1
        self.ww = F.matmul(F.matmul(self.a, self.ga) + F.matmul(self.cw, 1.0 - self.ga), self.gw) # N * 1
        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

        self.p = (1.0 - F.matmul(Variable(np.ones((self.N,1)).astype(np.float32)), F.reshape(F.sum(self.ww),(1,1)))) \
                  * self.p + self.ww # N * 1
        self.wwrep = F.matmul(self.ww, Variable(np.ones((1, self.N)).astype(np.float32))) # N * N
        self.L = (1.0 - self.wwrep - F.transpose(self.wwrep)) * self.L + F.matmul(self.ww, F.transpose(self.p)) # N * N
        self.L = self.L * (np.ones((self.N, self.N)) - np.eye(self.N)) # force L[i,i] == 0   

        self.fo = F.matmul(self.L, self.wr) # N * R
        self.ba = F.matmul(F.transpose(self.L), self.wr) # N * R

        self.cr_list = [0] * self.R
        for i in range(self.R):
            self.cr_list[i] = C(self.M, F.reshape(self.kr[i,:],(1, self.W)),
                                F.reshape(self.betar[0,i],(1, 1))) # N * 1
        self.cr = F.concat(self.cr_list) # N * R

        self.bacrfo = F.concat((F.reshape(F.transpose(self.ba),(self.R,self.N,1)),
                               F.reshape(F.transpose(self.cr),(self.R,self.N,1)),
                               F.reshape(F.transpose(self.fo) ,(self.R,self.N,1)),),2) # R * N * 3
        self.pi = F.reshape(self.pi, (self.R,3,1)) # R * 3 * 1
        self.wr = F.transpose(F.reshape(F.batch_matmul(self.bacrfo, self.pi), (self.R, self.N))) # N * R
            
        self.r = F.reshape(F.matmul(F.transpose(self.M), self.wr),(1, self.R * self.W)) # W * R (-> 1 * RW)
        
        self.y = self.l_Wr(self.r) + self.nu # 1 * Y
        return self.y
    def reset_state(self):
        self.l_dl.reset_state()
        self.u = Variable(np.zeros((self.N, 1)).astype(np.float32))
        self.p = Variable(np.zeros((self.N, 1)).astype(np.float32))
        self.L = Variable(np.zeros((self.N, self.N)).astype(np.float32))                           
        self.M = Variable(np.zeros((self.N, self.W)).astype(np.float32))
        self.r = Variable(np.zeros((1, self.R*self.W)).astype(np.float32))
        self.wr = Variable(np.zeros((self.N, self.R)).astype(np.float32))
        self.ww = Variable(np.zeros((self.N, 1)).astype(np.float32))
        # any variable else ?

X = 5
Y = 5
N = 10
W = 10
R = 2
mdl = DNC(X, Y, N, W, R)
opt = optimizers.Adam()
opt.setup(mdl)
datanum = 100000
loss = 0.0
acc = 0.0
for datacnt in range(datanum):
    lossfrac = np.zeros((1,2))
    # x_seq = np.random.rand(X,seqlen).astype(np.float32)
    # t_seq = np.random.rand(Y,seqlen).astype(np.float32)
    # t_seq = np.copy(x_seq)

    contentlen = np.random.randint(3,6)
    content = np.random.randint(0,X-1,contentlen)
    seqlen = contentlen + contentlen
    x_seq_list = [float('nan')] * seqlen
    t_seq_list = [float('nan')] * seqlen    
    for i in range(seqlen):
        if (i < contentlen):
            x_seq_list[i] = onehot(content[i],X)
        elif (i == contentlen):
            x_seq_list[i] = onehot(X-1,X)
        else:
            x_seq_list[i] = np.zeros(X).astype(np.float32)
            
        if (i >= contentlen):
            t_seq_list[i] = onehot(content[i-contentlen],X)    
    
    mdl.reset_state()
    for cnt in range(seqlen):
        x = Variable(x_seq_list[cnt].reshape(1,X))
        if (isinstance(t_seq_list[cnt], np.ndarray)):
            t = Variable(t_seq_list[cnt].reshape(1,Y))
        else:
            t = []
            
        y = mdl(x)
        if (isinstance(t,chainer.Variable)):
            loss += (y - t)**2
            print y.data, t.data, np.argmax(y.data)==np.argmax(t.data)
            if (np.argmax(y.data)==np.argmax(t.data)): acc += 1
        if (cnt+1==seqlen):
            mdl.cleargrads()
            loss.grad = np.ones(loss.data.shape, dtype=np.float32)
            loss.backward()
            opt.update()
            loss.unchain_backward()
            print '(', datacnt, ')', loss.data.sum()/loss.data.size/contentlen, acc/contentlen
            lossfrac += [loss.data.sum()/loss.data.size/seqlen, 1.]
            loss = 0.0
            acc = 0.0

In this paper, the details of the model are firmly written in Methods and Supplementary Material, and I felt very kind. In particular, the Supplementary contains all the variables and the formulas in the model, and I was able to complete the model simply by "porting" all the formulas written here into the code in order from the top. (Variable names in the above code almost correspond to variable names in Supplementary formulas.) Also, in the process of writing the code, I was able to get some details about the various functions that operate on Chainer's Variables. (I may summarize if I have the opportunity.)

The above code is a code that lets DNC learn a very simple task (a task that echoes a short symbol string with a delay). It works without any error. You can learn with about 1000 data. (* This task is not in the paper.) However, since it was implemented in a hurry, there is no certainty that the DNC was implemented correctly as in the paper. If you find any mistakes, please let us know.

In the future, I would like to apply the DNC created this time to various learning tasks. I want to see DeepMind's dissertation and try to solve puzzles.

Change log

2016/10/30 …… I found that the slicing of Variable can be written concisely, so I modified some of the code.

Recommended Posts

Overview of DNC (Differentiable Neural Computers) + Implementation by Chainer
Rank learning using neural network (Implementation of RankNet by Chainer)
Implementation of "blurred" neural network using Chainer
Bayesian optimization implementation of neural network hyperparameters (Chainer + GPyOpt)
Simple neural network implementation using Chainer
Summary of basic implementation by PyTorch
Implementation of a two-layer neural network 2
Python learning memo for machine learning by Chainer Chapter 13 Basics of neural networks
Implementation of 3-layer neural network (no learning)
Implementation of dialogue system using Chainer [seq2seq]
Implementation of SVM by stochastic gradient descent
[Chainer] Document classification by convolutional neural network
Recognition of handwritten numbers by multi-layer neural network
Deep learning learned by implementation (segmentation) ~ Implementation of SegNet ~