[Paper Review] A Biologically Inspired Visual Working Memory for Deep Networks

Ethan Harris, Mahesan Niranjan+1 more|arXiv (Cornell University)|Jan 9, 2019

Visual Attention and Saliency Detection35 references4 citations

TL;DR

This paper proposes a biologically inspired visual working memory module called STAWM (Short Term Attentive Working Memory), which uses a Hebb-Rosenblatt memory mechanism to store and update latent representations through transformational attention. The model achieves competitive classification accuracy on MNIST and CIFAR-10, enables unsupervised image reconstruction with disentangled parts-based representations, and provides an interpretable visual sketchpad that reveals model reasoning.

ABSTRACT

The ability to look multiple times through a series of pose-adjusted glimpses is fundamental to human vision. This critical faculty allows us to understand highly complex visual scenes. Short term memory plays an integral role in aggregating the information obtained from these glimpses and informing our interpretation of the scene. Computational models have attempted to address glimpsing and visual attention but have failed to incorporate the notion of memory. We introduce a novel, biologically inspired visual working memory architecture that we term the Hebb-Rosenblatt memory. We subsequently introduce a fully differentiable Short Term Attentive Working Memory model (STAWM) which uses transformational attention to learn a memory over each image it sees. The state of our Hebb-Rosenblatt memory is embedded in STAWM as the weights space of a layer. By projecting different queries through this layer we can obtain goal-oriented latent representations for tasks including classification and visual reconstruction. Our model obtains highly competitive classification performance on MNIST and CIFAR-10. As demonstrated through the CelebA dataset, to perform reconstruction the model learns to make a sequence of updates to a canvas which constitute a parts-based representation. Classification with the self supervised representation obtained from MNIST is shown to be in line with the state of the art models (none of which use a visual attention mechanism). Finally, we show that STAWM can be trained under the dual constraints of classification and reconstruction to provide an interpretable visual sketchpad which helps open the 'black-box' of deep learning.

Motivation & Objective

• To develop a biologically inspired visual working memory that integrates short-term memory with visual attention in deep networks.
• To enable deep networks to learn disentangled, parts-based representations through sequential glimpsing and memory updating.
• To create an interpretable visual sketchpad that reveals what the model 'sees' and why it makes certain predictions.
• To demonstrate that memory-enhanced attention improves performance beyond efficiency, enabling functions not achievable with attention alone.
• To explore dual-task training for classification and reconstruction under shared memory representation.

Proposed method

• The Hebb-Rosenblatt memory is implemented as the weight space of a neural network layer, where synaptic plasticity is modeled via Hebbian learning rules.
• Transformational attention is used to generate sequential glimpses, with each glimpse transformed via spatial transformations to model pose variation.
• The model uses a differentiable memory update mechanism that allows backpropagation through the memory state, enabling end-to-end training.
• A Bernoulli sampling method is applied to generate explicit attention masks at each glimpse, enabling unsupervised reconstruction and segmentation.
• The STAWM model jointly optimizes for classification and reconstruction by summing scaled losses, allowing shared memory representation across tasks.
• The canvas for image reconstruction is updated step-by-step, with each glimpse refining the latent representation toward a complete image.

Experimental results

Research questions

• RQ1 Can a biologically inspired visual working memory improve deep network performance beyond attention efficiency?
• RQ2 Can a memory mechanism enable unsupervised disentanglement of foreground and background in image reconstruction?
• RQ3 Can the memory representation be used to generate interpretable visual sketches that reflect model reasoning?
• RQ4 Does joint training on classification and reconstruction lead to better generalization and more robust representations?
• RQ5 How does the sequential updating of a latent canvas through glimpses reflect human visual processing?

Key findings

• The STAWM model achieves a test classification error of 1.0% on MNIST when jointly trained for classification and reconstruction, demonstrating strong performance.
• On CIFAR-10, the model achieves competitive classification accuracy, significantly outperforming previous attention-only models.
• In unsupervised image reconstruction on CelebA, the model learns to draw salient facial features first, producing a parts-based representation that separates face from background.
• The final attention mask from the last glimpse, when multiplied with the ground truth, reveals learned saliency, indicating unsupervised segmentation capability.
• For misclassified examples, the drawn canvas often resembles the predicted class more than the true target, indicating that the sketchpad reflects the model’s internal reasoning.
• The model’s ability to project different queries through the memory layer enables multiple tasks (classification, reconstruction) using a shared, interpretable latent space.