MEMO: A Deep Network for Flexible Combination of Episodic Memories
Authors: Andrea Banino
Adrià Puigdomènech Badia
Raphael Köster
Martin J. Chadwick
Vinicius Zambaldi
Demis Hassabis
Caswell Barry
Matthew Botvinick
Dharshan Kumaran
Charles Blundell
Venue:  (2020)
Keywords: Memory Augmented Neural Networks
Deep Learning

Abstract: Recent research developing neural network architectures with external memory have often used the benchmark bAbI question and answering dataset which provides a challenging number of tasks requiring reasoning. Here we employed a classic associative inference task from the human neuroscience literature in order to more carefully probe the reasoning capacity of existing memory-augmented architectures. This task is thought to capture the essence of reasoning -- the appreciation of distant relationships among elements distributed across multiple facts or memories. Surprisingly, we found that current architectures struggle to reason over long distance associations. Similar results were obtained on a more complex task involving finding the shortest path between nodes in a path. We therefore developed a novel architecture, MEMO, endowed with the capacity to reason over longer distances. This was accomplished with the addition of two novel components. First, it introduces a separation between memories/facts stored in external memory and the items that comprise these facts in external memory. Second, it makes use of an adaptive retrieval mechanism, allowing a variable number of ‘memory hops’ before the answer is produced. MEMO is capable of solving our novel reasoning tasks, as well as all 20 tasks in bAbI.

Published as a conference paper at ICLR 2020
MEMO: A DEEP NETWORK FOR FLEXIBLE COMBINA- TION OF EPISODIC MEMORIES
Andrea Banino 12,
Adrià Puigdomènech Badia
Raphael Köster'
Martin J. Chadwick
Vinicius Zambaldi!
Demis Hassabis
Caswell Barry
Matthew Botvinick
Dharshan Kumaran
Charles Blundell
ABSTRACT
Recent research developing neural network architectures with external memory have often used the benchmark bAbl question and answering dataset which provides a challenging number of tasks requiring reasoning. Here we employed classic associative inference task from the memory based reasoning neuroscience litera- ture in order to more carefully probe the reasoning capacity of existing memory- augmented architectures. This task is thought to capture the essence of reasoning the appreciation of distant relationships among elements distributed across multiple facts memories. Surprisingly, found that current architectures struggle reason over long distance associations. Similar results were obtained on more complex task involving finding the shortest path between nodes in path. We therefore developed MEMO, an architecture endowed with the capacity to reason over longer distances. This was accomplished with the addition of two novel com- ponents. First, it introduces separation between memories/facts stored in external memory and the items that comprise these facts in external memory. Second, it makes use of an adaptive retrieval mechanism, allowing variable number of 'memory hops' before the answer is produced. MEMO is capable of solving our novel reasoning tasks, well match state of the art results in bAbl.
INTRODUCTION
During our every day life we need to make several judgments that require connecting facts which were not experienced together, but acquired across experiences at different points in time. For instance, imagine walking your daughter to coding summer camp and encountering another little girl with a woman. You might conclude that the woman the mother of the little girl. Few weeks later, you are at a coffee shop near your house and you see the same little girl, this time with a man. Based on these two separated episodes you might infer that there This flexible recombination of single experiences in novel ways infer unobserved relationships is called inferential reasoning and is supported by the hippocampus (Zeithamova et al. 2012).
relationship between the woman and the man.
Interestingly, it has been shown that the hippocampus is storing memories independently of each other through a process called pattern separation Yassaa & Stark 2011 Marr et al. 1991). The reason hippocampal memories are kept separated is minimize interference between experiences. which allows us to recall specific events in the form of "episodic" memories (Eichenbaum & Cohen 2004 'Squire et al., 2004). Clearly, this separation is in conflict with the above mentioned role of the hippocampus in generalisation - i.c. how can separated memories be chained together? Interestingly, recent line of research (Kumaran & McClelland 2012 [anino al., 2016 Schapiro ct al. 2017 Koster et al. 2018) sheds lights on this tension by showing that the integration of separated experiences emerges at the point of retrieval through a recurrent mechanism. This allows
DeepMind, London, UK. COMPLEX, University College London, London, UK. Department of Cell and Developmental Biology, University College London, London, UK. Correspondece should he sent to abanino google.com, adrian@ google.oom chlundell@ google.com These authors contributed equally.

---

Published as a conference paper at ICLR 2020
multiple pattern separated codes interact, and therefore support inference. In this paper we rely on these findings to investigate how we can take inspiration from neuroscience models to investigate and enhance inferential reasoning in neural networks.
Neural networks augmented with external memory, like the Differential Neural Computer (Grves et al., 2016] DNC), and end to end memory networks (Sukhbaatar et al. 2015 EMN) have shown remarkable abilities to tackle difficult computational and reasoning tasks. Also, more powerful attention mechanisms (Vaswani et al. 2017: Dehghani et al.) 2018) or the use of context (Seo et al., 2016) have recently allowed traditional neural networks to tackle the same set of tasks. However, some of these tasks- - e.g. bAbI (Weston et al. 2015 present repetitions and commonalities between the train and the test set that neural networks can exploit to come up with degenerate solutions. To overcome this limitation we introduced new task, called Paired Associative Inference (PAI see below), which is derived from the neuroscientific literature (Bunsey & Hichenbaum, 1996 [Banino] jet al. 2016). This task is meant to capture the essence of inferential reasoning i.e. the appreciation of distant relationships among elements distributed across multiple facts or memories. PAI is fully procedurally generated and so is designed to force neural networks to learn abstractions to solve previously unseen associations.
We then use the PAI task, followed by a task involving finding the shortest path and finally bAbi to investigate what kind of memory representations effectively support memory based reasoning. The EMN and other similar models (Sukhbaatar et al., 2015 [Santor et al. 2017 Pavez et al.) (018)) have used fixed memory representations based on combining word embeddings with a positional encoding transformation. A similar approach has been recently implemented by current state of the art language model /Vaswani al. 2017 Devlin et al. 2018). By contrast our approach, called MEMO, retains the full set of facts into memory, and then learns a linear projection paired with a powerful recurrent attention mechanism that enable greater flexibility in the use of these memories. MEMO is based on the same basic structure of the external memory presented in EMN (Sukhbaatar et al. 2015), but its new architectural components can potentially allow for flexible weighting of individual elements in memory and so supporting the form of the inferential reasoning outlined above.
Next, we tackle the problem of prohibitive computation time. In standard neural networks, the computation grows as function of the size of the input, instead of the complexity of the problem being learnt. Sometimes the input is padded with a fixed number of extra values to provide greater computation (Graves al. 2016). in other cases, input values are systematically dropped to reduce the amount of computation (e.g., frame dropping in reinforcement learning (Mnih et al.] 2016)).
Critically, these values normally hand tuned by the experimenter; instead, here we are interested in adapting the amount of compute time to the complexity of the task. To do so we drawn inspiration from model of human associative memory called REMERGE (Kumaran & McClelland 2012)) In this model, the content retrieved from memory is recirculated back as the new query, then the difference between the content retrieved at different time steps in the re-circulation process used to calculate if the network has settled into fixed point, and if so this process terminates.
To implement this principle in a neural network, we were inspired by techniques such as adaptive computation time (Graves (016). In our architecture, the network outputs an action (in the reinforce- ment learning sense) that indicates whether it wishes to continue computing and querying its memory, or whether it is able to answer the given task. We call this the halting policy as the network learns the termination criteria of a fixed point operator. Like ACT, the network outputs a probability of halting, but unlike ACT, the binary halting random variable is trained using REINFORCE (Williams) 1992). Thus we use reinforcement learning to adjust weights based upon the counterfactual problem: what would be the optimal number of steps of computation, given particular number of steps was taken this time? The use of REINFORCE to perform variable amount of computation has been investigated already (e.g. (Shen et al. 2017: L017:oouizos an extra term to the REINFORCE loss that, by exploiting the mathematical properties of binary random variables, naturally minimizes the expected number of computation steps. Thus we directly encourage our network to explicitly prefer representations and computation that minimize the amount of required computation.
Louizos et al.. 2017) however our approach differs in that we added
To sum up, our contributions are:
1. A new task that stresses the essence of reasoning - i.. the appreciation of distant relation- ships among elements distributed across multiple facts.
2

---

Published as a conference paper at ICLR 2020
2. An in depth investigation of the memory representation that support inferential reasoning, and extensions to existing memory architectures that show promising results on these reasoning tasks.
3. A REINFORCE loss component that learn the optimal number of iterations required to learn to solve task.
4. Significant empirical results on three tasks demonstrating the effectiveness of the above two contributions: paired associative inference, shortest path finding, and bAbI (Weston et al. 2013
METHODS
2.1 RECAPITULATING END-TO-END MEMORY NETWORKS
We begin by describing End-to-End Memory Networks (Sukhbaatar et al. 2015 EMN), as a reminder of this architecture, to introduce notation and nomenclature, and also contrast work. We will focus on the multilayer, tied weight variant of EMN as this most closely resembles our architecture.
The set up used for the rest of the paper is as follows: given a set of knowledge inputs (zi)/=1 =
(?), 92. {s} E IR., the network must predict the answer I represents the length of the knowledge input sequence, and S is the length of each input sentence; for instance, in bAbl (Weston et al.] (015). 7 will be the number of stories and S is the number of words in each sentence in a story, Tis will be the word in position s in the sentence, in the ith story and will be O-dimensional one hot vector encoding one of O possible input words.
(an, T12, E12,..., F | and query or question
EMN embeds each word and sums the resulting vectors:
h; Σ le WA
(1)
U; Σ Wotii
(2)
Qo Wa4
(3)
where WA, W, E Rdxo Wg E R are embedding matrices for the key, values and query. respectively. Also le is a positional encoding column vector (as defined in (Sukhbaatar et al. 2015)), represent element wise multiplication and O is the size of the vocabulary.
At each step f, EMN calculates the vector of weights over the memory elements k; and produces the output. Let K be the I x d matrix formed by taking k; as its rows, and similarly V formed by as rows, then:
W; (oft)
(4)
4++1 V + W It
(5)
a, = Noftmax(Waage
(6)
where W; / are weights over the memory slots, Wie, Wa [gdxd is a linear mapping relating the query at the previous step to the current one, 4++1 is the query to be used at the next step, and a is the answer (usually only produced right at the end). EMN is trained via cross entropy loss on a, at
the final step.
2.2 MEMO
MEMO embeds the input differently. First, a common embedding C, of size x des is derived for each input matrix x; E -
W.
where W. E xxde Then each of these embeddings is adapted to either be a key or a value. However, contrary EMN, we do not use hand-coded positional embeddings, instead the words in
3

---

(N) Keeping learn how weight each of these items when Published as a conference paper at ICLR 2020
each sentence and their one-hot encoding in 24, embedded Ct, combined and then this vector is passed trough a linear projection followed by an attention mechanism (explained in detail below). This allows G; capture flexibly any part of the input sentence in z;- MEMO uses multiple heads to attend to the memory following (Vaswani et al. 2017). Each head has a different view of the same common inputs G₁. Let H denote the total number of heads, and h index the particular head, then for each h E {l,.... H} } we have:
{,10) = W (b)
veo(c;
(8)
(h) W!
(9)
40
(h) W
(10)
where (((h), w((h) E gaxSde and w( R are embedding matrices for the key, values and query respectively. vec(c) means flattening the matrix e into a vector with the same number of elements the original matrix, and ve) is the reverse operation of a vector into matrix such that vec-I (vec()))) = c. The result is three d-dimensional vectors (.(h), (P) and go each item separated into memory allow perform a memory lookup. This contrasts with the hand-coded positional embeddings used in EMN (Sukhbaatar al., 2013) and updated recently in Vaswani et al. (2017). and proved critical for enabling the flexible recombination of the stored items.
The attention mechanism used by MEMO also differs from that shown above for EMN, Firstly, it is adapted to use multi-head attention. Secondly, we use DropOut (Srivastava et al. 2014) and LayerNorm (Ba al. 2016 to improve generalisation and learning dynamics. Let K
R
denote the matrix formed by taking each ((()) as a row, and (b) be the matrix formed by taking each "i
(h) row. In contrast, let Q: H be the matrix formed by taking each g the rows. The
attention mechanism then becomes:
Vd WnK (1) (h)
(11)
w (ropOuussoftaa
(h)
(12)
(h) (h) v(h)
4+1
(13)
Q+11 LayerNorm (vec- (vec-\(W, (W,vec(Q(+1))) Q.)
(14)
0r softmax (W (roDtrOOu vec(Qr+1))))
(15)
where W, E RRIXI W, E R Hd Hd are matrices for transforming the logits and queries respectively. W. E Oxxd and W R.d. xd.xHd are the matrices of the output MLP that produces the answer a;- It is worth noting, that even though our attention mechanism uses some of the feature implemented in Vaswani et al. (2017) i.e. normalization factor sqr((d) and multiheading, differs from because rather than doing self-attention it preserves the query separated from the keys and the values. This aspect is particularly important in terms of computational complexity, in that MEMO is lincar with respect to the number of sentences of the input, whereas methods relying on self-attention (e.g. Dehghani et al, 2018) have quadratic complexity (see Appendix E).
2.3 THE HALTING POLICY
In the previous sections we described how MEMO can output sequence of potential answers to an input query, here we describe how to learn the number of computational steps hops required to effectively answer a. To make this decision, we collect some information at every step and we use this to create an observation 8t. That observation is then processed by gated recurrent units (GRUs) (Chung et al. 2015) followed by a MLP which defines a binary policy π (a) (s,, 0) and approximate its value function V (st, 0). The input s; to this network is formed by the Bhattacharyya distance (Bhattacharyya 1946) between the attention weights of the current time steps W2 and the ones at the previous time step. W (both W, and W taken after the softmax), and the number of steps taken so far as a one-hot vector t. The idea behind the way we build s, is that if the attention is focused on the same slot of memory for too many consecutive steps than there is reason to keep
4

---

Published as a conference paper at ICLR 2020
querying the memory because the information retrieved will be the same ie. the network has settled into a fixed point.
2, = GRUR( d(W, t))
(16)
Vt W2 MTPp((=)) h1 = (
(17)
(18)
This network is trained using REINFORCE (Williams 1992). More concretely, the parameters
are adjusted using n-step look ahead values, Rg = -E==o Y're+i + 7" (S 0), where y the discount factor. The objective function of this network is minimize: Hep--Net
and Lhop -Beroa (-(s,, 0)]. Interestingly, Loop is term that directly follows from the fact that is a binary policy. Specifically, the expectation of binary random variable is its probability and the expectation of their sum is the sum of the expectation. Consequently, the new term that introduce in the loss, CHop allows us to directly minimize the expected number of hops. This term, similar in motivation to (Louizos al. 2017) (although differs mathematically), directly encourages prefer representations and computation that minimise the amount of required computation. is also worth noting that an argument against using REINFORCE when training discrete random variables is that the variance can be prohibitively high (Sutton & Barto 2018). Interestingly. in the case of a binary halting random variable, the variance is just p p) where p the probability of halting and the variance is bounded by 1/4 which we find is not too large for learning to proceed successfully in practice.
Lar afy + _CHop where: Lx -BE [ fv
(R=-V(ss,))
network to explicitly
Finally, the reward structure is defined by the answer
f1, if = {0, otherwise
where a is the target answer associate with the imput and is the prediction from the network. The final layer of MLPR was initialized with biasyitt in order to increase the chances that π produces probability of (i.e. do more hop). Finally, we maximum number of hops, N, that the network could take. If N was reached, the network stopped performing additional hops. Critically, there gradient sharing between the hop network and the main MEMO network explained above. All model hyperparameters are reported in appendix D.
RELATED WORK
3.1 MEMORY-AUGMENTE NEURAL NETWORKS
In recent years there has been increasing interest in the development of memory-augmented networks, due to their potential for solving abstract and relational reasoning tasks. Alongside the EMN (described in detail above), another influential early model deploying memory-auggenntee the Differential Neural Computer (Graves et al., 2014 2016 DNC). The DNC operates sequentially on the inputs, and learns to read and write memory store. The model proved capable of solving a range of algorithmic problems, but was difficult to scale to higher-dimensional problem domains. more recent extension incorporated sparsity into the DNC(Rae et al., al., (016)) allowing the model to perform well at larger-scale tasks, including the bAbl task suite (Weston et al. 2015). Since these initial models were published. number of alternative
was
architectures have been developed (Kumar et al. 2016 Henaff et al., 2016: Pavez et al. 2018). The Dynamic Memory Network (Kumar et al. 2016) shares some similarities with EMNs, but operates on sequential inputs rather than parallel. The Recurrent Entity Network (Henaff et al. 2016) has similarities with the DNC, but uses a parallel architecture, enabling simultaneous updates across several memory locations. The Working Memory Network (Pavez et al. 2018) is closely based on EMNs, but additionally incorporates a working memory buffer, and a RelationNet (Santoro et al. 2017). These components
eeemory-augmeeee
thereby enable relational reasoning over the retrieved contents of memory. Each of these new models has proven to perform well at various reference tasks, including the bAbl task suite (Weston et al. 2015).
5

---

Published as a conference paper at ICLR 2020
3.2 ADAPTIVE COMPUTATION TIME
In general, the time required to solve a problem is expected to increase with the complexity of the task. However, most of the machine learning algorithms do not adapt their computational budget based on the complexity of the task. One approach to this problem is represented by Adaptive Computation Time (ACT) (Graves) 2016). ACT is mechanism for learning scalar halting probability, called the "ponder time". dynamically modulate the number of computational steps needed for each input. An alternative approach is represented by Adaptive Early Exit Networks (Bolukbasi et al., 2017). which give the network the ability to exit prematurely i.e. not computing the whole hierarchy of layers if no more computation is needed. Another approach to conditional computation is the use of REINFORCE (Williams [992) to learn discrete latent variables which dynamically adjust the number of computation steps. This has been applied to recurrent neural networks where each layer decides whether or not to activate the next one (Chung et al. 2016). REINFORCE has also been used to Icarn how many steps to "jump" in sequence, and so reducing the total number processed inputs YYu et al.) (017). This jump technique has also been applied recurrent neural network without the need of REINFORCE (Campos Camunez et al. 2018 A similar idea has also been applied to neural network augmented with external memeory (Shen al., (017). but it differs from ours both in the REINFORCE loss and in the fact that our method introduce the idea of using the distance between attention weights at different time steps (hop) as a proxy to check more information can be retrieved from memory i.e. another hop is needed or whether the network has settled and it has enough information to correctly answer the query.
3.3 GRAPH NEURAL NETWORKS
Graph Neural Networks (GNNs) (Scarselli et al., 2008 Gori et al. 2005) consist of an iterative
message passing process which propagates node and edge embeddings throughout a graph. Neural networks are then used to aggregate or learn functions over graph components to perform supervised. sumi-supervised, representation and reinforcement learning tasks (Kipf Welling 2016 Gilmer et al. 2017 Hamilton et al.) 2017 Zambaldi et al.) 2018). The message passing process implements
similar computation to attention mechanisms (Veličković et al., 2017: Battaglia et al. 2018) and as limit case, self-attention can be viewed as fully-connected GNN. Our work differs from GNNs in two fundamental ways: even though GNNs may exhibit a recurrent component (Li et al., 2015] 2018), its implementation is based in unrolling the recurrence for fixed number of steps and use backpropagation through time in the learning process, whereas our method performs an adaptive computation modulate the number of message passing steps; model does not require message passing between memories-input queries attend directly to the memory slots.
4 PAIRED ASSOCIATIVE INFERENCE TASK
B1
C2
C1
A/
AJBa
B
By
B:Cr
B2
C1
C2
Figure 1: Paired associative inference. The panel on the left illustrates a memory store filled with random pairs of images. The panels to the right illustrate (from left to right) two 'direct' queries (AB and BC) where no inference is require, and an 'indirect' query (AC) where inference is required
One contribution of this paper is to introduce task, derived from neuroscience, to carefully probe the reasoning capacity of neural networks. This task is thought to capture the essence of reasoning - the appreciation of distant relationships among elements distributed across multiple facts memories. This process is formalized in a prototypical task widely used to study the role of the hippocampus in generalization the paired associative inference (PAT) task (Bunsey & Fichenbaum. 1996. Banino

---

Published as a conference paper at ICLR 2020
et al. 2016 see Fig. 1). Here, two images randomly associated together. For example, analogous to seeing a little girl with a woman as in the example in the main text, in the PAI task the agent (human participant
artificial neural network) would be presented with an image A, e.g. woman, and image B, e.g. a girl, side by side. Later, in a separate event, the agent would be exposed to second pair, the image B again, but this time paired with new image C. e.g. the other person. This is analogous to seeing the little girl a second time with a different person. During test time two types of query can be asked: direct and indirect queries. Direct queries are test of episodic memory as the answer relies on retrieving an episode that was experienced. In contrast, indirect queries require inference across multiple episodes (see appendix A.1 for further details). Here the network is presented with CUE, image A, and two possible choices: image C, the MATCH, that was originally paired with B; or another image C', the LURE, which was paired with B i.e. forming different triplet A' - B - C'. The right answer, C, can only be produced by appreciating that A and C are linked because they both were paired with B. This is analogous to the insight that the two people walking the same little girl are likely to have some form of association. For the specific details on how the batch was created please refer to appendix A.
RESULTS
5.1 BASELINES
We compared MEMO with two other memory-augmented architectures: End to End Memory Networks (EMN) (Sukhbaatar et al. 2015) and DNC (Graves et al. (2016). We also compared MEMO
with the Universal Transformer (UT) (Dehghani et al. 2018), the current state of the art model the
bAbl task suite (Weston et al. 2015 See the appendix for details about the baselines.
5.2 PAIRED ASSOCIATIVE INFERENCE
Table Treports the summary of results of our model (MEMO) and the other baselines on the hardest inference query for each of the PAI tasks (full results are reported in the Appendix A.2). On the smaller set i.e. A-B-C MEMO was able to achieve the highest accuracy together with DNC, whereas EMN, even with 4 or 10 hops, wasn't able achieve the same level of accuracy, also UT was not able to accurately solve this inference test. For longer sequences -i.e. length and MEMO was the only architecture which successfully answered the most complex inference queries.
To investigate how these results were achieved, we run further analysis on the length 3 PAI task. Interestingly, to solve this task DNC required 10 pondering steps to get the same accuracy as MEMO, which instead converged to 3 hops (see Fig. 3b in Appendix A.3). To analyze how MEMO approached this task we then analyzed the attention weights of an inference query. where the goal was to associate CUE with the MATCH and avoid the interference of the LURE (see Appendix A. for task details). For clarity we report here the original sequence A B C. respectively composed by the following class IDs: 611 191- 840 (however this sequence was not directly experienced together by the network, as the two associations A - B and B - C' where stored in slot 10 and 25, respectively). As depicted in Figure 2, in the first hop MEMO retrieved the memory in slot 10, which contained the CUE, ID 611, and the associated item, ID 191, which form an A B association. Then in the following hop this slot partially active, but most of the mass was placed on slot 16, which contained the memory association B - C: that is, ID 191 and ID 840, the MATCH. Interestingly, slot 13 which was associated with the LURE, ID 943, got bit of mass associated with it. So in this second hop MEMO assigned appropriate probability masses to all the slots needed to support a correct inference decision, which was then confirmed in the last hop. This sequence of memories activation is reminiscent of the one predicted by computational model of the hippocampus (Kumaran & McClelland, 2012; Schapiro al.) 2017) and observed in neural data (Koster al. 2018). Moreover, another instance of MEMO which used 7 hops, to solve the task with the same level of accuracy, presented a very different pattern of memories activation (see Fig. Appendix A.3.2). This is indication of the fact that the the algorithm used to solve the inference problem depends on how many hops the network takes. This aspect could also be related to knowledge distillation in neural networks (Hinton et al., 20155] Frankle & Carbin 2018), whereby many hops are used to initially solved the task (i.c. pwe-peraaettiittiion and then these are automatically reduced to use less computation (see Fig. 3b in Appenix A.3).
7

---

Published as a conference paper at ICLR 2020
We also ran set of ablation experiments MEMO (see Table in Appendix A.3.1). This analysis confirmed that it is the combination of the specific memory representations (i.e. facts kept separated) and the recurrent attention mechanism that supports successful inference i.c. employing these two components individually was not enough. Interestingly, this conclusion was valid only for inference queries, not for direct queries (see Fig. 3c,d in Appendix A.3). Indeed, by definition, direct queries are a pure test of episodic memory and so can be solved with single memory look-up. Finally, we also compared our adaptive computation mechanism with ACT (Graves) 2016) and we found that, for this task, our method was more data efficient (see Fig. 5 in Appendix A.3.3).
Table 1: Inference queries
Length
EMN
DNC
UT
MEMO
3 items (set: A-B-C accuracy A-C)
61.01
96.85
85.60
98.26(0.67)
4 items (set: A-B-C-D accuracy on A-D)
48.66
51.56
44.16
97.22(0.13)
5 items (set: A-B-C-D-E- accuracy on A-E)
45.13
62.61
47.93
84.54(5.72)
Test results for the best 5 hyper- parameters (chosen by validation loss) for MEMO. For EMN and DNC and UT results on the best run (chosen by validation loss)
QUERY
H 0
HOP
2OP
MATCH LURE
PROB 2
PROB
MEMORY
179
304
SLOT ID
212
Figure 2: Weights analysis of an inference query in the length 3 PAI task. An example of memory content and related inference query is reported in the first column on the left. For clarity we report image class ID. Cue and Match are images from the same sequence e.g. A1u Clu, where 10 is the slot ID. The lure is an images presented in the same memory store, but associated with a different sequence. e.g. C13. The 3 most right columns report the weights associated with the 3 hops used by the network, for each probability mass we report the associated retrieved slot.
5.3 SHORTEST PATH ON RANDOMLY GENERATED GRAPHS
We then turn a of synthetic reasoning experiments on randomly generated graphs (see Appendix B.I for details). Table 2 shows the accuracy of the models on the task related to finding the shortest path between two nodes. On a small graph with 10 nodes, with a path length of 2 and 2 outgoing edges per node. DNC, the Universal Transformer, and MEMO had perfect accuracy in predicting the intermediate shortest path node. However, on more complicated graphs (20 nodes, separated outgoing edges), with a path length of 3, MEMO outperformed EMN in predicting the first node of the path (by more than 50%), and, similarly to DNC, almost completely solved the task. Additionally, MEMO outperformed DNC in more complicated graphs with a high degree of connectivity (5 out- degree), being better by more than 20% at predicting both nodes in the shortest path. This showed the great scalability that MEMO holds; the model was able to iterate and considered more paths as the number of hops increases. Finally. Universal Transformer had a different performance in predicting the first node versus the second one of the shortest path. In the latter case, the results showed that
8

---

Published as a conference paper at ICLR 2020
UT achieved slightly lower performance than MEMO, showing its great capability of computing operations that require direct reasoning.
Table 2: Undirected graph shortest path
Prediction of
Prediction of
Gruph Structure
First Node
Second Node
Nodes
Out-degree
Path length
EMN
UT
DNC
MEMO
EMN
DNC
MEMO
94,99
100.00
100.00
100.00(0.00)
n/a
1/a
n/a
20
31.99
39.00
97.00
94.40(0.002)
66.00
84.80
98.00
93.00(0.03)
20
23.99
28.00
30.00
69.20(0.07)
43.00
61.51
40.99
68.80(0.09)
Test results for the best 5 paper-parameters (chosen by training loss) for MEMO, mean and corresponding standard deviation are reported. For EMN, UT, and DNC we report results from the best run.
5.4 QUESTION ANSWERING ON THE BABI TASKS
Finally, we turn our attention to the bAbl question answering dataset (Weston et al. 2015). which consists of 20 different tasks. In particular we trained our model on the joint 10k training set (specifics of training are reported in the appendix C.I).
Table reports the averaged accuracy of our model (MEMO) and the other baselines bAbl (the accuracy for each single task averaged across the best set of paper-parameters is reported in Table
10 in the Appendix C.3). In the 10k training regime MEMO was able to solve all tasks, thereby matching the number tasks solved by (Seo et al. 2016 Dehghani et al. 2018), but with lower
error (for single task results refer to Appendix C.3).
We also ran an extensive set of ablation experiments to understand the contribution of each architec- tural component the final performance (see Appendix C.2 for details). As observed previously was the combination of memory representations and the powerful recurrent attention that was critical to achieve state of the art performance on bAbl. Finally the use of layernorm (Ba et al., (016) in recurrent attention mechanism was critical to achieve more stable training regime and so better performance.
Table 3: bAbl joint training
Length
Memory Networks
DNC
Universal Transformer
Memo
10k
4.2 (17/20)
3.8 (18/20) 0.29 (20/20)
0.21 (20/20)
Test results for the best run (chosen by validation loss) on the bAbl task. DNC results from (Graves et al. 2016), Universal Transformer results from (Dehghani et al.] (018). In parenthesis we report the number tasks solved, with 20 being the maximum.
CONCLUSIONS
In this paper we conducted an in-depth investigation of the memory representations that support inferential reasoning and we introduce MEMO, an extension to existing memory architectures, that shows promising results on these reasoning tasks. MEMO showed state-of-the-an results in new proposed task, the paired associative inference, which had been used in the neuroscience literature to explicitly test the ability to perform inferential reasoning. On both this task, and challenging graph traversal task, MEMO was the only architecture to solve long sequences. Also, MEMO was able to solve the 20 tasks of the bAbI dataset, thereby matching the performance of the current state-of-the-arr results. Our analysis also supported the hypothesis that these results achieved by the flexible weighting of individual elements in memory allowed by combining together the separated storage of single facts in memory with powerful recurrent attention mechanism.

---

Published as a conference paper at ICLR 2020
REFERENCES
J Ba, J Kiros, and G Hinton. Layer normalization. arxiv. 2016.
Andrea Banino, Raphael Koster, Demis Hassabis, and Dharshan Kumaran. Retrieval-based model accounts for striking profile of episodic memory and generalization. Scientific Reports, 6, 2016.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Goonalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: 180.00121, 2018.
A. Bhattacharyya. On a measure of divergence between two multinomial populations. Sankhyā; The Indian Journal of Statistics (9993-9660)) 7(4):401 1946. ISSN 00364452. http:
77 wwww. 1stor.
Tolga Bolukbasi, Joseph Wang, Ofer Dekel, and Venkatesh Saligrama. Adaptive neural networks for efficient inference. In Proceedings of the 34th International Conference on Machine Leurning- Volume 70. pp. 527-536. JMLR. org, 2017.
M Bunsey and H Eichenbaum. Conservation of hippocampal memory function in rats and humans. Nature, (79(656)):255, 1996.
Victor Campos Camunez, Brendan Jou, Xavier Giró Nieto, Jordi Torres Viñals, and Shih-Fu Chang. Skip mn: learning to skip state updates in recurrent neural networks. In Sixth International Conference on Learning Representations: Monday April 30- 30-Thursday May 03, 2018, Vancouver Convention Center, Vancouver.(proceedinnss) pp. 1-17, 2018.
Junyoung Chung. Caglar Gulcehre, Kyunghyun Cho, and Yoshua Bengio. Gated feedback recurrent neural networks. In International Conference on Machine Learning, pp. 2067-2075, 2015.
Junyoung Chung. Sungjin Ahn, and Yoshua Bengio. Hicrarchical multiscale recurrent neural networks. arXiv preprint 11900:160,,01744, 2016.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Eukasz Kaiser. Universal transformers. arXiv preprint / 2018.
J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei. ImageNet: A Large-Scale Hierarchical Image Database. In CVPR09, 2009.
Jacob Devlin, Ming-Was Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
Howard Eichenbaum and Neal J Cohen. From conditioning to conscious recollection: Memory systems of the brain. Number 35. Oxford University Press on Demand, 2004.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint 10Xiv:1803.03635, 2018.
Justin Gilmer, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In Proceedings of the 34th International Conference on Searning-Volume Volume 70, pp.
Machine Learning-
JMLR. org, 2017.
Marco Gori, Gabriele Monfardini, and Franco Scarselli. A new model for learning in graph domains. In Proceedings. 2005 IEEE International Joint Conference on Neural Networks, 2005., volume 2, pp. 729-734. IEFE, 2005.
Alex Graves. Adaptive computation time for recurrent neural networks. arXiv preprint 1rXiv:1603.08983, 2016.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint 110:5.1400.5401, 2014.
10

---

Published as a conference paper at ICLR 2020
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwiñska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using neural network with dynamic external memory. Nature, 538(7666):47,, 2016.
Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. In Advances in Neural Information Processing Systems, pp. 1024-1034, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 770-778, 2016.
Mikael Henaff, Jason Weston, Arthur Szlam, Antoine Bordes, and Yann LeCun. Tracking the world state with recurrent entity networks. arXiv preprint urXiv:16/2.03969, 2016.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint orXiv:1503.0255 2015.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 14Xiv:1112.69880 2014.
Thomas Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv:1609 160.00699 2016.
Raphael Koster, Martin J Chadwick, Yi Chen, David Berron, Andrea Banino, Emrah Düzzl,, Demis Hassabis, and Dharshan Kumaran. Big-loop recurrence within the hippocampal system supports integration of information across episodes. Neuron, 99(6):1342-1354 2018.
Ankit Kumar, Ozan Insoy, Peter Ondruska, Mohit lyyer, James Bradbury, Ishaan Gulrajani, Victor Zhong, Romain Paulus, and Richard Socher. Ask me anything: Dynamic memory networks for natural language processing. In Proceedings of The 33rd International Conference Machine Learning, pp. 1378-1387, 2016.
Dharshan Kumaran and James L McClelland. Generalization through the recurrent interaction of episodic memories: a model of the hippocampal system. Psychological Review, 119(3):573, 2012.
Yujia Li, Daniel Tarlow, Mare Brockschmidt, and Richard Zemel. Gated graph sequence neural networks. arXiv preprint 1 2015.
Yujia Li, Oriol Vinyals, Chris Dyer, Razvan Pascanu, and Peter Battaglia. Learning deep generative models of graphs. arXiv preprint 18X3::18003033244, 2018.
Christos Louizos, Max Welling, and Diederik P Kingma. Learning sparse neural networks through 1_0 regularization. arXiv preprint arXiv:1712.01312. 2017.
David Marr, David Willshaw, and Bruce McNaughton. Simple memory: theory for archicortex. In From the Retina to the Neocortex, pp. 59-128. Springer, 1991.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy P Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In Proceedings the 33nd International Conference Machine Learning, ICM62016, New York City, NY, USA, June 19-24, 2016, pp. 1928-1937, 2016.
Juan Pavez, Héctor Allende, and Héctor Allende-Cid Working memory networks: Augmenting memory networks with a relational reasoning module. In ACI., 2018.
Jack Rae, Jonathan J Hunt, Ivo Danihelka, Timothy Harley, Andrew W Senior, Gregory Wayne, Alex Graves, and Timothy Lillicrap. Scaling memory-augmented neural networks with sparse reads and writes. In Advances in Neural Information Processing Systems, pp. 3621-3629, 2016.
Adam Santoro, David Raposo, David G Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap. simple neural network module for relational reasoning. In Advances in neural information processing systems, pp. 4967-4976, 2017.
11

---

Published as a conference paper at ICLR 2020
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE Transactions on Neural Networks, 20(1)x61-800 2008
Anna C Schapiro, Nicholas B Turk-Browne, Matthew M Botvinick, and Kenneth A Norman. Com- plementary learning systems within the hippocampus: neural network modelling approach reconciling episodic memory with statistical learning. Philosophical Transactions of the Royal 2017.
Society B: Biological Sciences, (77((111):201600)))
Minjoon Seo, Sewon Min, Ali Farhadi, and Hannaneh Hajishirzi. Query-reduction networks for question answering, arXiv preprint (eXiv:1606004580 2016.
Yelong Shen, Po-Sen Huang, Jianfeng Gao, and Weizhu Chen. Reasonet: Learning stop reading in machine comprehension. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 1047-1055, ACM, 2017.
Larry R Squire, Craig EL Stark, and Robert E Clark. The medial temporal lobe. Annu. Rev. Neurosci., 27:279-306, 2004.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way to prevent neural networks from overfitting. The Journal of Machine Learning Research, (9l):1992-19988 2014.
Sainbayar Sukhbaatar, Jason Weston, Rob Fergus, et al. End-to-end memory networks. In Advances in Neural Information Processing Systems, pp. 2440-2448, 2015.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
T. Tieleman and G. Hinton. Lecture 6.5 35-RmsProp: Divide the gradient by running average of its recent magnitude. COURSERA: Neural Networks for Machine Learning, 2012.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pictro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint arXiv:1710.10903 2017.
Jason Weston, Antoine Bordes, Sumit Chopra, Alexander Rush, Bart van Merrinboer, Armand Joulin, and Tomas Mikolov. Towards ai-complete question answering: A set of prerequisite toy tasks. arXiv preprint arXiv: 500. 556905698, 2015.
Ronald J Williams. Simple statistical gralient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3-4):229-25 256, 1992.
Michael A Yassa and Craig EL Stark. Pattern separation in the hippocampus. Trends in neurosciences, 34(10):515-522 2011.
Adams Wei Yu, Hongrae Lee, and Quoe V Le. Learning to skim text. arXiv preprint arXiv:1704.06877, 2017.
Vinicius Zambaldi, David Raposo, Adam Santoro, Victor Bapst, Yujia Li. Igor Babuschkin, Karl Tuyls, David Reichert, Timothy Lillierap, Edward Lockhart, et al. Relational deep reinforcement learning. arXiv preprint arXiv: 0806.0006// 2018.
Dagmar Zeithamova, Margaret L Schlichting, and Alison R Preston. The hippocampus and inferential reasoning: building memories navigate future decisions, Frontiers in human neuroscience, 6:70, 2012.
ACKNOWLEDGMENTS
The authors would like thank Adam Santoro, and many other colleagues DeepMind for useful discussions and feedback on this work.
12

---

Published as a conference paper at ICLR 2020
A APPENDIX
A.1 PAIRED ASSOCIATIVE INFERENCE TASK
(Please refer to Figure 1 in the main text)
To make this task challenging for neural network we started from the ImageNet dataset (Deng et al. 2009). We created three sets, training, validation and test which used the images from the respective three sets of ImageNet to avoid any overlapping. All images were embedded using a pre-trained ResNet (He et al. 2016). We generated 3 distinct datasets with sequences of length three (i.e. A A-AA C), four (i.e. A B C D) and five (i.e. A - B-C C D - E)) items. Each dataset contains le6 training images, Ie5 evaluation images and 2c5 testing images. Each sequence was randomly generate with no repetition in each single dataset.
To explain how the batch was built let's refer to sequences of length, S. being equal to 3. Each batch entry is composed by a memory, a query and a target. In order to create a single entry in the batch we selected N sequences from the pool, with N = 16.
First, we created the memory content with all the possible pair wise association between the items in the sequence, e.g. A₁ and B2C,, A2 and B,C2, ..., Bn and BN For S = 3, this resulted in a memory with 32 rows.
Then generated all the possible queries. Each query consist of 3 images: the cuc, the match and the lure. The cue is an image from the sequence (e.g. A1), as is the match (e.g. C1). The lure is an image from the same memory set but from different sequence (e.g. C7). There are two types of queries 'direct' and 'indirect'. In 'direct' queries the cue and the match can be found in the same memory slot, so no inference is require. For example, the sequence A B Ci produces the pairs A B1 and B1 C₁ which are stored different slots in memory. An example of a direct test trail would be A1 (cue)- B1 (match) B3 (lure). Therefore, 'direct' queries are test of episodic memory as the answer relies on retrieving an episode that was experienced. In contrast, 'indirect' queries require inference across multiple episodes. For the previous example sequence, the inference trail would be A₁ (cue)- C₁ (match)- C3 (lure). The queries are presented to the network as concatenation of three image embedding vectors (the cue, the match and the lure). The cue is always in the first position in the concatenation, but to avoid any degenerate solution, the position of the match and lure are randomized, is worth noting that the lure image always has the same position in the sequence (e.g. if the match image is C the lure is also C) but it is randomly drawn from different sequence that is also present in the current memory. This way the task can only be solved by appreciating the correct connection between the images, and this need to be done by avoiding the interference coming for other items in memory. For each entry in the batch we generated all possible queries that the current memory store could support and then one was selected at random. Also the batch was balanced, i.c. half of the elements were direct queries and the other half was indirect. The targets that the network needed to predict are the class of the matches.
It is worth mentioning that longer sequences provide more 'direct' queries, but also multiple 'indirect' queries that require different levels of inference, e.g. the sequence An B, C, - D,, En produces the 'indirect" trial A (cue)- C₁ (target)- C3 (lure) with 'distance' 1 (one pair apart) and A (cue)- E, (target)- Eg (lure) with 'distance' (4 pairs apart). The latter trial required more inference steps and requires appreciate the overlapping images of the entire sequence.
Finally we use the inputs as follows:
For EMN and MEMO, memory and query are used as their naturally corresponding inputs in their architecture.
In the case of DNC (section G).
we embed stories and query in the same way it is done for MEMO. Memory and query are presented in sequence to the model (in that order), followed by blank inputs as pondering steps to provide a final prediction.
For UT, embed stories and query in the same way it is done for MEMO. Then we use the encoder of UT with architecture described in Section We use its output as the output of the model.
13

---

Published as a conference paper at ICLR 2020
A.2 PAI QUERIES WISE RESULTS
The result repoted below are from the evaluation set at the end of training. Each evaluation set contains 600 items.
Table 4:
Paired Associative length 3: A-B-C
Trial
EMN
DNC
UT
MEMO
Type
A-B
98.19
98.58
97.43
99.82(0.30)
B-C
97.93
99.34
98.28
99.76(0.38)
A-C
61.01
96.85
85.60
98.26(0.67)
Table 5:
Paired Associative length 4: A-B-C-D
Trial
EMN
DNC
UT
MEMO
Type
A-B
96.31
94.26
99.32
99.57(0.20)
B-C
97.57
84.94
88.31
99.33(0.13)
C-D
96.59
95.68
93.37
99.58(0.13)
A-C
48.71
49.38
54.87
98.93(0.15)
B-D
47.42
49.89
51.92
99.14(0.11)
A-D
48.63
58.63
44.16
97.22(0.13)
14

---

Published as a conference paper at ICLR 2020
Table 6:
Paired Associative length 4: A-B-C-D-E
Trial
EMN
DNC
UT
MEMO
Type
A-B
95.68
98.88
96.94
99.20(0.43)
B-C
95.82
94.60
92.63
98.93(0.17)
C-D
95.43
95.20
89.99
97.27(0.21)
D-E
95.16
95.98
97.27
95.06(0.12)
A-C
48.68
48.66
41.85
87.33(0.12)
B-D
45.75
46.87
39.62
86.65(1.27)
C-E
49.46
49.51
35.87
87.08(0.92)
A-D
52.08
50.32
52.38
86.12(0.57)
B-E
46.69
52.27
43.27
86.37(0.77)
A-E
48.31
48.79
47.93
84.54(5.72)
For EMN and DNC the results shown are from the hyper-parameter with the lower loss on the validation For MEMO the results are the average and relative standard deviation (reported in parenthesis), obtained by averaging the 5 hyper-parameters with the lower loss on the validation set.
A.3 PAIRED ASSOCIATIVE INFERENCE FURTHER ANALYSIS ON LENGTH 3 EXPERIMENTS.
DIRECT TRIALS
100
100
90
80
90
ACCURACY (%)
70
60
ACCURACY
80
50
40
70
20
60
50
d
INFERENCE TRIALS
20
100
16
90
NUMBER OF HOPS
ACCURACY
12
60
70
60
50
3000
SEPARATE MEMORIES
JOINT MEMORIES
EPOCH
ATTENTION
ATTENTION
Figure 3: Analysis of length PAI task. Evaluation accuracy the inference trial A-C; b. Number of hops taken during training; Distribution of evaluation accuracy obtained by averaging direct queries (A-B and B-C). This was obtained over 100 different but on the inference queries (A-C)
and seeds; d. same as
15

---

Published as a conference paper at ICLR 2020
A.3.1 ABLATIONS
Table 7: PAI - Ablations sequence of length 3: A-B-C
MEMO Network Architecture
A-C inference trial
Positional encoding as in Vaswani et al. 2017
Memories
Recurrent attention
kept separated
Accuracy
w/ Layernorm
X
X
57.59(10.11)
X
52.79(3.12)
X
73.26(15.86)
97.59(1.85)
Results for the best run (chosen by validation set) on the PAI task. X= not present; /= present
A.3.2 ATTENTION WEIGHTS ANALYSIS
HOP
CUE
LURE MATCH
PROB
PROB
MEMORY
622
HOP3 3
HOP
HOP
HOP 6
PROB
PROB
Figure 4: Attention weights analysis of length 3 PAI task, in the case where the network converged to 7 hops. In this case the network uses the first two hops to retrieve the slot where the cue is present and the the hops number 3, 4 and 5 to retrieve the slot with the match. The weights are sharp and they focus only on single slot.
16

---

Published as a conference paper at ICLR 2020
A.3.3 ADAPTIVE COMPUTATION
100
NBTO IIACT
80
60
50
EPOCH
Figure 5: Comparison between MEMO + REINFORCE and MEMO ACT on length 3 PAI task. MEMO wit REINFORCE shows more data efficiency that the one where the adaptive computation done with ACT.
B SHORTEST PATH TASK
B. TRAINING DETAILS
Graph generation In the shortest path experiments, generate the graph in the same fashion as Graves et al. (2016): the graphs used to train the networks are generated by uniformly sampling set of two-dimensional points from unit square, each point corresponding each node, the K nearest neighbours in the square are used as the K outbound connections, with K independently sampled from a uniform range for each node.
node in the graph. For
Graph representation We represent task in three parts: a graph description, query, and the larget. The graph description is presented sequence of tuples of integers that represent a connection between nodes, holding a token for the source node, and another for the destination node. The query is also represented as a tuple of integers, although, in that case, source and destination are simply the beginning and end of the path find. The target is the sequence node IDS that constitute the path between source and destination of the query.
When training, we sample a mini-batch of 64 graphs, with associated queries and target paths. Following our description above, queries are represented as matrix of size 64 x 2, targets are of size 64 x (T. - 1), and graph descriptions are of size 64 x M 2, where is the length of the shortest path, and .M is the number of maximum nodes we allow graph description to have. In our experiments, we fix the upper bound M to be the maximum number nodes that we have multiplied by the out-degree of the nodes in the graph.
All networks were trained for 2e4 epochs, each one formed by 100 batch updates.
For EMN and MEMO, we set the graph description to be the contents of their memory, and we use the query as input. In order answer the sequence of nodes that is used as target, we keep the keys 1.1h) and values v. fixed, and we proceed to use our algorithm as described for each answer, with independent numbers of hops for each one. The model then predict the answers for nodes sequentially: the first node is predicted before the second, However, one important difference between MEMO and EMN that for EMN we use the ground truth answer of the first node as the query for the second node, whereas for MEMO we used the answer predicted by the model for the first node as the query for the second node. This was done to enhance the performance of EMN while testing the real capabilities of MEMO to sequentially reasoning over multiple steps problems. The weights that are used for each answer are not shared.
For the Universal Transformer, also embed the query and graph description done for EMN and MEMO. After that, we concatenate the embeddings query and graph description and use the encoder of the UT architecture (with specific description in Section H). We use
17

---

Published as a conference paper at ICLR 2020
its output answer. After providing answer, that answer is provided as initial query for the following round of hops. The weights that are used for each answer are not shared.
For DNC, we also embed the query and graph description as done for EMN and MEMO. Since it is naturally sequential model, the information is presented differently: the tuples the graph description are presented first, and after that the query tuple is presented. After that, the pondering steps are used to be able to output the sequence of nodes that constitute the proposed shortest path.
The output of the models is trained using Adam using cross-entropy loss against all the sampled larget sequences. Training is done for fixed number of steps, detailed in Appendix Section D.2
For evaluation, we sample batch of 600 graph descriptions, queries, and targets. We evaluate the mean accuracy over all the nodes of the target path. We report average values and standard deviation over the best hyper parameters we used.
It is worth noting that given this training regime:
DNC and UT have a 'global view' on the problem in order to provide an answer for the second node. This means that, answer the second node in the path, they can still reason and work backwards from the end node, and while still having information about the initial node in the path. This makes it intuitive for them to achieve better performance in the second node, as it is closest to the end node of the path, so less reasoning is needed to achieve good performance.
On the contrary, MEMO has 'local view' on the problem, the answer the second node depends on the answer about the first node. Therefore, it cannot do better than chance if the answer to the first node is not correct.
B.2 RESULTS
To better compare the performance of MEMO versus EMN we also run another experiment where we tested the models in two conditions:
The ground truth answer of the first node was used the query for the second node.
The answer predicted by the model for the first node used as the query for the second node.
The results are summarized in Table In the case of 20 Nodes with outbound edges, we can see that if we give MEMO the ground truth for node 1 as query for node 2 the performance increases from the one related to the prediction of the first node (0.055 vs. 69.20%(0.077 Interestingly, if we use for EMN the same training regime used for MEMO i.e. the prediction is used to query the second node then EMN perform almost at chance level (22.30%). The same results are also confirmed in the simpler scenario with 20 nodes and outbound edges.
Table 8: Undirected graph shortest path- Comparing results on the second node based on using ground truth or predicted answer of the first node.
Prediction of First Node
Prediction of
Graph Structure
Sccond Node
EMN ground truth
EMN
MEMO
EMN
Nodes
Out degree
Path length
EMN
MEMO
predicted answer
ground truth 00.002
predicted annswer
20
31.99
94. 44.00.222
26.00
93.660 68. 8000.0)
20
3
23.99
29.20(0.07)
43,00
22.30
85.38(0.05)
Test results for the best 5 dlyeer-paameee (chosen by training loss) for MEMO, mean and corresponding standard deviation are reported. For EMN we report results from the best run.
18

---

Published as a conference paper at ICLR 2020
C
BABI
C.I TRAINING AND EVALUATION DETAILS
For this experiment we used the English Question Answer dataset Weston et al. (2015). We use the training and test datasets that they provide with the following pre-prccessing
All text is converted to lowercase.
Periods and interrogation marks were ignored.
Blank spaces are taken as word separation tokens.
Commas only appear in answers, and they are not ignored. This means that, e.g. for the path finding task, the answer 'n,s' has its own independent label from the answer 'n,w'. This also implies that every input (consisting of "qurry" and "stories') corresponds single answer throughout the whole dataset.
All the questions are stripped out from the text and put separately (given as "queries" to our system).
At training time, we sample a mini-batch of 128 queries from the test dataset, as well as its corre- sponding stories (which consist of the text prior to the question). As result, the queries are matrix of 128 x 11 tokens, and sentences are size 128 x 320 x 11, where 128 the batch size, 320 is the max number of stories, and 11 is the max sentence size. We pad with zeros every query and group of stories that do not reach the max sentence and stories size.
For EMN and MEMO, stories and query used their naturally corresponding inputs in their architecture.
In the case of DNC, we embed stories and query in the same way it is done for MEMO. Stories and query are presented in sequence to the model (in that order), followed by blank inputs as pondering steps provide final prediction.
For UT, we embed stories and query in the same way it is done for MEMO. Then, we use the encoder of UT with architecture described in Section We use its output as the output of the model.
After that mini-batch is sampled, we perform one optimization step using Adam for all the models that we have run in experiments, with hyper parameters detailed in Appendix Section D.2 We stop after a fixed number of time-steps, as also detailed in D.2
Many of the tasks in bAbl require added column vector to the memory store with a time encoding derived from Vaswani et al. 2017).
notion of temporal context, to account for this in MEMO
All networks were trained for 2c4 epochs, each one formed by 100 batch updates.
For evaluation, we sample a batch of 10, 000 elements from the dataset and compute the forward pass in the same fashion done in training. With that, we compute the mean accuracy over those examples, as well as the accuracy per task for each of the tasks of bAbl. We report average values and standard deviation over the best 5 hyper parameters used.
19

---

Published as a conference paper at ICLR 2020
C.2 BABI ABLATIONS
Table 9: bAbl - Ablations
MEMO Network Architecture
bAbI 10K
Positional encoding in Vaswani al. 2017)
Memories
Recurrent attention
Number of solved tasks
keept separated
Layernorm
Average error
X
X
11/20
14.81
X
X
X
14/20
10,43
AXX X
X
V
17/20
4.80
V
X
18/20
5.00
20/20
0.21
Results for the best run (chosen by validation set) on the bAbl task. The model was trained and tested jointly on all tasks, All tasks received approximately equal training resources. X= not present; /= present
20

---

Published as a conference paper at ICLR 2020
C.3 TASK TASK-WISE RESULTS
Table 10: bAbl Results average over 5 hyper-parameters with lower loss on the valiation set
Trial
MEMO
MEMO
Type
top 5 seeds
1 Single Supporting Fact
100.00
100.00(0.00)
2 - Two Supporting Facts
100.00
19.13(1.78)
3 Three Supporting Facts
97.05
94.15(6.35)
4 Two Arg. Relations
100.00
100.00(0.00)
5 Three Arg. Relations
100.00
100.00(0.00)
6 - Yes/No Questions
100.00
100.00(0.00)
7 Counting
100.00
96.69(3.57)
8 Lists/Sets
100.00
99. 191.94)
9 - Simple Negation
100.00
100.00(0.00)
10 Indefinite Knowledge
100.00
99.35(1 44.
11 Basic Coreference
100.00
100.00(0.00)
12- Conjunction
100.00
100.00(0.00)
13 Compound Coref
100.00
100.00(0.00)
14 Time Reasoning
100.00
100.00(0.00)
15 Basic Deduction
100.00
100.00(0.00)
16 Basic Induction
98.75
95.05(5.12)
17 Positional Reasoning
100.00
100.00(0.00)
18- Size Reasoning
100.00
99.13(1.94)
19 Path Finding
100.00
100.00(0.00)
20 Agent's Motivations
100.00
100.00(0.00)
Mcan error
0.21
0.86(1.11)
Solved tasks (>95% accuracy)
20/20
n/a
(Mean and standard deviation of test errors for the best 5 hyperparameters (chosen according the validation loss).
D MEMO TRAINING DETAILS AND HYPER-PARAMETERS
D. TRAINING DETAILS
To train MEMO network parameters we use Adam (Kingma & Ba 2014) with polynomial learning rate decay, starting at 1, 1,tartmemo value, and batch size always equal 64 for the PAI and shortest path and 128 for bAbl. In all the three tasks MEMO was trained using a cross entropy loss, and the network had to predict the class ID in the paired associative inference task, the node ID in the shortest path problem and the word ID in bAbi.
The halting policy network parameters with learning rate haalt-
updated using RMSProp TTeleman & Hinton 2012)
The other parameters reported in Table 11 and 12
21

---

Published as a conference paper at ICLR 2020
D.2
FIXED SEPEE-PARRAETER USED ACROSS TASKS
Table 11: Fixed hyper-parameterr used across tasks
Tasks
Parameter
PAT
PAT
Shortest path
Shortest path
Shortest path
bAbl
name
length 3
length 4
length 5
10-2-2
20-3-3
20-5-3
32
48
64
20
60
100
320
S
3
3
3
2
11
O
1000
1000
1000
1000
1000
1000
de
128
128
128
128
128
128
128
d
256
256
256
512
512
512
512
da
128
128
128
128
256
256
256
DropOuta
0.1
0.1
0.1
0.1
0.1
0.1
0.1
DropOut,,
0
0.5
GRUR hidden size MLPR number of layers
256
256
256
256
256
1
M LPR
64
64
64
64
64
64
64
hidden size
D.3 RANGE OF HYPER-PARAMEEN PARAMETERS USED IN SWEEPS
Table 12: Range of hyper-parameters used in sweeps
Tasks
Parameter
PAT
PAT length
Shortest path
Shortest path
Shortest path
bAbl
namne
length
length
10-2-2
20.3-3
20-5-3
[3
[5, 20)
[0]
(max number of hops)
[0.85, 0.9]
[0.9]
[0.85, 0.9]
[le-4, 1e-2, 0.1]
[le-2, 0.1]
[le-4, 1e-2, 0.1]
[1e-3, 1e-2, 0.1] [].
[le-2, 0.1]
[1e-3, 1e-2, 0.1]
bias
10]
(2.
H Ltarlmemo
1
[4.
[4, 8]
[7e-4, Se-4, le:4]
Ihatt
[1e44, Se-5]
22

---

Published as a conference paper at ICLR 2020
E МЕMО COMPLEXITY ANALYSIS
In terms of temporal complexity, MEMO has a complexity of O(,, A. H. d), where n, is the number of samples we process with our network, A is the number of answers, N is the upper bound the number of hops we can take, H the number of heads used. is the number stories, and S is the number of words in each sentence. This is due to the fact that, for each sample, we do the hopping procedure for every answer, taking number of hops. For each hop we query our memory by interacting with all its slots I, for all its size S d. For all our experiments, all parameters A, N, H, 7, S, d are fixed to constants.
Further, it is worth noting that MEMO is linear with respect the number of sentences of the input, whereas the Universal Transformer has quadratic complexity.
With respect spatial complexity, MEMO holds information of all the weights constant, apart from all the context information that needs to be used to answer particular query. Since the context information is the only one that is input dependent, the spatial complexity is in this case O S d). which is the size of our memory. In all our experiments, such size is fixed.
F ACT DESCRIPTION
We implement ACT as specified in Graves 2016). Based our implementation of MEMO, start by defining the halting unit h as the following:
h4 (
(19)
where Ny the binary policy of MEMO. This slightly different than the original ACT which represents such unit with:
h4 = o + 6n)
(20)
where Wa and by are trainable weights and biases respectively, and 8, is the previous observed state. We argue that this slight change increases the fairness of the comparison, for two reasons: firstly, Te st. f) depends on st. but it uses several non-linearities to do so, rather than being a simple linear projection, so it should enable more powerful representations. Secondly, this makes it much more similar to our model while still being able to evaluate the feasibility of this halting mechanism.
From this point we proceed as in the original work by defining the halting probability:
Pt =
R i- hr
(21)
otherwise
where
min
h4 1>=1-1
(22)
where e, as Graves 2016) is fixed to be 0.01 in all experiments. The reminder R is defined as:
R R= Σ h
(23)
Finally, the answer provided by MEMO+ACT is defined as:
Praz
(24)
where a, corresponds to the answer that MEMO has provided at hop
G DNC ARCHITECTURE AND HYPERPARAMETERS
We use the same architecture as described Graves et al. (2016), with exact sizes of each layer described on Table [3
We also performed a search on hyperparameters to find such hyperparameters, with ranges reported on Table 14
23

---

Published as a conference paper at ICLR 2020
Parameter name
Value
Optimizer algorithm Learning rate
Adam
0.0005
Input embedding size
128
Input pondering steps Controller type
10
L.TM
Controller hidden size Memory number of read heads
128
3
Memory word size
384
Memory output size
128
Table 13: Hyperparameters used on all tasks trained with DNC.
Parameter name
Value {0.0001, 0.0005)
Learning rate Input pondering steps
{0, 5, 10) (1,3)
Memory number of read heads
Table 14: Hyperparameters ranges used to search over with DNC.
H UNIVERSAL TRANSFORMER ARCHITECTURE AND HYPERPARAMETERS
We use the same architecture as described in Dehghani et al. 2018 More concretely, we the implementation and hyperparamettess described uunivrssal stapaaatttaasfomrrssmeeeee that is available at http://s/wthhaaccc//////o//////a
tensor2tensor mode
py
For complete-
ness, we describe the hyperpameters used on Table 15
We also performed search hyperparameters train our tasks, with ranges reported Table
16
Parameter name
Value
Optimizer algorithm Learning rate
Adam
5e-4
Input embedding size
128
Attention type
Dot product
Attention hidden size Attention number of heads Transition function Transition hidden size Number of hidden layers Recurrence type Number of recurrent steps Attention dropout rate ReLU dropout rate
512
Fully connected neural network
128
'RNN' (not LSTM)
0.1
0.1
Layer pre/postprocess dropout
0.1
Table 15: Hyperparameters used for all experiments for UT.
24

---

Published as a conference paper at ICLR 2020
Parameter name
Value
Learning rate
[le-3, 5e-4, lc-4]
Number of hidden layers Attention hidden size Transition hidden size Attention number of heads
(2, 6
[128, 512
(512, 128)
[4, 8)
Table 16: Hyperparameters ranges used to search over with UT.
25


Reviews:
[{"review_id": "BkxAILESqH", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nThis paper proposes two main changes to the End2End Memory Network (EMN) architecture: a separation between facts and the items that comprise these facts in the external memory, policy to learn the number of memory-hops to reason. The paper also introduces a new Paired Associative Inference (PAI) task inspired by neuroscience and shows that most of the existing models including transformers struggle to solve this task while the proposed architecture (called MEMO) solves it better. MEMO also works well in the shortest path finding tasks and bAbI tasks.\n\nMy comments:\n\nOverall, I see this paper as an improvement over EMN. The proposed PAI task can be seen as an example task where Transformers struggle while recurrent architectures learn better. Interestingly, the authors use a separate halting policy network to reduce computation time.\n\n1. Section 2.1 requires more clarity. There is a confusion in the usage of I and S. I represents the number of stories or the number of sentences in the stories?\n2. Scaling up NTM/DNC to larger memory work was done in Sparse Access Memory (SAM) by Rae et al. 2016. This needs to be included in section 3.1\n3. In Table 2, why is the prediction of the second node easier for the model than the prediction of the first node? I see this trend only for EMN, UT, DNC. Not in MEMO.\n4. Are the authors willing to release the code and data to reproduce their results?\n\nMinor comments:\n\n1. Page 2, second para: ENM should be EMN.\n2. Vec inverse in Eqn 14 was never introduced.\n3. Table 3: The notation of (20/20) was never introduced. I can guess what it means. But please be explicit.\n\n===============\nAfter rebuttal: Authors have addressed all my questions. I  recommend  \"Accept\".\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "BJlbQ9mwqS", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nThis paper presents a new task (paired associate inference), drawn from cognitive psychology, which requires linking many pieces of information together to make inferences with long range dependencies. Experimental results show that standard memory architectures fail on these tasks. To redress this, the paper proposes a new memory architecture with several new features that allow for much better performance on the paired associate task. Finally, the paper undertakes systematic experiments on more traditional domains like shortest path problems, showing that the new architecture achieves modest improvements.\n\nMajor comments:\n\nOverall this is an interesting and useful work which uses a task from cognitive psychology to illuminate reasoning limitations in prior memory augmented neural networks. It then proposes several architectural and algorithmic fixes to improve performance. The resulting solution appears to work based on the experimental evaluation, and although this does not lead to substantial improvements in the state of the art on the bAbI dataset, it seems likely that these improvements could be useful for even harder task settings. \n\nThe new algorithmic improvements are justified mainly by their performance in experiments, which is completely acceptable if the experiments have been done to a high standard. Nevertheless, it would be wonderful if more insight could be gained about why exactly this configuration of architectural changes was selected and what they contribute. The ablation experiments go some way to addressing this question and perhaps the paper would benefit from a greater discussion of these results in the main text.\n\nThe experiments appear to have been done to a high standard, and it is promising that the method’s advantages are more pronounced as the task gets harder (eg by demanding longer chains of inference).\n\nThe paper could also be improved by more clearly describing how hyper parameters were selected for each experiment (including the different ablation studies). Why were hyperparamters sometimes chosen based on training loss and other times chosen based on validation loss? In general it would be useful to explain the selection procedure very carefully for each result, since it appears to have differed between evaluations.\n\nThe paper is reasonably easy to follow, but could be streamlined to enable more discussion of material that is currently in the appendix.\n\nTypos:\n\nENM->EMN on pg2\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors introduce a new associative inference task from cognitive psychology, show shortcomings of current memory-augmented architectures, and introduce a new memory architecture that performs better with respect to the task. The reviewers like the motivation and thought the experimental results were strong, although they also initially had several questions and pointed to areas of the paper which lacked clarity. The authors updated the paper in response to the reviewer's questions and increased the clarity of the paper. The reviewers are satisfied and believe the paper should be accepted.
Decision: Accept (Poster)