Task-Optimized Artificial Neural Networks Align with Human Brain Activity in a Visual Working Memory Task

Brain encoding using Artificial Neural Networks (ANNs) has recently emerged as a powerful tool to uncover which features of complex stimuli drive brain activity in a fully data-driven manner. While previous research has primarily focused on perception (vision, audition), this study explores the application of recurrent neural networks (RNNs) to model visual working memory through a 2-back task. Participants viewed sequences of images and identified whether each image was seen two steps back in the sequence. We trained Long-Short Term Memory (LSTM) networks on this task using pixel-level representations of complex images from the ImageNet database and utilized the learned representations to predict functional brain activity recorded from subjects performing the 2-back task designed by the Human Connectome Project. Results demonstrate that ANNs can perform the 2-back task with high accuracy, and sparse attention mechanisms were beneficial for a fixed network capacity. The artificial representations of images effectively encoded individual brain activity during 2-back trials. High encoding accuracy was consistently observed in regions such as the dorsal and visual dorsal streams, and frontal cortices.