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. 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 a coding summer camp and encountering another little girl with a woman. You might conclude that the woman is 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 is a relationship between the woman and the man. This flexible recombination of single experiences in novel ways to infer unobserved relationships is called inferential reasoning and is supported by the hippocampus (Zeithamova et al., 2012) . Interestingly, it has been shown that the hippocampus is storing memories independently of each other through a process called pattern separation (Yassa & Stark, 2011; Marr et al., 1991) . The reason hippocampal memories are kept separated is to 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.e. how can separated memories be chained together? Interestingly, a recent line of research (Kumaran & McClelland, 2012; Banino et al., 2016; Schapiro et 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 multiple pattern separated codes to 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 (Graves 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 a new task, called Paired Associative Inference (PAI -see below), which is derived from the neuroscientific literature (Bunsey & Eichenbaum, 1996; Banino et 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 it 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; Santoro et al., 2017; Pavez et al., 2018) 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 et 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 a 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 et 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 are 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 a 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 is used to calculate if the network has settled into a 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, 2016) . In our architecture, the network outputs an action (in the reinforcement 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 a 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; Louizos et al., 2017) however our approach differs in that we added 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. To sum up, our contributions are: 1. A new task that stresses the essence of reasoning -i.e. the appreciation of distant relationships among elements distributed across multiple facts. 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 a 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., 2015) . 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-art results in a 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 a 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-art results. Our analysis also supported the hypothesis that these results are achieved by the flexible weighting of individual elements in memory allowed by combining together the separated storage of single facts in memory with a powerful recurrent attention mechanism. To make this task challenging for a 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 − B − C), four (i.e. A − B − C − D) and five (i.e. A − B − C − D − E) items. Each dataset contains 1e6 training images, 1e5 evaluation images and 2e5 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 1 B 1 and B 1 C 1 , A 2 B 2 and B 2 C 2 , ..., A N B N and B N C N . For S = 3, this resulted in a memory with 32 rows. Then we generated all the possible queries. Each query consist of 3 images: the cue, the match and the lure. The cue is an image from the sequence (e.g. A 1 ), as is the match (e.g. C 1 ). The lure is an image from the same memory set but from a different sequence (e.g. C 7 ). 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 1 -B 1 -C 1 produces the pairs A 1 -B 1 and B 1 -C 1 which are stored different slots in memory. An example of a direct test trail would be A 1 (cue) -B 1 (match) -B 3 (lure). Therefore, 'direct' queries are a 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 1 (cue) -C 1 (match) -C 3 (lure). The queries are presented to the network as a 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. It is worth noting that the lure image always has the same position in the sequence (e.g. if the match image is a C the lure is also a C) but it is randomly drawn from a 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.e. 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 A n -B n -C n -D n -E n produces the 'indirect' trial A 1 (cue) -C 1 (target) -C 3 (lure) with 'distance' 1 (one pair apart) and A 1 (cue) -E 1 (target) -E 3 (lure) with 'distance' 4 (4 pairs apart). The latter trial required more inference steps and requires to 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, 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 H. We use its output as the output of the model.
A memory architecture that support inferential reasoning. This paper proposes changes to the End2End Memory Network architecture, introduces a new Paired Associative Inference task that most existing models struggle to solve, and shows that their proposed architecture solves the task better. A new task (paired associate inference) drawn from cognitive psychology, and proposal for a new memory architecture with features that allow for better performance on the paired associate task.