A Cognitive Model for Learning Abstract Relational Structures from Memory-based Decision-Making Tasks
Authors: Haruo Hosoya
Venue: ICLR 2024 poster (2024)
Keywords: Brain-inspired model; hippocampus; entorhinal cortex; memory; relational representation
Primary area: applications to neuroscience & cognitive science

Abstract: Motivated by a recent neuroscientific hypothesis, some theoretical studies have accounted for neural cognitive maps in the rodent hippocampal formation as a representation of the general relational structure across task environments.  However, despite their remarkable results, it is unclear whether their account can be extended to more general settings beyond spatial random-walk tasks in 2D environments.  To address this question, we construct a novel cognitive model that performs memory-based relational decision-making tasks, inspired by previous human studies, for learning abstract structures in non-spatial relations.  Building on previous approaches of modular architecture, we develop a learning algorithm that performs reward-guided search for representation of abstract relations, while dynamically maintaining their binding to concrete entities using our specific memory mechanism enabling content replacement.  Our experiments show (i) the capability of our model to capture relational structures that can generalize over new domains with unseen entities, (ii) the difficulty of our task that leads previous models, including Neural Turing Machine and vanilla Transformer, to complete failure, and (iii) the similarity of  performance and internal representations of our model to recent human behavioral and fMRI experimental data in the human hippocampal formation.

Published as a conference paper at ICLR 2024
A COGNITIVE MODEL FOR LEARNING ABSTRACT RELATIONAL STRUCTURES FROM BEMED
DECISION-MAKING TASKS
Haruo Hosoya?.2
1 Department of Dynamic Brain Imaging, ATR International, Japan 2 WebLab, Graduate School of Engincering, The University of Tokyo, Japan hosoyaßatr.
ABSTRACT
Motivated by a recent neuroscientific hypothesis, some theoretical studies have accounted for neural cognitive maps in the rodent hippocampal formation as representation of the general relational structure across task environments. How- ever, despite their remarkable results, it is unclear whether their account be extended to more general settings beyond spatial random-walk tasks in 2D envi- ronments. To address this question, we construct novel cognitive model that performs memory-basee relational decision-making tasks, inspired by previous human studies, for learning abstract structures in non-spatial relations. Building on previous approaches of modular architecture, we develop a learning algorithm that performs -eward-guided search for representation of abstract relations, while dy- namically maintaining their binding to concrete entities using our specific memory mechanism enabling content replacement. Our experiments show (i) the capability of our model to capture relational structures that can generalize over domains with unseen entities, (ii) the difficulty of our task that leads previous models, in- cluding Neural Turing Machine and vanilla Transformer, complete failure, and (iii) the similarity of performance and internal representations of our model to recent human behavioral and fMRI experimental data in the human hippocampal formation.
INTRODUCTION
In everyday human cognition, we often find relationships among entities. Sometimes, we discern common relational structures across various domains and lift it to general knowledge (Figure abb) For example, ordering can naturally be found not only among numbers, but also among objects, among individuals, etc. Thus, after repeated experience of such finding, spontaneously come up with the abstract notion ordering independent concrete entities, which can be key to quick understanding of ordering in unknown domains. This is just one example of abstract relational structure; other examples include tree-like structure and cyclic structure.
In recent neuroscience, it has been hypothesized that the hippocampal formation may play central role abstract relational representation (Eichenbaum and Cohen 2014). Indeed, although relational representations have been long known in this neural subsystem (Bunsey and Eichenbaum 1996
Dusek and Eichenbaum 1997: Constantinescu et al. 2016 Bao et al. 2019 Park et al. 2020: 2021
newer evidence suggests that such representations
be abstract, i.c., independent of concrete entities (Kumaran et al. 2012 Garvert et al. 2017: Liu et al. 2019). On the theoretical side, Whittington el al.
(018|20200 have remarkably shown, their proposal of
Machine (TEM), that
place- and grid-cell properties in the rodent hippocampus and entorhinal cortex (Moser et al.) 2008) can emerge as result of learning the general relational structure of two-dimensional geometry during spatial random-walk tasks. Thus, abstract relational learning is becoming fascinating, unified view of the hippocampal computation, potentially explaining previously considered multiple functions such episodic memory and navigation by single principle.
'Tee hippocampal formation refers the hippocampus plus its neighbor, the entorhinal cortex.

---

Published as a conference paper at ICLR 2024
(b)
(e)
abstract relations
inference
abstract relations
s
domain A
p(a
domain B
reward
W
(c)
(d)
memory module
update
image pair presented
answer relation
get
reward
epoch
spoch
Figure 1: Learning of abstract relational structure. (a) One-dimensional hierarchy (order structure) and its two instantiations for different domains. (b) Two-dimensional hierarchy, with many-to-many relations. (c.d) Relational Decision-mmkinn task (trial and epochs). (e) Architecture of our proposed model (ARDMO) with relational and memory modules.
In this study, we push this direction further by proposing a novel cognitive model that performs relational derision-maaing. tasks for learning abstract relational structures. Our tasks, inspired by previous human experiments (Kumaran et al.) 2012 Park et al., 2020 2021), require to repeat pairwise relational inference one domain after another and eventually yy learn the common relational structure across domains (Figure To solve the tasks, building on previous approaches of modular architectures for relational representation (Whittington et al., 2018 2020), we develop a memory- based learning algorithm that performs ewaard-guided search for disentangled representations of domain generic relations and Nooainn-speciie entities (Figure [b). Importantly, we observe a task nature that requires careful choice of memory mechanism, key technique to maintain dynamic binding between relations and entities, which leads us to specific memory update rule that enables content replacement. We call our proposed model Abstract Relational Decision-making MOdel (ARDMO).
Our experiments yielded the following novel findings. First, model trained on two types of hierarchical relation successfully learned representations that are generalizable to new domains with completely unknown entities and usable in down-stream tasks such transitive inference. Second, several existing models including Neural Turing Machine (NTM) (Graves et al.) 2014). Differentiable Neural Computer (DNC) (Graves et al. (016). vanilla Transformer (Vaswani al. 2017). and version of TEM /Whittington et al. 22188 (020) could not solve our tasks despite their powerful capabilities. Third, our model exhibited task performance and internal representations that were compatible with behavioral and fMRI data from previous human experiments (Kumaran et al. 2001) Park et al. 2021), providing the first theoretical account, to our knowledge, for a cognitive map representation specific to the human hippocampal formation.
Our contributions can summarized as follows. neuroscience, we provided new theory that the principle of abstract relational structure account for the neural representation in the human hippocampal formation during relational decision -aainng tasks. In Al research, we identified a novel relations from small dataset and thus
learning problem that requires drastic generalization challenges existing powerful models, but can be solved by our proposed model at the human level.
BACKGROUND: RELATIONAL. DECISION-MAKING TASKS
To motivate our specific model design, we briefly review the tasks used in previous human experimen- tal studies on relational learning. In Kumaran et al. (2012), the task is to learn "one-dimensional hierarchy" or sequential ordering among 7 images of human faces or objects (Figure [T]. The true relations not made known the subject. During the learning session. in each trial, the subject presented with a pair of randomly chosen images that are adjacent in the ordering and required to make decision the relation between them, either higher lower; the subject gets rewarded
2

---

Published as a conference paper at ICLR 2024
if the answer is correct (Figure Tt.) After a certain number of trials, the subject's acquisition of the relation is assessed in a transitive inference task. This requires the subject to iteratively use the learned relation to relate given pair of non-adjacent images. The study reported nearly perfect score for this task by the human subjects (Kumaran et al. (012)).
In Park et al. 22021) the task is to learn "two-dimensional hierarchy" among 16 images of human faces (Figure [p). The images are assumed to be organized in a 4 x 4 grid, which gives two types of sequential ordering among the individuals corresponding
the two axes ("competence and popularity"). The true relations are again unknown to the subject. The task proceeds similarly to the mie-dimensional case, except that, each trial, the subject is hinted with one of the two axes in which to infer the relation between a given pair of images (multi-axis) axii). Also, unlike the one-dimensional case, each image can be related with up to four other images (many-to-many relation). After the learning session, the study conducted an fMRI experiment, which revealed distance coding and hexagonal modulation properties in the hippocampus and entorhinal cortex (Park al. [20TT.
FRAMEWORK
3.1
TASK FORMALIZED
Inspired by the human experiments overviewed in Section 2 we formalize our task design for training and evaluation of our proposed model. First, we assume relationship A consisting of finite set of relations. We also assume a domain D = (Ep, RD) with a finite set &D C _ of concrete entities or simply entities and a set RD consisting of tuples (I,, ag x/) with F,, x E &D and a₃ E A. In the sequel, we generally consider set of domains in which a common relational structure is imposed on the entities. For example, the "ordering" relationship with two relations < (prior to) and > (next to), as in Figure [b. has the following general structure. The two relations are converse to each other: if a domain has (22)) then it also has (I22, , 11), and vice versa; in the sequel, we sometimes write 21 < x2 instead of (I1, <, ).2. In addition, every domain has "minimal" entity æ for which there is no x' such that (x', X, x1 E RD; similarly, every domain has a "maximal" entity æt; we call these terminal entities.
Our goal is to learn the general relational structure hidden in given set of domains by performing the following relational inference task over multiple domains. The entire task undergoes series of epochs (Figure [t).. Each epoch, given some domain D (unknown to the model), repeats the following process for T times: at time (i) receive (random adjacent) pair of entities 22, 2/ (ii) infer their relation az and (iii) obtain an immediate reward r, = if the inference is correct, i.e., (Iz, a,,x,) x/) E RD. or no reward rt 0 otherwise (Figure [ In the training phase, the model learns relational representation by performing the above process while maximizing total rewards for training domains. In the test phase, the model performs a similar process for held-out test domains, which have new entities but with the same relational structure; the total rewards give the performance score. Note that it is the test phase that corresponds to the human tasks described in Section 2 the training phase would correspond to the subject's all experience prior to the experiments. Our framework adopts the same task design for both phases for simplicity. Note also that the task structure is different from random-w -aakking (Whittington et al.) 2018 2020), where receive entity 2, and a relation a, in each step and predict the next entity It+1 such that (Z,, at, I++1) E RD.
3.2
MODEL STRUCTURE
Our model, ARDMO, has an architecture with two modules, one representing abstract relations and the other representing their correspondence to concrete entities (Figure Te). Although such modular architecture is somewhat similar to TEM (Whitrington et al.! 2018 2020) well other
memory-based models (Graves et al.] 2014 2016 Webb et al.] 2020 our main novelty lies in the
learning algorithm described in Section 3.3
To represent abstract relations, we assume abstract entities or states s E RRS We define the "relational" probability p(s' (ol1s, a) that state s' is related with given state s by given relation a E A:
(I) where R2 E SxS is a relation matrix specific to relation a and a is a (global) scalar; o is an activation function, for which we the Izznormalization p(s) = >/||>]]. We also define the "prior"
p(s'ls. a) NV(p(( (p(Ras))
3

---

Published as a conference paper at ICLR 2024
probability p(als) that given state has relation (with some other state):
p(a s) (s(s) where Σ ga(ss)
(2)
where (s()) is non-linear function specific to a. The prior distribution is particularly important to represent terminal entities, e.g., p(- R 0 for maximal entities.
To represent binding between abstract states and concrete entities, we introduce a memory mechanism. Assuming key matrix W E R#x8 and memory matrix M E ExH, we refer to an entity corresponding to a state s by the following function:
read (s) = Mh where h softmax(W's
(3)
Here, the intermediate variable h softly represents an address pointing to the content of an entity stored in the memory. As discussed below, we take the key matrix W a parameter, but the memory matrix M as a hidden variable. The latter allows for dynamical updates of the binding, which is erucial for discovering domain- general relational structure.
3.3 LEARNING ALGORITHM
Given a set Di,..., Dn of domains, our learning procedure for ARDMO runs series of epochs described as follows (Figure In each epoch, we start with randomly selecting a domain Da and randomly initializing the memory matrix M1. At each time t. given a pair entities 24, x/z we first obtain the corresponding states s, inferM,,w (xz) and s infer w (x{) using the following
"contentt-bsed" inference function:
infer (r) = WTh where h = softmax (Mxx)
After this, we make decision on the relation between these by sampling:
de prpp(aasss s))
(5)
where the "posterior" distribution p(a, se. st)) can be obtained by Bayes' rule using equations 1 and 2 The reward given for the decision is r, = 1 (I',, deg xx) E RD. or r2 otherwise. Note that use the memory to estimate both states (s, and s() and use the relational representation to infer their relation (a,), unlike TEM (Whittington et al., 20188] 2020) or other ordinary recurrent networks, which use relational representation to estimate the next state (st) from given previous state (st-1)
and relation (at).
In the rewarded case, adjust the memory relational representations and thereby reflect the inference result. For this, we first re-postulate that the second state comes from the distribution for the inferred relation:
integrate the given input entities and the current
3) p(s de)
(6)
(but leave the first state 8t as it is for simplicity). We then recall the current memory contents E = readas,w w (8t) and x/ = read and simultaneously update the memory:
M41 Mr + (xe 2 softtmaaxx(ss (w)) ST)T + (a) -2) (W T
In the update, of the memory contents from the old ones (z, and 2!) to the new ones (I, and x²). Note that our memory mechanism is rather different from Hopfield-ty tyye auto-associative memory like Ba et al. (2016), which is used TEM |Whittington et al. 2018 2020). where the update rule only stores new associations, not erasing old ones. In this sense, our memory is more similar to Graves et al. (2014) 2016), although the old content is estimated directly from the memory in our update rule (equation 7 whereas it is computed by the controller network in their case. As shown in the experiment in Section 4. the choice of memory mechanism essential in our setting. This is probably because our task induces crucial interaction between states and memory and thus earlier inaccurate state-entity associations must be replaced later. (No such problem seems to arise TEM since current states can often be inferred directly from previous states in the random-wall task.)
relatively large coefficient, e.g., 0.7, which results in a drastic change
4

---

Published as a conference paper at ICLR 2024
Then, we define the following one-step loss function, which encourages correct inferences and discourage incorrect ones:
C. = p(eelse $())
(8)
Our goal is to minimize this steps and epochs with respect to the parameters 4 = (Ray Laga)la A} (W. J. where o(.) is the set of parameters used in the given function. Note that the loss function depends on the states, which in turn depend on the memory, which further depends the previous memory, inputs, etc. Therefore optimizing the loss necessarily causes pack-propagation through time Note that recurrent computation here occurs primarily through memory rather than through states, which is somewhat unique to our model. (Our algorithm can thus be seen to optimize the procedure of how to update the memory, thus "learning learn" the relations.)
Lastly, we incorporate a simple regularization that enforces the prior to follow the actual occurrences of the relations.
Crrior (rr p((sss()))))
In experience, without this regularization, the learned prior tends to be uniform. This makes the model fail to capture terminal entities and thereby degrade performance in transitive inference (Section 4. The entire learning procedure is summarized pseudo-code in Figure
3.4 TRANSITIVE INFERENCE
Transitive inference is an excellent task to test the learned relational representation. this, given pair æ² of entities,
ask to infer the relation whose more iterate, or transitive closure, relates them: (ar, o, x11 ), (I1, a, 222)..... (Im-1, a, E for some m 2 1 (where the subscript is for iteration, for time).
To solve this, we first obtain the states 8 corresponding to x, x'. We then calculate the following two probabilities pt (s"ls, a, m) (that state s' is related with given state s by the m-th iterate of given relation a) and p (a s, m)) (that given state has m-th iterate of relation a):
=)())) (t
II p(a)va(s)) k=
(10)
where x(p() and n(s) infer m, w (read w (s)); thth means the m-th iteration of function vb. Using these, we next obtain the following probability p* (als, s') (that given states and s' are related by the transitive closure of relation a):
M Σ
( (alss
(alas,
(11)
M
E. pt (s'ls, a', m)p (a' s.
(assuming that only the same relation is iterated at most M times). Finally, we answer a arg max, 1= (a)s,,') for the relation question. For later evaluation, we also use the confidence value ct = max, px (als,
The above approach can actually be derived as an approximation a certain probabilistic framework (Appendix A). The approach works in the case of one-to-one relations used in experiment in Section 4.1 Note that, for successful transitive inference, it is crucial for the model to learn both distributions p(s"|s, a)) and p(als) precisely; otherwise, it may make wrong judgment, e.g., that state has the m-th iterate even does not.
4 EXPERIMENTS
In this section, we present the results of experiments on two tasks. Additional details on the experiments are given in Appendices B and
5

---

Published as a conference paper at ICLR 2024
(a) Total reward
(b) Relational inference
(c) Transitive inference
3050
0.9
2520
WTM tmnin]
LETN tual
2210
NIM losi
BBC mmnn)
Infrence accore
2.0-
2000
UNC loaul
0.7
human (d5)
human (42)
FW train
102 rprs
1500
FW 2e53
spox
human (41)
502 spss
huran (42) nagu
1250
0.4
1000 opcs
spca
1.2-
regu
125
opochel 0og)
300
1.0
100
200
Stspe
300
400
500
600
300
200
stspe
300
400
500
600
Figure 2: Results on the one ineedimennional hierarchy task. (a) Traces of total rewards for ARDMO for training and test data, contrasted with LSTM, NTM. DNC, Transformer (TF), and Fast- sess--eights-baid based ARDMO (FW); FW can be seen as version of TEM. The horizontal axis is log-scaled. (NTM and DNC took very long training time ~ 40 days). (b) Performance ramp the relational inference within epoch (up to 600 steps) at different training stages, together with the corresponding human data (for two domains). (c) Inference score in the transitive inference for the models trained with and without regularization, together with human data. The human data are replotted from Kumaran et al. (2012). Error bars show standard errors over domains (models) or subjects (humans).
4.1 ONE-DIMENSIONAL IHERARCIY
The first experiment here simulates the previous human experiment on a simple one- -ie--imennional hierarchy (Kumaran et al., 2012) outlined Section 2 Using CIFAR 100, we formed each domain consisting of 7 images as entities, randomly selected and ordered from a single, randomly chosen object class. We prepared 600 training domains and 100 test domains so that these included images from disjoint classes. We assumed two (ordering) relations < and > introduced in Section 3.1 Recall that these relations are defined only for adjacent pairs; there are minimal and maximal entities. Note also the small-size nature of the task: only 12 possible pairs per domain. In the training phase, starting with a randomly initialized instance of ARDMO (H = 50; S = 20), we ran 8000 epochs each performing 3000 training steps for a randomly chosen domain. The training took about days. We then proceeded to the test phase, which was similar to training, but without optimization with respect to the model parameters. Below, we present the results.
First, Figure 2 shows the trace of per-epoch total rewards during training and test for model instance. The training succeeded with almost full rewards. Importantly, the test followed a similar trace to the training. indicating that the model successfully generalized the relational structure to unseen entities. Figure 2 also gives comparison with other existing models trained on the same task (see Appendix [..2.2 for details). (1) LSTM (Hochreiter and Schmidhuber] 1997). a conventional recurrent model, showed no such generalization, which is no surprise since it has no updatable memory. (2) NTM (Graves et al., 2014) and DNC (Graves er al., 2016). recurrent models with memory mechanism, also showed no generalization. despite their complex and powerful capability. This can be because that these models do not decouple well abstract and concrete representations (e.g.. values to write in to memory are determined only by the recurrent network). (3) vanilla Transformer (Vaswani et al. 2017) similarly failed to generalize; we tried various architecture and training settings but got similar results. We consider that the model structure is too complex to match our task nature with very small dataset and high-dimensional input (4) A modified version of ARDMO using Fast Weights, Hopfield-tye memory mechanism (Ba et al, 2016), failed even in training. The model can be seen as a version of TEM (Whitiigton et al. 2018 2020) replacing the learning algorithm
NNote that the learning procedure involves discrete sampling (equation 5) which prevents some gradients
from being propagated. Empirically, however, our algorithm stably optimizes (Section 4 possibly using imprecisely computed gradients. We observed notable improvement when using straight-through estimator (Bengio al.] [2013] Jang et al. (2017).
Whittington al. (2022) shows that Transformer can accommodate a version of TEM, for which no such optimization problem happens prohably because they restrict weight matrices in particular way.

---

Published as a conference paper at ICLR 2024
with ours (to match the task format), but retaining the memory mechanism. The shown result thus implies that the appropriate choice of memory mechanism highly depends on the task structure. In sum, our comparisons above highlight the particular difficulty of our task that completely defeats these powerful models.
Second, we inspected how our model behaves within epoch test data. Figure 2 shows the trace of per-block test performance (probability of reward in a block of 25 steps) in the relational inference within an epoch at different stages of training. In each epoch, the model started perform badly but gradually became better. This behavior is expected since the model initially knew nothing about new entities but got aequainted with their association with the learned relation during the epoch. Indeed, as the representation got improved during training, the time to gain full rewards (thus identify the correct relations) became shorter and shorter. After completion of training (8000 epochs), the peak performance was comparable with the human data (Kumaran et al. (012)) replotted in Figure 2h. (We are not concerned here about the speed of ramp as human brains can clearly do much more complex operations in single trial, such replays and some kind of inferences the model can at best simulate human at the abstract level.)
Third, we tested how well our model could exploit the learned representation for transitive inference task, using the scheme described in Section 3.4 For the sake of comparison with human data, we calculated the inference score, namely, the performance multiplied by the confidence value ([2.99c+ + 1.). Figure 2 shows that the inference score ramped in accordance with the relational inference in Figure 2b. Again, the peak score was comparable to the corresponding human data (Kumaran et al. 2012), replotted in Figure Z In addition, Figure 2 shows that, when a model was trained without the regularization described in Section 3.3 and therefore did not precisely learn the prior distribution, the performance became significantly poorer. This result suggests that, in the human experiment, although the task did not explicitly require to infer the terminalities of entities, humans might have recognized these implicitly and thereby achieved the high task score. This result can stand as testable prediction.
Finally, we confirmed that the results were robust across model instances trained under the same condition or under some hyperparameter variations, though the choice of high update rate (a) was crucial. Appendix B.3 summarizes the results.
4.2 WWO-DIMENSIONAL HIERARCHY
The second experiment simulates the previous human fMRI experiment on the two-dimensional hierarchy task (Park et al. 2021), as outlined in Section 2 The main aim of the human experiment was to investigate the neural representation of non-spatial 2D map in the human hippocampal formation. In particular, they reasoned that, if the 2D map is represented by neural population with hexagonal grid fields (Hafting et al.) (005).. the aggregate neural responses measured by fMRI should overall increase in particular directions in the map and decrease in other directions and there should be equally spaced six directions with increase (Figure 4): this property is called hexagonal modulation.
Our interest is whether hexagonal modulation emerges in our proposed model. Thus, we formed training and test domains each consisting of 16 images as entities, analogously to the 1D case. As noted in Section 2] the task involved multi-axis, many-to-mant
relations. That is, the 16 entities intendedly formed 4 x 4 grid, with ordering in two axes: X1 and > in axis 1, and <2 and >2 in axis 2 (Figure [). Also, in one axis, each entity can be related with up to four entities; e.g. x X1 x' can hold for entities that are adjacent in axis 1, but potentially non-adjacent in axis 2. We trained an instance of ARDMO (H = 500, 5 = 20) for 6000 epochs, with 2000 steps per epoch. In each epoch, a target axis i was randomly chosen and given to the model so that it could infer between -i and >;. The training conditions were otherwise similar to Section 4.1 For comparison, we trained another model instance with smaller memory (H- = 50) and a null model (trained with the same condition but random rewarding) [] The training of each model took about 2 days. As a basic assessment, we ran the two trained models in the relational inference task in test domains and confirmed their performance ramps to a reasonably high score for both models (Figure Bab).
"Since no other model so far has succeeded in solving our task and neural plausibility of a non-working model would he meaningless, we provide here no answer whether or not the similarity with the human FMRT data is unique ARDMO.
7

---

Published as a conference paper at ICLR 2024
(a) Hexagonal modulation
(b) Periodicity
(c) Distance coding (model)
9,15
503
9.20
nil
11:35
9.55
0.00
1981
0.00
220
1W02
m1
-2.2
200
200
30
D7A4
@PodSon Megg
Figure 3: Results on the two-dimensional hierarchy task. (a) The hexagonal modulation property. (b) The estimated magnitude at each periodicity. (c) Distance coding in the state (left) or intermediate (right) representations. The shown results are for ARDMOs with larger (H 500) or smaller memory (H = 50) in addition to null model and human fMRI data replotted from Park et al. (2021)
(Fig. 4b). Error bars show standard errors over domains (models) or participants (humans).
To see if the learned model showed hexagonal modulation, we performed the following analyses after running the model in each test domain (Appendix C.I for details). To simulate the protocol Park al. (202T), where they measured the average brain activities during consecutive presentation of image stimuli at each pair of grid positions, we obtained the state vectors (x) corresponding to the images (using equation 4 and took their mean. The (mean) state values, when plotted against the grid direction (angle) between the two positions, overall showed periodic pattern of increase and decrease in the cycle of 60 deg: Figure 3a plots the state values that were phase-aligned (across dimensions) and averaged (over dimensions and domains, within bin of 30 deg width), which indicates hexagonal modulation. To quantify the magnitude of the modulation, we calculated the mean differences between the maximal and the minimal of average state values (6-fold periodicity) and compared with the cases of 4-, 5-, 7-, and 8-fold periodicities; the 6-fold periodicity gave the largest magnitude (Figure 3p). The results resemble the human fMRI data (Park al., (02T) (replotted in Figure 3bb).
Figure 4: Mod- ulatory map di- rections
The model with smaller memory also gave similar results but with a lower magnitude.
In addition, Park et al. )2021) observed particular distance coding of 2D map in the hippocampus and entorhinal cortex, where representations tended to be more similar between closer grid positions. To conduct a compatible analysis on our learned model, we obtained state and intermediate vectors (s and h, using equation ( corresponding to the images at each pair of grid positions and computed the dissimilarity (Euclidean distance) between the vectors (Appendix C.I). Repeating these for different domains, the dissimilarity had either increasing or decreasing tendencies against the (Euclidean) distance in the grid space. Figure 3 summarizes the number of domains with increase or decrease in the intermediate state representation. In both representations, increasing tendencies were rarely found in the model with smaller memory, but much more often with larger memory, thus more similar to the human fMRI data reporting only increasing tendencies (Park et al.) (021). These results robust across model instances trained in the same condition (Appendix C.2p.
As mentioned, hexagonal modulation is expected as mixed neural signals from cells in the human entorhinal cortex if they have hexagonal grid fields as in the rodent (Hafting et al. 2005). In this sense, our result reproducing hexagonal modulation is related to previous model studies reproducing the grid cell property such as Whittington et al. 2018: 2020). However, note that neither result is directly predictable from the other since the task and model structures are rather different. Also, our study only offers high-level (abstract) account of the entorhinal cortex representation and therefore how it connects the cell-level grid-cell property remains open.
RELATED WORK
Recently, learning models for the cognitive map representations in the hippocampal formation have been drawing attention. Among others, WWittington et al.] (2018) 2020) proposed TEM, which accounts for place and grid cells in rodents as an abstract relational structure of 2--gcometric environments while performing spatial random-walking tasks. Our study has been much influenced
8

---

Published as a conference paper at ICLR 2024
by theirs, adopting similar modular architecture with separate relation and memory representations. However, the different task goal, relational decision- decision-making, necessitates our model construction to use different approaches in (1) the learning algorithm that maximizes rewards from memory-based inference and (2) the memory mechanism with the update rule that can replace old contents. In particular, showed that the Hopfield-type memory proposed by Ba et al. (2016) does not work in our setting (Section 4.1). which is quite striking, given that TEM was successful by using this mechanism. Thus, simplistic adaption of TEM our task fails. Later. Whittington et al. (2022) simplified TEM and clarified a formal relationship with Transformers (Vaswani et al. (017). A further refinement of these models revealed minimal constraints to reproduce grid-cell properties (Dorrell et al. 2022). A similar modular architecture for learning general 2D structure has also been proposed by Uria et al. (2022). explaining a variety hippocampal neural properties.
Different approaches can also explain cognitive-map properties from sequential learning. Clone- structured cognitive graphs allow for learning graph representations that are abstracted from concrete observations (George et al. 2021 Raju al. 2022 Guntupalli et al. 2023), somewhat similarly to TEM. The learned representations, in conjunction with planning through probabilistic inference, reproduce various complex hippocampal properties like remapping. However, the model offers no account for entorhinal cortex properties due to the finite, discrete nature of state space. Other accounts of the hippocampal formation have also been proposed using intermediate representations of recurrent networks (Banino et al.] 2018 Sorscher et al. 2019), successor representations of reinforcement learning (Stachenfeld et al.) 2017). 7, or simple unsupervised learning (Dordek al. 2016).
Memory-augmenttd recurrent models have been used for solving general, complex tasks requiring relational reasoning. Some such models use memory for connecting abstract relations with concrete entities to discover abstract rules in input sequences and thereby solve symbol-processing tasks (Webb et al. 2020 Chen et al. [021]. However, despite the apparent similarity, their tasks to find out common rules hidden within fixed-length sequences are not compatible with our tasks to find out general structure in a set of binary relations. [Grvess et al. (2014 2016) have proposed recurrent networks with external memory, highly influenced by von Neumann machines, that can learn to solve list and graph problems. Although these models are powerful, their way of coupling the recurrent network and the memory module seems to deteriorate disentangled representation of abstract relations and concrete domains (Section 4.1.. Santoro et al. (201)) has presented a different approach for relational reasoning by using head attention that allows for interaction between memory slots.
Recent studies have incorporated memory mechanisms in classifier neural network (Ba et al. 2016).
in reinforcement learning (Pritzel et al. 20177; Hansen et al.] 2018 Fortunato al. (019). in generative models (Li et al. 2016 Bornschein et al. 2017 Wu et al. 2018), in meta-learning (Santoro al. 2016; Munkhdalai and Yu (207)) and so Although these models have some technical commonalities with ours, they use memory to efficiently recall past experience and thereby increase specific task performance; therefore the goals are largely different from ours.
CONCLUSION
In this study, we have proposed novel memory-based cognitive model for learning abstract relational structure from decision-making tasks. The results showing good match with previous human behav- ioral and fMRI data indicate that the principle of abstract relational structure might account for the hippocampal representations not only for spatial random-walking but also for more general relational tasks. In another view, the present study addressed the challenging AI problem of cross-doman generalization from a few examples, highlighting the crucial role of memory in abstract conception. Various future directions relational structures and learning with delayed or implicit rewards; in neuroscience, accounts for other hippocampal phenomena like replays, interaction with other brain areas like the prefrontal cortex, and connection with rodent data.
conceivable: in AI research, representations of general graph-like
ACKNOWLEDGMENTS
This work has been supported by dation (22-06), New Energy and Industrial Technology Development Organization (P20006), and Corporate Sponsored Research Programs for World Model and Simulator of The University of Tokyo.
Grants-in-Aie for Scientific Research (21K19812), Okawa foun-

---

Published as a conference paper at ICLR 2024
REFERENCES
Jimmy Ba, Geoffrey Hinton, Volodymyr Mnih, Joel Z. Leibo, and Catalin Ionescu. Using fast weights to attend the recent past. Advances Neural Information Processing Systems, pages 4338 4346, 2016.
Andrea Banino, Caswell Barry, Benigno Uria, Charles Blundell, Timothy Lillicrap, Piotr Mirowski, Alexander Pritzel, Martin J Chadwick, Thomas Degris, Joseph Modayil, Greg Wayne, Hubert Soyer, Fabio Viola, Brian Zhang, Ross Goroshin, Neil Rabinowitz, Razvan Pascanu, Charlie Beattie, Stig Petersen, Amir Sadik, Stephen Gaffney, Helen King, Koray Kavukcuoglu, Demis Hassabis, Raia Hadsell, and Dharshan Kumaran, Vector-based navigation using grid-like representations in artificial agents. Nature, 557(7705):429 433, may 2018.
Xiaojun Bao, Eva Gjorgieva, Laura K. Shanahan, James D. Howard, Thorsten Kahnt, and Jay A. Gottfried. Grid-like Neural Representations Support Olfactory Navigation of a Two-Dimensional Odor Space. Neuron, 10665):066 2019.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation. arXiv, 1308.3432, 2013.
Jörg Bornschein, Andriy Mnih, Daniel Zoran, and Danilo Rezende. Variational memory addressing in generative models. Advances in Neural Information Processing Systems, 2017-Decem:3921- 3930, 2017.
M. Bunsey and H. Eichenbaum. Conservation of hippocampal memory function in rats and humans. Nature, 37966562):255-257, jan 1996.
Catherine Chen, Qihong Lu, Andre Beukers, Christopher Baldassano, and Kenneth A. Norman. Learning to perform role-filler binding with schematic knowledge. PeerJ, 9:1-27, 2021.
Alexandra O Constantinescu, Jill X. O'Reilly, and Timothy E.J. Behrens. Organizing conceptual knowledge in humans with a gridlike code. Science, 15246622)11444-1444 2016.
Yedidyah Dordek, Daniel Soudry, Ron Meir, and Dori Derdikman. Extracting grid cell characteristies from place cell inputs using non-negative principal component analysis. eliie,, 5(MARCH2016): 1-36, 2016.
William Dorrell, Peter F. Latham, Timothy E. J. Behrens, and James C. Whittington. Actionable Neural Representations: Grid Cells from Minimal Constraints. International Conference Learning Representations, 2022.
Jeffery A. Dusek and Howard Eichenbaum. The hippocampus and memory for orderly stimulus relations. Proceedings of the National Academy of Sciences of the United States of America, 94 (13):7109-7114, 1997.
Howard Fichenbaum and Neal J. Cohen. Can We Reconcile the Declarative Memory and Spatial Navigation Views on Hippocampal Function? Neuron, 8344)7764-770, 2014.
Moire Meire Fortunato, Melissa Tan, Ryan Faulkner, Steven Sleven Hansen, Adrift Puifidomenech Badla, Gavin Buttlmore, Charles Charlie Deck, Joel Z. Lelbo, Charles Blundell, Adrià Puig- domènech Badia, Gavin Buttimore, Charles Charlie Deck, Joel Z Leibo, and Charles Blundell. Generalization of Reinforcement Learners with Working and Episodic Memory. In H Wallach, H Larochelle,, A Beygelzimer, F d'Alché-Bue, E Fox, and R Garnett, editors, Advances in Neural Information Processing Systems, volume 32, pages 1-10, 2019.
Mona M. Garvert, Raymond J. Dolan, and Timothy E.J. Behrens. A map of abstract relational knowledge in the human hippocampal--nntrtinne cortex. eLife, 6:1-20, 2017.
Dileep George, Rajeev V. Rikhye, Nishad Gothoskar, J. Swaroop Guntupalli, Antoine Dedieu, and Miguel diizaro-Gredilly Clone-structured graph representations enable flexible learning and vicarious evaluation of cognitive maps. Nature Communications, 12(2392), dec 2021.
Alex Graves, Greg Wayne, and Ivo Danihelka, Neural Turing Machines. arXiv, 1410.5401, oct 2014.
10

---

Published as a conference paper at ICLR 2024
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwinska, Sergio Gómez Colmenarejo, Edward Grefenstette. Tiago Ramalho, John Agapiou, Adria 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 neural network with dynamic external memory. Nature, 538(7626): 471-476, 2016.
J. Swaroop Guntupalli, Rajkumar Vasudeva Raju, Shrinu Kushagra, Carter Wendelken, Danny Sawyer, Ishan Deshpande, Guangyao Zhou, Miguel Lázaro-Gredilla, and Dilcep George. Graph schemas as abstractions for transfer learning, inference, and planning. arXiv, 2302.07350, 2023.
Torkel Hafting, Marianne Fyhn, Sturla Molden, May Britt Moser, and Edvard Moser. Microstructure of a spatial map in the entorhinal cortex. Nature, 436(7052):001-806, 2005.
Steven S. Hansen, Pablo Sprechmann, Alexander Pritzel, André Barreto, and Charles Blundell. Fast deep reinforcement learning using online adjustments from the past. Advances in Neural Information Processing Systems, pages 10567-10577, 2018.
Hochreiter and Schmidhuber. Long short-term memory. Neural Computation, jan 1997.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel- International Conference Learning Representations, 2017.
Kingma and J Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations, jan 2015.
Dharshan Kumaran, Hans Ludwig Melo, and Emrah Duzel. The Emergence and Representation of Knowledge about Social and Nonsocial Hierarchies. Neuron, 76(3):653-666, 2012.
Chongxuan Li, Jun Zhu, and Bo Zhang. Learning to Generate with Memory. In Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 1177-1186, 2016.
Yunzhe Liu, Raymond J. Dolan, Zeb Kurth-Nelson, and Timothy E.J. Behrens. Human Replay Spontaneously Reorganizes Experience. Cell, 178(3):640-6522 2019.
Edvard I. Moser, Emilio Kropff, and May Britt Moser. Place cells, grid cells, and the brain's spatial representation system. Annual Review of Neuroscience, 31:69-89, 2008.
Tsendsuren Munkhdalai and Hong Yu. Meta Networks. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages
2017.
Seongmin A. Park, Douglas S. Miller, Hamed Nili, Charan Ranganath, and Erie D. Boorman. Map Making: Constructing, Combining, and Inferring on Abstract Cognitive Maps. Neuron, 107(6): 1226-1238.co 1238.c8, 2020.
Seongmin A. Park, Douglas S. Miller, and Erie D. Boorman. Inferences on a multidimensional social hierarchy use grid-like code. Nature Neuroscience, 24(9):1292-133 1301, 2021.
Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adrià Puigdomënech Badia, Oriol Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell. Neural Episodic Control. In Proceedings of the 34th International Conference on Machine Learning, volume 70 Proceedings of Machine Learning Research, pages 2827-2836, 2017.
Rajkumar Vasudeva Raju, Swaroop Guntupalli, Guangyao Zhou, Miguel Lázaro-Grediin and Dileep George. Space is a latent sequence: Structured sequence learning as a unified theory of representation in the hippocampus. arXiv, 2212.01508, 2022.
Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy Lillicrap. Meta- Learning with Memory-Aummented Neural Networks. In Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 1852-1850, 2016.
11

---

Published as a conference paper at ICLR 2024
Adam Santoro, Ryan Faulkner, David Raposo, Jack Rae, Mike Chrzanowski. Théophane Weber, Daan Wierstra, Oriol Vinyals, Razvan Pascanu, and Timothy Lillicrap. Relational recurrent neural networks. Advances in Neural Information Processing Systems, pages 7299-7310, 2018.
Ben Sorscher, Gabriel C. Mel, Surya Ganguli, and Samuel A. Ocko. A unified theory for the origin of grid cells through the lens of pattern formation. Advances in Neural Information Processing Systems, 2019.
Kimberly L Stachenfeld, Matthew M Botvinick, and Samuel Gershman. The hippocampus as predictive map. Nature Neuroscience, 20(11):1643-1553 2017.
Benigno Uria, Borja Ibarz, Andrea Banino, Vinicius Zambaldi, Dharshan Kumaran, Demis Hassabis, Caswell Barry, and Charles Blundell. A model of egocentric to allocentric understanding in mammalian brains. bioRriv, page 1020.11.11.378144 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Advances in Neural Information Processing Systems 30, pages 5998-69699 6008. 2017.
Taylor W. Webb, Ishan Sinha, and Jonathan D. Cohen. Emergent Symbols through Binding in External Memory. International Conference on Learning Representations, 2020.
James C.R. Whittington, Timothy H. Muller, Caswell Barry, Shirley Mark, and Timothy E.J. Behrens. Generalisation of structural knowledge in the hippocampal-entorhinal system. Advances in Neural Information Processing Systems, pages 8484-8495, 2018.
James C.R. Whittington, Timothy H. Muller, Shirley Mark, Guifen Chen, Caswell Barry, Neil Burgess, and Timothy E.J. Behrens. The Tolman-Fichenbaun Machine: Unifying Space and Relational Memory through Generalization in the Hippocampal Formation. Cell, 1249 2020.
James C.R. Whittington, Joseph Warren, and Timothy E.J. Behrens. Relating Transformers To Models and Neural Representations of the Hippocampal Formation. International Conference on Learning Representations, 2022.
Yan Wu, Greg Wayne, Alex Graves, and Timothy Lillicrap. The kanerva machine: generative distributed memory. International Conference on Learning Representations, 2018.
A DERIVATION OF THE APPROXIMATION IN SECTION 3.4
For solving the transitive inference task, we consider the probability y= (8'18,, a, m) that state s' related with given state by the m-th iterate of given relation
x (s'1s,a,
lf
/
m--
II p(s si a) p(s' (5m-1,a)ds;dee ... dsm-1
(12)
and the probability pt (als, m) that given state s has m-th iterate of relation a:
(a(ss = /f
2
II p((141,)(s)) (ss dsm-1
(13)
The above definitions are intractable in general. However, assuming one-to-one relations, we use the following approximation:
p² (s'ls, m) p(s'lu"" (s)) (s), (a)s,,m)
II p(alv 1=0
(14)
where (s)s) s(..s). (Empirically, iteration, and therefore use vi (s(s) = x()(R,))) with the function x(s) = infer M.w (read (s)). which recalls the concrete entity of a given state vector and then mapping it back state.)
find it more useful to "clean up" the state after each
12

---

Published as a conference paper at ICLR 2024
To derive the above, crudely approximate the Gaussian distribution in equation | by delta function: p(s'|s, a) zo öis' = p(Ras) To derive the approximation p (s'[s, a, m) in equation 14 first let:
m 2
A
II
p(5++1);;,0
(3m-- ... 1-m-1-k (15)
For 0 m 2, by the approximation assumption p(s'ls, a) ōهs' (a(s)), we obtain:
AB 2 |f... /
3-
p(8 84,
61sm-1-k (0
t=
(sm- a)ds ds2 ... dsm-1-k (16)
I'm 3-
II
p( (1841111,00
a) p(s (5m-2 adds dsm-2
(17)
k+1 = m,
(18)
Thus, noting A ppt [s'[s, a, m), the result follows.
To derive the approximation of p (als, m) in equation 14 let:
2-h
Bk m,
-... II p(si+1, als;)
(mal dsm-1-k (19)
o
For m 2, by the same approximation assumption, we obtain:
II 3-
B !f
p($14,
8[sm-1-k (0 ()p 8m-
I plalws (e [(tawv((mm--)))) dssde.... d8m-1-k
(20)
m 3-h
p(8+11,
II ==
plalv's ( 99m-2-n) 2-x)) dsids2 ddm-dmm-- 2-k (21)
==
B1++1
(22)
Noting 8 pt (als, m), the result follows.
B ADDENDUM ON THE EXPERIMENT IN SECTION 4.
B. DATASETS
In each experiment, we prepared a dataset using the CIFAR100 image database. To form each domain, randomly chose object class from which randomly selected 7 images entities and then ordered them randomly. We formed 600 training domains and 100 test domains, with no overlapping object classes between the two. Each image compressed 700 dimensions by PCA and normalized to unit norm.
B.2 TRAINING
B.2.1 ARDMO
For the architecture, apart from the description in Section 3 we used a two-layer perceptron to implement (s) (for prior) where the intermediate layer had 10 units with the sigmoid activation function and the output layer was fed to the softmax function. The pseudo-code of the learning algorithm is given in Figure
13

---

Published as a conference paper at ICLR 2024
1: procedure LEARN(D1, D2...., Dn, M)
2: initialize - = [Ra. 0, (( ) A} (
3: for all n ,,..., do
4: randomly select E []
5: initialize M
6:
for all t t= 1....... T do
7:
observe (æ,, x in RDs
8:
8: infer (zx)
9:
si = inferM,.w (x1)
10:
sample di p(an St, s(
11:
Ti + get get_reward (Tt, +4), RD.) if r. then
12:
13:
sample sy p( p(5/1s),, dt)
14:
z = read (8t)
15:
z read w (s/)
M411 t M1 a(x; so)ssmaaax((ss a(42 so)smmaa
17:
C= E (r: - pt, + + - p(āā())))
18:
miminize w.r.t.
Figure Learning algorithm for ARDMO
To train the model, we used hyper-parameterr a = 0.7 and 7 = 1, and Adam optimizer (Kingma and Ba 2015) with mini-batch size 5. We also used the same 0.7 7 for test. To boost the learning, we used truncated back- propagatoogt strategy. More precisely, we ran optimization after every 25 steps in epoch. That is, after every 25 steps, first initialized the accumulated loss function and stopped gradient propagation for the memory matrix M.
Generally, from observation of the loss function, training tended proceed follows. Initially, learning started with lengthy plateau (where the loss scarcely changed), then suddenly shifted to rapid drop (where the loss quickly improved), and thereafter fell into another plateau. This stairway- like course continued for few times and finally reached a very long and slow slope for convergence. In some cases, the initial plateau was so long that was unclear whether the optimization simply failed, in which case we manually stopped it after around 2000 epochs and started it over.
B.2.2 BASELINE MODELS
We adopted the following baseline models.
Long short-term memory (LSTM) (Hochreiter and Schmidhuber, 1997) well-known recur- rent model. Input dimensions: 700. Number of layers: Hidden state dimensions: 70. Batch size:
Neural Turing Machine (NTM) (Graves et al., 2014) A recurrent model with updatable external memory. Input dimensions: 700. Hidden state dimension: 20. Memory height (number of slots): 50. Memory width (content dimensions): 100. Batch size: 5.
Differentiable Neural Computer (DNC) (Graves al. 2016) An extension of NTM. Input di- mensions: 700. Hidden state dimension: 20. Memory height: 50. Memory width: 100. Number of write heads: 1. Number of read heads: Batch size: 5.
vanilla Transformer (Vaswani et al. 2017) well-known language model incorporating self- attention with (additive) sine/cosine position-encoding and temporal masking. Input dimen- sions: 100. Number of layers: 2 (repeated). Number of heads: 5. Hidden state dimension: 2048. Drop-out rate: 0.1. Context size: 100. Batch size: 100.
Fast-Weight-bassed ARDMO A modified version of ARDMO using the auto-associative memory proposed in Ba al. (016)). The encoding/decoding method of state and content described
14

---

Published as a conference paper at ICLR 2024
Table Summary of results of relational inference (performance) and transitive inference (inference score) on one-dimensional hierarchy (at 1000 steps in each epoch). The mean and s.d. of scores over 8 model instances are shown.
ARDMO
ARDMO w/o
1.STM
TF
FW
Human (dl)
Human (42)
Relational
1.00 0.00
1.00 + 0.00
0.51 0.02
0.50 0.01
0.51 0.01
0.97
Transitive Inf.
2.86 0.07
2,09 + 0.14
2.95
2.87
in Whittington al. 2020 Section 4) was used. Input dimensions: 100. (Other architecture parameters were the same as ARDMO.)
We used the default Pytorch implementations of LSTM and Transformer and a third party implemen- tation of NTM and D
In all of these, we input the concatenation of two given images to the model at each step and the correct relation target. For NTM and DNC, we appended the correct relation in one step the input in the next step, so that the model could take the correctness of the previous inference into account. For Transformer, we gave the concatenated image pair and the target alternately in the imput and let the model infer the target from each image pair, whose summed squired eπтоr became the loss function. For Fast- Best-Weigtt-baset model, we chose the reduced input dimensions ( (100) since the required memory size, quadratic in the input dimensions, would become prohibitive for the original dimensions (700); we chose the fixed update rate 7) = 0.5 following Ba et al. (2016).
B.3 ADDITIONAL RESULTS
Table summarizes the results from model instances for each task and for each model (except NTM and DNC for which we could train only one instance due to the very long training time), together with the human data. Figure shows the traces of performance in relational inference and inference score in transitive inference for all models in the same format Figure 2pe (showing model #1).
We also varied two hyperparameters, namely the memory update rate a and the state dimension S, and measured the total rewards (for test data). For this, we additionally trained 4 model instances for a = 0.01, 0.1, 0.3, 0.5, 0.9 and for S = 5, 10, 50, 100. Figure 6 plots the effect of varying a, indicating that the performance stayed high for relatively large value of (e.g.. 2 0.5) but dropped significantly for a smaller value. Figure 6 plots the effect of varying S, showing stable performance the tested range. However, using further larger state dimension like 200 often led to failure of training.
For Transformer, tried number of other architecture parameter settings than described in Appendix B.2.2] using more layers (5), less heads (1 or 2), or larger context sizes (300). The results were similar: all models failed generalize. This result was somewhat surprising since the task initially appeared to be easy for Transformer. Indeed, one can imagine Transformer that uses the input directly as a query and key and then uses position encoding find the corresponding relation found in the next entry in the sequence. However, optimization never found out such solution but the weights seemed excessively adapt to the input images, thus leading to overfitting. We also observed that even training did not succeed for higher input dimensions (700) or smaller batch size (20).
Lastly, we examined the ideal observer, which performs well possible in given epoch, assuming that it knows perfectly the ID task structure. We implemented the ideal observer such that it memorizes all past pairs in the epoch and, each step, simply searches for the given pair or its flipped pair; it finds, then can give the right answer (reward 1) or otherwise a random answer (reward 0.5). As a result, we found that both the model and human were much slower than the ideal observer, which reached nearly 1.0 around 50 steps in our calculation.
15

---

Published as a conference paper at ICLR 2024
(a) Relational inference
Transitive inference
(c) Effect of
(d) Effect of S
2300
2752
2752
2.5
2752
2003
1360
155
3352
125
160
0.4
3.6
92
300
77008
31.00 07901803
Figure Additional results on one-dimensional hierarchy. (a.b) The traces of performance in relational inference and inference score in transitive inference for 8 ARDMO instances. The error bars are standard errors over the domains. (c,d) The total rewards (test data) by ARDMO instances trained with different update rates or different state dimensions. The error bars are standard errors over the model instances.
C ADDENDUM ON THE EXPERIMENT IN SECTION
C.1 ANALYSIS METHOD
Our aim here is analyze the internal representations in trained model in way comparable to the brain analysis described in Park et al. 2021
First, in the fMRT experiment (Park al., 2021 J, they analyzed the hexagonal modulation property of the brain signals from the entorhinal cortex during two consecutive visual stimuli corresponding two grid positions. To simulate this, after running the model in given test domain, we obtained the state vector corresponding to the image at each grid position by using inference through the memory (equation 4 and took the average over the two state vectors. The values in state values are normalized as z-scores. By plotting the state values in each dimension against the grid direction, ie.. the angle between the grid positions (Figure 7 we obtained function over angles. Repeating this for different dimensions, we obtained a set of such functions. To see if those functions had periodicity in a given cycle deg, deg, 60 deg, 51.4 deg. and 45 deg corresponding to 4-, 5-, 6-, 7-, and 8-fold periodicities), with that each function may have a different phase. In Park et al. (2021). to avoid contamination of phase and periodicity, they estimated the phases by using another dataset obtained from separate experiment with the same task condition. We simulated this by using the state vectors obtained through the model at 25 steps prior to the completion of rumning in the test domain (which has different memory matrix). To estimate the phase, following the method in Park et al. (2021), we fit state values in each dimension against the sin and cos of the angle and took the arc-tangent of the regression coefficients. We then calculated the circular mean over the obtained phases. Finally, we took the then averaged the resulting functions for all domains.
visualize the average over the functions, where needed to deal
average of the functions for all dimensions separately for each domain, and
Second, in the same experiment Park al., (202T), they analyzed the distance coding in the brain signals from the hippocampus and entorhinal cortex. To simulate this, we similarly obtained the intermediate or state vectors corresponding to the images at each pair of grid positions (equation 4 and calculated the distance between these their dissimilarity. Although (Park et al.) 2021 used Mahalanobis distance, we used Euclidean distance since our case uses the deterministic state inference. By plotting the dissimilarity against the grid distance, i.c., Euclidean distance between the grid positions (Figure we determined the increase or decrease tendency by the correlation coefficient (positive negative).
C.2 ADDITIONAL RESULTS
We trained ARDMO instances with larger (H 500) smaller (H 30) memory. Figure summarizes the results for these in the same format as Figure 3 (showing model #1). These show overall robustness of the results across model instances.
16

---

Published as a conference paper at ICLR 2024
direction
distance
Figure Grid coding
(a) Performance ramp (H 500)
(b) Performance ramp (H = 50)
0.0
1-500
1-500
2.3
1-52,
0.5
1-500 #5
3.0
4-15
1-35.
8.7
1-500.
3.6-
05
Mw Maw. yesev
1956
1500
2039
1000
1002
2006
Saper
(c) Hexagonal modulation (H 500)
(d) Hexagonal modulation
50)
0.30
1-000
5-500
5.30
1-50
F-50 23
man won.atto
0.85
--00
0.05
1e50 84
0.80
5.00
-2.05
-3.35
PM)
30
300
120
200
250
302
50
200
:30
200
220
300
discten
drscton loog)
(c) Periodicity (H 500)
(f) Periodicity (H 50)
--00
1-50 #1
2000
0.00
0.26
5.06
50
P450
11-200
1-:50 80
nall
0.00
nut
-2.02
-0.02
(g) Distance coding (H = 500)
(h) Distance coding (H 50)
D407
042
0402
100
Figure 8: Results on the -wo-dimensional hierarchy task, analogous Figure from ARDMOs with larger (H = 500) or smaller (H = 50) memory.
17

---

Published as a conference paper at ICLR 2024
D LIMITATIONS
Currently, we know the following technical limitations of ARDMO.
The model struggles in learning a larger scale of relational structures. For example, the performance notably drops for 20 entities or larger in the ID case, and for 5 x 5 or larger in the 2D case.
The model can handle the case with 3 or more relations (without axis hint like in our 2D case) but is limited to small scale of relation structures. For example, in the case with relations, learning becomes significantly slower for 4 entities and fails for larger cases.
E PLATFORM DETAILS
We used computers with the following specifications: (1) 56 core CPUs (256G memory) with 4 V100 GPUs (16G memory each), (2) 56 core CPUs (256G memory) with 4 V100 GPUs (32G memory each), and (3) 96 core CPUs (256G memory) with 4 A100 GPUs (40G memory cach). All code is implemented with Python (3.9.13) Pytorch (1.12.1).
18


Reviews:
[{"review_id": "LweKBM1yLv", "reviewer": "Reviewer_RHcS", "summary": "This paper proposes a new cognitive model for performing memory-based decision-making tasks. The main contribution is a learning algorithm that allows the model to learn abstract relationships from reward-guided relational inference tasks, while maintaining dynamic binding between these abstract relations and concrete entities in a given task using a memory mechanism. The experiments demonstrate the model's ability to capture relational structures in one-dimensional and two-dimensional hierarchies. The authors also show that the model exhibits both performance and internal representations that bear resemblance to human behavioral and fMRI experimental data.", "questions": "- What properties does the relation matrix possess, could you offer insights into them?\n- Is there a separate MLP for each $g_a$? If yes, how do you ensure the probabilities sum to 1?\n- In equation 10, does $m$ in $\\psi_a^{m-1}$ correspond to the $m^{\\rm th}$ power? Or is it the $m^{\\rm th}$ iteration? If the latter, how is  $\\psi_a^{m-1}$ updated? \n- In the formula for the inference score in section 4.1, what is $c^{ti}$?\n- How was the value $\\alpha=0.7$ chosen?\n- In the caption of Figure 4, you mention that the for NTM and DNC the horizontal axis is 50-times reduced for readability. Does this mean that they used 50-times more epochs?\n- How were the hyperparameters selected for the baseline methods in Figure 4a? \n- In Figure 4b, why is the performance, after approximately 150 steps, slightly worse after 8000 epochs compared to the performance after 1000 epochs?\n- What is the effect of S (the length of the state vector) on the results?\n- In section 4.2, you discuss the learned intermediate representation $h$. I'd like to clarify the definition, since there appear to be two distinct uses of $h$ in equations 3 and 4.\n- Could you please clarify how the state values are computed in section 4.2? \n- Are there any thoughts about how this approach can be extended to more general relational graphs or to scenarios with sparse rewards?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "This paper introduces an interesting cognitive model for acquiring abstract relationships through reward-guided relational inference tasks. The experiments showcase the model's ability in learning relational structures that exhibit generalization across novel domains, featuring previously unseen entities. Notably, it significantly outperforms baseline models, such as LSTMs and standard Transformers. Further, the authors reveal an intriguing alignment between the model's behavior and fMRI data from humans performing the same tasks.", "weaknesses": "While the overall presentation of the paper is good, there are a couple of sections that are not easy to follow. The results on the two-dimensional hierarchy (section 4.2) can be challenging to understand for someone not very familiar with the findings in Park et al. (2021). Additionally, the notation in the section on transitive inference (3.4) can be a bit confusing (please see questions below).\n\nThe paper also misses references to related works on models for cognitive maps [1, 2]. Notably, [2] provides a unifying explanation for multiple hippocampal observations, while [3] presents an interesting approach for the reuse of learned abstractions in the form of graph schemas. It would be helpful to discuss the relationship between the approach in this work and these previous works.\n\nMinor: There are several grammatical errors and a few typos (e.g., 'Maharanobis' on page 15) scattered throughout the paper.\n\n\n[1] George, D., et al., 2021. Clone-structured graph representations enable flexible learning and vicarious evaluation of cognitive maps. Nature communications, 12(1), p.2392.\n\n[2] Raju, R.V., et al.., 2022. Space is a latent sequence: Structured sequence learning as a unified theory of representation in the hippocampus.\n\n[3] Guntupalli, J.S., et al., 2023. Graph schemas as abstractions for transfer learning, inference, and planning. arXiv preprint arXiv:2302.07350.", "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": "", "extra_scores": {}}, {"review_id": "vVchYyza6T", "reviewer": "Reviewer_rY1U", "summary": "This paper describes a new neural model that learns abstract relations for one-dimensional and two-dimensional set orderings. The model is trained and tested on sets from different domains, such that  the testing phase is identical to human experiments on learning transitive relations. \nThe model is demonstrated to generalize to unseen domains, while several stat-of-the art methods are not able to solve their task. This generalization is made possible by the model's architecture with two components - (1) a set matrixes for learning abstract relations (one per relation), and set of  memory matrixes for binding concrete tokens from a specific domain to the abstract relations. The model is cognitively plausible, as its performance during the test phase appears to be similar to human performance in relation learning experiments, while other state-of-the-art models fail to complete the task.", "questions": "The section analyzing hexagonal modulation within the model was not entirely clear to me -- it wasn't clear why the authors used the specific method of averaging state vectors.  Is there a citation, or maybe some explanation rationalizing this method? This can be included in the supplement. What is the significance of this hexagonal modulation emerging, given that the model's architecture is fundamentaly different from biological brains? Why would the authors expect hexagonal modulation to emerge?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. This paper makes a novel contribution, significantly improving on existing models.  \n\n2. Model testing based on simulating previous human studies strongly supports cognitive plausibility of the model.", "weaknesses": "I found the presentation to be poorly readable in places - this is not a big deal, but I would suggest editing for clarity.", "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": "", "extra_scores": {}}, {"review_id": "86xBcUtm3u", "reviewer": "Reviewer_k38n", "summary": "This paper presents a computational model for learning continuous spaces.  The primitives of the model include relationships, a, as well as entities x (just vectors in R^N).  To learn a one-dimensional space, the model is given two relations (analogous to greater than and lesser than); to learn a 2-dimensional space the model is given four relations (analogous to left vs right and above vs below).  The model's given a set of training data using near neighbors and then asked to generalize to pairs that were not observed (transitive inference).  And after learning a set of relationships, the model can generalize in that after learning a 1-D relationship among one set of stimuli, the model can more rapidly learn the analogous relationship among a second set of stimuli.  Notably, other approaches that are widely used in computational neuroscience (Tolman Eichenbaum machine, neural Turing machine, LSTM etc) not only fail to show these properties but can't learn the problem in the first place.", "questions": "The observation that transformers (for instance) don't learn these tasks is interesting.  Presumably, though this model is ill-suited for, say, language modeling.  What other problems can this computational approach solve (preferably in the general field of AI/ML)?   \n\nHow does this approach scale?  As the number of continuous dimensions goes up how does it behave?  Suppose you chose a different way to tile the plane.  Rather than placing items at grid coordinates, suppose each item had N near neighbors (or that the exemplars were irregularly scattered).  This would mean that the number of relations a has to grow.  How sensitive is this model to the number of relations (controlling for the dimension of the space)?  Does it depend on a regular tiling of the space with entities?\n\nAn alternative approach is to simply assume that the brain is organized to represent low dimensional spaces and the learning problem is to determine how to map those internal spaces onto the external world. E.g.,\nhttps://doi.org/10.1109/IJCNN54540.2023.10190998\nhttps://doi.org/10.1109/IJCNN54540.2023.10191578", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "As far as I'm aware this is a completely novel approach.   \n\nThe construction of ``internal spaces'' is an absolutely fundamental in computational cognitive neuroscience. \n\nThe idea of separating entities from relationships could be on the right track.", "weaknesses": "I wonder if it's possible to get TEM/transformer/NTM/etc to do something like this task if it's presented differently.  \n\nI found the connection to neuroscience very indirect, notwithstanding the observation that the similarity of internal states exhibits a distance effect and there's evidence for 60 degree symmetry.  This model is very abstract.", "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": "", "extra_scores": {}}, {"review_id": "DI0QT6hLrX", "reviewer": "Reviewer_NM51", "summary": "This work builds a model that learns generalizable abstract relational structure on a decision making task where one has to answer the relation between pairs of stimuli for a reward. The stimuli can have \"many to many\" relations in that each stimuli can have up to 4 relations with other stimuli. The model builds on works like TEM and NTM where an explicit external memory module is written to/read from. The authors tried a one-dimensional version of the task where stimuli only had two relations and a two-dimensional version of the task where there can be more relations. The authors then saw that the model reproduced some neural phenomenon on such relational tasks such as distance coding, hexagonal modulation, and distance coding.", "questions": "* Would it be possible for the model to learn these relations implicitly without specifically being rewarded for them? Sometimes for humans, abstract relational structure can often be learned in service of doing a specific task, rather than being trained specifically on finding the correct relations.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "* Very comprehensive review of past work. \n\n* Experiments are rigorous (multiple seeds, etc) and well-done. Showing the reproduction of the neural phenomenon is pretty nice. \n\n* Model is written clearly.", "weaknesses": "* There is extensive discussion of previous work, but I was left wondering what exact contributions this model makes over other models in this space like TEM. The paper discusses numerous differences, but I would like to see explicit discussion of what this model brings to the table, what specific phenomenon that this model produces that other models don't, etc. There are maybe some signs of this throughout the paper, but I didn't see any explicit discussion on it. \n\n* There's no limitations section/paragraph, which is an important part of any iclr paper.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)