PaperlayerView exact revision on arXiv

Source-linked AI summary

MEMO: A Deep Network for Flexible Combination of Episodic Memories

Andrea Banino, Adrià Puigdomènech Badia, Raphael Köster, Martin J. Chadwick, Vinicius Zambaldi, Demis Hassabis, Caswell Barry, Matthew Botvinick, Dharshan Kumaran, Charles Blundell

arXiv:2001.10913v1cs.LGcs.AI

TL;DR

Existing memory-augmented architectures struggle to reason over long-distance associations, motivating a more demanding associative-inference evaluation. MEMO separates memories from their constituent items and adaptively retrieves information over a variable number of hops, solving the new reasoning tasks and matching state-of-the-art bAbI results.

  • Problem

    Existing memory-augmented architectures struggle with long-distance associations, while bAbI can contain train–test repetitions and commonalities that permit degenerate solutions.

  • Method

    MEMO separates stored facts from their constituent items and uses an adaptive retrieval mechanism that learns how many memory hops to take before answering.

  • Results

    MEMO solves long sequences in paired associative inference and graph traversal, and matches state-of-the-art results across the 20 bAbI tasks.

  • Takeaways & Limitations

    Flexible item weighting and variable computation support inferential reasoning across distant relationships in the evaluated tasks.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper targets inferential reasoning across temporally separated facts, arguing that existing memory architectures need representations and computation that support long-distance associations. It introduces MEMO’s separated item representations and adaptive computation, evaluating them on new reasoning tasks and bAbI.

  • Motivation: Hippocampal pattern separation preserves distinct episodic memories while their interaction can support inference across experiences.The paper uses this neuroscience perspective to motivate investigating memory representations for neural-network reasoning.
  • Contributions: MEMO combines separated representations of facts’ constituent items with adaptive computation to support flexible long-distance inferential reasoning.The architecture separates memories from their constituent items and learns how many memory hops to take before answering.
  • Motivation: The new task stresses reasoning by requiring appreciation of distant relationships among elements distributed across multiple facts.This addresses the limitation that standard benchmarks can contain train–test repetitions and commonalities enabling degenerate solutions.
  • Contributions: A REINFORCE loss component learns the optimal number of computation iterations required for each task.The halting policy chooses whether to continue querying memory or answer, rather than using a fixed computation budget.
  • Evaluation: MEMO is evaluated on paired associative inference, shortest-path finding, and bAbI to test memory-based reasoning across three tasks.The contribution list identifies these experiments as evidence for the effectiveness of the proposed representations and adaptive computation.

2 METHODS

The methods compare MEMO with end-to-end memory networks and implement memory lookup through flexible item representations, multi-head attention, and learned halting. A REINFORCE-based policy selects a variable number of hops while directly penalizing expected computation.

  • 2.1 RECAPITULATING END-TO-END MEMORY NETWORKS: The baseline EMN embeds each input sentence into key and value memory vectors, attends over memory slots, updates the query, and predicts the answer at the final step.The setup represents knowledge inputs as a sequence of sentences and a query, with EMN using embedding matrices and positional encodings.
  • 2.2 MEMO: MEMO combines sentence-level embeddings with learned linear projections and multi-head attention while keeping the query separate from keys and values.Each attention head provides a different view of common input embeddings; dropout and layer normalization are also used.
  • 2.2 MEMO: MEMO retains full facts in memory, separates their constituent items, and learns flexible weighting during memory lookup.This contrasts with EMN’s fixed word embeddings combined with hand-coded positional encodings.
  • 2.3 THE HALTING POLICY: The halting policy observes each retrieval step and decides whether to continue querying memory or stop when computation has reached a fixed point.Repeated attention to the same memory slot indicates that further retrieval may return the same information.
  • 2.3 THE HALTING POLICY: REINFORCE trains the binary halting policy, while an additional loss term directly minimizes the expected number of hops.The binary halting variance is bounded by 1/4, and the implementation imposes a maximum hop count with no gradient sharing between the hop and main MEMO networks.

3 RELATED WORK

Related work includes external-memory networks, adaptive computation methods, conditional early exiting, REINFORCE-based discrete computation, and graph neural networks. These approaches address memory access, computation allocation, or iterative graph information propagation.

  • Memory-augmented networks: External-memory models such as EMN and DNC motivate architectures for abstract and relational reasoning by storing or accessing information beyond the immediate input.The paper positions MEMO as an extension of existing memory architectures rather than an unrelated architecture family.
  • Adaptive computation: Adaptive Computation Time learns a scalar halting probability to dynamically modulate computational steps according to each input.This provides a precedent for MEMO’s variable-computation strategy.
  • Adaptive computation: Adaptive Early Exit Networks stop computing higher layers when additional computation is unnecessary, while REINFORCE can learn discrete computation decisions.These methods represent alternative conditional-computation mechanisms discussed alongside MEMO’s halting policy.
  • Graph neural networks: Graph neural networks iteratively propagate node and edge embeddings through message passing before aggregating graph information for learning tasks.This iterative propagation is related to reasoning over graph structure, including the paper’s shortest-path setting.

4 PAIRED ASSOCIATIVE INFERENCE TASK

The paired associative inference task tests whether a model can infer a relationship between items that were never directly paired but are linked through a shared intermediate item.

  • Direct queries retrieve a directly experienced pair, whereas indirect queries require inferring the link between A and C through shared item B.For an A-B-C sequence, A-C is the indirect query and a lure from another sequence must be rejected.
  • PAI presents paired images across separate episodes and tests whether the model can connect distant relationships across those memories.The task is modeled on associative inference in neuroscience and uses sequences such as A-B-C.

5 RESULTS

MEMO outperformed competing architectures on long-distance associative and graph reasoning while matching state-of-the-art bAbI task coverage with lower error.

  • 5.2 PAIRED ASSOCIATIVE INFERENCE: MEMO was the only architecture to successfully answer the most complex length-4 and length-5 PAI inference queries.On length-3 PAI, MEMO matched DNC’s highest accuracy, while EMN and UT did not reach the same level.
  • 5.2 PAIRED ASSOCIATIVE INFERENCE: A length-3 PAI analysis showed that MEMO converged to 3 hops, whereas DNC required 10 pondering steps to reach the same accuracy.MEMO’s attention shifted from the cue-associated memory to the memory containing the matching continuation while avoiding the lure.
  • 5.2 PAIRED ASSOCIATIVE INFERENCE: Ablations showed that successful inference required combining separated fact representations with recurrent attention, while either component alone was insufficient.This combination was necessary for inference queries but not direct episodic-memory queries.
  • 5.3 SHORTEST PATH ON RANDOMLY GENERATED GRAPHS: MEMO outperformed EMN by more than 50% on complex shortest-path graphs and exceeded DNC by more than 20% on highly connected graphs.On smaller graphs, DNC, UT, and MEMO achieved perfect accuracy for the intermediate node.
  • 5.4 QUESTION ANSWERING ON THE BABI TASKS: MEMO solved all 20 bAbI tasks in the 10k training regime, matching the number solved by prior state-of-the-art models but with lower error.The result is reported for joint training across the bAbI task suite.

6 CONCLUSIONS

The paper introduces MEMO to support inferential reasoning over separated facts and variable-length reasoning chains, evaluating it on synthetic PAI and graph tasks and on bAbI.

  • 6 CONCLUSIONS: MEMO separates individual facts in external memory and flexibly weights their elements through recurrent attention.The architecture is intended to support inference across multiple stored associations rather than only direct memory retrieval.
  • A.1 PAIRED ASSOCIATIVE INFERENCE TASK: The evaluation uses ImageNet-derived sequences of length three, four, and five, with queries spanning direct and increasingly distant indirect associations.The task constructs memories from pairwise associations and varies the required inference distance.
  • A.1 PAIRED ASSOCIATIVE INFERENCE TASK: MEMO receives memory and query as corresponding architectural inputs, while DNC additionally uses blank pondering inputs and UT uses its encoder output.These input conventions define how the compared architectures are applied to the task.

A.2 PAI - QUERIES WISE RESULTS

The PAI results examine query-wise performance across increasing sequence lengths and analyze how MEMO retrieves facts, allocates hops, and compares against alternative architectures.

  • PAI query-wise results: The section reports PAI results for sequences A-B-C, A-B-C-D, and A-B-C-D-E, including query-wise tables for each length.
  • PAI query-wise results: Figure 3 analyzes length-3 PAI accuracy, training hop counts, and accuracy distributions for direct and inference queries across hyper-parameters and seeds.The analysis averages direct queries A-B and B-C separately from inference queries A-C.
  • PAI query-wise results: For a seven-hop solution, MEMO used early hops to retrieve the cue slot and later hops to retrieve the matching slot with sharp, single-slot attention.
  • PAI query-wise results: MEMO with REINFORCE was more data-efficient than MEMO using ACT for adaptive computation on the length-3 PAI task.

B.2 RESULTS

The shortest-path experiments compare models under ground-truth and predicted intermediate-node queries, highlighting how sequential prediction affects MEMO and EMN.

  • B.2 RESULTS: The experiment explicitly compared using the ground-truth first node versus the model-predicted first node as the second-node query.
  • B.2 RESULTS: In 20-node graphs with five outgoing edges, MEMO reached 85.38%(0.05) with the ground-truth first node but 69.20%(0.07) when using its predicted first node.The comparison concerns performance on the second node based on the query supplied to it.
  • B.2 RESULTS: When EMN used its predicted first node as the query for the second node, performance fell to 22.30%, with the same pattern confirmed in the simpler 20-node, three-outgoing-edge setting.

C.2 BABI ABLATIONS

The bAbI ablation results report task-level outcomes for jointly trained models and document the training objective used for MEMO.

  • C.2 BABI ABLATIONS: The listed bAbI result rows include 11/20, 14/20, 17/20, 18/20, and 20/20 task outcomes with corresponding values 14.81, 10.43, 4.80, 5.00, and 0.21.
  • C.2 BABI ABLATIONS: The model was trained and tested jointly on all bAbI tasks, with approximately equal training resources allocated to each task.
  • C.2 BABI ABLATIONS: The bAbI results table reports outcomes averaged over five hyper-parameters with the lowest validation loss.
  • C.2 BABI ABLATIONS: MEMO was trained with Adam and cross-entropy loss to predict class IDs for PAI, node IDs for shortest path, and word IDs for bAbI.

D.2 FIXED HYPER-PARAMETERS USED ACROSS TASKS

This section identifies the fixed hyper-parameters and the ranges used for hyper-parameter sweeps across the experiments.

  • D.2 FIXED HYPER-PARAMETERS USED ACROSS TASKS: Table 11 lists the fixed hyper-parameters used across tasks.
  • D.2 FIXED HYPER-PARAMETERS USED ACROSS TASKS: Together, the tables distinguish settings held fixed across tasks from values explored during hyper-parameter sweeps.
  • D.2 FIXED HYPER-PARAMETERS USED ACROSS TASKS: Table 12 lists the ranges of hyper-parameters used in sweeps.

E MEMO COMPLEXITY ANALYSIS

MEMO’s complexity analysis separates computational cost from memory size and compares its scaling with Universal Transformer and ACT-style halting.

  • MEMO complexity: MEMO has temporal complexity O(ns · A · N · H · I · S · d), reflecting answer-wise memory querying across multiple hops and attention heads.
  • MEMO complexity: MEMO is linear in the number of input sentences, whereas the Universal Transformer has quadratic complexity.
  • MEMO complexity: O(I · S · d) spatial complexity depends on the input memory size, while the stored weights remain constant across queries.The memory size is fixed in all experiments.
  • ACT comparison: The ACT comparison replaces ACT’s linear halting projection with MEMO’s policy-dependent unit to make the comparison more representationally similar.The authors argue that the modified unit can use more powerful nonlinear representations while remaining suitable for evaluating the halting mechanism.
  • Adaptive halting: MEMO’s adaptive halting policy chooses whether to continue querying memory or answer, learning termination criteria with a REINFORCE-trained binary variable.The policy outputs a halting probability and supports a variable number of computation steps.

G DNC ARCHITECTURE AND HYPERPARAMETERS

The experiments use established DNC and Universal Transformer architectures, with fixed hyperparameters and separate reported search ranges for each model family.

  • DNC architecture and hyperparameters: DNC experiments use the architecture described by Graves et al. (2016), with layer sizes specified in Table 13.
  • DNC architecture and hyperparameters: DNC hyperparameters were selected through a search over ranges reported in Table 14.
  • Universal Transformer architecture and hyperparameters: Universal Transformer experiments use the architecture and hyperparameters of Dehghani et al. (2018), specifically the available universal_transformer_small implementation.
  • Universal Transformer architecture and hyperparameters: Universal Transformer training also used a hyperparameter search with ranges reported in Table 16.
Loading 2001.10913v1…