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

External-memory architectures are often evaluated on bAbI, where train–test repetitions can permit degenerate solutions, leaving distant-relation reasoning in need of direct testing. MEMO combines full-fact memory representations with adaptive retrieval; it was the only architecture reported to solve long sequences on paired associative inference and graph traversal, while matching current state-of-the-art results on bAbI.

  • Problem

    Whether memory-augmented networks can infer distant relations across facts is insufficiently tested by bAbI, whose train–test overlap may permit degenerate solutions.

    Source: §1 · p. 2

  • Method

    MEMO stores full facts and uses learned projection, recurrent attention, and a halting policy to decide whether to retrieve again or answer.

    Source: §1 · p. 2

  • Results

    MEMO was the only architecture to solve long sequences on paired associative inference and graph traversal, achieved state-of-the-art results on paired associative inference, and matched current state-of-the-art results on all 20 bAbI tasks.

    Source: §6 · p. 9

  • Takeaways & Limitations

    The findings show MEMO handling long-sequence reasoning on associative and graph tasks while matching current state-of-the-art performance on bAbI.

    Source: §6 · p. 9

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 investigates how neural networks can recombine separated experiences to infer distant relationships. It introduces a reasoning task and MEMO, which changes memory representations and adapts computation to the task.

1 INTRODUCTION

  • Inferential reasoning recombines separate experiences to infer unobserved relationships, a function supported by the hippocampus.The paper draws on findings that separated memories can interact at retrieval through a recurrent mechanism to support inference.Source: §1 · pp. 1–2
  • Paired Associative Inference (PAI) tests distant relationships across facts and uses procedural generation to require generalization to unseen associations.It was introduced to address the possibility that networks exploit train–test repetitions and commonalities in bAbI.Source: §1 · p. 2
  • MEMO retains full facts in memory and learns a linear projection with recurrent attention, enabling flexible weighting of individual memory elements.The paper investigates memory representations for inferential reasoning and extends existing memory architectures.Source: §1 · pp. 2–3
  • A learned halting policy uses REINFORCE to choose whether to continue querying memory or answer, with an added loss term that encourages fewer computation steps.This adapts computation to task complexity rather than relying on experimenter-selected amounts; retrieved content is recirculated as the next query, following REMERGE.Source: §1 · pp. 2–3
  • The authors report significant results across paired associative inference, shortest-path finding, and bAbI.The introduction gives no numerical results or benchmark comparisons.Source: §1 · p. 3

2 METHODS

The methods recap EMN (Sukhbaatar et al., 2015), then describe MEMO’s flexible item-level memory attention and a learned policy for selecting the number of memory hops.

2.1 RECAPITULATING END-TO-END MEMORY NETWORKS

  • EMN (Sukhbaatar et al., 2015) embeds input sentences as key/value memories, attends over I slots, updates the query across steps, and trains the final answer with cross-entropy.Source: §2.1 · p. 3

2.2 MEMO

  • MEMO embeds each sentence into shared item representations, adapts them as keys or values, and uses multi-head attention to flexibly recombine items rather than fixed positional embeddings.Source: §2.2 · pp. 3–4
  • MEMO’s multi-head attention keeps the query separate from keys and values, uses DropOut and LayerNorm, and scales linearly with sentence count rather than quadratic self-attention (Dehghani et al., 2018).Source: §2.2 · p. 4

2.3 THE HALTING POLICY

  • A GRU–MLP binary policy chooses whether to continue hopping from changes in attention weights and the hop count, using a fixed-point intuition for stopping.REINFORCE training includes a term that minimizes expected hops; binary halting variance is bounded by 1/4. A maximum N caps hops, and the hop policy shares no gradients with MEMO.Source: §2.3 · pp. 4–5

3 RELATED WORK

Related work surveys memory-augmented networks, adaptive computation, and graph neural networks. The paper distinguishes its approach through adaptive memory hops and direct attention from queries to memory slots.

3.1 MEMORY-AUGMENTED NEURAL NETWORKS

  • Memory-augmented networks include sequential and parallel designs, with models such as DNC extensions, Dynamic Memory Networks, Recurrent Entity Networks, and Working Memory Networks performing well on reference tasks including bAbI.The DNC reads and writes sequentially but was difficult to scale; sparse DNC extensions addressed larger tasks. Other models differ in input processing, memory updates, or added working-memory and relational-reasoning components.Source: §3.1 · p. 5

3.2 ADAPTIVE COMPUTATION TIME

  • ACT learns a halting probability to adjust computation steps (Graves, 2016), while early exits and REINFORCE-based methods also condition computation; this method uses attention-weight distances to decide whether another memory hop is needed.The paper says its approach also differs from prior external-memory work (Shen et al., 2017) in its REINFORCE loss and its attention-distance criterion for deciding whether retrieval has settled.Source: §3.2 · p. 6

3.3 GRAPH NEURAL NETWORKS

  • Unlike GNNs, which unroll message passing for a fixed number of steps, this method adapts the number of steps and lets input queries attend directly to memory slots.GNN message passing propagates node and edge embeddings, and is related to attention; self-attention can be viewed as a fully connected GNN.Source: §3.3 · p. 6

4 PAIRED ASSOCIATIVE INFERENCE TASK

The paired associative inference task probes whether a network can infer distant relationships by linking items distributed across multiple memories. It tests inference beyond retrieval of a directly experienced pair.

4 PAIRED ASSOCIATIVE INFERENCE TASK

  • The paper introduces paired associative inference (PAI), a neuroscience-derived task for probing reasoning across multiple memories (Bunsey & Eichenbaum, 1996; Banino et al., 2016).The task is presented as capturing appreciation of distant relationships among elements distributed across facts or memories.Source: §4 · pp. 6–7
  • Agents learn image pairs A–B and B–C in separate events, then choose between MATCH C and LURE C′ after cue A.The LURE was paired with B′, forming a different triplet from the A–B–C associations.Source: §4 · p. 7
  • Direct AB and BC queries test episodic retrieval, while indirect AC queries require inferring the connection through shared B.The correct C is linked to A because both were paired with B.Source: Figure 1 · p. 6, §4 · p. 7

5 RESULTS

MEMO solved the paper’s paired-association and shortest-path reasoning tasks, and matched prior models on the number of bAbI tasks solved with lower error. Its inference performance depended on combining separated fact representations with recurrent attention, while its adaptive retrieval used a variable number of hops.

5.1 BASELINES

5.2 PAIRED ASSOCIATIVE INFERENCE

  • For length-3 inference, MEMO converged to 3 hops while DNC required 10 pondering steps for the same accuracy, and MEMO’s attention shifted from the cue association to the match association across hops.The second hop also assigned some probability mass to the lure; a different MEMO instance using 7 hops showed a different activation pattern, indicating that the inference strategy depended on hop count.Source: §5.2 · p. 7, Figure 2 · p. 8
  • For inference, successful MEMO performance required both separated fact representations and recurrent attention; layer normalization in recurrent attention was critical for more stable training and better bAbI performance.Using the components individually was insufficient for inference, but direct queries could be solved with a single memory lookup; MEMO’s adaptive computation was more data-efficient than ACT on the paired-association task.Source: §5.2 · p. 8, §5.4 · p. 9
  • MEMO alone successfully answered the most complex length-4 and length-5 paired-association queries; on A-B-C, it tied DNC for highest accuracy while EMN and UT did not reach that level.The comparison concerns the hardest inference query for each task.Source: §5.2 · p. 7, Table 1 · p. 8

5.3 SHORTEST PATH ON RANDOMLY GENERATED GRAPHS

  • On 20-node graphs with path length 3, MEMO outperformed EMN by more than 50% and nearly solved the task like DNC; at out-degree 5, it beat DNC by more than 20% on both path nodes.On the smaller 10-node graph, DNC, UT, and MEMO had perfect accuracy predicting the intermediate shortest-path node; UT was slightly below MEMO when predicting the second node in the more complex setting.Source: §5.3 · pp. 8–9, Table 2 · p. 9

5.4 QUESTION ANSWERING ON THE BABI TASKS

6 CONCLUSIONS · A APPENDIX · A.1 PAIRED ASSOCIATIVE INFERENCE TASK

MEMO solved long-sequence associative-inference and graph-traversal tasks, matched state-of-the-art results on bAbI, and uses separated fact storage with recurrent attention. The appendix describes an ImageNet-based paired-associative-inference task with direct and inference-dependent queries.

6 CONCLUSIONS

  • MEMO achieved state-of-the-art results on paired associative inference, alone solved long sequences on that and graph traversal, and matched state-of-the-art results on all 20 bAbI tasks.The conclusion attributes these results to flexible weighting of individual memory elements.Source: §6 · p. 9
  • The authors attribute MEMO’s results to flexible weighting of individual memory elements through separated fact storage and a recurrent attention mechanism.Source: §6 · p. 9

A.1 PAIRED ASSOCIATIVE INFERENCE TASK

  • The task uses nonoverlapping ImageNet training, validation, and test images embedded with a pretrained ResNet, with sequence lengths of three, four, or five and 1e6 training, 1e5 evaluation, and 2e5 test images per dataset.Source: §A.1 · p. 13
  • For length-three sequences, each batch entry selects N = 16 sequences and stores all adjacent item pairs, yielding 32 memory rows.Source: §A.1 · p. 13
  • Queries test either direct episodic retrieval from one memory slot or indirect inference across episodes, with the network predicting the match’s class.Each query contains a cue, a match, and a lure drawn from another sequence in the same memory set.Source: §A.1 · p. 13
  • Match and lure positions are randomized, and half of each batch’s queries are direct and half indirect, requiring correct connections despite interference from other memories.Source: §A.1 · p. 13
  • Longer sequences yield indirect queries at different inference distances, requiring the model to integrate overlapping items across the sequence.The described examples include distances of one and four pairs.Source: §A.1 · p. 13

A.2 PAI - QUERIES WISE RESULTS … B.1 TRAINING DETAILS

The appendix reports query-wise paired-associative results and length-3 analyses, then details the shortest-path task’s training setup and model-specific reasoning procedures.

A.2 PAI - QUERIES WISE RESULTS

  • Query-wise PAI evaluation results are tabulated for A-B-C, A-B-C-D, and A-B-C-D-E; each evaluation set contains 600 items.EMN and DNC results use the hyperparameter with the lowest validation loss, while MEMO reports the average and relative standard deviation across its five lowest-validation-loss hyperparameters.Source: Table 4 · p. 14, Table 5 · p. 14, §A.2 · pp. 14–15

A.3 PAIRED ASSOCIATIVE INFERENCE FURTHER ANALYSIS ON LENGTH 3 EXPERIMENTS.

  • Figure 3 analyzes length-3 PAI inference accuracy and hop counts, plus accuracy distributions for direct and inference queries across 100 hyperparameters and seeds.The direct-query distribution averages A-B and B-C accuracy; the inference-query distribution uses A-C.Source: Figure 3 · p. 15

A.3.1 ABLATIONS

  • The length-3 PAI ablation table reports best-validation-run results, with symbols distinguishing absent and present components.The supplied passage does not include the table’s cell values.Source: Table 7 · p. 16, §A.3.1 · p. 16

A.3.2 ATTENTION WEIGHTS ANALYSIS

  • In a seven-hop length-3 PAI run, MEMO uses hops one and two to retrieve the cue slot, then hops three through five to retrieve the matching slot, with sharp attention on one slot.The figure describes this attention pattern for a case that converged to seven hops.Source: Figure 4 · p. 16

A.3.3 ADAPTIVE COMPUTATION

  • On length-3 PAI, MEMO with REINFORCE is more data-efficient than MEMO with ACT.The comparison is presented in Figure 5.Source: Figure 5 · p. 17

B.1 TRAINING DETAILS

  • Shortest-path graphs connect sampled two-dimensional points to their K nearest neighbors, and graph edges, endpoint queries, and target paths are represented as node tuples or sequences.K is sampled independently for each node; target sequences list the nodes on the path between the query endpoints.Source: §B.1 · p. 17
  • Models train for 2e4 epochs of 100 batch updates, using 64-graph minibatches and Adam cross-entropy; evaluation reports mean node accuracy and standard deviation over 600 examples and the five best hyperparameters.The evaluation score averages accuracy over all nodes in each target path.Source: §B.1 · pp. 17–18
  • EMN uses the ground-truth first node as the next query, whereas MEMO feeds forward its prediction and cannot exceed chance on node two after an incorrect first answer; DNC and UT retain a global view.Answer weights are not shared across successive predictions.Source: §B.1 · pp. 17–18

B.2 RESULTS · C BABI · C.1 TRAINING AND EVALUATION DETAILS

On the shortest-path task, MEMO performed better when the first node’s ground-truth answer was used to query the second node than when its prediction was used; the accompanying bAbI subsection specifies training and evaluation procedures.

B.2 RESULTS

  • In 20-node graphs with 5 outbound edges, MEMO scored 85.38%(0.05) with the ground-truth first-node answer versus 69.20%(0.07) with its prediction; EMN scored 22.30% when using a prediction.The passage says the same pattern was confirmed in the simpler 20-node, 3-outbound-edge scenario.Source: Table 8 · p. 18, §B.2 · p. 18

C.1 TRAINING AND EVALUATION DETAILS

  • For bAbI, the authors used the English Question Answer dataset (Weston et al. (2015)), lowercased text, ignored periods and question marks, separated words by spaces, and extracted questions as queries.Stories were the text before each question.Source: §C.1 · p. 19
  • Commas were retained in answers, so distinct path labels such as ’n,s’ and ’n,w’ remained separate, and each query-and-stories input had one answer label.Source: §C.1 · p. 19
  • Training sampled 128 queries and their stories, represented as 128 × 11 query tokens and 128 × 320 × 11 story tokens, with zero padding to the maximum sizes.Source: §C.1 · p. 19
  • EMN and MEMO received stories and queries as their architecture specifies; DNC processed them in sequence with blank pondering inputs, while UT used its encoder output.DNC and UT inputs were embedded as for MEMO.Source: §C.1 · p. 19
  • All models used Adam and trained for 2e4 epochs of 100 batch updates; MEMO also received a time-encoding column in its memory store.The time encoding was derived from Vaswani et al. (2017); training stopped after a fixed number of time-steps specified in Appendix Section D.2.Source: §C.1 · p. 19
  • Evaluation sampled 10,000 examples, computed mean accuracy and accuracy on each of bAbI’s 20 tasks, and reported averages and standard deviations across the best five hyperparameters.Source: §C.1 · p. 19

C.2 BABI ABLATIONS … D.1 TRAINING DETAILS

The supplied section records bAbI ablation and task-wise evaluation summaries, then specifies MEMO’s optimization, prediction targets, and halting-policy training.

C.2 BABI ABLATIONS

  • The bAbI ablation lists five configurations with paired values from 11/20 and 14.81 to 20/20 and 0.21.The extracted configuration markers are unclear, so their individual settings cannot be identified.Source: §C.2 · p. 20, Table 9 · p. 20
  • The bAbI ablation’s best run was selected by validation set and trained and tested jointly on all tasks with approximately equal training resources per task.Source: §C.2 · p. 20

C.3 TASK-WISE RESULTS

  • Task-wise bAbI results summarize test-error means and standard deviations over the five hyperparameters with lowest validation loss.Source: Table 10 · p. 21, §C.3 · p. 21

D.1 TRAINING DETAILS

  • MEMO uses Adam with polynomial learning-rate decay and batch sizes of 64 for PAI and shortest path and 128 for bAbI.The decay starts at lstartmemo.Source: §D.1 · p. 21
  • Across PAI, shortest path, and bAbI, MEMO is trained with cross-entropy loss to predict a class ID, node ID, and word ID, respectively.Source: §D.1 · p. 21
  • The halting-policy network parameters are updated with RMSProp using learning rate lhalt.Source: §D.1 · p. 21
  • Other MEMO parameters are reported in Tables 11 and 12.Source: §D.1 · p. 21

D.2 FIXED HYPER-PARAMETERS USED ACROSS TASKS · D.3 RANGE OF HYPER-PARAMETERS USED IN SWEEPS

The appendices identify hyper-parameters fixed across tasks and the ranges examined in sweeps.

D.2 FIXED HYPER-PARAMETERS USED ACROSS TASKS

  • Table 11 presents the fixed hyper-parameters used across tasks.Source: Table 11 · p. 22

D.3 RANGE OF HYPER-PARAMETERS USED IN SWEEPS

  • Table 12 presents the hyper-parameter ranges used in sweeps.Source: Table 12 · p. 22

E MEMO COMPLEXITY ANALYSIS · F ACT DESCRIPTION

MEMO’s temporal and spatial complexity depend on input size through the number of sentences and the size of its memory, respectively. The ACT description adapts the halting unit to MEMO’s policy and defines a halting step, remainder, and final answer from per-hop answers.

E MEMO COMPLEXITY ANALYSIS

  • MEMO has temporal complexity O(ns · A · N · H · I · S · d), with A, N, H, I, S, and d fixed constants in the experiments.The factors represent samples, answers, hop bound, heads, stories, words per sentence, and embedding size, respectively.Source: E MEMO COMPLEXITY ANALYSIS · p. 23
  • MEMO is linear in the number of input sentences, whereas the Universal Transformer has quadratic complexity.Source: E MEMO COMPLEXITY ANALYSIS · p. 23
  • MEMO’s spatial complexity is O(I · S · d), the size of its memory, and this size is fixed in the experiments.Weights are constant; only query-dependent context information affects spatial complexity.Source: E MEMO COMPLEXITY ANALYSIS · p. 23

F ACT DESCRIPTION

  • MEMO+ACT follows Graves (2016), fixes ϵ at 0.01, defines a halting step T and remainder R, and specifies a final answer using MEMO’s answers at each hop.The supplied passage identifies a_t as the answer MEMO provides at hop t, but does not include the final-answer formula.Source: F ACT DESCRIPTION · p. 23
  • MEMO+ACT defines its halting unit as h_t = σ(π_t), using MEMO’s binary policy rather than ACT’s trainable linear projection of the previous state.The original ACT unit is given as h_t = σ(W_h s_t + b_h).Source: F ACT DESCRIPTION · p. 23
  • The authors argue this halting-unit change makes comparison fairer because MEMO’s policy depends on state through several nonlinearities and resembles their model.They also say the change still permits evaluation of the feasibility of the halting mechanism.Source: F ACT DESCRIPTION · p. 23

G DNC ARCHITECTURE AND HYPERPARAMETERS · H UNIVERSAL TRANSFORMER ARCHITECTURE AND HYPERPARAMETERS

The DNC and Universal Transformer sections specify the architectures and hyperparameter configurations used, and report the ranges searched for each model. Their tables identify the layer sizes or hyperparameter settings and search ranges, but the supplied passages do not give the table values.

G DNC ARCHITECTURE AND HYPERPARAMETERS

H UNIVERSAL TRANSFORMER ARCHITECTURE AND HYPERPARAMETERS

Related papers

Loading 2001.10913v1…