MEMO: A Deep Network for Flexible Combination of Episodic Memories

TL;DR

The paper tackles the problem of inferential reasoning over distributed episodic memories and shows that existing memory-augmented networks struggle with long-range associations. It introduces MEMO, which separates memory contents from their components and employs an adaptive, multi-hop retrieval mechanism guided by a REINFORCE halting policy. The authors validate MEMO on three tasks—the neuroscience-inspired Paired Associative Inference, a shortest-path graph task, and the bAbI question-answering suite—demonstrating state-of-the-art or competitive performance and data-efficient inference. These results highlight the potential of memory-structure design and adaptive computation for scalable, robust memory-based reasoning in neural networks.

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 memory-based reasoning 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 MEMO, an architecture 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 match state of the art results in bAbI.

Figures (5)

• 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
• 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. $A_{10}-C_{10}$, 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. $C_{13}$. 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.
• Figure 3: Analysis of length 3 PAI task. a. Evaluation accuracy on the inference trial A-C; b. Number of hops taken during training; c. Distribution of evaluation accuracy obtained by averaging direct queries (A-B and B-C). This was obtained over 100 different hyper-parameters and seeds; d. same as c, but on the inference queries (A-C)
• 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 1 single slot.
• 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 is done with ACT.