Multigrid Neural Memory

1 INTRODUCTION . Memory , in the form of generic , high-capacity , long-term storage , is likely to play a critical role in expanding the application domain of neural networks . A trainable neural memory subsystem with such properties could be a transformative technology—pushing neural networks within grasp of tasks traditionally associated with general intelligence and an extended sequence of reasoning steps . Development of architectures for integrating memory units with neural networks spans a good portion of the history of neural networks themselves ( e.g . from LSTMs ( Hochreiter & Schmidhuber , 1997 ) to the recent Neural Turing Machines ( NTMs ) ( Graves et al. , 2014 ) ) . Yet , while useful , none has elevated neural networks to be capable of learning from and processing data on size and time scales commensurate with traditional computing systems . Recent successes of deep neural networks , though dramatic , are focused on tasks , such as visual perception or natural language translation , with relatively short latency— e.g . hundreds of steps , which is often also the depth of the network itself . We present a design for neural memory subsystems that differs dramatically from prior architectures in the aspect of read/write interface , and , as a consequence , facilitates parameter-efficient scaling of memory capacity . Our design draws upon both established ideas , as well a key recent advance : multigrid convolutional networks implicitly capable of learning dynamic routing mechanisms and attentional behavior ( Ke et al. , 2017 ) . LSTMs serve as a fine-grained component , but are wrapped within a larger multigrid connection topology , from which novel capabilities emerge . In contrast to NTMs and the subsequent Differentiable Neural Computer ( DNC ) ( Graves et al. , 2016 ) , we intertwine memory units throughout the interior of a deep network . Memory is a first-class citizen , rather than a separate data store accessed via a special controller . We introduce a new kind of network layer—a multigrid memory layer—and use it as a stackable building block to create deep memory networks . Contrasting with simpler LSTMs , our memory is truly deep ; accessing an arbitrary memory location requires passing through several layers . Figure 1 provides a visualization of our approach ; we defer the full details to Section 3 . There are major benefits to this design strategy , in particular : • Memory scalability . Distributing storage over a multigrid hierarchy allows us to instantiate large amounts while remaining parameter-efficient . Read and write operations are similarly distributed . The low-level mechanism underlying these operations is essentially convolution , and we inherit the parameter-sharing efficiencies of convolutional neural networks ( CNNs ) . Our parameterized filters simply act across a spatially organized collection of memory cells , rather than the spatial extent of an image . Increasing feature channels stored in each spatial cell costs parameters , but adding more cells incurs no such cost , decoupling memory size from the number of parameters . The multigrid layout of our memory network provides an information routing mechanism that is efficient with respect to overall network depth . We can grow the spatial extent of memory exponentially with depth , while guaranteeing there is a connection pathway between the network input and every memory unit . This allows experimentation with substantially larger memories . • Unification of compute and storage . Our memory layers incorporate convolutional and LSTM components . Stacking such layers , we create not only a memory network , but also a generalization of both a CNN and an LSTM . Our memory networks are standard networks with additional capabilities . Though not our primary experimental focus , Section 4 shows they can learn tasks which require performing classification alongside storage and recall . A further advantage to this unification is that it opens a wide design space for connecting our memory networks to each other , as well as standard neural network architectures . For example , within a larger system , we could easily plug the internal state of our memory into a standard CNN—essentially granting that CNN read-only memory access . Sections 3 and 4 develop and experimentally validate two such memory interface approaches . A diverse array of synthetic tasks serves as our experimental testbed . Mapping and localization , an inherently spatial task with relevance to robotics , is one focus . However , we want to avoid only experimenting with tasks naturally fit to the architecturally-induced biases of our memory networks . Therefore , we also train them to perform the same kind of algorithmic tasks used in analyzing the capabilities of NTMs and DNCs . Throughout all experimental settings , DNC accuracy serves as a comparison baseline . We observe significant advantages for multigrid memory , including : • Long-term large-capacity retention . On spatial mapping tasks , our architecture retains large , long-term memory . It correctly remembers observations of an external environment collected over paths consisting of thousands of time steps . Visualizing internal memory unit activations actually reveals the representation and algorithmic strategy our network learns in order to solve the problem . The DNC , in contrast , fails to master this category of task . • Generality . On tasks decoupled from any notion of spatial geometry , such as associative recall or sorting , our memory networks prove equally as capable as DNCs . Section 4 further elaborates experimental results , while Section 5 summarizes implications . Before turning to technical details ( Section 3 ) , we review the history of coupling neural networks to memory . 2 RELATED WORK . There is an extensive history of work that seeks to grant neural networks the ability to read to and write from memory ( Das et al. , 1992 ; 1993 ; Mozer & Das , 1993 ; Zeng et al. , 1994 ; Hölldobler et al. , 1997 ) . Das et al . ( 1992 ) propose neural network pushdown automaton , which performs differential push and pop operations on external memory . Schmidhuber ( 1992 ) uses two feedforward networks : one that produces context-dependent weights for the second network , whose weights may change quickly and can be used as a form of memory . Schmidhuber ( 1993 ) proposes memory addressing in the form of a “ self-referential ” recurrent neural network that is able to modify its own weights . Recurrent neural networks with Long Short-Term Memory ( LSTMs ) ( Hochreiter & Schmidhuber , 1997 ) have enabled significant progress on a variety of sequential prediction tasks , including machine translation ( Sutskever et al. , 2014 ) , speech recognition ( Graves et al. , 2013 ) , and image captioning ( Donahue et al. , 2017 ) . LSTMs are Turing-complete ( Siegelmann & Sontag , 1995 ) and are , in principle , capable of context-dependent storage and retrieval over long time periods ( Hermans & Schrauwen , 2013 ) . However , a network ’ s capacity for long-term read-write is sensitive to the training procedure ( Collins et al. , 2017 ) and is limited in practice . Grid LSTMs ( Kalchbrenner et al. , 2015 ) arrange LSTM cells in a 2D or 3D grid , placing recurrent memory links along all axes of the grid . This sense of grid differs from our usage of multigrid , as the latter refers to links across a multiscale spatial layout . In Kalchbrenner et al . ( 2015 ) ’ s terminology , our multigrid memory networks are not Grid LSTMs , but are a variant of Stacked LSTMs ( Graves et al. , 2013 ) . In an effort to improve the long-term read-write abilities of recurrent neural networks , several modifications have been recently proposed . These include differentiable attention mechanisms ( Graves , 2013 ; Bahdanau et al. , 2014 ; Mnih et al. , 2014 ; Xu et al. , 2015 ) that provide a form of content-based memory addressing , pointer networks ( Vinyals et al. , 2015 ) that “ point to ” rather than blend inputs , and architectures that enforce independence among the neurons within each layer ( Li et al. , 2018 ) . A number of methods augment the short- and long-term memory internal to recurrent networks with external “ working ” memory , in order to realize differentiable programming architectures that can learn to model and execute various programs ( Graves et al. , 2014 ; 2016 ; Weston et al. , 2015 ; Sukhbaatar et al. , 2015 ; Joulin & Mikolov , 2015 ; Reed & de Freitas , 2015 ; Grefenstette et al. , 2015 ; Kurach et al. , 2015 ) . Unlike our approach , these methods explicitly decouple memory from computation , mimicking a standard computer architecture . A neural controller ( analogous to a CPU ) interfaces with specialized external memory ( e.g. , random-access memory or tapes ) . The Neural Turing Machine ( NTM ) augments neural networks with a hand-designed attention-based mechanism to read from and write to external memory in a differentiable fashion . This enables the NTM to learn to perform various algorithmic tasks , including copying , sorting , and associative recall . The Differential Neural Computer ( Graves et al. , 2016 ) improves upon the NTM with support for dynamic memory allocation and additional memory addressing modes . Without a sparsifying approximation , DNC runtime grows quadratically with memory due to the need to maintain the temporal link matrix . Our architecture has no such overhead or approximations , nor does it require maintaining an auxiliary state . Other methods enhance recurrent layers with differentiable forms of a restricted class of memory structures , including stacks , queues , and dequeues ( Grefenstette et al. , 2015 ; Joulin & Mikolov , 2015 ) . Gemici et al . ( 2017 ) augment structured dynamic models for temporal processes with various external memory architectures ( Graves et al. , 2014 ; 2016 ; Santoro et al. , 2016 ) . Similar memory-explicit architectures have been proposed for deep reinforcement learning ( RL ) tasks . While deep RL has been applied successfully to several challenging domains ( Mnih et al. , 2015 ; Hausknecht & Stone , 2015 ; Levine et al. , 2016 ) , most approaches reason over short-term representations of the state , which limits their ability to deal with partial observability inherent in many RL tasks . Several methods augment deep RL architectures with external memory to facilitate long-term reasoning . Oh et al . ( 2016 ) maintain a fixed number of recent states in memory and then read from the memory using a soft attention-based read operation . Parisotto & Salakhutdinov ( 2018 ) propose a specialized write operator , together with a hand-designed 2D memory structure , both specifically designed for navigation in maze-like environments . Rather than learn when to write to memory ( e.g. , as done by NTM and DNC ) , Pritzel et al . ( 2017 ) continuously write the experience of an RL agent to a dictionary-like memory module queried in a key-based fashion ( permitting large memories ) . Building on this framework , Fraccaro et al . ( 2018 ) augment a generative temporal model with a specialized form of spatial memory that exploits privileged information , including an explicit representation of the agent ’ s position in the environment . Though we experiment with RL , our memory implementation contrasts with this past work . Our multigrid memory architecture jointly couples computation with memory read and write operations , and learns how to use a generic memory structure rather than one specialized to a particular task . 3 MULTIGRID MEMORY ARCHITECTURES . To endow neural networks with long-term memory , we craft an architecture that generalizes modern convolutional and recurrent designs , embedding memory cells within the feed-forward computational flow of a deep network . Convolutional neural networks and LSTMs ( specifically , the convolutional LSTM variety ( Xingjian et al. , 2015 ) ) exist as strict subsets of the full connection set comprising our multigrid memory network . We even encapsulate modern residual networks ( He et al. , 2016 ) . Though omitted from diagrams for the sake of clarity , in all experiments we utilize residual connections linking the inputs of subsequent layers across the depth ( not time ) dimension of our memory networks . Implementing memory addressing behavior is the primary challenge when adopting our design philosophy . If the network structure is uniform , how will it be capable of selectively reading from and writing to only a sparse , input-dependent subset of memory locations ? A common approach is to build an explicit attention mechanism into the network design . Such attention mechanisms , independent of memory , have been hugely influential in natural language processing ( Vaswani et al. , 2017 ) . NTMs ( Graves et al. , 2014 ) and DNCs ( Graves et al. , 2016 ) construct a memory addressing mechanism by explicitly computing a soft attention mask over memory locations . This naturally leads to a design reliant on an external memory controller , which produces and then applies that mask when reading from or writing to a separate memory bank . Ke et al . ( 2017 ) recently proposed a multigrid variant of both standard CNNs and residual networks ( ResNets ) . While their primary experiments concern image classification , they also present a striking result on a synthetic image to image transformation task : multigrid CNNs ( and multigrid ResNets ) are capable of learning to emulate attentional behavior . Their analysis reveals that the network ’ s multigrid connection structure is both essential to and sufficient for enabling this phenomenon . The underlying cause is that bi-directional connections across a scale-space hierarchy ( Figure 1 , left ) create exponentially shorter signalling pathways between units at different locations on the spatial grid . Information can be efficiently routed from one spatial location to any other location by traversing only a few network layers , flowing up the scale-space hierarchy , and then back down again . We convert the inherent attentional capacity of multigrid CNNs into an inherent capacity for distributed memory addressing by replacing convolutional subcomponents with convolutional LSTMs ( Xingjian et al. , 2015 ) . Grid “ levels ” no longer correspond to operations on a multiresolution image representation , but instead correspond to accessing smaller or larger storage banks within a distributed memory hierarchy . Dynamic routing across scale space ( in the multigrid CNN ) now corresponds to dynamic routing into different regions of memory , according to a learned strategy . Appendix A provides an analysis of this dynamic routing capability . This routing mechanism is the foundation that enables the addressing capability in multigrid memory architectures .

Summary: The paper proposes a multigrid memory architecture by introducing multigrid CNN [1] into convolutional LSTM network [2]. The method extends the convolutional LSTM with bigger memory capacity in forms of multigrid CNNs. Some specific designs such as multiple threads and encoder-decoder are also proposed. The model is validated with synthetic tasks such as spatial mapping, associative recall and priority sort.