Learning Associative Inference Using Fast Weight Memory
Authors: Imanol Schlag
Tsendsuren Munkhdalai
Jürgen Schmidhuber
Venue: ICLR 2021 Poster (2021)
Keywords: memory-augmented neural networks
tensor product
fast weights

Abstract: Humans can quickly associate stimuli to solve problems in novel contexts. Our novel neural network model learns state representations of facts that can be composed to perform such associative inference. To this end, we augment the LSTM model with an associative memory, dubbed \textit{Fast Weight Memory} (FWM). Through differentiable operations at every step of a given input sequence, the LSTM \textit{updates and maintains} compositional associations stored in the rapidly changing FWM weights. Our model is trained end-to-end by gradient descent and yields excellent performance on compositional language reasoning problems, meta-reinforcement-learning for POMDPs, and small-scale word-level language modelling.

Humans can quickly associate stimuli to solve problems in novel contexts. Our novel neural network model learns state representations of facts that can be composed perform such associative inference. To this end, we augment the LSTM model with an associative memory, dubbed Fast Weight Memory (FWM). Through differentiable operations at every step a given input sequence, the LSTM updates and maintuins compositional associations stored in the rapidly changing FWM weights. Our model is trained end-to-end by gradient descent and yields excellent performance on compositional language reasoning problems, for POMDPs, and small-scale word-level language Published as a conference paper at ICLR 2021
LEARNING ASSOCIATIVE INFERENCE USING FAST WEIGHT MEMORY
Imanol Schlag
Tsendsuren Munkhdalai
The Swiss Lab IDSIA/ / /SI/ SUPSI imanol 1 @ Ridsia ch
Microsoft Research
tsendsuren munkhda l @mi crosoft.com
Jürgen Schmidhuber
The Swiss AI Lab IDSIA / USI/ SUPSI juergeneissia ch
ABSTRACT
modelling!
INTRODUCTION
Humans continually adapt in order to understand new situations changing environments. One important adaptive ability is associative inference for composing features extracted from distinct experiences and relating them to each other (Schlichting & Preston 2015 Gershman et al. 2015). Suppose Alice has shared with you pictures of her toddler. Later, at the office party, you see man carrying the depicted toddler. Since the toddler yields a shared feature two different contexts, it may be plausible to infer that the man is Alice's partner, without ever sceing him and Alice to- gether. The ability to rapidly associate and bind together novel stimuli can help to derive knowledge systematically, in addition to the knowledge gained directly from observation.
Virtually all modern cognitive architectures applied to challenging artificial intelligence problems based on deep artificial neural networks (NNs). Despite their empirical successes and theoretical generality, NNs tend to struggle to generalise in situations similar to the given example (Lake et al.) 2017 Phillips 1995 Lake & Baroni 2017). This weakness becomes even more severe if the training and test data exhibit systematic differences (Atzmon et al. 2016 Agrawal et al. 2017). For example, during training, the man's 's representation might never be associated with the toddler's. but during testing, this association might be necessary to make useful prediction. In problems where humans excel, this sort of inference is likely ubiquitous since data is often combinatorially complex a way that observations used during training will likely cover just a small fraction of all possible compositions. Such a lack of productivity and systematicity is a long-standing argument against the use of NNs substrate of an artificial cognitive architecture (Fodor & Pylyshyn 1988 Hadley 1994 McLaughlin 2009).
The hidden state of neural model is a learned representation of the task-relevant information tracted from the input. To generalise to never-seen-befowd compositions of stimuli, the function which produces the state representation must be able to systematically construct all possible states. This requires a general and preferrably differentiable method, such as the Tensor Product Repre- sentation (TPR: Smolensky (1990)) TPRs provide a general and differentiable method for embed-
Source code and data used in this paper is available at alggiiub.cmm/sshhggaatttaaeeee wwwphhhhueeaaatteiieeeeeerrrrrrr

---

Published as a conference paper at ICLR 2021
ding symbolic structures in vector spaces. A TPR state representation is constructed via the tensor product (i.e. the generalised outer- puter-product) of learned component representations. Under certain constraints, such mechanism guarantees unique representation for every possible combination of components (Smolensky 1990 2012).
In this work, we augment recurrent NN (RNN) with additional TPR-like memory representa- tion. To facilitate the learning of multi-step associative inference, the TPR memory can be queried multiple times in a row, allowing the model chain together various independent associations. In contrast to previous work on fast weights, we apply our memory-augment sequences. This requires the model to update its associative memory. Furthermore, we demonstrate the generality of our method by applying it to meta-reinforremmnt learning and small scale language modelling problems.
RNN to much longer
In the next section, we cover related memory-augma auggeeneed NNs. Section describes the FWM in detail. Section demonstrates the generality of our method through experiments in the supervised, self-supervised, and meta-reinforceemet learning setting. The supervised-leaminng experiments in subsection consisss consist of or catbAbl. The meta-reinforcement learning experiment in section 4.2 demonstrates the FWM's ability to learn to explore a partially observable environment through its ability to perform associa- tive inference. Finally, the self- self-supervised experiments in subsection 4.3
more challenging version of the bAbl dataset dubbed concatenated-ba
demonstrate that the FWM can compete with the state-of-the word-level language models on small benchmark datasets.
RELATED WORK
RNNs such as the Long Short-Term Memory (LSTM; Hochreiter & Schmidhuber (1997); y Gers ct al. (2000) are in theory capable of implementing any algorithm (Siegelmann & Sontag 1991). However, the linear growth of the hidden state of fully connected RNN leads quadratic growth in the number of trainable weights. Early work addressed this issue through the use of additional memory (Das et al. 1992, Mozer& & Das 1993) and differentiable fast weights (Schmidhuber) 1992 NNs have solved algorithmic toy problems (Graves et al,, 2014 2016) as well as reasoning and inference problems in synthetic and natural language (Weston et al.
1993). Recently,
2015b Xiong et al. 2016).
Inspired by the random-access access memory of computer architectures. a common approach is incor- porate a soft and differentiable lookup table into the NN model. Such slot-based memory matrices have shown be difficult to train (Munkhdalai & Yu 2017b) and require sophisticated mechanisms for the allocation and deallocation of memory (Csordas & Schmidhuber, (011)). The Transformer- Dai
XL (TXL; Dai et al. 2019)), an autoregressive language model variant of the Transformer [Vaswani ct al. 2017), can be understood pushed into an immutable queue of finite size. Although it is recurrent, the layers of a transformer architecture
slot-based memory-aaugmented RNN where every new state is
strictly forced to use inputs from a lower layer which limits its generality. Neverthe- less, a sufficiently deep and well regularised TXL model has achieved state-of-the-arr performance in large scale language modelling tasks.
A biologically more plausible alternative of increasing the memory capacity of NNs are fast- changing weights, i.e. stateful weights that can adapt as a function of its input. Non-differentiable fast weights or "dynamic links" have been published since 1981 (von der Malsburg 1981 Feldman
1982 Hinton & Plaut 1987). Subsequent work showed that a regular network can be trained by
gradient descent to control the fast weights of separate network (Schmidhuber 1992) or of it- self (Schmidhuber 1993) in an end-to-end differentiable fashion. Recently, fast weights have made comeback and achieved good results in small toy problems where regular NNs fall short (Ba et al.) 2016a [chlaa & Schmidhuber) 2017] Munkhdalai & Yu! 201 Pritzel et al. 2017 Ha et 2017
Zhang & Zhou 2017 Miconi et al. 2018 2019 SSchlag & Schmidhuber 2018 18 Munkhdalai et al. 2019 Bartunov et al. 2020).
Most memory-augmenteed NNs are based on content-based or key-based lookup mechanisms. An alternative to the storage of patterns in a lookup table is the idea that patterns are reconstructed through the implicit iterative minimisation of an energy function, such in the classical Hopfield network (Steinbuch] [961] Willshaw et al. 1969 HHofield
1982] Kanerva, 1988) or the modern
Hopfield network (Krotov & Hopfield 2016 Demircigil al. 2017 Ramsauer et al. 2020). This is
2

---

Published as a conference paper at ICLR 2021
often described an auto-associatiive type of memory as it reconstructs a previously stored pattern that mostly resembles the current pattern. A much less studied variation is the hetero associatve memory (see e.g. Kosko (1988) where the retrieved pattern is different from the imput pattem. This is more relevant for our use case. We aim to train an LSTT to construct, maintain, and edit its associative memory. The ability edit Hopfield networks partially is not very well studied. For this reason, we employ a simple (multi-)linear hetero-assoose associavvie memory as it is more closely related to the theory of TPRs (whose manipulation is well understood) and because the association is retrieved in single step.
Our work directly builds on two examples of differentiable fast weight memories: the -PR-RNN by [Schlag& & Schmidhuber (2018) and the Metalearned Neural Memory (MNM) by Munkhdalai jet al. (2019). The TPR-RNN is a sentence-level model for reasoning on text. It achieves excellent results on the regular bAbl tasks but it underperforms on word-level bAbl (Schlag et al. 2019) or algorithmic toy problems (Le et al., 2020). In contrast, the MNM word-level model which augments the LSTM with a fully-connected multi-layer feed-forward network its memory and trains it using a meta-learning objective. Both, MNM and TPR-RNN were developed on the regular bAbl dataset which only contains short sequences and does not require the model to remove depre- cated associations from its memory. In this work, we train on an infinite sequence of bAbl stories where our FWM achieves excellent performance and improves over MNM. We further demonstrate strong performance in ssall-scale language modelling and meta reinforcement-leaniin which demonstrates the generality of our contribution.
3 PROPOSED METHOD
Our FWM is a fast-changing, LSTM. The slow weights of the LSTM are regular NN weights which are updated during training by gradient descent. In contrast, the fast weights of the FWM are updated by the LSTM at every step of the imput sequence through Hebb-like differentiable mechanism. This allows the FWM function to change rapidly even during testing-henne weights, the LSTM also generates a memory query which is used to retrieve information that was previously stored. The retrieved information then becomes part of the model's output.
map which is controlled by a slowly- non-linear
the name fast weights. Along with updating the fast
3.1 THE FAST WEIGHT MEMORY
Given a sequence of tokens x = (#1...., Fr- from vocabulary V, the task of language modelling is to train a model which maximizes the joint probability p(x) which we factorize autoregressively p(xiT)
XN,
F
FWM
F
II (1) where To is an artificial start to- ken In this work, we train an RNN model to encode the input sequence X141 into h2 the hidden state of the LSTM, and F, the fast weight tensor of the FWM, to maximize the probability of the next token T++1-
LSTM
h
A D
At step t of the input sequence, the input to- ken 2; is embedded in dg-dimensional vector space using a lookup table e, = cmbedding(az). An LSTM with diSTM hidden units encodes the sequence of embedded tokens into a fixed size vector representation hr
Figure simplified illustration of our proposed method where A refers to the write mechanism described in section 3.1.I F. are the recurrent weights of the FWM which have been generated by the LSTM. The LSTM is a regular slow RNN. The residual connection between the FWM and the LSTM is not depicted.
LSTM(e,, h; h(-1). The probability distribution over the next token 4.+1 = Noftmax(W((s) (h+ + FWM(h, Ft)) where F. E x are the fast weights of the FWM at step t and ((+) E R|VixdISMM, Note that the fast weight matrix F, is a reshaped third-order tensor F; E Redona x dews x dewn, This allows us to describe third-order tensor operations using matrix multiplications. We'll now describe in detail the FWM function and how its fast weights are updated.
We use the notation X1: to refer to the sequence (01, Tt).
3

---

Published as a conference paper at ICLR 2021
3.1.1 WRITING
The FWM is updated every step using the write mechanism described in this section. To this end, we extract from the hidden state he: the write strength 8 (a scalar bounded by and I using the sigmoid function σ), the two key vectors k1, k2, and the new value v.
kk. k2, v) = 6 W ww) 3 = a WWhh))
(1)
(2)
The purpose of writing to memory is to learn context-specific association between the input pattern k @ l22 and the output pattern v. The usage of the tensor-pie munsoor-rrouctttcc the input pattern factorises the the representational space which guarantees unique orthogonal vector representations for novel key pairs. A specific example of such is given and demonstrated by |Shhlag & Schmidhuber (2018) where the first key learns to represent entity and the second key specific action, thereby, learning representational space that generalises to never seen entity and action compositions.
In stark contrast the complex memory operations of the TPR-RNE and word-level operation which is closely related to the perceptron learning rule (Rosenblatt) 1958). It allows the model to replace the previous association Vold with convex combination of the old and new value f0 + (1 B) Vold- With the scalar ,3 the LSTT controls the new association fully replaces the previous value (3 or if the information of both are mixed together. Our fast weight update works as follows: First, the current value Vold that is associated with k @ k2 is retrieved. Second, we remove the old association from the map by subtracting vec(k1 & k22 @ Val from our memory, where vec vectorises the matrix. Third, we add vec(k) & k2) @ 23v++ (1 V)) ). All three steps can be achieved at once using the following update rule (see appendix section B for the proof):
employ a single, simple,
F F 1 3 vec(k1 @ k2) & (v Vola).
To prevent the fast weights from potentially growing endlessly, we scale down the fast weights whenever || 1. This is achieved through the following clement- scaling.
F
F m1 F
(4)
3.1.2 READING
For each step of the input sequence, the model queries the memory in order to retrieve previously stored value. Due to the keys and values being generated separately, the network can retrieve values which are informationally independent from their keys. In order to perform more complex associa- tive inference, like c.g. transitive inference (a 6, b + c, therefore, reads where we use the retrieved value as one of the keys in the next query (see equation 7).
c), employ multiple
(0) = (WW,)))
(5)
e,
(1) o(W! h.), 1is,,,, 6 1) 2 siissM
(6)
n,
(1) LN(F, vec(n,
(7)
FWM(h, F2) W.n (N.)
(8)
Here LN refers to layernorm without the learned element-wise affine map (Ba et al. 2016b), vec reshapes the matrix into a vector, @ is the hyperbolic tangent function, and the matrices W., W! E RddwsscoLSsI, i E {1.NN,} and W. IR x are regular slow weights trained by gradient de- scent which allow us to decouple the dimensionality of the LSTM from the dimensionality of the FWM. In eq. 7 F; is the multi-linear map which we query using the LSSTM-genertte "input" e(i) and the previous retrieval (5-1) (except for the first query where both keys are L.SM--genraate)).
EXPERIMENTS
4.1 CONCATENATED-BABI
The bAbl tasks is a popular toy dataset to benchmark neural networks with memory augmenta- tions and reasoning capabilities (Weston al. 2015a). It consists of set of short stories with
4

---

Published as a conference paper at ICLR 2021
questions embedded in the text. The stories were generated by simulating multiple entities in vir- tual environment and cover different contexts in which entities change their state on their own or through an interaction. Each story- belongs to one of 20 different tasks that the authors of the dataset considered important for intelligent dialogue agents. The tasks contain questions which re- quire reasoning capabilities like deduction, coreference, or counting. All tasks require some level of symbolic reasoning, and the first neural and non-neural baselines demonstrated poor generalisation performance on test data (Weston al. 2015a).
We aim to improve the bAbl benchmark as a means of developing intelligent dialogue agents. To this end, we propose concatenated bAbl (catbAbl); an infinite sequence of bAbl stories. catbAbl is generated from the bAbI dataset and during training, a random sample/story from any task is drawn without replacement and concatenated to the ongoing story. The preprocessing for catbAbl ad- dresses several issues: it removes the supporting facts, leaves the questions embedded in the story, inserts the correct answer after the question mark, and tokenises the full sample into a single se- quence of words. As such, catbAbl is designed to be trained in autoregressive way and analogous to closed-book question answering.
catbAbl models can be trained in two different ways: language modelling mode (LM-mode) or cression-answeringg mode (QA-mode). In LM-mode, the catbAbl models are trained like autoregressive word-level language models. In QA-mode, the catbAbl models are only trained to predict the tokens that are answers to questions makinn it more similar to regular bAbl. QA-mode is simply implemented by masking out losses non-answer predictions. In both training modes, the model performance is solely measured by its accuracy and perplexity when answering the questions. Performance on non-answers is irrelevant on catbAbl because the tokens are either very predictive or inherently unpredictable, and there is nothing appealing to be learned. Despite measuring performance only for answers, we argue that LM-mode is interesting for three reasons. First, -mooe removes the bias of knowing which words would benefit from symbolic inference mechanism. Second, L.-modee trains the model on sequence with tokens which are inherently unpredictable. Such tokens could also appear in natural language and might harm the model's ability to learn a useful representation of the story. Indeed, in the next section, we will give evidence for such generalisation gap. Third, the LM-mode setting allows us to directly compare our method with state-of sfate-of-the-art the-art language models.
4.1.1 RESULTS
We compare our FWM directly with the current state-of-the-art on word-level bAbl: Metalearned Neural Memory (MNM; Munkhdalai et al. (.19)) We also include two strong autoregressive word-level language models as baselines: a regularized LSTM (Merity et al. 2018 Melis et al.) 2017) and a regularized Transformer-X XL (TXL; Dai et al. (2019)). Lastly, we also evaluate Ba's Fast Weights which attend to the recent past (JBFW; Ba et al. (2016a)) but were unable to find hyperparameters that converged. We truncate backpropagation through time (tBPTT) to 200 tokens for all models and limited the amount of GPU memory to -16GB for practical reasons. For every model, we performed a hyperparameter search in QA mode over the first 3k steps of which smaller selection was trained for 30-60k steps. For all models,
adopt the best QA mode hyperparameters for the LM mode results. Table [ lists the best accuracy and perplexity of each model over three seeds while figure 2 shows the learning curves of the best seeds. Further hyperparameter search results can be found in the appendix section F
Our experiments on catbAbl show that regularized, 4-layer deep, and residual LSTM, and a 3- layer deep TXL with attention over the last 1400 tokens, achieve strong performance on catbAbl. MNM. on the other hand, suffered a ~10% drop in QA mode accuracy compared to its performance on bAbl which demonstrates the increased difficulty of catbAbl. The JBFW model is not able to make meaningful predictions on catbAbl which may be due to its inability of removing previous associations and fixed fast weight memory decay. Our FWM achieves an excellent accuracy on catbAbl while being by far the smallest in parameter count and weight to activation ratio. The performance gap between FWM and MNM suggests the importance of our fast weight memory mechanism. In figure 3 we visualise how the FWM can chain memories from different points in time to perform transitive inference.
We chose to include the TXL model in our comparison due to its autoregressive nature and strong performance in large-scale language modelling benchmarks. However, we point out that the TXLs
5

---

Published as a conference paper at ICLR 2021
Table Accuracy and perplexity on test data over three seeds each model's best hyperparameters setting according to our hyperparameter search. Detailed hyperparameters and results can be found the appendix section F
Mode
JBFW
LSTM
TXI.
MNM
FWM
QA
13.22 + 0.0
80.88% + 0.30
87.66% + 2.82
88.97% + 6.28
96.75% + 0.05
QA ppl
31.19 + 8.8
1.93 + 0.11
1.50 + 0.14
2.50 + 1.07
1.36 + 0.06
LM acc
0.0 0.0
80.15% + 0.40
90.23% + 1.01
69.30 % = 16.60
93.04% % + 0.62
LM ppl
160.3 24.3
1.84 0.02
1.39 + 0.03
2.60 1.02
1.45 + 0.14
weights
548k
8M
10.5M
1.1M
694k
activations
263k
4096
4.3M
30.5k
33.3k
1.000
0.975
0.950
0.925
0.900
0.875
0.650
0.825
0.800
accuraey
0.775
0.750
0.700
0.675
0.650
0.625
0.600
0.575
Transformer-XI
0.550
FWM
0.525
0.500
10000
20000
30000
40000
50000
60000
step
Figure 2: QM model validation accuracy the best-over-all seeds of each model over training steps.
context window is larger than the average bAbl story. In this case, due to the shortness of the stories, catbAbl becomes more of an open-book problem for the TXL model since it has the capability of looking up representations of its previous input whereas the RNN models do not. This fundamentally limits the TXL model as it can only condition its prediction on information that is no longer than its attention window to past states. The RNN models, which are general function approximators, for better or for worse, are instead forced to learn to carry the necessary information through time.
4.2 CETTTA-RENOOR LEARNING
Meta reinforcement learning (Meta-RI.) applies meta-learning (Schmidhuber) 1987 Hochreiter al. [200T] Finn al. 2017) the field of reinforcement learning (Schmidhuber 1994). An agent is trained on multiple environments (or tasks) and receives environmental feedback as part of its input. maximise its total reward in an environment, the agent has to leverage the feedback signals and adapt. A successful agent is capable of maximising its reward in novel environments that it has been exposed during training. Recent work achieved notable progress this domain (Santoro et al. 2016 Mishra al. 2018; Kirsch et al. 2020). We experiment with tasks drawn randomly from large set of partially observable Markov decision processes (POMDPs). this set, every environment consists of precisely five states and three actions. Globally, every environment can be viewed as a sparse directed graph where nodes are locations, and the directed edges are one-way modes of transportation- similar metro transit map of city (Graves et al. 2016). To generate
Bigger JBFW models did not improve performance, See appendix section F.5
'The number of immutable activations is 512 2 x 3 x (1200 + 199) while the number of mutable activations merely 512 2 3 3072. Only the TXL. model maintains immutable activations.

---

Published as a conference paper at ICLR 2021
-200
-100
100
2
3
zue
Playe
ro
asnow
aip
Preueo
sangen
daays
Pieye
ashou
sanyos
aue
Braye
daays
Anua
St
daays
st seya
Spripa Spripe
Bieye
vo
Figure 3: A visualisation of the FWMs ability to chain independent associations to perform transitive reasoning on the catbAbl validation data. The colour of each grid cells represent the dot product refers to
(k1 @ k2 @ e) where k1, k2 are the write keys of each previous position while the respective queries generated at "?" (second position from the right) for each of the N, = 3 memory reads. The first query matches most with the keys at the recent positions where the input was "gertrude" and "afraid" (first row of grid cells). The second query, which partially consists of the value retrieved from the first query, matches with the "getrude is a mouse" section. The third query, which partially consists of the value retrieved from the second query, matches with the "mice are afraid of wolves" section. Finally, the FWM correctly outputs the next word and answer to the question: wolf (not seen). This likely completes the deduction: gertrude is mouse, mice are afraid of wolves, therefore, gertrude is afraid of wolves.
(1.)
(0.)
9
(1,
2)
3
(0.)
(0,
(2.)
(2.)
(1.)
1)
Figure 4: Two randomly generated environments with the agent's location coloured in green and the reward location coloured in yellow. Edge labels indicate the set of valid actions (0, 1, or 2) to transition along that arrow. Invalid actions are not visualised. The graph and the locations of the agent and reward are set randomly at the beginning of the experiment. If the agent reaches the reward location or did not reach it after six steps, both are randomly reset.
a new environment, we sample the adjacency matrix of the graph such that actions are deterministic, and every location is reachable from any other location (see figure 4 We sample graphs such that there are no actions that lead to the same location, and such that not every action is always valid way of transitioning. We added the exact algorithm to generate graphs, as well as further details, to the appendix section[]
The agent's goal is to reach the reward location. Upon arrival, the agent receives the reward, followed by a random reset of the agent's and reward's location. Whenever the agent takes an action that does not lead to a new location, receives penalty. At every step, the agent receives as an input: its current location, the reward location, its last action, and the reward received so far.
We run our experiment for 30 steps and compare our FWM to an LSTM baseline. Both methods are trained on the same training set of 600 graphs and tested on 600 novel graphs. We optimise our agent with the Advantage
Crittc (A2C) algorithm, a non- non-asynchrooous version of the A3C method (Mnih et al.) 2016). In our experiments, the SSMM-baee agent requires more episodes, bigger network, and eventually overfit to the training graphs. The WWM-based agent however trains faster and generalises to randomly sampled graphs.
7

---

Published as a conference paper at ICLR 2021
FATM 32 16
train environments
test environments
LSTM-512 (x73)
LSTM-1026 (x2C7)
LSTM-2048 (x1207)
LSTM 4096 (x4614)
340
140
369 tetal reward
120
120
300
100
60
60
40
40
500
1000
1500
2000
2500
500
1000
1500
2000
2500
steps 50
steps 50
Figure 5: Average total reward of the agent when trained on 600 random graphs (left plot) and tested on 600 different graphs (right plot). The FWM agent (blue) has slow LSTM with 32 hidden units and fast weight memory of size 16 x 162. We compare to LSTM agents with different sized hidden states. The largest LSTM has 4096 hidden units (red) which roughly matches the number of temporal variables of the FWM. The FWM has 14k trainable weights which is by far the lowest. The largest LSTM has 67. 37.4M weights which roughly 4814 times more than the FWM. The relative factor of each LSTM is added the legend. All LSTMs take longer to train and eventually overfit on the training data. Due to the overfitting, the LSTM does not have to explore, which results in a higher total reward training environments but lower total reward test environments.
We argue that the bAbl stories and the episodes on the graphs are similar in the following three ways. First, in both problems, the network has to construct useful and context-specific representation from its ongoing input. Second, as part of its input, the network repeatedly receives an objective (the reward location versus the question) which requires the exploitation of the context-specific information. Third, the model has to produce a discrete sequence (actions in the environment in RL and reasoning steps in catbAbl) to optimise its training signal (high reward versus low uncertainty).
4.3 LANGUAGE MODELLING
Comparing FWM autoregressive language models catbAbl begs the question: how does FWM perform on popular word-level language modelling datasets such as Penn Treebank (PTB; Mikolov al. (2010) or WikiText-2 (WT2: Merity et al. (2117)) It is unclear to which extend a symbolic inference mechanism is beneficial for language modelling. PTB and WT2 contain virtually no ques- tions and are constructed from Wikipedia and news articles which are designed be easily parsed by the reader. Nevertheless, in figure 6 we show how our FWM exploits recurring subject names to reduce its uncertainty. Not many memory augmented NNs have been able to bridge from small and toy reasoning tasks to general language models-an and those which did, underperformed (Paperno) ctt al., 2016; Sukhbaatar et al. 2015). We use the regularized 3-layer AWD-LSTM (Merity et al.
Table 2: Best perplexity on the test data of Penn Treebank (PTB) and WikiText-2 (WT2) from three seeds. Detailed results can be found in the appendix table All PTB models have roughly 24M parameters and all WT2 models have roughly 37M parameters. The AWD-TXL. is the Transformer-XI. architecture as reported by Dai al. (2019) with the necessary AWD-style regularisation, model averaging, and softmax temperature tuning (see appendix section
PTB
WT2
Model
Validation
Test
Validation
Test
AWD-LSTM (Merity et al. 2018)
60.0
57.3
68.6
65.8
AWD-TXL Daa et al. 2019
54.52
AWD-TXL (ours)
59.39
56.50
65.73
63.11
AWD-FWM (ours)
56.76
54.48
63.98
61.65
8

---

Published as a conference paper at ICLR 2021
dfr
Moraago
in
-3g
Figure 6: Loss comparison between the LSTM and our FWM on a section of the PTB test set. The colour of the grid cells in the first row stands for the cross-entropy error of the LSTM and FWM model. The second row, for their respective difference. Our FWM sometimes shows a lower error on rare subject words such as names of companies and people once they have been introduced. As seen in the red circles, the initial mentioning of "phillips" has similar uncertainty between the LSTM and FWM but shortly after that the subject of the sentences is more predictable and the FWM is more certain (4.3 bits difference) whereas the LSTM's uncertainty remains roughly on the same level (12.8 bits).
2018 as the slow RNN in our FWM model to minimize further hyperparameter search. The experi- mental results in table 2 demonstrate a relative improvement over the AWD-LSTM baselines, which suggest the benefit of our FWM even in language modelling benchmarks. However, in contrast to catbAbl, all three models achieve very similar results which might indicate that PTB and WT2 do not benefit strongly from associative reasoning capacity. We added the experimental details to the appendix section
Since the publication of AWD-LSTM (Merity al. 2018 various extensions (some which are orthogonal to our memory augmentation) have been proposed (Krause et al. 2018 Merity et al.) 2018 Yang et al., 2018). In this work, we are not primarily interested in beating the state-of-the thethaarrt in language modelling and leave it for future work to explore the possible synergies between these methods.
DISCUSSION
An order-three memory tensor is a computationally demanding method for constructing compo- sitional state representations. With vector components in R", the tensor product computation alone has space and time complexity of O(n)). For practical reasons, this forces the FWM to remain small, relative to the slow NN, which limits the number of associations that can be maintained at once. Previous work has proposed approximations of such memory tensors in variance-optimal way (Schlag et al., 2019). In our ablation experiments in section we show on catbAbl that concatenating the keys results in performance accuracy drop of -5% We also experiment with fewer read operations (smaller N,) which also results in a performance degradation (appendix figure 77. However, further improvements might not come from scaling up but from more general symbolic manipulations. We address the capacity of the FWM and the necessity of the tensor product from linear hetero-associative memory perspective in section A of the appendix. Finally, our fast weight memory can be thought of as primitive "working memory" of the analogous the working memory in the human brain (Spalding al. 2018). This idea is supported by recent work which proposes cognitive model of the human brain that is based on such higher-order tensors (Tresp & Ma 2017).
CONCLUSION
Our new FWM is fast weights architecture capable of learning from synthetic data to answer questions which require various symbolic reasoning skills. To improve generality, we overcome issues of the popular bAbl dataset by introducing more general and more difficult variation dubbed catbAbl. We report excellent performance on catbAbl and compare with strong baselines based on state-of-the-art language models, as well as, the previous state-of-the-art in word-level bAbl. We also apply the FWM in a challenging mett-reiinooremmnnt learning environment where the agent generalises to novel environments by learning from its observations and actions. Finally, in a self- supervised setting, we apply the FWM to word-level language modelling on PTB and WT2 where beats the AWD-LSTM and AWD-Transformer: XL. baselines.

---

Published as a conference paper at ICLR 2021
ACKNOWLEDGEMENTS
We thank NVIDIA Corporation for donating several DGX machines, and IBM for donating Minsky machine. This research was supported by an European Research Council Advanced Grant (no:
742870).
REFERENCES
Aishwarya Agrawal, Aniruddha Kembhavi, Dhruv Batra, and Devi Parikh. C-vqa: A compositional split of the visual question answering (vqa) v1.0 dataset. ArXiv, abs/1704.08243. 2017.
Yuval Atzmon, Jonathan Berant, Vahid Kezami, Amir Globerson, and Gal Chechik. Learning to generalize to new compositions in image understanding. arXiv preprint arXiv:
2016.
Jimmy Ba, Geoffrey Hinton, Volodymyr Mnih, Joel Z Leibo, and Catalin Ionescu. Using fast
weights to attend to the recent past. In Advances In Neural Information Processing Systems, pp- 4331 1331-4339, 201
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: 1607.06450, 2016b.
Sergey Bartunov, Jack Rae, Simon Osindero, and Timothy Lillicrap. Meta-learning deep energy- based memory models. In International Conference on Learning Representations, 2020. URL
https::/
net/ wow.://ooommmiddyyyyye..
Robert Csordas and Juergen Schmidhuber. Improving differentiable neural computers through mem- ory masking, de-allocation, and link distribution sharpness control. In International Confer- ence on Learning Representations, 2019. URL HyGEM3C9KQ
net. wtw.eroouwweeeoooummmmmmmm forum?id
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdi- nov. Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint arXiv:1901.02860 2019.
S. Das, C.L. Giles, and G.Z. Sun. Learning context-free grammars: Capabilities and limitations of a neural network with an external stack memory. In Proceedings of the The Fourteenth Annual Conference of the Cognitive Science Society, Bloomington, 1992.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations, 2019. URL https:
Mete Demircigil, Judith Heusel, Matthias Löwe, Sven Upgang, and Franck Vermet. On model of associative memory with huge storage capacity. Journal of Statistical Physics, 168(2):288-299, 2017.
Jerome A Feldman. Dynamic connections in neural networks. Biological cybernetics, 46(1):27-39, 1982.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 1126-1135. JMLR. org, 2017.
Jerry A Fodor and Zenon W Pylyshyn. Connectionism and cognitive architecture: critical analy- sis. Cognition, 28(1-22 28(1-2):3-71. 1988.
A. Gers, J. Schmidhuber, and F. Cummins. Learning to forget: Continual prediction with LSTM. Neural Computation, 12(10):2451-4777 2000.
Samuel Gershman, Kenneth A Norman, and Yael Niv. Discovering latent causes in reinforcement learning. Current Opinion in Behavioral Sciences, 5:43-50, 2015.
10

---

Published as a conference paper at ICLR 2021
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv:1410.4011 2014.
arXiv preprint
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwinska, Sergio Gomez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, AdriÀ Puigdomenech Badia, Karl Moritz Hermann, Yori Zwols, Georg Ostrovski, Adam Cain, Helen King, Christopher Summerfield, Phil Blunsom, Koray Kavukcuoglu, and Demis Hassabis. Hybrid computing using a neural network with dynamic external memory. Nature, 538(7626): 471-476, 2016.
David Ha, Andrew Dai, and Quoe V Le. Hyperettorks.. ICLR 2017, 2017.
Robert F Hadley. Systematicity in connectionist language learning. Mind & Language, 9(3):247- 272, 1994.
Geoffrey E Hinton and David C Plaut. Using fast weights to deblur old memories. In Proceedings of the ninth annual conference of the Cognitive Science Society, pp. 177-186, 1987.
S. Hochreiter and J. Schmidhuber. Long Short-Term Memory. Neural Computation, 18)) 1780, 1997. Based on TR FKI-207-95, TUM (1995).
S. Hochreiter, S. Younger, and P. R. Conwell. Learning to learn using gradient descent. In Lecture Notes on Comp. Sei. 2 2130, Proc. Intl. Conf. on Artificial Neural Networks (ICANN-2001), pp. 87-
94. Springer: Berlin, Heidelberg, 2001.
John J Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 7955):2554-2558, 1982.
Daniel D. Johnson. Learning graphical state transitions. International Conference on Learning Representations, 2017. URLhtttss htttts:://opeereeeeeeeeeeeeeeeee /ttp///wwpp./...ep///ppp. wwt. wttp::eeeetoornnooooooUoo.
Pentti Kanerva. Sparse distributed memory. MIT press, 1988.
Louis Kirsch, Sjoerd van Steenkiste, and Juergen Schmidhuber. Improving generalization in meta
reinforcement learning using learned objectives. In International Conference on Learning Repre- sentations, 2020. htttlp/ps.ppppeeeeeeeee/e/////tr
net www.:eteeeooummiiidss1eeeeeeerre
Tamara G Kolda and Brett W Bader. Tensor decompositions and applications. SIAM review, 51(3): 455-500, 2009.
Bart Kosko. Bidirectional associative memories. IEEE Transactions on Systems, man, and Cyber- netics, 18(1):49-66 60, 1988.
Ben Krause, Emmanuel Kahembwe, lain Murray, and Steve Renals. Dynamic evaluation of neural sequence models. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 2766-2775, Stockholmsmässsan, Stockholm Sweden, 10-15 Jul 2018. PMLR. URL http: //proceedings mmr...ssm.......m......... mlr www.spresssseee... /v80 hatr///rruuellaa.. html
Dmitry Krotov and John J Hopfield. Dense associative memory for pattern recognition. In Advances in neural information processing systems, pp. 1172-1180, 2016.
Joseph B Kruskal. Rank, decomposition, and uniqueness for 3-way and n-way arrays. Multiway data analysis, pp- 7-18, 1989.
Ankit Kumar, Ozan Irsoy, Peter Ondruska, Mohit Tyyer, James Bradbury, Ishaan Guirajani, Vic- tor Zhong, Romain Paulus, and Richard Socher. Ask me anything: Dynamic memory networks for natural language processing. In Maria Florina Balcan and Kilian Q. Weinberger (eds.), Pro- ceedings of The 33rd International Conference Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 1378-1387, New York, New York, USA, 20-22 Jun 2016. PMLR. wwt.httpp//prrceeeee ngs ml press/ / v48/ /a/kkmmaaaa.. 6 html
Brenden M. Lake and Marco Baroni. Still not systematic after all these years: On the compositional skills of eoquencc-to-sequeeee recurrent networks. CoRR, abs/1711.00355 2017. URL http: 7/arx v. hwtpa//wwwrrg/bb///11.o.o 00350
11

---

Published as a conference paper at ICLR 2021
Brenden M Lake, Tomer D Ullman, Joshua Tenenbaum, and Samuel Gershman. Building machines that learn and think like people. Behavioral and Brain Sciences, 40, 2017.
Hung Le, Truyen Tran, and Svetha Venkatesh. Self-attentive associative memory. arXiv preprint vrXiv:2002.03519, 2020.
Yujia Li, Daniel Tarlow, Mare Brockschmidt, and Richard Zemel. Gated graph sequence neural networks. In International Conference Learning Representations, 2016. URL https://
arxiv. arxiv. org/abs 1511.05493
Brian McLaughlin. Systematicity redux. Synthese, 170(2):551-774, 2009.
Gábor Melis, Chris Dyer, and Phil Blunsom. On the state of the art of evaluation in neural language models. CoRR, abs/1 (0s/1707.05589, 2017. URL ULL/tttp//aaaav..o///a/////O/ 17589.05899999
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mix-
ture models. In International Conference Learning Representations, 2017. URL https
//openrevies ew wwweeeewwweieeeeeeeeeeeeeeeeeeee 2
Stephen Merity, Nitish Shirish Keskar, and Richard Socher. Regularizing and optimizing LSTM
language models. In International Conference Learning Representations, 2018. URL https:
//openeeview.neeeee wwt.::eee//ooooeee d-yyyPP0T2
Thomas Miconi, Kenneth Stanley, and Jeff Clune. Differentiable plasticity: training plastic neu- ral networks with backpropagation. In Jennifer Dy and Andreas Krause (eds.). Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 3559-35666 3568, Stockholmsmääsan, Stockholm Sweden, 10-15 Jul 2018. PMLR. URL htt.h:ttpppcceeeinngge
mlr wwwppppwsssseeeeeeeepepp.peee hwtp////micnnnaa..mm .tmll
Thomas Miconi, Aditya Rawal, Jeff Clune, and Kenneth O. Stanley. Backpropamine: training self-modifying neural networks with differentiable neuromodulated plasticity. In International Conference Learning Representations, 2019. URL/https://openeriiie ?
Ad=rilrAiA5Ye
Tomáš Mikolov, Martin Karafiát, Lukáš Burget, Jan Černocký, and Sanjeev Khudanpur. Recurrent neural network based language model. Eleventh annual conference of the international speech communication association, 2010.
Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Pieter Abbeel. A simple neural attentive leta-learner. In International Conference on Learning Representations, 2018. URL https:
//openzeview. www.:neee/tooommeooomommmomm DmUzWAW
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. International conference on machine learning, pp. 1928-1937, 2016.
Michael C Mozer and Sreerupa Das. connectionist symbol manipulator that discovers the structure of context-free languages. In Advances in neural information processing systems, pp. 863-870, 1993.
Tsendsuren Munkhdalai and Hong Yu. Meta networks. In Proceedings of the 34th International Conference on Machine Scarning-Volume 70, pp. 2554-2563. JMLR. org, 2017a.
Tsendsuren Munkhdalai and Hong Yu. Neural semantic encoders. In Proceedings of the confer- ence. Association Computational Linguistics. Meeting, volume I, pp. 397. NIH Public Access, 2017b.
Tsendsuren Munkhdalai, Alessandro Sordoni, Tong Wang, and Adam Trischler. Metalearned neural memory. In Advances in Neural Information Processing Systems, pp. 13310-13321, 2019.
12

---

Published as a conference paper at ICLR 2021
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring broad discourse context. In Proceedings of the 54th An- nual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp- 1525-1534, Berlin, Germany, August 2016. Association for Computational Linguistics. doi: 10.18653/v1/PD 144. URLhttps:: /wwp:///wwwww/////o/////////// htt.w//w..r//aahhoogg//////
Steven Andrew Phillips. Connectionism and the problem of systematicity. PhD thesis, University of Queensland, 1995.
Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adria Puigdomenech Badia, Oriol Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell. Neural episodic control. In Proceedings of the 34th International Conference on Machine Learning- Scarning-Volume 70, pp. 2827-2836. JMLR. org, 2017.
Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Lukas Gru- ber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, Victor Greift, et al. Hopfield networks is all you need. arXiv preprint 20Xiv:2008.02217, 2020.
Frank Rosenblatt. The perceptron: probabilistic model for information storage and organization in the brain. Psychological review, 65(6):386, 1958.
Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy Lillicrap. Meta- learning with aanemory--uggeneen neural networks. In International conference on machine learn- ing. pp. 1842-1850, 2016.
Imanol Schlag and Jürgen Schmidhuber. Gated fast weights for on-the-fly neural program genera- tion. In NIPS Metalearning Workshop, 2017.
Imanol Schlag and Jürgen Schmidhuber. Learning to reason with third order tensor products. In Advances in Neural Information Processing Systems (NeurIPS), pp. 9981-9993, 2018.
Imanol Schlag, Paul Smolensky, Roland Fernandez, Nebojsa Jojic, Jürgen Schmidhuber, and Jian- feng Gao. Enhancing the transformer with explicit relational encoding for math problem solving. 2019.
arXiv preprint
Margaret L Schlichting and Alison R Preston. Memory integration: neural mechanisms and impli- cations for behavior. Current opinion in behavioral sciences, 1:1-8, 2015.
J. Schmidhuber. Evolutionary principles in self-referential learning, or on learning how to learn: the meta-meta-... hook. Diploma thesis, Inst. f. Inf., Tech. Univ. Munich, 1987.
J. Schmidhuber. Learning to control fast-weight memories: An alternative to recurrent nets. Neural Computation, 4(1):131-139, 1992.
J. Schmidhuber. On decreasing the ratio between learning complexity and number of time-varying variables fully recurrent nets. In Proceedings of the International Conference on Artificial Neural Networks, Amsterdam, pp. 460-463. Springer, 1993.
J. Schmidhuber. On learning how to learn learning strategies. Technical Report FKI-198-94, Fakultät für Informatik, Technische Universität München, 1994. Revised 1995.
H. T. Siegelmann and E. D. Sontag. Turing computability with neural nets. Applied Mathematics Letters, 4(6):77-00, 1991.
P. Smolensky. Tensor product variable binding and the representation of symbolic structures in connectionist systems. Artif. Intell., 46(1-2):159-116 November 1990. ISSN 0004-3702. doi: 10. 101670004-7702999990000- URL wtttp://eeee doi, 00..://......000.01110000/0000.0 3702 (90) 90007-M
Paul Smolensky. Symbolic functions from neural computation. Phil. Trans. R. Soc. A, 370(1971): 3543-3569, 2012.
13

---

Published as a conference paper at ICLR 2021
Kelsey N Spalding, Margaret L Schlichting, Dagmar Zeithamova, Alison R Preston, Daniel Tranel, Melissa C Duff, and David E Warren. Ventromedial prefrontal contex necessary for normal associative inference and memory integration. Journal (
Karl Steinbuch. Die lernmatrix. Kybernetik, 1(1):36-45, 1961.
Sainbayar Sukhbaatar, Jason Weston, Rob Fergus, al. End-to-end memory networks. In Advances in neural information processing systems, pp. 2440-2448, 2015.
Volker Tresp and Yunpu Ma. The tensor memory hypothesis. ArXiv, abs/ 1oss//7880,, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Christoph von der Malsburg. The correlation theory of brain function (internal report 81-2). Goet- tingen: Department of Neurobiology, Max Planck Intitute for Biophysical Chemistry, 1981.
Jason Weston, Antoine Bordes, Sumit Chopra, and Tomas Mikolov. Towards ai-complete question answering: A set prerequisite toy tasks. CoRR, abs/1502.05698, 2015a. URL http://
arxiv. org. htgg/bbs/10 1502. 05698
Jason Weston, Sumit Chopra, and Antoine Bordes. Memory networks. In International Conference on Learning Representations, 2015b. URL
David J Willshaw, O Peter Buneman, and Hugh Christopher Longuct- Non-holographic associative memory. Nature, 222(5197):960- 1969.
Caiming Xiong, Stephen Merity, and Richard Socher. Dynamic memory networks for visual and textual question answering. In International Conference Machine Learning, pp. 2397-2406, 2016.
Zhilin Yang, Zihang Dai, Ruslan Salakhutdinov, and William W. Cohen. Breaking the softmax bottleneck: A high-rank RNN language model. In International Conference Learning Repre- sentations, 2018. URL htt.::/e/eeeeeeeeeeeeeeeeee/eeea
net Ewr.m/idummiiioeiee HwwzsiiG--wW--WCC
Wei Zhang and Bowen Zhou. Learning to update auto-associative memory in recurrent neural net- works for improving sequence memorization. arXiv preprint arXiv: 1709.0649, 2017.
14

---

Published as a conference paper at ICLR 2021
CONTENTS (APPENDIX)
Further Discussion
16
Derivation of the Update Rule
16
Comment the Regular bAbl Dataset and Previous Work
16
D Concatenated-bAbl Details
17
Ablation
18
Hyperparameter search for catbAbl
F.1
Fast Weight Memory
20
F2
Metalearned Neural Memory
F3
Transformer-LL
22
F4
LSTM
23
F.5
Attention to the Recent Past Fast Weights
24
Best catbAbl Runs Broken Down by Task
25
Language Modelling
27
Results
27
Meta Reinforcement Learning
28
15

---

Published as a conference paper at ICLR 2021
A FURTHER DISCUSSION
One way of assessing the capacity of the third-order tensor memory is its rank (which is analogous to the rank of a matrix). However, there exists general algorithm to determine the rank of given higher-order tensor A E R There exists only loose upper bound described by rank (A) -
min IK, JK3 Kruskal 1989 Kolda & Bader 2009).
might be tempting to simplify the FWM by replacing the outer-product of the input with a con- catenation as a means to reduce the space and time complexity. However, in highly compositional domains, the concatenated input will suffer from interference between memories. Consider a prob- lem which, from set of 10 symbols, requires the association of any three symbols represented by the vectors 8, r, & E R100 In the case of a concatenation, one rank of the fast weight memory is [s; r] @ t where we refer to [s; r) as the key representation. The read vectors s', 3'r' R are then concatenated and matrix multiplied to retrieve the previous association f = F(s'; :"). Here we refer to [s';r'] as the query representation. Since there are ten distinct symbols of which any two can behave as a key representation, there exist 102 100 unique key patterns. To guarantee noise-free retrieval in any context, the vectors of the key representations have to be orthogonal. However, [s': re:r"] is only a 20 dimensional space which means that certain key representations cannot be used simultaneously without interference. The tensor product, on the other hand, is capable of noise-free retrieval because it represents the key as s @ r E IR 10 which allows for 100 orthogonal keys and as such the possibility of noise-free retrieval. We conclude that if the problem is highly composi- tional, in a sense that every component can be composed with any other component, then the tensor product will be better suited than concatenation. Experimentally we evaluate concatenated keys in section E The results show that concatenated keys will result in a slightly worse performance (see figure [8). As an alternative, non-linear memory. e.g. through the use of softmax, would not require orthogonality in it's keys to be free of interference and could result in a larger storage capacity.
B DERIVATION OF THE UPDATE RULE
Theorem B.1. Given two key vectors k1, k2 RR and two value vectors Volds Unew Rd with d E Zo a mixxng coefficient 3 E (0, 1), and fast weight memory Fou vec(k) @k2) Vold where vec refers to the vectorisation the higher-order tensor, then the (recurrent) fast weight update rule given by Fola vec k1 @ k2) (Unew Vold) results in Faew vec(k1 @ k2)2 [[1 (011--)yar Boodd + peww).
Proof.
Fnow
Fold + 3 vec(k)
@ k2) @ (Unew Vold)
(9)
veo(k, 69 kg) Vold vec(k; k2) s9 (Bu (Bola))
(10)
( @ k2) @ [Uold + Beee [Woaa]]
(11)
veo(A:
2 k22) [[1 3)uold [Uee]]
(12)
C COMMENT ON THE REGULAR BABI DATASET AND PREVIOUS WORK
The bAbl tasks is popular toy dataset to benchmark neural networks with memory augmentations and reasoning capabilities (Weston et al. 2015a). It consists of a set of short stories with questions embedded in the text. The stories were generated by simulating multiple entities in a virtual envi- ronment and cover different contexts in which entities change their state or interact with each other. Each story-sample belongs to one of 20 different tasks that the authors of the dataset considered important for intelligent dialogue agents. The tasks contain questions which require reasoning capa- bilities like deduction, coreference, or counting. All tasks require some level symbolic reasoning, and the first neural and non-neural baselines demonstrated poor generalisation performance test data /Weston et al. 2015a). In addition to the story sentences, the questions, and the answers, the dataset also included supporting facts which demarcated question-relevant sentences in the story. The stories often follow multiple parallel plots where each new sentence is advancing one of the plots by single fact.
16

---

Published as a conference paper at ICLR 2021
The bAbl dataset did not include strict experimental protocol which resulted in several variations that differed slightly. Early methods achieved good results by relying on the supporting facts (We-) ston et al. 22015b Kumar et al. 2016) other supervised training signals (see e.g. Johnson 2011 7);
Li et al. (2016)).
Some researchers achieved great results by reformatting the data such that the question is read before the story or, similarly, by giving the model the capacity to lookup parts of the story, e.g. through some attentional mechanism, after the question has been read (Sukhbaatar al. 2015 Xiong al.) 2016 Dehghani et al., 2019). Such methods have shown to be useful for answering questions while maintaining access to the full story. We argue that this is similar to open-book question answering. In such a setting, the model is incentivised to look up information instead of capturing the useful bits of the data it has seen. The advantage of the latter becomes more evident in different scenario: imagine the model is processing book where user can ask a question about the content at any time. An open-book approach will have to store all previous sentences in its memory and apply its answer-search mechanism to all of the data. Instead, closed-book approach would store compressed version of the story, or the question-relevant information of the story.
It is essential to acknowledge that the sentences in the bAbl stories of all tasks are short and simplis- tic. Virtually every sentence contains precisely one fact. Because of that, it might be that sentence- level models have advantage over word-level models. Indeed, previous sentence-level model has reported poor performance in the word-level setting (Schlag & Schmidhuber 2018). This limits their generality since sentences in natural language are often not limited to a single fact.
Lastly, even though the bAbl dataset was initially designed with the questions embedded in the story, virtually all methods so far preprocess the dataset such that sample with four questions is split into four samples with one question each (Weston et al. 2015b). This arguably simplifies the problem because the model does not need to maintain the state of other entities which are not relevant the question once it is read. However, it remains to be tested if this would result in inferior performance.
D CONCATENATED-BAN DETAILS
Concatenated-bAbll (catbAbl) is a preprocessing and experimental procedure to evaluate autore- gressive models in their capability of predicting words which require certain reasoning skills (here answers of questions). In this work we only focused on the 10k samples per task version of bAbl but all our scripts can be applied the 1k version as well. We used the same train/test/valid split of the data as in regular bAbl. In contrast to previous work, we do not split the stories to contain only one question. We remove the sentence indecies and concatenate the sentences with answers following a question mark into one long sequence of words. The preprocessed data is shuffled list of samples. Each sample comes with its task id for diagnosis. All answers are preceeded by a question mark.
To ensure that stories do not overlap and become ambiguous, add special end-of-story token before concatenating the new story. For each word, the preprocessing script provides its task id to measure the performance on different tasks. Similarly, it also provides special answer token which signifies if the current word an answer or not. Naturally, the task id and answer information are not provided to the model an input. The validation and test data are processed likewise, but for a proper comparison of various models, validation and test data are shuffled only once 5 During training and evaluation, the validation and test stories are drawn deterministically.
Table 3: Statistics of the catbAbl dataset based on our preprocessing of the regular bAbI data.
subset
number of tokens
number of stories
number of questions
train
-5M
56,376
179.909
valid
-560k
6,245
19,907
test
-560k
6,247
19,910
During training we uniformly sample stories without replacement and concatenate them into a long sequence. Since a question mark is not always the end of story we resolve any ambiguity by
'We provide the preprocessed catbAbl data together with our code so future work can compare using the validation and test sequence.
17

---

Published as a conference paper at ICLR 2021
separating the stories with special end-of-story token. The model is trained on this long sequence in an autoregressive way with truncated backpropagation. At the end the epoch, we fill the batch with padding symbols if the sequences in the batch have different lengths.
In LM-mode we mask padding tokens and in QA-mode we mask everything except the steps with question mark input. At the end of the epoch carry over the hidden states the epoch. Reseting all hidden states to the same or to zeros had a weak negative effect on final performance but was explored thouroghly. For evaluation first batch element used. Evaluation on valid is done throughout training with large batch-size maintain speed. Evaluation on test is done with batch-size of one. During evaluation on valid and test the samples are picked sequentially to ensure that all models are evaluated on the same valid and test sequence of bAbl stories.
valid and test splits a copy of the hidden state of the
E ABLATION
We evaluate the FWM model with different number of recurrent steps. Experiments in figure indicate that just step is already achieving 95% accuracy but more inference steps help rarer but harder tasks. We also test a FWM version where the read and query keys are concatenated instead of multiplied through the tensor product. this version, the FWM results in a weight matrix The results in figure & indicate drop in performance.
with R 2dewx deWM instead of R xdewx
label
bs=64.0 I0=000000 1g=64.0_rr=0.c0 mr=
ba bo.04.0 00--64.0.1r-.000 Ir=6.001
1.0
bs 61-.640.0-0.oo0ooreea0
1.000
0.950
0.925
bask
0.900
0.875
0.850
0.825
0.800
0.775
0.700
10000
13000
20200
25000
$3000
step
Figure 7: Comparison of the FWM with the same seed but with different N
18

---

Published as a conference paper at ICLR 2021
1.000
0.975
0.950
0.925
0.900
0.875
acc task 0
0.850
0.825
10.800
0.775
0.750
label
0.725
concat seectt_seed=111
0.700
concat secc=t33
0.675
concat seed=222
0.650
0.625
best FWM
0.600
10000
20000
30000
40000
50000
60000
step
1.0
1.0
0.9
0.9
0.8
0.8
0.7
0.7
0.6
0.6
ra acc
0.5
label
0.4
0.4
concat_seed=11
0.3
e 0.3
concat =eed=333 333
concat seed=222 best FWM
0.1
0.1
0.0
0.0
10000
20000
30000
40000 50000 60000
10000 20000
30000 40000 50000 60000 step
step
1.0
0.9
0.9
0.8
0.8
0.7
task
0.6
0.6
0.5
0.5
0.4
0.3
0.3
0.2
0.2
0.1
0.1
0.0
0.0
10000 20000 30000 40000 50000 60000
10000 20000 30000 40000 50000 60000 step
step
Figure 8: FWM model with a concatenated keys compared with the tensor product of the keys. With a concatenation of the respective keys and queries the Fast Weight tensor has a squared space and compute complexity (dffws) but performs difference is limited to more complex tasks such as 3, 14, 16, 19 (bottom figures).
average (top figure). The performance
19

---

Published as a conference paper at ICLR 2021
F HYPERPARAMETER SEARCH FOR CATBABI
Since catbAbl is an ongoing sequence of stories, backpropagation through time (BPTT) infeasable for all models which why truncate BPTT to the last 200 tokens. Hyperparameters such that they fit roughly on one GPU with 16GB of memory. All models use token embedding size of 256 and the Adam optimizer. We exclusively tuned the hyperparameters for the QM setting and transfer only the best to the LM setting. We run a grid search over the batch-size, learning rate, and various model specific parameters such as dropout rates or number of layers on top of additional manually chosen settings. For computational reasons we run two rounds of grid-search: round of 3,000 steps of which the best are moved to the second round where we train them for 30,000 or 60,000 steps. In the following subsections give further details for each model seperately.
chosen
initial
F.1 FAST WEIGHT MEMORY
We set dLSTM 256, drwm 32, .N,- 3 and searched experimented with two seeds for batch sizes 64, 128 and learning rates 0.0001, 0.00025, 0.0005, 0.001, 0.002.
test accuracy all tasks
1000
0.90
0.85
0.75
label
bs-128 00000_mmm 3__r0 002
0
0.65
bs=128 0 I-= 001
bs-128 Ir_ 0.001
0.55
bs=128 10=00==00000000 00025
bs-64 0 Ir-0 001
0.45
bx 0 In-0 0001
bs= 0 Im=0 00025
0.40
500
1000
1500
2000
2500
3000
FWM test accuracy on all tasks
0.950
0995
0900
0.80
0.95
0.800
0.775
0.725
0.700
0.675
0.650
0.600
0575
0.50
label
0475
bs=64.0
001
seed-111
0.400
bs-64.0
Ir-0
001
seed-333
04455
bs=64.00
0 Jm-0 0.000000111000000 seed-222
0.400
10000
20000
30000
40000
50000
60000
step
Figure 9: Top: Hyperparameter search runs for different batch sizes and learning rates of the FWM model the QM setting with the average accuracy on all tasks. Bottom: FWM performance over 60,000 steps with three seeds.
20

---

Published as a conference paper at ICLR 2021
F.2 METALEARNED NEURAL MEMORY
We only experimented with the plastic version of MNM it was reported to be the best. We used the same hyperparameters for the fast weights as reported by Munkhdalai et al. 20099: layer of fast weights with dimensionality of 100. We searched over the batch sizes 64, 128: learning rates 0.00025, 0,0005, 0,001, 0.002; and meta-objective coefficient (reg) 1.0, 2.0. In the first 3,000 steps the MNM didn't I show any instability but for longer runs the MNM would sometimes result in NaNs or become unstable.
test accuracy on all tasks
label
bs=64 0 ltt_:0000000eeeeee reg-11 0 heads=3.0
1,00
ba=64.0_ 0000:00000000_p_110..eee..e
0.95
bd-64.0. ha_r:_____00000___________eeeeee
bs=64 0_rr=-000000_rrg-___0__
0_heads==3.0
bx=64 0 ir000000111110000000000000000000
ptt_rrgg_...hhaa-33.. heads=3.0
bs=128.0 0 Ir=0 Ir-0.0005 000000055.000 reg=1.0 heads=3.0
0.85
bs=126.0 1ew.:000reeee111_1
2agg=10_heads=30 3.0
0.80
bs=64.00 0 0000000r=0000001000000 reg-e.eeeccc heeds=3.0
0.70
0.65
0.60
0.55
0.45
0.40
500
1000
1500
2000
2500
3000
step
accuracy on tasks
1.000
0.975
0.950
0.900
0.875
087
0.800
0.750
0.725
label
0.700
0.655
k===4400r=00000255ree:000
heaa==30___edd1111...
0.650
0
0.600
10=12880_r=0000025__eggg-.
0.575
heads=3.
bs=64.0
0.555
bs=64.0
0==64.0_1==0000_reee:111 heads=3.0 0.ee0330e
0.500
1a=64.0_Ir=00
0.450
0.425
100.00_rr=00000220 te==640_1r-0000055_eeee:0...33.0 0ttp::=0.._haada33.000ees
0.400
10000
20000
30000
40000
50000
60000
step
Figure 10: Top: Hyperparameter search runs for different batch sizes and learning rates of the MNM model in the QM setting with the average accuracy all tasks. Bottom: MNM model with three different seeds, batch size 64, and learning rate 0.001 in the QM setting. Reported accuracy the average all tasks.
21

---

Published as a conference paper at ICLR 2021
F.3 TRANSFORMER XL
We ported the official Transformer-XL. implementation to our own codebase; fully reusing the model code for our catbAbl experiments. We employ a linear learning-rate warm-up schedule over the first 1000 steps and run grid search over batch size, learning rate, number of layers, and memory length with some additional manual selected parameters. Our best setting uses a learning rate of 0.00025, memory width of 1200, a hidden state size of dmodee 512, inner dimension of the fully connected part of dimner = 2048, and 3 transformer layers. Several long runs can be seen in figure [2] Our experiments show how various seeds eventually become unstable and overfit. Some settings also resulted in NaNs which we have removed from figure The best performing models and most stable where layer models with large memory and a small learning rate (see figure
test sccurscy on tooks
tect accurocy on all tooks
15
315
3.20
tom12000 (ewd 2000 ele mom-
jwod Loo L 4:30
220
280
0.70
270
00025
4.45
atel
lowod Jowe UAZD, Derem
2 2.10
Mosie fra memoces i
2.15
3.0
225
0.2
22
0.10
200
1200
2000 2003
2000
1000 1500
2000
2033
2000
dep
dep
Figure 11: Hyperparameter search runs for different batch sizes and learning rates of the Transformer-XX XL in the QM setting with the average accuracy on all tasks. Left graph varies number of layers and memory length. Right graph varies batch size and learning rate for 7 layers.
test tooks
0.5
50.4
8-M00 kotki Laed
650
peodeets
MoAAARA sied pnweeee 2000 Lendee111
TAAT WOS tadit
Tact NIE Ialr al tadis
020
023
050
Figure 12: Long hyperparameter search runs for TXL with various layers and memory sizes. The experiments are grouped based on the number of layers. Many runs begin to diverge late into the training process.
'Source: trans PYY
22

---

Published as a conference paper at ICLR 2021
test accuracy tasks
1.00
0.95
0.90
0.85
0.80
BEC talk
0.75
10.70
0.65
0.60
0.55
bc--128.0.11000:0025mmm
1:000225_n333______1 m01.3000mmmmm20000p..0mm 1atp-1220000_ed11111 0_s00d-111
bs-64.0
r-0.00025
nl-3.0 mem-1200.0 Geed-111
0.50
bs-64.0,
/tm.mmm0000mmmmm...mmmmm -30d=-333
ts=b4.0
0.45
bs-64.0
0.00025 ml03.0 mem- 1200.0 seed-222
bs-64.0 r=0.0025
nl-3.0 ml-3.
mem-1200.0 404d-444
0.40
10000
30000
40000
50000
60000
stap
Figure 13: Various seeds for the best Transformer- XL hyperparameters: 3 3-layers, memory windows of 1200 tokens, a learning rate of 0.00025, and a batch size 64.
F.4 LSTM
We heavily regularize a four-layer stack of residually connected LSTM cells, each with 512 hidden units. Inspired by AWD-LSTM (Merity et al., 2018), we use dropout in four different ways to regularize the model. We dropout the tokens of the input sequence, elements of the embedding vector, elements of the recurrent weight matrix, and elements of the of the hidden representation between LSTM layers.
scoraay Tasio
accunny on
2.0
Nevi and CR Soi--.0 R1B to.
Sropee 253503
281
you eoau
sud out PAENE
Figure 14: Hyperparameter search runs for different batch sizes and learning rates of the LSTM in the QM setting with the average accuracy all tasks.
23

---

Published as a conference paper at ICLR 2021
test accuracy on all tasks
label
bs=64.0 1r= 1r=000011 1=4 1tt.op=00110..om /0t.://////////...... 3/0.5 3eed=3
1.00
bs=64.0 |__ 4 http:pppp......////331....3
0.95
bs-64.0 Ir-00000 00000000000000000000000000000000
0.90
0.85
0.80
acc task
0
10000
20000
30000
40000
50000
60000
step
Figure 15: Average accuracy of three seeds of the best LSTM settings over all tasks on the catbAbl
QM-mode dataset.
F.5 ATTENTION TO THE RECENT PAST FAST WEIGHTS
We evaluate our own implementation Fast Weights as introduced by Ba et al. (2016a). They propose an RNN augmented with fast weights which modulate the slow weights of Elman RNN using a fixed fast weight learning and decay rate (JBFW). Our hyperparameter search did not result any model performing over 15% on the test data.
10
label
0.9
1o=6440_rr=0.000hh=256
0 d_deaa==0099
0=-=1880_Jr0000005__=25555 http:/shhssssffff______eeeeeee
0.8
3g=6640..==000000hh22555 flr=0.5 decaa=00.99
0.7
0.5
hs==64.0.rr00002_h==220
0.4
0 fir= 5_decy==0.05055
bs=118 0 Ir 001 h=256 0.5 0.9
0.2
0.0
500
1000
1500
2000
2500
3000
step
Figure 16: Hyperparameter search for the Fast Weights attending to the recent past by Ba et al. 2016a).
24

---

Published as a conference paper at ICLR 2021
G
BEST CATBABI RUNS BROKEN DOWN BY TASK
100
0.9
0.9
0.8
0.5
0.5
0.4
Iabel
0.4
label
0.3
MNH
0.3
MNM
Transformer-XL
Transforme-XX
FWM
LSTH
LSTM
10000
20000
30000
40000
50000
60000
10000
20000
30000
40000
50000
00000
step
1.0
0.9
0.9
0.8
0.5
0.4
label
8 0.4
label
0.3
MNM
0.3
MNH
Transfomee-XL
Transfomer-XL
LSTM
LSTH
0.0
10000
20000
30000
40000
50000
60000
10000
20000
30000
40000
50000
60000
100p
1.0
0.9
0.9
0.8
0.5
0.4
label
0.4
label
0.3
MNM
0.3
MNM
Transformer33
Transforee-XL
LSTM
LSTH
0.0
10000
20000
30000
40000
50000
60000
10000
20000
30000
40000
50000
60000
step
1.0
0.9
0.9
0.8
0.5
0.4
8 0.4
label
0.3
MNM
0.3
MNM
Transforer-XI
Transformer-0L
0.
LSTM
LSTM
0.0
10000
20000
30000
40000
50000
60000
10000
20000
30000
40000
50000
60000
2tep
1.0
0.9
0.9
0.8
0.5
0.4
8 0.4
label
0.3
MNM
030
MNM
Transforer-
Transformer-0
LSTM
LSTM
0.0
10000
20000
30000
40000
50000
60000
10000
20000
30000
40000
50000
60000
step
step
Figure 17: Per-task test performance comparison of the best catbAbl runs (first part).
25

---

Published as a conference paper at ICLR 2021
0.8
0.8
0.7
0.7
0.6
Iabel
label
0.3
MNM
Transfomer-XL
Transformer-X
02
FWM
02
FWM
0.1
LSTM
LSTM
10000
20000
30000
40000
50000
60000
10000
20000
30000
40000
50000
60000
0.8
0.8
0.7
0.7
0.6
0.4
Iabel
8
label
0.3
MNN
8
0.3
Transforee-XL
Transformer-30
02
FWM
0.2
FWM
0.1
LSTH
LSTH
10000
20000
30000
40000
50000
60000
10000
20000
30000
40000
50000
60000
0.8
0.7
0.7
0.6
label
8
label
0.3
0.3
Transfomer-
Transforee--x
02
FWM
02
FWM
0.1
0.0
10000
40000
50000
60000
10000
20000
30000
40000
50000
60000
step
0.9
0.8
0.8
0.7
0.7
label
0.6
MNM
Transformer-J1
label
0
0.3
8 03
LSTT
Transforme..t XL
02
FWM
02
0.1
10000
20000
40000
50000
20000
30000
50000
60000
step
0.8
0.8
30
0.7
0.7
0.6
label
label
0.3
0.
Transformer
Transformer
02
FWM
0.2
FWM
0.1
10000
40000
50000
60000
10000
20000
30000
50000
60000
step
Figure 18: Per-task test set performance comparison of the best catbAbl runs (second part).
26

---

Published as a conference paper at ICLR 2021
H LANGUAGE MODELLING
The code of our language modelling experiments is forked from Uber AI Lab's (github.com/uber. http://oieeettiiilloaaaaatttttts
which itself forked from the Sales- (manSeeeritaaaw--tmmmmn The FWM uses the same three force Language model toolkit (github.com//Serityaaaae
layer LSTM as the slow RNN with the same optimisations done by Merity et al. (2018). alternative which we do not explore here is to use multiple FWM-layers each with one LSTM cell and one FWM. We trained our model for 1000 epochs on PTB and 1600 epochs on WT2. Similar to Merity et al. (0018) () we switched from Adam to Averaged Stochastic Gradient Descent (ASGD) after 9T6 epochs and 1372 epochs for PTB and WT2 models respectively. We tune the dropout pa- rameters on the validation set and, after training, also tune the softmax temperature (tuning the softmax temperature results in -1 ppl of improvement). The embedding layers were initialized ran- domly from uniform distribution, uniform(-0.25 0.25), which erucial in our FWM language models. The hyperparameters used for all reported results are table [
The Transformer-X XL PTB results were based using the authors official code and hyperparameter setting (see zihangdai. waatuu.io/misc/ppb.... aging, and softmax tuning. The WT2 results are based the same code using the best hyperparam- eters found by Tim Dettmers (see httpl////w...ttttttttat//aeeeeen
which includes AWD-style regularisation, model aver-
Table 4: Best hyperparameters of the FWM for our language modelling experiments
dataset
droupout
dropoute
dropouth
dropouti
wdrop
batch size
ADAM
ASGD
PTB
0.4
0.1
0.3
0.5
0.66
20
0.001
2.0
WT2
0.4
0.1
0.25
0.7
0.61
80
0.001
0.5
H.1
RESULTS
Table 5: The detailed evaluation results of the FWM and Transformer-Xl. language model for all data partitions of the PTB and WT2 datasets using batch size of Experiment logs can found in our git repository.
loss
ppl valid
bits per word valid
model
dataset
seed
train
valid
test
train
test
train
141
2.82
4.04
4.00
16.77
56.76
54.48
4.068
5.827
5,768
PTB
142
2.66
4.05
4.01
14.26
57.43
55.17
3.834
5.844
5,786
143
3.16
4.08
4.04
23.66
59.31
56.90
4.564
5.890
5.830
FWM
1881
3.32
4.23
4.18
27.80
68.74
65.07
4.797
6.103
6.024
WT2
1882
2.81
4.16
4.12
16.66
63.98
61.65
4.058
6.000
5.942
1883
3.28
4.23
4.17
26.60
68.39
64.91
4.733
6.096
6.020
2.87
4.09
4.04
17.62
59.71
56.63
4.139
5,900
5.824
PTB
3
2.88
4.08
4.03
17.84
59.39
56.50
4.157
5,892
5,820
2.86
4.09
4.03
17.52
59.73
56.53
4.131
5.900
5.821
TXL
444
2.61
4.19
4.15
13.60
65.71
63.28
13.599
65.706
63.283
WT2
555
2.61
4.19
4.15
13.66
65.83
63.40
13.660
65.830
63,400
666
2.61
4.14
4.19
13.62
65.73
63.11
13.622
65.725
63.109
27

---

Published as a conference paper at ICLR 2021
1 META REINFORCEMENT LEARNING
The meta reinforcement learning experiments trains an agent in training POMDPs and evaluates on test POMDPs. The environments are directed graphs with labeled edges. As part of the data generating process, novel graphs are sampled according the python algorithm in listing [ Actions and states are one-hot encoded. The agent receives 17 dimensional input: the reward location, the current location, the previous action, a fixed bit, the fractional progress as current step and the current total steps
reward sum. Getting to the reward location gives reward of 10. Choosing an invalid action gives a penalty of 0.05. We use discounting factor of 0.9 and a value coefficient of 0.1. The entropy coefficient of A2C is set to 0.03.
The agent and reward locations are randomly selected the beginning of the episode. With only 5 states, the reward is reachable in most 5 steps. As elaborated in section 4.2] such optimal behaviour is only possible once the agent has learned the graphs from its experience. Whenever the reward is placed in the environment reset timer set to O. When the agent reaches the reward, or after 6 unsuccessful steps, the reset timer is set to 0 and the reward and agent are randomly placed in the environment. We train with batch size of 600 agents and optimize the average step loss using the Adam optimizer.
import numpy as np
def sample sample_adjacent adfacency ttattatt a_states, random_stater while True:
A = np. zeros (n_actions, n_states, n_states))
every state has to be leavable by at least one action for from_state in range (r_ssttte) ([i for i in range (n (n_states) if != from_state]) action random_stated crrdom_statte.rnddnntt (0, n_actions)
to_state
Aaactoon, fron_state, to_state] =
every state has to be reachable by one or more from-states for to_state in range (n_states):
# only select states which don't have any neighbours given an action app.wheee sum
fetion_list, from_list # remove self from the selection
0)
options np. ap.asarray (list (filter (lambda x: x [(0] != to_state, 2ip (from list, action list)))
indecies np ar.arang (aptions. [hape[[]]]
chosen idx random_state choice (indecies)
from from_state, action = options [chosen_idx] A[action, from_state, to_state]
reject if they are not all connected 0 = A. sum (0)
010 0]
for in range (n_states):
Q np matmul (Q, Q)
if (Q = 0) sum () == 0:
return
Listing 1: Python3 code to sample new environments such that any state is reachable by any other state.
28


Reviews:
[{"review_id": "paOWDCZfttG", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "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": "The solution proposed is the combination of an RNN (LSTM) and Fast Weighted Memory (FWM). The LSTM produces a query to the memory used to retrieve information from the memory and be presented at the model output. It also controls the memory through fast weights that are updated through a Hebbian mechanism. The FWM is based on Tensor Product Representations (TPR). The FWM is differentiable and builds upon the work of TPR-RNN from Schlag and Schmidhuber and Metalearned Neural Memory (MNM) by Munkhdalai et al. In the experimental section, the authors propose a concatenated version of the bAbI dataset to test their model with language modeling and question answering. Further the model is trained on a meta-learning task over POMDPs on graphs, and on language modeling on the PennTree Bank dataset. They show that the LSTM-FWM model generalizes better than without memory and similar models and with smaller capacity.\n\n======================================\n\nIndeed, the FWM model is relevant to this community and involves current scientific discussion and challenges. The paper is clear and is enjoyable to read. Math derivations and experimental results seem sound. Nevertheless, there are some clarity issues with the PTB language modeling task.\n\n======================================\n\nWould appreciate if the authors can answer to the following questions:\n\nHow is the FWM (tensor $\\mathbf{F}_t$) initialized? How does the initialization influence training and performance?\n\nHow is Nr selected?\n\nWhat is the vocabulary size in catbAbI? Is the embedding layer learned or pre-trained?\n\n“The experimental results in table 2 demonstrate a relative improvement over the AWD- LSTM baselines, which suggest the benefit of our FWM.” It is unclear what is the benefit in the PTB dataset. The results show that the LSTM model has slightly better perplexity (60.0 / 57.3) than the LSTM-FWM (61.39 / 59.37). Please, could you clarify the above note versus the numbers?\n\nDoes Figure 2 have missing details? The caption doesn’t seem to match the figure or it is unclear what authors are referring to.\n\nFigure 3 can benefit from using a bigger font for the node and edge values.\n\n======================================\n\nI'm inclined to accepting this paper. I found the idea simple but yet effective, and tested correctly in the experimental sections. Would appreciate it if the authors can improve the clarity surrounding Figure 2, and explain the misleading comment regarding the PTB task.\n\n======================================\n\nMinor issues:\n\n-Page 2: “An biologically” -> “A biologically”\n\n-Page 2: “pattern is is different” -> “pattern is different”\n\n-Page 5: Please correct with the missing number “suffered a TODO% drop”\n\n-Page 5: “figure 4.1.1” -> “Figure 2”\n\n-Page 6: “noteable” -> “notable”\n\n==================================\nUPDATE\n\nThank you for replying to my questions and clarifying in the document.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "nydafL1iHL0", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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\nThe authors present a working memory model composed of a recurrent neural network trained via gradient descent and an associative memory based on the approach taken by Ba et al. (2016) in \"Using Fast Weights to Attend to the Recent Past\".  The model consists of an LSTM to which takes the input and its own state from the previous step to produce an output (or new state) which is then passed to a fast weight memory (FWM) module.\n\nThe application of the fast weights are decomposed into two steps: read and write.  The write step composes the fast weight matrix update where new information is written into F given the LSTM hidden state and the fast weight matrix from the last step.  The read step consists of potentially several  recurrent \"inference steps\" over the FWM producing an output (e.g. a next step prediction or encoding). \n\nThe authors evaluate the model over two separate datasets.  The first is a modified version of the bAbI dataset which concatenates separate bAbI stories together and can be trained and evaluated in either a language modelling (LM) mode or question-answering (QA) mode where knowledge about past facts must be utilized.\n\nStrengths & Weaknesses:\n\nThe problem itself is well motivated since associative inference is useful in solving problems that require an accurate working memory.  Fast weight approaches allow us to learn to produce good state representations of the input sequence via slow weights (h_t) and where fast weights provide the associative mechanism to make important links across time. \n\nThe authors propose a new model that combines a novel read-write mechanism that relies on a number of inference steps over the fast weights allowing a nice disentanglement of read/write operations taking advantage of the associative inference to both add new relevant associative information (v) while also filtering stale data (v_old). \n\nThat said the overall form of the model doesn't seem fundamentally different from what is proposed by Ba et al. (2016) who also used fast weights as a way to attend over past hidden states in combination with a \"slow\" weighted RNN trained via gradient descent optimization albeit some of the details differ.\n\nFurther it would be helpful if the authors could clarify more around the rationale around why particular architectural choices were made.  For instance, why are two keys are generated in the write operation? \n\nResults for both catbAbI don't seem to exceed the performance of the TransformerXL when comparing perplexities in both QA and LM mode and don't exceed TrXL accuracy in LM mode.  However it is noted that the FWM model is in fact much smaller.  It may have been useful to investigate the gated transformer XL which is known to exhibit stronger stability for RL.  Figure 2 is nice though, is there any intuition why the reads vary among strong negative or positive activations as it seems to indicate?\n\nAs for the meta-RL problem it would have been nice to see comparisons to baselines other than an LSTM.  For instance, Ritter et al. (2020) in \"Rapid Task Solving in Novel Environments\" introduce a model that combines an episodic memory with self attention to meta-learn how to explore and exploit navigation to goals in connected graphs. \n\n\nOther points:\n\nThe labelling for the edges Figure 3 isn't really clear.\nThere's a missing reference in second to last paragraph on page 5: \"... QA-mode suffered a TODO% drop in accuracy ...\" \n\n\nRecommendation:\n\nI don't think there's enough here to recommend acceptance.  For starters, I don't think there's quite enough in justification around the architectural choices of the model and exactly what distinguishes this from the model proposed by Ba et al. which also used fast weights in combination with a \"slow\" weighted RNN.  Next, the results are not strong enough and additional or stronger baselines would have helped paint a better picture of the potential benefits of this approach.  For the results in general, while I think that these results point in the possible direction of the utility of FWM I don't believe the paper in its current form demonstrate that FWM exceeds state of the art in the chosen domains in which it was evaluated.  That said, I believe this is a promising line of research and encourage the authors to try to address the issues raised.\n\n\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "LBfmwKZemZm", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "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": "This paper presents a new method called Fast Weights Memory (FWM) to add an associative memory to an LSTM. \n\nModel:\n* FWM updates its fast weights through a differentiable perceptron like update at every step of an input sequence. The slow weights of the LSTM are instead updated only during training using gradient descent.\n* FWM is based on previous work: TPR (Tensor Product Representation). TPR is a mechanism that uses tensor products to generate unique representations of combination of components.\n* For long sequences FWM also has specialized components that allow it to update deprecated associations.\n\nFWM is related to:\n* TPR-RNN: a sentence-level model for reasoning on text, achieving excellent results on bAbI.\n* MNM (Metalearned Neural Memory): a word-level model which augments an LSTM with a FFNN as its memory, trained with a meta-learning objective.\n\nThe authors propose the new task \"catbAbI\", a variation of the existing task \"bAbI\".  catbAbI seems to be mostly just a concatenation of the stories, questions and answers in bAbI into a single textual sequence. It's unclear how much harder catbAbI is compare to bAbI in principle.\n\nTPR-RNN and MNM are only trained for short sequences and so will have a hard time on catbAbI. The authors show that MNM in particular does poorly on the long sequences in catbAbI.\n\nResults:\n* good performance on catbAbI (language reasoning) -- but this is a new task, so no real baselines in other papers.\n* good results meta-reinforcement-learning for POMDPs compared to LSTMs.\n* good results on PTB language models, better than other published models, but not state of the art.\n\nLimitations:\n* FWM requires an order 3 tensor, which scales poorly in both time and space computational complexity. This limits this work to relatively small models.\n\nQuestions:\n* catbAbi simply converts bAbI into a single sequence of tokens. Does this really increase the true difficulty of the task, or is it rather a way of artificially limiting the class of models used to solve the task to simple LM-like models? Is it possible to reconstruct bAbI from catbAbI with simple heuristics?\n* Could you report results for FWM on bAbI? It’s pretty unclear at the moment how to compare the results on bAbI of FMW to the ones e.g. in cited “Metalearned Neural Memory” paper. Or at least results on a version of bAbI where predictions are run for each story separately, so that MNM is not as penalized for not being able to deal with long sequences of text.\n* In figure 2, what does the color represent?\n\nTypos:\n* Page 2:\n   * *A biologically more plausible\n   * *stateful weights that can adapt\n   * Most memory-augmented NNs are based *on content-based…\n* Page 3: \n   * becomes a part of the *model's output.\n   * Figure 1: A simplified illustration of *our proposed method\n   * third-order tensor operations using *matrix multiplications\n* Page 4:\n   * Wq → Wn in equation (1)\n* Page 5:\n   * there is TODO left\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "KBrk68bhvck", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "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": "This looks like an interesting paper with an original proposal. The empirical results on synthetic tasks are also good. The main problem that I am having is with the proposed network, specifically equation 3. I do not see why it makes sense to consider an outer product of $n$ and $e$ as an argument to FWM. As is mentioned in the paper (in appendix A) it would make more sense (both conceptually and from the perspective of complexity) to concatenate those two inputs, or even consider two separate inputs for the associative memory module. The authors argue that in that case the memories would interfere with each other. This is true if a weak associative memory, like the one considered in this work is used. However, if the authors used a modern Hopfield network such an interference would not be a problem. Specifically, consider the situation when after applying the FWM weights to $n$ and $e$ the results are passed through a steep non-linear activation function, like in Ref [1] (see for instance formula 10). This would suppress the interference between the memories and provide a nice memory recovery. Additionally, with these “stronger” models of associative memory the key vectors do not have to be orthogonal.  \n\nExperimental results look fine, however, I think the work would benefit from some comparisons with other proposals for fast changing weights models, for example Ref [2]. \n\nI am not sure I understand the last paragraph on page 2. It is very easy to convert Modern Hopfield Networks from the autoassociative to heteroassociative type. One just needs to introduce additional matrices for queries, keys and values, like it is done in Ref [3] when comparing Modern Hopfield Networks with attention. Also, when referring to Modern Hopfield Networks, the reference for the original work, Ref [1], is missing.  \n\nA couple of presentational suggestions: \n\n1. Figure 1 seem to be inaccurate. In order to generate x_{t+1} one needs to take into account both the output of FWM and current state h_t. Only the first arrow is shown in figure 1.\n\n2. After equation 4, what is W_n? Looks like a misprint - should it be W_q? Also in the second line after equation 4 there are some misprints in the formulas. \n\nI also have some questions: \n\n1. Typically most associative memory models converge to a fixed point if one runs them for a long time. It is not obvious to me if dynamical rules described by equations 1-3 converge to a fixed point after a sufficiently large number Nr of iterations. Do they converge to a fixed point or not? \n\n2. It looks to me that the results reported in table 2 indicate that LSTM without FWM have lower perplexity than LSTM with FWM on that task. At the same time, the authors seem to say in the text (second paragraph on page 8) the opposite. Could the authors please clarify this? \n\nI am willing to increase the scores for this submission if the questions/comments above are addressed. \n\nReferences: \n\n[1] Krotov and Hopfield, NeurIPS 2016. Dense associative memory for pattern recognition, arXiv:1606.01164.\n\n[2] Ba, et al, NeurIPS 2016, Using fast weights to attend to the recent past, arXiv:1610.06258.  \n\n[3] Ramsauer, et al., 2020. Hopfield networks is all you need, arXiv:2008.02217.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposed a way to combine LSTMs with Fast weights for associative inference.

While reviewers had concerns about comparison with Ba et al., and experimental results, the authors addressed all the concerns and convinced the reviewers. The revision strengthened the paper significantly. I recommend an accept.
Decision: Accept (Poster)