# Understanding AI Cognition: A Neural Module for Inference Inspired by Human
Memory Mechanisms

Xiangyu Zeng, Jie Lin, Piao Hu, Ruizheng Huang, Zhicheng Zhang
University of Electronic Science and Technology of China
<EMAIL_ADDRESS>
<EMAIL_ADDRESS>Corresponding author

###### Abstract

How humans and machines make sense of current inputs for relation reasoning
and question-answering while putting the perceived information into context of
our past memories, has been a challenging conundrum in cognitive science and
artificial intelligence. Inspired by human brain’s memory system and cognitive
architectures, we propose a PMI framework that consists of perception, memory
and inference components. Notably, the memory module comprises working and
long-term memory, with the latter endowed with a higher-order structure to
retain more accumulated knowledge and experiences. Through a differentiable
competitive write access, current perceptions update working memory, which is
later merged with long-term memory via outer product associations, averting
memory overflow and minimizing information conflicts. In the inference module,
relevant information is retrieved from two separate memory origins and
associatively integrated to attain a more comprehensive and precise
interpretation of current perceptions. We exploratively apply our PMI to
improve prevailing Transformers and CNN models on question-answering tasks
like bAbI-20k and Sort-of-CLEVR datasets, as well as relation calculation and
image classification tasks, and in each case, our PMI enhancements
consistently outshine their original counterparts significantly. Visualization
analyses reveal that memory consolidation, along with the interaction and
integration of information from diverse memory sources, substantially
contributes to the model effectiveness on inference tasks.

##  1 INTRODUCTION

Cognitive science, neuroscience and AI (artificial intelligence) collectively
advance our grasp of intelligence, defined as the general mental abilities of
perception, memory and reasoning, each with a unique role in human cognition.
To construct more human-like intelligent systems, often referred to as the
standard model of the mind (Laird et al., 2017), it is imperative to delve
into the interactions among perception, memory, and reasoning in a unified
system. Recently, scholars have uncovered a significant flaw in previous deep
learning architectures: the absence of dedicated memory module that is
critical for long-term information retention and relational reasoning. This
drawback becomes evident when considering the constraints of many intelligent
systems, which either exclusively concentrate on perception and reasoning or
intricately interweave computation with implicit memory. Therefore, many
memory-based studies have emerged, mainly focusing on designing item-based
memory models with recurrent neural networks (RNNs) (Hopfield, 1982;
Hochreiter & Schmidhuber, 1997; Dai et al., 2019; Ramsauer et al., 2020;
Schlag et al., 2021) and memory-augmented neural networks (MANNs) (Graves et
al., 2014; 2016a; Le et al., 2018; Liang et al., 2023).

Nonetheless, existing approaches expose four limitations: (i) Implicit memory
(hidden state) may gradually lose previous information as the model constantly
updates its weights to accommodate new inputs, which prevents reusing the
precomputed relations in sequential tasks (Vaswani et al., 2017; Santoro et
al., 2017; Devlin et al., 2018). (ii) The memory system is configured in one
of two forms: either as a singular memory unit without hierarchical
construction or as multiple separate memory components with identical data
structures, both of which struggle to align with human memory traits and
achieve robust generalization (Goyal et al., 2022; Dai et al., 2019; Jaegle et
al., 2021; Wu et al., 2022; Kang et al., 2023; Liang et al., 2023). (iii) The
memory-memory relation is either crude, expressed as weighted summation via
neural networks or dot product attention, or it undergoes intricate memory
transformation algorithms. (Vaswani et al., 2017; Santoro et al., 2018). (iv)
Memory exploitation is confined to rudimentary retrieval, whether it’s
content-based addressing (Wu et al., 2020; Goyal et al., 2022; Kang et al.,
2023) or explicit address (Graves et al., 2016b; Liang et al., 2023).
Arguably, modern MANNs have yet to develop general architectural frameworks
for learning both diverse memory components and how they should interact
internally and externally.

Multiple Memory Systems Theory (MMS) asserts that working memory (WM) and
long-term memory (LTM) stand as pivotal subassemblies of human cognitive
processes (Atkinson & Shiffrin, 1968; Baddeley & Hitch, 1974; Eichenbaum &
Cohen, 2004) , where the former serves to temporarily buffers and processes
data for current tasks, while the latter is responsible for the retention of
enduring knowledge and experiences. Additionally, the Global Workspace Theory
(GWT) (Baars, 1993; Dehaene et al., 2021) suggests a communication and
coordination scheme, in which disparate cognitive units write information into
a shared workspace that is broadcast to all modules, along with the notion
that write access is restricted.

Inspired by the MMS, GWT and cognitive theories, we assume that optimizing the
structure of memory module and its internal and external correspondence
mechanisms holds great promise in surmounting the extant restrictions.
Accordingly, we hypothesize a cognitive framework called PMI that consists of
perception, memory and reasoning modules, wherein memory is posited as a dual-
layer memory block featuring distinct communion principles of the inner and
outer. More concretely, in terms of its structure, WM exists separately from
LTM, with the latter possessing a higher-order structure to preserve intricate
patterns and relations. When it comes to interactions, there are two exterior
procedures: perception-based competitive writing and inference-oriented
information service, alongside one inner channel—designed to establish
heterogeneous associations among the two memory units to facilitate efficient
information filtering, storage, and knowledge consolidation. We apply modern
different neural network architectures like Transformers attention-based
(Vaswani et al., 2017; Brown et al., 2020) and convolutional networks (He et
al., 2016), which all equipped with our dual-memory module, to multifarious
tasks that may require both WM and LTM: text and visual question-answering,
visual relations calculation and image classification. Across all these tasks,
models integrated with our memory block consistently outperform their original
counterparts.

##  2 Method

###  2.1 Overview

An overview of our PMI framework is illustrated in Fig. 1(a), which contains
three pivotal components: perception, memory and inference (both potentially
learned). Given an input $X$ (e.g., text, an image, or an audio signal), it is
processed through a series of computational stages indexed by $t$ to derive
the cognitive understanding $U$ of the current perception, as outlined below:

  1. 1.

P component: (Perception) — Convert the incoming input $X$ to an internal
feature representation $H=\mathcal{P}(X)$.

  2. 2.

M component: (Memory) — Update old memories given the input representation
$H^{t-1}$: $M^{t}=\mathcal{M}(H^{t-1},M^{t-1})$.

  3. 3.

I component: (Inference) — Reason (interpret) the current content given the
updated memory: $U=\mathcal{I}(H^{t-1},M^{t})$.

In this framework, trainable parameters are learned through backpropagation,
while memory blocks are updated solely through the forward process, which
constitute the process of memory precipitation through multiple iterations. A
more elaborate description of our method is presented as follows.

(a) PMI framework

(b) Memory update

(c) Inference component

Figure 1: Model overview and the process of grasping the current input at
calculation step $t$. (a) The memory module consists of WM $M_{w}$ and LTM
$M_{l}$, each characterized by distinct data structures. (b) WM is updated by
current perception via a differentiable and constrained write access, which is
then integrated into LTM through outer product association. (c) The inference
component retrieves pertinent data from both WM and LTM using content-based
addressing MHC and MHSC. Subsequently, through integration steps, it
consolidates info from these sources to generate fresh insights into the
input, which is used for next rounds of inference or to directly support the
decision-making process.

###  2.2 perception

The perceptual operation maps the original input data to internal entity
representations. Focusing on the prevailing models and taking Transformers as
an example, text inputs undergoes embedding and positional encoding to yield
initial feature representation $h^{0}\in\mathbb{R}^{T\times D}$, where $T$ is
the sequence length of $D$ dimension. In the ViT model, a $2D$ image $x$
$\in\mathbb{R}^{H\times W\times C}$ is split into $N$ patches
$x_{p}\in\mathbb{R}^{N\times(P^{2}\cdot C)}$, each of which is linearly
embedded. Then positional embeddings are added to obtain the final embedding
vector $h_{0}$ = [ $x_{class}$ ; $x_{p}^{1}$ $E$; $x_{p}^{2}$ $E$; $\cdots$ ;
$x_{p}^{N}$ $E$]+ $E_{pos}$ , $E$ $\in\mathbb{R}^{(P^{2}\cdot C)\times D}$,
$E_{pos}$ $\in$ $\mathbb{R}^{(N+1)\times D}$, which contains fundamental
feature information of the image, such as color and texture, where $(H,W)$ is
the resolution of the original image, $C$ is the number of channels and
$(P,P)$ is the resolution of each image patch. This process resembles the
human perceptual system that receives external information and converts it
into understandable internal representations, laying a foundation for
subsequent memory and reasoning.

###  2.3 Global Shared Dual-Level Memory

This section provides a detailed exposition of the proposed dual-layer memory
module and the internal and external communication mechanisms utilized for its
update, as illustrated in Fig. 1(b). We posit that the memory module should be
globally shared, including the working memory $M_{w}$, which temporarily
stores and processes information required for the current task, and the long-
term memory $M_{l}$, which persistently stores knowledge and experiences. We
opt for $M_{w}\in\mathbb{R}^{N\times D_{m}}$ as a form of WM, which is indexed
based on $m_{i}$ slot. Here, $N$ is the number of slots, each with a dimension
of $D_{m}$. While the LTM is represented as a $3D$ structure
$M_{l}\in\mathbb{R}^{C\times N\times D_{m}}$, with $C$ denoting memory
fragments, which is biologically plausible (Marr & Thach, 1991).

####  2.3.1 External channel: $\mathcal{M}_{w}$-Write

External communication serves to update the contents of WM via two pivotal
steps: competitive writing and forgetting, which informed by a fundamental
aspect of the human memory system — our inclination not to retain a permanent
record of every perception but rather to discerningly preserve essential
elements within memory. Collectively, these processes guarantee the storage of
the most critical information pertinent to the ongoing task, an indispensable
facet in tasks involving reasoning.

Write with competition This process aims to selectively inscribe perceived
inputs into $M_{w}$ with finite capacity, also inspired by Miller’s Law, which
states that the number of information units that human WM can handle
simultaneously is limited, often around $7\pm 2$. We use a multi-headed sparse
cross-attention mechanism (MHSC) for this execution, as expressed in Eq. 1, 2.
Cognate to the MH mechanism used in Transformers, but MHSC exhibits two
distinctive aspects: (i) it necessitates separate origins for Q and K and (ii)
it introduces a sparsity-inducing operation on the attention weight matrix.
Specifically, the result of the $t-1$ step $h^{t-1}\in\mathbb{R}^{T\times D}$
is projected into keys and values, along with
$M_{w}^{t-1}\in\mathbb{R}^{N\times D_{m}}$ that is projected into queries. The
current inputs compete to write through our MHSC, in conjunction with some
other operations to yield the intermediate state $\widetilde{M}_{w}^{t}$ . The
whole formulas are as follows:

$\displaystyle
s_{k}=softmax\left(\frac{M_{w}^{t-1}W^{Q}(h^{t-1}W^{K})^{T}}{\sqrt{d^{K}}}\right)$
(1) $\displaystyle\widetilde{M}_{w}^{t}=s_{k}^{*}h^{t-1}W^{V}$ (2)
$\displaystyle\widetilde{M}_{w}^{t}={LN}_{1}(\widetilde{M}_{w}^{t}+M_{w}^{t-1})$
(3)
$\displaystyle\widetilde{M}_{w,i}^{t}={ReLU}({MLP}_{i}(\widetilde{M}_{w,i-1}^{t})),\quad
i\in\\{1,\ldots,k\\}$ (4)
$\displaystyle\widetilde{M}_{w}^{t}={LN}_{2}(M_{w}^{t-1}+\widetilde{M}_{w,k}^{t})$
(5)

It’s noteworthy that the input needs to be linearly projected to the same
dimension $D_{m}$ as $M_{w}^{t-1}$ (following the traditional practice of
$D=D_{m}$). $W^{Q}$, $W^{K}$ and $W^{V}$ are weight matrices.
$s_{k}\in\mathbb{R}^{N\times(T+N)}$ is the attention weight scores of
$M_{w}^{t-1}$ and $h^{t-1}$. Unlike the standard soft competition, we use a
top-k softmax (Ke et al., 2018) to select a fixed number of entities for
updating the $M_{w}$. $s_{k}^{*}$ denotes the pre-softmax value, please
consult Algorithm 1 for details. ${LN}_{1}$ and ${LN}_{2}$ signify different
LayerNorms, employed to uphold memory stability over prolonged time steps.
$ReLU$ is the $ReLU$ function, ${MLP}_{i}$ is the $i^{th}$ multilayer
perceptron and $\widetilde{M}_{{w},k}^{t}$ is the intermediate output through
$k$ multilayer perceptrons.

Forgetting Memory forgetting entails the elimination or reduction of
previously stored data to make space for new info, optimizing memory
performance. It is reasonable to adopt the gating mechanism since it emulates
the biological memory process and effectively alleviates information
conflicts. This is implemented in Eq. 6, where $I_{t}$ and $F_{t}$ indicate
the input and forget gates respectively, as proposed in RMC (Santoro et al.,
2018). Further details can be found in Appendix C.1.

$M_{w}^{t}=F_{t}(M_{w}^{t-1},h^{t-1})\odot
M_{w}^{t-1}+I_{t}(M_{w}^{t-1},h^{t-1})\odot\widetilde{M}_{w}^{t}$ (6)

####  2.3.2 Internal channel: $\mathcal{M}_{l}$-Write

The internal channel is utilized to update LTM, which boasts a larger capacity
to accommodate more info. As illustrated in Eq. 7, we conduct an outer product
calculation between the updated $M_{w}^{t}$ and the previous-step LTM
$M_{l}^{t-1}\in\mathbb{R}^{C\times N\times D_{m}}$ to merge novel vital info
into the current LTM $M_{l}^{t}$. In contrast to scalar product computation
that only yield a numerical value, the outer product operation is used to
capture relations and interactions between vectors, which not only enhances
higher-order representational capacity but also contributes to information
precipitation and memory reinforcement.

$M_{l}^{t}={LN}_{3}\left((M_{w}^{t}\otimes M_{l}^{t-1})+M_{l}^{t-1}\right)$
(7)

Here, ${LN}_{3}$ denotes LayerNorm, and $\otimes$ signifies the outer product
operation.

###  2.4 Inference

Inference component, guided by the updated memories, provides insights of
current perceptions. Our interpretation of the inference is that it stems from
an assumption on the form of the joint distribution between perceptual inputs
and current memory. To mimic human-like ability to focus on crucial details of
the ongoing task while leverage extensive knowledge and experience to navigate
complex situations, we use content-based addressing MHC that is equivalent to
MHSC without sparsity and MHSC to retrieve relevant memories from $M_{w}^{t}$
and $M_{l}^{t}$ based on current input $h^{t-1}$, getting $h_{w}^{t}$ and
$h_{l}^{t}$ respectively, as shown in Eq. 8-10.

$\displaystyle U_{w}^{t}$
$\displaystyle=\textit{MHC}\left(h^{t-1}\widetilde{W}^{Q},M_{w}^{t}\widetilde{W}^{K},M_{w}^{t}\widetilde{W}^{V}\right)$
(8) $\displaystyle\widehat{M}_{l}^{t}$
$\displaystyle=\frac{1}{C}\sum_{i=1}^{C}M_{l}^{t}[i,:,:]\quad
where\quad\widehat{M}_{l}^{t}\in\mathbb{R}^{N\times D_{m}}$ (9) $\displaystyle
U_{l}^{t}$
$\displaystyle=\textit{MHSC}\left(h^{t-1}\widehat{W}^{Q},\widehat{M}_{l}^{t}\widehat{W}^{K},\widehat{M}_{l}^{t}\widehat{W}^{V}\right)$
(10)

Subsequently, the understanding $U_{l}^{t}$ from LTM serves to further revise
and supplement the understanding $U_{w}^{t}$ from WM via the MHC mechanism,
where $U_{l}^{t}$ creates queries that match with keys and values from
$U_{w}^{t}$ to generate a richer representation $U_{wl}^{t}$. Then a linear
combination of $U_{w}^{t}$ and $U_{wl}^{t}$ is conducted with a hyper-
parameter $\alpha$ to yield the final cognition $U^{t}$, as shown in Eq. 12.
This process of multiple correlation and fusion of various information sources
contributes to extract richer and more valuable insights that support higher-
level decision-making and reasoning.

$\displaystyle U_{wl}^{t}$
$\displaystyle=\textit{MHC}\left(U_{l}^{t}\bar{W}^{Q},U_{w}^{t}\bar{W}^{K},U_{w}^{t}\bar{W}^{V}\right)$
(11) $\displaystyle U^{t}$ $\displaystyle=\alpha
U_{w}^{t}+(1-\alpha)U_{wl}^{t}$ (12)

##  3 RELATED WORK

Cognitive Science In cognitive neuroscience, memory studies endeavor to
unravel the intricacies of information storage, organization and retrieval in
brains, and their profound impact on thinking, cognition and behavior,
building on the pioneering work of Ebbinghaus (1885) and Bartlett & Bartlett
(1932). Afterwards, Atkinson & Shiffrin (1968) proposed a multi-store model
including sensory, short-term and LTM, which contributes to our insights of
different memory types and stages. The successor Baddeley & Hitch (1974)
further refined and delineated this model by substituting short-term memory
with WM—a transient storage that can interact with LTM. Sigma (Rosenbloom et
al., 2016) and Soar (Laird, 2019) are canonical cognitive frameworks of recent
advancements, both of which employ a similar memory system comprising WM and
LTM that play crucial roles in complex reasoning and problem-solving tasks.
Moreover, the Global Workspace Theory (Baars, 1993) put forward a coordination
and collaboration mechanism with restricted write access, which sheds light on
the interaction of diverse cognitive components.

Memory networks Semi-parametric MANNs, as a form of using implicit knowledge
to perform complex reasoning tasks, are a persistent theme in neural network
research. Today MANNs typically rely on explicit memory and attention
mechanisms, with pioneering models like Memory Networks (Weston et al., 2014)
and Neural Turing Machines (NTMs) (Graves et al., 2014), both of which are
equipped with a storage for vector representations accessible via attention.
Memory Networks use addressable memory to execute tasks through a series of
read operations. In contrast, NTMs also utilize addressable content storage,
but unlike Memory Networks, which pre-load memories using all the inputs, NTMs
write and read the memory one input at a time. Following this are
Differentiable Neural Computers (DNC) (Graves et al., 2016a) and Sparse DNC
(Rae et al., 2016), which are realized as recurrent neural networks (RNNs)
capable of read and write operations on memory over time and are trained via
BPTT (Werbos, 1990). A parallel research path involves enhancing RNNs like
LSTM by incorporating data structures such as lists, stacks or queues (Joulin
& Mikolov, 2015; Grefenstette et al., 2015).

Transformers with memory extensions Memory is a topic of active exploration in
diverse Transformer studies. Transformer-XL (Dai et al., 2019) and its
successors, Compressive Transformer (Rae et al., 2019), RMT (Bulatov et al.,
2022) and Scaling Transformer (Bulatov et al., 2023) re-introduce the notion
of memory and recurrence by caching self-attention hidden states from each
layer into a fixed-size queue and reusing them in subsequent attention
computations, with the difference that Compressive Transformer utilizes a
compression network to further compress its memories into fewer vectors. In
addition, various forms of global representations are introduced as a model
memory that learns to gather information from input sequence tokens. Notable
examples of these approaches include Set Transformers (Lee et al., 2019), ETC
(Ainslie et al., 2020), Longformer (Beltagy et al., 2020) and TR+HSW (Goyal et
al., 2022), all of which redesign the self-attention mechanism to reduce
computational complexity. Memory modules, with their read-write global memory
operations, have recently attracted attention for their potential to remember
prior information, driving a movement towards more structured models. For
instance, Memformer (Wu et al., 2020) proposes a dedicated external dynamic
memory module after the primitive self-attention layer and interacts with it
through memory reader and writer components to store previous hidden states in
concise representations for efficient sequence modeling. More recently, DT-Mem
(Kang et al., 2023) introduces a WM that contains N memory slots between the
Transformer module and the MLP to store and retrieve information through an
attention-based approach, where the Transformer module is similar to the GPT-2
(Radford et al., 2019) module without the feedforward layer. Most pertinent to
our work, Goyal et al. (2022), taking cues from the GWT theory, replace
Transformers’ pairwise interactions with a shared workspace featuring
constrained write access—a concept equivalent to our WM that can read and
write. While these endeavors are closely related to explicit memories, their
memory structures are monolithic, which leads to boundaries in representing
certain higher-order information or relations. Hence, one takeaway from our
work is that it may be prospective to revisit previous memory enhancement
methods in light of insights from cognitive science into memory structures.

##  4  EXPERIMENTS

To assess the efficacy of the PMI module in discovering and learning inference
entities and their relations, we conduct a preliminary exploration by
incorporating it as a replacement for the pairwise self-attention layers in
Transformers and ViT (Dosovitskiy et al., 2020), where memory components are
shared globally. This modified architecture, called MITR, are then applied to
a diverse range of tasks, including visual QA, text-based QA, image
classification and visual relations calculation. Readers can refer to the
Appendix D E for full details on each task and details on hyperparameter
settings for the model.

###  4.1 Relational Reasoning : Sort-of-CLEVR

Sort-of-CLEVR (Santoro et al., 2017) is a dataset similar to CLEVR, designed
specifically for research on relational reasoning. Each 2D image in Sort-of-
CLEVR is of size 75 × 75 and comes with 6 randomly placed geometric shapes of
6 possible colors and 2 possible shapes. There are 10 non-relational and 20
relational questions that are equally divided into binary and ternary types
per image, along with corresponding answers (details in Appendix E.2). Given
the bounded answer space, this task is treated as a classification task. Each
image is partitioned into a sequence of uniform patches and then encoded as in
ViT. Subsequently, we concatenate the image embedding with its corresponding
question embedding as input into our MITR, in line with Goyal et al. (2022).

For this task we evaluated our MITR with the following five baselines:
Transformers [TR] (Vaswani et al., 2017), Set transformer [ISAB]: Transformers
where self-attention is replaced by ISAB module (Lee et al., 2019),
Transformers with Shared Workspace with top-k competition [TR+HSW] (Goyal et
al., 2022) and High Capacity Transformers [TR+HC]: Same as TR but with
different parameters across layers.

(a) Unary Accuracy

(b) Binary Accuracy

(c) Ternary Accuracy

Figure 2: Test accuracy vs training iterations for the Sort-of-CLEVR task.

The test accuracy curves over 200 training epochs are illustrated in Fig. 2.
We observe that transformers equipped with our global shared memory module
converge faster compared to all other baselines, demonstrating superior
performance on both relational and non-relational tasks. In contrast, the
TR+HSW model excels in addressing non-relational questions but struggles with
relational problems. We conjecture this might be because non-relational
problems frequently demand the model to tackle only small amounts of
information about individual objects. Interestingly, the single-level memory
slots in the global workspace (similar to our WM) possess the capability to
store and process scant present information, allowing them to handle these
issues with ease. However, relational questions regularly necessitate multi-
step reasoning to obtain an answer, such as extracting object attributes
followed by relational analysis. The introduction of LTM enables the model to
deposit crucial information during the learning process. Consequently, it can
retrieve pertinent knowledge from this memory module in upcoming reasoning
steps, going beyond its reliance solely on the current input. This contributes
to a more comprehensive understanding and handling of relational questions.

###  4.2 Text-based QA : bAbI

Table 1: Test error rates: mean ± std. (in %) on the 20 bAbI tasks for models
trained using 10k examples and best error over 10 runs. † is reported from
Dehghani et al. (2018)

Model | Error
---|---
Mean | Best
LSTM (Hochreiter & Schmidhuber, 1997) | 27.3±0.8 | 25.2
TR† (Vaswani et al., 2017) | 22.1 | N/A
DNC (Graves et al., 2016b) | 12.8±4.7 | 3.8
H-Mem (Limbacher & Legenstein, 2020) | 10.8 | N/A
NUTM (Le et al., 2020a) | 5.6±1.9 | 3.3
MemNet (Dou & Principe, 2023) | 5.6 | N/A
TR+HSW (Goyal et al., 2022) | 3.6±0.46 | 3.25
MITR (ours) |  2.55±0.11 | 2.32

BAbI is a pure text-based QA dataset (Weston et al., 2015) that is widely used
to assess the ability of MANNs, attention mechanisms and other types of models
to remember and reason on textual information. This dataset contains 20
challenging tasks, each corresponding to a particular type of reasoning, such
as logical deduction, counting, pathfinding and induction, which possibly
require both WM and LTM. Each question is associated with a set of supporting
facts. For example, the facts “John journeyed to the office” and “John left
the milk” support the question “Where is the milk?” (answer: “office”) (more
in Appendix E.1). Following Le et al. (2020b), each story is preprocessed into
a sentence-level sequence, which is fed into our MITR model as the input
sequence. A model succeeds on a task if its performance surpasses 95%. We
compare our model with recent memory networks and report the results in Table
1 (more in Appendix F.2).

###  4.3 Detecting Equilateral Triangles

Figure 3: Detecting Equilateral Triangles. This figure compares the
performance of Transformers with our PMI [MITR] against other Transformer
baselines.

In this binary classification task, our goal is to determine whether a 64 × 64
sized image contains an equilateral triangle composed of three randomly placed
point clusters (Ahmad & Omohundro, 2009). For equilateral triangles, the
midpoints of these clusters are equidistant from each other. To feed an image
into our MITR, we adopt the same methodology as employed in ViT (Dosovitskiy
et al., 2020). Specifically, each image is divided into equally sized 4 × 4
patches, which are then utilized as distinct input positions for the MITR. In
order to make precise judgments, this task requires the model to adeptly
comprehend and memorize the spatial relations between disparate point
clusters, embodying the relative positions and distances among them. By
incorporating our PMI module, with shared WM and LTM across all layers, the
model can preserve decisive info concerning each point cluster for subsequent
inference procedures. Moreover, the constrained capacity of WM compels the
model to selectively inscribe crucial information into the memory module,
which coincides favorably with the inherent sparsity intrinsic to the task.

The results in Fig. 3 reveal that MITR outperforms the standard TR model in
terms of both convergence speed and accuracy, converging faster (curves
available in Appendix F.1) and achieving an impressive accuracy of 97.9%, an
8.1% improvement. Additionally, our approach surpasses other baselines,
further confirming the efficacy of the PMI module. Here, STR denotes
Transformers with sparse factorizations of the attention matrix (Child et al.,
2019), [MITR+S] is a variant of the MITR without top-k sparsity, and other
baselines are detailed in experiment 4.1.

###  4.4 Image Classification : Cifar-10

CIFAR-10 is benchmark image dataset commonly used in the field of computer
vision, which consist of 50k training and 10k test images of resolution 32 ×
32 with the total number of classes 10. For this task, in addition to MITR
mentioned above, we also incorporate PMI into the convolutional series model.
Explicitly, the original images, after four convolutional layers, serve as
perceptions into our PMI module to obtain understandings, which then undergo
linear and softmax transformations to yield final classification results. The
performance of the best models on test sets is reported in Table 2, where
CNN_MI w  / o refers to CNN_MI without guidance from LTM. It’s obvious that
our MITR and CNN_MI models both exhibit superior performance, achieving
accuracies of 79.12% and 78.69% in CIFAR-10, respectively, with an improvement
of 2.94% (compared to TR) and 0.08% (compared to CNN_MLP). These results
further underscore the universality of our PMI module.

Table 2: Results of different models on CIFAR-10

Models | Trans. | Conv.
---|---|---
ViT | ISAB | TR+HSW | MITR (ours) | CNN_MLP | CNN_MI (ours) | CNN_MI w/o (ours)
Acc (%) | 76.18 | 76.39 | 76.28 | 79.12 | 78.61 | 78.69 | 78.63
Params (M) | 0.75 | 2.21 | 2.01 | 2.0 | 0.11 | 1.75 | 1.68

###  4.5 More Explorations of Memory Module

####  4.5.1 Memory Attributes and Communication Modes

In this section, we conduct a more qualitative exploration of memory
properties and the modes of internal and external communication on bAbI and
Sort-of-CLEVR tasks. The study on memory properties aims to address three
questions: (i) How does the capacity of WM and LTM impact model performance?
(ii) How does the global shared memory (persistence of memory) affect model
performance? (iii) Does knowledge consolidation occur in the shared LTM?
Experiments with communication modes aim to determine whether competitive
writing, as well as the correction and supplement of LTM-derived data to
relevant info obtained from WM is effective. To tackle these questions, we set
up three models of distinct sizes $MITR_{s}(l=4,h=4)$, $MITR_{m}(l=8,h=8)$ and
$MITR_{l}(l=12,h=16)$, and run them on various combinations of $N$, $M$ and
$k$, where $l$ and $h$ are the number of layers and heads in MITR,
respectively, considering their critical roles in model performance.

The results are reported in Table 3, where $MITR_{m}{w/o}_{1}$ denotes MITR
without memory sharing among its layers, while $MITR_{m}{w/o}_{2}$ indicates
that info retrieved from LTM is directly aggregated with data from WM via
$\alpha$ without correction step. soft is a standard soft competition mode,
not a top-k strategy. We can derive following key findings. For memory
properties, firstly, greater memory capacity doesn’t necessarily equate to
better performance. The optimal results are achieved at $N=8$ and $M=5$,
aligning with discoveries in cognitive neuroscience. Secondly, memory
persistence markedly improves the performance and speed of convergence in
relational inference tasks, especially in binary and ternary relations,
respectively, by 7.32%, 6.88% over non-globally shared cases (more in Appendix
B.1). Notably, independent memory modules result in an eightfold increase in
trainable parameters. Regarding the communication mode, constrained writing
exhibits heightened sensitivity in binary and ternary inference tasks, albeit
with contrasting effects. We speculate that this divergence may be attributed
to the larger volume of info storage required for ternary problems, thus
necessitating a slightly larger $k$ value. Moreover, without the guidance of
LTM, as an erudite scholar, there is a minor uptick in error rates for bAbI
task under the same setup, and the three types of Sort-of-CLEVR task exhibit
respective decreases of 0.31% (unary), 7.34% (binary) and 4.54% (ternary) in
accuracy, underscoring the constructive influence of previously accumulated
knowledge on relational reasoning.

Table 3: Results of ablation studies on memory properties and communication
modes.

Model N M Top-$k$ bAbI Sort-of-CLEVR Params Err% Params Unary% Binary%
Ternary% $MITR_{s}$ 6 3 5 2.00M 2.81 2.03M 99.14 77.12 61.29 8 5 5 2.27M 2.72
2.29M 99.45 86.06 62.85 8 5 7 2.27M 2.73 2.29M 99.50 82.84 64.35 10 7 9 2.53M
2.78 2.55M 99.19 80.24 59.48 $MITR_{m}$ 6 3 5 2.07M 2.61 2.09M 99.40 80.13
65.93 8 5 5 2.33M 2.55 2.36M 99.34 87.61 62.45 8 5 7 2.27M 2.57 2.36M 99.19
81.93 60.89 10 7 9 2.59M 2.62 2.62M 99.40 80.18 65.83 $MITR_{l}$ 6 3 5 2.20M
2.73 2.22M 99.40 81.92 64.21 8 5 5 2.46M 2.58 2.49M 99.14 84.73 65.52 8 5 7
2.46M 2.59 2.49M 99.29 80.68 66.94 10 7 9 2.73M 2.71 2.75M 99.09 81.47 65.01
$MITR_{m}{w\scalebox{0.8}{/}o}_{1}$ 8 5 5 16.48M 2.84 16.5M 99.14 80.29 60.06
$MITR_{m}{w\scalebox{0.8}{/}o}_{2}$ 8 5 5 2.33M 2.91 2.36M 99.19 79.96 62.40
$MITR_{m}$ 8 5 $soft$ 2.33M 2.75 2.36M 99.15 79.64 61.87

####  4.5.2 Visualizations of Attention Patterns

To explore whether knowledge accumulates in LTM, we use visualizations of
attention patterns between current perceptions and LTM on the bAbI task, shown
in Fig. 4. Here, current inputs act as queries, and the LTM matrix serves as
keys and values for cross-attention computation. As the depth increases, a
clear trend emerges in the heatmaps: more colored regions appear that
gradually stabilize and resemble, implying a growing correlation between
inputs and LTM that gradually converges (more explanations in Appendix B.2).
This may indicate that richer knowledge is accumulated in LTM, leading to a
more consistent grasp of different elements within the input data across these
layers.

(a) Attention patterns between inputs and WM

(b) Attention patterns between inputs and LTM

Figure 4: Attention patterns between perceptions and memories across different
layers of the MITR.

##  5 Conclusion

Inspired by multiple memory systems and global workspace theory in cognitive
neuroscience, we propose the PMI module containing perception, dual-layer
memory and inference components, to generate a more comprehensive and accurate
understanding of current inputs while depositing vital contents into memory to
cope with more complex situations. We have explored PMI’s dual utility: as an
alternative to self-attention layers in Transformers and as a complement to
convolutional networks. Extensive experiments on images and texts provide
compelling evidence for the effectiveness and adaptability of PMI, meaning it
could be a core component in diverse architectures. We look forward to a
broader application of our method, including its integration into various
frameworks and its extension to a wide range of tasks across varying
modalities.

## Ethics Statement

The authors do not foresee any negative social impacts of this work, but of
course the accumulation of improvements in ML could be misused as it may give
more power to nefarious agents.

## References

  * Ahmad & Omohundro (2009) Subutai Ahmad and Stephen M. Omohundro.  Equilateral triangles: A challenge for connectionist vision.  2009\.  URL https://api.semanticscholar.org/CorpusID:28176342.
  * Ainslie et al. (2020) Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang.  Etc: Encoding long and structured inputs in transformers.  _arXiv preprint arXiv:2004.08483_ , 2020.
  * Atkinson & Shiffrin (1968) Richard C Atkinson and Richard M Shiffrin.  Human memory: A proposed system and its control processes.  In _Psychology of learning and motivation_ , volume 2, pp. 89–195. Elsevier, 1968.
  * Baars (1993) Bernard J Baars.  _A cognitive theory of consciousness_.  Cambridge University Press, 1993.
  * Baddeley & Hitch (1974) Alan D Baddeley and Graham Hitch.  Working memory.  In _Psychology of learning and motivation_ , volume 8, pp. 47–89. Elsevier, 1974.
  * Bartlett & Bartlett (1932) Frederic Charles Bartlett and Frederic C Bartlett.  _Remembering: A study in experimental and social psychology_.  Cambridge university press, 1932.
  * Beltagy et al. (2020) Iz Beltagy, Matthew E Peters, and Arman Cohan.  Longformer: The long-document transformer.  _arXiv preprint arXiv:2004.05150_ , 2020.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Bulatov et al. (2022) Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev.  Recurrent memory transformer.  _Advances in Neural Information Processing Systems_ , 35:11079–11091, 2022.
  * Bulatov et al. (2023) Aydar Bulatov, Yuri Kuratov, and Mikhail S Burtsev.  Scaling transformer to 1m tokens and beyond with rmt.  _arXiv preprint arXiv:2304.11062_ , 2023.
  * Child et al. (2019) Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever.  Generating long sequences with sparse transformers.  _arXiv preprint arXiv:1904.10509_ , 2019.
  * Dai et al. (2019) Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov.  Transformer-xl: Attentive language models beyond a fixed-length context.  _arXiv preprint arXiv:1901.02860_ , 2019.
  * Dehaene et al. (2021) Stanislas Dehaene, Hakwan Lau, and Sid Kouider.  What is consciousness, and could machines have it?  _Robotics, AI, and Humanity: Science, Ethics, and Policy_ , pp. 43–56, 2021.
  * Dehghani et al. (2018) Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Łukasz Kaiser.  Universal transformers.  2018\.  URL https://arxiv.org/abs/1807.03819.
  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Dosovitskiy et al. (2020) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_ , 2020.
  * Dou & Principe (2023) Ran Dou and Jose Principe.  Universal recurrent event memories for streaming data.  In _2023 International Joint Conference on Neural Networks (IJCNN)_ , pp. 1–10. IEEE, 2023.
  * Ebbinghaus (1885) Hermann Ebbinghaus.  _Über das gedächtnis: untersuchungen zur experimentellen psychologie_.  Duncker & Humblot, 1885.
  * Eichenbaum & Cohen (2004) Howard Eichenbaum and Neal J Cohen.  _From conditioning to conscious recollection: Memory systems of the brain_.  Number 35. Oxford Psychology, 2004.
  * Goyal et al. (2022) Anirudh Goyal, Aniket Rajiv Didolkar, Alex Lamb, Kartikeya Badola, Nan Rosemary Ke, Nasim Rahaman, Jonathan Binas, Charles Blundell, Michael Curtis Mozer, and Yoshua Bengio.  Coordination among neural modules through a shared global workspace.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=XzTtHjgPDsT.
  * Graves et al. (2014) Alex Graves, Greg Wayne, and Ivo Danihelka.  Neural turing machines.  _arXiv preprint arXiv:1410.5401_ , 2014.
  * Graves et al. (2016a) Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska-Barwińska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al.  Hybrid computing using a neural network with dynamic external memory.  _Nature_ , 538(7626):471–476, 2016a.
  * Graves et al. (2016b) Alex Graves, Greg Wayne, Malcolm Reynolds, Tom Harley, Ivo Danihelka, Agnieszka Grabska-Barwińska, Sergio Gomez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al.  Hybrid computing using a neural network with dynamic external memory.  _Nature_ , 538(7626):471–476, 2016b.
  * Grefenstette et al. (2015) Edward Grefenstette, Karl Moritz Hermann, Mustafa Suleyman, and Phil Blunsom.  Learning to transduce with unbounded memory.  _Advances in neural information processing systems_ , 28, 2015.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 770–778, 2016.
  * Hochreiter & Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  _Neural computation_ , 9(8):1735–1780, 1997.
  * Hopfield (1982) John J Hopfield.  Neural networks and physical systems with emergent collective computational abilities.  _Proceedings of the national academy of sciences_ , 79(8):2554–2558, 1982.
  * Jaegle et al. (2021) Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira.  Perceiver: General perception with iterative attention.  In _International conference on machine learning_ , pp. 4651–4664. PMLR, 2021.
  * Joulin & Mikolov (2015) Armand Joulin and Tomas Mikolov.  Inferring algorithmic patterns with stack-augmented recurrent nets.  _Advances in neural information processing systems_ , 28, 2015.
  * Kang et al. (2023) Jikun Kang, Romain Laroche, Xindi Yuan, Adam Trischler, Xue Liu, and Jie Fu.  Think before you act: Decision transformers with internal working memory.  _arXiv preprint arXiv:2305.16338_ , 2023.
  * Ke et al. (2018) Nan Rosemary Ke, Anirudh Goyal ALIAS PARTH GOYAL, Olexa Bilaniuk, Jonathan Binas, Michael C Mozer, Chris Pal, and Yoshua Bengio.  Sparse attentive backtracking: Temporal credit assignment through reminding.  _Advances in neural information processing systems_ , 31, 2018.
  * Kingma & Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Laird (2019) John E Laird.  _The Soar cognitive architecture_.  MIT press, 2019.
  * Laird et al. (2017) John E Laird, Christian Lebiere, and Paul S Rosenbloom.  A standard model of the mind: Toward a common computational framework across artificial intelligence, cognitive science, neuroscience, and robotics.  _Ai Magazine_ , 38(4):13–26, 2017.
  * Le et al. (2018) Hung Le, Truyen Tran, and Svetha Venkatesh.  Dual memory neural computer for asynchronous two-view sequential learning.  In _Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining_, pp. 1637–1645, 2018.
  * Le et al. (2020a) Hung Le, Truyen Tran, and Svetha Venkatesh.  Neural stored-program memory.  In _International Conference on Learning Representations_ , 2020a.
  * Le et al. (2020b) Hung Le, Truyen Tran, and Svetha Venkatesh.  Self-attentive associative memory.  In _International Conference on Machine Learning_ , pp. 5682–5691. PMLR, 2020b.
  * Lee et al. (2019) Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh.  Set transformer: A framework for attention-based permutation-invariant neural networks.  In _International conference on machine learning_ , pp. 3744–3753. PMLR, 2019.
  * Liang et al. (2023) Yao Liang, Hongjian Fang, Yi Zeng, and Feifei Zhao.  A brain-inspired memory transformation based differentiable neural computer for reasoning-based question answering.  _arXiv preprint arXiv:2301.02809_ , 2023.
  * Limbacher & Legenstein (2020) Thomas Limbacher and Robert Legenstein.  H-mem: Harnessing synaptic plasticity with hebbian memory networks.  _Advances in Neural Information Processing Systems_ , 33:21627–21637, 2020.
  * Loshchilov & Hutter (2017) Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.
  * Marr & Thach (1991) David Marr and W Thomas Thach.  A theory of cerebellar cortex.  _From the Retina to the Neocortex: Selected Papers of David Marr_ , pp. 11–50, 1991.
  * Radford et al. (2019) Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8):9, 2019.
  * Rae et al. (2016) Jack Rae, Jonathan J Hunt, Ivo Danihelka, Timothy Harley, Andrew W Senior, Gregory Wayne, Alex Graves, and Timothy Lillicrap.  Scaling memory-augmented neural networks with sparse reads and writes.  _Advances in Neural Information Processing Systems_ , 29, 2016.
  * Rae et al. (2019) Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap.  Compressive transformers for long-range sequence modelling.  _arXiv preprint arXiv:1911.05507_ , 2019.
  * Ramsauer et al. (2020) Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, et al.  Hopfield networks is all you need.  _arXiv preprint arXiv:2008.02217_ , 2020.
  * Rosenbloom et al. (2016) Paul S Rosenbloom, Abram Demski, and Volkan Ustun.  The sigma cognitive architecture and system: Towards functionally elegant grand unification.  _Journal of Artificial General Intelligence_ , 7(1):1, 2016.
  * Santoro et al. (2017) Adam Santoro, David Raposo, David G Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap.  A simple neural network module for relational reasoning.  _Advances in neural information processing systems_ , 30, 2017.
  * Santoro et al. (2018) Adam Santoro, Ryan Faulkner, David Raposo, Jack Rae, Mike Chrzanowski, Theophane Weber, Daan Wierstra, Oriol Vinyals, Razvan Pascanu, and Timothy Lillicrap.  Relational recurrent neural networks.  _Advances in neural information processing systems_ , 31, 2018.
  * Schlag et al. (2021) Imanol Schlag, Tsendsuren Munkhdalai, and Jürgen Schmidhuber.  Learning associative inference using fast weight memory.  In _9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021_. OpenReview.net, 2021.  URL https://openreview.net/forum?id=TuK6agbdt27.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Werbos (1990) Paul J Werbos.  Backpropagation through time: what it does and how to do it.  _Proceedings of the IEEE_ , 78(10):1550–1560, 1990.
  * Weston et al. (2014) Jason Weston, Sumit Chopra, and Antoine Bordes.  Memory networks.  _arXiv preprint arXiv:1410.3916_ , 2014.
  * Weston et al. (2015) Jason Weston, Antoine Bordes, Sumit Chopra, Alexander M Rush, Bart Van Merriënboer, Armand Joulin, and Tomas Mikolov.  Towards ai-complete question answering: A set of prerequisite toy tasks.  _arXiv preprint arXiv:1502.05698_ , 2015.
  * Wu et al. (2020) Qingyang Wu, Zhenzhong Lan, Kun Qian, Jing Gu, Alborz Geramifard, and Zhou Yu.  Memformer: A memory-augmented transformer for sequence modeling.  _arXiv preprint arXiv:2010.06891_ , 2020.
  * Wu et al. (2022) Yuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, and Christian Szegedy.  Memorizing transformers.  In _The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022_. OpenReview.net, 2022.  URL https://openreview.net/forum?id=TrjbxzRcnf-.

Appendix

##  Appendix A Pseudo Codes

Notation: Given an original input sequence of length $T$ and an embedding
dimension of $D$ for the proposed model. Let $h^{l}$ denote the output of the
$l^{th}$ layer. The working memory is represented as a matrix
${M}_{w}\in\mathbb{R}^{N\times D_{m}}$ with distinct compartmentalized
memories for each row, where $m_{w,i}$ represents the state of slot $i$ (the
total number of slots is $N$). The long-term memory is represented as 3D
matrix ${M}_{l}\in\mathbb{R}^{C\times N\times D_{m}}$, where $C$ represents
the number of memory segments.

Initialization: Convert the raw input $X\in\mathbb{R}^{T\times{vocab\\_size}}$
to ${h_{0}}={p}_{x}\oplus{e}_{x}\oplus{cls}_{token}$, where
$h_{0}\in\mathbb{R}^{T\times D}$, ${e}_{x}$ signifies the positional encoding,
and ${e}_{x}$ corresponds to the embedding of $X$. Initialize the working
memory matrix ${M}_{w}$ and the long-term memory matrix ${M}_{l}$,
respectively, guaranteeing their universal sharing across all layers
encompassed within our framework.

Input to the layer l: ${h}^{l-1}$ having shape $\mathbb{R}^{T\times D}$ Step
1: The current perceptions competes with each other to be selected to update
the ${M}_{w}$

  * •

$Q={M}_{w}^{l-1}{W}^{q}$

  * •

$s_{k}={softmax}\left(\frac{{Q}\left({h}^{l-1}{W}^{k}\right)^{T}}{\sqrt{d^{k}}}\right)$,
where $s_{k}\in{R}^{N\times T}$

  * •

Construct a set $\mathcal{F}_{t}$ containing the indices of the $k$ selected
specialists with the top-$k$ largest values of $s_{k}$.

$s_{k}^{*}=\begin{cases}s_{[n,t]},&t\in\mathcal{F}_{t},\\\
0,&t\notin\mathcal{F}_{t}\end{cases}\quad\text{for all
}n\in\\{1,\ldots,N\\},t\in\\{1,\ldots,T\\}$

Step 2: External communication: $\mathcal{M}_{w}$-Write

  * •

$residual=M_{w}^{l-1}$

  * •

$\widetilde{M}_{w}^{l}={LN}_{1}(s_{k}^{*}h^{l-1}W^{v}+residual)$

  * •

$\widetilde{M}_{w,i}^{l}=\text{ReLU}\left(\text{MLP}_{i}(\widetilde{M}_{w,i-1}^{l})\right)\quad
i\in\\{1,\ldots,k\\}$, $k$ signifies the number of layers in the MLP.

  * •

$\widetilde{M}_{w}^{l}=LN_{2}(\widetilde{M}_{w,k}^{l}+residual)$

  * •

$M_{w}^{l}=F_{l}(M_{w}^{l-1},h^{l-1})\odot
M_{w}^{l-1}+I_{l}(M_{w}^{l-1},h^{l-1})\odot\tanh(\widetilde{M}_{w}^{l})$

$F_{l}$ refers to the forget gate and $I_{l}$ represents the input gate of
$l^{th}$ layer.

Step 3: Internal communication: $\mathcal{M}_{l}$-Write with the updated
${M}_{w}$

  * •

$M_{l}^{l}={LN}_{3}((M_{w}^{l}\otimes M_{l}^{l-1})+M_{l}^{l-1})$

Step 4: Make sense of the current perceptions based on the updated ${M}_{w}$
and ${M}_{l}$: $\mathcal{M}_{w}$-Read and $\mathcal{M}_{l}$-Read

  * •

$\bm{\mathit{{h}_{w}^{l}}}=\text{softmax}\left(\frac{{h^{l-1}\widetilde{W}^{q}\left(M_{w}^{l}\widetilde{W}^{k}\right)^{T}}}{\sqrt{d^{k}}}\right)M_{w}^{l}\widetilde{W}^{v}$

  * •

$\widehat{M}_{l}^{l}=\frac{1}{C}\sum_{i=1}^{C}M_{l}^{l}[i,:,:]$ where
$\widehat{M}_{l}^{l}\in\mathbb{R}^{N\times D_{m}}$

  * •

$c_{k}={\text{softmax}}\left(\frac{{h^{l-1}\widehat{W}^{q}\left(\widehat{M}_{l}^{l}\widehat{W}^{k}\right)^{T}}}{\sqrt{d^{k}}}\right)$

  * •

$\bm{\mathit{{h}_{l}^{l}}}=c_{k}^{*}\widehat{v}\quad\text{where}\quad\widehat{v}\in\widehat{M}_{l}^{l}\widehat{W}^{v}$

Similar to $s_{k}^{*}$, $c_{k}^{*}$ represents the outcome of the top-$k$
largest sparse attention scores associated with $c_{k}$.

  * •

$\bm{\mathit{{h}_{wl}^{l}}}=\textit{MHC}\left(h_{l}^{l}\bar{W}^{Q},h_{w}^{l}\bar{W}^{K},h_{w}^{l}\bar{W}^{V}\right)$

  * •

$\bm{\mathit{{h}^{l}}}=\alpha h_{w}^{l}+(1-\alpha)h_{wl}^{l}$

Algorithm 1 MITR Algorithm

##  Appendix B Experiments on Memory Consolidation

###  B.1 More Ablation Results

In this section, we present the results of ablation experiments concerning the
global sharing of memory components. The plot in Fig. 5 illustrates the change
in testing accuracy over training iterations. Here, the MITR model denotes the
usage of a globally shared working memory module, while $MITR_{w}/o$ signifies
the absence of global sharing. TR represents the conventional Transformers
model. The results demonstrate several key findings. Firstly, across all three
tasks, both our MITR and $MITR_{w}/o$ models exhibit faster convergence
compared to the traditional TR model. Secondly, when compared to the baseline
TR model and the $MITR_{w}/o$ model, MITR achieves the highest accuracy and
convergence rates, especially in complex binary and ternary relation tasks,
showcasing the effectiveness of the adopted global sharing strategy in complex
relational reasoning tasks. However, for non-relational tasks, their
performance remains relatively on par, potentially due to the fact that non-
relational tasks often emphasize extracting direct patterns and information
from input data, and thus, Transformer models’ self-attention mechanisms can
effectively capture these correlations.

(a) Unary Accuracy

(b) Binary Accuracy

(c) Ternary Accuracy

Figure 5: Test accuracy vs. training iterations for the Sort-of-CLEVR task.
Results of the ablation experiments on the memory module’s global sharing or
persistence.

###  B.2 Visualization

We present four Q&A instances from the bAbI dataset, illustrating the
attention matrix of perception towards both working memory (Fig. 6(a)) and
long-term memory (Fig. 6(b)) during the reasoning process within eight layers
of the MITR model. Each row corresponds to distinct Q&A examples, with columns
representing the layers of our MITR. Moreover, in each heatmap, the vertical
axis is a representation of perceptual information, functioning as queries,
while the horizontal axis refers to long-term memory or working memory,
serving as keys in the cross-attention mechanism.

As an illustration, let’s consider a specific problem with its set of stories
and corresponding answer. In the first line of the working memory attention
patterns, each $Si$ represents a distinct sentence in the narrative. For
instance, $S1$ corresponds to ’the hallway is east of the bathroom’, and $S2$
corresponds to ’the bedroom is west of the bathroom’, and so forth until
$S10$, where he maximum $S_{max}$ for each task may vary. In cases where there
are fewer than $S_{max}$ story sentences, placeholder sentences (all zeros)
are used as filler input. The letter ’Q’ represents the question: ’What is the
bathroom east of?’, and ’CLS’ is the classification header.

(a) Attention patterns between perception and working memory across different
layers of the MITR for 4 examples from the bAbI dataset. Firstly, it is a
well-known fact that the magnitude of attention weights is reflected through
color intensity, where the deeper the color, the higher the weight. During the
inference process, with an increase in layers within the MITR model, a clear
pattern emerges in the heatmaps: the presence of colored blocks markedly
diminishes, yet the intensity of color within the remaining blocks steadily
deepens. This implies a growing focus of working memory on essential segments
of current inputs, especially those indispensable for question- answering.
This result is consistent with the restrained write access elucidated in
Section 2.3.1, which encourages working memory to concentrate on the utmost
critical elements of the task.

(b) Attention patterns between perception and long-term memory across
different layers of the MITR for 4 examples from the bAbI dataset. For each
problem (per row), as layers increase (from left to right), the heat maps show
a clear trend opposite to the working memory pattern, with more and more
colored areas but less intensity, and the last several heat maps tend to be
stable and approximate. This suggests that long-term memory gradually
accumulates more knowledge as computation steps $t$ increase, and thus
exhibits moderate correlation with a broader input area (unlike working memory
with limited capacity, which focuses only on the most important info). Similar
attention patterns in the final layers imply convergence of the three-
dimensional long-term memory matrix within finite calculation steps, resulting
in a more consistent understanding of various input aspects.

Figure 6: Attention patterns between perception and memory.

##  Appendix C Implementation Details

###  C.1 Gating Mechanism

To update the working memory, we employ the gating mechanism introduced by
Santoro et al. (2018), which consists of an input gate and a forget gate. Let
$X^{t-1}=[X^{t-1}_{1},X^{t-1}_{2},\cdots,X^{t-1}_{T}]\in\mathbb{R}^{B\times
T\times D}$ represent the perceptual input, and
$M_{w}^{t-1},M_{w}^{t}\in\mathbb{R}^{B\times N\times D_{m}}$ represent the
previous and updated working memory, respectively. $\widetilde{M}_{w}^{t}$ is
the intermediate result described by Eq. 5 in Section 2.3.1. The gating
mechanism can be formulated as follows.

$\displaystyle\bar{X}$
$\displaystyle=\frac{1}{T}\sum_{i=1}^{T}\text{relu}(X_{i}\times W^{I})$
$\displaystyle K$ $\displaystyle=\bar{X}+\tanh(M_{w}^{t-1})\times W^{F}$
$\displaystyle I_{t}$ $\displaystyle=\text{sigmoid}(K+b_{i})$ $\displaystyle
F_{t}$ $\displaystyle=\text{sigmoid}(K+b_{f})$ $\displaystyle M_{w}^{t}$
$\displaystyle=I_{t}\times\tanh(\widetilde{M}_{w}^{t})+F_{t}\times
M_{w}^{t-1}$

Here, $I_{t}$ and $F_{t}$ represent the input and forget gates of the current
calculation step t, with corresponding weight matrices $W^{I}$ and $W^{F}$.
The biases for the input and forget gates are denoted as $b_{i}$ and $b_{f}$.
In practice, we set $b_{i}=0$, $b_{f}=1$, $D=D_{m}$.

##  Appendix D Hyperparameters Setting

The hyperparameter settings of the MITR model on the four tasks are shown in
Table 4, where Adam and AdamW were proposed by (Kingma & Ba, 2014) and
(Loshchilov & Hutter, 2017), respectively. The other baseline models remain
the same configuration for the corresponding tasks.

Table 4: The hyperparameter setting of MITR model on four tasks Parameters | Tasks
---|---
bAbI |  Sort-of-CLEVR |  Triangle |  Cifar-10
Top-k |  5 |  5 |  5 |  5
Number of layers |  8 |  4 |  2 |  4
Number of attention heads |  8 |  4 |  4 |  4
Embedding dimensions |  256 |  256 |  128 |  256
Optimizer |  Adam |  Adam |  Adam |  AdamW
Weight decay |  N/A |  N/A |  N/A |  0.09
Learning rate |  0.0002 |  0.0001 |  0.0001 |  0.0002
Batch size |  64 |  64 |  100 |  64
Inp Dropout |  0.1 |  0.1 |  0.1 |  0.1
Seed |  1 |  1 |  1 |  1
Number of working memory slots ($N$) |  8 |  8 |  8 |  8
Number of long-term memory segments ($M$) |  5 |  5 |  5 |  5
Size of each working memory slot ($D_{m}$) |  256 |  256 |  256 |  256
Size of each long-term memory segment |  $5\times 256$ |  $5\times 256$ |  $5\times 256$ |  $5\times 256$
Number of MLP layers in attention |  4 |  4 |  5 |  4
Memory attention heads |  8 |  4 |  1 |  1
Gate style |  ’unit’ |  ’unit’ |  ’unit’ |  ’unit’
Initial $\alpha$ value |  0.7 |  0.75 |  0.7 |  0.55

##  Appendix E Description of the Tasks/Datasets

###  E.1 bAbI

The bAbI-20k dataset consists of 20 distinct text-based QA tasks, each
presenting a unique reasoning challenge, including counting, deduction and
induction. Each task is divided into training, validation and test datasets,
with 9k, 1k and 1k questions respectively. They are presented in the form of
short stories or text passages, comprising narratives, questions, answers and
supporting facts. Narratives introduces entities, actions and contextual info
relevant to the question, and the answer is substantiated by facts from the
narratives. Four detailed examples are shown below.

Task 1: Single Supporting Fact1\. John travelled to the hallway. 2\. Mary
journeyed to the bathroom. 3\. Daniel went back to the bathroom. 4\. John
moved to the bedroom. Where is Mary? A: bathroom S: 2Task 8: Lists/Sets1\.
Mary got the milk there. 2\. John moved to the bedroom. 3\. John picked up the
football there. 4\. John journeyed to the bathroom. What is John carrying? A:
football S: 3

Task 12: Conjunction1\. Mary and Daniel travelled to the bathroom. 2\. John
and Daniel travelled to the office. 3\. Sandra and Daniel moved to the
kitchen. 4\. Sandra and John moved to the garden. Where is Sandra? A: garden
S: 4Task 16: Basic Induction1\. Julius is a lion. 2. Lily is a rhino. 3\.
Bernhard is a swan. 4. Lily is white. 5\. Bernhard is green. 6. Greg is a
rhino. 7\. Greg is gray. 8. Julius is white. 9\. Brian is a lion. What color
is Brian? A: white S: 9 1 8

###  E.2 Sort-of-CLEVR

Ternary questionsQ: How many objects are in the rectangle formed by the
centers of the red and yellow objects?Answer: 1 Q: Are there any objects on
the line formed by the centers of the red and grey objects?Answer: yes Q: How
many objects form an obtuse triangle with a blue object and a yellow
object?Answer: 3 Binary questions:Q: What is the shape of the object that is
furthest from the green object? Answer: orange Q: What is the color of the
object that is closest to the red object? Answer: green Q: How many objects
have the shape of the yellow object? Answer: 3  Figure 4: An instance from the
sort-of-clevr dataset. Unary questions:Q: What is the shape of the red object?
Answer: circle Q: Is the green object on the left or right of the image?
Answer: left Q: Is the grey object on the top or bottom of the image? Answer:
top

###  E.3 Detecting Equilateral Triangles

(a)

(b)

(c)

(d)

Figure 7: An illustration of the equilateral triangle detection task. Here,
(a) and (b) refer to non-equilateral triangle, (c) and (d) refer to
equilateral triangle.

##  Appendix F Additional Experimental Results

###  F.1 Test Curves for Detecting Equivalent Triangles

As shown in Fig. 8, we observe that our proposed models, MITR and its variant,
MITR+S without top-k sparsity, show significantly faster convergence compared
to other models. It is worth noting that the MITR model achieves slightly
higher accuracy compared to the MITR+S model, highlighting the usefulness of
our competitive writing mechanism.

Figure 8: Comparison of convergence speeds on equilateral triangle detection
task.

###  F.2 bAbI Detailed Results

Detailed results are shown in the table 5

Task |  run1 |  run2 |  run3 |  run4 |  run5 |  run6 |  run7 |  run8 |  run9 |  run10 |  Mean±std
---|---|---|---|---|---|---|---|---|---|---|---
1:Single Supporting Fact |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0
2:Two Supporting Facts |  0.21 |  0.08 |  0 |  0.10 |  0 |  0.52 |  0.56 |  0.73 |  0.35 |  0.62 |  0.32±0.28
3:Three Supporting Facts |  1.35 |  0.35 |  0.03 |  0.15 |  0.56 |  1.88 |  0.92 |  0.04 |  0.15 |  1.19 |  0.66±0.64
4:Two Arg. Relations |  0 |  0 |  0 |  0 |  0 |  0.21 |  0 |  0 |  0.31 |  0 |  0.05±0.11
5:Three Arg. Relations |  0.52 |  0.46 |  0.52 |  0.73 |  0.52 |  0.52 |  0.73 |  0.73 |  0.83 |  0.42 |  0.60±0.14
6:Yes/No Questions |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0
7:Counting |  0.31 |  0.25 |  0.04 |  0.73 |  0.35 |  0.35 |  0.35 |  0.46 |  0.15 |  0.04 |  0.30±0.21
8:Lists/Sets |  0 |  0 |  0.21 |  0.21 |  0.31 |  0.10 |  0.52 |  0.10 |  0.10 |  0.21 |  0.18±0.16
9:Simple Negation |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0
10:Indefinite Knowledge |  0 |  0 |  0.31 |  0.10 |  0 |  0 |  0 |  0 |  0.10 |  0 |  0.05±0.10
11:Basic Coreference |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0
12:Conjunction |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0
13:Compound Coref. |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0
14:Time Reasoning |  0 |  0 |  0 |  0 |  0 |  0 |  0.10 |  0 |  0 |  0 |  0.01±0.03
15:Basic Deduction |  0 |  0 |  0.21 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0.02±0.07
16:Basic Induction |  48.21 |  47.12 |  44.96 |  47.92 |  48.04 |  48.63 |  47.35 |  50.81 |  47.27 |  48.56 |  47.89±1.47
17:Positional Reasoning |  0.94 |  0.03 |  0.03 |  0.31 |  0.58 |  1.67 |  0.83 |  0.60 |  0.77 |  0.38 |  0.61±0.49
18:Size Reasoning |  0.42 |  0.01 |  0.00 |  0.52 |  0.21 |  0.21 |  0.21 |  0.31 |  0.52 |  0.21 |  0.26±0.18
19:Path Finding |  0.00 |  0 |  0 |  0.21 |  0 |  0 |  0.10 |  0 |  0.10 |  0 |  0.04±0.07
20:Agent’s Motivations |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0 |  0
Average |  2.60 |  2.42 |  2.32 |  2.55 |  2.53 |  2.70 |  2.58 |  2.69 |  2.53 |  2.58 |  2.55±0.11
Failed task(>5%) |  1 |  1 |  1 |  1 |  1 |  1 |  1 |  1 |  1 |  1 |
Table 5: Results from 10 test runs of the MITR model on 20 bAbI tasks, each
consisting of 10k samples, following 200 epochs of joint training. Bold
denotes the best run.