SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure A2: Ablation experiments. (A) We show the top-1 accuracy of the LSTM-256 model as a function of the serial position of the memory items in the VSR task. Here, the model was trained only on the VSR task. In (B), we show the behavioral accuracy as a function of list lengths in LSR task for LSTM-256 model. In (C), we show the plot for an LSTM-256 model trained only on the LSR task with longer sequence lengths (32 time steps).
A3.5 Fine-grained retention interval testing in VIR task
We evaluate the LSTM-1024 model on more retention intervals in the VIR task. Specifically, we test our model on 0, 2, 4, 5, and 6 retention intervals. We show these results in Fig. A3. We observe similar results as Fig. 4B, and we do not find any statistical significance in working memory performance differences across various short retention intervals.
A4 Additional Results and Analysis
A4.1 Behavioral Performance Comparison across Architectures and Memory Capacities
We analyzed various memory characteristics in Fig. 4. Here, we expand over different model architectures and model capacities in Fig A7, A8, A9, A10, A11, A12. We found that the performances in top-1 accuracy vary across architectures and memory capacities. In general, the larger the memory capacity, the better the behavioral performance. However, there are a few exceptions where the models with smaller capacities yield higher accuracy (e.g. Fig. A11, J versus K). We also notice that some tasks are extremely difficult for all WM models to learn regardless of the
Figure A3: Short retention interval testing in VIR task. We show the results for LSTM-1024 on the VIR task with multiple retention intervals including 0, 2, 4, 5, and 6. The chance level is 0.5 for all the above plots.
Models	Hidden size	Number of parameters
RNN-96	96	1.8M
RNN-256	256	2M
RNN-1024	1024	3.5M
GRU-96	96	2M
GRU-256	256	2.5M
GRU-1024	1024	6.5M
LSTM-96	96	2M
LSTM-256	256	2.5M
LSTM-1024	1024	8M
TRF-96	96	2.5M
TRF-256	256	4.5M
TRF-1024	1024	20M
Table A4: Different computational models of varying memory capacity with their number of parameters.
architectures. For example, in Fig. A9, almost all the WM models fail. Even with a large memory capacity of 1024, which proves to be sufficient for performing other WM tasks, it still performs badly in the CS task.
For a fair comparison, we can either compare different architectures with an equal number of hidden sizes or with an equal number of network parameters. To compare different architectures with the same hidden size, we compare RNN-1024, GRU-1024, LSTM-1024, and TRF-1024. These models are jointly trained on all the working memory tasks. We show the joint average accuracy, i.e., the average accuracy of these models over all tasks in Fig. A4. We observe that LSTM-1024 achieves the best joint average accuracy, followed by GRU-1024, TRF-1024, and RNN-1024 respectively.
Similarly, to compare architectures with the same number of parameters, we can compare RNN-256, GRU-256, LSTM-256, and TRF-96 (see Tab. A4 for the number of parameters in different models). In Fig. A4, we again observe LSTM-256 is the best performing model, followed by GRU-256, TRF-96, and RNN-256 respectively. In either comparison, we can note that recurrent networks like LSTMs and GRUs perform better than transformers. While recurrent architectures (GRUs, LSTMs) generally perform better in top-1 accuracy and exhibit the primacy and recency effect (Sec. 5A), the transformer architectures fail to replicate such phenomenon; and their top-1 accuracy is also lower than the recurrent architectures.
Figure A4: Training curves for different model architectures and capacities. The y-axis shows the averaged validation accuracy across all tasks and the x-axis shows epochs.
A4.2 Further probing into the mechanism attributing to primacy effects
In Fig. 4A and Fig. A2, we observe consistent primacy effects. Our hypothesis for the primacy effect was that the first stimulus stored in the working memory acts as the first principal component, steering the memory representations of the subsequent stimulus to align with the representation of the first stimulus. Hence, the first item was forgotten less; thus, the network exhibits the primacy effect. To validate this hypothesis, in every trial of the Visual Serial Recall (VSR) task, we first compute the hidden states at time step 1, when the first item is presented, and then compute its cosine similarity with the hidden state representations over subsequent time steps  $t$  when subsequent memory items are presented. We report the results in Fig. A5. Based on our hypothesis, the cosine similarity between the  $1^{st}$  time step with  $n^{th}$  time step should be higher than the cosine similarity between  $2^{nd}$  time step and  $n^{th}$  time step. However, that is not the case in Fig. A5 (see along column), suggesting that our hypothesis is incorrect. Future work should explore other plausible hypotheses for the primacy effect in recurrent models.
Figure A5: Cosine similarity of memory representations between different serial positions. We show the cosine similarity between the memory representation of different serial positions, averaged across all trials in the VSR task. We consider the LSTM-256 model for the above plot.
A4.3 Visualizing Neural Correlates in Task Switching
In the STS experiment, we first take the hidden representations at each  $\mathbf{t}_{probe}$  for all the test trials within the task and then perform t-SNE [58] to project these hidden representations to 2D.
Next, we introduce two types of color schemes to label hidden representations. First, on the left panel in Fig. 5A, we color-code the points based on the ground truth at that particular  $\mathbf{t}_{probe}$ . We assign a blue color to label those hidden representations where the ground truth is 0 i.e. when the marker is either at the left or top of the grid and we assign an orange color to label those hidden representations where the ground truth is 1 i.e. when the marker is either at the right or bottom of the grids. In essence, we color-code the points based on the ground truth response for that particular  $\mathbf{t}_{probe}$ .
Second, on the right panel in Fig. 5A, we color-code the hidden representations based on the task that was supposed to be performed at  $\mathrm{t}_{probe}$ . We assign a green color to label those hidden representations when the task at hand is to discriminate top versus bottom and we assign a red color to label those hidden representations when the task at hand is to discriminate left versus right.
By color-coding the neural representations based on the task to be performed at  $\mathbf{t}_{probe}$ , we observed the emergence of two distinct clusters corresponding to the two tasks involved, namely top-versus-bottom and left-versus-right. This observation highlights how the model's neural representations encode and maintain task identity over time, enabling accurate performance in task-switching experiments. Importantly, it should be noted that at  $\mathbf{t}_{probe}$ , the model does not receive any explicit information regarding the task to be performed at that specific time step.
A4.4 Visualizing Task-Specialized Neural Clusters
Task Variance (TV) [63] is a scalar that quantifies the selectivity of a hidden unit within a recurrent model towards a specific task, where higher values indicate greater selectivity. To obtain the results
in Fig. 5B, C in the main text, we first compute the TV for each hidden unit for each task, resulting in a TV matrix of size  $10 \times 256$  for LSTM-256. We normalize the task variance of each hidden unit such that the maximum normalized TV across all tasks was 1. We then perform t-SNE along the TV dimensions and project the TV vector of each hidden unit to 2D. Subsequently, we identify 8 clusters among the hidden units and assign unique colors to each cluster (Fig 5B). In line with [63], we chose the optimal number of clusters based on the highest silhouette score. We then group the hidden units based on their clusters. The sorted TV matrix based on the clusters is shown in Fig. 5C in the main text, where the intensity denotes the magnitude of normalized TV values.
Given a pair of tasks A and B, the fractional task variance (FTV) [63] defines the preference of a hidden unit for one task over the other. For a particular hidden unit  $i$  in a WM model, the FTV with respect to task A and task B is defined as  $FTV_{i}(A,B) = \frac{TV_{i}(A) - TV_{i}(B)}{TV_{i}(A) + TV_{i}(B)}$ , where  $TV_{i}(A)$  and  $TV_{i}(B)$  are the TVs for tasks A and B, respectively. FTV(A, B) ranges from -1 to 1, with 1 being more selective to task A over B, and vice versa. FTV = 0, when the hidden unit is neutral about both tasks and stays equally active during both tasks. To obtain the results in Fig. 5D in the main text, for each pair of tasks, we compute the FTV for each hidden unit and plot the histogram of the total number of hidden units based on their FTV values. In the histogram, the x-axis denotes the FTV whereas the y-axis denotes the proportion of hidden units. If the histogram distribution is skewed towards FTV = 1, it means that the majority of the hidden units are more selective to Task A and vice versa if the distribution is skewed toward FTV = -1. If the distribution is unimodal in the center, it implies that the majority of the hidden units do not have task preference and are equally active in both tasks. In the case of the bimodal distribution where the two modes are located at both ends, this implies that there are two separate populations of neurons, dedicated to two individual tasks respectively.
A4.5 Causal Analysis of Neural Population
In Fig A6, we show the drop in performance for different tasks, when specific neural clusters from the network (shown in Fig. 5C) are lesioned. To lesion a neural cluster, we select neurons within that cluster and zero out their projection weights to task-specific classifiers. As anticipated, we note the corresponding causal evidence when distinct clusters are lesioned. For instance, when cluster 6 is deactivated, we observe the maximum decline in performance for the VSR and VSRec tasks. Similarly, when cluster 3 is lesioned, the maximum drop in performance is for the CD task.
A4.6 Visualizing Task Similarity Matrix
We also visualize the task similarity matrix based on the learned task embeddings for each WM model of different memory capacities in Fig. A13. As expected, we observe the brightest values along
Figure A6: Effect of lesioning neural clusters on task performance. The y-axis shows the different tasks and the x-axis shows the neural cluster which is lesioned. See Sec. 5 and Fig. 5 in the main text for neural clusters and analysis.
the diagonal of the task similarity matrix as the embeddings from one task compare against itself. Interestingly, we found that in recurrent architectures, task similarities between different conditions within the CD task are high Fig. A13A-I. However, we did not observe such strong similarities in transformers Fig. A13J-L.
Figure A7: Behavioral accuracy as a function of list lengths in VSR task. In addition to Fig. 4A in the main text, we show the results of three recurrent architectures (first 3 rows) with three memory capacities (3 columns) and the transformer architecture with two memory capacities (J-L). The chance level is 0.11 for all the above plots.
Figure A8: Behavioral accuracy as a function of retention interval in VIR task. In addition to Fig. 4B in the main text, we show the results of four architectures with different memory capacities. The layout interpretations follow Fig. A7. The chance level is 0.5 for all the above plots.
Figure A9: Behavioral accuracy as a function of memory domain conflicts in CS task. In addition to Fig. 4C in the main text, we show the results of four architectures with different memory capacities. The layout interpretations follow Fig. A7. The chance level is 0.04 for all the above plots.
Figure A10: Behavioral accuracy as a function of set size in SMU task. In addition to Fig. 4D in the main text, we show the results of four architectures with different memory capacities. The layout interpretations follow Fig. A7. The chance level is 0.11 for the above plot.
Figure A11: Behavioral accuracy as a function of memory resolution  $n$  in VSRec task. In addition to Fig. 4E in the main text, we show the results of four architectures with different memory capacities. The layout interpretations follow Fig. A7. The chance level is 0.5 for all the above plots.
Figure A12: Behavioral accuracy as a function of features and conjunctions in CD task. In addition to Fig. 4F in the main text, we show the results of four architectures with different memory capacities. The layout interpretations follow Fig. A7. The chance level is 0.5 for all the above plots.
Figure A13: Visualization of similarity matrix based on learned task embeddings from WM models with different memory capacities. We take the task embedding learned by each WM model with different memory capacities and compute the cosine similarity between a pair of task embedding vectors from two corresponding tasks. We present the similarity matrix between tasks. Each row of the similarity matrix is normalized such that the maximum similarity is 1 and the minimum is 0. See the color bar on the right for the normalized similarity values. The brighter values indicate that the learned task embeddings are more similar for the two corresponding tasks.
Figure A14: Training curve of individual tasks for LSTM-1024. The y-axis shows the validation accuracy for the tasks and the x-axis shows epochs. The training curves for other models show similar trends; thus, they are omitted for simplicity. Here, we plot the model with the highest joint validation accuracy, i.e., LSTM-1024.
A5 Datasheet for WorM Dataset
A5.1 Motivation
1. For what purpose was the dataset created? Was there a specific task in mind? Was there a specific gap that needed to be filled? Please provide a description.
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. 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.
2. Who created the dataset (e.g., which team, research group) and on behalf of which entity (e.g., company, institution, organization)?
The dataset was created by Ankur Sikarwar and Mengmi Zhang from Deep NeuroCognition Lab in affiliation with Center for Frontier AI Research (CFAR), and Institute for Infocomm Research (I2R), from Agency for Science, Technology, and Research (A*STAR), Singapore.
3. Who funded the creation of the dataset? If there is an associated grant, please provide the name of the grantor and the grant name and number
This creation of the dataset is supported by the National Research Foundation, Singapore under its AI Singapore Programme (AISG Award No: AISG2-RP-2021-025) and its NRFF award NRF-NRFF15-2023-0001.
4. Any other comments?
A5.2 Composition
1. What do the instances that comprise the dataset represent (e.g., documents, photos, people, countries)? Are there multiple types of instances (e.g., movies, users, and ratings; people and interactions between them; nodes and edges)? Please provide a description.
The dataset consists of 10 distinct working memory tasks. Each data instance in the dataset represents a trial that is time-based and includes stimuli and corresponding responses for various time steps within the trial.
2. How many instances are there in total (of each type, if appropriate)?
The dataset consists of 10 working memory (WM) tasks. One of these tasks i.e. the CD task contains 5 separate task conditions, each of which has separate training, validation, and testing trials. For each of these task conditions within the CD task and 9 remaining WM tasks, there are 86,400 trials for training, 9,600 trials for validation, and 9,600 trials for testing.
3. Does the dataset contain all possible instances or is it a sample (not necessarily random) of instances from a larger set? If the dataset is a sample, then what is the larger set? Is the sample representative of the larger set (e.g., geographic coverage)? If so, please describe how this representativeness was validated/verified. If it is not representative of the larger set, please describe why not (e.g., to cover a more diverse range of instances, because instances were withheld or unavailable).
Since the dataset is synthetic, it is possible to generate an arbitrary number of samples with the provided code. For consistency, we limit the total number of generated samples, including training, validation, and testing, to 105,600 trials for each task.
4. What data does each instance consist of? "Raw" data (e.g., unprocessed text or images) or features? In either case, please provide a description.
Every instance in the dataset includes raw image stimuli corresponding to each time step within the trial, along with the corresponding responses for that particular trial.
5. Is there a label or target associated with each instance? If so, please provide a description.
The label for each instance is the ground-truth response corresponding to the task and the specific trial.
6. Is any information missing from individual instances? If so, please provide a description, explaining why this information is missing (e.g., because it was unavailable). This does not include intentionally removed information, but might include, e.g., redacted text.
There is no missing information; the data is complete.
7. Are relationships between individual instances made explicit (e.g., users' movie ratings, social network links)? If so, please describe how these relationships are made explicit. N/A
8. Are there recommended data splits (e.g., training, development/validation, testing)? If so, please provide a description of these splits, explaining the rationale behind them.
We propose our own training, validation, and testing splits for each task that is available on the GitHub repository.
9. Are there any errors, sources of noise, or redundancies in the dataset? If so, please provide a description.
We generate trials through random sampling. While it is theoretically possible to generate the same trial more than once, the probability of this occurrence is extremely low due to the vast number of possible combinations of stimuli across multiple time steps.
10. Is the dataset self-contained, or does it link to or otherwise rely on external resources (e.g., websites, tweets, other datasets)? If it links to or relies on external resources, a) are there guarantees that they will exist, and remain constant, over time; b) are there official archival versions of the complete dataset (i.e., including the external resources as they existed at the time the dataset was created); c) are there any restrictions (e.g., licenses, fees) associated with any of the external resources that might apply to a dataset consumer? Please provide descriptions of all external resources and any restrictions associated with them, as well as links or other access points, as appropriate.
11. Does the dataset contain data that might be considered confidential (e.g., data that is protected by legal privilege or by doctor-patient confidentiality, data that includes the content of individuals' non-public communications)? If so, please provide a description.
N/A. The dataset is synthetic
12. Does the dataset contain data that, if viewed directly, might be offensive, insulting, threatening, or might otherwise cause anxiety? If so, please describe why.
N/A. The dataset is synthetic.
A5.3 Collection Process
N/A. No data collection process was involved. The dataset has been synthetically generated and we make our dataset generation code available on GitHub. The human data was collected from the existing literature. We directly borrowed their human data and plot them side by side for comparisons with WM models.
A5.4 Preprocessing/cleaning/labeling
1. Was any preprocessing/cleaning/labeling of the data done (e.g., discretization or bucketing, tokenization, part-of-speech tagging, SIFT feature extraction, removal of instances, processing of missing values)? If so, please provide a description. If not, you may skip the remaining questions in this section.
No preprocessing was done.
2. Was the "raw" data saved in addition to the preprocessed/cleaned/labeled data (e.g., to support unanticipated future uses)? If so, please provide a link or other access point to the "raw" data.
3. Is the software that was used to preprocess/clean/label the data available? If so, please provide a link or other access point. N/A
4. Any other comments? N/A
A5.5 Uses
1. Has the dataset been used for any tasks already? If so, please provide a description. The dataset is being used for the first time.
2. Is there a repository that links to any or all papers or systems that use the dataset? If so, please provide a link or other access point. N/A
3. What (other) tasks could the dataset be used for?
The dataset itself contains ten well-defined WM tasks and is intended to be used for those specific tasks.
4. Is there anything about the composition of the dataset or the way it was collected and preprocessed/cleaned/labeled that might impact future uses? For example, is there anything that a dataset consumer might need to know to avoid uses that could result in unfair treatment of individuals or groups (e.g., stereotyping, quality of service issues) or other risks or harms (e.g., legal risks, financial harms)? If so, please provide a description. Is there anything a dataset consumer could do to mitigate these risks or harms?
No, there are no risks or harms of the dataset. The dataset contains no human data.
5. Are there tasks for which the dataset should not be used? If so, please provide a description.
6. Any other comments?
A5.6 Distribution
1. Will the dataset be distributed to third parties outside of the entity (e.g., company, institution, organization) on behalf of which the dataset was created? If so, please provide a description.
The dataset is publicly available.
2. How will the dataset will be distributed (e.g., tarball on website, API, GitHub)? Does the dataset have a digital object identifier (DOI)?
The dataset and the supplementary code is available on the GitHub repository.
3. When will the dataset be distributed?
The dataset is available from June 2023.
4. Will the dataset be distributed under a copyright or other intellectual property (IP) license, and/or under applicable terms of use (ToU)? If so, please describe this license and/or ToU, and provide a link or other access point to, or otherwise reproduce, any relevant licensing terms or ToU, as well as any fees associated with these restrictions.
The dataset is released under the MIT License.
5. Have any third parties imposed IP-based or other restrictions on the data associated with the instances? If so, please describe these restrictions, and provide a link or other access point to, or otherwise reproduce, any relevant licensing terms, as well as any fees associated with these restrictions.
6. Do any export controls or other regulatory restrictions apply to the dataset or to individual instances? If so, please describe these restrictions, and provide a link or other access point to, or otherwise reproduce, any supporting documentation.
7. Any other comments?
A5.7 Maintenance
1. Who will be supporting/hosting/maintaining the dataset?
Ankur Sikarwar is supporting and maintaining the dataset.
2. How can the owner/curator/manager of the dataset be contacted (e.g., email address)?  
ankursikarwards@gmail.com
3. Is there an erratum? If so, please provide a link or other access point.
The GitHub repository will reflect any changes or improvements made to the dataset.
4. Will the dataset be updated (e.g., to correct labeling errors, add new instances, delete instances)? If so, please describe how often, by whom, and how updates will be communicated to dataset consumers (e.g., mailing list, GitHub)?