SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ankur Sikarwar and Mengmi Zhang
Deep NeuroCognition Lab, Institute for Infocomm Research (I2R), A*STAR, Singapore Center for Frontier AI Research (CFAR), Agency for Science, Technology, and Research (A*STAR), Singapore School of Computer Science and Engineering, Nanyang Technological University, Singapore Address correspondence to mengmi@i2r.a-star.edu.sg
Working memory (WM), a fundamental cognitive process facilitating the temporary storage, integration, manipulation, and retrieval of information, plays a vital role in reasoning and decision-making tasks. Robust benchmark datasets that capture the multifaceted nature of WM are crucial for the effective development and evaluation of AI WM models. Here, we introduce a comprehensive Working Memory (WorM) benchmark dataset for this purpose. WorM comprises 10 tasks and a total of 1 million trials, assessing 4 functionalities, 3 domains, and 11 behavioral and neural characteristics of WM. We jointly trained and tested state-of-the-art recurrent neural networks and transformers on all these tasks. For comparison, we also include human behavioral benchmarks. Note that all computational models were never trained with any human data or behavioral biases; yet, these models remarkably replicate some characteristics of WM in biological brains, such as primacy and recency effects. Moreover, we performed neural population analysis on these models and identified neural clusters specialized for different domains and functionalities of WM. Not all computational models exhibit a strong alignment with all human behaviors. Our experimental results also reveal several limitations in existing models to match with working memory capabilities of humans. This dataset serves as a valuable resource for communities in cognitive psychology, neuroscience, and AI, offering a standardized framework to compare and enhance WM models, investigate WM's neural underpinnings, and develop WM models with human-like capabilities. Our source code and data are available at: link.
Working memory (WM) defines a core cognitive process enabling temporary storage, integration, manipulation, and recall of information. It is vital for many downstream tasks involving reasoning [57, 46, 34], decision-making [50, 26], and language comprehension [5, 9, 1]. Understanding and modeling WM has significant implications for both neuroscience and AI research. In neuroscience and cognitive science, extensive investigations have been conducted to unravel the underlying mechanisms of WM. The non-exhaustive list of works involves (1) identifying brain regions such as the prefrontal cortex [15, 51, 37], hippocampus [6, 55], and parietal cortex [28, 32, 47] as key players in WM processes; (2) exploring the role of WM capacity in cognitive processes such as selective attention [17, 48], and executive control [19, 7]; and (3) investigating how WM abilities differed in individuals of different ages [20, 53], intelligence levels [12], and neurological disorders [2, 23].
However, despite significant progress in individual subdomains, there remains a notable gap in the broader study of WM and its relevance for the AI community. To date, research efforts in neuroscience, and cognitive science have mostly focused on specific aspects of WM or individual
memory tasks. There is a lack of systematic, integrative, and quantitative exploration that covers multiple tasks, encompasses various functionalities and domains of WM, and provides a systematic benchmark for evaluating both humans and AI models in these memory tasks.
In this work, we aim to address this gap by establishing a general framework to study WM. To capture the multifaceted nature of WM due to its complexity, flexibility, and variability, we include 10 WM tasks and curate 1 million trials. Each of the state-of-the-art recurrent neural networks and transformer models is jointly trained and tested on these tasks with the following goals: (1) to assess model behaviors, such as set size effect and retention interval; (2) to investigate neural populations, such as task-specialized neural clusters and neural correlates of executive control; and (3) to examine and compare performance across multiple WM models when performing different memory functions: storage, integration, manipulation, and supervision. We include human performance on the WM tasks for behavioral comparison with AI models. In the experiments, we observe that AI models replicate some human-like characteristics of WM; however, we also identify several limitations in the existing WM models, such as disparities between human and model behaviors and inferior model performances in tasks that humans excel at. Through this interdisciplinary exploration, we strive to enhance our understanding of WM processes, uncover novel insights into the cognitive architecture of human WM, and propel the development of AI systems with more robust and human-like WM capabilities.
Main contributions are highlighted:
1. We introduce WorM, a comprehensive large-scale WM benchmark dataset, covering 10 WM tasks and containing 1 million trials, encompassing 4 functionalities and 3 domains of WM.
2. We introduce evaluation metrics for all our tasks and establish a general methodology to jointly train, study, and benchmark WM models on all tasks.
3. We explore 8 WM behavioral benchmarks and find that recurrent neural networks are more closely aligned with humans, while transformers exhibit limitations in replicating human behavior.
4. We examine the neural population of WM models. Our analysis reveals task-specialized neural clusters and uncovers neural correlates of executive control in WM models.
Working memory (WM) has been a topic of interest in cognitive psychology and neuroscience for decades. Numerous studies [40, 24, 39, 35, 61, 3, 8] have investigated different behavioral effects and their underlying mechanisms and neural correlates. There has also been an interest [14, 42, 36, 43] in developing computational models of WM that can explain and approximate human behavior on WM tasks. Several WM tasks [16, 54, 13, 35, 41, 31] have been proposed to evaluate WM performance of humans. For example, the N-back task is a widely used WM task that requires participants to remember a sequence of stimuli and respond when a stimulus matches the one presented N trials earlier. It is important to emphasize that the majority of these prior works have typically focused on isolated aspects of working memory or evaluated a single working memory model with a restricted range of working memory tasks. In contrast, our contribution involves the development of a comprehensive and quantitative benchmark that encompasses 10 distinct working memory tasks, spanning four functionalities and three domains of working memory, for the evaluation of both human performance and AI models.
In the field of AI, WM has also attracted considerable attention. AI researchers have aimed to develop computational models that can emulate and augment human-like WM capabilities. Various memory architectures [10, 27, 25, 62, 56] have been proposed, such as neural networks with memory cells, recurrent neural networks, and memory-augmented neural networks. These models have demonstrated promising results in specific WM tasks, such as copying, sorting, and memory recalls, as well as real-world applications [22, 33], such as navigation in naturalistic environments, video recognition, and language processing.
Recently, [45] established a comprehensive working memory (WM) benchmark to aid in the development of WM theories by collecting empirical findings from previous human studies in a wide range of WM tasks. Different from their work, in this paper, we present a large-scale WM
A Spatial Free Recall (SFR)
B Change Detection (CD)
C Visual Item Recognition (VIR)
D Visual Serial Recall (VSR)
E Visual Serial Recognition (VSRec)
F Complex Span (CS)
G Spatial Coordination (SC)
H Spatial Integration (SI)
CS - Visual + Spatial
Spatial Memory Updating (SMU)
Figure 1: Schematic illustration of all 10 working memory tasks. Each panel illustrates the schematic for a different task. In each panel, the arrow pointing from left to right denotes the progression of time. Overlapping frames indicate multiple time steps whereas single frames represent a single time step in the trial. In F, we further divide the task into four types. Two types are shown here. See Sec. 3 for a detailed description of each task.
J Spatial Task Switching (STS)
benchmark dataset to bridge the gap between cognitive psychology and the AI field, allowing modern AI models to be trained on such WM tasks for direct behavioral and neural comparison.
3 Psychophysics Experiments in WorM
Phase definitions in WM experiments. In this section, we refer to our 10 WM tasks as 10 experiments. A WM experiment often consists of several but not necessarily all phases (Fig. 1) described below: * Presentation phase ( $\mathbf{P}_{\text{present}}$ ): A sequence of stimuli to remember is presented sequentially. * Retention phase ( $\mathbf{P}_{\text{retent}}$ ): A grey blank image is presented for a period of time, during which the remembered patterns in  $\mathbf{P}_{\text{present}}$  have to be maintained in the memory. * Probing phase ( $\mathbf{P}_{\text{probe}}$ ): Responses to a series of questions asking to recall certain aspects of the memory content are required. Sometimes, there will be no questions but free recalls are required. * Distractor phase ( $\mathbf{P}_{\text{distractor}}$ ): A series of other distracting tasks are presented in this phase to interfere with the original WM experiments. * Memory updating phase ( $\mathbf{P}_{\text{update}}$ ): A series of instructions to update the memory content are given. * Task switching phase ( $\mathbf{P}_{\text{switch}}$ ): A cue image is presented to signal the transition between a pair of tasks.
Time steps in WM experiments. The time duration of each phase in human psychophysics experiments is often measured in milliseconds (ms). However, computational WM models do not have the notion of time. We establish a mapping between the time step  $t$  in WM models and ms in human psychophysics experiments for all the phases of all the experiments (see Appendix A2.2 for mapping details). For consistency, we describe all the main text with time steps  $t$ . We introduce  $T$  to denote the total time steps for each experiment and  $T_{phase}$  to denote the total time steps of each phase, where phase could be from any of the phases introduced above. For example,  $T_{present} = 3$
(a) Memory domains
Figure 2: WorM covers a wide spectrum of working memory (WM) tasks encompassing 3 domains, 4 functionalities, and 11 characteristics. For all ten tasks, we categorize them based on their objectives in studying different (a) memory domains, and (b) memory functions. See Fig. 1 for task acronyms. In (c), we present word clouds representing the studied memory characteristics. Larger fonts indicate that more tasks (the exact number in brackets) produce results in this category.
(b) Memory functionalities
(c) Memory characteristics
indicates that there are 3 time steps for the presentation phase and  $\mathrm{t}_{\text{present}} = 1$  indicates the first time step of the presentation phase.
Nomenclatures in cognitive science and neuroscience. In human psychophysics experiments [4, 13, 29], researchers often use "length list", denoted as  $L$ , to indicate the total number of stimuli presented during each  $P_{\text{present}}$ , and "serial position" to indicate the position of the  $l$ th stimulus in  $L$ . For WM models, each stimulus is presented at each time step of  $P_{\text{present}}$ . Thus,  $L = T_{\text{present}}$  and  $t_{\text{present}} = l$  denotes the  $l^{th}$  serial position during  $P_{\text{present}}$ . Following psychology and neuroscience literature [35], we also define "set size", denoted as  $S$ , to be the total number of items that need to be held and manipulated in working memory at a given time step. For all the experiments, we use these terms interchangeably, and we also use "agents" to indicate that the test subjects can be both humans and computational WM models.
Introduction to ten psychophysics experiments
We introduce 10 psychophysics experiments on Working Memory (WM) from cognitive science and neuroscience [13, 35, 29, 4, 60, 41, 52, 44, 38]. To facilitate comparisons with computational models, we directly copy the human behavioral data collected from these existing works to our paper. Based on WM functionalities these experiments entail (Fig. 2b), we categorize them into four functional groups: (A) memory storage, (B) memory integration, (C) memory manipulation, and (D) memory supervision. We also show the overview of the three WM domains (visual, spatial, and temporal) that each experiment involves in Fig. 2a. See Appendix A1 for experimental details.
Experiment A1 - Spatial Free Recall (SFR) [13] (Fig. 1A). At every  $t_{\text{present}}$ , a randomly chosen square on a green grid of size  $10 \times 10$  turns red. In  $P_{\text{probe}}$ , a blank grey image is shown. The task for the agent is to remember all the previous spatial locations where the squares turned red and provide a 100-way multi-label classification response.
Experiment A2 - Change Detection (CD) [35] (Fig. 1B). In  $\mathrm{P}_{\text {present }}$ , a memory array is shown, consisting of  $S$  bars with varying attributes such as color, orientation, size, and the presence/absence of a gap in each bar. Followed by a sequence of blank grey images in  $\mathrm{P}_{\text {retent }}$ , the agents are tested with a probe array in  $\mathrm{P}_{\text {probe }}$ , which can be identical to the memory array or have a single bar that differs along one feature dimension. The task for the agents is to make a binary prediction about whether there is a change or not given the test array.
Experiment A3 - Visual Item Recognition (VIR) [29] (Fig. 1C). In  $\mathrm{P}_{\text {present }}$ , a sequence of distinct  $6 \times 6$  matrix patterns is displayed. In  $\mathrm{P}_{\text {retent }}$ , a series of blank grey images is shown. In  $\mathrm{P}_{\text {probe }}$ , a probe image is presented, consisting of two patterns side-by-side. One pattern matches a previously shown pattern, while the other is a distractor. The agents must perform a binary classification task, where 0 represents the left pattern and 1 represents the right pattern.
Experiment A4 - Visual Serial Recall (VSR) [4] (Fig. 1D). In  $\mathrm{P}_{\text {present }}$ , a sequence of  $L$  matrix patterns is sequentially presented. These patterns consist of a  $6 \times 6$  grid with half of the cells filled in green. In  $\mathrm{P}_{\text {probe}}$ , a probe image is displayed, containing all  $L$  patterns presented during  $\mathrm{P}_{\text {present }}$ . Agents perform a 9-way classification task at each  $\mathrm{t}_{\text {probe}}$ , recalling the matrix patterns seen at  $\mathrm{t}_{\text {present }}$  in a serial order. See Appendix A3.1 for a language-based variant of this experiment.
Experiment A5 - Visual Serial Recognition (VSRec) [4] (Fig. 1E).  $\mathrm{P}_{\text {present }}$  is exactly the same as  $\mathrm{P}_{\text {present }}$  in the VSR experiment. Differences were introduced during  $\mathrm{P}_{\text {probe }}$ . At each  $\mathrm{t}_{\text {probe }}$ , a target pattern from  $\mathrm{P}_{\text {present }}$  and a distractor pattern were presented together. Distractor patterns were distinct from the initially presented matrix patterns and differed from the target pattern by  $n$  cells, where  $n$  denotes the pattern-distractor difference. Agents performed a binary classification task, with 0 indicating the left pattern and 1 indicating the right pattern.
Experiment A6 - Complex Span (CS) [60] (Fig. 1F). The experiment includes two types of  $\mathrm{P}_{\text {present }}$  and  $\mathrm{P}_{\text {probe}}$ , as well as two types of  $\mathrm{P}_{\text {distractor}}$ . (I) The first type involves memorizing visual patterns and recalling them in order (visual). (II) The second type involves memorizing spatial patterns of ball movements and recalling their directions (spatial). Two types of  $\mathrm{P}_{\text {distractor}}$  are introduced to interfere with memory. (III) color discrimination tasks (visual) and (IV) spatial symmetry discrimination tasks (spatial) are used as distractors. Four variations of experiment conditions are considered: visual storage + visual distractor, spatial storage + visual distractor, visual storage + spatial distractor, and spatial storage + spatial distractor.
Experiment B1 - Spatial Coordination (SC) [41] (Fig. 1G).  $\mathrm{P}_{\text {present }}$  involves a  $10 \times 10$  grid where one cell is highlighted in green at each  $\mathrm{t}_{\text {present }}$ . In  $\mathrm{P}_{\text {probe}}$ , agents are presented with a blank gray image and asked to perform a binary discrimination task to determine whether the pattern formed by integrating all the patterns in  $\mathrm{P}_{\text {present }}$  is symmetric about the vertical axis or not.
Experiment B2 - Spatial Integration (SI) [52] (Fig. 1H).  $\mathrm{P}_{\text {present }}$  involves sequentially presenting partial line drawings on a  $4 \times 4$  grid, eventually completing a multi-segment figure. The number of line segments in the partial drawings determines the number of integration operations required, also equivalent to  $L$ . In  $\mathrm{P}_{\text {probe }}$ , agents perform a binary classification task to determine if the given figure matches the mentally integrated figure from  $\mathrm{P}_{\text {present }}$ .
Experiment C1 - Spatial Memory Updating (SMU) [44] (Fig. 1I). In  $\mathrm{P}_{\text {present }}$ , a stimulus with green squares is presented, each containing a red marker in one of the nine possible locations on a  $3 \times 3$  grid.  $\mathrm{P}_{\text {update }}$  involves arrows indicating the direction for mentally updating the marker's location. In  $\mathrm{P}_{\text {probe}}$ , agents perform a 9-way classification task to recall the marker's location in the highlighted square.
Experiment D1 - Spatial Task Switching (STS) [38] (Fig. 1J). In  $\mathrm{P}_{\text{switch}}$ , a cue image indicates the discrimination task for subsequent  $\mathrm{P}_{\text{probe}}$ . There are two tasks across which to switch: top-versus-bottom and left-versus-right. At every  $\mathrm{P}_{\text{probe}}$ , a red marker appears in a  $2 \times 2$  grid, and agents switch between the two tasks based on cues presented during  $\mathrm{P}_{\text{switch}}$ .
Experimental trial splits for computational WM models. For every WM psychophysics experiment introduced above, we generate 86,400 trials for training, 9,600 trials for validation, and 9,600 trials for testing by following a ratio of 9:1:1. These trials are distributed uniformly among various experimental conditions for individual experiments.
4 Computational Models of Working Memory
Joint training and testing regime. There exist multiple approaches to train and evaluate computational working memory (WM) models using our dataset, such as "train on many; test on one" and "train on one; test on one". However, our aim here is not to exhaustively benchmark WM models across all possible training and testing regimes. Instead, we aim to lay the foundation for a systematic and quantitative methodology to comprehensively study the multi-facets of WM. In this regard, we propose a joint learning paradigm where we simultaneously train and test different models on all conditions of all ten WM tasks. This approach allows us to explore the intricate interactions and inter-dependencies among the different WM tasks, capturing the complex nature of WM performance. For comparisons with the joint training regime, we also conduct single-task training on the VSR task. See Appendix A3.3 for more details.
Architectures for working memory models. We investigate several state-of-the-art variations of recurrent neural networks, including vanilla recurrent neural networks (RNNs), Gated Recurrent Units (GRUs), and Long Short-Term Memory (LSTM) networks [10, 27]. Furthermore, considering the recent success of Transformers [59], we also include vanilla transformer-based encoders (TRF). Below, we introduce details of our model architectures (Fig. 3).
Figure 3: Overview of Working Memory (WM) models. All WM models take a batch of stimuli in sequences from  $I_{t=1}$  to  $I_{t=T}$  as inputs and predict the task-specific responses at every single time step  $t$ . These responses are compared against the batched and masked ground truths with classification losses. The gradients are only back-propagated based on the actual target values. For time steps without responses required (denoted as "NaN"), the predicted responses from the linear classifiers are not computed, and thus, no gradients are back-propagated. See legends for box notations. "shared" refers to learnable parameters shared jointly over tasks except for "task-specific shared" in linear classifiers, which refers to learnable weight parameters shared across time steps within the same task but different over tasks. See Sec. 4 for model details.
Feature Extraction. At each time step  $t$  of a trial, the visual stimulus  $I_{t}$  comes as a  $32 \times 32 \times 3$  tensor which is then fed to a 4-layer 2D-convolutional network (2D-ConvNet). See Appendix A2.1 for the 2D-ConvNet architecture. The output feature maps from 2D-ConvNet are then flattened to a vector of dimension  $K$  and linearly projected to a feature representation  $F_{t}$  of dimension  $D$  using weight matrix  $\in \mathbb{R}^{K \times D}$ . Different trials can be of different lengths, therefore we pad all the trials with blank images to ensure that all trials have a consistent length of 20.
Task embeddings. In addition to extracting stimuli features, we also introduce learnable task-specific embeddings  $M \in \mathbb{R}^{14 \times 14}$ , informing the network about the task identity during the joint training. At every  $t$  of each task, the corresponding task embedding is concatenated with the stimulus feature representation  $F_{t}$  from  $I_{t}$ , resulting in a (14+D)-dimensional representation  $A_{t}$ .  $A_{t}$  is further fed to various WM models introduced below for further processing. See Appendix A2.1 for details.
Encoding Time-dependent Representations. As we have RNN-based networks and transformer-based networks to process stimuli sequences, we introduce their model designs separately. For RNNs, GRUs, and LSTMs, we use only one recurrent layer. The recurrent networks take (14+D)-dimensional representation  $A_{t}$  above as inputs for each  $t$ . We define the capacity of these models  $C$  as the number of hidden units in this recurrent layer. At every  $t$ , the recurrent networks output a hidden representation  $h_{t}$  of dimension  $C$  used for predicting task-specific responses.
For TRFs, we introduce extra learnable positional embeddings  $P \in \mathbb{R}^{20 \times (14 + D)}$  which is shared across all tasks. Each row of  $P$  indicates a "time step". At  $t^{th}$  time step,  $P_{t}$  is incorporated into  $A_{t}$  via element-wise addition. We denote this positional embedding-modulated stimulus feature representation as  $A_{TRF,t}$  which is then fed to the TRF. We include two standard transformer encoder blocks in TRF. We define the model's capacity  $C_{TRF}$  as the dimension  $14 + D$  of the stimulus feature vector  $A_{t}$ . Additionally, to simulate a more realistic scenario where the TRF model can only rely on past context to make predictions, we apply masking in the self-attention mechanism to prevent the model from accessing future information. We take the output  $h_{t,TRF}$  of TRF corresponding to the  $t^{th}$  time step and use it for predicting task-specific responses at  $t$ .