Distributed Associative Memory Network with Association Reinforcing Loss

1 INTRODUCTION . The essential part of human intelligence for understanding the story and predicting unobserved facts largely depends on the ability of memorizing the past and reasoning for relational information based on the pieces of memory . In this context , research on artificial intelligence has focused on designing human like associative memory network which can easily store and recall both events and relational information from a part of information . In neural network research , many approaches generally model sequential data with memory systems , such as Long Short Term Memory ( LSTM ) ( Hochreiter & Schmidhuber , 1997 ) and memoryaugmented neural networks ( MANN ) . Especially , recent approach in MANN constructs an associative memory with a content-based addressing mechanism and stores both input data and its relational information to a single external memory . MANN has already proven to be an essential component on many tasks which need long term context understanding ( Weston et al. , 2014 ; Sukhbaatar et al. , 2015 ; Graves et al. , 2014 ; 2016 ; Gulcehre et al. , 2018 ) . Also , compared to recurrent neural networks , it can store more information from sequential input data and correctly recall desired information from memory with a given cue . However , even with its promising performance on a wide range of tasks , MANN still has difficulties in solving complex relational reasoning problems ( Weston et al. , 2015 ) . Since content-based addressing model implicitly encodes data item and its relational information into one vector representation , they often result in a lossy representation of relational information which is not rich enough for solving relational reasoning tasks . To address such weakness , some researches find relational information by leveraging interaction between memory entities with attention ( Palm et al. , 2018 ; Santoro et al. , 2018 ) . Others focus on long sequence memorization performance of memory ( Trinh et al. , 2018 ; Le et al. , 2019 ; Munkhdalai et al. , 2019 ) . Another attempts apply a self-attention to memory contents and explicitly encode relational information to a separate external memory ( Le et al. , 2020b ) . However , all those models need to explicitly find relational information among memory entities with highly computational attention mechanism and have to repeatedly recompute it on every memory update . In this research , we approach the same problem in a much simpler and efficient way which do not require any explicit relational computation , such as self-attention , among memory entities . We hypothesize that lossy representation of relational information ( Le et al. , 2020b ) in MANN is caused by both a single memory based representation and long-temporal data association performance . Although MANN learns to correlate sequential events across time , its representation is not rich enough to reflect complex relational information existing in input data . Therefore , for the enhanced relation learning , we focus on the richness of representation which implicitly embeds associations existing in input data . For this purpose , we introduce a novel Distributed Associative Memory ( DAM ) architecture which is inspired by how the biological brain works ( Lashley , 1950 ; Bruce , 2001 ) . In DAM , we replace the single external memory with multiple smaller sub-memory blocks and update those memory blocks simultaneously and independently . The basic operations for each associative memory block are based on the content-based addressing mechanism of MANN , but its parallel memory architecture allows each sub-memory system to evolve over time independently . Therefore , similar to the underlying insight of multi-head attention ( Vaswani et al. , 2017 ) , our memory model can jointly attend to information from different representation subspaces at different sub-memory blocks and is able to provide a more rich representation of the same common input data . To retrieve rich information for relational reasoning , we apply a soft-attention based interpolation to the diverse representations distributed across multiple memories . Moreover , to enrich long-term relational information in the memory , we introduce a novel association reinforcing loss ( ARL ) which fortifies data associations of the memory and generally enhances the memorization capacity of MANN . The ARL forces the memory network to learn to reproduce the number of stochastically sampled input data only based on the stored memory contents . As if , other associated pieces of memory are reminded together whenever a person recalls a certain event in his memory , the data reproducing task enables MANN to have better association and memorization ability for input data . It is designed to reproduce a predefined percentage of input representations in the memory matrix on average and , while optimizing two different tasks at the same time , keep the balance between ARL and target objective loss by dynamically re-weighting each task ( Liu & Zhou , 2006 ; Cui et al. , 2019 ) . By combining the above two approaches , DAM , and ARL , our architecture provides rich representation which can be successfully used for tasks requiring both memorization and relational reasoning . We apply our architecture to Differential Neural Computer ( DNC ) ( Graves et al. , 2016 ) , which is one of the representative content-based addressing memory , to construct novel distributed associative memory architecture with ARL . DNC has promising performance on diverse tasks but also known to be poor at complex relational reasoning tasks . In experiments , we show that our architecture greatly enhances both memorization and relation reasoning performance of DNC , and even achieves the state-of-the-art records . 2 DIFFERENTIABLE NEURAL COMPUTER . We first briefly summarize DNC architecture which is a baseline model for our approaches . DNC ( Graves et al. , 2016 ) is a memory augmented neural network inspired by conventional computer architecture and mainly consists of two parts , a controller and an external memory . When input data are provided to the controller , usually LSTM , it generates a collection of memory operators called as an interface vector ξt for accessing an external memory . It consists of several keys and values for read/write operations and constructed with the controller internal state ht as ξt = Wξht at each time step t. Based on these memory operators , every read/write operation on DNC is performed . During writing process , DNC finds a writing address , wwt ∈ [ 0 , 1 ] A , where A is a memory address size , along with write memory operators , e.g . write-in key , and built-in functions . Then it updates write-in values , vt ∈ RL , in the external memory , Mt−1 ∈ RA×L , along with erasing value , et ∈ [ 0 , 1 ] L , where L is a memory length size as follows : Mt =Mt−1 ◦ ( E −wwt e > t ) +wwt v > t ( 1 ) where ◦ denotes element-wise multiplication and E is 1A×L . In the reading process , DNC searches a reading address , wr , it ∈ [ 0 , 1 ] A , for R read heads , along with read memory operators , e.g . read-out key . Then , it reads out information from the external memory : rit =Mtw r , i t > ( 2 ) Finally , the output is computed as yt = Wy [ ht ; rt ] ∈ Rdo , where rt = { rit ∈ RL ; 1 ≤ i ≤ R } . Through these operations , DNC can learn how to store input data and utilize stored information to solve a given task . These whole mechanisms make DNC suitable for a general purposed memory augmented neural network . 3 PROPOSED METHOD . In this section , we introduce two methods that improve both memorization and relational reasoning ability of conventional DNC , a distributed associative memory architecture , and an ARL function . For a clear explanation , we illustrate DAM mechanism with a single read head case . For R read head cases of DAM , the details are in the Appendix . 3.1 DISTRIBUTED ASSOCIATIVE MEMORY ARCHITECTURE . The distributed associative memory architecture consists of a controller network and K associative memory blocks where each memory block is a content addressable memory similar to the original DNC ( Graves et al. , 2016 ) . Figure 1 ( a ) shows the overall read/write process of the proposed DAM . For the writing operation , the controller of DAM produces multiple writing operator vectors for multiple memory blocks . Each writing operator vector is used for the content-based addressing of one of the multiple memory blocks , and it is independent of other memory blocks . Since it is produced based on the current input and previous hidden states of the controller , it can independently store its own representation of the same input contents . This writing process enables DAM to store the diverse representations of the same input data to multiple memory blocks with much flexibility . Furthermore , for the reading process , all memory blocks are read at the same time , and read values are interpolated with soft attention to produce single read-out information . Through this attention-based reading process , DAM retrieves the most suitable information for the current task from representations distributed in the multiple memory blocks . Based on these read/write operations , DAM learns how to store and retrieve the diverse representations of input data for different purposed tasks . The following sections detail the main operations . 3.1.1 CONTROLLER FOR MULTIPLE ASSOCIATIVE MEMORY BLOCKS . At each time step t , the controller receives an external input , it , read-out of the previous time step , rt−1 , and previous hidden state of controller , ht−1 , to update its current hidden state , ht . After layer normalization , it produces an interface vector , ξt ∈ RK∗ ( L∗R+3L+3R+3 ) , which includes read and write parameters for multiple memory access . 3.1.2 WRITE INTO MULTIPLE SUB-MEMORY BLOCKS . The multiple memory writing processes in our architecture is based on the content-based memory accessing mechanism of DNC . A single memory block is addressed and updated with the same procedure of DNC and such single memory block updating is applied to all blocks independently at the same time . As shown in Eq . ( 3 ) , each memory block has its own interface vector relevant weight Wξ,1 , · · · , Wξ , k , where k ∈ { 1 , · · · , K } . Theses weights are multiplied with a controller hidden state vector , ht , and used for memory operations of each independent memory block as following . ξt = [ ξt,1 , · · · , ξt , K , ĝatt ] = [ Wξ,1 , · · · , Wξ , K , Wξ , at ] ht ( 3 ) where ξt , k is a interface vector for each memory block and ĝatt is an attentive gate at time t. Based on a writing operator obtained from ξt , k , DAM updates input information into each memory block , Mt−1 , k , independently and simultaneously , following Eq . ( 1 ) . That independent and simultaneous writing procedures of sub-memory blocks allow that our DAM learns to construct diverse representations for the same common input data . The following attention-based reading process is designed to integrate representations distributed across sub-memory blocks , and it contributes to enrich representation for relational reasoning tasks . 3.1.3 READ FROM MULTIPLE SUB-MEMORY BLOCKS . As in the writing process , DAM obtains a reading operator from ξt , k , and computes reading address , wrt , k ∈ [ 0 , 1 ] A , for each memory block . Based on those addresses , DAM reads values from each memory block and derives read-out value , rt ∈ RL , from them , using processed attentive gate , gatt ∈ [ 0 , 1 ] K , as follows : rt = K∑ k=1 gatt , kM > t , kw r t , k where g at t , k = Softmax ( ĝ at t , k ) for k = 1 , · · · , K. ( 4 ) Compared to Eq . ( 2 ) of DNC , this reading process integrates representations stored in multiple memory blocks with attention gate and enables DAM to learn to provides the most appropriate distributed representation for a target task .

Summary: The authors propose a distributed memory architecture which shares some interface with the Differentiable Neural Computer however crucially segments memory into a collection of K units. The authors show that by increasing K the model learns to use its memory for algorithmic tasks such as copying and associative recall and learn faster. The authors also propose an auxiliary loss to improve memory representations, which involves reconstructing inputs from the representations in memory.