Title: ADAPTIVE MEMORY NETWORKS

Abstract: Real-world Question Answering (QA) tasks often consist of thousands of words that represent many facts and entities. Existing models based on LSTMs require a large number of parameters to support external memory and do not generalize efficiently for long sequence inputs. Memory networks address these limitations by storing information to an external memory module but must examine all inputs in the memory. Hence, for longer sequence inputs, the intermediate memory components proportionally scale in size, resulting in poor inference times. We present Adaptive Memory Networks (AMN) that process input-question pairs to dynamically construct a network architecture optimized for lower inference times. AMN creates multiple memory banks that store entities from the input story to answer the questions. The model learns to reason important entities from the input text based on the question and concentrates these entities within a single memory bank. At inference, one or few banks are used, creating a tradeoff between accuracy and performance. AMN is enabled by first, a novel bank controller that makes discrete decisions with high accuracy and second, the capabilities of a dynamic framework (such as PyTorch) that allow for dynamic network sizing and efficient variable mini-batching. In our results, we demonstrate that our model learns to construct a varying number of memory banks based on task complexity and achieves faster inference times for standard bAbI tasks, and modified bAbI tasks. We solve all bAbI tasks with an average of 48% fewer entities on tasks containing excess, unrelated information.

1. INTRODUCTION:

Question Answering (QA) tasks are gaining significance due to their widespread applicability to recent commercial applications such as chatbots, voice assistants and even medical diagnosis (Goodwin & Harabagiu (2016) ). Furthermore, many existing natural language tasks can also be re-phrased as QA tasks. Providing faster inference times for QA tasks is crucial. Consumer device based question-answer services have hard timeouts for answering questions. For example, Amazon Alexa, a popular QA voice assistant, allows developers to extend the QA capabilities by adding new "Skills" as remote services (Amazon (2017) ). However, these service APIs are wrapped around hard-timeouts of 8 seconds which includes the time to transliterate the question to text on Amazon's servers and the round-trip transfer time of question and the answer from the remote service, and sending the response back to the device. Furthermore, developers are encouraged to provide a list of questions ("utterances") apriori at each processing step to assist QA processing (Amazon (2017) ).

Modeling QA tasks with LSTMs can be computationally expensive which is undesirable especially during inference. Memory networks, a class of deep networks with explicit addressable memory, have recently been used to achieve state of the art results on many QA tasks. Unlike LSTMs, where the number of parameters grows exponentially with the size of memory, memory networks are comparably parameter efficient and can learn over longer input sequences. However, they often require accessing all intermediate memory to answer a question. Furthermore, using focus of attention over the intermediate state using a list of questions does not address this problem. Soft attention based models compute a softmax over all states and hard attention models are not differentiable and can be difficult to train over a large state space. Previous work on improving inference over memory networks has focused on using unsupervised clustering methods to reduce the search space (Chandar et al. (2016) ; Rae et al. (2016) ). Here, the memory importance is not learned and the performance of nearest-neighbor style algorithms is often comparable to a softmax operation over memories.

List of questions:

Strength GRU Entities from the story To provide faster inference for long sequence-based inputs, we present Adaptive Memory Networks (AMN), that constructs a memory network on-the-fly based on the input. Like past approaches to addressing external memory, AMN constructs the memory nodes dynamically. However, distinct from past approaches, AMN constructs a memory architecture with network properties that are decided dynamically based on the input story. Given a list of possible questions, our model computes and stores the entities from the input story in a memory bank. The entities represent the hidden state of each word in the story while a memory bank is a collection of entities that are similar w.r.t the question. As the number of entities grow, our network learns to construct new memory banks and copies entities that are more relevant towards a single bank. Entities may reside in different bank depending on their distance from the question. Hence, by limiting the decoding step to a dynamic number of constructed memory banks, AMN achieves lower inference times. AMN is an end-to-end trained model with dynamic learned parameters for memory bank creation and movement of entities.

Figure 1 demonstrates a simple QA task where AMN constructs two memory banks based on the input. During inference only the entities in the left bank are considered reducing inference times. To realize its goals, AMN introduces a novel bank controller that uses reparameterization trick to make discrete decisions with high accuracy while maintaining differentiability. Finally, AMN also models sentence structures on-the-fly and propagates update information for all entities that allows it to solve all 20 bAbI tasks.

2. RELATED WORK:

Memory Networks: Memory networks store the entire input sequence in memory and perform a softmax over hidden states to update the controller (Weston et al. (2014) ; Sukhbaatar et al. (2015) ). DMN+ connects memory to input tokens and updates them sequentially (Xiong et al. (2016) ). For inputs that consist of large number of tokens or entities, these methods can be expensive during inference. AMN stores entities with tied weights in different memory banks. By controlling the number of memory banks, AMN achieves low inference times with reasonable accuracy. Nearest neighbor methods have also been explored over memory networks. For example, Hierarchical Memory Networks separates the input memory into groups using the MIPS algorithm (Chandar et al. (2016) ) . However, using MIPS is as slow as a softmax operation, so the authors propose using an approximate MIPS that gives inferior performance. In contrast, AMN is end to end differentiable, and reasons which entities are important and constructs a network with dynamic depth.

Neural Turing Machine (NTM) consists of a memory bank and a differentiable controller that learns to read and write to specific locations (Graves et al. (2014) ). In contrast to NTMs, AMN memory bank controller is more coarse grained and the network learns to store entities in memory banks instead of specific locations. AMN uses a discrete bank controller that gives improved performance for bank controller actions over NTM's mechanisms. However, like NTMs, our design is consistent with the modeling studies of working memory by (Hazy et al. (2006) ) where the brain performs robust memory maintenance and may maintain multiple working representations for individual working tasks. Sparse access memory uses approximate nearest neighbors (ANN) to reduce memory usage in NTMs (Rae et al. (2016) ). However, ANNs are not differentiable. AMN, uses a input specific memory organization that does not create sparse structures. This limits access during inference to specific entities reducing inference times.

Graph-based networks, (GG-NNs, Li et al. (2015) and GGT-NNs, Johnson (2017)) use nodes with tied weights that are updated based on gated-graph state updates with shared weights over edges. However, unlike AMN, they require strong supervision over the input and teacher forcing to learn the graph structure. Furthermore, the cost of building and training these models is expensive and if every edge is considered at every time-step the amount of computation grows at the order of O(N 3 ) where N represents the number of nodes/entities. AMN does not use strong supervision but can solve tasks that require transitive logic by modeling sentence walks on the fly. EntNet constructs dynamic networks based on entities with tied weights for each entity (Henaff et al. (2017) ). A key-value update system allows it to update relevant (learned) entities. However, Entnet uses soft-attention during inference to attend to all entities that incur high inference costs. To summarize, majority of the past work on memory networks uses softmax over memory nodes, where each node may represent input or an entity. In contrast, AMN learns to organize memory into various memory banks and performs decode over fewer entities reducing inference times.

Conditional Computation & Efficient Inference: AMN is also related to the work on conditional computation which allows part of networks to be active during inference improving computational efficiency (Bengio et al. (2015) ). Recently, this has been often accomplished using a gated mixture of experts (Eigen et al. (2013) ; Shazeer et al. (2017) ). AMN conditionally attends to entities in initial banks during inference improving performance. For faster inference using CNNs, pruning (Le Cun et al. (1989) ; Han et al. (2016) ), low rank approximations (Denton et al. ( 2014)), quantization and binarization (Rastegari et al. (2016) ) and other tricks to improve GEMM performance (Vanhoucke et al. (2011) ) have been explored. For sequence based inputs, pruning and compression has been explored (Giles & Omlin (1994) ; See et al. (2016) ). However, compression results in irregular sparsity that reduces memory costs but may not reduce computation costs. Adaptive computation time (Graves (2016) ) learns the number of steps required for inferring the output and this can also be used to reduce inference times (Figurnov et al. (2016) ). AMN uses memory networks with dynamic number of banks to reduce computation costs.

Dynamic networks: Dynamic neural networks that change structure during inference have recently been possible due to newer frameworks such as Dynet and PyTorch. Existing work on pruning can be implemented using these frameworks to reduce inference times dynamically like dynamic deep network demonstrates (Liu & Deng (2017) ). AMN utilizes the dynamic architecture abilities to construct an input dependent memory network of variable memory bank depth and the dynamic batching feature to process a variable number of entities. Furthermore, unlike past work that requires an apriori number of fixed memory slots, AMN constructs them on-the-fly based on the input. The learnable discrete decision-making process can be extended to other dynamic networks which often rely on REINFORCE to make such decisions (Liu & Deng (2017) ).

Neuroscience: Our network construction is inspired by work on working memory representations.:

There is sufficient evidence for multiple, working memory representations in the human brain (Hazy et al. (2006) ). Semantic memory (Tulving et al. (1972) ), describes a hierarchical organization starting with relevant facts at the lowest level and progressively more complex and distant concepts at higher levels. AMN constructs entities from the input stories and stores the most relevant entities based on the question in the lowest level memory bank. Progressively higher level memory banks represent distant concepts (and not necessarily higher level concepts for AMN). Other work demonstrates organization of human memory in terms of "priority structure" where attention is a gate-keeper of working memory-guided by executive control's goals, plans, and intentions as in Watzl (2017) , similar in spirit to AMN's question guided network construction.

3. DIFFERENTIABLE ADAPTIVE MEMORY MODULE:

In this section, we describe the design process and motivation of our memory module. Our memory network architecture is created during inference time for every story. The architecture consists of different memory banks and each memory bank stores entities from the input story. Hence, a memory entity represents the hidden state of each entity (each word in our case) from the input story while a memory bank is a collection of entities. Intuitively, each memory bank stores entities that have a similar distance score from the question.

At a high level, entities are gradually and recurrently copied through memory banks to filter out irrelevant nodes such that in the final inference stage, fewer entities are considered by the decoder. Note that the word filter implies a discrete decision and that recurrence implies time. If we were to perform a strict cut off and remove entities that appear to be irrelevant at each time step, learning the reasoning logic that requires previous entities that were cut off would not be possible. Thus, smoothed discretization is required.

We design filtering to be a two-stage pseudo-continuous process to simulate discrete cut offs (Π move , Π new ), while keeping reference history. The overall memory (M ) consists of multiple memory banks. A memory bank is a collection or group of entities (m 0...l ), where m 0 denotes the initial and most general bank and m l denotes the most relevant bank. Note that |l| is input dependent and learned. First, entities are moved from m 0 gradually towards m l based off of their individual relevance to the question and second, if m l becomes too saturated, m l+1 is created. Operations in the external memory allowing for such dynamic restructuring and entity updates are described below.

Note that these operations still maintain end to end differentiability.

1. Memory bank creation (Π new ), which creates a new memory bank depending on the current states of entities m i . If the entropy, or information contained (explained below), of m i is too high, Π new (m i ) will learn to create a new memory bank m i+1 to reduce entropy.

2. Moving entities across banks (Π move ), which determines which entities are relevant to the current question and move such entities to further (higher importance) memory banks.

3. Adding/Updating entities in a bank (Π au ), which adds entities that are not yet encountered to the first memory bank m 0 or if the entity is already in m 0 , the operation updates the entity state.

4. Propagating changes across entities (Π prop ), which updates the entity states in memory banks based on node current states Π prop (M ) and their semantic relationships. This is to communicate transitive logic.

Both Π new , Π move require a discrete decision (refer to section 4.2.1.), and in particular, for Π new we introduce the notion of entropy. That is to say if m i contains too many nodes (the entropy becomes too high), the memory module will learn to create a new bank m i+1 and move nodes to m i+1 to reduce entropy. By creating more memory banks, the model spreads out the concentration of information which in turn better discretizes nodes according to relevance.

4. ADAPTIVE MEMORY NETWORKS:

A high-level overview is shown in Figure 2 , followed by a mathematical detail of the model's modules.

Our model adopts the encoder-decoder framework with an augmented adaptive memory module. For an overview of the algorithm, refer to Section A.1.

Notation and Problem Statement: Given a story represented by N input sentences (or statements), i.e., (l 1 , • • • , l N ), and a question q, our goal is to generate an answer a. Each sentence l is a sequence of N words, denoted as (w 1 , • • • , w Ns ), and a question is a sequence of N q words denoted as (w 1 , • • • , w Nq ). Throughout the model we refer to entities; these can be interpreted as a 3-tuple of e w = (word ID wi, hidden state w, question relevance strength s). Scalars, vectors, matrices, and dot products are denoted by lower-case letters, boldface lower-case letters and boldface capital letters, and angled brackets respectively.

4.1. ENCODER:

The input to the model, starting with the encoder, are story-question input pairs. On a macro level, sentences l 1...N are processed. On a micro level, words w 1...Ns are processed within sentences. For each w i ∈ l i , the encoder maps w i to a hidden representation and a question relevance strength ∈ [0, 1]. The word ID of w i is passed through a standard embedding layer and then encoded through an accumulation GRU. The accumulation GRU captures the entity states through time by adding the output of each GRU time step to its respective word, stored in a lookup matrix. The initial states of e w are set to this GRU output. Meanwhile, the question is also embedded and encoded in the same manner sans accumulation.

In the following, the subscripts i, j are used to iterate through the total number of words in a statement and question respectively, D stores the accumulation GRU output, and w i is a GRU encoding output.

The last output of the GRU will be referred to as w N , w Nq for statements and questions.

ui, uj = EM BED(wii), EM BED(wij

EQUATION

EQUATION

To compute the question relevance strength s ∈ [0, 1] for each word, the model uses GRU-like equations. The node strengths are first initialized to Xavier normal and the inputs are the current word states w in , the question state w Nq , and when applicable, the previous strength. Sentences are processed each time step t.

z t = σ(Uzw in + WzwN q + Xzs t-1 ) (5)

EQUATION

)

EQUATION

EQUATION

In particular, equation ( 6) shows where the model learns to lower the strengths of nodes that are not related the question. First, a dot product between the current word states and question state are computed for similarity (high correlation), then it is subtracted from a 1 to obtain the dissimilarity. We refer to these operations as SGRU (Strength GRU) in Algorithm 1.

4.2. ADAPTIVE MEMORY MODULE:

The adaptive memory module recurrently restructures entities in a question relevant manner so the decoder can then consider fewer entities (namely, the question relevant entities) to generate an answer.

The following operations are performed once per sentence.

4.2.1. MEMORY BANK CONTROLLER:

As mentioned earlier, discrete decisions are difficult for neural networks to learn so we designed a specific memory bank controller Π ctrl for binary decision making. The model takes ideas from the reparameterization trick and uses custom backpropagation to maintain differentiability.

In particular, the adaptive memory module needs to make two discrete decisions on a Here, q (eq 9) is a fully connected layer and the input is the concatenation of all the current memory bank m i 's entity states [w 0 ...w i ] ∈ R 1,n|ew| . Intuitively, q will learn a continuous decision that is later discretized by eq 10 based on entity states and the number of entities. Note this is only performed for the last memory bank.

Π new ([w 0 ...w i ]) = M.new() if 1(Π ctrl ([w 0 ...w i ])) else pass (11)

2. Moving entities through memory banks: Similar to Π new , individual entities' relevance scores are passed into the bank controller to determine H as the input. The relevance score is computed by multiplying an entity state by its respective relevance ∈ R n,|ew| . Here, q has a slight modification and is the identity function. Note that this operation can only be performed if there is a memory bank to move nodes to, namely if m i+1 exists. Additionally, each bank has a set property where it cannot contain duplicate nodes, but the same node can exist in two different memory banks.

Π move (s i * w i ) = m.add(1(Π ctrl (s i * w i ))) ∀i ∈ m (12)

3. Adding/Updating entities in a bank: Recall that entities are initially set to the output of D. However, as additional sentences are processed, new entities and their hidden states are observed. In the case of a new entity e w , the entity is added to the first memory bank m 0 . If the entity already exists in m 0 , then e w 's corresponding hidden state is updated through a GRU. This procedure is done for all memory banks.

EQUATION

4. Propagating updates to related entities: So far, entities exist as a bag of words model and the sentence structure is not maintained. This can make it difficult to solve tasks that require transitive reasoning over multiple entities. To track sentence structure information, we model semantic relationships as a directed graph stored in adjacency matrix A. As sentences are processed word by word, a directed graph is drawn progressively from w 0 ...w i ...w N . If sentence l k 's path contains nodes already in the current directed graph, l k will include said nodes in its path. After l k is added to A, the model propagates the new update hidden state information a i among all node states using a GRU. a i for each node i is equal to the sum of the incoming edges' node hidden states. Additionally, we add a particular emphasis on l k 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. Again, this is done within a memory bank for all memory banks. For entities that have migrated to another bank, the update for these entities is a no-op but propagation information as per the sentence structure is maintained. A single iteration is shown below:

EQUATION

When nodes are transferred across banks, A is still preserved. If intermediate nodes are removed from a path, a transitive closure is drawn if possible.

After these steps are finished at the end of a sentence, namely, the memory unit has reasoned through how large (number of memory banks) the memory should be and which entities are relevant at the current point in the story, all entities are passed through the strength modified GRU (4.1, eq 5-8) to recompute their question relevance (relevance score).

4.3. DECODE:

After all sentences l 1...N are ingested, the decode portion of the network learns to interpret the results from the memory banks. The network iterates through the memory banks using a standard attention mechanism. To force the network to understand the question importance weighting, the model uses an exponential function d to weight important memory banks higher. C m are the hidden states contained in memory m, s m are the relevance strengths of memory bank m, w Nq is the question hidden state, ps is the attention score, r, h are learned weight masks, g are the accumulated states, and l is the final logits prediction. During inference, fewer memory banks are considered.

Cm = sm • [w0, ...wi] ∀i ∈ m (16) ps = Sof tmax( Cm, wN q ) (17) g += d( Cm, ps ) (18) L = r(PReLU(h(g) + wN q ) if m is last (19) 4.4 LOSS

Loss is comprised of two parts, answer loss, which is computed from the given annotations, and secondary loss (from Π new , Π move ), which is computed from sentence and story features at each sentence time step l 0...N . Answer loss is standard cross entropy at the end of the story after l N is processed.

L p ( L) = CrossEntropy( L, L)

After each sentence l i , the node relevance s li is enforced by computing the expected relevance E[s li ]. E[s] is determined by nodes that are connected to the answer node a in a directed graph; words that are connected to a are relevant to a. They are then weighted with a deterministic function of distance from a.

L s (s) = D KL (s li ||E[s li ])

Additionally, bank creation is kept in check by constraining p li w.r.t. the expected number of memory banks. The expected number of memory banks can be thought of as a geometric distribution ∼ Geometric(p li ) parameterized by pli , a hyperparameter. Typically, at each sentence step p is raised to the inverse power of the current sentence step to reflect the amount of information ingested.

Intuitively, this loss ensures banks are created when a memory bank contains too many nodes. On the other hand, the learned mask q (eq. 9) enables the model to weight certain nodes a higher entropy to prompt bank creation. Through these two dependencies, the model is able to simulate bank creation as a function of the number of nodes and the type of nodes in a given memory bank.

L b (p li ) = D KL (p li ||p 1 β|l i | )

All components combined, the final loss is given in the following equation

L total = L p ( L) + |ln| i=1 (L i s (s) + L i b (p))

5. EVALUATION:

In this section, we evaluate AMN accuracy and inference times on the bAbI dataset Weston et al. (2015) and extended bAbI tasks dataset. We compare our performance with Entnet (Henaff et al. (2017) ), which recently achieved state of the art results on the bAbi dataset. For accuracy measurements, we also compare with DMN+ and encoder-decoder methods. Finally we discuss the time trade offs between AMN and current SOTA methods. The portion regarding inference times are not inclusive of story ingestion. We summarize our experiments results as follows:

• We are able to solve all bAbi tasks using AMN. Furthermore, AMN is able to reason important entities and propagate them to the final memory bank allowing for 48% fewer entities examined during inference.

• We construct extended bAbI tasks to evaluate AMN behavior. First, we extend Task 1 for multiple questions in order to gauge performance in a more robust manner. For example, if a reasonable set of questions are asked (where reasonable means that collectively they do not require all entities to answer implying entities can be filtered out), will the model still sufficiently reason through entities. We find that our network is able to reason useful entities for both tasks and store them in the final memory bank. Furthermore, we also scale bAbI for a large number of entities and find that AMN provides additional benefits at scale since only relevant entities are stored in the final memory bank.

5.1. EXPERIMENT SETTINGS:

We implement our network in PyTorch (Paszke et al. (2017) ). We initialize our model using Xavier initialization, and the word embeddings utilize random uniform initialization ranging from -√ 3 to √ 3. The learning rate is set as 0.001 initially and updated with a learning rate scheduler. E[s] contains nodes in the connected components of A containing the answer node a which has relevance scores sampled from a Gaussian distribution centered at 0.75 with a variance of 0.05 (capped at 1). Nodes that are not in the connected component containing a are similarly sampled from a Gaussian centered from 0.3 with a variance of 0.1 (capped at 0). pli is initially set to 0.8 and β varies depending on the story length from 0.1 ≤ β ≤ 0.25. Note that for transitive tasks, pli is set to 0.2. We train our models using the Adam optimizer (Kingma & Ba, 2014) .

5.2. BABI DATASET:

The bAbI task suite consists of 20 reasoning tasks that include deduction, induction, path finding etc. Results are from the following parameters: ≤ 200 epochs, best of 10 runs. Table 1 shows the accuracy and Table 4 shows the inference performance in terms of the number of entities examined. A task is considered passed if the error rate is less than 5%.

We find that AMN creates 1 -6 memory banks for different tasks. We also find that 8 tasks can be solved by looking at just one memory bank and 14 tasks can be solved with half the total number of memory banks. Lastly, all tasks can be solved by examining less than or equal the total number of entities (e ∈ M ≤ |V | + )1 . Tasks that cannot be solved in fewer than half the memory banks either require additional entities due to transitive logic or have multiple questions. For transitive logic, additional banks could be required as an relevant nodes may be in a further bank. However, this still 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 memory bank behavior for certain tasks.

Inference performance Table 4 shows the number of banks created and required to solve a task, as well as the ratio of entities examined to solve the task. Table 3 shows the complexity of AMN and other SOTA models. Entnet uses an empirically selected parameter, typically set to the number of vocabulary words. GGT-NN uses the number of vocabulary words and creates new k new nodes intermittently per sentence step.

For tasks where nodes are easily separable where nodes are clearly irrelevant to the question(s), AMN is able to successfully reduce the number of nodes examined. However for tasks that require more information, such as counting (Task 7), the model is still able to obtain the correct answer without using all entities. Lastly, transitive logic tasks where information is difficult to separate due to dependencies of entities, the model creates very few banks (1 or 2) and uses all nodes to correctly generate an answer. We note that in the instance where the model only creates one bank, it is very sparse, containing only one or two entities.

Because variations in computation times in text are minute, the number of entities required to construct an answer are of more interest as they directly correspond to the number of computations required. Additionally, due to various implementations of current models, their run times can significantly vary. However, for the comparison of inference times, AMN's decoder and EntNet's decoder are highly similar and contain roughly the same number of operations.

5.3. EXTENDED BABI TASKS:

We extend the bAbI tasks by adding additional entities and sentences and adding multiple questions for a single story, for Task 1.

Scaled Task 1: We increase the the number of entities to 100 entities in the task generation system instead of existing 38. We also extend the story length to 90 to ensure new entities are referenced. We find that AMN creates 6 memory banks and the ratio of entities in the final banks versus the overall entities drops to 0.13 given the excess entities that are not referenced in the questions.

Multiple questions: We also augment the tasks with multiple 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-generating the network for every question. To do so, we modify bAbi to generate several questions per story for tasks that do not currently have multiple questions.

For single supporting fact (Task 1), the model creates 3 banks and requires 1 bank to successfully pass the task. Furthermore, the ratio of entities required to pass the task only increases by 0.16 for a total of 0.38.

6. CONCLUSION AND FUTURE WORK:

In this paper, we present Adaptive Memory Network that learns to adaptively organize the memory to answer questions with lower inference times. Unlike NTMs which learn to read and write at individual memory locations, Adaptive Memory Network demonstrates a novel design where the learned memory management is coarse-grained that is easier to train.Through our experiments, we demonstrate that AMN can learn to reason, construct, and sort memory banks based on relevance over the question set.

AMN architecture is generic and can be extended to other types of tasks where the input sequence can be separated into different entities. In the future, we plan to evaluate AMN over such tasks to evaluate AMN generality. We also plan to experiment with larger scale datasets (beyond bAbI, such as a document with question pairs) that have a large number of entities to further explore scalability.

Method Complexity Entnet (Henaff et al. (2017) ) 

O(|V |) GGT-NN (Johnson (2017)) O(|V | + kS) AMN (ours) O(α|V |)) : 0 < α < 1 +

A APPENDIX:

A.1 ALGORITHM We describe our overall algorithm in pseudo-code in this section. We follow the notation as described in the paper.

Algorithm 1 AMN(S, q, a)

1: M ← Ø 2: for sentence s ∈ S do 3:

for word w ∈ s do 4:

D ← ENCODE(w, q)

5:

end for 6:

n mi ← SGRU(D) 7:

for memory bank m i ∈ M do 8: end for 17: end for 18: â ← DECODE(M, q)

m i Π au (m i , D) 9: m i ← Π prop (m i ) 10: m i+1 ← Π move (m i ,

A.2 DECODE OVERHEAD We compare the computations costs during the decode operation during inference for solving the extended bAbi task. We compute the overheads for AMN Entnet (Henaff et al. (2017) ) and GGT-NN. Table 3 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 and saves inference times. In this section, we understand memory bank behavior of AMN. Figure 3 shows the memory banks and the entity creation for a single story example, for some of the tasks from bAbI. Depending upon the task, and distance from the question AMN creates variable number of memory banks. The heatmap demonstrates how entities are copied across memory banks. Grey blocks indicate absence of those banks. x-axis shows the task IDs (refer to Table 1 for task details for each ID.) 

t = 0 t = 0 t = 1 t = 1 t = 2 t = 2 A A A 2 A 2 A 3 A 3

A.4 PROPAGATION EXAMPLE:

In this section, we explain propagation with an example. Figure 4 shows how propagation happens after every time step. The nodes represent entities corresponding to words in a sentence. As sentences are processed word by word, a directed graph is drawn progressively from w 0 ...w i ...w N . If sentence l k 's path contains nodes already in the current directed graph, l k will include said nodes in the its path. After l k is added to A, the model propagates the new update hidden state information a i among all node states using a GRU. a i for each node i is equal to the sum of the incoming edges' node hidden states. Additionally, we add a particular emphasis on l k 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.

:

The entities used to construct an answer and pass the task 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% more entities on some experiments, and thus we included an term.

Figures and Tables:

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).

Figure 2: Adaptive memory networks.

Figure 3: Heat map showing distribution of entities across various memory banks for simple bAbI tasks. The

Figure 4: Propagation in AMN (shown for a single memory bank across time). A.4 PROPAGATION EXAMPLE

ctrl and has two instances: one for memory bank creation Π new and one for moving entities across banks Π move . In equation 9, depending on what Π ctrl is used for, q is a polymorphic function and will take on a different operation and * will be a different input. Examples of such are given in their respective sections (4.2.2.1, 4.2.2.2). Memory bank creation Π new : To determine when a new memory bank is created, in other words, if the current memory bank becomes too saturated, the memory bank controller (4.2.1.) will make a discrete decision to create a new memory bank.

Table content: <table><tr><td>p = q( * )</td><td>(9)</td><td>H = Bernoulli(p)</td><td>(10)</td></tr><tr><td>4.2.2 MEMORY BANK OPERATIONS</td><td/><td/><td/></tr><tr><td>1.</td><td/><td/><td/></tr></table>

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.

Table content: <table><tr><td>Task</td><td colspan="5">AMN Entnet DMN+ MemN2N EncDec</td></tr><tr><td>1 -Single Supporting Fact</td><td>0.0</td><td>0.0</td><td>0.0</td><td>0.0</td><td>52.0</td></tr><tr><td>2 -Two Supporting Facts</td><td>2.1</td><td>0.1</td><td>0.3</td><td>0.3</td><td>66.1</td></tr><tr><td>3 -Three Supporting Facts</td><td>4.7</td><td>4.1</td><td>1.1</td><td>2.1</td><td>71.9</td></tr><tr><td>4 -Two Arg. Relations</td><td>0.0</td><td>0.0</td><td>0.0</td><td>0.0</td><td>29.2</td></tr><tr><td>5 -Three Arg. Relations</td><td>2.7</td><td>0.3</td><td>0.5</td><td>0.8</td><td>14.3</td></tr><tr><td>6 -Yes/No Questions</td><td>3.1</td><td>0.2</td><td>0.0</td><td>0.1</td><td>31.0</td></tr><tr><td>7 -Counting</td><td>0.0</td><td>0.0</td><td>2.4</td><td>2.0</td><td>21.8</td></tr><tr><td>8 -Lists/Sets</td><td>0.0</td><td>0.5</td><td>0.0</td><td>0.9</td><td>27.6</td></tr><tr><td>9 -Simple Negation</td><td>1.3</td><td>0.1</td><td>0.0</td><td>0.3</td><td>36.4</td></tr><tr><td>10 -Indefinite Knowledge</td><td>1.2</td><td>0.6</td><td>0.0</td><td>0.0</td><td>36.4</td></tr><tr><td>11 -Basic Coreference</td><td>2.7</td><td>0.3</td><td>0.0</td><td>0.1</td><td>31.7</td></tr><tr><td>12 -Conjunction</td><td>2.2</td><td>0.0</td><td>0.0</td><td>0.0</td><td>35.0</td></tr><tr><td>13 -Compound Coref.</td><td>4.6</td><td>1.3</td><td>0.0</td><td>0.0</td><td>6.80</td></tr><tr><td>14 -Time Reasoning</td><td>2.1</td><td>0.0</td><td>0.2</td><td>0.1</td><td>67.2</td></tr><tr><td>15 -Basic Deduction</td><td>1.8</td><td>0.0</td><td>0.0</td><td>0.0</td><td>62.2</td></tr><tr><td>16 -Basic Induction</td><td>4.2</td><td>0.0</td><td>45.3</td><td>51.8</td><td>54.0</td></tr><tr><td>17 -Positional Reasoning</td><td>4.3</td><td>0.5</td><td>4.2</td><td>18.6</td><td>43.1</td></tr><tr><td>18 -Size Reasoning</td><td>2.0</td><td>0.3</td><td>2.1</td><td>5.3</td><td>6.60</td></tr><tr><td>19 -Path Finding</td><td>2.4</td><td>2.3</td><td>0.0</td><td>2.3</td><td>89.6</td></tr><tr><td>20 -Agents Motivations</td><td>0.0</td><td>0.0</td><td>0.0</td><td>0.0</td><td>2.30</td></tr><tr><td>No. of failed tasks (&gt;5%)</td><td>0</td><td>0</td><td>5</td><td>6</td><td>20</td></tr><tr><td>Task</td><td colspan="6">Created Banks (Rounded Average) Required Banks Ratio ( e∈M |V | )</td></tr><tr><td>1 -Single Supporting Fact</td><td/><td/><td>3</td><td/><td>1</td><td>0.22</td></tr><tr><td>2 -Two Supporting Facts</td><td/><td/><td>5</td><td/><td>1</td><td>0.41</td></tr><tr><td>4 -Two Arg. Relations</td><td/><td/><td>2</td><td/><td>1</td><td>0.70</td></tr><tr><td>7 -Counting</td><td/><td/><td>5</td><td/><td>2</td><td>0.81</td></tr><tr><td>10 -Indefinite Knowledge</td><td/><td/><td>1</td><td/><td>1</td><td>1.00</td></tr><tr><td>11 -Basic Coreference</td><td/><td/><td>3</td><td/><td>1</td><td>0.43</td></tr><tr><td>12 -Conjunction</td><td/><td/><td>2</td><td/><td>1</td><td>0.37</td></tr><tr><td>14 -Time Reasoning</td><td/><td/><td>3</td><td/><td>1</td><td>0.60</td></tr><tr><td>15 -Basic Deduction</td><td/><td/><td>1</td><td/><td>1</td><td>1.00</td></tr><tr><td>16 -Basic Induction</td><td/><td/><td>2</td><td/><td>2</td><td>1.06</td></tr><tr><td>17 -Positional Reasoning</td><td/><td/><td>1</td><td/><td>1</td><td>1.00</td></tr><tr><td>18 -Size Reasoning</td><td/><td/><td>3</td><td/><td>2</td><td>0.82</td></tr><tr><td>19 -Path Finding</td><td/><td/><td>2</td><td/><td>2</td><td>1.05</td></tr><tr><td>20 -Agents Motivations</td><td/><td/><td>2</td><td/><td>1</td><td>0.26</td></tr><tr><td>Extended bAbi</td><td/><td/><td/><td/><td/></tr><tr><td>1 -Single Supporting Fact, 100 Entities</td><td/><td/><td>6</td><td/><td>1</td><td>.13</td></tr><tr><td>1 -Single Supporting, Multiple Questions</td><td/><td/><td>3</td><td/><td>1</td><td>.38</td></tr></table>

Memory bank analysis of indicative tasks.

Table content: <table/>

Comparison of decode complexity for AMN, Entnet and GGT-NN.

Table content: <table/>

n mi )

Table content: <table><tr><td>11:</td><td/></tr><tr><td>14:</td><td>Repeat 8 to 11 once</td></tr><tr><td>15:</td><td>end if</td></tr><tr><td>16:</td><td/></tr></table>

Memory bank wall clock times for representative tasks for 1000 examples (time in secs).

Table content: <table><tr><td>Task</td><td colspan="4">Created Banks Required Banks Baseline (All banks) AMN (Required banks)</td></tr><tr><td>1 -Single Supporting Fact</td><td>3</td><td>1</td><td>2.15 s</td><td>0.6 s</td></tr><tr><td>2 -Two Supporting Facts</td><td>5</td><td>1</td><td>15.8 s</td><td>3.2 s</td></tr><tr><td>7 -Counting</td><td>5</td><td>2</td><td>21 s</td><td>6.0 s</td></tr><tr><td colspan="2">A.3 MEMORY BANK BEHAVIOR</td><td/><td/><td/></tr></table>

Peer review (iclr2018):
 Summary: This paper proposes a dynamic memory augmented neural network for question answering. The proposed model iteratively creates a shorter list of relevant entities such that the decoder can look at only a smaller set of entities to answer the given question. Authors show results in bAbi dataset. My comments: 1. While the proposed model is very interesting, I disagree with the claim that AMN has lower inference times. The memory creation happens only after reading the question and hence the entire process can be considered as part of inference. So it is not clear if there is a huge reduction in the inference time when compared to other models that the authors compare. However, the proposed model looks like a nice piece of interpretable reasoning module. In that sense, it is not any better than EntNet based on the error rate since EntNet is doing better than AMN in 15 out of 20 tasks. So it is not very clear what is the advantage of AMN over EntNet or other MANN architectures. 2. Can you explain equation 9 in detail? What is the input to the softmax function? What is the output size of the softmax? I assume q produces a scalar output. But what is the input size to the q function? 3. In the experiment, when you say “best of 10 runs”, is it based on a separate validation set? Please report the mean and variance of the 10 runs. It is sad that people just report best of multiple runs in the bAbi tasks and not report the variance in the performance. I would like to see the mean and variance in the performance. 4. What happens when number of entities is large? Can you comment about how this model will be useful in situations other than reading comprehension style QA? 5. Are the authors willing to release the code for reproducing the results? Minor comments: 1. Page 2, second line: “Networks(AMN)” should be “Networks (AMN). 2. In page 3, first line: “Hazy et al. (2006)” should be “(Hazy et al. 2006)”. 3. In page 3, second para, first line, both references should use .citep instead of .citet. 4. In page 4, fourth para, Vanhoucke et al should also be inside .citep. 5. In page 4, notations paragraph: “a question is a sequence of N_q words” - “of” is missing. 6. In page 5, first paragraph is not clear. 7. In page 6, point 4, 7th line: “nodes in the its path” should be “nodes in its path”. 8. In page 9, section 5.3, multiple questions, 2nd line: “We extend the our model” should be “We extend our model”.