Improving Differentiable Neural Computers Through Memory Masking, De-allocation, and Link Distribution Sharpness Control

1 INTRODUCTION . Although Recurrent Neural Networks ( RNNs ) such as LSTM ( Hochreiter & Schmidhuber , 1997 ; Gers et al. , 2000 ) are in theory capable of solving complex algorithmic tasks ( Siegelmann & Sontag , 1992 ) , in practice they often struggle to do so . One reason is the large amount of time-varying memory required for many algorithmic tasks , combined with quadratic growth of the number of trainable parameters of a fully connected RNN when increasing the size of its internal state . Researchers have tried to address this problem by incorporating an external memory as a useful architectural bias for algorithm learning ( Das et al. , 1992 ; Mozer & Das , 1993 ; Graves et al. , 2014 ; 2016 ) . Especially the Differentiable Neural Computer ( DNC ; Graves et al . ( 2016 ) ) has shown great promise on a variety of algorithmic tasks – see diverse experiments in previous work ( Graves et al. , 2016 ; Rae et al. , 2016 ) . It combines a large external memory with advanced addressing mechanisms such as content-based look-up and temporal linking of memory cells . Unlike approaches that achieve state of the art performance on specific tasks , e.g . MemNN ( Sukhbaatar et al. , 2015 ) or Key-Value Networks ( Miller et al. , 2016 ) for the bAbI dataset ( Weston et al. , 2015 ) , the DNC consistently reaches near state of the art performance on all of them . This generality makes the DNC worth of further study . Three problems with the current DNC revolve around the content-based look-up mechanism , which is the main memory addressing system , and the temporal linking used to read memory cells in the same order in which they were written . First , the lack of key-value separation negatively impacts the accuracy of content retrieval . Second , the current de-allocation mechanism fails to remove deallocated data from memory , which prevents the network from erasing outdated information without explicitly overwriting the data . Third , with each write , the noise from the write address distribution accumulates in the temporal linking matrix , degrading the overall quality of temporal links . Here we propose a solution to each of these problems . We allow for dynamic key-value separation through a masking of both look-up key and data that is more general than a naive fixed key-value memory , yet does not suffer from loss of accuracy in addressing content . We propose to wipe the content of a memory cell in response to a decrease of its usage counter to allow for proper memory de-allocation . Finally , we reduce the effect of noise accumulation in the temporal linking matrix through exponentiation and re-normalization of the links , resulting in improved sharpness of the corresponding address distribution . These improvements are orthogonal to other previously proposed DNC modifications . Incorporation of the differentiable allocation mechanism Ben-Ari & Bekker ( 2017 ) or certain improvements to memory usage and computational complexity ( Rae et al. , 2016 ) might further improve the results reported in this paper . Certain bAbI-specific modifications Franke et al . ( 2018 ) are also orthogonal to our work . We evaluate each of the proposed modifications empirically on a benchmark of algorithmic tasks and on bAbI ( Weston et al. , 2015 ) . In all cases we find that our model outperforms the DNC . In particular , on the bAbI task we observe a 43 % relative improvement in terms of mean error rate . We find that improved de-allocation together with sharpness enhancement leads to zero error and 3x faster convergence on the large repeated copy task , while DNC is not able to solve it at all . Section 2 provides a brief overview of the DNC . Section 3 discusses identified problems and proposed solutions in more detail . Section 4 analyzes these modifications one-by-one , demonstrating their positive effects . 2 DIFFERENTIBLE NEURAL COMPUTER . Here we provide a brief overview of the Differentiable Neural Computer ( DNC ) . More details can be found in the original work of Graves et al . ( 2016 ) . The DNC combines a neural network ( called controller ) with an external memory that includes several supporting modules ( subsystems ) to do : read and write memory , allocate new memory cells , chain memory reads in the order in which they were written , and search memory for partial data . A simplified block diagram of the memory access is shown in Fig . 1 . External memory A main component is the external fixed 2D memory organized in cells ( Mt ∈ RN×W , where N is the number of cells , W is the cell length ) . N is independent of the number of trainable parameters . The controller is responsible for producing the activations of gates and keys controlling the memory transactions . The memory is accessed through multiple read heads and a single write head . Cells are addressed through a distribution over the whole address space . Each cell is read from and written to at every time step as determined by the address distributions , resulting in a differentiable procedure . Memory addressing The DNC uses three addressing methods . The most important one is contentbased look-up . It compares every cell to a key ( k∗t ∈ RW ) produced by the controller , resulting in a score , which is later normalized to get an address distribution over the whole memory . The second is the temporal linking , which has 2 types : forward and backward . They show which cell is being written after and before the one read in the previous time step . They are useful for processing sequences of data . A so-called temporal linkage matrix ( Lt ∈ RN×N ) is used to project any address distribution to a distribution that follows ( f it ∈ [ 0 , 1 ] N ) or precedes it ( bit ∈ [ 0 , 1 ] N ) . The third is the allocation mechanism , which is used only for write heads , and used when a new memory cell is required . Memory allocation Memory allocation works by maintaining usage counters for every cell . These are incremented on memory writes and optionally decremented on memory reads ( de-allocation ) . When a new cell is allocated , the one with the lowest usage counter is chosen . De-allocation is controlled by a gate , which is based on the address distribution of the previous read and decreases the usage counter of each cell . Read / Write The memory is first written to , then read from . A write address is generated as a weighted average of the write content-based look-up and the allocation distribution . The update is done in a gated way by erasing vector et ∈ [ 0 , 1 ] W . Parallel to the write , the temporal linkage matrix is also updated . Finally the memory is read from . The final read address distribution ( wr , it ∈ [ 0 , 1 ] N ) is generated as the weighted average of the read content-based look-up distribution and forward and backward temporal links . Memory cells are averaged based on this address , resulting in a single vector , which is the retrieved data . This data is combined with the output of the controller to produce the model ’ s final output . 3 METHOD . 3.1 MASKED CONTENT-BASED ADDRESSING . The goal of content-based addressing is to find memory cells similar to a given key . The query key contains partial information ( it is a partial memory ) , and the content-based memory read completes its missing ( unknown ) part based on previous memories . However , controlling which part of the key vector to search for is difficult because there is no key-value separation : the entire key and entire cell value are compared to produce the similarity score . This means that the part of the cell value that is unknown during search time and should be retrieved is also used in the normalization part of the cosine similarity , resulting in an unpredictable score . The shorter the known part and the longer the part to be retrieved , the worse the problem . This might result in less similar cells having higher scores , and might make the resulting address distribution flat because of the division by the length of the data before the softmax acts as an increased temperature parameter . Imagine a situation when the network has to tag some of the data written into the memory ( for example because it is the start of the sequence in the repeated copy task , see Section 4 ) . For tagging it can use only a single element of the memory cell , using the rest for the data that has to be stored . Later , when the network has to find this tag , it searches for it , specifying only that single element of the key . But the resulting score is also normalized by the data that is stored along with the tag , which takes most of the memory cell , hence the resulting score will almost completely be dominated by it . The problem can be solved by providing a way of explicitly masking the part that is unknown and should not be used in the query . This is more general than key-value memory , since the key-value separation can be controlled dynamically and does not suffer from the incorrect score problem . We achieve this by producing a separate mask vector m∗t ∈ [ 0 − 1 ] W through the controller , and multiplying both the search key and the memory content by it before comparing ( β is the write key strength controlling the temperature of the softmax ) : C ( M , k , β , m ) [ i ] = softmax ( D ( k m , M 1mT ) β ) cwt = C ( Mt−1 , k w t , β w t , m w t ) c r , i t = C ( Mt , k r , i t , β r , i t , m r , i t ) Fig . 2 shows how the masking step is incorporated in the address generation of the DNC . 3.2 DE-ALLOCATION AND CONTENT-BASED LOOK-UP . The DNC tracks allocation states of memory cells by so-called usage counters which are increased on memory writes and optionally decreased after reads . When allocating memory , the cell with the lowest usage is chosen . Decreasing is done by element-wise multiplication with a so-called retention vector ( ψt ) , which is a function of previously read address distributions ( w r , i t−1 ) and scalar gates . Vector ψt indicates how much of the current memory should be kept . The problem is that it affects solely the usage counters and not the actual memory Mt . But memory content plays a vital role in both read and write address generation : the content based look-up still finds de-allocated cells , resulting in memory aliasing . Consider the repeated copy task ( Section 4 ) , which needs repetitive allocation during storage of a sequence and de-allocation after it was read . The network has to store and repeat sequences multiple times . It has to tag the sequence beginning to know from where the repetition should start from . This could be done by content-based look-up . During the repetition phase , each cell read is also de-allocated . However when the repetition of the second sequence starts , the search for the tagged cell can find both the old and the new marked cell with equal score , making it impossible to determine which one is the correct match . We propose to zero out the memory contents by multiplying every cell of the memory matrix Mt with the corresponding element of the retention vector . Then the memory update equation becomes : Mt = Mt−1 ψt1T ( E −wwt e ᵀ t ) +w w t v ᵀ t ( 1 ) where is the element-wise product , 1 ∈ RN is a vector of ones , E ∈ RN×W is a matrix of ones . Note that the cosine similarity ( used for comparing the key to the memory content ) is normalized by the length of the memory content vector which would normally cancel the effect of Eq . 1 . However , in practice , due to numerical stability , cosine similarity is implemented as D ( u , v ) = u · v |u||v|+ where is a small constant . In practice , free gates f it tend to be almost 1 , so ψt is very close to 0 , making the stabilizing constant dominant with respect to the norm of the erased memory content vector . This will assign a low score to the erased cell in the content addressing : the memory is totally removed .

Summary: This paper is built on the top of DNC model. Authors observe a list of issues with the DNC model: issues with deallocation scheme, issues with the blurring of forward and backward addressing, and issues in content-based addressing. Authors propose changes in the network architecture to solve all these three issues. With toy experiments, authors demonstrate the usefulness of the proposed modifications to DNC. The improvements are also seen in more realistic bAbI tasks.