Adaptive Memory Networks Daniel Li 1 Asim Kadav 2 Abstract APIs are wrapped around hard-timeouts of 8 seconds which includes the time to transliterate the question to text on Ama- We present Adaptive Memory Networks (AMN) zon’s servers and the round-trip transfer time of question that processes input-question pairs to dynamically and the answer from the remote service, and sending the construct a network architecture optimized for response back to the device. Furthermore, developers are en- lower inference times for Question Answering couraged to provide a list of questions (“utterances”) apriori (QA) tasks. AMN processes the input story to at each processing step to assist QA processing (Amazon, extract entities and stores them in memory banks. 2017). We propose that apart from helping transliteration Starting from a single bank, as the number of in- these questions can also provide hints for reducing inference put entities increases, AMN learns to create new times for QA tasks based on large knowledge bases. banks as the entropy in a single bank becomes too high. Hence, after processing an input-question(s) Modeling QA tasks with LSTMs can be computationally pair, the resulting network represents a hierarchi- expensive which is undesirable during inference. Memory cal structure where entities are stored in different networks, a class of deep networks with explicit addressable banks, distanced by question relevance. At in- memory, have recently been used to achieve state of the ference, one or few banks are used, creating a art results on many QA tasks. Unlike LSTMs, where the tradeoff between accuracy and performance. number of parameters grows exponentially with the size of memory, memory networks are comparably parameter AMN is enabled by dynamic networks that allow efficient and can learn over longer input sequences. How- input dependent network creation and efficiency ever, they often require accessing all intermediate memory in dynamic mini-batching as well as our novel to answer a question. Furthermore, using focus of attention bank controller that allows learning discrete deci- over the intermediate state using a list of questions does sion making with high accuracy. In our results, we not address this problem. Soft attention based models com- demonstrate that AMN learns to create variable pute a softmax over all states and hard attention models are depth networks depending on task complexity and not differentiable and can be difficult to train over a large reduces inference times for QA tasks. state space. Previous work on improving inference over memory networks has focused on using unsupervised clus- tering methods to reduce the search space (Chandar et al., 1. Introduction 2016; Rae et al., 2016). Here, the memory importance is Question Answering (QA) tasks are gaining significance not learned and the performance of nearest-neighbor style due to their widespread applicability to recent commercial algorithms is often comparable to a softmax operation over applications such as chatbots, voice assistants and even med- memories. arXiv:1802.00510v1 [cs.AI] 1 Feb 2018 ical diagnosis (Goodwin & Harabagiu, 2016). Furthermore, To provide faster inference for long sequence-based inputs, many existing natural language tasks can also be re-phrased we present Adaptive Memory Networks (AMN), that con- as QA tasks. Providing faster inference times for QA tasks structs a memory network on-the-fly based on the input. is crucial. Consumer device based question-answer services Like past approaches to addressing external memory, AMN have hard timeouts for answering questions. For example, constructs the memory entity nodes dynamically. However, Amazon Alexa, a popular QA voice assistant, allows devel- distinct from past approaches, AMN stores the entities from opers to extend the QA capabilities by adding new “Skills” the input story in variable number of memory banks. The as remote services (Amazon, 2017). However, these service entities represent the hidden state of each word in the story 1University of California-Berkeley, Berkeley, USA, Work done while a memory bank is a collection of entities that are as a NEC Labs Intern. 2NEC Laboratories America, Princeton, similar w.r.t the question. As the number of entities grow, USA. Correspondence to: Daniel Li <[email protected]>, and the entropy within a single bank becomes too high, our Asim Kadav <[email protected]>. network learns to construct new memory banks and copies entities that are more relevant towards a single bank. Hence, Adaptive Memory Networks by limiting the decoding step to a dynamic number of rele- limits access during inference to specific entities reducing vant memory banks, AMN achieves lower inference times. inference times. AMN is an end-to-end trained model with dynamic learned Graph-based networks, (GG-NNs (Li et al., 2015) and GGT- parameters for memory bank creation and movement of NNs (Johnson, 2017)) use nodes with tied weights that are entities. updated based on gated-graph state updates with shared Figure1 demonstrates a simple QA task where AMN con- weights over edges. However, unlike AMN, they require structs two memory banks based on the input. During in- strong supervision over the input and teacher forcing to ference only the entities in the left bank are considered learn the graph structure. Furthermore, the cost of building reducing inference times. To realize its goals, AMN intro- and training these models is expensive and if every edge is duces a novel bank controller that uses reparameterization considered at every time-step the amount of computation trick to make discrete decisions with high accuracy while grows at the order of O(N 3) where N represents the num- maintaining differentiability. Finally, AMN also models ber of nodes/entities. AMN does not use strong supervision sentence structures on-the-fly and propagates update infor- but can solve tasks that require transitive logic by model- mation for all entities that allows it to solve all 20 bAbI ing sentence walks on the fly. EntNet constructs dynamic tasks. networks based on entities with tied weights for each en- tity (Henaff et al., 2017). A key-value update system allows 2. Related Work it to update relevant (learned) entities. However, Entnet uses soft-attention during inference to attend to all entities that Memory Networks: Memory networks store the entire incur high inference costs. To summarize, majority of the input sequence in memory and perform a softmax over hid- past work on memory networks uses softmax over memory den states to update the controller (Weston et al., 2014; nodes, where each node may represent input or an entity. Sukhbaatar et al., 2015). DMN+ connects memory to input In contrast, AMN learns to organize memory into various tokens and updates them sequentially (Xiong et al., 2016). memory banks and performs decode over fewer entities For inputs that consist of large number of tokens or en- reducing inference times. tities, these methods can be expensive during inference. AMN stores entities with tied weights in different mem- Conditional Computation & Efficient Inference: ory banks. By controlling the number of memory banks, AMN is also related to the work on conditional computation AMN achieves low inference times with reasonable accu- which allows part of networks to be active during inference racy. Nearest neighbor methods have also been explored improving computational efficiency (Bengio et al., 2015). over memory networks. For example, Hierarchical Memory Recently, this has been often accomplished using a gated Networks separates the input memory into groups using mixture of experts (Eigen et al., 2013; Shazeer et al., 2017). the MIPS algorithm (Chandar et al., 2016). However, us- AMN conditionally attends to entities in initial banks during ing MIPS is as slow as a softmax operation, so the authors inference improving performance. For faster inference propose using an approximate MIPS that gives inferior per- using CNNs, pruning (Le Cun et al., 1989; Han et al., 2016), formance. In contrast, AMN is end to end differentiable, low rank approximations (Denton et al., 2014), quantization and reasons which entities are important and constructs a and binarization (Rastegari et al., 2016) and other tricks network with dynamic depth. to improve GEMM performance (Vanhoucke et al., 2011) Neural Turing Machine (NTM) consists of a memory bank have been explored. For sequence based inputs, pruning and a differentiable controller that learns to read and write to and compression has been explored (Giles & Omlin, 1994; specific locations (Graves et al., 2014). In contrast to NTMs, See et al., 2016). However, compression results in irregular AMN memory bank controller is more coarse grained and sparsity that reduces memory costs but may not reduce the network learns to store entities in memory banks instead computation costs. Adaptive computation time (Graves, of specific locations. AMN uses a discrete bank controller 2016) learns the number of steps required for inferring that gives improved performance for bank controller actions the output and this can also be used to reduce inference over NTM’s mechanisms. However, our design is consistent times (Figurnov et al., 2016). AMN uses memory networks with the modeling studies of working memory by (Hazy with dynamic number of banks to reduce computation costs. et al., 2006) where the brain performs robust memory main- tenance and may maintain multiple working representations Dynamic networks: Dynamic neural networks that for individual working tasks. Sparse access memory uses change structure during inference have recently been possi- approximate nearest neighbors (ANN) to reduce memory ble due to newer frameworks such as Dynet and PyTorch. usage in NTMs (Rae et al., 2016). However, ANNs are Existing work on pruning can be implemented using these not differentiable. AMN, uses a input specific memory frameworks to reduce inference times dynamically like dy- organization that does not create sparse structures. This namic deep network demonstrates (Liu & Deng, 2017). Adaptive Memory Networks Mary moved to the bathroom.