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
Mary moved to the bathroom. List of Strength GRU questions John went to the hallway. Bank Controller Daniel went back to the hallway.
Sandra moved to the garden. Sandra John moved to the office.
Sandra journeyed to the bathroom.
Mary moved to the hallway.
Daniel travelled to the office.
John went back to the garden. Memory Bank 1 Memory Bank 2 State of entities
John moved to the bedroom. Entities from the Where is Sandra? story
Figure 1. Overview of Adaptive memory networks. Multiple memory banks are created based on the story and input entities are moved in them based on their relevance to the question. Inference is performed on a single (or less than all) banks most relevant to the question(s). AMN utilizes the dynamic architecture abilities to construct bank is a collection of entities. Intuitively, each memory an input dependent memory network of variable memory bank stores entities that have a similar distance score from bank depth and the dynamic batching feature to process a the question. variable number of entities. Furthermore, unlike past work At a high level, entities are gradually and recurrently copied that requires an apriori number of fixed memory slots, AMN through memory banks to filter out irrelevant nodes such constructs them on-the-fly based on the input. The learn- that in the final inference stage, fewer entities are considered able discrete decision-making process can be extended to by the decoder. Note that the word filter implies a discrete other dynamic networks which often rely on REINFORCE decision and that recurrence implies time. If we were to to make such decisions (Liu & Deng, 2017). perform a strict cut off and remove entities that appear to be irrelevant at each time step, learning the reasoning logic Neuroscience: Our network construction is inspired by that requires previous entities that were cut off would not work on working memory representations. There is suf- be possible. Thus, smoothed discretization is required. ficient evidence for multiple, working memory represen- We design filtering to be a two-stage pseudo-continuous tations in the human brain (Hazy et al., 2006). Semantic process to simulate discrete cut offs (Π , Π ), while memory (Tulving et al., 1972), describes a hierarchical or- move new keeping reference history. The overall memory (M) consists ganization starting with relevant facts at the lowest level and of multiple memory banks. A memory bank is a collection progressively more complex and distant concepts at higher or group of entities (m ), where m denotes the initial levels. AMN constructs entities from the input stories and 0...l 0 and most general bank and m denotes the most relevant stores the most relevant entities based on the question in the l bank. Note that |l| is input dependent and learned. First, lowest level memory bank. Progressively higher level mem- entities are moved from m gradually towards m based off ory banks represent distant concepts (and not necessarily 0 l of their individual relevance to the question and second, if higher level concepts for AMN). Other work demonstrates m becomes too saturated, m is created. Operations in organization of human memory in terms of “priority struc- l l+1 the external memory allowing for such dynamic restructur- ture” where attention is a gate-keeper of working memory- ing and entity updates are described below. Note that these guided by executive control’s goals, plans, and intentions as operations still maintain end to end differentiability. in (Watzl, 2017), similar in spirit to AMN’s question guided network construction. 1. Memory bank creation (Πnew), which creates a new memory bank depending on the current states of en- 3. Differentiable adaptive memory module tities mi. If the entropy, or information contained (explained below), of m is too high, Π (m ) will In this section, we describe the design process and motiva- i new i learn to create a new memory bank m to reduce tion of our memory module. Our memory network architec- i+1 entropy. ture is created during inference time for every story. The architecture consists of different memory banks and each 2. Moving entities across banks (Πmove), which deter- memory bank stores entities from the input story. Hence, mines which entities are relevant to the current question a memory entity represents the hidden state of each entity and move such entities to further (higher importance) (each word in our case) from the input story while a memory memory banks. Adaptive Memory Networks
3. Adding/Updating entities in a bank (Πau), which adds In the following, the subscripts i, j are used to iterate entities that are not yet encountered to the first memory through the total number of words in a statement and ques- bank m0 or if the entity is already in m0, the operation tion respectively, D stores the accumulation GRU output, updates the entity state. and wi is a GRU encoding output. The last output of the
GRU will be referred to as wN , wNq for statements and 4. Propagating changes across entities (Πprop), which up- questions. dates the entity states in memory banks based on node current states Πprop(M) and their semantic relation- ui, uj = EMBED(wii),EMBED(wij ) (1)
ships. This is to communicate transitive logic. wi = GRU(ui, wi−1) (2)
D[i] += wi (3)
Both Πnew, Πmove require a discrete decision (refer to sec- wj = GRU(uj , wj−1) (4) tion 4.2.1.), and in particular, for Πnew we introduce the no- tion of entropy. That is to say if mi contains too many nodes To compute the question relevance strength s ∈ [0, 1] for (the entropy becomes too high), the memory module will each word, the model uses GRU-like equations. The node learn to create a new bank mi+1 and move nodes to mi+1 strengths are first initialized to Xavier normal and the inputs in to reduce entropy. By creating more memory banks, the are the current word states w , the question state wNq , model spreads out the concentration of information which and when applicable, the previous strength. Sentences are in turn better discretizes nodes according to relevance. processed each time step t.
t in t−1 z = σ(Uzw + WzwNq + Xzs ) (5) 4. Adaptive Memory Networks t t−1 r = 1 − σ(Urhs , wNq i) (6) A high-level overview is shown in Figure2, followed by t in t t−1 es = σ(Whw + Uh(r s )) (7) a mathematical detail of the model’s modules. Our model st = zt st−1 + (1 − zt) st (8) adopts the encoder-decoder framework with an augmented e adaptive memory module. For an overview of the algorithm, In particular, equation (6) shows where the model learns to refer to Section A.4. lower the strengths of nodes that are not related the question. First, a dot product between the current word states and Notation and Problem Statement: Given a story rep- question state are computed for similarity (high correlation), resented by N input sentences (or statements), i.e., then it is subtracted from a 1 to obtain the dissimilarity. (l , ··· , l ), and a question q, our goal is to generate an an- 1 N We refer to these operations as SGRU (Strength GRU) in swer a. Each sentence l is a sequence of N words, denoted Algorithm1. as (w1, ··· , wNs ), and a question is a sequence of Nq words denoted as (w , ··· , w ). Throughout the model we refer 1 Nq 4.2. Adaptive Memory Module to entities; these can be interpreted as a 3-tuple of ew = (word ID wi, hidden state w, question relevance strength s). The adaptive memory module recurrently restructures enti- Scalars, vectors, matrices, and dot products are denoted by ties in a question relevant manner so the decoder can then lower-case letters, boldface lower-case letters and boldface consider fewer entities (namely, the question relevant enti- capital letters, and angled brackets respectively. ties) to generate an answer. The following operations are performed once per sentence. 4.1. Encoder 4.2.1. MEMORY BANK CONTROLLER The input to the model, starting with the encoder, are story- question input pairs. On a macro level, sentences l1...N are As mentioned earlier, discrete decisions are difficult for processed. On a micro level, words w1...Ns are processed neural networks to learn so we designed a specific mem- within sentences. ory bank controller Πctrl for binary decision making. The model takes ideas from the reparameterization trick and uses For each w ∈ l , the encoder maps w to a hidden represen- i i i custom backpropagation to maintain differentiability. tation and a question relevance strength ∈ [0, 1]. The word ID of wi is passed through a standard embedding layer and In particular, the adaptive memory module needs to make then encoded through an accumulation GRU. The accumula- two discrete decisions on a {0, 1} basis, one in Πnew to tion GRU captures the entity states through time by adding create a new memory bank and the other in Πmove to move the output of each GRU time step to its respective word, nodes to a different memory bank. The model uses a scalar stored in a lookup matrix. The initial states of ew are set to p ∈ {0, 1} to parameterize a Bernoulli distribution where this GRU output. Meanwhile, the question is also embedded the realization H, is the decision the model makes. How- and encoded in the same manner sans accumulation. ever, backpropagation through a random node is intractable, Adaptive Memory Networks
Encoder Semantic memory module Decoder Question Strength GRU Bank Controller
GRU
GRU Attention Answer
GRU
GRU Memory Banks Story
Figure 2. Adaptive memory networks. so the model detaches H from the computation graph and it cannot contain duplicate nodes, but the same node introduces H as a new node. Finally, H is used as a mask to can exist in two different memory banks. zero out entities in the discrete decision. Meanwhile, p is kept in the computation graph and has a special computed loss (Section 4.4). The operations below will be denoted as Πmove(si ∗ wi) = m.add(1(Πctrl(si ∗ wi))) ∀i ∈ m Πctrl and has two instances: one for memory bank creation (12) Πnew and one for moving entities across banks Πmove. In equation 9, depending on what Πctrl is used for, q is a poly- morphic function and will take on a different operation and 3. Adding/Updating entities in a bank: Recall that en- ∗ will be a different input. Examples of such are given in tities are initially set to the output of D. However, as their respective sections (4.2.2.1, 4.2.2.2). additional sentences are processed, new entities and their hidden states are observed. In the case of a new entity ew, the entity is added to the first memory bank p = q(∗) (9) H = Bernoulli(p) (10) m0. If the entity already exists in m0, then ew’s corre- sponding hidden state is updated through a GRU. This 4.2.2. MEMORY BANK OPERATIONS procedure is done for all memory banks.
1. Memory bank creation Πnew: To determine when a new memory bank is created, in other words, if the cur- rent memory bank becomes too saturated, the memory ( i i m0.add(ew) if ew 6∈ m0 else bank controller (4.2.1.) will make a discrete decision Πau([w0...wi]) = t+1 t w = GRU(wN , w ) ∀m ∈ M to create a new memory bank. Here, q (eq 9) is a i i (13) fully connected layer and the input is the concatena- tion of all the current memory bank mi’s entity states 1,n|ew| 4. Propagating updates to related entities: So far, en- [w0...wi] ∈ R . Intuitively, q will learn a contin- uous decision that is later discretized by eq 10 based tities exist as a bag of words model and the sentence on entity states and the number of entities. Note this is structure is not maintained. This can make it difficult only performed for the last memory bank. to solve tasks that require transitive reasoning over multiple entities. To track sentence structure informa- tion, we model semantic relationships as a directed A ( graph stored in adjacency matrix . As sentences are M.new() if 1(Πctrl([w0...wi])) else processed word by word, a directed graph is drawn Πnew([w0...wi]) = pass progressively from w0...wi...wN . If sentence lk’s path (11) contains nodes already in the current directed graph, lk will include said nodes in its path. After lk is added to 2. Moving entities through memory banks: Similar to A, the model propagates the new update hidden state information ai among all node states using a GRU. ai Πnew, individual entities’ relevance scores are passed into the bank controller to determine H as the input. for each node i is equal to the sum of the incoming The relevance score is computed by multiplying an edges’ node hidden states. n,|ew| entity state by its respective relevance ∈ R . Here, Additionally, we add a particular emphasis on lk to q has a slight modification and is the identity function. simulate recency. At face value, one propagation step Note that this operation can only be performed if there of A will only have a reachability of its immediate is a memory bank to move nodes to, namely if mi+1 neighbor, so to reach all nodes, A is raised to a consec- exists. Additionally, each bank has a set property where utive power r to reach and update each intermediate Adaptive Memory Networks
node. r can be either the longest path in A or a set by nodes that are connected to the answer node a in a di- parameter. Again, this is done within a memory bank rected graph; words that are connected to a are relevant to for all memory banks. For entities that have migrated a. They are then weighted with a deterministic function of to another bank, the update for these entities is a no-op distance from a. but propagation information as per the sentence struc- ture is maintained. A single iteration is shown below: Ls(s) = DKL(sli ||E[sli ]) Additionally, bank creation is kept in check by constraining r T t t−1 a = (A ) [w0...wi] (14) w = GRU(a, w ) (15) pli w.r.t. the expected number of memory banks. The expected number of memory banks can be thought of as When nodes are transferred across banks, A is still a geometric distribution ∼ Geometric(ˆp ) parameterized preserved. If intermediate nodes are removed from a li by pˆ , a hyperparameter. Typically, at each sentence step path, a transitive closure is drawn if possible. li pˆ is raised to the inverse power of the current sentence step to reflect the amount of information ingested. Intuitively, After these steps are finished at the end of a sentence, this loss ensures banks are created when a memory bank namely, the memory unit has reasoned through how large contains too many nodes. On the other hand, the learned (number of memory banks) the memory should be and mask q (eq. 9) enables the model to weight certain nodes a which entities are relevant at the current point in the story, all higher entropy to prompt bank creation. Through these two entities are passed through the strength modified GRU (4.1, dependencies, the model is able to simulate bank creation eq 5-8) to recompute their question relevance (relevance as a function of the number of nodes and the type of nodes score). in a given memory bank.
1 4.3. Decode β|li| Lb(pli ) = DKL(pli ||pˆ ) After all sentences l1...N are ingested, the decode portion of All components combined, the final loss is given in the the network learns to interpret the results from the memory following equation banks. The network iterates through the memory banks using a standard attention mechanism. To force the network |ln| ˆ X i i to understand the question importance weighting, the model Ltotal = Lp(L) + (Ls(s) + Lb(p)) uses an exponential function d to weight important memory i=1 banks higher. Cm are the hidden states contained in memory m, sm are the relevance strengths of memory bank m, wNq 5. Evaluation is the question hidden state, ps is the attention score, r, h are learned weight masks, g are the accumulated states, and l is In this section, we evaluate AMN accuracy and inference the final logits prediction. During inference, fewer memory times on the bAbI dataset (Weston et al., 2015) and extended banks are considered. bAbI tasks dataset. We compare our performance with Ent- net (Henaff et al., 2017), which recently achieved state of the art results on the bAbi dataset. For accuracy measure- ments, we also compare with DMN+ and encoder-decoder Cm = sm · [w0, ...wi] ∀i ∈ m (16) methods. Finally we discuss the time trade offs between ps = Softmax(hC , w i) (17) m Nq AMN and current SOTA methods. The portion regarding g += d(hC , psi) (18) m inference times are not inclusive of story ingestion. We ˆ L = r(PReLU(h(g) + wNq ) if m is last (19) summarize our experiments results as follows:
4.4. Loss • We are able to solve all bAbi tasks using AMN. Fur- Loss is comprised of two parts, answer loss, which is com- thermore, AMN is able to reason important entities and puted from the given annotations, and secondary loss (from propagate them to the final memory bank allowing for 48% fewer entities examined during inference. Πnew, Πmove), which is computed from sentence and story features at each sentence time step l0...N . Answer loss is • We construct extended bAbI tasks to evaluate AMN standard cross entropy at the end of the story after lN is behavior. First, we extend Task 1 for multiple questions processed. in order to gauge performance in a more robust manner. For example, if a reasonable set of questions are asked Lp(Lˆ) = CrossEntropy(Lˆ, L) (where reasonable means that collectively they do not
After each sentence li, the node relevance sli is enforced by require all entities to answer implying entities can be computing the expected relevance E[sli ]. E[s] is determined filtered out), will the model still sufficiently reason Adaptive Memory Networks
through entities. We find that our network is able to Inference performance Table4 shows the number of reason useful entities for both tasks and store them in banks created and required to solve a task, as well as the the final memory bank. Furthermore, we also scale ratio of entities examined to solve the task. Table 3 shows bAbI for a large number of entities and find that AMN the complexity of AMN and other SOTA models. Entnet provides additional benefits at scale since only relevant uses an empirically selected parameter, typically set to the entities are stored in the final memory bank. number of vocabulary words. GGT-NN uses the number of vocabulary words and creates new k new nodes intermit- tently per sentence step. 5.1. Experiment settings For tasks where nodes are easily separable where nodes are We implement our network in PyTorch (Paszke et al., 2017). clearly irrelevant to the question(s), AMN is able to suc- We initialize our model using Xavier initialization, and the cessfully reduce the number of nodes examined. However word embeddings√ utilize√ random uniform initialization rang- for tasks that require more information, such as counting ing from − 3 to 3. The learning rate is set as 0.001 (Task 7), the model is still able to obtain the correct answer initially and updated with a learning rate scheduler. E[s] without using all entities. Lastly, transitive logic tasks where contains nodes in the connected components of A containing information is difficult to separate due to dependencies of the answer node a which has relevance scores sampled from entities, the model creates very few banks (1 or 2) and uses a Gaussian distribution centered at 0.75 with a variance of all nodes to correctly generate an answer. We note that in 0.05 (capped at 1). Nodes that are not in the connected the instance where the model only creates one bank, it is component containing a are similarly sampled from a Gaus- very sparse, containing only one or two entities. sian centered from 0.3 with a variance of 0.1 (capped at Because variations in computation times in text are minute, 0). pˆli is initially set to 0.8 and β varies depending on the story length from 0.1 ≤ β ≤ 0.25. Note that for transitive the number of entities required to construct an answer are of more interest as they directly correspond to the num- tasks, pˆli is set to 0.2. We train our models using the Adam optimizer (Kingma & Ba, 2014). ber of computations required. Additionally, due to various implementations of current models, their run times can sig- 5.2. bAbi dataset nificantly vary. However, for the comparison of inference times, AMN’s decoder and EntNet’s decoder are highly sim- The bAbI task suite consists of 20 reasoning tasks that in- ilar and contain roughly the same number of operations. We clude deduction, induction, path finding etc. Results are provide inference wall-clock time performance and memory from the following parameters: ≤ 200 epochs, best of 10 bank behavior in supplementary behavior for representative runs. Table1 shows the accuracy and Table4 shows the tasks. inference performance in terms of the number of entities examined. A task is considered passed if the error rate is 5.3. Extended bAbi tasks less than 5%. We extend the bAbI tasks by adding additional entities and We find that AMN creates 1 − 6 memory banks for different sentences and adding multiple questions for a single story, tasks. We also find that 8 tasks can be solved by looking at for Task 1. just one memory bank and 14 tasks can be solved with half the total number of memory banks. Lastly, all tasks can be Scaled Task 1: We increase the the number of entities to solved by examining less than or equal the total number of 100 entities in the task generation system instead of existing entities (e ∈ M ≤ |V | + )1. Tasks that cannot be solved in . We also extend the story length to to ensure new fewer than half the memory banks either require additional 38 90 entities are referenced. We find that AMN creates memory entities due to transitive logic or have multiple questions. 6 banks and the ratio of entities in the final banks versus the For transitive logic, additional banks could be required as an overall entities drops to given the excess entities that relevant nodes may be in a further bank. However, this still 0.13 are not referenced in the questions. avoids scanning all banks. In the case of multiple questions, all nodes may become necessary to construct all answers. We provide additional evaluation in Appendix to examine Multiple questions: We also augment the tasks with mul- memory bank behavior for certain tasks. tiple questions to understand if AMN can handle when a story has multiple questions associated with it. We extend our model to handle multiple questions at once to limit re- 1 The entities used to construct an answer and pass the task generating the network for every question. To do so, we are examined as the sum of all entities across the M which is usually O(|V |). However, this is within an error margin of 6% modify bAbi to generate several questions per story for tasks more entities on some experiments, and thus we included an that do not currently have multiple questions. For single term. supporting fact (Task 1), the model creates 3 banks and Adaptive Memory Networks Task AMN Entnet DMN+ MemN2N EncDec 1 - Single Supporting Fact 0.0 0.0 0.0 0.0 52.0 2 - Two Supporting Facts 2.1 0.1 0.3 0.3 66.1 3 - Three Supporting Facts 4.7 4.1 1.1 2.1 71.9 4 - Two Arg. Relations 0.0 0.0 0.0 0.0 29.2 5 - Three Arg. Relations 2.7 0.3 0.5 0.8 14.3 6 - Yes/No Questions 3.1 0.2 0.0 0.1 31.0 7 - Counting 0.0 0.0 2.4 2.0 21.8 8 - Lists/Sets 0.0 0.5 0.0 0.9 27.6 9 - Simple Negation 1.3 0.1 0.0 0.3 36.4 10 - Indefinite Knowledge 1.2 0.6 0.0 0.0 36.4 11 - Basic Coreference 2.7 0.3 0.0 0.1 31.7 12 - Conjunction 2.2 0.0 0.0 0.0 35.0 13 - Compound Coref. 4.6 1.3 0.0 0.0 6.80 14 - Time Reasoning 2.1 0.0 0.2 0.1 67.2 15 - Basic Deduction 1.8 0.0 0.0 0.0 62.2 16 - Basic Induction 4.2 0.0 45.3 51.8 54.0 17 - Positional Reasoning 4.3 0.5 4.2 18.6 43.1 18 - Size Reasoning 2.0 0.3 2.1 5.3 6.60 19 - Path Finding 2.4 2.3 0.0 2.3 89.6 20 - Agents Motivations 0.0 0.0 0.0 0.0 2.30 No. of failed tasks (>5%) 0 0 5 6 20
Table 1. Performance comparison of various models in terms of test error rate (%) and the number of failed tasks on the bAbI dataset. The bold task scores are where AMN can solve the task using only 1 memory bank. e∈M Task Created Banks (Rounded Average) Required Banks Ratio ( |V | ) 1 - Single Supporting Fact 3 1 0.22 2 - Two Supporting Facts 5 1 0.41 4 - Two Arg. Relations 2 1 0.70 7 - Counting 5 2 0.81 10 - Indefinite Knowledge 1 1 1.00 11 - Basic Coreference 3 1 0.43 12 - Conjunction 2 1 0.37 14 - Time Reasoning 3 1 0.60 15 - Basic Deduction 1 1 1.00 16 - Basic Induction 2 2 1.06 17 - Positional Reasoning 1 1 1.00 18 - Size Reasoning 3 2 0.82 19 - Path Finding 2 2 1.05 20 - Agents Motivations 2 1 0.26 Extended bAbi 1 - Single Supporting Fact, 100 Entities 6 1 .13 1 - Single Supporting, Multiple Questions 3 1 .38
Table 2. Memory bank analysis of indicative tasks. requires 1 bank to successfully pass the task. Furthermore, AMN presents a new design paradigm in memory networks. the ratio of entities required to pass the task only increases Unlike NTMs, where the network learns where to read/write by 0.16 for a total of 0.38. to fine-grained address information, AMN only learns to write input entities to coarse-grained banks. As a result, 6. Conclusion and Future Work AMN is easier to train (e.g. does not require curriculum learning like NTMs) and does not require a separate spar- In this paper, we present Adaptive Memory Network that sification mechanism like approximate nearest neighbors learns to adaptively organize the memory to answer ques- for inference efficiency. Apart from saving inference times, tions with lower inference times. Unlike NTMs which learn AMN can learn to reason which specific entities contribute to read and write at individual memory locations, Adaptive towards the final answer improving interpretability. Memory Network demonstrates a novel design where the learned memory management is coarse-grained that is eas- References ier to train.Through our experiments, we demonstrate that AMN can learn to reason, construct, and sort memory banks Amazon. Alexa skills kit. In https://developer. based on relevance over the question set. amazon.com/alexa-skills-kit, 2017. Adaptive Memory Networks
Bengio, Emmanuel, Bacon, Pierre-Luc, Pineau, Joelle, Kingma, Diederik and Ba, Jimmy. Adam: A method for and Precup, Doina. Conditional computation in stochastic optimization. arXiv preprint arXiv:1412.6980, neural networks for faster models. arXiv preprint 2014. arXiv:1511.06297, 2015. Le Cun, Yann, Denker, John S, and Solla, Sara A. Optimal Chandar, Sarath, Ahn, Sungjin, Larochelle, Hugo, Vincent, brain damage. In NIPS, 1989. Pascal, Tesauro, Gerald, and Bengio, Yoshua. Hierarchi- cal memory networks. arXiv preprint arXiv:1605.07427, Li, Yujia, Tarlow, Daniel, Brockschmidt, Marc, and Zemel, 2016. Richard. Gated graph sequence neural networks. arXiv preprint arXiv:1511.05493, 2015. Denton, Emily L, Zaremba, Wojciech, Bruna, Joan, LeCun, Yann, and Fergus, Rob. Exploiting linear structure within Liu, Lanlan and Deng, Jia. Dynamic deep neural networks: convolutional networks for efficient evaluation. In Ad- Optimizing accuracy-efficiency trade-offs by selective vances in Neural Information Processing Systems, pp. execution. arXiv preprint arXiv:1701.00299, 2017. 1269–1277, 2014. Paszke, Adam, Gross, Sam, and Chintala, Soumith. Pytorch, Eigen, David, Ranzato, Marc’Aurelio, and Sutskever, Ilya. 2017. Learning factored representations in a deep mixture of Rae, Jack, Hunt, Jonathan J, Danihelka, Ivo, Harley, Timo- experts. arXiv preprint arXiv:1312.4314, 2013. thy, Senior, Andrew W, Wayne, Gregory, Graves, Alex, Figurnov, Michael, Collins, Maxwell D, Zhu, Yukun, Zhang, and Lillicrap, Tim. Scaling memory-augmented neural Li, Huang, Jonathan, Vetrov, Dmitry, and Salakhutdinov, networks with sparse reads and writes. In Advances in Ruslan. Spatially adaptive computation time for residual Neural Information Processing Systems, pp. 3621–3629, networks. arXiv preprint arXiv:1612.02297, 2016. 2016. Giles, C Lee and Omlin, Christian W. Pruning recurrent Rastegari, Mohammad, Ordonez, Vicente, Redmon, Joseph, neural networks for improved generalization performance. and Farhadi, Ali. Xnor-net: Imagenet classification us- IEEE transactions on neural networks, 5(5):848–851, ing binary convolutional neural networks. In European 1994. Conference on Computer Vision, pp. 525–542. Springer, 2016. Goodwin, Travis R and Harabagiu, Sanda M. Medical question answering for clinical decision support. In Pro- See, Abigail, Luong, Minh-Thang, and Manning, Christo- ceedings of the 25th ACM International on Conference on pher D. Compression of neural machine translation mod- Information and Knowledge Management, pp. 297–306. els via pruning. Proceedings of the SIGNLL Conference ACM, 2016. on Computational Natural Language Learning (CoNLL), 2016. Graves, Alex. Adaptive computation time for recurrent neural networks. arXiv preprint arXiv:1603.08983, 2016. Shazeer, Noam, Mirhoseini, Azalia, Maziarz, Krzysztof, Davis, Andy, Le, Quoc, Hinton, Geoffrey, and Dean, Graves, Alex, Wayne, Greg, and Danihelka, Ivo. Neural Jeff. Outrageously large neural networks: The turing machines. arXiv preprint arXiv:1410.5401, 2014. sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538, 2017. Han, Song, Mao, Huizi, and Dally, William J. Deep Com- pression: Compressing Deep Neural Networks with Prun- Sukhbaatar, Sainbayar, szlam, arthur, Weston, Jason, and ing, Trained Quantization and Huffman Coding. In ICLR, Fergus, Rob. End-to-end memory networks. pp. 2440– 2016. 2448, 2015.
Hazy, Thomas E, Frank, Michael J, and OReilly, Randall C. Tulving, Endel et al. Episodic and semantic memory. Orga- Banishing the homunculus: making working memory nization of memory, 1:381–403, 1972. work. Neuroscience, 139(1):105–118, 2006. Vanhoucke, Vincent, Senior, Andrew, and Mao, Mark Z. Henaff, Mikael, Weston, Jason, Szlam, Arthur, Bordes, An- Improving the speed of neural networks on cpus. In Proc. toine, and LeCun, Yann. Tracking the world state with Deep Learning and Unsupervised Feature Learning NIPS recurrent entity networks. Proceedings of the Interna- Workshop, volume 1, pp. 4, 2011. tional Conference on Learning Representations, 2017. Watzl, Sebastian. Structuring mind: The nature of attention Johnson, Daniel D. Learning graphical state transitions. In and how it shapes consciousness. Oxford University ICLR, 2017. Press, 2017. Adaptive Memory Networks
Weston, Jason, Chopra, Sumit, and Bordes, Antoine. Mem- arXiv:1502.05698, 2015. ory networks. arXiv preprint arXiv:1410.3916, 2014. Xiong, Caiming, Merity, Stephen, and Socher, Richard. Weston, Jason, Bordes, Antoine, Chopra, Sumit, Rush, Dynamic memory networks for visual and textual ques- Alexander M, van Merrienboer,¨ Bart, Joulin, Armand, tion answering. In International Conference on Machine and Mikolov, Tomas. Towards AI-complete question an- Learning, pp. 2397–2406, 2016. swering: A set of prerequisite toy tasks. arXiv preprint Adaptive Memory Networks
Method Complexity t =0 t =1 t =2 Entnet ((Henaff et al., 2017)) O(|V |) GGT-NN ((Johnson, 2017)) O(|V | + kS) AMN (ours) O(α|V |)) : 0 < α < 1 +
Table 3. Comparison of decode complexity for AMN, Entnet and GGT-NN.
A A2 A3 A. Appendix Figure 3. Propagation in AMN (shown for a single memory bank A.1. Decode overhead across time). We compare the computations costs during the decode oper- ation during inference for solving the extended bAbi task. We compute the overheads for AMN Entnet (Henaff et al., 2017) and GGT-NN. Table3 gives the decode comparisons between AMN, Entnet and GGT-NN. Here, |V | represents to the total number of entities for all networks. GGT-NN can dynamically create nodes and k k is hyper parameter the new nodes created for S sentences in input story. α is the percent of entities stored in the final bank w.r.t to the total entities for AMN. We compare the wall clock execution times for three tasks within bAbI for 1000 examples/task. We compare the wall- clock times for three tasks. We compare the inference times of considering all banks (and entities) versus the just looking at the passing banks as required by AMN. We find that AMN requires fewer banks and as a consequence fewer entities Figure 4. Heat map showing distribution of entities across various and saves inference times. memory banks for simple bAbI tasks. The x-axis shows the task IDs (refer to Table 1 for task details for each ID.) A.2. Propagation Example tion AMN creates variable number of memory banks. The In this section, we explain propagation with an example. heatmap demonstrates how entities are copied across mem- Figure3 shows how propagation happens after every time ory banks. Grey blocks indicate absence of those banks. step. The nodes represent entities corresponding to words in a sentence. As sentences are processed word by word, a A.4. Algorithm directed graph is drawn progressively from w0...wi...wN . We describe our overall algorithm in pseudo-code in this If sentence lk’s path contains nodes already in the current section in Algorithm1. We follow the notation as described directed graph, lk will include said nodes in the its path. in the paper. After lk is added to A, the model propagates the new update hidden state information ai among all node states using a GRU. ai for each node i is equal to the sum of the incoming edges’ node hidden states. Additionally, we add a particular emphasis on lk to simulate recency. At face value, one propagation step of A will only have a reachability of its immediate neighbor, so to reach all nodes, A is raised to a consecutive power r to reach and update each intermediate node. r can be either the longest path in A or a set parameter.
A.3. Memory bank behavior In this section, we understand memory bank behavior of AMN. Figure4 shows the memory banks and the entity cre- ation for a single story example, for some of the tasks from bAbI. Depending upon the task, and distance from the ques- Adaptive Memory Networks
Task Created Banks Required Banks Baseline (All banks) AMN (Required banks) 1 - Single Supporting Fact 3 1 2.15 s 0.6 s 2 - Two Supporting Facts 5 1 15.8 s 3.2 s 7 - Counting 5 2 21 s 6.0 s
Table 4. Memory bank wall clock times for representative tasks for 1000 examples (time in secs).
Algorithm 1 AMN(S, q, a)
1: M ← Ø 10: mi+1 ← Πmove(mi, nmi )
2: for sentence s ∈ S do 11: nmi ← SGRU(D, nmi ) 3: for word w ∈ s do 12: if i = |M| and Πnew(mi) then 4: D ← ENCODE(w, q) 13: M, p ← [M, mi+1] 5: end for 14: Repeat 8 to 11 once
6: nmi ← SGRU(D) 15: end if 7: for memory bank mi ∈ M do 16: end for 8: mi ← Πau(mi, D) 17: end for 9: mi ← Πprop(mi) 18: aˆ ← DECODE(M, q)