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 proposes modifications to the original Differentiable Neural Computer architecture in three ways. First by introducing a masked content-based addressing which dynamically induces a key-value separation. Second, by modifying the de-allocation system by also multiplying the memory contents by a retention vector before an update. Finally, the authors propose a modification in the link distribution, through renormalization. They provide some theoretical motivation and empirical evidence that it helps avoiding memory aliasing.