SlowGuess/NeurIPS_2023_OCR huggingface dataset:
domains of WM. Experiment schematic: The experiment is similar to the VSR experiment above in terms of the setup. The experiment consists of  $\mathbf{P}_{\text {present }}$  and  $\mathbf{P}_{\text {probe }}$ .  $\mathbf{P}_{\text {present }}$  is exactly the same as  $\mathbf{P}_{\text {present }}$  in the VSR experiment. The only differences happen at  $\mathbf{P}_{\text {probe }}$ . During  $\mathbf{P}_{\text {probe }}$ , a series of two-alternative recognition tests are presented in consecutive time steps to assess memory for the presented patterns in forward serial order. Essentially, at every  $\mathbf{t}_{\text {probe }}$ , the target pattern from corresponding  $\mathbf{t}_{\text {present }}$  and a distractor pattern are presented side by side. The agents perform a binary classification task to indicate the target pattern, where 0 corresponded to the left pattern and 1 corresponded to the right pattern. The distractor patterns were generated from the target pattern by changing the values of n cells such that the total number of filled cells remained constant. We denote  $n$  as the pattern-distractor difference. We made sure that the distractor patterns were not among the initially presented matrix patterns. Experiment conditions: We vary list length  $L = T_{\text {present }}$  from 2 to 9 and the pattern-distractor difference  $n$  from the range of 2, 4, 6, 8, and 10. See Fig. A1E for an example trial.
Experiment A6 - Complex Span (CS) [60]. This experiment aims to investigate the interference between the visual and spatial domains of WM, as well as the effect of cognitive load on recall performance. Moreover, in this experiment, agents are required to perform a serial recall of information. Hence, this experiment encompasses all three domains of WM: visual, spatial, and temporal. Experiment schematics: The experiment consists of two types of  $\mathbf{P}_{\text{present}}$  and two corresponding types of  $\mathbf{P}_{\text{probe}}$ , and two types of  $\mathbf{P}_{\text{distractor}}$ . (I) The first type of  $\mathbf{P}_{\text{present}}$  and  $\mathbf{P}_{\text{probe}}$  corresponds to visual storage and involves memorizing a sequence of visual patterns in  $\mathbf{P}_{\text{present}}$  and recalling the exact visual patterns in a serial order in  $\mathbf{P}_{\text{probe}}$ . The visual patterns are represented as 4 × 4 grid, where half of the cells are in red. (II) The second type of  $\mathbf{P}_{\text{present}}$  and  $\mathbf{P}_{\text{probe}}$  corresponds to spatial storage and involves memorizing a sequence of spatial patterns involving ball movements in  $\mathbf{P}_{\text{present}}$  and recalling the exact directions of ball movements in a serial order in  $\mathbf{P}_{\text{probe}}$ . To interfere with the spatial/visual memories in  $\mathbf{P}_{\text{present}}$ , two types of  $\mathbf{P}_{\text{distractor}}$  are introduced: (III) The first type of  $\mathbf{P}_{\text{distractor}}$  corresponds to the visual domain and involves color discrimination tasks where agents have to classify the color of the given panel as blue or red. (IV) The second type of  $\mathbf{P}_{\text{distractor}}$  corresponds to the spatial domain and involves symmetry discrimination tasks where agents have to classify whether the given pattern is symmetric about the vertical axis or not. Experiment conditions: We consider four different variations:  $\mathrm{I}(\mathrm{P}_{\text{present}}) + \mathrm{III}(\mathrm{P}_{\text{distractor}}) + \mathrm{I}(\mathrm{P}_{\text{probe}})$ ;  $\mathrm{II}(\mathrm{P}_{\text{present}}) + \mathrm{III}(\mathrm{P}_{\text{distractor}}) + \mathrm{II}(\mathrm{P}_{\text{probe}})$ ;  $\mathrm{I}(\mathrm{P}_{\text{present}}) + \mathrm{IV}(\mathrm{P}_{\text{distractor}}) + \mathrm{I}(\mathrm{P}_{\text{probe}})$ ; and  $\mathrm{II}(\mathrm{P}_{\text{present}}) + \mathrm{IV}(\mathrm{P}_{\text{distractor}}) + \mathrm{II}(\mathrm{P}_{\text{probe}})$ ; namely, visual + visual, spatial+visual, visual+spatial, and spatial + spatial respectively. Additionally, we change the cognitive load by varying  $T_{\text{distractor}}$  from 0, 1, 3, and 5. See Fig. A1F for an example trial.
Experiment B1 - Spatial Coordination (SC) [41]. The objective of this experiment is to assess the agents' working memory (WM) capacity in the spatial domain, specifically focusing on their ability to coordinate and integrate spatial information across various time steps. Experiment schematics: The experiment consists of  $\mathrm{P}_{\text {present }}$  and  $\mathrm{P}_{\text {probe }}$  in the respective order. At every  $\mathrm{t}_{\text {present }}$ , one cell in a  $10 \times 10$  grid is highlighted in green. In  $\mathrm{P}_{\text {probe }}$ , the agents are probed using a blank grey image to perform a binary discrimination task to indicate whether the pattern integrated mentally over all  $\mathrm{T}_{\text {present }}$  is symmetric about the vertical axis or not. Experimental conditions: We vary the list lengths  $L = T_{\text {present}}$  from 10, 12, 14, 16, and 18. Refer Fig. A1G for an example trial.
Experiment B2 - Spatial Integration (SI) [52]. This experiment focuses on the spatial domain and aims to investigate the agents' abilities to integrate spatial information. Experiment schematics: The experiment consists of  $\mathbf{P}_{\text {present }}$  and  $\mathbf{P}_{\text {probe }}$ . In  $\mathbf{P}_{\text {present }}$ , there is a sequential presentation of partial line drawings at every  $\mathbf{t}_{\text {present }}$ . In  $\mathbf{P}_{\text {probe }}$ , where  $\mathbf{T}_{\text {probe}} = 1$ , a complete line drawing containing 12 line segments is presented. The task is to mentally integrate the partial drawings shown during  $\mathbf{P}_{\text {present }}$  and compare them to the complete line drawing shown during  $\mathbf{P}_{\text {probe }}$ . Essentially, the agents perform a binary classification to indicate whether the figure shown in  $\mathbf{P}_{\text {probe }}$  is identical to the mentally integrated figure from  $\mathbf{P}_{\text {present }}$ . All line drawings were generated by connecting neighboring points within an imaginary  $4 \times 4$  grid. Experiment conditions: We vary the number of line segments in the partial drawings from 12, 6, 4, 3, 2, and 1, resulting in list lengths  $L = T_{\text {present }}$  of 1, 2, 3, 4, 6, and 12 respectively. As the list length increases, the number of integrations needed to solve the trial also increases. See Fig. A1H for an example trial.
Experiment C1 - Spatial Memory Updating (SMU) [44]. This experiment assesses the capabilities of memory manipulation in the spatial domain. Experiment schematics: The experiment consists of  $\mathbf{P}_{\text {present }}$ ,  $\mathbf{P}_{\text {update }}$ , and  $\mathbf{P}_{\text {probe }}$  in the respective order. In  $\mathbf{P}_{\text {present }}$  where  $\mathbf{T}_{\text {present }} = 1$ , a stimulus is
presented, consisting of  $S$  green squares, arranged in an imaginary circular formation. Within each square, there is a red marker located in one of the nine possible locations within an invisible  $3 \times 3$  grid. Agents are required to memorize the spatial locations of the markers within each square. Next, at every  $t_{update}$  in  $\mathbf{P}_{update}$ , an arrow is presented at the center of one of the squares, indicating the direction for mentally updating the marker's new location. The arrows can be oriented vertically, horizontally, or diagonally, with the condition that the marker never exits the square. A sequence of a total of 8 memory update operations is presented in a clockwise sequence through the  $S$  squares. Hence,  $T_{update} = 8$ . Finally, at every  $t_{probe}$  of  $\mathbf{P}_{probe}$ , one out of the  $S$  squares is highlighted in red. This red square acts as a cue for the agents to recall and indicate the marker's location within that particular square through a 9-way classification task. The red marker locations for all  $S$  squares are probed in random order until  $t_{probe} = S$ . Experiment conditions: We vary  $S$  from 1 to 8. Refer Fig. A1I for an example trial.
Experiment D1 - Spatial Task Switching (STS) [38]. This experiment involves the spatial domain of WM and investigates the agent's capacity to flexibly transition between two spatial discrimination tasks based on supervision cues. Experiment schematics: The experiment consists of two types of phases  $\mathbf{P}_{\text {switch }}$  and  $\mathbf{P}_{\text {probe }}$  presented in random order. In  $\mathbf{P}_{\text {switch }}$ , a cue image is presented to the agents to indicate the discrimination task to be performed for subsequent  $\mathbf{P}_{\text {probe }}$ . There are two binary discrimination tasks: top-versus-bottom and left-versus-right. For the left-versus-right discrimination task, the cue image contains dash markers placed on the left and right sides of the image. In contrast, for the top-versus-bottom discrimination task, the dash markers are positioned on the top and bottom sides of the image. At every  $t_{\text {probe }}$  in  $\mathbf{P}_{\text {probe }}$ , a red marker is presented in any of the four locations within a  $2 \times 2$  grid. The agents are required to switch between the two discrimination tasks during every  $\mathbf{P}_{\text {switch }}$  and respond with 0 for top/left positions and 1 for bottom/right positions. Experiment conditions: We randomly vary the number of task switches i.e.  $\mathbf{P}_{\text {switch }}$  in different trials as well as the time step at which task switches occur. See Fig. A1J for an example trial.
Experimental trial splits. In each WM psychophysics experiment, we generate a total of 86,400 trials for training, 9,600 trials for validation, and 9,600 trials for testing. Importantly, note that for the CD experiment, there are five distinct conditions, each of which consists of its own set of 86,400 trials for training, 9,600 trials for validation, and 9,600 trials for testing.
Time
Figure A1: Detailed schematic of each experiment. We expand the overviews of all ten experiments introduced in Fig. 1 and Sec. 3 with detailed experiment schematics. Each column represents the schematic of one experiment. The time goes from the top to the bottom. Each row presents the example stimulus at that time step  $t$ . The legend and interpretations of this figure follow Fig. 1.
A2 Computational Models of Working Memory
A2.1 Additional Implementation and Training Details
Feature Extractor. In all our models, we employ a 4-layer 2D-convolutional network with 64, 128, 256, and 512 channels in the successive layers. We also experimented with ViT [21] pre-trained on ImageNet, where we freeze the model and use it for feature extraction. However, features from natural images of ImageNet [18] do not generalize to synthetic images of our dataset.
Task Embeddings. We learn task-specific embeddings  $M \in \mathbb{R}^{14 \times 14}$ , which informs the model of which task needs to be performed. There are a total of 10 WM tasks. Although, the Change Detection (CD) task comprises five distinct task conditions, namely color, orientation, size, gap, and conjunction, each of which necessitates a unique task identifier. Consequently, there are a total of 14  $(9 + 5)$  unique task embeddings.
Training Details. During training, we sample 10 trials from each task, obtain the model's responses for all the sampled trials, and compute task-specific losses for each respective task. These individual losses are then aggregated and the joint loss is utilized to perform back-propagation. In essence, the model parameters are optimized jointly over all tasks. Note that we formulate the SFR task as multi-label classification and the other tasks as multi-class or binary classification. Consequently, the task-specific losses are computed using either cross-entropy or binary cross-entropy, depending on the nature of the tasks.
Implementation Details. All our models were trained from scratch. Hyperparameter tuning was done using grid search. We show the different hyperparameters for our models in Tab. A1. We employed the Adam optimizer [30] with a starting learning rate of 1e-4 and the first moment of 0.9. The learning scheduler reduces the learning rate by a factor of 0.8 when the validation loss doesn't improve for 3 epochs. We report the average top-1 accuracy and standard error for various model architectures and capacities, computed across test trials from different conditions. For all our experiments, we used 16 NVIDIA RTX A5000 GPUs, each equipped with 24 GB of memory.
Data Generation Pipeline. Given any condition in a task, we follow the three steps below to synthesize stimulus for every time step. For an easy illustration of the three steps, we use the Visual Serial Recognition task (Appendix A1) as an example. First, depending on the list length of the trial, we generate multiple unique  $6 \times 6$  matrix patterns by sampling grid cells to be filled. These matrix patterns constitute the memory and distractor items for the trial. Second, we generate the ground truth labels for  $\mathrm{P}_{probe}$  based on the sampled matrix patterns. Third, given all the patterns sampled in Step 1, we synthesize the stimulus using the Pillow library [11].
Our systematic process ensures the creation of task-specific trials with multiple controlled variations and difficulty. We provide detailed documentation and made our code easy to use and customize for future studies, such as generalization tests of working memory models with longer task sequences, and larger set sizes. Our code and data are publicly available: link.
A2.2 Time mapping between WM models and human psychophysics experiments
There could be multiple ways of establishing the time mapping between WM models and human psychophysics experiments (Fig. 1). Here, for the WM models, we stick to 1 stimulus per time step rule. In Tab. A2, we report the mapping of a single time step in WM models to the duration of the corresponding stimulus presentation in human experiments. For instance, in the SFR human
Hyperparameters	RNN	GRU	LSTM	TRF
Input Size	512	512	512	-
Hidden / Embedding Size	[96, 256, 1024]	[96, 256, 1024]	[96, 256, 1024]	[96, 256, 1024]
TRF Encoder Layers	-	-	-	2
TRF Attention Heads	-	-	-	8
Batch Size	[140, 70]	[140, 70]	[140, 70]	[140, 70]
Learning Rate	[0.00008, 0.0001, 0.0003]	[0.00008, 0.0001, 0.0003]	[0.00008, 0.0001, 0.0003]	[0.00008, 0.0001, 0.0003]
Epochs	200	200	200	200
Table A1: Different hyperparameters considered for models trained on the WorM dataset. The best hyperparameters for each model are in bold. Here, batch size refers to the total number of trials from all tasks within one batch.
TASK	Ppresent	Pretent	Pprobe	Pdistractor	Pswitch	Pupdate
SFR	750 ms	*	-	*	*	*
CD	100 ms	150 ms	2000 ms	*	*	*
VIR	1000 ms	1000 ms	-	*	*	*
VSR	1550 ms	*	-	*	*	*
VSRec	1550 ms	*	-	*	*	*
CS	500 ms	*	-	1700 ms	*	*
SC	1000 ms	*	-	*	*	*
SI	-	*	-	*	*	*
SMU	-	*	-	*	*	-
STS	*	*	1500 ms	*	203 ms	*
Table A2: Time mapping between WM models and Human psychophysics experiments. We show the mapping of a single time step in WM models to ms in different phases of the corresponding human psychophysics experiments. Asterisks (*) indicate the absence of a specific phase in the experiment. Dashes (-) indicate that either humans were provided with an indefinite amount of time during that phase, or the duration of stimulus presentation was variable and not standardized.
experiment, the duration of each stimulus presentation during  $\mathrm{P}_{\text {present }}$  was  $750~\mathrm{ms}$ , whereas the WM models were exposed to each stimulus for only 1 time step. Therefore, within the  $\mathrm{P}_{\text {present }}$  phase of the SFR experiment, a single time step of the WM models corresponds to a duration of  $750~\mathrm{ms}$  in the human experiments.
A2.3 Quantitative comparison with human behaviors
To quantitatively compare computational models and humans, we introduce two types of scores: the slope difference score  $(A)$  and the average accuracy difference score  $(B)$ . First, we define the slope difference score  $(A)$ , as  $|(M_A - H_A) / H_A|$ , where  $M_A$  and  $H_A$  are slopes for computational models and humans respectively. We use the slope difference score to capture similarity in the trend of accuracy against different conditions like serial position, set size, and cognitive load. Second, we define the average accuracy difference score  $(B)$ , as  $|(M_B - H_B) / H_B|$ , where  $M_B$  and  $H_B$  are average accuracy for computational models and humans respectively across different conditions. For example,  $A_{\text{patterndiff}}$  in VSRec task (Fig. 4E) for LSTM-256 model is calculated by first computing  $M_A$  i.e. the slope of the best-fit linear line for the model's top 1 accuracy as a function of serial positions in the plot. Likewise, we compute  $H_A$  and hence,  $A_{\text{patterndiff}}$ . Similarly, to calculate  $M_B$ , we first compute the average accuracy of the model over all the serial positions for each condition of pattern difference and then compute the slope of the best-fit linear line for the model's top 1 accuracy as a function of pattern differences in the plot. Similar calculations follow afterward to obtain  $H_B$  and subsequently  $B_{\text{patterndiff}}$ . Note that in the VSR task (Fig. 4A), we calculate two different scores for the primacy  $(A_{\text{primacy}})$  and recency  $(A_{\text{recency}})$  effect. For this, we divide the list at the point of minimum top-1 accuracy, into two lists, i.e. primacy and recency lists.
With these evaluation metrics, we report the quantitative measure of behavioral consistency with humans in Tab. A3. Ideally, if the model performs the same as humans, the slope difference would be zero. To test whether the results are significantly different from the chance model, we computed the p-value using two-tailed t-tests. From the results, we found that the majority of the models are more consistent with humans than the chance model; however, there still exists performance gaps between these models and humans.
R-96	R-256	R-1024	G-96	G-256	G-1024	L-96	L-256	L-1024	T-96	T-256	T-1024
VSR Task
Aprimacy	3.300*	3.264*	3.112*	8.108*	6.748*	6.859*	3.363*	6.458*	6.538*	3.307*	3.287*	3.480*
Recency	1.164*	1.183*	1.138*	1.449*	1.327*	1.276*	1.134*	1.384*	1.328*	1.170*	1.124*	1.102*
Blistrength	0.692*	0.697*	0.705*	0.467*	0.387*	0.400*	0.696*	0.348*	0.364*	0.693*	0.693*	0.693*
VIR Task
Retentioninterval	0.324*	0.729*	0.958*	0.801*	2.244*	0.729*	3.021*	0.951*	0.554*	0.995*	0.968*	1.438*
Retentioninterval	0.245*	0.272*	0.261*	0.164*	0.126*	0.068*	0.265*	0.057*	0.140*	0.446*	0.441*	0.262*
CS Task
Cognitiveload	0.995*	0.997*	0.996*	0.995*	0.996*	0.989*	0.999*	0.998*	0.991*	0.998*	0.999*	0.995*
Cognitiveload	0.754*	0.763*	0.718*	0.750*	0.741*	0.350*	0.708*	0.594*	0.340*	0.218*	0.511*	0.773*
SMU Task
Setsize	0.347*	1.059*	0.006*	1.001*	1.001*	1.000*	0.985*	0.993*	0.999*	1.049*	1.000*	1.255*
Setsize	0.680*	0.396*	0.310*	0.160*	0.160*	0.161*	0.159*	0.160*	0.161*	0.155*	0.161*	0.133*
VSRec Task
Patterndiff	1.103*	1.382*	0.979*	0.150*	0.569*	0.910*	0.531*	0.097*	0.625*	0.977*	0.923*	1.196*
Patterndiff	0.302*	0.318*	0.347*	0.030*	0.271*	0.311*	0.080*	0.220*	0.287*	0.356*	0.358*	0.323*
CD Task
Setsize	0.999*	0.999*	0.799*	1.003*	0.999*	0.999*	0.999*	0.999*	0.999*	0.999*	0.999*	1.003*
Setsize	0.110*	0.110*	0.154*	0.110*	0.110*	0.110*	0.110*	0.110*	0.110*	0.110*	0.110*	0.110*
Table A3: Quantitative scores for comparing humans and computational models on different behavioral effects. We conduct two-tailed t-tests against chance and compute corresponding p-values. Asterisk (*) indicates p-value  $< 0.01$ . See Sec. A2.3 for introduction to slope difference score  $A$  and average accuracy difference score  $B$ .
A3 Ablation Experiments
A3.1 Experiments for language modality
We introduce a new experiment in language modality, with a paradigm similar to that of Visual Serial Recall (VSR) Sec. 3. Here, we investigate the behavioral characteristics of the computational models in language modality instead of the visual modality.
Language Serial Recall (LSR). This experiment assesses the ability to recall words accurately as well as the order in which they were presented. Experiment schematic: The experiment consists of  $\mathrm{P}_{\text {present }}$  followed by  $\mathrm{P}_{\text {probe }}$ .  $\mathrm{P}_{\text {present }}$  involves the presentation of  $L$  words in a sequential manner. We randomly sample these words from the Google News dataset. This dataset contains pre-trained vectors for around 3 million words and phrases. At each time step during  $\mathrm{P}_{\text {present }}$ , we first map the word to the corresponding 300-dimensional word embedding followed by zero padding this vector to produce  $L \times 300$  dimensional representation which is then fed to the WM model. In  $\mathrm{P}_{\text {probe }}$ , at each time step, all the  $L$  words are presented simultaneously. We collect the embeddings for all  $L$  words and append them in a random order to generate  $L \times 300$  dimensional representation. At each  $\mathfrak{t}_{\text {probe }}$ , based on this representation, the agents perform a 9-way classification task to indicate the first word presented during  $\mathrm{P}_{\text {present }}$ , followed by the second word, and so on. Experiment conditions: We vary list length  $L = T_{\text {present }}$  from 2 - 9.
Result analysis: From the experimental results shown in Fig. A2B, we observed consistent performances as the ones from the VSR experiment. First, as list length increases, there is a corresponding increase in the memory load. Consequently, we observe a decrease in overall accuracy across various list lengths. We also observe pronounced primacy effects across varying list lengths, alongside a recency effect that becomes apparent in longer list lengths. These results suggest that the behavioral findings generalize across different modalities.
A3.2 Training with Longer List Lengths
To verify if our behavioral findings are valid for longer list lengths, we train our model on the LSR Task (Appendix A3.1) with trials of list lengths up to 16 instead of 9. Therefore, in this case, trials consist of 32 time steps in total instead of 20 time steps. We report our results in Fig. A2C. Here, we observe similar primacy and recency effects as we observed in the LSR task with shorter list lengths (compare with Fig. A2A). We notice similar trends where primacy effects are prominent in shorter list lengths and the recency effects start emerging in longer list lengths with higher memory load. This result suggests that the primacy-recency effect is a general phenomenon even for longer list lengths.
A3.3 Individual task training
To isolate the effect of the joint training paradigm, we also conduct individual training on the VSR task. Specifically, we train the LSTM-256 model only on the VSR task. We report the results in Fig. A2A. We make similar observations as in the joint training paradigm. First, with increasing list length, the overall accuracy of the model decreases. Next, the primacy-recency effect is more prominent in cases of longer list lengths with higher memory load. Finally, the primacy effect is more prominent in smaller list lengths while the recency effect is observed for longer list lengths. This indicates that the primacy-recency effect is a more general phenomenon agnostic of the training paradigm.
A3.4 Increased difficulty in CD task
In order to introduce additional difficulty in the CD task, we also experimented with more complex variants where we presented random noise images during the retention phase, instead of simple gray images. These random noise images are generated by randomly sampling a pixel value from 0-255 for every pixel of RGB channels of the image. This also aligns with the backward masking technique in neuroscience, which is often used for disrupting recurrent connections in the biological brains [49]. We train and test the LSTM-96 model on this new variant of the CD task. We observe that while LSTM-96 was successful in completely solving the initial version of the CD task with  $100\%$  accuracy, it completely failed to learn this modified variant.