SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Response Generation. The output of a WM network  $h_t$  or  $h_{t,TRF}$  is fed to a task-specific linear classifier with weight matrix  $O_{task}^{(14 + D)\times U_{task}}$  for generating final probability distributions over a fixed set of all possible response choices in the corresponding task. Note that  $O_{task}$  is shared across all  $T$  time steps within the same task, but it is independent over different tasks. For example, in the VSR task, during  $\mathrm{P}_{probe}$ , the WM model has to predict which pattern among the 9 shown patterns is the target pattern. In this case,  $U_{task} = 9$ , and the linear classifier outputs a probabilistic distribution of dimension 9, and the choice with the largest probability is the final response selected by the model.
Figure 4: Performance benchmarks and behavioral analysis for working memory models and humans. We present the behavioral accuracy as a function of (A) list lengths  $L$  for LSTM-256 in VSR task, (B) retention interval in VIR task, (C) memory domain conflicts for LSTM-1024 in CS task, (D) set sizes  $S$  in SMU task, (E) memory resolution  $n$  in VSRec task, and (F) number of features or conjunctions of features per item in CD task. See Sec. 3 for all the task introductions. See Sec. 5 for the analysis of these results.
Training Details. Practically, the WM models output responses for all  $T$  steps. However, training supervision with corresponding losses is only applied on the set of time steps where actual responses are required. All model parameters, including the feature extractors and WM networks, are initialized using Xavier initialization and trained from scratch. See Appendix A2.1 for more training details and Tab. A1 for hyper-parameter configurations.
Comparison with humans. To quantitatively compare computational models and humans on behavioral benchmarks, we introduce two types of scores: the slope difference score  $(A)$  and the average accuracy difference score  $(B)$ . We presented the experimental details and the results in Appendix A2.3. We conducted statistical tests on the results. We found that the majority of the models are more consistent with humans than the chance model; however, there still exists performance discrepancies among different models and performance gaps between these models and humans.
We analyze multiple memory characteristics (Fig. 2c) individually for the corresponding tasks for different WM models. We use "model-  $C$  " to refer to the WM model of memory capacity  $C$ . For example, LSTM-128 refers to the model with LSTM of memory capacity  $C = 128$  as the backbone. As benchmarks, we copy the exact human performances here from [13, 35, 29, 4, 60, 41, 52, 44, 38] for direct comparison with WM models. Importantly, the WM models were not trained on human data or influenced by human biases. Therefore, our focus is not on comparing absolute performance disparities between humans and WM models. Instead, we aim to examine their qualitative trends across different conditions. Due to the extensive combinations of models and tasks, we present a condensed selection of results for some models in this section. See Appendix A4 for the expanded result analysis with all model variations in these experiments. For fair comparisons among different model architectures in terms of overall WM performance, we compare different architectures with either the same memory capacity sizes  $C$  or the same number of network parameters. In both comparisons, we observe that recurrent networks like LSTMs and GRUs perform better than transformers on different working memory tasks (refer Appendix A4 for more details).
A. Primacy and recency effect across list lengths is an emergent phenomenon. We report the top-1 accuracy as a function of serial positions across various list lengths  $L$  for LSTM-256 in the VSR task (Sec. 3, Fig. 4A). We make several observations. Firstly, as  $L$  increases, the memory load in WM also increases, resulting in a decrease in overall accuracy across different list lengths (indicated by the varying colors of the lines, from blue to brown). Secondly, our findings in WM models align with the well-known cognitive science and neuroscience phenomena of the primacy and recency effects. These effects demonstrate that items presented at the beginning (primacy) and end (recency)
of a trial are better retained than those presented in the middle. This pattern holds true across different  $L$  values, indicating that the primacy and recency effects are an emergent phenomenon in our WM model. Thirdly, we observed an asymmetry in the strength of the primacy and recency effects, with the primacy effect being more prominent than the recency effect. This is evident when comparing accuracy at positions 1 and 9 for  $L = 9$ . Furthermore, the primacy and recency effects become more pronounced with longer list lengths  $L$ , suggesting that these effects are more prominent when the memory load is high. For instance, the effects are minimal for  $L = 3$  but become significant for  $L = 7$ . Fifthly, we compared the behaviors of WM models with humans and found that, though the models were never trained on human data, they approximate human-like qualitative trends in the primacy and recency effects across different serial positions and  $L$  values.
Furthermore, we conduct three ablations on VSR task to validate if these findings generalize across modalities, longer list lengths, and training paradigms. First, we report results for LSTM-256 in the Language Serial Recall (LSR) task, which is a language-based variant of the VSR task (Appendix A3.1, Fig. A2B). Secondly, we experiment with the LSR task containing trials with even longer list lengths (Appendix A3.2, Fig. A2C). Thirdly, instead of a joint training regime on all the tasks, we conduct individual training of a model only on the VSR task (Appendix A3.3, Fig. A2A). In all ablations, we observe similar primacy and recency effects suggesting that these findings are general across different modalities, longer list lengths, and various training paradigms. To further explore the underpinning mechanism attributing to intriguing and emergent primacy and recency behaviors in computational models, we expanded our discussions in Appendix A4.2.
B. Short retention intervals have minimal effects on working memory. We report the accuracy as a function of serial positions under the conditions of different retention intervals in the VIR task (Sec. 3, Fig. 4B). Aligning with human results, we observe that short retention intervals have minimal effects on WM performances, as indicated by the small performance gap between  $T_{\text{retent}}$  of 0 and 5 (blue versus orange lines). It is possible that the effect might be more dramatic with longer retention intervals. Moreover, same as Fig. 4A, here, we noted that there also exists a small recency effect for both models and humans. In particular, the effect is slightly stronger in  $T_{\text{retent}} = 5$  than  $T_{\text{retent}} = 0$ . However, this effect is less prominent than the ones observed in Fig. 4A, probably due to the short  $T_{\text{retent}} = 5$  or the small list length.
Moreover, we experimented on retention intervals of 2, 4, and 6 (shown in Appendix A3.5). We observed similar results as Fig. 4B with no significant performance differences across various short retention intervals (Fig. A3). In addition to the duration of retention intervals, we also studied the effect of the stimulus complexity during  $\mathrm{P}_{\text{retent}}$  on the working memory performances. We introduced a variant of the Change Detection (CD) task where we replaced gray images with random noise images (see Appendix A3.4 for more details).
C. Domain conflicts and increased cognitive loads impair working memory performances. To investigate whether domain conflicts would impair WM performance, we report accuracy in the CS task as a function of cognitive loads in different conditions involving combinations of visual and spatial domains (Sec. 3, Fig. 4C). In humans, regardless of the domain of the distractor task, we observe a monotonic decrease in accuracy with increasing cognitive load. This suggests that WM performance is impaired due to the cognitive load imposed by the distractor tasks, rather than the specific domain of the distractor task itself. However, in the WM model, we observe a decrease in accuracy only when the distractor task is of the spatial domain, implying that our spatial distractor task introduces a higher cognitive load than the visual distractor task, thereby hurting WM performance.
D. Monotonic decrease in accuracy with increasing set size. We report the agent's accuracy as a function of set sizes  $S$  in SMU task (Sec. 3) in Fig. 4D. For recurrent models, the recall accuracy monotonically decreases with increasing set sizes. For different recurrent backbones with the same memory capacities, GRU outperforms LSTMs and RNNs (compare GRU-96 vs RNN-96 vs LSTM-96). Moreover, for the same recurrent backbone, larger memory capacity enhances overall recall accuracy across all set sizes (RNN-96 vs RNN-256). Humans show a similar qualitative set size effect as recurrent models. In contrast, TRF-96 shows an opposite trend to humans and recurrent models, with increased accuracy as a function of increasing set size initially. Also, note that TRF-96 outperforms other RNN models in recall accuracy at set size 8 even with smaller memory capacity (e.g. RNN-256).
E. Working memory stores fine-grained details of visual patterns. We report the accuracy as a function of serial positions over various pattern-distractor differences  $n$  in VSRec (Sec. 3, Fig. 4E).
Figure 5: Visualization of neural correlates and task-specialized neural clusters (A) For LSTM-1024 in STS task, the visualization results of neural clusters based on response behaviors (left) and task identities (right) are presented. See Sec. 5G for details. In (B), we present the t-SNE visualization of neural clusters based on Task Variance (TV) of all hidden units in LSTM-256 over all 10 tasks. In (C), we present the neural selectivity of all hidden units for all the tasks. The neural selectivity is defined as TV values. All hidden units from left to right are sorted based on the clusters in (B). See the colorbar for TV values. In (D), we present the matrix of histogram plots for any pairs of tasks, where each histogram indicates the number of hidden units more selective to one task over the other. The selectivity for each hidden unit over a pair of tasks is defined as the fractional task variance (FTV). The x-axis and the y-axis of the histogram denote the FTV and the frequency of hidden units respectively. See Sec. 5H for details.
Intuitively, as it becomes easier to discern the differences between the correct pattern and the distractor, the more accurate it is to recall the correct patterns. Indeed, we see an increase in overall accuracy for bigger  $n$ . This trend is similar to human results. Moreover, as also observed in Fig. 4A, we see a primacy effect on both WM models and humans, where the accuracy is highest at position 1 and slowly decreases over subsequent positions. Interestingly, despite the difficulty of VSRec, for both humans and the models, the accuracy of recalling 36 cells correctly from the distractors with only 2 cell differences is still way above chance. This suggests that both humans and WM models are very good at memorizing fine-grained details of complex patterns.
F. Memory capacity is independent of the features or conjunctions of features. We study the capacity for storing different features and conjunctions in the CD task based on the agent's accuracy over set sizes under different feature conditions (Sec. 3, Fig. 4F). Humans display minimal accuracy differences across feature conditions, suggesting that memory capacity is independent of specific features or feature combinations. Similarly, WM models exhibit comparable accuracy regardless of feature conditions. However, caution is advised in interpreting model results, as accuracy saturates at  $100\%$  regardless of set size, likely due to overfitting. To draw stronger conclusions, further investigation in a low-data regime using our trials is needed.
G. Neural population carries active task representation in task-switching experiment. In STS task (Sec. 3), we present the visualization results of clusters based on neural activation of all hidden units across all trials and all conditions for LSTM-1024 in Fig. 5A. Specifically, we first take the hidden representation during  $\mathbf{P}_{probe}$  for all the trials within the task and then perform t-SNE [58] to project the hidden representation to 2D. On the left panel in Fig. 5A, we color-code the points based on the ground truth at that particular  $\mathbf{t}_{probe}$ . In this case, we did not observe distinct colored clusters whereas when we color-coded the points based on the task that was supposed to be performed at  $\mathbf{t}_{probe}$ , we observe two distinct colored clusters. In other words, the neural representations in the recurrent model encode task identities within a task-switching paradigm, thus playing a supervisory role. See Appendix A4.3 for more details.
H. Responses from neural populations reveal task selectivity during joint training. To study the task-specialized neural clusters in recurrent models, we present the visualization results in Fig. 5B-D. See Appendix A4.4 for the steps to obtain and interpret the visualization results. From Fig. 5B, C, we observed 8 distinct neural clusters and each neural cluster is selective to a different set of tasks.
For example, cluster 4 specializes in the SFR task, while cluster 6 is active in temporal domain tasks like VSR and VSRec. We also note that clusters of neurons mostly emerge in networks of lower capacity with memory constraints. From Fig. 5D, we show diverse neural populations where some have preferences for one task over the other, while some are equally active in both tasks. For instance, the entry in row 1, and column 5 indicates that the majority of hidden units are more selective to SC task compared with STS, as the distribution of neural population shifts to the right with the majority of the neurons having FTV = 1. Additionally, we conduct experiments where we lesion specific neural clusters to investigate causal evidence. As expected, we observe the corresponding performance drops in the accuracy of specific tasks depending on the neural cluster that has been lesioned (shown in Fig. A6). See Appendix A4.5 for details.
Despite significant research progress in studying individual aspects of working memory (WM), there remains a huge gap in the broader study of WM. We take initial steps in this direction by establishing a systematic and quantitative methodology to study the multi-facets of WM. However, it is important to recognize that we are still at the preliminary stage in a broader exploration of WM. The complexity, flexibility, and variability of WM present numerous avenues for future research. These avenues may involve refining the temporal alignment between human performance and model predictions, fine-tuning task difficulty across experiments by introducing variations in experimental parameters, experimenting with different training and testing methods, investigating the models' generalization abilities across a broader spectrum of tasks and modalities, and extending benchmarking experiments to include bio-inspired working memory models.
In our work, we introduce a comprehensive benchmark dataset consisting of 10 tasks, and 1 million trials covering 4 functionalities and 3 domains of WM. Moreover, we benchmark WM models and humans on the 10 tasks and compare their 11 memory characteristics at the behavioral, performance, and neural levels with a set of newly introduced evaluation metrics. We obtained insights about the underlying mechanism of WM from these tasks and identified the limitations in existing WM models.
It is worth noting that all the models were never trained with any human data or behavioral biases. Yet, these models still remarkably exhibit interesting alignments with human behaviors, such as in dependency on set sizes, pattern similarities, and primacy and recency effects. We report the top-1 accuracy for each condition in all the experiments; however, caution is needed when comparing the absolute performances between WM models and humans for multiple reasons: (1) differences in the concept of time, with humans perceiving time continuously in seconds while AI operates in discrete time steps; (2) varied learning mechanisms, including Hebbian learning in biological systems and back-propagation in AI; (3) distinct neuronal mechanisms between artificial and biological neurons, such as synaptic plasticity and structural plasticity; (4) disparities in the visual diets fed to humans and AI models, with humans exposed to video streams and AI models working with random static images; and (5) constraints on memory capacity varying between humans and AI models.
Our work contributes a suite of tools with detailed documentation for stimulus synthesis and parameter configurations in psychophysics experiments, training and testing state-of-the-art working memory models, evaluating model behaviors and neural responses, and comparing behavioral consistency between humans and models. This methodology serves as a valuable resource for comprehensive investigations into working memory, fostering advancements in both biological and artificial intelligence systems.
This research is supported by the National Research Foundation, Singapore under its AI Singapore Programme (AISG Award No: AISG2-RP-2021-025), its NRFF award NRF-NRFF15-2023-0001, Mengmi Zhang's Startup Grant from Agency for Science, Technology, and Research (A*STAR), and Early Career Investigatorship from Center for Frontier AI Research (CFAR), A*STAR. The authors declare that they have no competing interests.
A1 Detailed Descriptions about Psychophysics Experiments in WorM
Experiment A1 - Spatial Free Recall (SFR) [13]. This experiment focuses on the spatial domain of WM and aims to assess the ability to remember spatial locations and engage in immediate free recall of that information. Experiment schematic: The experiment begins with  $\mathbf{P}_{\text {present }}$  followed by  $\mathbf{P}_{\text {probe }}$ . In each trial, the stimulus image in  $\mathbf{P}_{\text {present }}$  contains 30 green squares randomly chosen from a  $10 \times 10$  grid. At every  $\mathbf{t}_{\text {present }}$ , one randomly selected square of the 30 green ones turns red. Finally, in  $\mathbf{P}_{\text {probe }}$ , a grey blank image is presented signaling the WM model to recall all the past spatial locations where the squares turned red. On the other hand, humans received an auditory cue to begin the recall. Humans indicate the highlighted squares by clicking on them in any order using a mouse cursor. The WM models, however, engage in a 100-way multi-label classification task with multi-hot encoded targets corresponding to the highlighted square locations. Afterward, we select the top- $L$  squares from the WM model's prediction based on the output probabilities, sort them in descending order, and generate the model's recall response, where  $L = T_{\text {present }}$  refers to the list length in the trial. Experiment conditions: We vary list length  $L = T_{\text {present }} = 1 - 8, 10, 12, 15, 18$ . Refer Fig. A1A for example trial.
Experiment A2 - Change Detection (CD) [35]. This experiment involves the visual domain of WM and aims to evaluate the agent's memory capacity for simple features and conjunctions. Experiment schematic: The experiment consists of  $\mathbf{P}_{\text {present }}$ ,  $\mathbf{P}_{\text {retent }}$ , and  $\mathbf{P}_{\text {probe }}$  in the respective order. In  $\mathbf{P}_{\text {present }}$  where  $\mathrm{T}_{\text {present }} = 1$ , a memory array, consisting of  $S$  bars with varying attributes in color (red or green), orientation (horizontal or vertical), size (small or big), and presence/absence of a gap in the middle of each bar, is presented. Here,  $S$  refers to the set size of the trial. During  $\mathbf{P}_{\text {retent }}$ , a series of blank grey images are presented. In  $\mathbf{P}_{\text {probe }}$ , where  $\mathrm{T}_{\text {probe }} = 1$ , agents are probed with a test array that can be either identical to the memory array or has one bar that differs along a single feature dimension. The agents have to make a binary prediction about whether or not there is a change given the test array. Experiment conditions: We formulate 5 main experiment conditions by varying the change of the test array in five aspects: color, orientation, bar size, presence or absence of the gap, and conjunctions of the four features where either of the four features could change. Additionally, we vary S from 2, 4, 6, 8, 10, 12 and also vary  $T_{\text {retent }}$  from 0, 6, 12, and 18. See Fig. A1B for an example trial from this experiment.
Experiment A3 - Visual Item Recognition (VIR) [29]. Here, we examine visual WM with different length lists and investigate the effect of retention interval on recognition performance. Experiment schematic: The experiment is structured with three phases presented in the following order:  $\mathbf{P}_{\text {present }}$ ,  $\mathbf{P}_{\text {retent }}$ , and  $\mathbf{P}_{\text {probe }}$ . In  $\mathbf{P}_{\text {present }}$ , a series of distinct matrix patterns are presented in sequence. These patterns are composed of  $6 \times 6$  square cells, with half of the cells being filled with red. In  $\mathbf{P}_{\text {retent }}$ , a series of grey blank images are presented. In  $\mathbf{P}_{\text {probe }}$ , where  $\mathbf{T}_{\text {probe}} = 1$ , a probe image is presented, consisting of two patterns presented side-by-side. One of these patterns is an exact match to one of the previously shown patterns, while the other pattern serves as a distractor. The distractor pattern is generated from the target pattern by changing 2 unfilled cells to filled ones and vice-versa. The agents have to perform a binary classification task, wherein 0 represents the left pattern and 1 represents the right pattern. Experiment conditions: We vary list length  $L = T_{\text {present }}$  from 4, 6, 8, and 10, and also vary  $T_{\text {retent }}$  from 0, 2, 4, 5, and 6. See Fig. A1C for an example trial.
Experiment A4 - Visual Serial Recall (VSR) [4]. This experiment evaluates the ability to accurately recall matrix patterns as well as the order in which they were presented. Therefore, this experiment encompasses both the visual and temporal domains of WM. Experiment schematic: The experiment consists of  $\mathbf{P}_{\text {present }}$  followed by  $\mathbf{P}_{\text {probe }}$ .  $\mathbf{P}_{\text {present }}$  involves the presentation of  $L$  matrix patterns in a sequential manner. These matrix patterns consist of square cells arranged in a  $6 \times 6$  grid, with half of the cells filled in green. In  $\mathbf{P}_{\text {probe }}$ , a probe image is displayed containing all the  $L$  patterns presented during  $\mathbf{P}_{\text {present }}$ . These  $L$  patterns are placed simultaneously at randomly chosen locations from a set of 9 evenly spaced locations in a  $3 \times 3$  grid. Locations that were not used remained blank. At each  $\mathbf{t}_{\text {probe }}$ , the agents perform a 9-way classification task to indicate the first pattern presented during  $\mathbf{P}_{\text {present }}$ , followed by the second pattern, and so on. Experiment conditions: We vary list length  $L = T_{\text {present }}$  from  $2 - 9$ . An example trial is shown in Fig. A1D.
Experiment A5 - Visual Serial Recognition (VSRec) [4]. In this experiment, agents are required to recall sequentially presented matrix patterns in a forward serial order within a recognition paradigm. Here, we also investigate the resolution of the recalled patterns by varying the difference between target patterns and distractor patterns. This experiment again involves both the visual and temporal